From 0cd4386f32f5f74412cfb7d10823cc498d7cb735 Mon Sep 17 00:00:00 2001
From: imaqtkatt <eduarda.so@proton.me>
Date: Fri, 28 Jun 2024 11:08:03 -0300
Subject: [PATCH] Reorder local defs lifting

---
 src/imp/lift_local_defs.rs                          | 12 +++++-------
 .../golden_tests/desugar_file/local_def_shadow.bend | 10 ++++++++++
 .../desugar_file__local_def_shadow.bend.snap        | 13 +++++++++++++
 tests/snapshots/desugar_file__main_aux.bend.snap    |  8 ++++----
 4 files changed, 32 insertions(+), 11 deletions(-)
 create mode 100644 tests/golden_tests/desugar_file/local_def_shadow.bend
 create mode 100644 tests/snapshots/desugar_file__local_def_shadow.bend.snap

diff --git a/src/imp/lift_local_defs.rs b/src/imp/lift_local_defs.rs
index 6e3b3eef..1353f1d6 100644
--- a/src/imp/lift_local_defs.rs
+++ b/src/imp/lift_local_defs.rs
@@ -24,17 +24,15 @@ impl Stmt {
     match self {
       Stmt::LocalDef { .. } => {
         let Stmt::LocalDef { mut def, mut nxt } = std::mem::take(self) else { unreachable!() };
-        let children = def.lift_local_defs(gen)?;
-        nxt.lift_local_defs(parent, defs, gen)?;
-
         let local_name = Name::new(format!("{}__local_{}_{}", parent, gen, def.name));
+        def.body.lift_local_defs(&local_name, defs, gen)?;
+        nxt.lift_local_defs(parent, defs, gen)?;
         *gen += 1;
 
         let (r#use, mut def, fvs) = gen_use(local_name.clone(), *def, nxt)?;
         *self = r#use;
         apply_closure(&mut def, fvs);
 
-        defs.extend(children);
         defs.insert(def.name.clone(), def);
         Ok(())
       }
@@ -122,8 +120,8 @@ fn gen_use(
 
 fn apply_closure(def: &mut fun::Definition, fvs: Vec<Name>) {
   let rule = &mut def.rules[0];
-  let mut n_pats = fvs.into_iter().map(|x| Pattern::Var(Some(x))).collect::<Vec<_>>();
+  let mut captured = fvs.into_iter().map(|x| Pattern::Var(Some(x))).collect::<Vec<_>>();
   let rule_pats = std::mem::take(&mut rule.pats);
-  n_pats.extend(rule_pats);
-  rule.pats = n_pats;
+  captured.extend(rule_pats);
+  rule.pats = captured;
 }
diff --git a/tests/golden_tests/desugar_file/local_def_shadow.bend b/tests/golden_tests/desugar_file/local_def_shadow.bend
new file mode 100644
index 00000000..94c4d4b8
--- /dev/null
+++ b/tests/golden_tests/desugar_file/local_def_shadow.bend
@@ -0,0 +1,10 @@
+def main:
+  def A():
+    def B():
+      return 0
+    return B()
+  def A():
+    def B():
+      return 1
+    return B()
+  return A()
diff --git a/tests/snapshots/desugar_file__local_def_shadow.bend.snap b/tests/snapshots/desugar_file__local_def_shadow.bend.snap
new file mode 100644
index 00000000..f1ff3f33
--- /dev/null
+++ b/tests/snapshots/desugar_file__local_def_shadow.bend.snap
@@ -0,0 +1,13 @@
+---
+source: tests/golden_tests.rs
+input_file: tests/golden_tests/desugar_file/local_def_shadow.bend
+---
+(main__local_0_A__local_0_B) = 0
+
+(main__local_1_A__local_1_B) = 1
+
+(main__local_1_A) = λa a
+
+(main__local_0_A) = λa a
+
+(main) = (main__local_1_A main__local_1_A__local_1_B)
diff --git a/tests/snapshots/desugar_file__main_aux.bend.snap b/tests/snapshots/desugar_file__main_aux.bend.snap
index 03a86542..d532f38f 100644
--- a/tests/snapshots/desugar_file__main_aux.bend.snap
+++ b/tests/snapshots/desugar_file__main_aux.bend.snap
@@ -2,10 +2,10 @@
 source: tests/golden_tests.rs
 input_file: tests/golden_tests/desugar_file/main_aux.bend
 ---
-(aux__local_0_aux) = λa λb (+ b a)
+(main__local_0_aux__local_0_aux__local_0_aux) = λa λb (+ b a)
 
-(aux__local_1_aux) = λa λb λc (a b c)
+(main__local_0_aux__local_0_aux) = λa λb λc (a b c)
 
-(main__local_2_aux) = λa λb λc λd (b a c d)
+(main__local_0_aux) = λa λb λc λd (a b c d)
 
-(main) = (main__local_2_aux aux__local_0_aux aux__local_1_aux 89 2)
+(main) = (main__local_0_aux main__local_0_aux__local_0_aux main__local_0_aux__local_0_aux__local_0_aux 89 2)