From 63cf35b47f21d9cf079d738785f57bc678e64d19 Mon Sep 17 00:00:00 2001
From: Calixte Denizet <calixte.denizet@gmail.com>
Date: Mon, 30 Mar 2026 16:30:33 +0200
Subject: [PATCH] Avoid expressions duplication in the ps AST and use a local
 instead when compiling to WASM

---
 src/core/postscript/ast.js           | 45 +++++++++++++++++++++++++++-
 src/core/postscript/wasm_compiler.js | 31 +++++++++++++++++--
 test/unit/postscript_spec.js         | 12 ++++++++
 3 files changed, 84 insertions(+), 4 deletions(-)

diff --git a/src/core/postscript/ast.js b/src/core/postscript/ast.js
index c4c36dcf1..086005d5d 100644
--- a/src/core/postscript/ast.js
+++ b/src/core/postscript/ast.js
@@ -623,7 +623,50 @@ class PSStackToTree {
       stack.push(new PsArgNode(i));
     }
     this._evalBlock(program.body, stack);
-    return this._failed ? null : stack;
+    if (this._failed) {
+      return null;
+    }
+    PSStackToTree.#markShared(stack);
+    return stack;
+  }
+
+  // Set node.shared / sharedCount on non-atomic nodes referenced more than
+  // once.  arg/const are excluded — they are cheap to re-emit inline.
+  static #markShared(outputs) {
+    const refCount = new Map();
+    const visit = node => {
+      if (!node || node.type === PS_NODE.arg || node.type === PS_NODE.const) {
+        return;
+      }
+      const prev = refCount.get(node) ?? 0;
+      refCount.set(node, prev + 1);
+      if (prev > 0) {
+        return;
+      }
+      switch (node.type) {
+        case PS_NODE.unary:
+          visit(node.operand);
+          break;
+        case PS_NODE.binary:
+          visit(node.first);
+          visit(node.second);
+          break;
+        case PS_NODE.ternary:
+          visit(node.cond);
+          visit(node.then);
+          visit(node.otherwise);
+          break;
+      }
+    };
+    for (const output of outputs) {
+      visit(output);
+    }
+    for (const [node, count] of refCount) {
+      if (count > 1) {
+        node.shared = true;
+        node.sharedCount = count; // remaining-use tracking in backends
+      }
+    }
   }
 
   _evalBlock(block, stack) {
diff --git a/src/core/postscript/wasm_compiler.js b/src/core/postscript/wasm_compiler.js
index f792a3e64..32e59d50f 100644
--- a/src/core/postscript/wasm_compiler.js
+++ b/src/core/postscript/wasm_compiler.js
@@ -260,6 +260,8 @@ class PsWasmCompiler {
     // Params 0..nIn-1 are automatically locals; extras start at _nextLocal.
     this._nextLocal = this._nIn;
     this._freeLocals = [];
+    // node → {local, remaining} for shared sub-expression caching (CSE).
+    this._sharedLocals = new Map();
   }
 
   // Wasm emit helpers
@@ -310,9 +312,30 @@ class PsWasmCompiler {
 
   /**
    * Emit Wasm instructions for `node`, leaving exactly one f64 on the Wasm
-   * operand stack.  Returns false if the node cannot be compiled.
+   * operand stack. Returns false if the node cannot be compiled.
    */
   _compileNode(node) {
+    if (node.shared) {
+      const entry = this._sharedLocals.get(node);
+      if (entry !== undefined) {
+        this._emitLocalGet(entry.local);
+        if (--entry.remaining === 0) {
+          this._releaseLocal(entry.local);
+        }
+        return true;
+      }
+      if (!this._compileNodeImpl(node)) {
+        return false;
+      }
+      const local = this._allocLocal();
+      this._sharedLocals.set(node, { local, remaining: node.sharedCount - 1 });
+      this._emitLocalTee(local);
+      return true;
+    }
+    return this._compileNodeImpl(node);
+  }
+
+  _compileNodeImpl(node) {
     switch (node.type) {
       case PS_NODE.arg:
         this._emitLocalGet(node.index);
@@ -642,11 +665,13 @@ class PsWasmCompiler {
   }
 
   _compileStandardBinaryNode(op, first, second) {
-    // Identical compound operands: compile once, reuse via local_tee/local_get.
+    // Identical non-atomic operands: compile once, tee/get.
+    // Skip when shared — _compileNode already handles that case.
     if (
       first === second &&
       first.type !== PS_NODE.arg &&
-      first.type !== PS_NODE.const
+      first.type !== PS_NODE.const &&
+      !first.shared
     ) {
       const tmp = this._allocLocal();
       try {
diff --git a/test/unit/postscript_spec.js b/test/unit/postscript_spec.js
index ee050125e..dafc90394 100644
--- a/test/unit/postscript_spec.js
+++ b/test/unit/postscript_spec.js
@@ -986,6 +986,18 @@ describe("PostScript Type 4 lexer, parser, and Wasm compiler", function () {
       expect(r).toBeCloseTo(0.25, 9);
     });
 
+    it("CSE: shared subexpression compiled once, correct result, one local", async function () {
+      // { 3 add dup mul } → (x+3)^2.  The "x+3" node is shared (dup), so CSE
+      // caches it in one local that is released and reused for later operands.
+      const source = "{ 3 add dup mul }";
+      const bytes = compilePostScriptToWasm(source, [0, 10], [0, 169]);
+      expect(bytes).not.toBeNull();
+      expect(getWasmLocalCount(bytes)).toBe(1);
+      // x=2 → (2+3)^2 = 25
+      const r = await compileAndRun(source, [0, 10], [0, 169], [2]);
+      expect(r).toBeCloseTo(25, 9);
+    });
+
     it("reuses temporary locals across sequential shared-subexpression codegen", function () {
       const bytes = compilePostScriptToWasm(
         "{ dup 1 add dup mul exch 2 add dup mul add }",