From 7367fff9cdef79d14a9f2b9167c402c8ea146e46 Mon Sep 17 00:00:00 2001
From: Ethan Reesor <ethan.reesor@gmail.com>
Date: Fri, 15 Nov 2024 19:32:57 -0700
Subject: [PATCH] Derive the correct path for module dependencies

Updates golang/vscode-go#3597
---
 package.json         |  8 ++++++--
 src/test/coverage.ts | 43 ++++++++++++++++++++++++++++++++-----------
 src/utils/util.ts    | 24 ++++++++++++++++++++++++
 yarn.lock            |  5 +++++
 4 files changed, 67 insertions(+), 13 deletions(-)

diff --git a/package.json b/package.json
index a943c3f..2db9f6d 100644
--- a/package.json
+++ b/package.json
@@ -139,7 +139,10 @@
           ]
         },
         "exp-vscode-go.testExplorer.codeLens": {
-          "type": ["string", "boolean"],
+          "type": [
+            "string",
+            "boolean"
+          ],
           "default": false,
           "scope": "resource",
           "markdownDescription": "Show code lenses for running and debugging tests",
@@ -327,6 +330,7 @@
     "vscode-languageserver-types": "^3.17.5"
   },
   "dependencies": {
+    "@streamparser/json": "^0.0.21",
     "axios": "^1.7.7",
     "chroma-js": "^3.1.1",
     "deep-equal": "^2.2.3",
@@ -335,4 +339,4 @@
     "node-html-parser": "^6.1.13",
     "tree-kill": "^1.2.2"
   }
-}
\ No newline at end of file
+}
diff --git a/src/test/coverage.ts b/src/test/coverage.ts
index 4522125..aef065b 100644
--- a/src/test/coverage.ts
+++ b/src/test/coverage.ts
@@ -4,6 +4,7 @@ import { promisify } from 'node:util';
 import { Location, Range, StatementCoverage, Uri } from 'vscode';
 import { Context } from './testing';
 import { Module, RootItem } from './item';
+import { parseJSONStream } from '../utils/util';
 
 /**
  * Parses a coverage file from `go test` into a map of
@@ -19,9 +20,19 @@ export async function parseCoverage(context: Context, scope: RootItem, coverageF
 		throw new Error('Failed to run "go env" as the "go" binary cannot be found in either GOROOT or PATH');
 	}
 
+	const modules: Record<string, string> = {};
+	parseJSONStream(await gets(binPath, scope, 'list', '-m', '-json', 'all'), (v) => {
+		const dep = v as {
+			Path: string;
+			Dir: string;
+		};
+		modules[dep.Path] = dep.Dir;
+	});
+
 	const env = {
-		GOROOT: await getEnv(binPath, 'GOROOT'),
-		GOMODCACHE: await getEnv(binPath, 'GOMODCACHE'),
+		modules,
+		GOROOT: await gets(binPath, scope, 'env', 'GOROOT'),
+		GOMODCACHE: await gets(binPath, scope, 'env', 'GOMODCACHE'),
 	};
 
 	const lines = Buffer.from(await context.workspace.fs.readFile(coverageFile))
@@ -48,14 +59,15 @@ export async function parseCoverage(context: Context, scope: RootItem, coverageF
 	return coverage;
 }
 
-async function getEnv(binPath: string, name: string) {
-	const { stdout } = await promisify(cp.execFile)(binPath, ['env', name]);
+async function gets(binPath: string, scope: RootItem, ...args: string[]) {
+	const { stdout } = await promisify(cp.execFile)(binPath, args, { cwd: scope.dir.fsPath });
 	return stdout.trim();
 }
 
 // Derived from https://golang.org/cl/179377
 
 interface Env {
+	modules: Record<string, string>;
 	GOROOT: string;
 	GOMODCACHE: string;
 }
@@ -126,13 +138,22 @@ function resolveCoveragePath(env: Env, scope: RootItem, filename: string) {
 		return path.join(scope.dir.fsPath, filename.substring(scope.path.length + 1));
 	}
 
-	// If the first segment of the path contains a dot, assume it's a module
-	const [first] = filename.split(/\\|\//);
-	if (first.includes('.')) {
-		// TODO: Resolve the version
-		return path.join(env.GOMODCACHE, filename);
+	// If the first segment does not contain a dot, assume it's a stdlib package
+	const parts = filename.split(/\\|\//);
+	if (!parts[0].includes('.')) {
+		return path.join(env.GOROOT, 'src', filename);
 	}
 
-	// If the first segment does not contain a dot, assume it's a stdlib package
-	return path.join(env.GOROOT, 'src', filename);
+	// If the first segment contains a dot, attempt to find a module that
+	// matches it
+	for (let i = parts.length - 1; i > 0; i--) {
+		const s = parts.slice(0, i).join('/');
+		if (s in env.modules) {
+			return path.join(env.modules[s], ...parts.slice(i));
+		}
+	}
+
+	// There's no matching module. This is guaranteed to fail but it's better
+	// than nothing as it will give the user some idea where to look.
+	return path.join(env.GOMODCACHE, filename);
 }
diff --git a/src/utils/util.ts b/src/utils/util.ts
index cb91041..beaa375 100644
--- a/src/utils/util.ts
+++ b/src/utils/util.ts
@@ -1,7 +1,9 @@
+/* eslint-disable @typescript-eslint/no-unused-vars */
 /* eslint-disable @typescript-eslint/no-explicit-any */
 import fs from 'node:fs';
 import os from 'node:os';
 import path from 'node:path';
+import { Tokenizer, TokenParser, ParsedElementInfo } from '@streamparser/json';
 
 // From vscode-go
 
@@ -96,3 +98,25 @@ export function cleanupTempDir() {
 	}
 	tmpDir = undefined;
 }
+
+type JsonValue = Exclude<ParsedElementInfo.ParsedElementInfo['value'], undefined>;
+
+export function parseJSONStream(s: string, onValue: (_: JsonValue) => void) {
+	const t = new Tokenizer();
+
+	let p = new TokenParser();
+	p.onValue = (x) => {
+		if (x.parent || !x.value) return;
+		onValue(x.value);
+	};
+	p.onEnd = () => {
+		const { onValue, onEnd } = p;
+		p = new TokenParser();
+		Object.assign(p, { onValue, onEnd });
+	};
+
+	t.onToken = (t) => {
+		p.write(t);
+	};
+	t.write(s);
+}
diff --git a/yarn.lock b/yarn.lock
index dea350e..ad72ebd 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -897,6 +897,11 @@
   dependencies:
     "@sinonjs/commons" "^3.0.0"
 
+"@streamparser/json@^0.0.21":
+  version "0.0.21"
+  resolved "https://registry.yarnpkg.com/@streamparser/json/-/json-0.0.21.tgz#8ce7b241d6fbe0a7f1907ca402f646baadbbfadf"
+  integrity sha512-v+49JBiG1kmc/9Ug79Lz9wyKaRocBgCnpRaLpdy7p0d3ICKtOAfc/H/Epa1j3F6YdnzjnZKKrnJ8xnh/v1P8Aw==
+
 "@tsconfig/node10@^1.0.7":
   version "1.0.11"
   resolved "https://registry.npmjs.org/@tsconfig/node10/-/node10-1.0.11.tgz"