compiler: implement pos-only mode for implicit spaces (#602)

Implement a "pos-only" mode for implicit space skipping. In pos-only
mode, we don't build a CST, do failure recording, etc. This results in a
significant performance improvement.

**Before**

```
$ bin/es5bench-wasm
Compile: 66ms
......
JS match: 3312.9ms ± 284.9ms (n=3)
Wasm match: 79.1ms ± 26.7ms (n=3)
Wasm vs JS match: 41.9x
Wasm memory: 174.0MB
```

**After**
```
$ bin/es5bench-wasm
Compile: 57ms
......
JS match: 3311.9ms ± 113.2ms (n=3)
Wasm match: 55.8ms ± 15.7ms (n=3)
Wasm vs JS match: 59.3x
Wasm memory: 87.0MB
```

Author: pdubroy

packages/compiler/runtime/ohmRuntime.ts

@@ -299,6 +299,8 @@ function useMemoizedResult(ruleId: i32, result: MemoEntry): ApplyResult {
 
 // Evaluate $spaces without allocating a CST node or pushing a binding.
 // Stores a sentinel in the memo table encoding just the match length.
+// Calls the pos-only compiled version of $spaces (at table index numRules + 1),
+// which skips all CST building internally.
 export function evalSpacesImplicit(): void {
   const memo = memoTableGet(pos, IMPLICIT_SPACES_RULE_ID);
   if (memo !== EMPTY) {
@@ -307,11 +309,8 @@ export function evalSpacesImplicit(): void {
     return;
   }
   const origPos = pos;
-  const origChunk = bindingsChunk;
-  const origIdx = bindingsIdx;
-  evalRuleBody(IMPLICIT_SPACES_RULE_ID);
+  call_indirect<RuleEvalResult>(numRules + 1);
   const matchLen = <i32>pos - <i32>origPos;
-  restoreBindings(origChunk, origIdx); // discard child bindings
   memoTableSet(origPos, IMPLICIT_SPACES_RULE_ID, (matchLen << 2) | MEMO_SPACES_FLAG);
 }
 

packages/compiler/src/Compiler.ts

@@ -254,6 +254,7 @@ class Assembler {
   _frameDepth: number;
   _savedAtDepthStack: Set<string>[];
   useChunkedBindings: boolean;
+  posOnlyMode: boolean;
 
   // In WebAssembly, function indices (funcidx) are a flat numbering: imports
   // come first, then locally-defined functions. The prebuilt module has its
@@ -281,6 +282,7 @@ class Assembler {
     this._frameDepth = 0;
     this._savedAtDepthStack = [];
     this.useChunkedBindings = useChunkedBindings;
+    this.posOnlyMode = false;
 
     this._typeMap = typeMap;
   }
@@ -627,6 +629,9 @@ class Assembler {
   }
 
   saveNumBindings(): SavedBindingsState {
+    if (this.posOnlyMode) {
+      return noopBacktrackPoint.bindings;
+    }
     if (!this.useChunkedBindings) {
       // Array mode: save a single length value.
       this.callPrebuiltFunc('getBindingsLength');
@@ -707,6 +712,7 @@ class Assembler {
   }
 
   maybeRecordFailure(origPosThunk: () => void, failureId: number): void {
+    if (this.posOnlyMode) return;
     this.globalGet('errorMessagePos');
     this.i32Const(0);
     this.emit(instr.i32.ge_s);
@@ -810,6 +816,10 @@ class Assembler {
   }
 
   newIterNode(saved: SavedBacktrackPoint, arity: number, isOpt = false): void {
+    if (this.posOnlyMode) {
+      this.i32Const(1);
+      return;
+    }
     saved.pos.get();
     this.globalGet('pos');
     saved.bindings.getChunk();
@@ -822,6 +832,10 @@ class Assembler {
   // Wrap the bindings accumulated since the last pushDepth() in a
   // nonterminal node.
   newNonterminalNode(saved: SavedBacktrackPoint, ruleId: number): void {
+    if (this.posOnlyMode) {
+      this.i32Const(1);
+      return;
+    }
     saved.pos.get();
     this.globalGet('pos');
     this.i32Const(ruleId);
@@ -833,6 +847,10 @@ class Assembler {
 
   // [startIdx: i32] -> [tagged: i32]
   newTerminalNode(): void {
+    if (this.posOnlyMode) {
+      this.i32Const(1);
+      return;
+    }
     this.localGet('postSpacesPos');
     this.globalGet('pos');
     this.callPrebuiltFunc('newTerminalNode');
@@ -854,17 +872,20 @@ class Assembler {
   }
 
   pushFluffySavePoint(): void {
+    if (this.posOnlyMode) return;
     this.callPrebuiltFunc('pushFluffySavePoint');
   }
 
   popFluffySavePoint(shouldMark: boolean): void {
+    if (this.posOnlyMode) return;
     this.i32Const(shouldMark ? 1 : 0);
     this.callPrebuiltFunc('popFluffySavePoint');
   }
 
   // Pop the fluffy save point, marking failures as fluffy only when
   // pos matches errorMessagePos. This mirrors ohm-js's scoped failure recording.
   popFluffySavePointIfAtErrorPos(): void {
+    if (this.posOnlyMode) return;
     this.globalGet('errorMessagePos');
     this.globalGet('pos');
     this.i32Eq();
@@ -875,6 +896,7 @@ class Assembler {
   // On success, mark failures as fluffy if pos is at errorMessagePos;
   // on failure, discard without marking.
   popFluffySavePointOnResult(): void {
+    if (this.posOnlyMode) return;
     this.localGet('ret');
     this.ifElse(
       w.blocktype.empty,
@@ -1262,10 +1284,11 @@ export class Compiler {
     const restoreFailurePos = name === '$spaces';
 
     const descriptionId = ruleInfo.description ? this._strings.add(ruleInfo.description) : -1;
-    const hasDescription = descriptionId >= 0;
+    const hasDescription = !asm.posOnlyMode && descriptionId >= 0;
 
+    const funcName = asm.posOnlyMode ? `$${name}_posOnly` : `$${name}`;
     this.beginLexContext(!ruleInfo.isSyntactic);
-    asm.addFunction(`$${name}`, paramTypes, [w.valtype.i32], () => {
+    asm.addFunction(funcName, paramTypes, [w.valtype.i32], () => {
       asm.addLocal('ret', w.valtype.i32);
       asm.addLocal('tmp', w.valtype.i32);
       asm.addLocal('postSpacesPos', w.valtype.i32);
@@ -1605,10 +1628,11 @@ export class Compiler {
       exports.push(w.export_(name, [0x03, prebuilt.globalidxByName[name]]));
     }
     // The module will have a table containing references to all of the rule eval functions,
-    // plus a compiler-generated $isRuleSyntactic dispatch function at the end.
-    // The rule ID can be used directly as the table index; $isRuleSyntactic is at index numRules.
+    // plus compiler-generated helper functions at the end:
+    //   [numRules]     = $isRuleSyntactic
+    //   [numRules + 1] = $spaces pos-only (for evalSpacesImplicit)
     const numRules = this.ruleIdByName.size;
-    const tableSize = numRules + 1; // +1 for $isRuleSyntactic
+    const tableSize = numRules + 2; // +1 for $isRuleSyntactic, +1 for $spaces pos-only
     const table = w.table(
       w.tabletype(w.elemtype.funcref, w.limits.minmax(tableSize, tableSize))
     );
@@ -1617,6 +1641,10 @@ export class Compiler {
     const isRuleSyntacticIdx = functionDecls.findIndex(f => f.name === '$isRuleSyntactic');
     assert(isRuleSyntacticIdx !== -1, 'No $isRuleSyntactic function found');
     tableData.push(w.funcidx(compilerFuncOffset + isRuleSyntacticIdx));
+    // Add $spaces pos-only as the next table entry.
+    const spacesPosOnlyIdx = functionDecls.findIndex(f => f.name === '$$spaces_posOnly');
+    assert(spacesPosOnlyIdx !== -1, 'No $$spaces_posOnly function found');
+    tableData.push(w.funcidx(compilerFuncOffset + spacesPosOnlyIdx));
     assert(tableSize === tableData.length, 'Invalid table size');
 
     // Determine the index of the start function.
@@ -1772,6 +1800,12 @@ export class Compiler {
     });
     ruleDecls.push(checkNotNull(asm._functionDecls.at(-1)));
 
+    // Compile a pos-only version of $spaces that skips all CST building.
+    // This goes into the table at index numRules + 1.
+    asm.posOnlyMode = true;
+    ruleDecls.push(this.compileRule('$spaces'));
+    asm.posOnlyMode = false;
+
     return ruleDecls;
   }
 
@@ -2009,7 +2043,7 @@ export class Compiler {
 
     const {asm} = this;
 
-    if (this.shouldInlineRule(exp.ruleName)) {
+    if (this.shouldInlineRule(exp.ruleName) || asm.posOnlyMode) {
       this.emitInlinedApply(exp);
       return;
     }

packages/compiler/test/test-wasm.js

@@ -212,6 +212,48 @@ test('cst: leadingSpaces children via lazy parsing', async t => {
   });
 });
 
+test('cst: leadingSpaces with custom spaces rule', async t => {
+  const g = await compileAndLoad(`G {
+    Start = word+
+    word = letter+
+    space += comment
+    comment = "//" (~"\\n" any)*
+  }`);
+  g.match('abc // yo\n def').use(r => {
+    t.true(r.succeeded());
+    const root = r.getCstRoot();
+
+    // The Plus list should contain two words.
+    const list = root.children[0];
+    t.is(list.children.length, 2);
+
+    const [word1, word2] = list.children;
+    t.is(word1.sourceString, 'abc');
+    t.is(word2.sourceString, 'def');
+
+    // word2 should have leadingSpaces covering the comment and whitespace.
+    const spaces = word2.leadingSpaces;
+    t.truthy(spaces);
+    t.is(spaces.ctorName, 'spaces');
+    t.is(spaces.sourceString, ' // yo\n ');
+    t.is(spaces.matchLength, 8);
+
+    // Lazy CST children should have correct source strings.
+    const starList = spaces.children[0];
+    t.true(starList.isList());
+    t.is(starList.children.length, 4);
+    t.is(starList.children[0].sourceString, ' ');
+    t.is(starList.children[1].sourceString, '// yo');
+    t.is(starList.children[2].sourceString, '\n');
+    t.is(starList.children[3].sourceString, ' ');
+
+    // The comment child should be a 'space' wrapping a 'comment'.
+    const commentSpace = starList.children[1];
+    t.is(commentSpace.ctorName, 'space');
+    t.is(commentSpace.children[0].ctorName, 'comment');
+  });
+});
+
 test('cst: lazy parsing survives memory.grow()', async t => {
   const g = await compileAndLoad('G { Start = "x" }');
   g.match('  x').use(r => {

packages/runtime/src/miniohm.ts

@@ -54,6 +54,20 @@ const UnicodeCategoryNames = [
 const utf8 = new TextDecoder('utf-8');
 const utf16le = new TextDecoder('utf-16le');
 
+function isSyntacticRuleName(ruleName: string): boolean {
+  const firstChar = ruleName[0];
+  return firstChar === firstChar.toUpperCase();
+}
+
+function getRuleId(ctx: MatchContext, ptr: number): number {
+  return ctx.view.getInt32(ptr + CST_TYPE_AND_DETAILS_OFFSET, true) >>> 2;
+}
+
+// Check whether the nonterminal at `ptr` is a syntactic rule (name starts with uppercase).
+function isSyntacticRule(ctx: MatchContext, ptr: number): boolean {
+  return ctx.ruleIsSyntactic[getRuleId(ctx, ptr)];
+}
+
 // Minimal implementation of Interval (for FailedMatchResult)
 export class Interval {
   startIdx: number;
@@ -199,6 +213,8 @@ export class Grammar {
   /** @internal */
   private _ruleNames: string[] = [];
   /** @internal */
+  private _ruleIsSyntactic: boolean[] = [];
+  /** @internal */
   private _input = '';
 
   /** @internal */
@@ -319,9 +335,11 @@ export class Grammar {
   private _extractStrings(module: WebAssembly.Module): void {
     assert(this._ruleNames.length === 0);
     assert(this._ruleIds.size === 0);
+    assert(this._ruleIsSyntactic.length === 0);
     for (const ruleName of parseStringTable(module, 'ruleNames')) {
       this._ruleIds.set(ruleName, this._ruleIds.size);
       this._ruleNames.push(ruleName);
+      this._ruleIsSyntactic.push(isSyntacticRuleName(ruleName));
     }
     for (const str of parseStringTable(module, 'strings')) {
       this._strings.push(str);
@@ -363,6 +381,7 @@ export class Grammar {
 
     const ctx: MatchContext = {
       ruleNames: this._ruleNames,
+      ruleIsSyntactic: this._ruleIsSyntactic,
       view: new DataView(buffer),
       input,
       getSpacesLenAt: exports.getSpacesLenAt,
@@ -417,6 +436,7 @@ export class Grammar {
     const {exports} = this._instance as any;
     ctx ??= {
       ruleNames: this._ruleNames,
+      ruleIsSyntactic: this._ruleIsSyntactic,
       view: new DataView(exports.memory.buffer),
       input: this._input,
       getSpacesLenAt: exports.getSpacesLenAt,
@@ -425,7 +445,7 @@ export class Grammar {
     };
     const spacesLen = Math.max(0, exports.getSpacesLenAt(0));
     const rootAddr = exports.bindingsAt(0);
-    const root = new CstNodeImpl(ctx, rootAddr, spacesLen);
+    const root = new CstNodeImpl(ctx, rootAddr, spacesLen, isSyntacticRule(ctx, rootAddr));
     if (spacesLen > 0) {
       root.leadingSpaces = new LazySpacesNode(ctx, 0, spacesLen);
     }
@@ -449,6 +469,7 @@ export class Grammar {
 
 export interface MatchContext {
   ruleNames: string[];
+  ruleIsSyntactic: boolean[];
   view: DataView;
   input: string;
   getSpacesLenAt?: (pos: number) => number;
@@ -532,11 +553,17 @@ class CstNodeImpl implements CstNodeBase {
   leadingSpaces?: NonterminalNode = undefined;
   source: {startIdx: number; endIdx: number};
 
-  constructor(ctx: MatchContext, ptr: number, startIdx: number) {
+  // Whether this node's children are in a syntactic context (i.e., have
+  // implicit space skipping between them). Nonterminals set this from
+  // their rule id; other node types inherit from their parent.
+  _syntactic!: boolean;
+
+  constructor(ctx: MatchContext, ptr: number, startIdx: number, syntactic?: boolean) {
     // Non-enumerable properties
     Object.defineProperties(this, {
       _ctx: {value: ctx},
       _children: {writable: true},
+      _syntactic: {value: syntactic ?? false, writable: true},
     });
     this._base = ptr;
     this.startIdx = startIdx;
@@ -663,15 +690,17 @@ class CstNodeImpl implements CstNodeBase {
     const children: (CstNodeImpl | TaggedTerminalNode)[] = [];
     let {startIdx} = this;
     const {getSpacesLenAt} = this._ctx;
+    // Only look up implicit spaces when we're in a syntactic context.
+    const doSpaceLookup = this._syntactic && !!getSpacesLenAt;
     for (let i = 0; i < this.count; i++) {
       const slotOffset = this._base + CST_CHILDREN_OFFSET + i * 4;
       const ptr = this._ctx.view.getUint32(slotOffset, true);
 
       if (isTaggedTerminal(ptr)) {
         // Tagged terminals always have parent-level space skipping.
         let spacesLen = 0;
-        if (getSpacesLenAt) {
-          spacesLen = Math.max(0, getSpacesLenAt(startIdx));
+        if (doSpaceLookup) {
+          spacesLen = Math.max(0, getSpacesLenAt!(startIdx));
           if (spacesLen > 0) startIdx += spacesLen;
         }
         const node = new TaggedTerminalNode(this._ctx, ptr, startIdx);
@@ -686,18 +715,24 @@ class CstNodeImpl implements CstNodeBase {
       // Only query spaces for terminals and nonterminals — not for
       // iteration (ITER_FLAG) or optional nodes, which handle space
       // skipping internally.
-      const type = (this._ctx.view.getInt32(ptr + CST_TYPE_AND_DETAILS_OFFSET, true) &
-        MATCH_RECORD_TYPE_MASK) as MatchRecordType;
+      const typeAndDetails = this._ctx.view.getInt32(ptr + CST_TYPE_AND_DETAILS_OFFSET, true);
+      const type = (typeAndDetails & MATCH_RECORD_TYPE_MASK) as MatchRecordType;
       let spacesLen = 0;
       if (
-        getSpacesLenAt &&
+        doSpaceLookup &&
         (type === MatchRecordType.NONTERMINAL || type === MatchRecordType.TERMINAL)
       ) {
-        spacesLen = Math.max(0, getSpacesLenAt(startIdx));
+        spacesLen = Math.max(0, getSpacesLenAt!(startIdx));
         if (spacesLen > 0) startIdx += spacesLen;
       }
 
-      const node = new CstNodeImpl(this._ctx, ptr, startIdx);
+      // Nonterminals determine syntactic context from their name;
+      // other node types (iter, opt) inherit from the parent.
+      const childSyntactic =
+        type === MatchRecordType.NONTERMINAL
+          ? this._ctx.ruleIsSyntactic[typeAndDetails >>> 2]
+          : this._syntactic;
+      const node = new CstNodeImpl(this._ctx, ptr, startIdx, childSyntactic);
       if (spacesLen > 0) {
         node.leadingSpaces = new LazySpacesNode(this._ctx, startIdx - spacesLen, spacesLen);
       }
@@ -713,13 +748,12 @@ class CstNodeImpl implements CstNodeBase {
 
   isSyntactic(): boolean {
     assert(this.isNonterminal(), 'Not a nonterminal');
-    const firstChar = this.ctorName[0];
-    return firstChar === firstChar.toUpperCase();
+    return this._syntactic;
   }
 
   isLexical(): boolean {
     assert(this.isNonterminal(), 'Not a nonterminal');
-    return !this.isSyntactic();
+    return !this._syntactic;
   }
 
   toString(): string {
@@ -824,7 +858,8 @@ class LazySpacesNode implements NonterminalNode {
     if (memory && this._ctx.view.buffer !== memory.buffer) {
       this._ctx.view = new DataView(memory.buffer);
     }
-    const fullNode = new CstNodeImpl(this._ctx, ptr, this._startIdx);
+    // The spaces rule is lexical, so pass syntactic=false.
+    const fullNode = new CstNodeImpl(this._ctx, ptr, this._startIdx, false);
     return fullNode.children;
   }