1import {
2  DefinitionProvider,
3  TextDocument,
4  CancellationToken,
5  Definition,
6  Location,
7  Position,
8  ProviderResult,
9  workspace,
10  Uri,
11} from "vscode";
12
13const NAMED_SNAPSHOT_ASSERTION: RegExp = /(?:\binsta::)?(?:assert(?:_\w+)?_snapshot!)\(\s*['"]([^'"]+)['"]\s*,/;
14const STRING_INLINE_SNAPSHOT_ASSERTION: RegExp = /(?:\binsta::)?(?:assert(?:_\w+)?_snapshot!)\(\s*['"]([^'"]+)['"]\s*,\s*@(r#*)?["']/;
15const UNNAMED_SNAPSHOT_ASSERTION: RegExp = /(?:\binsta::)?(?:assert(?:_\w+)?_snapshot!)\(/;
16const INLINE_MARKER: RegExp = /@(r#*)?["']/;
17const FUNCTION: RegExp = /\bfn\s+([\w]+)\s*\(/;
18const TEST_DECL: RegExp = /#\[test\]/;
19const FILENAME_PARTITION: RegExp = /^(.*)[/\\](.*?)\.rs$/;
20const SNAPSHOT_FUNCTION_STRIP: RegExp = /^test_(.*?)$/;
21const SNAPSHOT_HEADER: RegExp = /^---\s*$(.*?)^---\s*$/ms;
22
23type SnapshotMatch = {
24  snapshotName: string | null;
25  line: number | null;
26  path: string;
27  localModuleName: string | null;
28  snapshotType: "inline" | "named";
29};
30
31type ResolvedSnapshotMatch = SnapshotMatch & {
32  snapshotUri: Uri;
33};
34
35export class SnapshotPathProvider implements DefinitionProvider {
36  /**
37   * This looks up an explicitly named snapshot (simple case)
38   */
39  private resolveNamedSnapshot(
40    document: TextDocument,
41    position: Position
42  ): SnapshotMatch | null {
43    const line =
44      (position.line >= 1 ? document.lineAt(position.line - 1).text : "") +
45      document.lineAt(position.line).text;
46
47    const snapshotMatch = line.match(NAMED_SNAPSHOT_ASSERTION);
48    if (!snapshotMatch) {
49      return null;
50    }
51    const snapshotName = snapshotMatch[1];
52    const fileNameMatch = document.fileName.match(FILENAME_PARTITION);
53    if (!fileNameMatch) {
54      return null;
55    }
56    const path = fileNameMatch[1];
57    const localModuleName = fileNameMatch[2];
58    return {
59      snapshotName,
60      line: null,
61      path,
62      localModuleName,
63      snapshotType: "named",
64    };
65  }
66
67  /**
68   * This locates an implicitly (unnamed) snapshot.
69   */
70  private resolveUnnamedSnapshot(
71    document: TextDocument,
72    position: Position,
73    noInline: boolean
74  ): SnapshotMatch | null {
75    function unnamedSnapshotAt(lineno: number): boolean {
76      const line = document.lineAt(lineno).text;
77      return !!(
78        line.match(UNNAMED_SNAPSHOT_ASSERTION) &&
79        !line.match(NAMED_SNAPSHOT_ASSERTION) &&
80        (noInline || !line.match(STRING_INLINE_SNAPSHOT_ASSERTION))
81      );
82    }
83
84    // if we can't find an unnnamed snapshot at the given position we bail.
85    if (!unnamedSnapshotAt(position.line)) {
86      return null;
87    }
88
89    // otherwise scan backwards for unnamed snapshot matches until we find
90    // a test function declaration.
91    let snapshotNumber = 1;
92    let scanLine = position.line - 1;
93    let functionName = null;
94    let isInline = !!document.lineAt(position.line).text.match(INLINE_MARKER);
95    console.log("inline", document.lineAt(position.line), isInline);
96
97    while (scanLine >= 0) {
98      // stop if we find a test function declaration
99      let functionMatch;
100      const line = document.lineAt(scanLine);
101      if (
102        scanLine > 1 &&
103        (functionMatch = line.text.match(FUNCTION)) &&
104        document.lineAt(scanLine - 1).text.match(TEST_DECL)
105      ) {
106        functionName = functionMatch[1];
107        break;
108      }
109      if (!isInline && line.text.match(INLINE_MARKER)) {
110        isInline = true;
111      }
112      if (unnamedSnapshotAt(scanLine)) {
113        // TODO: do not increment if the snapshot at that location
114        snapshotNumber++;
115      }
116      scanLine--;
117    }
118
119    // if we couldn't find a function or an unexpected inline snapshot we have to bail.
120    if (!functionName || (noInline && isInline)) {
121      return null;
122    }
123
124    let snapshotName = null;
125    let line = null;
126    let path = null;
127    let localModuleName = null;
128
129    if (isInline) {
130      line = position.line;
131      path = document.fileName;
132    } else {
133      snapshotName = `${functionName.match(SNAPSHOT_FUNCTION_STRIP)![1]}${
134        snapshotNumber > 1 ? `-${snapshotNumber}` : ""
135      }`;
136      const fileNameMatch = document.fileName.match(FILENAME_PARTITION);
137      if (!fileNameMatch) {
138        return null;
139      }
140      path = fileNameMatch[1];
141      localModuleName = fileNameMatch[2];
142    }
143
144    return {
145      snapshotName,
146      line,
147      path,
148      localModuleName,
149      snapshotType: isInline ? "inline" : "named",
150    };
151  }
152
153  public findSnapshotAtLocation(
154    document: TextDocument,
155    position: Position,
156    token: CancellationToken,
157    noInline: boolean = false
158  ): Thenable<ResolvedSnapshotMatch | null> {
159    const snapshotMatch =
160      this.resolveNamedSnapshot(document, position) ||
161      this.resolveUnnamedSnapshot(document, position, noInline);
162    if (!snapshotMatch) {
163      return Promise.resolve(null);
164    }
165
166    if (snapshotMatch.snapshotType === "inline") {
167      return Promise.resolve({
168        snapshotUri: document.uri,
169        ...snapshotMatch,
170      });
171    }
172
173    const getSearchPath = function (
174      mode: "exact" | "wildcard-prefix" | "wildcard-all"
175    ): string {
176      return workspace.asRelativePath(
177        `${snapshotMatch.path}/snapshots/${mode !== "exact" ? "*__" : ""}${
178          snapshotMatch.localModuleName
179        }${mode === "wildcard-all" ? "__*" : ""}__${
180          snapshotMatch.snapshotName
181        }.snap`
182      );
183    };
184
185    function findFiles(path: string): Thenable<Uri | null> {
186      return workspace
187        .findFiles(path, "", 1, token)
188        .then((results) => results[0] || null);
189    }
190
191    // we try to find the file in three passes:
192    // - exact matchin the snapshot folder.
193    // - with a wildcard module prefix (crate__foo__NAME__SNAP)
194    // - with a wildcard module prefix and suffix (crate__foo__NAME__tests__SNAP)
195    // This is needed since snapshots can be contained in submodules. Since
196    // getting the actual module name is tedious we just hope the match is
197    // unique.
198    return findFiles(getSearchPath("exact"))
199      .then((rv) => rv || findFiles(getSearchPath("wildcard-prefix")))
200      .then((rv) => rv || findFiles(getSearchPath("wildcard-all")))
201      .then((snapshot) =>
202        snapshot ? { snapshotUri: snapshot, ...snapshotMatch } : null
203      );
204  }
205
206  public provideDefinition(
207    document: TextDocument,
208    position: Position,
209    token: CancellationToken
210  ): ProviderResult<Definition> {
211    return this.findSnapshotAtLocation(document, position, token, true).then(
212      (match) => {
213        if (!match) {
214          return null;
215        }
216        return workspace.fs.readFile(match.snapshotUri).then((contents) => {
217          const stringContents = Buffer.from(contents).toString("utf-8");
218          const header = stringContents.match(SNAPSHOT_HEADER);
219          let location = new Position(0, 0);
220          if (header) {
221            location = new Position(header[0].match(/\n/g)!.length + 1, 0);
222          }
223          return new Location(match.snapshotUri, location);
224        });
225      }
226    );
227  }
228}
229