Collect arithmetic command expansions during parse, drop tree walk

parseArithmeticExpression now pushes ArithmeticCommandExpansion nodes
into a module-level pending list as they're created. Callers drain the
list and parse inner scripts directly instead of walking the arith tree.

Eliminates resolveArithmeticExpansions and the per-word tree-walk loops
in parts.ts. Arith-heavy fixtures see ~9% faster `parse + parts`; other
fixtures within noise.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: webpro

src/arithmetic.ts

@@ -106,7 +106,16 @@ function opRightAssoc(op: string): boolean {
   }
 }
 
+let pendingArithCmdExps: ArithmeticCommandExpansion[] | null = null;
+
+export function drainArithCmdExps(): ArithmeticCommandExpansion[] | null {
+  const out = pendingArithCmdExps;
+  pendingArithCmdExps = null;
+  return out;
+}
+
 export function parseArithmeticExpression(src: string, offset: number = 0): ArithmeticExpression | null {
+  pendingArithCmdExps = null;
   let pos = 0;
   const len = src.length;
 
@@ -410,14 +419,16 @@ export function parseArithmeticExpression(src: string, offset: number = 0): Arit
         }
         const text = src.slice(start, pos);
         const inner = text.slice(2, -1); // remove "$(" and ")"
-        return {
+        const node: ArithmeticCommandExpansion = {
           type: "ArithmeticCommandExpansion",
           pos: start + offset,
           end: pos + offset,
           text,
           inner,
           script: undefined,
-        } satisfies ArithmeticCommandExpansion;
+        };
+        (pendingArithCmdExps ??= []).push(node);
+        return node;
       }
     } else if (c === CH_LBRACE) {
       // ${ parameter expansion }

src/lexer.ts

@@ -1,5 +1,6 @@
 // oxlint-disable unicorn/no-thenable
 import type {
+  ArithmeticCommandExpansion,
   DeferredCommandExpansion,
   DoubleQuotedChild,
   ExtGlobOperator,
@@ -8,7 +9,7 @@ import type {
   Word,
   WordPart,
 } from "./types.ts";
-import { parseArithmeticExpression } from "./arithmetic.ts";
+import { parseArithmeticExpression, drainArithCmdExps } from "./arithmetic.ts";
 import { WordImpl } from "./word.ts";
 import {
   CH_TAB,
@@ -325,6 +326,7 @@ export class Lexer {
   private hasPeek: boolean;
   private pendingHereDocs: PendingHereDoc[];
   private collectedExpansions: DeferredCommandExpansion[];
+  private collectedArithCmdExps: ArithmeticCommandExpansion[] | null = null;
   _errors: ParseError[] | null = null;
   _buildParts = false;
 
@@ -351,6 +353,10 @@ export class Lexer {
     return this.collectedExpansions;
   }
 
+  getCollectedArithCmdExps(): ArithmeticCommandExpansion[] | null {
+    return this.collectedArithCmdExps;
+  }
+
   getPos(): number {
     return this.pos;
   }
@@ -1738,6 +1744,11 @@ export class Lexer {
     this._resultHasExpansion = false;
     if (this._buildParts) {
       const expr = parseArithmeticExpression(body) ?? undefined;
+      const drained = drainArithCmdExps();
+      if (drained) {
+        if (this.collectedArithCmdExps) this.collectedArithCmdExps.push(...drained);
+        else this.collectedArithCmdExps = drained;
+      }
       this._resultPart = { type: "ArithmeticExpansion", text, expression: expr };
     } else {
       this._resultPart = undefined;

src/parser.ts

@@ -39,7 +39,7 @@ import type {
   Word,
 } from "./types.ts";
 import { LexContext, Token, Lexer, TokenValue } from "./lexer.ts";
-import { parseArithmeticExpression } from "./arithmetic.ts";
+import { parseArithmeticExpression, drainArithCmdExps } from "./arithmetic.ts";
 import { computeWordParts, computeHereDocBodyParts } from "./parts.ts";
 import { WordImpl } from "./word.ts";
 
@@ -62,7 +62,7 @@ class ArithmeticCommandImpl implements ArithmeticCommand {
   get expression(): ArithmeticExpression | undefined {
     if (this.#expression === null) {
       this.#expression = parseArithmeticExpression(this.body, this.pos + 2) ?? undefined;
-      if (this.#expression) resolveArithmeticExpansions(this.#expression);
+      resolveDrainedArithCmdExps();
     }
     return this.#expression;
   }
@@ -112,10 +112,8 @@ class ArithmeticForImpl implements ArithmeticFor {
     if (this.#initialize === null) {
       if (this.#initStr) {
         const expr = parseArithmeticExpression(this.#initStr);
-        if (expr) {
-          offsetArith(expr, this.#initPos);
-          resolveArithmeticExpansions(expr);
-        }
+        if (expr) offsetArith(expr, this.#initPos);
+        resolveDrainedArithCmdExps();
         this.#initialize = expr ?? undefined;
       } else {
         this.#initialize = undefined;
@@ -131,10 +129,8 @@ class ArithmeticForImpl implements ArithmeticFor {
     if (this.#test === null) {
       if (this.#testStr) {
         const expr = parseArithmeticExpression(this.#testStr);
-        if (expr) {
-          offsetArith(expr, this.#testPos);
-          resolveArithmeticExpansions(expr);
-        }
+        if (expr) offsetArith(expr, this.#testPos);
+        resolveDrainedArithCmdExps();
         this.#test = expr ?? undefined;
       } else {
         this.#test = undefined;
@@ -150,10 +146,8 @@ class ArithmeticForImpl implements ArithmeticFor {
     if (this.#update === null) {
       if (this.#updateStr) {
         const expr = parseArithmeticExpression(this.#updateStr);
-        if (expr) {
-          offsetArith(expr, this.#updatePos);
-          resolveArithmeticExpansions(expr);
-        }
+        if (expr) offsetArith(expr, this.#updatePos);
+        resolveDrainedArithCmdExps();
         this.#update = expr ?? undefined;
       } else {
         this.#update = undefined;
@@ -209,29 +203,14 @@ function offsetArith(node: ArithmeticExpression, base: number): void {
   }
 }
 
-export function resolveArithmeticExpansions(expr: ArithmeticExpression): void {
-  switch (expr.type) {
-    case "ArithmeticBinary":
-      resolveArithmeticExpansions(expr.left);
-      resolveArithmeticExpansions(expr.right);
-      break;
-    case "ArithmeticUnary":
-      resolveArithmeticExpansions(expr.operand);
-      break;
-    case "ArithmeticTernary":
-      resolveArithmeticExpansions(expr.test);
-      resolveArithmeticExpansions(expr.consequent);
-      resolveArithmeticExpansions(expr.alternate);
-      break;
-    case "ArithmeticGroup":
-      resolveArithmeticExpansions(expr.expression);
-      break;
-    case "ArithmeticCommandExpansion":
-      if (expr.inner !== undefined) {
-        expr.script = parse(expr.inner);
-        expr.inner = undefined;
-      }
-      break;
+function resolveDrainedArithCmdExps(): void {
+  const list = drainArithCmdExps();
+  if (!list) return;
+  for (const node of list) {
+    if (node.inner !== undefined) {
+      node.script = parse(node.inner);
+      node.inner = undefined;
+    }
   }
 }
 

src/parts.ts

@@ -1,6 +1,6 @@
 import type { DeferredCommandExpansion, Word, WordPart } from "./types.ts";
 import { Lexer } from "./lexer.ts";
-import { parse, resolveArithmeticExpansions } from "./parser.ts";
+import { parse } from "./parser.ts";
 
 /**
  * Compute the structural parts of a word by re-scanning the source.
@@ -12,19 +12,7 @@ export function computeWordParts(source: string, word: Word): WordPart[] | undef
   const lexer = new Lexer(source);
   const parts = lexer.buildWordParts(word.pos);
   if (!parts) return undefined;
-
-  // Resolve command expansions: parse inner scripts
-  for (const exp of lexer.getCollectedExpansions()) {
-    resolveExpansion(exp);
-  }
-
-  // Resolve command substitutions inside arithmetic expressions
-  for (const part of parts) {
-    if (part.type === "ArithmeticExpansion" && part.expression) {
-      resolveArithmeticExpansions(part.expression);
-    }
-  }
-
+  resolveCollected(lexer);
   return parts;
 }
 
@@ -37,20 +25,23 @@ export function computeHereDocBodyParts(source: string, word: Word): WordPart[]
   const lexer = new Lexer(source);
   const parts = lexer.buildHereDocParts(word.pos, word.end);
   if (!parts) return undefined;
+  resolveCollected(lexer);
+  return parts;
+}
 
-  // Resolve command expansions: parse inner scripts
+function resolveCollected(lexer: Lexer): void {
   for (const exp of lexer.getCollectedExpansions()) {
     resolveExpansion(exp);
   }
-
-  // Resolve command substitutions inside arithmetic expressions
-  for (const part of parts) {
-    if (part.type === "ArithmeticExpansion" && part.expression) {
-      resolveArithmeticExpansions(part.expression);
+  const arithCmdExps = lexer.getCollectedArithCmdExps();
+  if (arithCmdExps) {
+    for (const node of arithCmdExps) {
+      if (node.inner !== undefined) {
+        node.script = parse(node.inner);
+        node.inner = undefined;
+      }
     }
   }
-
-  return parts;
 }
 
 function resolveExpansion(e: DeferredCommandExpansion) {