From b402e35683f36d7accefc78f1a3ca1a5722a456d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=A9bastien=20Villemot?= <sebastien@dynare.org>
Date: Fri, 5 Jun 2020 17:11:29 +0200
Subject: [PATCH] In the driver file, provide mapping between model local
 variables and indices in the temporary terms vector

The structures are M_.model_local_variables_{dynamic,static}_tt_idxs.

Ref. dynare#1722
---
 src/DynamicModel.cc | 8 ++++++++
 src/ModelTree.cc    | 2 +-
 src/ModelTree.hh    | 2 +-
 src/StaticModel.cc  | 8 ++++++++
 4 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/src/DynamicModel.cc b/src/DynamicModel.cc
index 629b3ec3..7558c70a 100644
--- a/src/DynamicModel.cc
+++ b/src/DynamicModel.cc
@@ -2575,6 +2575,14 @@ DynamicModel::writeDriverOutput(ostream &output, const string &basename, bool bl
       for (size_t i = 0; i < temporary_terms_derivatives.size(); i++)
         output << temporary_terms_derivatives[i].size() + (i == 0 ? temporary_terms_mlv.size() : 0) << "; ";
       output << "];" << endl;
+
+      /* Write mapping between model local variables and indices in the temporary
+         terms vector (dynare#1722) */
+      output << modstruct << "model_local_variables_dynamic_tt_idxs = {" << endl;
+      for (auto [mlv, value] : temporary_terms_mlv)
+        output << "  '" << symbol_table.getName(mlv->symb_id) << "', "
+               << temporary_terms_idxs.at(mlv)+1 << ';' << endl;
+      output << "};" << endl;
     }
 
   // Write equation tags
diff --git a/src/ModelTree.cc b/src/ModelTree.cc
index 5fbd7fd8..a958bf08 100644
--- a/src/ModelTree.cc
+++ b/src/ModelTree.cc
@@ -70,7 +70,7 @@ ModelTree::copyHelper(const ModelTree &m)
 
   // Temporary terms
   for (const auto &it : m.temporary_terms_mlv)
-    temporary_terms_mlv[f(it.first)] = f(it.second);
+    temporary_terms_mlv[dynamic_cast<VariableNode *>(f(it.first))] = f(it.second);
   for (const auto &it : m.temporary_terms_derivatives)
     temporary_terms_derivatives.push_back(convert_temporary_terms_t(it));
   for (const auto &it : m.temporary_terms_idxs)
diff --git a/src/ModelTree.hh b/src/ModelTree.hh
index a5cc30c8..1b3eb22d 100644
--- a/src/ModelTree.hh
+++ b/src/ModelTree.hh
@@ -123,7 +123,7 @@ protected:
 
   //! Used model local variables, that will be treated as temporary terms
   /*! See the comments in ModelTree::computeTemporaryTerms() */
-  map<expr_t, expr_t, ExprNodeLess> temporary_terms_mlv;
+  map<VariableNode *, expr_t, ExprNodeLess> temporary_terms_mlv;
 
   //! Temporary terms for residuals and derivatives
   /*! Index 0 is temp. terms of residuals, index 1 for first derivatives, ... */
diff --git a/src/StaticModel.cc b/src/StaticModel.cc
index ad5fd35d..1c347c3b 100644
--- a/src/StaticModel.cc
+++ b/src/StaticModel.cc
@@ -1728,6 +1728,14 @@ StaticModel::writeDriverOutput(ostream &output, bool block) const
     output << temporary_terms_derivative.size() << "; ";
   output << "];" << endl;
 
+  /* Write mapping between model local variables and indices in the temporary
+     terms vector (dynare#1722) */
+  output << "M_.model_local_variables_static_tt_idxs = {" << endl;
+  for (auto [mlv, value] : temporary_terms_mlv)
+    output << "  '" << symbol_table.getName(mlv->symb_id) << "', "
+           << temporary_terms_idxs.at(mlv)+1 << ';' << endl;
+  output << "};" << endl;
+
   if (!block)
     return;