Parse command substitution in arithmetic expressions (#2)

Author: jdiamond

src/arithmetic.ts

@@ -1,4 +1,4 @@
-import type { ArithmeticExpression } from "./types.ts";
+import type { ArithmeticCommandExpansion, ArithmeticExpression } from "./types.ts";
 import {
   CH_TAB,
   CH_NL,
@@ -400,14 +400,24 @@ export function parseArithmeticExpression(src: string, offset: number = 0): Arit
         }
       } else {
         // $( command substitution )
-        pos++;
+        pos++; // skip (
         let depth = 1;
         while (pos < len && depth > 0) {
           const ch = src.charCodeAt(pos);
           if (ch === CH_LPAREN) depth++;
           else if (ch === CH_RPAREN) depth--;
           pos++;
         }
+        const text = src.slice(start, pos);
+        const inner = text.slice(2, -1); // remove "$(" and ")"
+        return {
+          type: "ArithmeticCommandExpansion",
+          pos: start + offset,
+          end: pos + offset,
+          text,
+          inner,
+          script: undefined,
+        } satisfies ArithmeticCommandExpansion;
       }
     } else if (c === CH_LBRACE) {
       // ${ parameter expansion }

src/parser.ts

@@ -61,6 +61,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);
     }
     return this.#expression;
   }
