Fix unused dup var linearization

src/term/transform/linearize.rs

@@ -111,11 +111,6 @@ fn count_var_uses_in_term(term: &Term, uses: &mut HashMap<Name, Val>) {
   }
 }
 
-/// We use a special named variable to only count the number of actual var uses on the resulting term
-fn used_var_counter(var_name: &Name) -> Name {
-  Name(format!("${var_name}$used"))
-}
-
 /// Var-declaring terms
 fn term_with_bind_to_affine(
   term: &mut Term,
@@ -123,17 +118,10 @@ fn term_with_bind_to_affine(
   var_uses: &mut HashMap<Name, Val>,
   let_bodies: &mut HashMap<Name, Term>,
 ) {
-  if let Some(nam_some) = nam {
-    if let Some(mut uses) = var_uses.get(nam_some).copied() {
+  if let Some(name) = nam {
+    if let Some(_) = var_uses.get(name).copied() {
       term_to_affine(term, var_uses, let_bodies);
-
-      // When a term is removed, the number of uses of a variable can change inside the term_to_affine call.
-      // We check the updated count instead of using the one we caulculated initially
-      if uses != 0 {
-        let name = used_var_counter(nam_some);
-        uses = var_uses.get(&name).copied().unwrap_or(0)
-      };
-
+      let uses = actual_var_use(var_uses, name);
       duplicate_lam(nam, term, uses);
       return;
     }
@@ -181,13 +169,23 @@ fn term_to_affine(term: &mut Term, var_uses: &mut HashMap<Name, Val>, let_bodies
         uses => {
           term_to_affine(val, var_uses, let_bodies);
           term_to_affine(nxt, var_uses, let_bodies);
-          duplicate_let(nam, nxt, uses, val);
+
+          let mut new_uses = actual_var_use(var_uses, nam);
+
+          // If the number of uses changed (because a term was linearized out):
+          //  We could redo the whole pass on the term.
+          //  Or we can just create extra half-erased-dups `let {nam_n *} = nam_m_dup; nxt` to match the correct labels
+          if uses != new_uses {
+            new_uses += 1
+          };
+
+          duplicate_let(nam, nxt, new_uses, val)
         }
       }
       *term = std::mem::take(nxt.as_mut());
     }
 
-    Term::Let { pat: Pattern::Var(None), val, nxt, .. } => {
+    Term::Let { pat: Pattern::Var(None), val, nxt } => {
       if val.has_unscoped() {
         term_to_affine(val, var_uses, let_bodies);
         term_to_affine(nxt, var_uses, let_bodies);
@@ -211,17 +209,16 @@ fn term_to_affine(term: &mut Term, var_uses: &mut HashMap<Name, Val>, let_bodies
 
     // Var-using terms
     Term::Var { nam } => {
-      let uses = var_uses[nam];
       *var_uses.get_mut(nam).unwrap() -= 1;
       if let Some(subst) = let_bodies.remove(nam) {
         *term = subst.clone();
       } else {
         let used_counter = used_var_counter(nam);
 
-        *nam = dup_name(nam, uses);
+        let actual_uses = var_uses.entry(used_counter).or_default();
+        *actual_uses += 1;
 
-        // Updates de actual var uses on the resulting term
-        *var_uses.entry(used_counter).or_default() += 1;
+        *nam = dup_name(nam, *actual_uses);
       }
     }
 
@@ -255,25 +252,56 @@ fn get_var_uses(nam: Option<&Name>, var_uses: &HashMap<Name, Val>) -> Val {
   if let Some(nam) = nam { *var_uses.get(nam).unwrap() } else { 0 }
 }
 
-fn make_dup_tree(nam: &Name, nxt: &mut Term, uses: Val, dup_body: Option<&mut Term>) {
-  // TODO: Is there a difference between a list of dups and a complete binary tree of dups
-  // Creates this: "dup x1 x1_dup = body; dup x2 x2_dup = x1_dup; dup x3 x4 = x2_dup; nxt"
+// TODO: Is it really necessary to `count_var_uses_in_term` before `term_to_affine`?
+// Can't it be unified in the same function? (given that the actual var uses in a term is already dynamic)
+//
+/// When a term is removed, the number of uses of a variable can change inside the term_to_affine call.
+/// This returns the updated count instead of using the one we calculated initially
+fn actual_var_use(var_uses: &HashMap<Name, Val>, name: &Name) -> Val {
+  var_uses.get(&used_var_counter(name)).copied().unwrap_or(0)
+}
+
+/// We use a special named variable to only count the number of actual var uses on the resulting term
+fn used_var_counter(var_name: &Name) -> Name {
+  Name(format!("${var_name}$used"))
+}
+
+// TODO: Is there a difference between a list of dups and a complete binary tree of dups?
+//
+/// # Example
+///
+/// ```txt
+/// let {x1 x1_dup} = body;
+/// let {x2 x2_dup} = x1_dup;
+/// let {x3 x4}     = x2_dup;
+/// nxt
+/// ```
+fn make_dup_tree(nam: &Name, nxt: &mut Term, uses: Val, mut dup_body: Option<&mut Term>) {
+  let free_vars = &mut nxt.free_vars();
+
+  if let Some(ref body) = dup_body {
+    free_vars.extend(body.free_vars())
+  };
+
+  let make_name = |uses| {
+    let dup_name = dup_name(nam, uses);
+    free_vars.contains_key(&dup_name).then(|| dup_name)
+  };
+
   for i in (1 .. uses).rev() {
-    let old_nxt = std::mem::replace(nxt, Term::Era);
     *nxt = Term::Dup {
       tag: Tag::Auto,
-      fst: Some(dup_name(nam, i)),
-      snd: if i == uses - 1 { Some(dup_name(nam, uses)) } else { Some(internal_dup_name(nam, uses)) },
+      fst: make_name(i),
+      snd: if i == uses - 1 { make_name(uses) } else { Some(internal_dup_name(nam, uses)) },
       val: if i == 1 {
-        if let Some(dup_body) = &dup_body {
-          Box::new((*dup_body).clone()) // TODO: don't clone here
-        } else {
-          Box::new(Term::Var { nam: nam.clone() })
+        match dup_body.as_deref_mut() {
+          Some(body) => Box::new(std::mem::take(body)),
+          None => Box::new(Term::Var { nam: nam.clone() }),
         }
       } else {
         Box::new(Term::Var { nam: internal_dup_name(nam, uses) })
       },
-      nxt: Box::new(old_nxt),
+      nxt: Box::new(std::mem::take(nxt)),
     };
   }
 }

tests/golden_tests/compile_file/unused_dup_var_linearization.hvm

@@ -0,0 +1,10 @@
+Main =
+  let var = @x x
+  let copy = var
+  let copy2 = var
+  let copy3 = var
+  let copy_copy = copy
+  let copy_copy2 = copy
+  let copy_copy3 = copy2
+  let copy_copy4 = copy2
+  (var copy_copy3)

tests/snapshots/compile_file__unused_dup_var_linearization.hvm.snap

@@ -0,0 +1,6 @@
+---
+source: tests/golden_tests.rs
+input_file: tests/golden_tests/compile_file/unused_dup_var_linearization.hvm
+---
+@main = (a a)
+

tests/snapshots/compile_file_o_all__exp.hvm.snap

@@ -2,5 +2,5 @@
 source: tests/golden_tests.rs
 input_file: tests/golden_tests/compile_file_o_all/exp.hvm
 ---
-@main = ({5 ({5 a b} c) (c {5 b d})} (a d))
+@main = ({5 (a {5 b c}) ({5 c d} a)} (d b))
 

tests/snapshots/compile_file_o_all__list_merge_sort.hvm.snap

@@ -5,7 +5,7 @@ input_file: tests/golden_tests/compile_file_o_all/list_merge_sort.hvm
 @F = (a (* a))
 @G = (* (a a))
 @J = (* @Nil)
-@L = {8 a {6 {4 @L {4 @J (b c)}} ({3 b (a d)} {4 {8 d {6 c e}} {4 * e}})}}
+@L = {8 a {6 {4 @L {4 @J (b c)}} ({3 (a d) b} {4 {8 d {6 c e}} {4 * e}})}}
 @Map = ({4 @L {4 @J a}} a)
 @Nil = {4 * {4 a a}}
 @Pure = (a {4 {8 a {6 @Nil b}} {4 * b}})
@@ -14,16 +14,16 @@ input_file: tests/golden_tests/compile_file_o_all/list_merge_sort.hvm
 @S = (a ({4 @Q {4 @R (a b)}} b))
 @V = {8 a {6 {4 @o {4 @p (b (a c))}} (b c)}}
 @W = (* @Nil)
-@b = {8 {15 a {17 b c}} {6 {19 {4 @b {4 @c (d (e (f g)))}} h} ({5 d {7 i (j (c {2 @F {2 @G ({4 {8 k {6 l m}} {4 * m}} ({4 {8 a {6 g n}} {4 * n}} o))}}))}} ({9 e {11 k j}} ({13 f {4 @d {4 @e (i ({4 {8 b {6 h p}} {4 * p}} l))}}} o)))}}
+@b = {8 {15 a {17 b c}} {6 {19 d {4 @b {4 @c (e (f (g h)))}}} ({5 (i (a {2 @F {2 @G ({4 {8 j {6 k l}} {4 * l}} ({4 {8 c {6 h m}} {4 * m}} n))}})) {7 o e}} ({9 i {11 j f}} ({13 {4 @d {4 @e (o ({4 {8 b {6 d p}} {4 * p}} k))}} g} n)))}}
 @c = (* @s)
 @d = {8 a {6 b (c ({4 @b {4 @c (c (a (b d)))}} d))}}
 @e = (* (a a))
-@i = {8 a {6 {4 @V {4 @W (b {4 @i {4 @j (c (d e))}})}} ({21 {23 b f} c} ({4 @d {4 @e (f (a d))}} e))}}
+@i = {8 a {6 {4 @V {4 @W (b {4 @i {4 @j (c (d e))}})}} ({21 c {23 f b}} ({4 @d {4 @e (f (a d))}} e))}}
 @j = (* (a a))
 @main = (a (b c))
 & @S ~ (a (d c))
 & @Map ~ (b (@Pure d))
-@o = {8 a {6 {4 @V {4 @W (b c)}} ({23 b d} ({4 @d {4 @e (d (a e))}} {4 {8 e {6 c f}} {4 * f}}))}}
+@o = {8 a {6 {4 @V {4 @W (b c)}} ({23 d b} ({4 @d {4 @e (d (a e))}} {4 {8 e {6 c f}} {4 * f}}))}}
 @p = (* @t)
 @s = (a (b {4 {8 a {6 b c}} {4 * c}}))
 @t = (a {4 {8 a {6 @Nil b}} {4 * b}})

tests/snapshots/compile_file_o_all__match_dup_and_reconstruction.hvm.snap

@@ -4,7 +4,7 @@ input_file: tests/golden_tests/compile_file_o_all/match_dup_and_reconstruction.h
 ---
 @6 = {4 a (b [b a])}
 @Boxed = (a {2 {4 a b} b})
-@Got = ({3 a {2 @6 (a b)}} b)
+@Got = ({3 {2 @6 (a b)} a} b)
 @main = a
 & @Got ~ (b a)
 & @Boxed ~ (#10 b)

tests/snapshots/compile_file_o_all__scrutinee_reconstruction.hvm.snap

@@ -6,7 +6,7 @@ input_file: tests/golden_tests/compile_file_o_all/scrutinee_reconstruction.hvm
 @8 = {4 * @A}
 @A = (a (* a))
 @None = {2 * {2 a a}}
-@Option.or = ({3 a {2 @8 {2 @7 (a b)}}} b)
+@Option.or = ({3 {2 @8 {2 @7 (a b)}} a} b)
 @Some = (a {2 {4 a b} {2 * b}})
 @main = a
 & @Option.or ~ (b (@None a))

tests/snapshots/compile_file_o_all__spacing.hvm.snap

@@ -2,5 +2,5 @@
 source: tests/golden_tests.rs
 input_file: tests/golden_tests/compile_file_o_all/spacing.hvm
 ---
-@main = ({3 a (a b)} b)
+@main = ({3 (a b) a} b)
 

tests/snapshots/compile_term__infer_dup.hvm.snap

@@ -2,4 +2,4 @@
 source: tests/golden_tests.rs
 input_file: tests/golden_tests/compile_term/infer_dup.hvm
 ---
-({3 a (a b)} b)
+({3 (a b) a} b)

tests/snapshots/compile_term__let_substitution.hvm.snap

@@ -2,4 +2,4 @@
 source: tests/golden_tests.rs
 input_file: tests/golden_tests/compile_term/let_substitution.hvm
 ---
-({3 a (a b)} b)
+({3 (a b) a} b)

tests/snapshots/compile_term__lets.hvm.snap

@@ -4,6 +4,6 @@ input_file: tests/golden_tests/compile_term/lets.hvm
 ---
 a
 & (b b) ~ (c (d (e a)))
-& (f f) ~ {9 g {11 h {13 i (i (h (g e)))}}}
-& (j j) ~ {5 k {7 l (l (k d))}}
-& (m m) ~ {3 n (n c)}
+& (f f) ~ {9 (g (h (i e))) {11 g {13 h i}}}
+& (j j) ~ {5 (k (l d)) {7 k l}}
+& (m m) ~ {3 (n c) n}