fix: handle case where subscript/index appears in expression position (#398)

Fixes #397

Author: chrispcampbell

models/subscript/subscript.dat

@@ -176,3 +176,12 @@ u[C4]
 0	4
 u[C5]
 0	5
+v[A1]
+0	0
+1	0
+v[A2]
+0	1
+1	1
+v[A3]
+0	0
+1	0

models/subscript/subscript.mdl

@@ -53,6 +53,12 @@ u[C3] = 3 ~~~:SUPPLEMENTARY|
 u[C4] = 4 ~~~:SUPPLEMENTARY|
 u[C5] = 5 ~~~:SUPPLEMENTARY|
 
+v[DimA] = IF THEN ELSE (DimA = A2, 1, 0)
+	~
+	~	dimension name and subscript/index name both referenced in an expression
+  ~:SUPPLEMENTARY
+	|
+
 ********************************************************
 	.Control
 ********************************************************~

packages/compile/src/generate/equation-gen.js

@@ -997,13 +997,18 @@ export default class EquationGen extends ModelReader {
         // Emit the size of the dimension in place of the dimension name.
         this.emit(`${sub(varName).size}`)
       } else {
-        // A subscript masquerading as a variable takes the value of the loop index var plus one
-        // (since Vensim indices are one-based).
+        // A dimension masquerading as a variable (i.e., in expression position) takes the
+        // value of the loop index var plus one (since Vensim indices are one-based).
         let s = this.rhsSubscriptGen([varName])
         // Remove the brackets around the C subscript expression.
         s = s.slice(1, s.length - 1)
         this.emit(`(${s} + 1)`)
       }
+    } else if (isIndex(varName)) {
+      // A subscript masquerading as a variable (i.e., in expression position) takes the
+      // numeric index value plus one (since Vensim indices are one-based).
+      const index = sub(varName).value
+      this.emit(`${index + 1}`)
     } else {
       this.varNames.push(varName)
       if (functionName === '_VECTOR_SELECT') {

packages/compile/src/generate/gen-equation-c.spec.ts

@@ -228,6 +228,34 @@ describe('generateEquation (Vensim -> C)', () => {
     expect(genC(vars.get('_y'))).toEqual(['_y = _IF_THEN_ELSE(!_z, 1.0, 0.0);'])
   })
 
+  it('should work for conditional expression with reference to dimension', () => {
+    const vars = readInlineModel(`
+      DimA: A1, A2 ~~|
+      x = 1 ~~|
+      y[DimA] = IF THEN ELSE(DimA = x, 1, 0) ~~|
+    `)
+    expect(vars.size).toBe(2)
+    expect(genC(vars.get('_x'))).toEqual(['_x = 1.0;'])
+    expect(genC(vars.get('_y'))).toEqual([
+      'for (size_t i = 0; i < 2; i++) {',
+      '_y[i] = _IF_THEN_ELSE((i + 1) == _x, 1.0, 0.0);',
+      '}'
+    ])
+  })
+
+  it('should work for conditional expression with reference to dimension and subscript/index', () => {
+    const vars = readInlineModel(`
+      DimA: A1, A2 ~~|
+      y[DimA] = IF THEN ELSE(DimA = A2, 1, 0) ~~|
+    `)
+    expect(vars.size).toBe(1)
+    expect(genC(vars.get('_y'))).toEqual([
+      'for (size_t i = 0; i < 2; i++) {',
+      '_y[i] = _IF_THEN_ELSE((i + 1) == 2, 1.0, 0.0);',
+      '}'
+    ])
+  })
+
   it('should work for data variable definition', () => {
     const extData: ExtData = new Map([
       [

packages/compile/src/model/equation-reader.js

@@ -22,10 +22,11 @@ import {
 import {
   extractMarkedDims,
   indexNamesForSubscript,
+  isDimension,
+  isIndex,
   normalizeSubscripts,
   separatedVariableIndex,
-  sub,
-  isDimension
+  sub
 } from '../_shared/subscript.js'
 import ModelReader from '../parse/model-reader.js'
 import { createParser } from '../parse/parser.js'
@@ -238,8 +239,8 @@ export default class EquationReader extends ModelReader {
     // Get the var name of a variable in a call and save it as a reference.
     let id = ctx.Id().getText()
     let varName = canonicalName(id)
-    // Do not add a dimension name as a reference.
-    if (!isDimension(varName)) {
+    // Do not add a dimension or index name as a reference.
+    if (!isDimension(varName) && !isIndex(varName)) {
       let fn = this.currentFunctionName()
       this.refId = varName
       this.expandedRefIds = []

packages/compile/src/model/read-equations.spec.ts

@@ -119,6 +119,38 @@ describe('readEquations', () => {
     ])
   })
 
+  it('should work for conditional expression with reference to dimension', () => {
+    const vars = readInlineModel(`
+      DimA: A1, A2 ~~|
+      x = 1 ~~|
+      y[DimA] = IF THEN ELSE(DimA = x, 1, 0) ~~|
+    `)
+    expect(vars).toEqual([
+      v('x', '1', {
+        refId: '_x',
+        varType: 'const'
+      }),
+      v('y[DimA]', 'IF THEN ELSE(DimA=x,1,0)', {
+        refId: '_y',
+        subscripts: ['_dima'],
+        references: ['_x']
+      })
+    ])
+  })
+
+  it('should work for conditional expression with reference to dimension and subscript/index', () => {
+    const vars = readInlineModel(`
+      DimA: A1, A2 ~~|
+      y[DimA] = IF THEN ELSE(DimA = A2, 1, 0) ~~|
+    `)
+    expect(vars).toEqual([
+      v('y[DimA]', 'IF THEN ELSE(DimA=A2,1,0)', {
+        refId: '_y',
+        subscripts: ['_dima']
+      })
+    ])
+  })
+
   it('should work for data variable definition', () => {
     const vars = readInlineModel(
       `
@@ -7693,6 +7725,10 @@ describe('readEquations', () => {
         subscripts: ['_c5'],
         varType: 'const'
       }),
+      v('v[DimA]', 'IF THEN ELSE(DimA=A2,1,0)', {
+        refId: '_v',
+        subscripts: ['_dima']
+      }),
       v('Time', '', {
         refId: '_time',
         varType: 'const'

packages/compile/src/model/read-variables.spec.ts

@@ -1350,6 +1350,9 @@ describe('readVariables', () => {
       v('u[C5]', '5', {
         subscripts: ['_c5']
       }),
+      v('v[DimA]', 'IF THEN ELSE(DimA=A2,1,0)', {
+        subscripts: ['_dima']
+      }),
       v('Time', '')
     ])
   })