From 2ad630de90964def36ac99b0476b842ac6aeb17f Mon Sep 17 00:00:00 2001
From: rvcas <x@rvcas.dev>
Date: Sat, 4 Jun 2022 00:29:30 -0400
Subject: [PATCH] feat: use a stack for scope

---
 crates/uplc/example/integer.uplc |  2 +-
 crates/uplc/src/ast.rs           |  8 +++----
 crates/uplc/src/debruijn.rs      | 40 ++++++++++++++++++++++----------
 3 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/crates/uplc/example/integer.uplc b/crates/uplc/example/integer.uplc
index b00bace0..86df1598 100644
--- a/crates/uplc/example/integer.uplc
+++ b/crates/uplc/example/integer.uplc
@@ -1,3 +1,3 @@
 (program 11.22.33
-    (lam x x)
+    (lam x (lam x y))
 )
\ No newline at end of file
diff --git a/crates/uplc/src/ast.rs b/crates/uplc/src/ast.rs
index a5cfefc2..f0f37663 100644
--- a/crates/uplc/src/ast.rs
+++ b/crates/uplc/src/ast.rs
@@ -86,15 +86,15 @@ pub struct NamedDeBruijn {
 }
 
 #[derive(Debug, Clone, PartialEq, Copy)]
-pub struct DeBruijn(u64);
+pub struct DeBruijn(usize);
 
-impl From<u64> for DeBruijn {
-    fn from(i: u64) -> Self {
+impl From<usize> for DeBruijn {
+    fn from(i: usize) -> Self {
         DeBruijn(i)
     }
 }
 
-impl From<DeBruijn> for u64 {
+impl From<DeBruijn> for usize {
     fn from(d: DeBruijn) -> Self {
         d.0
     }
diff --git a/crates/uplc/src/debruijn.rs b/crates/uplc/src/debruijn.rs
index 16455e9f..624525e9 100644
--- a/crates/uplc/src/debruijn.rs
+++ b/crates/uplc/src/debruijn.rs
@@ -3,18 +3,18 @@ use std::collections::HashMap;
 use crate::ast::{DeBruijn, Name, NamedDeBruijn, Term, Unique};
 
 #[derive(Debug, Copy, Clone)]
-struct Level(u64);
+struct Level(usize);
 
 pub(crate) struct Converter {
     current_level: Level,
-    levels: HashMap<Unique, Level>,
+    levels: Vec<HashMap<Unique, Level>>,
 }
 
 impl Converter {
     pub(crate) fn new() -> Self {
         Converter {
             current_level: Level(0),
-            levels: HashMap::new(),
+            levels: vec![HashMap::new()],
         }
     }
 
@@ -41,11 +41,15 @@ impl Converter {
                     index,
                 };
 
-                self.with_scope();
+                self.start_scope();
+
+                let body = self.name_to_named_debruijn(*body)?;
+
+                self.end_scope();
 
                 Term::Lambda {
                     parameter_name: name,
-                    body: Box::new(self.name_to_named_debruijn(*body)?),
+                    body: Box::new(body),
                 }
             }
             Term::Apply { function, argument } => Term::Apply {
@@ -69,20 +73,32 @@ impl Converter {
     }
 
     fn get_index(&mut self, unique: Unique) -> anyhow::Result<DeBruijn> {
-        if let Some(found_level) = self.levels.get(&unique) {
-            let index = self.current_level.0 - found_level.0;
+        for scope in self.levels.iter().rev() {
+            if let Some(found_level) = scope.get(&unique) {
+                let index = self.current_level.0 - found_level.0;
 
-            Ok(index.into())
-        } else {
-            anyhow::bail!("Free unique: {}", isize::from(unique));
+                return Ok(index.into());
+            }
         }
+
+        anyhow::bail!("Free unique: {}", isize::from(unique))
     }
 
     fn declare_unique(&mut self, unique: Unique) {
-        self.levels.insert(unique, self.current_level);
+        let scope = &mut self.levels[self.current_level.0];
+
+        scope.insert(unique, self.current_level);
     }
 
-    fn with_scope(&mut self) {
+    fn start_scope(&mut self) {
         self.current_level = Level(self.current_level.0 + 1);
+
+        self.levels.push(HashMap::new());
+    }
+
+    fn end_scope(&mut self) {
+        self.current_level = Level(self.current_level.0 - 1);
+
+        self.levels.pop();
     }
 }