From 02e06ff93aecbc75f36ea83aa55876a56f121c34 Mon Sep 17 00:00:00 2001
From: Chris Campbell <ccampbell@climateinteractive.org>
Date: Tue, 2 Feb 2021 13:11:17 -0800
Subject: [PATCH] fix: record variants of a subscripted variable in
 removeUnusedVariables

---
 models/prune/prune.mdl       | 18 ++++++++++++++++++
 models/prune/prune_check.sh  |  4 ++++
 models/prune/prune_spec.json | 23 ++++++++++++-----------
 src/Model.js                 | 15 +++++++++++----
 4 files changed, 45 insertions(+), 15 deletions(-)

diff --git a/models/prune/prune.mdl b/models/prune/prune.mdl
index 6bb4a837..7203ff05 100644
--- a/models/prune/prune.mdl
+++ b/models/prune/prune.mdl
@@ -97,6 +97,24 @@ With Look1 at t1 = WITH LOOKUP ( 1, ([(0,0)-(2,2)],(0,0),(1,1),(2,2)) )
 With Look2 at t1 = WITH LOOKUP ( 1, ([(0,0)-(2,2)],(0,0),(1,1),(2,2)) )
 	~~|
 
+Constant Partial 1 = 1
+	~~|
+
+Constant Partial 2 = 2
+	~~|
+
+Initial Partial[C1] =
+  INITIAL( Constant Partial 1 )
+	~~|
+
+Initial Partial[C2] =
+  INITIAL( Constant Partial 2 )
+	~~|
+
+Partial[C2] =
+  Initial Partial[C2]
+	~~|
+
 ********************************************************
 	.Control
 ********************************************************~
diff --git a/models/prune/prune_check.sh b/models/prune/prune_check.sh
index 1a4b6a61..b3173483 100755
--- a/models/prune/prune_check.sh
+++ b/models/prune/prune_check.sh
@@ -39,6 +39,10 @@ expect_present "__lookup1"
 expect_present "_look1"
 expect_present "_look1_value_at_t1"
 expect_present "_with_look1_at_t1"
+expect_present "_constant_partial_1"
+expect_present "_constant_partial_2"
+expect_present "_initial_partial"
+expect_present "_partial"
 
 # Verify that unreferenced variables do not appear in the generated C file
 expect_not_present "_input_3"
diff --git a/models/prune/prune_spec.json b/models/prune/prune_spec.json
index f7325b46..78dc38ee 100644
--- a/models/prune/prune_spec.json
+++ b/models/prune/prune_spec.json
@@ -1,17 +1,18 @@
 {
   "name": "prune",
-  "inputVars": [
-    "_input_1",
-    "_input_2"
+  "inputVarNames": [
+    "Input 1",
+    "Input 2"
   ],
-  "outputVars": [
-    "_time",
-    "_input_1_and_2_total",
-    "_simple_totals",
-    "_a_totals",
-    "_b1_totals",
-    "_look1_value_at_t1",
-    "_with_look1_at_t1"
+  "outputVarNames": [
+    "Time",
+    "Input 1 and 2 Total",
+    "Simple Totals",
+    "A Totals",
+    "B1 Totals",
+    "Look1 Value at t1",
+    "With Look1 at t1",
+    "Partial[C2]"
   ],
   "externalDatfiles": [
     "prune_data.dat"
diff --git a/src/Model.js b/src/Model.js
index 96b3c972..97f75450 100644
--- a/src/Model.js
+++ b/src/Model.js
@@ -245,7 +245,7 @@ function removeUnusedVariables(spec) {
   // Keep track of all variable names that are referenced somewhere.  Note that we
   // don't attempt to track specific "ref ids" (e.g. `_some_variable[_subscript]`)
   // but instead just track generic variable names (e.g. `_some_variable`).  This
-  // ensure that we include all subscripts for a variable, which might mean we
+  // ensures that we include all subscripts for a variable, which might mean we
   // include some subscripts that aren't needed, but it is safer than trying to
   // eliminate those and possibly omit something that is needed.
   const referencedVarNames = []
@@ -259,7 +259,7 @@ function removeUnusedVariables(spec) {
   }
 
   // Add the given variable to the list of referenced variables, and do the same for
-  // some special things (i.e., lookups) that it might reference.a
+  // some special things (i.e., lookups) that it might reference.
   const recordUsedVariable = v => {
     // Add the variable to the list of referenced variables
     recordUsedVarName(v.varName)
@@ -291,8 +291,15 @@ function removeUnusedVariables(spec) {
   // that it references) as being "used".
   const referencedRefIds = []
   const recordRefsOfVariable = v => {
-    let refs = v.references.concat(v.initReferences)
-    for (const refId of refs) {
+    // If this variable is subscripted, we need to record all subscript variants;
+    // `refIdsWithName` will return those.  We also need to record all variables
+    // that are referenced by this variable, either directly (`v.references`) or
+    // in an "INITIAL" expression (`v.initReferences`).  It's OK if we end up with
+    // duplicates in this list, because we will examine each reference only once.
+    let refIds = refIdsWithName(v.varName)
+    refIds = refIds.concat(v.references)
+    refIds = refIds.concat(v.initReferences)
+    for (const refId of refIds) {
       if (!referencedRefIds.includes(refId)) {
         referencedRefIds.push(refId)
         const refVar = varWithRefId(refId)