Merge pull request #1376 from ccnmtl/js-quote-unquote-compiler

Make a "JIT JS compiler" for math.js expressions.

Author: drewyoungren

media/src/objects/Field.svelte

@@ -5,7 +5,7 @@
 <script>
     import { onMount, onDestroy, untrack } from 'svelte';
     import * as THREE from 'three';
-    import { create, all, e } from 'mathjs';
+    import { create, all, e, xgcd } from 'mathjs';
 
     import M from '../M.svelte';
     import ObjHeader from './ObjHeader.svelte';
@@ -17,6 +17,7 @@
     import InputChecker from '../form-components/InputChecker.svelte';
     import PlayButtons from '../form-components/PlayButtons.svelte';
     import FlowArrowMesh from './FlowArrowMesh';
+    import { mathToJSFunction } from './mathutils';
 
     const config = {};
     const math = create(all, config);
@@ -314,20 +315,20 @@
         render();
     };
 
-    let fieldF;
-
-    const updateField = function () {
+    let fieldF = $derived.by(() => {
         const { p, q, r } = params;
 
-        const [P, Q, R] = [p, q, r].map((x) => math.parse(x).compile());
-
-        fieldF = (x, y, z, vec) => {
-            const args = { x, y, z };
-            vec.set(P.evaluate(args), Q.evaluate(args), R.evaluate(args));
+        const [P, Q, R] = [p, q, r].map((x) =>
+            mathToJSFunction(x, ['x', 'y', 'z']),
+        );
 
+        return (x, y, z, vec) => {
+            vec.set(P(x, y, z), Q(x, y, z), R(x, y, z));
             return vec;
         };
-    };
+    });
+
+    const updateField = function () {};
 
     let maxLength = 2;
     flowArrows.name = uuid;

media/src/objects/Function.svelte

@@ -10,7 +10,7 @@
     import { LineMaterial } from 'three/addons/lines/LineMaterial.js';
     import { LineSegmentsGeometry } from 'three/addons/lines/LineSegmentsGeometry.js';
 
-    import { create, all, neutronMassDependencies } from 'mathjs';
+    import { create, all } from 'mathjs';
 
     import M from '../M.svelte';
     import ObjHeader from './ObjHeader.svelte';
@@ -37,6 +37,9 @@
         blockGeometry,
         checksum,
     } from '../utils.js';
+
+    import { mathToJSFunction } from './mathutils';
+
     import { flashDance } from '../sceneUtils';
 
     let {
@@ -209,10 +212,12 @@
     scene.add(point);
 
     // Compile main function
-    let func = $derived.by(() => {
-        const z = math.parse(params.z).compile();
-        return (x, y, t) => z.evaluate({ x, y, t });
-    });
+    // let func = $derived.by(() => {
+    //     const z = math.parse(params.z).compile();
+    //     return (x, y, t) => z.evaluate({ x, y, t });
+    // });
+
+    let func = $derived(mathToJSFunction(params.z, ['x', 'y', 't']));
 
     const tangentVectors = function () {
         // const arrowParams = {
@@ -721,8 +726,8 @@
                 ymax: D,
                 level: lev,
                 zLevel: 0,
-                nX: data.nX,
-                nY: data.nX,
+                nX: 200,
+                nY: 200,
             });
 
             if (points.length > 1) {

media/src/objects/Surface.svelte

@@ -34,6 +34,9 @@
         checksum,
         ParametricGeometry,
     } from '../utils.js';
+
+    import { mathToJSFunction } from './mathutils';
+
     import { flashDance } from '../sceneUtils';
     import InputChecker from '../form-components/InputChecker.svelte';
     import ColorBar from '../settings/ColorBar.svelte';
@@ -180,33 +183,34 @@
         }
     });
 
