feat: Contextual dot resolution

src/ast.ts

@@ -1,9 +1,15 @@
 export type CompilerAST = CompilerAST[] | CompilerNode | string | null;
 export type CompilerSpan = { name: 'Pos'; args: [string, string, string] };
 
+/** Opaque expression object from the Motoko compiler - do not construct manually */
+export type RawExp = unknown & { readonly __brand: 'RawExp' };
+
+const rawExpSymbol = Symbol('rawExp');
+
 export interface CompilerNode {
     name: string;
     args: CompilerAST[];
+    rawExp?: RawExp;
 }
 
 export type Span = [number, number];
@@ -25,6 +31,19 @@ export interface Node extends Partial<Source> {
     args?: AST[];
 }
 
+/**
+ * Safely retrieves the raw expression from a node.
+ * 
+ * The raw expression is stored internally using a Symbol to prevent accidental access,
+ * which can break property descriptors on child nodes (e.g., when logging).
+ * 
+ * @param node The node to get the raw expression from
+ * @returns The raw expression, or undefined if not available
+ */
+export function getRawExp(node: Node): RawExp | undefined {
+    return (node as any)[rawExpSymbol];
+}
+
 export function asNode(ast: AST | undefined): Node | undefined {
     if (ast && typeof ast === 'object' && !Array.isArray(ast)) {
         return ast;
@@ -103,6 +122,8 @@ export function simplifyAST(ast: CompilerAST, parent?: Node | undefined): AST {
     const node: Node = {
         name: ast.name,
     };
+    // Store rawExp using a Symbol to prevent accidental access, e.g. in logs which destroys the type annotations in the AST...
+    (node as any)[rawExpSymbol] = ast.rawExp;
     Object.defineProperty(node, 'parent', {
         value: parent,
         enumerable: false,

src/index.ts

@@ -1,4 +1,4 @@
-import { CompilerNode, Node, simplifyAST } from './ast';
+import { CompilerNode, Node, RawExp, simplifyAST } from './ast';
 import { Scope, file } from './file';
 import {
     Package,
@@ -14,6 +14,9 @@ export type Motoko = ReturnType<typeof wrapMotoko>;
 
 type Compiler = any; // TODO: generate from `js_of_ocaml`?
 
+/** Opaque scope object from the Motoko compiler - do not construct manually */
+export type RawScope = unknown & { readonly __brand: 'RawScope' };
+
 export type Diagnostic = {
     source: string;
     range: {
@@ -63,8 +66,8 @@ export default function wrapMotoko(compiler: Compiler) {
             throw new Error(
                 result.diagnostics
                     ? result.diagnostics
-                          .map(({ message }: Diagnostic) => message)
-                          .join('; ')
+                        .map(({ message }: Diagnostic) => message)
+                        .join('; ')
                     : '(no diagnostics)',
             );
         }
@@ -98,6 +101,7 @@ export default function wrapMotoko(compiler: Compiler) {
         ast: Node;
         type: Node;
         immediateImports: string[];
+        scope: RawScope;
     };
     function parseMotokoTypedWithScopeCache(
         paths: string,
@@ -114,12 +118,12 @@ export default function wrapMotoko(compiler: Compiler) {
         scopeCache: Map<string, Scope>,
         enableRecovery?: boolean,
     ): [
-        (
-            | ParseMotokoTypedWithScopeCacheResult
-            | ParseMotokoTypedWithScopeCacheResult[]
-        ),
-        Map<string, Scope>,
-    ] {
+            (
+                | ParseMotokoTypedWithScopeCacheResult
+                | ParseMotokoTypedWithScopeCacheResult[]
+            ),
+            Map<string, Scope>,
+        ] {
         if (enableRecovery === undefined) {
             enableRecovery = false;
         }
@@ -142,15 +146,18 @@ export default function wrapMotoko(compiler: Compiler) {
                     ast,
                     typ,
                     immediateImports,
+                    scope,
                 }: {
                     ast: CompilerNode;
                     typ: CompilerNode;
                     immediateImports: string[];
+                    scope: RawScope;
                 }) => {
                     return {
                         ast: simplifyAST(ast),
                         type: simplifyAST(typ),
                         immediateImports,
+                        scope,
                     };
                 },
             ),
@@ -278,6 +285,17 @@ export default function wrapMotoko(compiler: Compiler) {
         },
         parseMotokoTyped,
         parseMotokoTypedWithScopeCache,
+        resolveDotCandidates(
+            scope: RawScope,
+            rawExp: RawExp,
+        ): {
+            name: string;
+            type: string;
+            moduleName?: string;
+        }[] {
+            // TODO: consider not exposing the RawScope nor RawExp outside of this library...
+            return invoke('resolveDotCandidates', false, [scope, rawExp]);
+        },
         resolveMain(directory: string = ''): string | undefined {
             return resolveMain(mo, directory);
         },