Fallback to WASM tree-sitter parser

.github/workflows/build-and-test.yml

@@ -66,6 +66,9 @@ jobs:
       - name: Test
         run: npm run test
 
+      - name: Test (WASM parser)
+        run: npm run test:wasm
+
   build-test-go:
     needs: [ get-configs ]
     runs-on: ${{ inputs.runs-on }}

assets/featureFlag/alpha.json

@@ -48,6 +48,10 @@
         "FileDb": {
             "enabled": true,
             "fleetPercentage": 100
+        },
+        "WasmParser": {
+            "enabled": false,
+            "fleetPercentage": 0
         }
     }
 }

assets/featureFlag/beta.json

@@ -47,6 +47,10 @@
         "FileDb": {
             "enabled": false,
             "fleetPercentage": 100
+        },
+        "WasmParser": {
+            "enabled": false,
+            "fleetPercentage": 0
         }
     }
 }

assets/featureFlag/prod.json

@@ -47,6 +47,10 @@
         "FileDb": {
             "enabled": false,
             "fleetPercentage": 0
+        },
+        "WasmParser": {
+            "enabled": false,
+            "fleetPercentage": 0
         }
     }
 }

package.json

@@ -25,6 +25,7 @@
         "watch": "rm -rf out && tsc -b -w",
         "test": "cross-env NODE_ENV=test vitest run",
         "test:integration": "cross-env NODE_ENV=test vitest run --config vitest.integration.config.ts",
+        "test:wasm": "cross-env NODE_ENV=test BUILD_TARGET=legacy vitest run --config vitest.integration.config.ts",
         "test:unit": "cross-env NODE_ENV=test vitest run --config vitest.unit.config.ts",
         "test:leaks": "cross-env NODE_ENV=test vitest run --pool=forks --logHeapUsage",
         "lint": "eslint --cache --cache-location node_modules/.cache/eslint --max-warnings 0 .",
@@ -83,6 +84,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.2.0",
         "vscode-languageserver": "9.0.1",
         "vscode-languageserver-textdocument": "1.0.12",

src/app/standalone.ts