-    let xyz;
-    let abcd;
-
-    // compile (in the math.js sense) each expression once they change
-    // and thus have been cleared
-    $effect(() => {
-        // console.log('parm compile');
+    let xyz = $derived.by(() => {
         const [x, y, z] = [params.x, params.y, params.z].map((f) =>
-            math.parse(f).compile(),
+            mathToJSFunction(f, ['u', 'v', 't']),
         );
-        xyz = (u, v, vec, t = 0) =>
-            vec.set(
-                x.evaluate({ u, v, t }),
-                y.evaluate({ u, v, t }),
-                z.evaluate({ u, v, t }),
-            );
-        // });
-        // $effect(() => {
+        return (u, v, vec, t = 0) =>
+            vec.set(x(u, v, t), y(u, v, t), z(u, v, t));
+    });
+    let abcd = $derived.by(() => {
         const [a, b] = [params.a, params.b].map((f) => math.evaluate(f));
-        const [c, d] = [params.c, params.d].map((f) => math.parse(f).compile());
+        const [c, d] = [params.c, params.d].map((f) =>
+            mathToJSFunction(f, ['u']),
+        );
 
-        // take uv on unit square to actual uv coords for parametricgeom
-        abcd = (u, v) => [
+        // take uv on unit square to actual uv coords for ParametricGeom
+        return (u, v) => [
             a + u * (b - a),
-            c.evaluate({ u: a + u * (b - a) }) * (1 - v) +
-                d.evaluate({ u: a + u * (b - a) }) * v,
+            c(a + u * (b - a)) * (1 - v) + d(a + u * (b - a)) * v,
         ];
+    });
+
+    // compile (in the math.js sense) each expression once they change
+    // and thus have been cleared
+    $effect(() => {
+        // console.log('parm compile');
+
+        // React on changes here.
+        const { a, b, c, d, x, y, z } = params;
+
         untrack(updateSurface);
     });
 
@@ -304,15 +308,15 @@
     });
 
     const updateSurface = function () {
-        const { t0, t1 } = params;
-        const time = t0
-            ? math.evaluate(t0) + tau * (math.evaluate(t1) - math.evaluate(t0))
-            : 0;
+        // const { t0, t1 } = params;
+        // const time = t0
+        //     ? math.evaluate(t0) + tau * (math.evaluate(t1) - math.evaluate(t0))
+        //     : 0;
 
         const geometry = new ParametricGeometry(
             (u, v, vec) => {
                 const [U, V] = abcd(u, v);
-                xyz(U, V, vec, time);
+                xyz(U, V, vec, tVal);
             },
             nX || 30,
             nX || 30,

media/src/objects/mathutils.js

@@ -0,0 +1,138 @@
+import {
+  parse,
+  simplifyConstant,
+  OperatorNode,
+  FunctionNode,
+  SymbolNode,
+  ConditionalNode,
+} from 'mathjs';
+
+const allowedMathFns = new Set([
+  'sin', 'cos', 'tan', 'asin', 'acos', 'atan', 'atan2',
+  'sqrt', 'pow', 'exp', 'log', 'abs',
+  'min', 'max', 'round', 'floor', 'ceil', 'sign', 'hypot',
+]);
+
+const disallowedSymbols = new Set([
+  '__proto__', 'constructor', 'prototype',
+  'Function', 'window', 'document', 'globalThis',
+]);
+
+const constantMap = {
+  pi: 'Math.PI',
+  e: 'Math.E',
+  LOG2E: 'Math.LOG2E',
+  LOG10E: 'Math.LOG10E',
+  LN2: 'Math.LN2',
+  LN10: 'Math.LN10',
+  SQRT2: 'Math.SQRT2',
+  SQRT1_2: 'Math.SQRT1_2',
+};
+
+/**
+ * Checks if a math.js AST contains a ConditionalNode
+ * @param {*} node 
+ * @returns boolean
+ */
+function containsConditional(node) {
+  let found = false;
+  node.traverse((child) => {
+    if (child.type === 'ConditionalNode') {
+      found = true;
+    }
+  });
+  return found;
+}
+
+
+// 🔁 Transform e^x → exp(x)
+function rewriteEToExp(ast) {
+  return ast.transform((node) => {
+    if (
+      node.type === 'OperatorNode' &&
+      node.fn === 'pow' &&
+      node.args.length === 2 &&
+      node.args[0].type === 'SymbolNode' &&
+      node.args[0].name === 'e'
+    ) {
+      return new FunctionNode('exp', [node.args[1]]);
+    }
+    return node;
+  });
+}
+
+function mathToJSFunction(expression, variableNames = null) {
+  const rawAST = parse(expression);
+  const rewrittenAST = rewriteEToExp(rawAST);
+  const ast = containsConditional(rewrittenAST) ? rewrittenAST : simplifyConstant(rewrittenAST);
+
+  // Infer variables if not given
+  const inferredVars = new Set();
+  ast.traverse((node) => {
+    if (node.isSymbolNode) {
+      const name = node.name;
+      if (disallowedSymbols.has(name)) {
+        throw new Error(`Disallowed symbol: ${name}`);
+      }
+      if (!allowedMathFns.has(name) && !constantMap[name] && !/^[0-9.]+$/.test(name)) {
+        inferredVars.add(name);
+      }
+    }
+  });
+
+  const vars = variableNames ?? Array.from(inferredVars);
+
+  // Compile AST to JS expression string
+  const jsExpr = ast.toString({
+    handler: (node, options) => {
+      switch (node.type) {
+        case 'OperatorNode':
+          if (node.fn === 'pow') {
+            return `Math.pow(${node.args[0].toString(options)}, ${node.args[1].toString(options)})`;
+          }
+          if (node.args.length === 2) {
+            return `(${node.args[0].toString(options)} ${node.op} ${node.args[1].toString(options)})`;
+          } else if (node.args.length === 1) {
+            return `(${node.op}${node.args[0].toString(options)})`;
+          }
+          break;
+
+        case 'FunctionNode':
+          const fnName = node.fn.name;
+          if (fnName === 'mod') {
+            const [a, b] = node.args;
+            return `(((${a.toString(options)} % ${b.toString(options)}) + ${b.toString(options)}) % ${b.toString(options)})`;
+          }
+
+          if (!allowedMathFns.has(fnName)) {
+            throw new Error(`Unsupported function: ${fnName}`);
+          }
+          const jsFn = fnName === 'atan2' ? 'Math.atan2' : `Math.${fnName}`;
+          return `${jsFn}(${node.args.map(arg => arg.toString(options)).join(', ')})`;
+
+        case 'ConstantNode':
+          return `${node.value}`;
+
+        case 'SymbolNode':
+          return constantMap[node.name] || node.name;
+
+        case 'ParenthesisNode':
+          return `(${node.content.toString(options)})`;
+
+        case 'ConditionalNode':
+          return `(${node.condition.toString(options)} ? ${node.trueExpr.toString(options)} : ${node.falseExpr.toString(options)})`;
+
+        default:
+          throw new Error(`Unsupported node type: ${node.type}`);
+      }
+    },
+    parenthesis: 'keep',
+    implicit: 'show',
+  });
+
+  const argList = vars.join(', ');
+  return new Function(argList, `return ${jsExpr};`);
+}
+
+
+export { mathToJSFunction };