@@ -110,7 +111,10 @@ class ArithmeticForImpl implements ArithmeticFor {
     if (this.#initialize === null) {
       if (this.#initStr) {
         const expr = parseArithmeticExpression(this.#initStr);
-        if (expr) offsetArith(expr, this.#initPos);
+        if (expr) {
+          offsetArith(expr, this.#initPos);
+          resolveArithmeticExpansions(expr);
+        }
         this.#initialize = expr ?? undefined;
       } else {
         this.#initialize = undefined;
@@ -126,7 +130,10 @@ class ArithmeticForImpl implements ArithmeticFor {
     if (this.#test === null) {
       if (this.#testStr) {
         const expr = parseArithmeticExpression(this.#testStr);
-        if (expr) offsetArith(expr, this.#testPos);
+        if (expr) {
+          offsetArith(expr, this.#testPos);
+          resolveArithmeticExpansions(expr);
+        }
         this.#test = expr ?? undefined;
       } else {
         this.#test = undefined;
@@ -142,7 +149,10 @@ class ArithmeticForImpl implements ArithmeticFor {
     if (this.#update === null) {
       if (this.#updateStr) {
         const expr = parseArithmeticExpression(this.#updateStr);
-        if (expr) offsetArith(expr, this.#updatePos);
+        if (expr) {
+          offsetArith(expr, this.#updatePos);
+          resolveArithmeticExpansions(expr);
+        }
         this.#update = expr ?? undefined;
       } else {
         this.#update = undefined;
@@ -198,6 +208,32 @@ 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;
+  }
+}
+
 // Lookup tables for O(1) token classification (replaces sequential comparisons)
 const listTerminators = new Uint8Array(37);
 listTerminators[Token.EOF] = 1;

src/parts.ts

@@ -1,6 +1,6 @@
 import type { DeferredCommandExpansion, Word, WordPart } from "./types.ts";
 import { Lexer } from "./lexer.ts";
-import { parse } from "./parser.ts";
+import { parse, resolveArithmeticExpansions } from "./parser.ts";
 
 /**
  * Compute the structural parts of a word by re-scanning the source.
@@ -25,6 +25,13 @@ export function computeWordParts(source: string, word: Word): WordPart[] | undef
     resolveExpansion(exp);
   }
 
+  // Resolve command substitutions inside arithmetic expressions
+  for (const part of parts) {
+    if (part.type === "ArithmeticExpansion" && part.expression) {
+      resolveArithmeticExpansions(part.expression);
+    }
+  }
+
   return parts;
 }
 

src/printer.ts

@@ -323,6 +323,8 @@ function arithExpr(e: ArithmeticExpression): string {
       if (arithNeedsParens(e.right, prec, ra ? true : false)) right = "(" + right + ")";
       return left + " " + e.operator + " " + right;
     }
+    case "ArithmeticCommandExpansion":
+      return e.text; // preserve original text for roundtrip
   }
 }
 

src/types.ts

@@ -94,7 +94,8 @@ export type ArithmeticExpression =
   | ArithmeticUnary
   | ArithmeticTernary
   | ArithmeticGroup
-  | ArithmeticWord;
+  | ArithmeticWord
+  | ArithmeticCommandExpansion;
 
 export interface ArithmeticBinary {
   type: "ArithmeticBinary";
@@ -137,6 +138,15 @@ export interface ArithmeticWord {
   value: string;
 }
 
+export interface ArithmeticCommandExpansion {
+  type: "ArithmeticCommandExpansion";
+  pos: number;
+  end: number;
+  text: string; // e.g., "$(cmd)"
+  inner: string | undefined; // e.g., "cmd" - cleared after resolution
+  script: Script | undefined; // set after resolution
+}
+
 export type DoubleQuotedChild =
   | LiteralPart
   | SimpleExpansionPart

test/arithmetic.test.ts

@@ -5,6 +5,7 @@ import { parse } from "../src/parser.ts";
 import { computeWordParts } from "../src/parts.ts";
 import type {
   ArithmeticBinary,
+  ArithmeticCommandExpansion,
   ArithmeticExpression,
   ArithmeticGroup,
   ArithmeticTernary,
@@ -439,3 +440,85 @@ test("arithmetic expressions parse without errors", () => {
     assert.ok(ast.commands.length > 0, `Failed: ${script}`);
   }
 });
+
+// --- Command substitution in arithmetic ---
+
+test("command substitution in arithmetic - raw parse", () => {
+  const e = parseArithmeticExpression("$(cmd) + 1")!;
+  assert.equal(e.type, "ArithmeticBinary");
+  assert.equal(bin(e).operator, "+");
+  const left = bin(e).left;
+  assert.equal(left.type, "ArithmeticCommandExpansion");
+  assert.equal((left as ArithmeticCommandExpansion).text, "$(cmd)");
+  assert.equal((left as ArithmeticCommandExpansion).inner, "cmd");
+  assert.equal((left as ArithmeticCommandExpansion).script, undefined);
+});
+
+test("command substitution with argument in arithmetic", () => {
+  const e = parseArithmeticExpression("$(echo hello) + x")!;
+  assert.equal(e.type, "ArithmeticBinary");
+  const left = bin(e).left as ArithmeticCommandExpansion;
+  assert.equal(left.type, "ArithmeticCommandExpansion");
+  assert.equal(left.text, "$(echo hello)");
+  assert.equal(left.inner, "echo hello");
+});
+
+test("nested command substitution in arithmetic", () => {
+  const e = parseArithmeticExpression("$(echo $(inner))")!;
+  assert.equal(e.type, "ArithmeticCommandExpansion");
+  assert.equal((e as ArithmeticCommandExpansion).text, "$(echo $(inner))");
+  assert.equal((e as ArithmeticCommandExpansion).inner, "echo $(inner)");
+});
+
+test("command substitution at start and end of expression", () => {
+  const e = parseArithmeticExpression("$(a) + $(b)")!;
+  assert.equal(e.type, "ArithmeticBinary");
+  const left = bin(e).left as ArithmeticCommandExpansion;
+  const right = bin(e).right as ArithmeticCommandExpansion;
+  assert.equal(left.type, "ArithmeticCommandExpansion");
+  assert.equal(right.type, "ArithmeticCommandExpansion");
+  assert.equal(left.text, "$(a)");
+  assert.equal(right.text, "$(b)");
+});
+
+test("command substitution resolved in arithmetic expansion", () => {
+  const ast = parse("echo $(( $(cmd) + 1 ))");
+  const parts = computeWordParts("echo $(( $(cmd) + 1 ))", getCmd(ast).suffix[0])!;
+  const arith = parts[0] as import("../src/types.ts").ArithmeticExpansionPart;
+  assert.equal(arith.type, "ArithmeticExpansion");
+  const binary = arith.expression as ArithmeticBinary;
+  assert.equal(binary.type, "ArithmeticBinary");
+  const left = binary.left as ArithmeticCommandExpansion;
+  assert.equal(left.type, "ArithmeticCommandExpansion");
+  assert.equal(left.inner, undefined); // cleared after resolution
+  assert.ok(left.script); // now populated
+  assert.equal(left.script!.commands[0].command.type, "Command");
+});
+
+test("command substitution in arithmetic command", () => {
+  const ast = parse("(( $(cmd) ))");
+  const arithCmd = ast.commands[0].command as import("../src/types.ts").ArithmeticCommand;
+  const expr = arithCmd.expression!;
+  assert.equal(expr.type, "ArithmeticCommandExpansion");
+  assert.ok((expr as ArithmeticCommandExpansion).script);
+});
+
+test("command substitution in arithmetic for loop", () => {
+  const ast = parse("for (( i = $(start); i < $(limit); i++ )); do echo $i; done");
+  const forLoop = ast.commands[0].command as ArithmeticFor;
+  assert.ok(forLoop.initialize);
+  const initBin = forLoop.initialize as ArithmeticBinary;
+  assert.equal(initBin.type, "ArithmeticBinary");
+  assert.equal(initBin.operator, "=");
+  const initRight = initBin.right as ArithmeticCommandExpansion;
+  assert.equal(initRight.type, "ArithmeticCommandExpansion");
+  assert.equal(initRight.text, "$(start)");
+  assert.ok(initRight.script);
+
+  assert.ok(forLoop.test);
+  const testBin = forLoop.test as ArithmeticBinary;
+  assert.equal(testBin.type, "ArithmeticBinary");
+  const testRight = testBin.right as ArithmeticCommandExpansion;
+  assert.equal(testRight.type, "ArithmeticCommandExpansion");
+  assert.ok(testRight.script);
+});

test/verify.ts

@@ -2,7 +2,7 @@
 // Walks AST, fills gaps from source, validates content fields against source.
 // Also verifies word parts: parts.map(p => p.text).join('') === source span.
 
-import type { Node, Script, WordPart, DoubleQuotedChild } from "../src/types.ts";
+import type { ArithmeticExpression, Node, Script, WordPart, DoubleQuotedChild } from "../src/types.ts";
 import { computeWordParts } from "../src/parts.ts";
 
 type AnyNode = { type: string; pos: number; end: number; [k: string]: any };
@@ -76,6 +76,9 @@ function checkContent(src: string, node: AnyNode) {
     case "ArithmeticWord":
       if (node.value !== span) fail(node, "value", span, node.value);
       break;
+    case "ArithmeticCommandExpansion":
+      if (node.text !== span) fail(node, "text", span, node.text);
+      break;
     case "ArithmeticCommand":
       // body is text between (( and )), source span starts with ((
       if (span.startsWith("((") && span.endsWith("))")) {
@@ -157,4 +160,39 @@ function verifyPartChildren(source: string, part: WordPart | DoubleQuotedChild)
       );
     }
   }
+
+  if (part.type === "ArithmeticExpansion" && part.expression) {
+    verifyArithExpansions(part.expression);
+  }
+}
+
+function verifyArithExpansions(e: ArithmeticExpression): void {
+  switch (e.type) {
+    case "ArithmeticCommandExpansion":
+      if (e.script) {
+        const innerSrc = e.text.slice(2, -1); // remove "$(" and ")"
+        const rebuilt = _verify(innerSrc, e.script as any);
+        if (rebuilt !== innerSrc) {
+          throw new Error(
+            `ArithmeticCommandExpansion inner script verify failed: expected ${JSON.stringify(innerSrc)}, got ${JSON.stringify(rebuilt)}`,
+          );
+        }
+      }
+      break;
+    case "ArithmeticBinary":
+      verifyArithExpansions(e.left);
+      verifyArithExpansions(e.right);
+      break;
+    case "ArithmeticUnary":
+      verifyArithExpansions(e.operand);
+      break;
+    case "ArithmeticTernary":
+      verifyArithExpansions(e.test);
+      verifyArithExpansions(e.consequent);
+      verifyArithExpansions(e.alternate);
+      break;
+    case "ArithmeticGroup":
+      verifyArithExpansions(e.expression);
+      break;
+  }
 }