fix: support syntax node as the first arg in local query

Author: huydo862003

packages/dbml-parse/__tests__/examples/compiler/syncDiagramView.test.ts

@@ -996,7 +996,7 @@ describe('syncDiagramView - entity name quoting', () => {
       },
     ]);
     const compiler = new Compiler();
-    compiler.setSource(newDbml);
+    compiler.setSource(DEFAULT_ENTRY, newDbml);
     expect(compiler.parse.errors()).toHaveLength(0);
   });
 

packages/dbml-parse/__tests__/snapshots/lexer/lexer.test.ts

@@ -36,7 +36,7 @@ describe('[snapshot] lexer', () => {
     const program = readFileSync(path.resolve(__dirname, `./input/${testName}.in.dbml`), 'utf-8');
 
     const compiler = new Compiler();
-    compiler.setSource(program);
+    compiler.setSource(DEFAULT_ENTRY, program);
 
     const lexer = new Lexer(program, DEFAULT_ENTRY);
 

packages/dbml-parse/__tests__/snapshots/nan/nan.test.ts

@@ -9,6 +9,7 @@ import {
   scanTestNames, toSnapshot,
 } from '@tests/utils';
 import Compiler from '@/compiler';
+import { DEFAULT_ENTRY } from '@/constants';
 import type { Database } from '@/core/types/schemaJson';
 import type Report from '@/core/types/report';
 
@@ -30,7 +31,7 @@ describe('[snapshot] nan', () => {
     const program = readFileSync(path.resolve(__dirname, `./input/${testName}.in.dbml`), 'utf-8');
 
     const compiler = new Compiler();
-    compiler.setSource(program);
+    compiler.setSource(DEFAULT_ENTRY, program);
     const report = compiler.parse._();
 
     it(testName, () => expect(serializeInterpreterResult(compiler, report)).toMatchFileSnapshot(path.resolve(__dirname, `./output/${testName}.out.json`)));

packages/dbml-parse/src/compiler/index.ts

@@ -19,7 +19,7 @@ import {
   type Internable, type Primitive, intern,
 } from '@/core/types/internable';
 import {
-  SyntaxNodeIdGenerator,
+  SyntaxNode, SyntaxNodeIdGenerator,
 } from '@/core/types/nodes';
 import {
   NodeSymbolIdGenerator, SymbolFactory,
@@ -94,14 +94,6 @@ const COMPUTING = Symbol('COMPUTING');
 type QuerySymbol = symbol;
 
 export default class Compiler {
-  // A global cache mapping from:
-  // (QuerySymbol, interned argument string) -> Result
-  private globalCache = new Map<QuerySymbol, Map<string, any>>();
-
-  // A local cache mapping from:
-  // (Filepath, QuerySymbol, interned argument string) -> Result
-  private localCache = new Map<FilepathId, Map<QuerySymbol, Map<string, any>>>();
-
   // Interners
   nodeIdGenerator = new SyntaxNodeIdGenerator();
 
@@ -132,30 +124,45 @@ export default class Compiler {
     this.globalCache.clear();
   }
 
+  // A global cache mapping from:
+  // (QuerySymbol, interned argument string) -> Result
+  private globalCache = new Map<QuerySymbol, Map<string, any>>();
+
+  // Turn a normal function into a Compiler's global query
+  // Input: A function that only accepts internable types | primitive types
+  // Output: A global query wrapping the function
   private globalQuery<Args extends (Primitive | Primitive[] | Internable<Primitive>)[], Return> (
     fn: (this: Compiler, ...args: Args) => Return,
   ): (...args: Args) => Return {
+    // 1. Generates a unique query symbol for the query
     const queryKey = Symbol();
+
     return ((...args: Args): Return => {
+      // Ensure there's a cache entry for this Query
+      if (!this.globalCache.has(queryKey)) {
+        this.globalCache.set(queryKey, new Map());
+      }
+
+      // 2. Intern all arguments and join into a single string
       const argKey = args.map((a) => intern(a)).join('\0');
-      let subCache = this.globalCache.get(queryKey);
-      if (subCache instanceof Map) {
-        if (subCache.has(argKey)) {
-          const cached = subCache.get(argKey);
-          if (cached === COMPUTING) {
-            throw new Error(`Cycle detected in query: ${fn.name}(${argKey})`);
-          }
-          return cached;
+
+      // 3. Check the cache using (QuerySymbol, interned argument string)
+      const subCache = this.globalCache.get(queryKey)!;
+      if (subCache.has(argKey)) {
+        const cached = subCache.get(argKey);
+        // 4. If the cache entry is being computed -> Throw query cycle error
+        if (cached === COMPUTING) {
+          throw new Error(`Cycle detected in query: ${fn.name}(${argKey})`);
         }
+        // 5. Cache hit, return
+        return cached;
       }
 
-      if (!(subCache instanceof Map)) {
-        subCache = new Map();
-        this.globalCache.set(queryKey, subCache);
-      }
+      // 6. Flag the current entry as computing
       subCache.set(argKey, COMPUTING);
 
       try {
+        // 7. Cache miss, recompute & cache
         const result = fn.apply(this, args);
         subCache.set(argKey, result);
         return result;
@@ -166,24 +173,45 @@ export default class Compiler {
     }) as (...args: Args) => Return;
   }
 
-  private localQuery<Args extends [Filepath, ...(Primitive | Primitive[] | Internable<Primitive>)[]], Return> (
+  // A local cache mapping from:
+  // (Filepath, QuerySymbol, interned argument string) -> Result
+  private localCache = new Map<FilepathId, Map<QuerySymbol, Map<string, any>>>();
+
+  // Turn a normal function into a Compiler's local query
+  // Input: A function whose first arg is a Filepath or SyntaxNode (from which we extract .filepath)
+  //        & accepts internable types | primitive types for the rest
+  // Output: A local query wrapping the function, cached per-file
+  private localQuery<Args extends [
+    Filepath | SyntaxNode, // filepath or node as first arg
+    ...(Primitive | Primitive[] | Internable<Primitive>)[],
+  ], Return> (
     fn: (this: Compiler, ...args: Args) => Return,
   ): (...args: Args) => Return {
+    // 1. Generates a unique query symbol for the query
     const queryKey = Symbol();
+
     return ((...args: Args): Return => {
+      // 2. Retrieve the filepath (from Filepath directly or from SyntaxNode.filepath)
       const [
-        filepath,
+        firstArg,
         ...rest
       ] = args;
+      const filepath = firstArg instanceof SyntaxNode ? firstArg.filepath : firstArg;
       const filepathId = filepath.intern();
-      const argKey = rest.map((a) => intern(a)).join('\0');
 
+      // 3. Intern all arguments (including first) and join into a single string
+      //    The first arg is interned too so different nodes from the same file
+      //    don't collide in the cache.
+      const argKey = args.map((a) => intern(a)).join('\0');
+
+      // Ensure that the local cache for the filepath exists
       let filepathCache = this.localCache.get(filepathId);
       if (!filepathCache) {
         filepathCache = new Map();
         this.localCache.set(filepathId, filepathCache);
       }
 
+      // Ensure that the local subcache for the (filepath, QuerySymbol) exists
       let subCache = filepathCache.get(queryKey);
       if (!subCache) {
         subCache = new Map();
@@ -192,14 +220,18 @@ export default class Compiler {
 
       if (subCache.has(argKey)) {
         const cached = subCache.get(argKey);
+        // 4. Query cycle, throw
         if (cached === COMPUTING) {
           throw new Error(`Cycle detected in query: ${fn.name}(${filepathId}, ${argKey})`);
         }
+        // 5. Cache hit, return
         return cached;
       }
 
+      // 6. Flag the current query as being computed
       subCache.set(argKey, COMPUTING);
       try {
+        // 7. Recompute the query result & cache
         const result = fn.apply(this, args);
         subCache.set(argKey, result);
         return result;
@@ -210,102 +242,105 @@ export default class Compiler {
     }) as (...args: Args) => Return;
   }
 
-  // A local query.
-  // Lex + parse a single file into an AST. Related: parseProject.
-  // (filepath: Filepath) => Report<FileParseIndex>
+  // A local query
+  // Lex + parse a single file into an AST
+  // Signature: (filepath: Filepath) => Report<FileParseIndex>
   parseFile = this.localQuery(parseFile);
-  // A global query.
-  // Lex + parse all entry-point files. Related: parseFile.
-  // () => Map<string, Report<FileParseIndex>>
+  // A global query
+  // Lex + parse all entry-point files
+  // Returns a map from the file's absolute path to FileParseIndex
+  // Signature: () => Map<string, Report<FileParseIndex>>
   parseProject = this.globalQuery(parseProject);
 
-  // A global query.
-  // Run the binder on a single AST node (dispatches to global modules). Related: bindFile, bindProject.
-  // (node: SyntaxNode) => Report<void> | Report<Unhandled>
+  // A local query
+  // Validate an AST node and return any compile errors
+  // Signature: (node: SyntaxNode) => Report<void> | Report<Unhandled>
+  validateNode = this.localQuery(validateNode);
+  // A local query
+  // Return the fully-qualified name segments of an AST node (e.g. ['public', 'users'])
+  // Signature: (node: SyntaxNode) => Report<string[] | undefined> | Report<Unhandled>
+  nodeFullname = this.localQuery(fullname);
+  // A local query
+  // Return the alias string of an AST node if one is declared
+  // Signature: (node: SyntaxNode) => Report<string | undefined> | Report<Unhandled>
+  nodeAlias = this.localQuery(nodeAlias);
+  // A local query
+  // Return the settings/options map of an AST node
+  // Signature: (node: SyntaxNode) => Report<Settings> | Report<Unhandled>
+  nodeSettings = this.localQuery(nodeSettings);
+
+  // A global query
+  // Run the binder on a single AST node
+  // Signature: (node: SyntaxNode) => Report<void> | Report<Unhandled>
   bindNode = this.globalQuery(bindNode);
-  // A global query.
-  // Bind a single file's AST (calls bindNode on the root). Related: bindNode, bindProject.
-  // (filepath: Filepath) => Report<void>
+  // A global query
+  // Bind a single file's AST (calls bindNode on the root)
+  // Signature: (filepath: Filepath) => Report<void>
   bindFile = this.globalQuery(bindFile);
-  // A global query.
-  // Bind all entry-point files in dependency order. Related: bindFile.
-  // () => Map<string, Report<void>>
+  // A global query
+  // Bind all entry-point files in dependency order
+  // Returns a map from the file's absolute path to bind errors
+  // Signature: () => Map<string, Report<void>>
   bindProject = this.globalQuery(bindProject);
 
-  // A global query.
-  // Interpret a single AST node into a SchemaElement (dispatches to global modules). Related: interpretFile, interpretProject.
-  // (node: SyntaxNode) => Report<SchemaElement | SchemaElement[] | undefined> | Report<Unhandled>
+  // A global query
+  // Interpret a single AST node into a SchemaElement
+  // Signature: (node: SyntaxNode) => Report<SchemaElement | SchemaElement[] | undefined> | Report<Unhandled>
   interpretNode = this.globalQuery(interpretNode);
-  // A global query.
-  // Interpret a single file's AST into a raw Database. Related: interpretNode, interpretProject.
-  // (filepath: Filepath) => Report<Readonly<Database> | undefined>
+  // A global query
+  // Interpret a single file's AST into a raw Database
+  // Signature: (filepath: Filepath) => Report<Readonly<Database> | undefined>
   interpretFile = this.globalQuery(interpretFile);
-  // A global query.
-  // Interpret all reachable files and merge into a MasterDatabase. Related: interpretFile.
-  // () => Report<MasterDatabase>
+  // A global query
+  // Interpret all reachable files and merge into a MasterDatabase (A database but for multifile)
+  // Signature: () => Report<MasterDatabase>
   interpretProject = this.globalQuery(interpretProject);
-  // A global query.
-  // Export a reconciled, alias-resolved Database for a single file (calls interpretProject internally). Related: interpretProject.
-  // (filepath: Filepath) => Report<Readonly<Database> | undefined>
+  // A global query
+  // Start from an entrypoint file, collects all imports and merges into the entrypoint
+  // Export a Database representing this merged view of the MasterDatabase
+  // Signature: (filepath: Filepath) => Report<Readonly<Database> | undefined>
   exportSchemaJson = this.globalQuery(exportSchemaJson);
 
-  // A global query.
-  // Resolve the NodeSymbol declared by a given AST node. Related: symbolMembers, nodeReferee.
-  // (node: SyntaxNode) => Report<NodeSymbol> | Report<Unhandled>
+  // A global query
+  // Return the NodeSymbol for a single SyntaxNode
+  // Signature: (node: SyntaxNode) => Report<NodeSymbol> | Report<Unhandled>
   nodeSymbol = this.globalQuery(nodeSymbol);
-  // A global query.
-  // Return all direct child symbols of a symbol. Related: nodeSymbol, lookupMembers.
-  // (symbol: NodeSymbol) => Report<NodeSymbol[]> | Report<Unhandled>
+  // A global query
+  // Return all direct child symbols of a symbol
+  // Signature: (symbol: NodeSymbol) => Report<NodeSymbol[]> | Report<Unhandled>
   symbolMembers = this.globalQuery(symbolMembers);
-  // A global query.
-  // Look up a member by kind and name inside a symbol or node scope. Related: symbolMembers.
-  // (symbolOrNode: NodeSymbol | SyntaxNode, targetKind: SymbolKind, targetName: string) => Report<NodeSymbol | undefined>
+  // A global query
+  // Look up a member by kind and name inside a symbol or node scope
+  // Signature: (symbolOrNode: NodeSymbol | SyntaxNode, targetKind: SymbolKind, targetName: string) => Report<NodeSymbol | undefined>
   lookupMembers = this.globalQuery(lookupMembers);
-  // A global query.
+  // A global query
   // Resolve what symbol a reference node points to. Related: nodeSymbol, symbolReferences.
-  // (node: SyntaxNode) => Report<NodeSymbol | undefined> | Report<Unhandled>
+  // Signature: (node: SyntaxNode) => Report<NodeSymbol | undefined> | Report<Unhandled>
   nodeReferee = this.globalQuery(nodeReferee);
-  // A global query.
+  // A global query
   // Find all reference nodes that point to a given symbol. Related: nodeReferee.
-  // (symbol: NodeSymbol) => Report<SyntaxNode[]>
+  // Signature: (symbol: NodeSymbol) => Report<SyntaxNode[]>
   symbolReferences = this.globalQuery(symbolReferences);
-  // A global query.
+  // A global query
   // Return all AliasSymbols across the project whose originalSymbol is the given symbol (transitive).
-  // (symbol: NodeSymbol) => Report<AliasSymbol[]>
+  // Signature: (symbol: NodeSymbol) => Report<AliasSymbol[]>
   symbolAliases = this.globalQuery(symbolAliases);
-  // A global query.
+  // A global query
   // Return all UseSymbols across the project whose originalSymbol is the given symbol (transitive).
-  // (symbol: NodeSymbol) => Report<UseSymbol[]>
+  // Signature: (symbol: NodeSymbol) => Report<UseSymbol[]>
   symbolUses = this.globalQuery(symbolUses);
 
-  // A global query.
-  // Validate an AST node and return any compile errors.
-  // (node: SyntaxNode) => Report<void> | Report<Unhandled>
-  validateNode = this.globalQuery(validateNode);
-  // A global query.
-  // Return the fully-qualified name segments of an AST node (e.g. ['public', 'users']).
-  // (node: SyntaxNode) => Report<string[] | undefined> | Report<Unhandled>
-  nodeFullname = this.globalQuery(fullname);
-  // A global query.
-  // Return the alias string of an AST node if one is declared. Related: nodeFullname.
-  // (node: SyntaxNode) => Report<string | undefined> | Report<Unhandled>
-  nodeAlias = this.globalQuery(nodeAlias);
-  // A global query.
-  // Return the settings/options map of an AST node.
-  // (node: SyntaxNode) => Report<Settings> | Report<Unhandled>
-  nodeSettings = this.globalQuery(nodeSettings);
-
-  // A local query.
+  // A local query
   // Return the direct import filepath IDs declared by use statements in a file. Related: reachableFiles.
-  // (filepath: Filepath) => Set<Filepath>
+  // Signature: (filepath: Filepath) => Filepath[]
   fileDependencies = this.localQuery(fileDependencies);
-  // A local query.
+  // A local query
   // BFS-traverse imports from an entry filepath and return all reachable files. Related: fileDependencies.
-  // (entry: Filepath) => Set<Filepath>
+  // Signature: (entry: Filepath) => Filepath[]
   reachableFiles = this.localQuery(reachableFiles);
-  // A global query.
+  // A global query
   // Return the importable members (non-schema, schema, reuses, uses) of a schema symbol or file. Related: topLevelSchemaMembers.
-  // (symbolOrFilepath: SchemaSymbol | Filepath) => Report<{ nonSchemaMembers, schemaMembers, reuses, uses }>
+  // Signature: (symbolOrFilepath: SchemaSymbol | Filepath) => Report<{ nonSchemaMembers, schemaMembers, reuses, uses }>
   usableMembers = this.globalQuery(usableMembers);
 
   // transform queries
@@ -343,6 +378,7 @@ export default class Compiler {
     return {
       definitionProvider: new DBMLDefinitionProvider(this),
       referenceProvider: new DBMLReferencesProvider(this),
+      // Trigger completion on these characters
       autocompletionProvider: new DBMLCompletionItemProvider(this, [
         '.',
         ',',