Improve import handling on autocompletion

package.json

@@ -2,7 +2,7 @@
     "name": "vscode-motoko",
     "displayName": "Motoko",
     "description": "Motoko language support",
-    "version": "0.6.3",
+    "version": "0.6.4",
     "publisher": "dfinity-foundation",
     "repository": "https://github.com/dfinity/vscode-motoko",
     "engines": {
@@ -133,7 +133,7 @@
         "mnemonist": "0.39.5",
         "motoko": "3.1.1",
         "prettier": "2.8.0",
-        "prettier-plugin-motoko": "0.2.4",
+        "prettier-plugin-motoko": "0.2.6",
         "url-relative": "1.0.0",
         "vscode-languageclient": "8.0.2",
         "vscode-languageserver": "8.0.2",

src/server/ast.ts

@@ -11,11 +11,16 @@ export interface AstStatus {
     outdated: boolean;
 }
 
+export interface AstImport {
+    path: string;
+    field?: string;
+}
+
 const globalCache = new Map<string, AstStatus>(); // Share non-typed ASTs across all contexts
 
 export default class AstResolver {
-    private _cache = globalCache;
-    private _typedCache = new Map<string, AstStatus>();
+    private readonly _cache = globalCache;
+    private readonly _typedCache = new Map<string, AstStatus>();
 
     clear() {
         this._cache.clear();
@@ -26,7 +31,7 @@ export default class AstResolver {
         const text = tryGetFileText(uri);
         if (!text) {
             this.delete(uri);
-            return false;
+            return true;
         }
         return this._updateWithFileText(uri, text, typed);
     }
@@ -51,20 +56,33 @@ export default class AstResolver {
         }
         try {
             const { motoko } = getContext(uri);
-            const ast = typed
-                ? motoko.parseMotokoTyped(resolveVirtualPath(uri)).ast
-                : motoko.parseMotoko(text);
+            const virtualPath = resolveVirtualPath(uri);
+            let ast: AST;
+            try {
+                ast = typed
+                    ? motoko.parseMotokoTyped(virtualPath).ast
+                    : motoko.parseMotoko(text);
+            } catch (err) {
+                throw new SyntaxError(String(err));
+            }
             status.ast = ast;
             const program = fromAST(ast);
             if (program instanceof Program) {
                 status.program = program;
+            } else {
+                console.log(`Unexpected AST node for URI: ${uri}`);
+                console.log(ast);
             }
             status.outdated = false;
             if (typed) {
                 console.log('Parsed typed AST');
             }
             return true;
         } catch (err) {
+            if (!(err instanceof SyntaxError)) {
+                console.error(`Error while parsing AST for ${uri}:`);
+                console.error(err);
+            }
             status.outdated = true;
             return false;
         }

src/server/context.ts

@@ -3,12 +3,23 @@ import * as baseLibrary from 'motoko/packages/latest/base.json';
 import ImportResolver from './imports';
 import AstResolver from './ast';
 
-export interface Context {
-    uri: string;
-    motoko: Motoko;
-    astResolver: AstResolver;
-    importResolver: ImportResolver;
-    error: string | undefined;
+/**
+ * A Motoko compiler context.
+ */
+export class Context {
+    public readonly uri: string;
+    public readonly motoko: Motoko;
+    public readonly astResolver: AstResolver;
+    public readonly importResolver: ImportResolver;
+
+    public error: string | undefined;
+
+    constructor(uri: string, motoko: Motoko) {
+        this.uri = uri;
+        this.motoko = motoko;
+        this.astResolver = new AstResolver();
+        this.importResolver = new ImportResolver(this);
+    }
 }
 
 const motokoPath = './motoko'; // Bundle generated by `esbuild`
@@ -52,19 +63,6 @@ function requestDefaultContext() {
 }
 requestDefaultContext(); // Always add a default context (provisional)
 
-/**
- * Create a new context with the given directory and compiler instance.
- */
-function createContext(uri: string, motoko: Motoko): Context {
-    return {
-        uri,
-        motoko,
-        astResolver: new AstResolver(),
-        importResolver: new ImportResolver(),
-        error: undefined,
-    };
-}
-
 /**
  * Reset all contexts (used to update Vessel configuration).
  */
@@ -89,7 +87,7 @@ export function addContext(uri: string): Context {
         return existing;
     }
     const motoko = requestMotokoInstance(uri);
-    const context = createContext(uri, motoko);
+    const context = new Context(uri, motoko);
     // Insert by descending specificity (`uri.length`) and then ascending alphabetical order
     let index = 0;
     while (index < contexts.length) {

src/server/dfx.ts

@@ -15,11 +15,11 @@ interface DfxConfig {
 type Cached<T> = T | undefined;
 
 export default class DfxResolver {
+    private readonly _findPath: () => string | null;
+
     private _path: Cached<string | null>;
     private _cache: Cached<DfxConfig | null>;
 
-    private _findPath: () => string | null;
-
     constructor(findPath: () => string | null) {
         this._findPath = findPath;
     }

src/server/imports.ts

@@ -3,6 +3,7 @@ import { getRelativeUri } from './utils';
 import { matchNode, Program } from './syntax';
 import { Node, AST } from 'motoko/lib/ast';
 import { pascalCase } from 'change-case';
+import { Context, getContext } from './context';
 
 interface ResolvedField {
     name: string;
@@ -11,24 +12,28 @@ interface ResolvedField {
 }
 
 export default class ImportResolver {
+    public readonly context: Context;
+
     // (module name -> uri)
-    private _moduleNameUriMap = new MultiMap<string, string>(Set);
+    private readonly _moduleNameUriMap = new MultiMap<string, string>(Set);
     // (uri -> resolved field)
-    private _fieldMap = new MultiMap<string, ResolvedField>(Set);
+    private readonly _fieldMap = new MultiMap<string, ResolvedField>(Set);
+
+    constructor(context: Context) {
+        this.context = context;
+    }
 
     clear() {
         this._moduleNameUriMap.clear();
     }
 
     update(uri: string, program: Program | undefined): boolean {
-        const motokoUri = getImportUri(uri);
-        if (!motokoUri) {
+        const info = getImportInfo(uri, this.context);
+        if (!info) {
             return false;
         }
-        const name = pascalCase(/([^/]+)$/i.exec(motokoUri)?.[1] || '');
-        if (name) {
-            this._moduleNameUriMap.set(name, motokoUri);
-        }
+        const [name, importUri] = info;
+        this._moduleNameUriMap.set(name, importUri);
         if (program?.export) {
             // Resolve field names
             const { ast } = program.export;
@@ -77,14 +82,15 @@ export default class ImportResolver {
     }
 
     delete(uri: string): boolean {
-        const motokoUri = getImportUri(uri);
-        if (!motokoUri) {
+        const info = getImportInfo(uri, this.context);
+        if (!info) {
             return false;
         }
+        const [, importUri] = info;
 
         let changed = false;
         for (const key of this._moduleNameUriMap.keys()) {
-            if (this._moduleNameUriMap.remove(key, motokoUri)) {
+            if (this._moduleNameUriMap.remove(key, importUri)) {
                 changed = true;
             }
         }
@@ -134,19 +140,46 @@ export default class ImportResolver {
     }
 }
 
-function getImportUri(uri: string): string | undefined {
+function getImportName(path: string): string {
+    return pascalCase(path);
+}
+
+function getImportInfo(
+    uri: string,
+    context: Context,
+): [string, string] | undefined {
     if (!uri.endsWith('.mo')) {
         return;
     }
     uri = uri.slice(0, -'.mo'.length);
-    const match = /\.vessel\/([^/]+)\/[^/]+\/src\/(.+)/.exec(uri);
-    if (match) {
-        // Resolve `mo:` URI for Vessel packages
-        const [, pkgName, path] = match;
-        uri = `mo:${pkgName}/${path}`;
-    } else if (/\.vessel\//.test(uri)) {
-        // Ignore everything else in `.vessel`
+    // Resolve package import paths
+    for (const regex of [
+        /\.vessel\/([^\/]+)\/[^\/]+\/src\/(.+)/,
+        /\.mops\/([^%\/]+)%40[^\/]+\/src\/(.+)/,
+        /\.mops\/_github\/([^%\/]+)%40[^\/]+\/src\/(.+)/,
+    ]) {
+        const match = regex.exec(uri);
+        if (match) {
+            if (getContext(uri) !== context) {
+                // Skip packages from other contexts
+                return;
+            }
+            const [, name, path] = match;
+            if (path === 'lib') {
+                // Account for `lib.mo` entry point
+                return [getImportName(name), `mo:${name}`];
+            } else {
+                // Resolve `mo:` URI for Vessel and MOPS packages
+                return [
+                    getImportName(/([^/]+)$/i.exec(uri)?.[1] || name),
+                    `mo:${name}/${path}`,
+                ];
+            }
+        }
+    }
+    if (uri.includes('/.vessel/') || uri.includes('/.mops/')) {
+        // Ignore everything else in Vessel and MOPS cache directories
         return;
     }
-    return uri;
+    return [getImportName(uri), uri];
 }