From edcc1a6fce080e6d0d8f30728f2e60901cf823ac Mon Sep 17 00:00:00 2001
From: Pranav Gaddamadugu <pranav@aleo.org>
Date: Wed, 11 Oct 2023 12:13:49 -0400
Subject: [PATCH] Update RenameTable

---
 .../passes/src/common/rename_table/mod.rs     | 29 +++++++++++++++----
 1 file changed, 23 insertions(+), 6 deletions(-)

diff --git a/compiler/passes/src/common/rename_table/mod.rs b/compiler/passes/src/common/rename_table/mod.rs
index f5a267ade6..d29fea3626 100644
--- a/compiler/passes/src/common/rename_table/mod.rs
+++ b/compiler/passes/src/common/rename_table/mod.rs
@@ -14,6 +14,7 @@
 // You should have received a copy of the GNU General Public License
 // along with the Leo library. If not, see <https://www.gnu.org/licenses/>.
 
+use leo_ast::{NodeID};
 use leo_span::Symbol;
 
 use indexmap::IndexMap;
@@ -24,29 +25,34 @@ pub struct RenameTable {
     /// The `RenameTable` of the parent scope.
     pub(crate) parent: Option<Box<RenameTable>>,
     /// The mapping from names in the original AST to new names in the renamed AST.
-    mapping: IndexMap<Symbol, Symbol>,
+    names: IndexMap<Symbol, Symbol>,
+    /// The mapping from symbols to node IDs.
+    /// These are used to ensure that newly introduced symbols reference the appropriate information
+    /// that has been previously indexed by node ID. e,g. `TypeTable`.
+    ids: IndexMap<Symbol, NodeID>,
 }
 
 impl RenameTable {
     /// Create a new `RenameTable` with the given parent.
     pub(crate) fn new(parent: Option<Box<RenameTable>>) -> Self {
-        Self { parent, mapping: IndexMap::new() }
+        Self { parent, names: IndexMap::new(), ids: IndexMap::new() }
     }
 
     /// Returns the symbols that were renamed in the current scope.
     pub(crate) fn local_names(&self) -> impl Iterator<Item = &Symbol> {
-        self.mapping.keys()
+        self.names.keys()
     }
 
     /// Updates `self.mapping` with the desired entry.
     /// Creates a new entry if `symbol` is not already in `self.mapping`.
-    pub(crate) fn update(&mut self, symbol: Symbol, new_symbol: Symbol) {
-        self.mapping.insert(symbol, new_symbol);
+    pub(crate) fn update(&mut self, symbol: Symbol, new_symbol: Symbol, id: NodeID) {
+        self.names.insert(symbol, new_symbol);
+        self.ids.insert(new_symbol, id);
     }
 
     /// Looks up the new name for `symbol`, recursively checking the parent if it is not found.
     pub(crate) fn lookup(&self, symbol: Symbol) -> Option<&Symbol> {
-        if let Some(var) = self.mapping.get(&symbol) {
+        if let Some(var) = self.names.get(&symbol) {
             Some(var)
         } else if let Some(parent) = &self.parent {
             parent.lookup(symbol)
@@ -54,4 +60,15 @@ impl RenameTable {
             None
         }
     }
+
+    /// Looks up the node ID for `symbol`, recursively checking the parent if it is not found.
+    pub(crate) fn lookup_id(&self, symbol: Symbol) -> Option<&NodeID> {
+        if let Some(id) = self.ids.get(&symbol) {
+            Some(id)
+        } else if let Some(parent) = &self.parent {
+            parent.lookup_id(symbol)
+        } else {
+            None
+        }
+    }
 }