Merge pull request #134 from HigherOrderCO/bug/sc-332/unscoped-lambdas-break-in-match-desugaring

[sc-332] Linearize unscoped vars in adt matches

src/term/mod.rs

@@ -1,5 +1,5 @@
 use hvmc::run::Val;
-use indexmap::IndexMap;
+use indexmap::{IndexMap, IndexSet};
 use shrinkwraprs::Shrinkwrap;
 use std::collections::{BTreeMap, HashMap};
 
@@ -361,6 +361,36 @@ impl Term {
     }
   }
 
+  /// Substitute the occurence of an unscoped variable with the given term.
+  pub fn subst_unscoped(&mut self, from: &Name, to: &Term) {
+    match self {
+      Term::Lnk { nam } if nam == from => {
+        *self = to.clone();
+      }
+      Term::Match { scrutinee, arms } => {
+        scrutinee.subst_unscoped(from, to);
+        arms.iter_mut().for_each(|(_, arm)| arm.subst_unscoped(from, to));
+      }
+      Term::List { els } => els.iter_mut().for_each(|el| el.subst_unscoped(from, to)),
+      Term::Chn { bod, .. } | Term::Lam { bod, .. } => bod.subst_unscoped(from, to),
+      Term::App { fun: fst, arg: snd, .. }
+      | Term::Let { val: fst, nxt: snd, .. }
+      | Term::Dup { val: fst, nxt: snd, .. }
+      | Term::Sup { fst, snd, .. }
+      | Term::Tup { fst, snd }
+      | Term::Opx { fst, snd, .. } => {
+        fst.subst(from, to);
+        snd.subst(from, to);
+      }
+      Term::Var { .. }
+      | Term::Lnk { .. }
+      | Term::Ref { .. }
+      | Term::Num { .. }
+      | Term::Str { .. }
+      | Term::Era => (),
+    }
+  }
+
   /// Collects all the free variables that a term has
   /// and the number of times each var is used
   pub fn free_vars(&self) -> IndexMap<Name, u64> {
@@ -437,6 +467,48 @@ impl Term {
     free_vars
   }
 
+  pub fn unscoped_vars(&self) -> (IndexSet<Name>, IndexSet<Name>) {
+    fn go(term: &Term, decls: &mut IndexSet<Name>, uses: &mut IndexSet<Name>) {
+      match term {
+        Term::Chn { tag: _, nam, bod } => {
+          decls.insert(nam.clone());
+          go(bod, decls, uses);
+        }
+        Term::Lnk { nam } => {
+          uses.insert(nam.clone());
+        }
+        Term::Match { scrutinee, arms } => {
+          go(scrutinee, decls, uses);
+          for (_, arm) in arms {
+            go(arm, decls, uses);
+          }
+        }
+        Term::List { els } => {
+          for el in els {
+            go(el, decls, uses);
+          }
+        }
+        Term::Let { val: fst, nxt: snd, .. }
+        | Term::App { fun: fst, arg: snd, .. }
+        | Term::Tup { fst, snd }
+        | Term::Dup { val: fst, nxt: snd, .. }
+        | Term::Sup { fst, snd, .. }
+        | Term::Opx { fst, snd, .. } => {
+          go(fst, decls, uses);
+          go(snd, uses, decls);
+        }
+        Term::Lam { bod, .. } => {
+          go(bod, decls, uses);
+        }
+        Term::Var { .. } | Term::Num { .. } | Term::Str { .. } | Term::Ref { .. } | Term::Era => (),
+      }
+    }
+    let mut decls = Default::default();
+    let mut uses = Default::default();
+    go(self, &mut decls, &mut uses);
+    (decls, uses)
+  }
+
   /// Creates a new [`Term::Match`] from the given terms.
   /// If `scrutinee` is not a `Term::Var`, creates a let binding containing the match in its body
   pub fn new_native_match(

src/term/transform/desugar_match_expressions.rs

@@ -57,6 +57,7 @@ pub enum MatchError {
   Repeated(Name),
   Missing(HashSet<Name>),
   LetPat(Box<MatchError>),
+  Linearize(Name),
 }
 
 impl std::error::Error for MatchError {}
@@ -86,6 +87,7 @@ impl std::fmt::Display for MatchError {
 
         Ok(())
       }
+      MatchError::Linearize(var) => write!(f, "Unable to linearize variable {var} in a match block."),
     }
   }
 }
@@ -118,7 +120,8 @@ impl Term {
           // For now, to prevent extraction we can use `let (a, b) = ...;`
           Type::Adt(_) | Type::Tup => {
             *match_count += 1;
-            let match_term = linearize_match_free_vars(self);
+            let match_term = linearize_match_unscoped_vars(self)?;
+            let match_term = linearize_match_free_vars(match_term);
             let Term::Match { scrutinee: box Term::Var { nam }, arms } = match_term else { unreachable!() };
             let nam = std::mem::take(nam);
             let arms = std::mem::take(arms);
@@ -399,3 +402,52 @@ fn linearize_match_free_vars(match_term: &mut Term) -> &mut Term {
     }
   }
 }
