Fallback to WASM tree-sitter parser

package.json

@@ -92,6 +92,7 @@
     "pyodide": "0.28.2",
     "tree-sitter": "0.22.4",
     "tree-sitter-json": "0.24.8",
+    "web-tree-sitter": "0.22.4",
     "ts-essentials": "10.1.1",
     "vscode-languageserver": "9.0.1",
     "vscode-languageserver-textdocument": "1.0.12",

src/context/syntaxtree/SyntaxTree.ts

@@ -1,9 +1,8 @@
-import YamlGrammar from '@tree-sitter-grammars/tree-sitter-yaml';
-import Parser, { Edit, Point, SyntaxNode, Tree, Language } from 'tree-sitter';
-import JsonGrammar from 'tree-sitter-json';
+import { Edit, Point, SyntaxNode, Tree } from 'tree-sitter';
 import { Position } from 'vscode-languageserver-textdocument';
 import { DocumentType } from '../../document/Document';
 import { createEdit } from '../../document/DocumentUtils';
+import { parserFactory } from '../../parser/ParserFactory';
 import { Measure } from '../../telemetry/TelemetryDecorator';
 import { TopLevelSection, TopLevelSections, IntrinsicsSet } from '../CloudFormationEnums';
 import { normalizeIntrinsicFunction } from '../semantic/Intrinsics';
@@ -15,20 +14,13 @@ import { NodeType } from './utils/NodeType';
 import { createSyntheticNode } from './utils/SyntheticEntityFactory';
 import { CommonNodeTypes, JsonNodeTypes, YamlNodeTypes } from './utils/TreeSitterTypes';
 
-// Optimization to only load the different language grammars once
-// Loading native/wasm code is expensive
-const JSON_PARSER = new Parser();
-JSON_PARSER.setLanguage(JsonGrammar as Language);
-
-const YAML_PARSER = new Parser();
-YAML_PARSER.setLanguage(YamlGrammar as Language);
-
 export type PropertyPath = ReadonlyArray<string | number>;
 export type PathAndEntity = {
     path: ReadonlyArray<SyntaxNode>; // All nodes from target to root
     propertyPath: PropertyPath; // Path like ["Resources", "MyBucket", "Properties"]
     entityRootNode?: SyntaxNode; // The complete entity definition (e.g., entire resource)
 };
