Expand refs in main, don't float combinators in main

CHANGELOG.md

@@ -1,9 +1,12 @@
 # Changelog
 
-## 0.2.32
+## 0.2.33
+- Added `expand_main`, a compilation pass that expands references in the entry point function. (#424)
+- Changed the `float_combinators` pass to not extract in the entry point function. (#424)
 
-- Added the built-in `Tree` datatype.
-- Added `![]` and `!` syntax for `Tree` literals.
+## 0.2.32
+- Added the built-in `Tree` datatype. (#528)
+- Added `![]` and `!` syntax for `Tree` literals. (#528)
 - Moved the builtins documentation to `/docs`.
 - Created the changelog.
 

Cargo.toml

@@ -2,7 +2,7 @@
 name = "bend-lang"
 description = "A high-level, massively parallel programming language"
 license = "Apache-2.0"
-version = "0.2.32"
+version = "0.2.33"
 edition = "2021"
 rust-version = "1.74"
 exclude = ["tests/"]

src/fun/transform/expand_main.rs

@@ -0,0 +1,121 @@
+use crate::{
+  fun::{Book, Name, Pattern, Term},
+  maybe_grow,
+};
+use std::collections::HashMap;
+
+impl Book {
+  /// Expands the main function so that it is not just a reference.
+  /// While technically correct, directly returning a reference is never what users want.
+  pub fn expand_main(&mut self) {
+    if self.entrypoint.is_none() {
+      return;
+    }
+
+    let main = self.defs.get_mut(self.entrypoint.as_ref().unwrap()).unwrap();
+    let mut main_bod = std::mem::take(&mut main.rule_mut().body);
+
+    let mut seen = vec![self.entrypoint.as_ref().unwrap().clone()];
+    main_bod.expand_ref_return(self, &mut seen, &mut 0);
+
+    let main = self.defs.get_mut(self.entrypoint.as_ref().unwrap()).unwrap();
+    main.rule_mut().body = main_bod;
+  }
+}
+
+impl Term {
+  /// Expands references in the main function that are in "return" position.
+  ///
+  /// This applies to:
+  /// - When main returns a reference.
+  /// - When main returns a lambda whose body is a reference.
+  /// - When main returns a pair or superposition and one of its elements is a reference.
+  ///
+  /// Only expand recursive functions once.
+  pub fn expand_ref_return(&mut self, book: &Book, seen: &mut Vec<Name>, globals_count: &mut usize) {
+    maybe_grow(|| match self {
+      Term::Ref { nam } => {
+        if seen.contains(nam) {
+          // Don't expand recursive references
+        } else {
+          seen.push(nam.clone());
+          let mut body = book.defs.get(nam).unwrap().rule().body.clone();
+          body.rename_unscoped(globals_count, &mut HashMap::new());
+          *self = body;
+          self.expand_ref_return(book, seen, globals_count);
+          seen.pop().unwrap();
+        }
+      }
+      Term::Fan { els, .. } | Term::List { els } => {
+        for el in els {
+          el.expand_ref_return(book, seen, globals_count);
+        }
+      }
+      Term::Lam { bod: nxt, .. }
+      | Term::With { bod: nxt, .. }
+      | Term::Open { bod: nxt, .. }
+      | Term::Let { nxt, .. }
+      | Term::Ask { nxt, .. }
+      | Term::Use { nxt, .. } => nxt.expand_ref_return(book, seen, globals_count),
+      Term::Var { .. }
+      | Term::Link { .. }
+      | Term::App { .. }
+      | Term::Num { .. }
+      | Term::Nat { .. }
+      | Term::Str { .. }
+      | Term::Oper { .. }
+      | Term::Mat { .. }
+      | Term::Swt { .. }
+      | Term::Fold { .. }
+      | Term::Bend { .. }
+      | Term::Era
+      | Term::Err => {}
+    })
+  }
+}
+
+impl Term {
+  /// Since expanded functions can contain unscoped variables, and
+  /// unscoped variable names must be unique, we need to rename them
+  /// to avoid conflicts.
+  fn rename_unscoped(&mut self, unscoped_count: &mut usize, unscoped_map: &mut HashMap<Name, Name>) {
+    match self {
+      Term::Let { pat, .. } | Term::Lam { pat, .. } => pat.rename_unscoped(unscoped_count, unscoped_map),
+      Term::Link { nam } => rename_unscoped(nam, unscoped_count, unscoped_map),
+      _ => {
+        // Isn't an unscoped bind or use, do nothing, just recurse.
+      }
+    }
+    for child in self.children_mut() {
+      child.rename_unscoped(unscoped_count, unscoped_map);
+    }
+  }
+}
+
+impl Pattern {
+  fn rename_unscoped(&mut self, unscoped_count: &mut usize, unscoped_map: &mut HashMap<Name, Name>) {
+    maybe_grow(|| {
+      match self {
+        Pattern::Chn(nam) => rename_unscoped(nam, unscoped_count, unscoped_map),
+        _ => {
+          // Pattern isn't an unscoped bind, just recurse.
+        }
+      }
+      for child in self.children_mut() {
+        child.rename_unscoped(unscoped_count, unscoped_map);
+      }
+    })
+  }
+}
+
+/// Generates a new name for an unscoped variable.
+fn rename_unscoped(nam: &mut Name, unscoped_count: &mut usize, unscoped_map: &mut HashMap<Name, Name>) {
+  if let Some(new_nam) = unscoped_map.get(nam) {
+    *nam = new_nam.clone();
+  } else {
+    let new_nam = Name::new(format!("{nam}%{}", unscoped_count));
+    unscoped_map.insert(nam.clone(), new_nam.clone());
+    *unscoped_count += 1;
+    *nam = new_nam;
+  }
+}

src/fun/transform/float_combinators.rs

@@ -33,6 +33,15 @@ impl Book {
     let mut ctx = FloatCombinatorsCtx::new(&book, max_size);
 
     for (def_name, def) in self.defs.iter_mut() {
+      // Don't float combinators in the main entrypoint.
+      // This avoids making programs unexpectedly too lazy,
+      // returning just a reference without executing anything.
+      if let Some(main) = self.entrypoint.as_ref() {
+        if def_name == main {
+          continue;
+        }
+      }
+
       let builtin = def.builtin;
       let body = &mut def.rule_mut().body;
       ctx.reset();

src/fun/transform/mod.rs

@@ -10,6 +10,7 @@ pub mod desugar_with_blocks;
 pub mod encode_adts;
 pub mod encode_match_terms;
 pub mod expand_generated;
+pub mod expand_main;
 pub mod fix_match_defs;
 pub mod fix_match_terms;
 pub mod float_combinators;

src/lib.rs

@@ -148,6 +148,8 @@ pub fn desugar_book(
     ctx.book.merge_definitions();
   }
 
+  ctx.book.expand_main();
+
   ctx.book.make_var_names_unique();
 
   if !ctx.info.has_errors() {

tests/golden_tests/run_file/unused_main_var.bend

@@ -0,0 +1,15 @@
+# Despite having an unused variable in main, `map(tt,l)` should not be extracted into a new definition.
+def map(fn, list):
+  fold list:
+    case List/Cons:
+      return List/Cons(fn(list.head), list.tail)
+    case List/Nil:
+      return []
+
+def tt(x):
+  return x*2
+
+def main():
+  l = [5,6,7,8]
+  k = [1,2,3,4]
+  return map(tt,l)

tests/snapshots/cli__compile_all.bend.snap

@@ -10,8 +10,6 @@ input_file: tests/golden_tests/cli/compile_all.bend
 
 @Pair/Pair = (a (b ((0 (a (b c))) c)))
 
-@main = b
-  & @Pair.get ~ (@main__C0 (a b))
-  & @Pair/Pair ~ (40 (2 a))
-
-@main__C0 = ($([+] $(a b)) (a b))
+@main = d
+  & @Pair.get ~ (($([+] $(a b)) (a b)) (c d))
+  & @Pair/Pair ~ (40 (2 c))

tests/snapshots/cli__desugar_bool_scott.bend.snap

@@ -2,7 +2,7 @@
 source: tests/golden_tests.rs
 input_file: tests/golden_tests/cli/desugar_bool_scott.bend
 ---
-(main) = Boolean/True
+(main) = λa λ* a
 
 (Boolean/True) = λa λ* a
 

tests/snapshots/cli__desugar_float_combinators.bend.snap

@@ -8,10 +8,6 @@ input_file: tests/golden_tests/cli/desugar_float_combinators.bend
 
 (get) = λa (a get__C0 0)
 
-(main) = (get main__C1)
+(main) = (get (S (S Z)))
 
 (get__C0) = λa (+ a 1)
-
-(main__C0) = (S Z)
-
-(main__C1) = (S main__C0)

tests/snapshots/cli__desugar_merge.bend.snap

@@ -8,4 +8,4 @@ input_file: tests/golden_tests/cli/desugar_merge.bend
 
 (F__M_Z) = λ* λa a
 
-(main) = F__M_Z
+(main) = λ* λa a

tests/snapshots/cli__no_check_net_size.bend.snap

@@ -27,15 +27,6 @@ input_file: tests/golden_tests/cli/no_check_net_size.bend
   &[email protected] ~ (a (b c))
   &[email protected] ~ (d (e f))
 
-@Main__C0 = a
-  & @Gen ~ (4 a)
-
-@Main__C1 = a
-  & @Reverse ~ (@Main__C0 a)
-
-@Main__C2 = a
-  & @Sort ~ (@Main__C1 a)
-
 @Map_/Both = (a (b ((@Map_/Both/tag (a (b c))) c)))
 
 @Map_/Both/tag = 2
@@ -173,5 +164,8 @@ input_file: tests/golden_tests/cli/no_check_net_size.bend
 
 @ToMap__C3 = (?((@Map_/Free @ToMap__C2) a) a)
 
-@main = a
-  & @Sum ~ (@Main__C2 a)
+@main = d
+  & @Sum ~ (c d)
+  & @Sort ~ (b c)
+  & @Reverse ~ (a b)
+  & @Gen ~ (4 a)

tests/snapshots/cli__run_pretty.bend.snap

@@ -4,9 +4,9 @@ input_file: tests/golden_tests/cli/run_pretty.bend
 ---
 Result:
 λa switch a = a {
-  0: λa switch a {
+  0: λb switch b = b {
       0: "ba";
-      _: λ* "ta";
+      _: "ta";
     };
   _: λ* "ata";
 }

tests/snapshots/compile_file__addition.bend.snap

@@ -2,8 +2,6 @@
 source: tests/golden_tests.rs
 input_file: tests/golden_tests/compile_file/addition.bend
 ---
-@main = a
-  & @main__C0 ~ (8 a)
-
-@main__C0 = (a b)
+@main = c
+  & (a b) ~ (8 c)
   & $(1 $([+] $(a b))) ~ [+1]

tests/snapshots/compile_file__church_one.bend.snap

@@ -2,7 +2,5 @@
 source: tests/golden_tests.rs
 input_file: tests/golden_tests/compile_file/church_one.bend
 ---
-@main = b
-  & @main__C0 ~ ((* (a a)) b)
-
-@main__C0 = ((a (b c)) ({(c d) a} (b d)))
+@main = f
+  & ((a (b c)) ({(c d) a} (b d))) ~ ((* (e e)) f)

tests/snapshots/compile_file__elif.bend.snap

@@ -2,17 +2,8 @@
 source: tests/golden_tests.rs
 input_file: tests/golden_tests/compile_file/elif.bend
 ---
-@main = d
-  & @main__C3 ~ (a (b (c d)))
-  & $(1 a) ~ [<2]
-  & $(2 b) ~ [>3]
-  & $(2 c) ~ [=2]
-
-@main__C0 = (?((0 (* 2)) a) a)
-
-@main__C1 = (a (?((@main__C0 (* (* 3))) (a b)) b))
-
-@main__C2 = (a (b (?((@main__C1 (* (* (* 4)))) (a (b c))) c)))
-
-@main__C3 = a
-  & $(2 ?((@main__C2 (* (* (* (* 1))))) a)) ~ [=1]
+@main = j
+  & $(2 ?(((d (e (?(((b (?(((?((0 (* 2)) a) a) (* (* 3))) (b c)) c)) (* (* (* 4)))) (d (e f))) f))) (* (* (* (* 1))))) (g (h (i j))))) ~ [=1]
+  & $(1 g) ~ [<2]
+  & $(2 h) ~ [>3]
+  & $(2 i) ~ [=2]

tests/snapshots/compile_file__fst_snd.bend.snap

@@ -2,12 +2,6 @@
 source: tests/golden_tests.rs
 input_file: tests/golden_tests/compile_file/fst_snd.bend
 ---
-@main = a
-  & @main__C2 ~ (@main__C1 a)
-
-@main__C0 = ((a *) a)
-
-@main__C1 = a
-  & @main__C0 ~ (((1 3) 2) a)
-
-@main__C2 = ((* a) a)
+@main = d
+  & ((* a) a) ~ (c d)
+  & ((b *) b) ~ (((1 3) 2) c)

tests/snapshots/compile_file__lets.bend.snap

@@ -2,11 +2,9 @@
 source: tests/golden_tests.rs
 input_file: tests/golden_tests/compile_file/lets.bend
 ---
-@main = @main__C0
-  & (a a) ~ *
-
-@main__C0 = n
-  & (a a) ~ (f (i (m n)))
-  & (d d) ~ {(e f) e}
-  & (c c) ~ {(g (h i)) {g h}}
-  & (b b) ~ {(j (k (l m))) {j {k l}}}
+@main = o
+  & (a a) ~ (g (j (n o)))
+  & (e e) ~ *
+  & (d d) ~ {(f g) f}
+  & (c c) ~ {(h (i j)) {h i}}
+  & (b b) ~ {(k (l (m n))) {k {l m}}}

tests/snapshots/compile_file__ref_to_ref.bend.snap

@@ -18,4 +18,4 @@ input_file: tests/golden_tests/compile_file/ref_to_ref.bend
 
 @C1 = @B4
 
-@main = ((@A1 @B1) @C1)
+@main = ((1 2) 2)

tests/snapshots/compile_file__switch_in_switch_arg.bend.snap

@@ -2,6 +2,4 @@
 source: tests/golden_tests.rs
 input_file: tests/golden_tests/compile_file/switch_in_switch_arg.bend
 ---
-@main = (?((0 (a a)) ?((0 @main__C0) b)) b)
-
-@main__C0 = ($([+1] a) a)
+@main = (?((0 (a a)) ?((0 ($([+1] b) b)) c)) c)

tests/snapshots/compile_file__vicious_circles.bend.snap

@@ -7,3 +7,5 @@ input_file: tests/golden_tests/compile_file/vicious_circles.bend
   Found term that compiles into an inet with a vicious cycle
 In compiled inet 'dup_self':
   Found term that compiles into an inet with a vicious cycle
+In compiled inet 'main':
+  Found term that compiles into an inet with a vicious cycle

tests/snapshots/compile_file_o_all__addition.bend.snap

@@ -2,8 +2,6 @@
 source: tests/golden_tests.rs
 input_file: tests/golden_tests/compile_file_o_all/addition.bend
 ---
-@main = a
-  & @main__C0 ~ (8 a)
-
-@main__C0 = (a b)
+@main = c
+  & (a b) ~ (8 c)
   & $(1 $([+] $(a b))) ~ [+1]