+
+fn linearize_match_unscoped_vars(match_term: &mut Term) -> Result<&mut Term, MatchError> {
+  let Term::Match { scrutinee: _, arms } = match_term else { unreachable!() };
+  // Collect the vars
+  let mut free_vars = IndexSet::new();
+  for (_, arm) in arms.iter_mut() {
+    let (decls, uses) = arm.unscoped_vars();
+    // Not allowed to declare unscoped var and not use it since we need to extract the match arm.
+    if let Some(var) = decls.difference(&uses).next() {
+      return Err(MatchError::Linearize(Name(format!("λ${var}"))));
+    }
+    // Change unscoped var to normal scoped var if it references something outside this match arm.
+    for var in uses.difference(&decls) {
+      arm.subst_unscoped(var, &Term::Var { nam: Name(format!("%match%unscoped%{var}")) });
+    }
+    free_vars.extend(uses);
+  }
+
+  // Add lambdas to the arms
+  for (_, body) in arms {
+    let old_body = std::mem::take(body);
+    let new_body = free_vars.iter().rev().fold(old_body, |body, var| Term::Lam {
+      tag: Tag::Static,
+      nam: Some(Name(format!("%match%unscoped%{var}"))),
+      bod: Box::new(body),
+    });
+    *body = new_body;
+  }
+
+  // Add apps to the match
+  let old_match = std::mem::take(match_term);
+  *match_term = free_vars.into_iter().fold(old_match, |acc, nam| Term::App {
+    tag: Tag::Static,
+    fun: Box::new(acc),
+    arg: Box::new(Term::Lnk { nam }),
+  });
+
+  // Get a reference to the match again
+  // It returns a reference and not an owned value because we want
+  //  to keep the new surrounding Apps but still modify the match further.
+  let mut match_term = match_term;
+  loop {
+    match match_term {
+      Term::App { tag: _, fun, arg: _ } => match_term = fun.as_mut(),
+      Term::Match { .. } => return Ok(match_term),
+      _ => unreachable!(),
+    }
+  }
+}

tests/golden_tests/encode_pattern_match/match_adt_unscoped_lambda.hvm

@@ -0,0 +1,6 @@
+data Maybe = None | (Some val)
+
+main = (match (Some 1) {
+  None: @$x *
+  (Some x): x
+} $x)

tests/golden_tests/encode_pattern_match/match_adt_unscoped_var.hvm

@@ -0,0 +1,13 @@
+data Maybe = None | (Some val)
+
+Foo = @$x match (Some 1) {
+  None: $x
+  (Some x): x
+}
+
+Bar = (match (Some 1) {
+  None: $x
+  (Some x): x
+} @$x *)
+
+main = *

tests/snapshots/encode_pattern_match__match_adt_unscoped_lambda.hvm.snap

@@ -0,0 +1,5 @@
+---
+source: tests/golden_tests.rs
+input_file: tests/golden_tests/encode_pattern_match/match_adt_unscoped_lambda.hvm
+---
+Unable to linearize variable λ$x in a match block.

tests/snapshots/encode_pattern_match__match_adt_unscoped_var.hvm.snap

@@ -0,0 +1,33 @@
+---
+source: tests/golden_tests.rs
+input_file: tests/golden_tests/encode_pattern_match/match_adt_unscoped_var.hvm
+---
+(Foo) = λ$x let %temp%scrutinee = (Some 1); (Foo$match$1 %temp%scrutinee $x)
+
+(Bar) = (let %temp%scrutinee = (Some 1); (Bar$match$1 %temp%scrutinee $x) λ$x *)
+
+(main) = *
+
+(None) = #Maybe λNone #Maybe λSome None
+
+(Some) = λval #Maybe λNone #Maybe λSome #Maybe.Some.val (Some val)
+
+(Foo$match$1) = λx #Maybe (x Foo$match$1$PNone Foo$match$1$PSome)
+
+(Bar$match$1) = λx #Maybe (x Bar$match$1$PNone Bar$match$1$PSome)
+
+(Foo$match$1$R0) = λ%match%unscoped%x %match%unscoped%x
+
+(Foo$match$1$R1) = λx λ%match%unscoped%x x
+
+(Foo$match$1$PNone) = Foo$match$1$R0
+
+(Foo$match$1$PSome) = #Maybe.Some.val λy0 (Foo$match$1$R1 y0)
+
+(Bar$match$1$R0) = λ%match%unscoped%x %match%unscoped%x
+
+(Bar$match$1$R1) = λx λ%match%unscoped%x x
+
+(Bar$match$1$PNone) = Bar$match$1$R0
+
+(Bar$match$1$PSome) = #Maybe.Some.val λy0 (Bar$match$1$R1 y0)