From 62e1ab1400d887aeccfab9ed406103483c66af95 Mon Sep 17 00:00:00 2001
From: Chris Campbell <ccampbell@climateinteractive.org>
Date: Sat, 2 Dec 2023 22:56:06 -0800
Subject: [PATCH] fix: change reduceConditionals to not resolve variable
 references outside of the condition expression (#409)

Fixes #408
---
 packages/parse/src/ast/reduce-expr.spec.ts | 34 ++++++++++++++++++++++
 packages/parse/src/ast/reduce-expr.ts      | 20 ++++++-------
 2 files changed, 43 insertions(+), 11 deletions(-)

diff --git a/packages/parse/src/ast/reduce-expr.spec.ts b/packages/parse/src/ast/reduce-expr.spec.ts
index a1268474..ffb2af13 100644
--- a/packages/parse/src/ast/reduce-expr.spec.ts
+++ b/packages/parse/src/ast/reduce-expr.spec.ts
@@ -232,4 +232,38 @@ describe('reduceConditionals', () => {
       binaryOp(y, '*', zero)
     )
   })
+
+  it('should try to resolve variable refs in the condition but not elsewhere', () => {
+    interface Variable {
+      refId: string
+      expr?: Expr
+    }
+
+    const xVar: Variable = { refId: '_x', expr: one }
+
+    function resolveVarRef(varRef: VariableRef): Expr | undefined {
+      if (varRef.varId === '_x') {
+        return xVar.expr
+      } else {
+        return undefined
+      }
+    }
+
+    // Note that unlike the `reduceExpr` tests above, the `reduceConditionals` function
+    // should not attempt to reduce `variable-ref` expressions outside of the condition
+    // expression, so the `x` variable references below are expected to remain unchanged
+    // (should not be reduced to or replaced with the constant value)
+    const opts: ReduceExprOptions = {
+      resolveVarRef
+    }
+    expect(reduceConditionals(varRef('x'), opts)).toEqual(varRef('x'))
+    expect(reduceConditionals(varRef('y'), opts)).toEqual(varRef('y'))
+
+    // In this example, the `x` used in the condition expression resolves to a constant
+    // value of 1, so the whole `IF THEN ELSE` is expected to be replaced with the "true"
+    // branch expression
+    expect(reduceConditionals(binaryOp(call('IF THEN ELSE', varRef('x'), varRef('x'), zero), '*', two), opts)).toEqual(
+      binaryOp(varRef('x'), '*', two)
+    )
+  })
 })
diff --git a/packages/parse/src/ast/reduce-expr.ts b/packages/parse/src/ast/reduce-expr.ts
index 9dbda936..8ea269fa 100644
--- a/packages/parse/src/ast/reduce-expr.ts
+++ b/packages/parse/src/ast/reduce-expr.ts
@@ -296,8 +296,7 @@ export function reduceExpr(expr: Expr, opts?: ReduceExprOptions): Expr {
     }
 
     default:
-      return expr
-    // assertNever(expr)
+      assertNever(expr)
   }
 }
 
@@ -319,14 +318,9 @@ export function reduceConditionals(expr: Expr, opts?: ReduceExprOptions): Expr {
       return expr
 
     case 'variable-ref':
-      if (opts?.resolveVarRef !== undefined) {
-        // Note that we assume the resolved expression has already been reduced by
-        // the callback, and don't attempt to reduce further
-        const resolvedExpr = opts.resolveVarRef(expr)
-        if (resolvedExpr) {
-          return resolvedExpr
-        }
-      }
+      // If we get here, it means that we are processing an expression that is not
+      // inside the condition expression for an `IF THEN ELSE`, so we do not attempt
+      // to resolve the variable reference or reduce it further
       return expr
 
     case 'unary-op': {
@@ -367,9 +361,13 @@ export function reduceConditionals(expr: Expr, opts?: ReduceExprOptions): Expr {
     case 'function-call': {
       if (expr.fnId === '_IF_THEN_ELSE') {
         // Note that (unlike all other places in this function) we use `reduceExpr` here
-        // to aggressively reduce the condition
+        // to aggressively reduce the condition expression (the first argument for the
+        // `IF THEN ELSE` call)
         const conditionExpr = reduceExpr(expr.args[0], opts)
         if (conditionExpr.kind === 'number') {
+          // The condition resolved to a simple numeric constant.  If it is non-zero,
+          // replace the `IF THEN ELSE` call with the "true" branch, otherwise replace
+          // it with the "false" branch.
           return conditionExpr.value !== 0
             ? reduceConditionals(expr.args[1], opts)
             : reduceConditionals(expr.args[2], opts)