@@ -14,9 +14,15 @@ async function onInitialize(params: ExtendedInitializeParams) {
     staticInitialize(params.clientInfo, params.initializationOptions?.['aws']);
 
     // Dynamically load these modules so that OTEL can instrument all the libraries first
+    const { syntaxTreeFactory } = await import('../context/syntaxtree/SyntaxTreeFactory');
+    await syntaxTreeFactory.ready;
+
     const { CfnInfraCore } = await import('../server/CfnInfraCore');
     const core = new CfnInfraCore(lsp.components, params);
 
+    // CfnInfraCore.initialize may switch to WASM based on feature flag
+    await syntaxTreeFactory.ready;
+
     const { CfnServer } = await import('../server/CfnServer');
     server = new CfnServer(lsp.components, core);
     return LspCapabilities;

src/context/syntaxtree/JsonSyntaxTree.ts

@@ -1,8 +1,10 @@
 import { DocumentType } from '../../document/Document';
+import { ParserFactory } from '../../parser/ParserFactory';
 import { SyntaxTree } from './SyntaxTree';
+import { ParserType } from './SyntaxTreeFactory';
 
 export class JsonSyntaxTree extends SyntaxTree {
-    constructor(content: string) {
-        super(DocumentType.JSON, content);
+    constructor(content: string, factory: ParserFactory, parserType: ParserType) {
+        super(DocumentType.JSON, content, factory, parserType);
     }
 }

src/context/syntaxtree/SyntaxTree.ts

@@ -1,12 +1,12 @@
-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';
+import { ParserType } from './SyntaxTreeFactory';
 import { extractEntityFromNodeTextYaml } from './utils/NodeParse';
 import { NodeSearch } from './utils/NodeSearch';
 import { NodeStructure } from './utils/NodeStructure';
@@ -15,36 +15,33 @@ 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 {
     protected tree: Tree;
     private readonly parser;
     private rawContent: string;
     private _lines: string[] | undefined;
+    private readonly parserType: ParserType;
 
     protected constructor(
         public readonly type: DocumentType,
         content: string,
+        factory: ParserFactory,
+        parserType: ParserType,
     ) {
+        this.parserType = parserType;
         if (type === DocumentType.YAML) {
-            this.parser = YAML_PARSER;
+            this.parser = factory.createYamlParser();
         } else {
-            this.parser = JSON_PARSER;
+            this.parser = factory.createJsonParser();
         }
         this.rawContent = content;
         this.tree = this.parser.parse(this.rawContent);
@@ -637,7 +634,7 @@ export abstract class SyntaxTree {
                     }
 
                     // Stop if we've reached the root node (avoid infinite loop)
-                    if (current === this.tree.rootNode) {
+                    if (current.id === this.tree.rootNode.id) {
                         break;
                     }
                     current = current.parent;

src/context/syntaxtree/SyntaxTreeFactory.ts

@@ -0,0 +1,70 @@
+import { DocumentType } from '../../document/Document';
+import { FeatureFlag } from '../../featureFlag/FeatureFlagI';
+import { ParserFactory, parserFactory, parserFactoryReady } from '../../parser/ParserFactory';
+import { WasmParserFactory } from '../../parser/WasmParserFactory';
+import { LoggerFactory } from '../../telemetry/LoggerFactory';
+import { ScopedTelemetry } from '../../telemetry/ScopedTelemetry';
+import { Telemetry } from '../../telemetry/TelemetryDecorator';
+import { JsonSyntaxTree } from './JsonSyntaxTree';
+import { SyntaxTree } from './SyntaxTree';
+import { YamlSyntaxTree } from './YamlSyntaxTree';
+
+export type ParserType = 'native' | 'wasm';
+
+const log = LoggerFactory.getLogger('SyntaxTreeFactory');
+const isLegacyLinux = process.env.BUILD_TARGET === 'legacy';
+
+export class SyntaxTreeFactory {
+    @Telemetry() private readonly telemetry!: ScopedTelemetry;
+
+    private factory: ParserFactory;
+    private type: ParserType;
+    private readyPromise: Promise<void>;
+
+    constructor(nativeFactory: ParserFactory = parserFactory) {
+        this.factory = nativeFactory;
+        this.type = isLegacyLinux ? 'wasm' : 'native';
+        this.readyPromise = parserFactoryReady;
+    }
+
+    /**
+     * Called once during server initialization to lock in the parser type
+     * for the lifetime of the session based on the feature flag state.
+     */
+    initialize(wasmFlag: FeatureFlag): void {
+        if (isLegacyLinux) {
+            return; // Already using WASM via parserFactory
+        }
+        if (wasmFlag.isEnabled()) {
+            log.info('WasmParser feature flag enabled, switching to WASM parser');
+            const wasm = new WasmParserFactory();
+            this.factory = wasm;
+            this.type = 'wasm';
+            this.readyPromise = wasm.initialize().catch((error: unknown) => {
+                log.error(error, 'WASM initialization failed, falling back to native');
+                this.factory = parserFactory;
+                this.type = 'native';
+            });
+        }
+    }
+
+    get parserType(): ParserType {
+        return this.type;
+    }
+
+    get ready(): Promise<void> {
+        return this.readyPromise;
+    }
+
+    createSyntaxTree(content: string, documentType: DocumentType): SyntaxTree {
+        this.telemetry.count('createSyntaxTree', 1, {
+            attributes: { 'parser.type': this.type },
+        });
+        if (documentType === DocumentType.JSON) {
+            return new JsonSyntaxTree(content, this.factory, this.type);
+        }
+        return new YamlSyntaxTree(content, this.factory, this.type);
+    }
+}
+
+export const syntaxTreeFactory = new SyntaxTreeFactory();

src/context/syntaxtree/SyntaxTreeManager.ts

@@ -3,9 +3,8 @@ import { CloudFormationFileType, DocumentType } from '../../document/Document';
 import { detectDocumentType } from '../../document/DocumentUtils';
 import { LoggerFactory } from '../../telemetry/LoggerFactory';
 import { Measure } from '../../telemetry/TelemetryDecorator';
-import { JsonSyntaxTree } from './JsonSyntaxTree';
 import { SyntaxTree } from './SyntaxTree';
-import { YamlSyntaxTree } from './YamlSyntaxTree';
+import { syntaxTreeFactory } from './SyntaxTreeFactory';
 
 const logger = LoggerFactory.getLogger('SyntaxTreeManager');
 
@@ -50,11 +49,11 @@ export class SyntaxTreeManager {
     }
 
     private createJsonSyntaxTree(uri: string, content: string) {
-        this.syntaxTrees.set(uri, new JsonSyntaxTree(content));
+        this.syntaxTrees.set(uri, syntaxTreeFactory.createSyntaxTree(content, DocumentType.JSON));
     }
 
     private createYamlSyntaxTree(uri: string, content: string) {
-        this.syntaxTrees.set(uri, new YamlSyntaxTree(content));
+        this.syntaxTrees.set(uri, syntaxTreeFactory.createSyntaxTree(content, DocumentType.YAML));
     }
 
     public getSyntaxTree(uri: string): SyntaxTree | undefined {

src/context/syntaxtree/YamlSyntaxTree.ts

@@ -1,8 +1,10 @@
 import { DocumentType } from '../../document/Document';
+import { ParserFactory } from '../../parser/ParserFactory';
 import { SyntaxTree } from './SyntaxTree';
+import { ParserType } from './SyntaxTreeFactory';
 
 export class YamlSyntaxTree extends SyntaxTree {
-    constructor(content: string) {
-        super(DocumentType.YAML, content);
+    constructor(content: string, factory: ParserFactory, parserType: ParserType) {
+        super(DocumentType.YAML, content, factory, parserType);
     }
 }

src/context/syntaxtree/utils/NodeSearch.ts

@@ -88,7 +88,7 @@ export class NodeSearch {
 
             const nearbyNode = rootNode.namedDescendantForPosition(nearbyPoint);
 
-            if (nearbyNode !== originalNode && predicate(nearbyNode)) {
+            if (nearbyNode.id !== originalNode.id && predicate(nearbyNode)) {
                 return nearbyNode;
             }
         }
@@ -161,6 +161,6 @@ export class NodeSearch {
             return false;
         }
 
-        return pair.parent !== mainMapping;
+        return pair.parent?.id !== mainMapping.id;
     }
 }

src/context/syntaxtree/utils/NodeStructure.ts

@@ -161,7 +161,7 @@ export class NodeStructure {
     ): void {
         while (contextPairs.length > 0) {
             const lastPair = contextPairs[contextPairs.length - 1];
-            const lastPairInfo = allPairs.find((p) => p.node === lastPair);
+            const lastPairInfo = allPairs.find((p) => p.node.id === lastPair.id);
             const lastIndentLevel = lastPairInfo?.indentLevel ?? -1;
 
             // If last pair has same or greater indentation, it's a sibling/child - remove it

src/featureFlag/FeatureFlagSupplier.ts

@@ -109,6 +109,7 @@ function featureConfigSupplier(
 const FeatureBuilders = {
     Constants: buildStatic,
     FileDb: buildLocalHost,
+    WasmParser: buildLocalHost,
 } as const satisfies Record<string, FeatureFlagBuilderType>;
 const TargetedFeatureBuilders = {
     EnhancedDryRun: (name: string, config?: FeatureFlagConfigType) => {