+
 const LARGE_NODE_TEXT_LIMIT = 200; // If a node's text is > 200 chars, we are likely not at the most specific node (indicating that it might be invalid)
 
 export abstract class SyntaxTree {
@@ -42,9 +34,9 @@ export abstract class SyntaxTree {
         content: string,
     ) {
         if (type === DocumentType.YAML) {
-            this.parser = YAML_PARSER;
+            this.parser = parserFactory.createYamlParser();
         } else {
-            this.parser = JSON_PARSER;
+            this.parser = parserFactory.createJsonParser();
         }
         this.rawContent = content;
         this.tree = this.parser.parse(this.rawContent);

src/parser/GrammarManager.ts

@@ -0,0 +1,136 @@
+import { join } from 'path';
+import Parser from 'web-tree-sitter';
+import { DocumentType } from '../document/Document';
+import { readBufferIfExists } from '../utils/File';
+
+export interface GrammarConfig {
+    yamlGrammarPath?: string;
+    jsonGrammarPath?: string;
+    maxRetries?: number;
+    retryDelay?: number;
+    wasmBasePath?: string;
+}
+
+export class GrammarManager {
+    private static instance: GrammarManager;
+    private initialized = false;
+    private readonly grammarCache = new Map<DocumentType, Parser.Language>();
+    private readonly loadingPromises = new Map<DocumentType, Promise<Parser.Language>>();
+    private readonly config: Required<GrammarConfig>;
+
+    private constructor(config: GrammarConfig = {}) {
+        const basePath = config.wasmBasePath ?? this.getDefaultWasmPath();
+
+        this.config = {
+            yamlGrammarPath: config.yamlGrammarPath ?? join(basePath, 'tree-sitter-yaml.wasm'),
+            jsonGrammarPath: config.jsonGrammarPath ?? join(basePath, 'tree-sitter-json.wasm'),
+            maxRetries: config.maxRetries ?? 3,
+            retryDelay: config.retryDelay ?? 100,
+            wasmBasePath: basePath,
+        };
+    }
+
+    private getDefaultWasmPath(): string {
+        // In bundled environment, WASM files are in the same directory as the bundle
+        if (typeof __dirname !== 'undefined') {
+            // __dirname points to the bundle directory, WASM files are in ./wasm/
+            return join(__dirname, 'wasm');
+        }
+        // Fallback for different environments
+        return './wasm';
+    }
+
+    public static getInstance(config?: GrammarConfig): GrammarManager {
+        if (!GrammarManager.instance) {
+            GrammarManager.instance = new GrammarManager(config);
+        }
+        return GrammarManager.instance;
+    }
+
+    private async ensureInitialized(): Promise<void> {
+        if (this.initialized) return;
+
+        await Parser.init({
+            locateFile: (scriptName: string) => {
+                if (scriptName === 'tree-sitter.wasm') {
+                    return join(this.config.wasmBasePath, '..', 'tree-sitter.wasm');
+                }
+                return scriptName;
+            },
+        });
+
+        this.initialized = true;
+    }
+
+    private async loadGrammarWithRetry(type: DocumentType): Promise<Parser.Language> {
+        const grammarPath = type === DocumentType.YAML ? this.config.yamlGrammarPath : this.config.jsonGrammarPath;
+
+        let lastError: Error | undefined;
+
+        for (let attempt = 1; attempt <= this.config.maxRetries; attempt++) {
+            try {
+                const wasmBuffer = readBufferIfExists(grammarPath);
+                return await Parser.Language.load(wasmBuffer);
+            } catch (error) {
+                lastError = error as Error;
+
+                if (attempt < this.config.maxRetries) {
+                    await new Promise((resolve) => setTimeout(resolve, this.config.retryDelay * attempt));
+                }
+            }
+        }
+
+        throw new Error(
+            `Failed to load ${type} grammar after ${this.config.maxRetries} attempts: ${lastError?.message}`,
+        );
+    }
+
+    public async loadGrammar(type: DocumentType): Promise<Parser.Language> {
+        // Return cached grammar if available
+        const cached = this.grammarCache.get(type);
+        if (cached) {
+            return cached;
+        }
+
+        // Return existing loading promise if in progress
+        const existingPromise = this.loadingPromises.get(type);
+        if (existingPromise) {
+            return await existingPromise;
+        }
+
+        // Start new loading process
+        const loadingPromise = this.loadGrammarInternal(type);
+        this.loadingPromises.set(type, loadingPromise);
+
+        try {
+            const grammar = await loadingPromise;
+            this.grammarCache.set(type, grammar);
+            return grammar;
+        } finally {
+            this.loadingPromises.delete(type);
+        }
+    }
+
+    private async loadGrammarInternal(type: DocumentType): Promise<Parser.Language> {
+        await this.ensureInitialized();
+        return await this.loadGrammarWithRetry(type);
+    }
+
+    public async preloadGrammars(types: DocumentType[] = [DocumentType.YAML, DocumentType.JSON]): Promise<void> {
+        const promises = types.map((type) => this.loadGrammar(type));
+        await Promise.all(promises);
+    }
+
+    public isGrammarLoaded(type: DocumentType): boolean {
+        return this.grammarCache.has(type);
+    }
+
+    public clearCache(): void {
+        this.grammarCache.clear();
+        this.loadingPromises.clear();
+    }
+
+    public getGrammarPath(type: DocumentType): string {
+        return type === DocumentType.YAML ? this.config.yamlGrammarPath : this.config.jsonGrammarPath;
+    }
+}

src/parser/ParserFactory.ts

@@ -0,0 +1,83 @@
+import TreeSitterYaml from '@tree-sitter-grammars/tree-sitter-yaml';
+import Parser from 'tree-sitter';
+import TreeSitterJson from 'tree-sitter-json';
+import { LoggerFactory } from '../telemetry/LoggerFactory';
+import { WasmParserFactory } from './WasmParserFactory';
+
+const log = LoggerFactory.getLogger('ParserFactory');
+
+export interface ParserFactory {
+    createYamlParser(): Parser;
+    createJsonParser(): Parser;
+    initialize?(): Promise<void>;
+}
+
+class NativeParserFactory implements ParserFactory {
+    private readonly yamlParser: Parser;
+    private readonly jsonParser: Parser;
+    private wasmFallback?: WasmParserFactory;
+    private readonly nativeFailed: boolean = false;
+
+    constructor() {
+        try {
+            this.yamlParser = new Parser();
+            this.yamlParser.setLanguage(TreeSitterYaml as unknown as Parser.Language);
+
+            this.jsonParser = new Parser();
+            this.jsonParser.setLanguage(TreeSitterJson as unknown as Parser.Language);
+
+            log.info('Native tree-sitter parsers initialized successfully');
+        } catch {
+            log.error('Native tree-sitter initialization failed, will use WASM fallback');
+            this.nativeFailed = true;
+            this.yamlParser = new Parser();
+            this.jsonParser = new Parser();
+            this.initializeWasmFallback();
+        }
+    }
+
+    private initializeWasmFallback(): void {
+        log.info('Initializing WASM fallback...');
+        this.wasmFallback = new WasmParserFactory();
+        this.wasmFallback.initialize().catch((error: unknown) => {
+            log.error(error, 'WASM fallback initialization failed');
+        });
+    }
+
+    createYamlParser(): Parser {
+        if (this.nativeFailed && this.wasmFallback) {
+            return this.wasmFallback.createYamlParser();
+        }
+        return this.yamlParser;
+    }
+
+    createJsonParser(): Parser {
+        if (this.nativeFailed && this.wasmFallback) {
+            return this.wasmFallback.createJsonParser();
+        }
+        return this.jsonParser;
+    }
+}
+
+// Environment detection and factory creation
+const shouldForceWasm = (): boolean => {
+    return process.env.CLOUDFORMATIONLSP_USE_WASM === 'true';
+};
+
+// Initialize the factory - async initialization happens in background
+let factoryInstance: ParserFactory;
+
+if (shouldForceWasm()) {
+    log.info('Forcing WASM tree-sitter implementation (CLOUDFORMATIONLSP_USE_WASM=true)');
+    const wasmFactory = new WasmParserFactory();
+    // eslint-disable-next-line unicorn/prefer-top-level-await
+    wasmFactory.initialize().catch((error: unknown) => {
+        log.error(error, 'Failed to initialize WASM parser factory');
+    });
+    factoryInstance = wasmFactory;
+} else {
+    log.info('Using native tree-sitter implementation with WASM fallback');
+    factoryInstance = new NativeParserFactory();
+}
+
+export const parserFactory: ParserFactory = factoryInstance;