From e6c812302d922686bcdc3e5cf86660e888eaac49 Mon Sep 17 00:00:00 2001
From: Eduardo Sandalo Porto <edu@sandalo.dev>
Date: Thu, 6 Jun 2024 00:24:26 -0300
Subject: [PATCH 01/11] Add first version of safety propagation

TODO:
- Document how the algorithm works along with complexities
- From the original bug that led to this, `ok.hvm` now works as expectes but `bad.hvm` still doesn't...
  - Investigate
---
 src/ast.rs  | 90 +++++++++++++++++++++++++++++++++++++++++++++++++++--
 src/main.rs | 10 +++---
 2 files changed, 92 insertions(+), 8 deletions(-)

diff --git a/src/ast.rs b/src/ast.rs
index 2d6c0370..493e3c1a 100644
--- a/src/ast.rs
+++ b/src/ast.rs
@@ -1,7 +1,7 @@
 use TSPL::{new_parser, Parser};
 use highlight_error::highlight_error;
 use crate::hvm;
-use std::{collections::BTreeMap, fmt::{Debug, Display}};
+use std::{collections::{btree_map::Entry, BTreeMap, BTreeSet}, fmt::{Debug, Display}};
 
 // Types
 // -----
@@ -488,6 +488,19 @@ impl Tree {
       },
     }
   }
+
+  pub fn direct_dependencies<'name>(&'name self) -> BTreeSet<&'name str> {
+    match self {
+      Tree::Ref { nam } => BTreeSet::from([nam.as_str()]),
+      Tree::Con { fst, snd } => &fst.direct_dependencies() | &snd.direct_dependencies(),
+      Tree::Dup { fst, snd } => &fst.direct_dependencies() | &snd.direct_dependencies(),
+      Tree::Opr { fst, snd } => &fst.direct_dependencies() | &snd.direct_dependencies(),
+      Tree::Swi { fst, snd } => &fst.direct_dependencies() | &snd.direct_dependencies(),
+      Tree::Num { val } => BTreeSet::new(),
+      Tree::Var { nam } => BTreeSet::new(),
+      Tree::Era => BTreeSet::new(),
+    }
+  }
 }
 
 impl Net {
@@ -511,7 +524,7 @@ impl Book {
     CoreParser::new(code).parse_book()
   }
 
-  pub fn build(&self) -> hvm::Book {
+  pub fn build(&self) -> (hvm::Book, BTreeMap<String, usize>) {
     let mut name_to_fid = BTreeMap::new();
     let mut fid_to_name = BTreeMap::new();
     fid_to_name.insert(0, "main".to_string());
@@ -523,6 +536,7 @@ impl Book {
       }
     }
     let mut book = hvm::Book { defs: Vec::new() };
+    let mut lookup = BTreeMap::new();
     for (fid, name) in &fid_to_name {
       let ast_def = self.defs.get(name).expect("missing `@main` definition");
       let mut def = hvm::Def {
@@ -535,7 +549,77 @@ impl Book {
       };
       ast_def.build(&mut def, &name_to_fid, &mut BTreeMap::new());
       book.defs.push(def);
+      lookup.insert(name.clone(), book.defs.len() - 1);
     }
-    return book;
+    self.propagate_safety(&mut book, &lookup);
+    return (book, lookup);
+  }
+
+  /// When calling this function, it is expected that definitions that are directly
+  /// unsafe already know so. We then propagate the unsafety to all definitions that
+  /// have references to unsafe definitions.
+  fn propagate_safety(&self, compiled_book: &mut hvm::Book, lookup: &BTreeMap<String, usize>) {
+    let rev_dependencies = self.direct_dependencies_reversed();
+    let mut visited: BTreeSet<&str> = BTreeSet::new();
+    let mut stack: Vec<&str> = Vec::new();
+
+    for (name, _) in self.defs.iter() {
+      let def = &compiled_book.defs[lookup[name]];
+      if !def.safe {
+        stack.push(&name);
+      }
+    }
+
+    while let Some(curr) = stack.pop() {
+      if visited.contains(curr) {
+        continue;
+      }
+      visited.insert(curr);
+
+      let def = &mut compiled_book.defs[lookup[curr]];
+      def.safe = false;
+
+      for &next in rev_dependencies[curr].iter() {
+        stack.push(next);
+      }
+    }
+  }
+
+  /// Calculates the dependencies of each definition but stores them reversed,
+  /// that is, if definition `A` requires `B`, `B: A` is in the return set.
+  /// This is used to propagate unsafe definitions to others that depend on it.
+  /// 
+  /// TODO: Verify
+  /// Complexity: O(n*h + m)
+  /// - `n` is the number of definitions in the book
+  /// - `m` is the number of direct references in each definition
+  /// - `h` is the accumulated height of each net's trees
+  fn direct_dependencies_reversed<'name>(&'name self) -> BTreeMap<&'name str, BTreeSet<&'name str>> {
+    let mut result = BTreeMap::new();
+    for (name, _) in self.defs.iter() {
+      result.insert(name.as_str(), BTreeSet::new());
+    }
+
+    let process = |tree: &'name Tree, name: &'name str, result: &mut BTreeMap<&'name str, BTreeSet<&'name str>>| {
+      for dependency in tree.direct_dependencies() {
+        match result.entry(dependency) {
+          Entry::Vacant(_) => panic!("global definition depends on undeclared reference"),
+          Entry::Occupied(mut entry) => {
+            // dependency => name
+            entry.get_mut().insert(name);
+          },
+        }
+      }
+    };
+
+    for (name, net) in self.defs.iter() {
+      process(&net.root, name, &mut result);
+      for (_, r1, r2) in net.rbag.iter() {
+        process(r1, name, &mut result);
+        process(r2, name, &mut result);
+      }
+    }
+
+    result
   }
 }
diff --git a/src/main.rs b/src/main.rs
index 3232a63e..61ea27d6 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -68,13 +68,13 @@ fn main() {
     Some(("run", sub_matches)) => {
       let file = sub_matches.get_one::<String>("file").expect("required");
       let code = fs::read_to_string(file).expect("Unable to read file");
-      let book = ast::Book::parse(&code).unwrap_or_else(|er| panic!("{}",er)).build();
+      let book = ast::Book::parse(&code).unwrap_or_else(|er| panic!("{}",er)).build().0;
       run(&book);
     }
     Some(("run-c", sub_matches)) => {
       let file = sub_matches.get_one::<String>("file").expect("required");
       let code = fs::read_to_string(file).expect("Unable to read file");
-      let book = ast::Book::parse(&code).unwrap_or_else(|er| panic!("{}",er)).build();
+      let book = ast::Book::parse(&code).unwrap_or_else(|er| panic!("{}",er)).build().0;
       let mut data : Vec<u8> = Vec::new();
       book.to_buffer(&mut data);
       //println!("{:?}", data);
@@ -89,7 +89,7 @@ fn main() {
     Some(("run-cu", sub_matches)) => {
       let file = sub_matches.get_one::<String>("file").expect("required");
       let code = fs::read_to_string(file).expect("Unable to read file");
-      let book = ast::Book::parse(&code).unwrap_or_else(|er| panic!("{}",er)).build();
+      let book = ast::Book::parse(&code).unwrap_or_else(|er| panic!("{}",er)).build().0;
       let mut data : Vec<u8> = Vec::new();
       book.to_buffer(&mut data);
       let run_io = sub_matches.get_flag("io");
@@ -104,7 +104,7 @@ fn main() {
       // Reads book from file
       let file = sub_matches.get_one::<String>("file").expect("required");
       let code = fs::read_to_string(file).expect("Unable to read file");
-      let book = ast::Book::parse(&code).unwrap_or_else(|er| panic!("{}",er)).build();
+      let book = ast::Book::parse(&code).unwrap_or_else(|er| panic!("{}",er)).build().0;
 
       // Gets optimal core count
       let cores = num_cpus::get();
@@ -135,7 +135,7 @@ fn main() {
       // Reads book from file
       let file = sub_matches.get_one::<String>("file").expect("required");
       let code = fs::read_to_string(file).expect("Unable to read file");
-      let book = ast::Book::parse(&code).unwrap_or_else(|er| panic!("{}",er)).build();
+      let book = ast::Book::parse(&code).unwrap_or_else(|er| panic!("{}",er)).build().0;
 
       // Generates the interpreted book
       let mut book_buf : Vec<u8> = Vec::new();

From bc870fc549f845813953bb8c43747fb8618bc3e9 Mon Sep 17 00:00:00 2001
From: Eduardo Sandalo Porto <edu@sandalo.dev>
Date: Thu, 6 Jun 2024 13:54:40 -0300
Subject: [PATCH 02/11] Clean up and better comments

---
 src/ast.rs | 24 ++++++++++++++----------
 1 file changed, 14 insertions(+), 10 deletions(-)

diff --git a/src/ast.rs b/src/ast.rs
index 493e3c1a..a1f5ad34 100644
--- a/src/ast.rs
+++ b/src/ast.rs
@@ -555,9 +555,10 @@ impl Book {
     return (book, lookup);
   }
 
+  /// Propagate unsafe definitions to those that reference them.
+  ///
   /// When calling this function, it is expected that definitions that are directly
-  /// unsafe already know so. We then propagate the unsafety to all definitions that
-  /// have references to unsafe definitions.
+  /// unsafe are already marked as such in the `compiled_book`.
   fn propagate_safety(&self, compiled_book: &mut hvm::Book, lookup: &BTreeMap<String, usize>) {
     let rev_dependencies = self.direct_dependencies_reversed();
     let mut visited: BTreeSet<&str> = BTreeSet::new();
@@ -586,10 +587,14 @@ impl Book {
   }
 
   /// Calculates the dependencies of each definition but stores them reversed,
-  /// that is, if definition `A` requires `B`, `B: A` is in the return set.
-  /// This is used to propagate unsafe definitions to others that depend on it.
+  /// that is, if definition `A` requires `B`, `B: A` is in the return map.
+  /// This is used to propagate unsafe definitions to others that depend on them.
   /// 
-  /// TODO: Verify
+  /// This solution has linear complexity on the number of definitions in the
+  /// book and the number of direct references in each definition, but it also
+  /// traverses each definition's trees entirely once. Assuming the tree traversals
+  /// are O(h), which they're not with `BTreeSet`s, we have:
+  ///
   /// Complexity: O(n*h + m)
   /// - `n` is the number of definitions in the book
   /// - `m` is the number of direct references in each definition
@@ -600,7 +605,7 @@ impl Book {
       result.insert(name.as_str(), BTreeSet::new());
     }
 
-    let process = |tree: &'name Tree, name: &'name str, result: &mut BTreeMap<&'name str, BTreeSet<&'name str>>| {
+    let mut process = |tree: &'name Tree, name: &'name str| {
       for dependency in tree.direct_dependencies() {
         match result.entry(dependency) {
           Entry::Vacant(_) => panic!("global definition depends on undeclared reference"),
@@ -613,13 +618,12 @@ impl Book {
     };
 
     for (name, net) in self.defs.iter() {
-      process(&net.root, name, &mut result);
+      process(&net.root, name);
       for (_, r1, r2) in net.rbag.iter() {
-        process(r1, name, &mut result);
-        process(r2, name, &mut result);
+        process(r1, name);
+        process(r2, name);
       }
     }
-
     result
   }
 }

From 55b6b4c0d85a7a81372c748c336d09d0a1d439aa Mon Sep 17 00:00:00 2001
From: Eduardo Sandalo Porto <edu@sandalo.dev>
Date: Thu, 6 Jun 2024 14:15:03 -0300
Subject: [PATCH 03/11] Remove unnecessary lookup table return from
 `Book::build`

---
 src/ast.rs  |  4 ++--
 src/main.rs | 10 +++++-----
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/src/ast.rs b/src/ast.rs
index a1f5ad34..72f82cca 100644
--- a/src/ast.rs
+++ b/src/ast.rs
@@ -524,7 +524,7 @@ impl Book {
     CoreParser::new(code).parse_book()
   }
 
-  pub fn build(&self) -> (hvm::Book, BTreeMap<String, usize>) {
+  pub fn build(&self) -> hvm::Book {
     let mut name_to_fid = BTreeMap::new();
     let mut fid_to_name = BTreeMap::new();
     fid_to_name.insert(0, "main".to_string());
@@ -552,7 +552,7 @@ impl Book {
       lookup.insert(name.clone(), book.defs.len() - 1);
     }
     self.propagate_safety(&mut book, &lookup);
-    return (book, lookup);
+    return book;
   }
 
   /// Propagate unsafe definitions to those that reference them.
diff --git a/src/main.rs b/src/main.rs
index 03529e0f..dee219ab 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -68,13 +68,13 @@ fn main() {
     Some(("run", sub_matches)) => {
       let file = sub_matches.get_one::<String>("file").expect("required");
       let code = fs::read_to_string(file).expect("Unable to read file");
-      let book = ast::Book::parse(&code).unwrap_or_else(|er| panic!("{}",er)).build().0;
+      let book = ast::Book::parse(&code).unwrap_or_else(|er| panic!("{}",er)).build();
       run(&book);
     }
     Some(("run-c", sub_matches)) => {
       let file = sub_matches.get_one::<String>("file").expect("required");
       let code = fs::read_to_string(file).expect("Unable to read file");
-      let book = ast::Book::parse(&code).unwrap_or_else(|er| panic!("{}",er)).build().0;
+      let book = ast::Book::parse(&code).unwrap_or_else(|er| panic!("{}",er)).build();
       let mut data : Vec<u8> = Vec::new();
       book.to_buffer(&mut data);
       #[cfg(feature = "c")]
@@ -87,7 +87,7 @@ fn main() {
     Some(("run-cu", sub_matches)) => {
       let file = sub_matches.get_one::<String>("file").expect("required");
       let code = fs::read_to_string(file).expect("Unable to read file");
-      let book = ast::Book::parse(&code).unwrap_or_else(|er| panic!("{}",er)).build().0;
+      let book = ast::Book::parse(&code).unwrap_or_else(|er| panic!("{}",er)).build();
       let mut data : Vec<u8> = Vec::new();
       book.to_buffer(&mut data);
       #[cfg(feature = "cuda")]
@@ -101,7 +101,7 @@ fn main() {
       // Reads book from file
       let file = sub_matches.get_one::<String>("file").expect("required");
       let code = fs::read_to_string(file).expect("Unable to read file");
-      let book = ast::Book::parse(&code).unwrap_or_else(|er| panic!("{}",er)).build().0;
+      let book = ast::Book::parse(&code).unwrap_or_else(|er| panic!("{}",er)).build();
 
       // Gets optimal core count
       let cores = num_cpus::get();
@@ -129,7 +129,7 @@ fn main() {
       // Reads book from file
       let file = sub_matches.get_one::<String>("file").expect("required");
       let code = fs::read_to_string(file).expect("Unable to read file");
-      let book = ast::Book::parse(&code).unwrap_or_else(|er| panic!("{}",er)).build().0;
+      let book = ast::Book::parse(&code).unwrap_or_else(|er| panic!("{}",er)).build();
 
       // Generates the interpreted book
       let mut book_buf : Vec<u8> = Vec::new();

From 77f8b7314b2da2d265f293522c356beb86d4c318 Mon Sep 17 00:00:00 2001
From: Eduardo Sandalo Porto <edu@sandalo.dev>
Date: Fri, 7 Jun 2024 15:56:37 -0300
Subject: [PATCH 04/11] Tidy up before pull request

---
 src/ast.rs                               | 20 +++++++++++++-------
 tests/snapshots/run__file@empty.hvm.snap |  2 +-
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/src/ast.rs b/src/ast.rs
index 72f82cca..25f3c041 100644
--- a/src/ast.rs
+++ b/src/ast.rs
@@ -1,7 +1,8 @@
 use TSPL::{new_parser, Parser};
 use highlight_error::highlight_error;
 use crate::hvm;
-use std::{collections::{btree_map::Entry, BTreeMap, BTreeSet}, fmt::{Debug, Display}};
+use std::fmt::{Debug, Display};
+use std::collections::{btree_map::Entry, BTreeMap, BTreeSet};
 
 // Types
 // -----
@@ -559,6 +560,12 @@ impl Book {
   ///
   /// When calling this function, it is expected that definitions that are directly
   /// unsafe are already marked as such in the `compiled_book`.
+  /// 
+  /// This does not completely solve the cloning safety in HVM. It only stops invalid
+  /// **global** definitions from being cloned, but local unsafe code can still be
+  /// cloned and can generate seemingly unexpected results, such as placing eraser
+  /// nodes in weird places. See HVM issue [#362](https://github.com/HigherOrderCO/HVM/issues/362)
+  /// for an example.
   fn propagate_safety(&self, compiled_book: &mut hvm::Book, lookup: &BTreeMap<String, usize>) {
     let rev_dependencies = self.direct_dependencies_reversed();
     let mut visited: BTreeSet<&str> = BTreeSet::new();
@@ -592,13 +599,12 @@ impl Book {
   /// 
   /// This solution has linear complexity on the number of definitions in the
   /// book and the number of direct references in each definition, but it also
-  /// traverses each definition's trees entirely once. Assuming the tree traversals
-  /// are O(h), which they're not with `BTreeSet`s, we have:
+  /// traverses each definition's trees entirely once.
   ///
-  /// Complexity: O(n*h + m)
-  /// - `n` is the number of definitions in the book
-  /// - `m` is the number of direct references in each definition
-  /// - `h` is the accumulated height of each net's trees
+  /// Complexity: O(d*t + r)
+  /// - `d` is the number of definitions in the book
+  /// - `r` is the number of direct references in each definition
+  /// - `t` is the number of nodes in each tree
   fn direct_dependencies_reversed<'name>(&'name self) -> BTreeMap<&'name str, BTreeSet<&'name str>> {
     let mut result = BTreeMap::new();
     for (name, _) in self.defs.iter() {
diff --git a/tests/snapshots/run__file@empty.hvm.snap b/tests/snapshots/run__file@empty.hvm.snap
index 953bfae9..37fe8b17 100644
--- a/tests/snapshots/run__file@empty.hvm.snap
+++ b/tests/snapshots/run__file@empty.hvm.snap
@@ -4,6 +4,6 @@ expression: rust_output
 input_file: tests/programs/empty.hvm
 ---
 exit status: 101
-thread 'main' panicked at src/ast.rs:527:41:
+thread 'main' panicked at src/ast.rs:542:41:
 missing `@main` definition
 note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace

From 2bfbd994a426b577343970f6603e27243bfa4c38 Mon Sep 17 00:00:00 2001
From: Eduardo Sandalo Porto <edu@sandalo.dev>
Date: Fri, 7 Jun 2024 15:58:09 -0300
Subject: [PATCH 05/11] Bump version

---
 Cargo.lock | 2 +-
 Cargo.toml | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 4583af30..70d59e85 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -163,7 +163,7 @@ checksum = "809e18805660d7b6b2e2b9f316a5099521b5998d5cba4dda11b5157a21aaef03"
 
 [[package]]
 name = "hvm"
-version = "2.0.19"
+version = "2.0.20"
 dependencies = [
  "TSPL",
  "cc",
diff --git a/Cargo.toml b/Cargo.toml
index 11e5ef1f..8f055f27 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -2,7 +2,7 @@
 name = "hvm"
 description = "A massively parallel, optimal functional runtime in Rust."
 license = "Apache-2.0"
-version = "2.0.19"
+version = "2.0.20"
 edition = "2021"
 rust-version = "1.74"
 build = "build.rs"

From 607426deaae1457ea050f05870385d95dd18c8cc Mon Sep 17 00:00:00 2001
From: Eduardo Sandalo Porto <edu@sandalo.dev>
Date: Fri, 7 Jun 2024 16:10:28 -0300
Subject: [PATCH 06/11] Add `@test-rust-only` option

---
 tests/run.rs | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tests/run.rs b/tests/run.rs
index d03f6a50..d19e31d8 100644
--- a/tests/run.rs
+++ b/tests/run.rs
@@ -45,6 +45,11 @@ fn test_file(path: &Path) {
   let rust_output = execute_hvm(&["run".as_ref(), path.as_os_str()], false).unwrap();
   assert_snapshot!(rust_output);
 
+  if contents.contains("@test-rust-only = 1") {
+    println!("only testing rust implementation for {path:?}");
+    return;
+  }
+
   println!("  testing {path:?}, C...");
   let c_output = execute_hvm(&["run-c".as_ref(), path.as_os_str()], false).unwrap();
   assert_eq!(c_output, rust_output, "{path:?}: C output does not match rust output");

From 0f44d07aad436146963b146c420bd68ef14c23f1 Mon Sep 17 00:00:00 2001
From: Eduardo Sandalo Porto <edu@sandalo.dev>
Date: Fri, 7 Jun 2024 16:10:49 -0300
Subject: [PATCH 07/11] Add test for safety checking

---
 tests/programs/safety-check.hvm               | 31 +++++++++++++++++++
 .../snapshots/run__file@safety-check.hvm.snap |  6 ++++
 2 files changed, 37 insertions(+)
 create mode 100644 tests/programs/safety-check.hvm
 create mode 100644 tests/snapshots/run__file@safety-check.hvm.snap

diff --git a/tests/programs/safety-check.hvm b/tests/programs/safety-check.hvm
new file mode 100644
index 00000000..801006bf
--- /dev/null
+++ b/tests/programs/safety-check.hvm
@@ -0,0 +1,31 @@
+@List/Cons = (a (b ((@List/Cons/tag (a (b c))) c)))
+
+@List/Cons/tag = 1
+
+@List/Nil = ((@List/Nil/tag a) a)
+
+@List/Nil/tag = 0
+
+@id = (a a)
+
+@list = c
+  & @List/Cons ~ (1 (b c))
+  & @List/Cons ~ (2 (@List/Nil b))
+
+@main = b
+  & @map ~ (@main__C0 (a b))
+  & @List/Cons ~ (@id (@List/Nil a))
+
+@main__C0 = (a b)
+  & @map ~ (a (@list b))
+
+@map = (a ((@map__C1 (a b)) b))
+
+@map__C0 = (* (a (d ({(a b) c} f))))
+  & @List/Cons ~ (b (e f))
+  & @map ~ (c (d e))
+
+@map__C1 = (?(((* @List/Nil) @map__C0) a) a)
+
+// Test flags
+@test-rust-only = 1
\ No newline at end of file
diff --git a/tests/snapshots/run__file@safety-check.hvm.snap b/tests/snapshots/run__file@safety-check.hvm.snap
new file mode 100644
index 00000000..fbc70b78
--- /dev/null
+++ b/tests/snapshots/run__file@safety-check.hvm.snap
@@ -0,0 +1,6 @@
+---
+source: tests/run.rs
+expression: rust_output
+input_file: tests/programs/safety-check.hvm
+---
+ERROR: attempt to clone a non-affine global reference.

From 0f675a8bea6bb5388a11b2baac8f48209d0aaa67 Mon Sep 17 00:00:00 2001
From: Eduardo Sandalo Porto <edu@sandalo.dev>
Date: Mon, 10 Jun 2024 15:03:40 -0300
Subject: [PATCH 08/11] Resolve suggestions

---
 src/ast.rs                               | 54 ++++++++++++------------
 tests/snapshots/run__file@empty.hvm.snap |  2 +-
 2 files changed, 29 insertions(+), 27 deletions(-)

diff --git a/src/ast.rs b/src/ast.rs
index 25f3c041..8b22f36e 100644
--- a/src/ast.rs
+++ b/src/ast.rs
@@ -2,7 +2,7 @@ use TSPL::{new_parser, Parser};
 use highlight_error::highlight_error;
 use crate::hvm;
 use std::fmt::{Debug, Display};
-use std::collections::{btree_map::Entry, BTreeMap, BTreeSet};
+use std::collections::{BTreeMap, BTreeSet};
 
 // Types
 // -----
@@ -491,16 +491,22 @@ impl Tree {
   }
 
   pub fn direct_dependencies<'name>(&'name self) -> BTreeSet<&'name str> {
-    match self {
-      Tree::Ref { nam } => BTreeSet::from([nam.as_str()]),
-      Tree::Con { fst, snd } => &fst.direct_dependencies() | &snd.direct_dependencies(),
-      Tree::Dup { fst, snd } => &fst.direct_dependencies() | &snd.direct_dependencies(),
-      Tree::Opr { fst, snd } => &fst.direct_dependencies() | &snd.direct_dependencies(),
-      Tree::Swi { fst, snd } => &fst.direct_dependencies() | &snd.direct_dependencies(),
-      Tree::Num { val } => BTreeSet::new(),
-      Tree::Var { nam } => BTreeSet::new(),
-      Tree::Era => BTreeSet::new(),
+    let mut stack: Vec<&Tree> = vec![self];
+    let mut acc: BTreeSet<&'name str> = BTreeSet::new();
+    
+    while let Some(curr) = stack.pop() {
+      match curr {
+        Tree::Ref { nam } => { acc.insert(nam); },
+        Tree::Con { fst, snd } => { stack.push(fst); stack.push(snd); },
+        Tree::Dup { fst, snd } => { stack.push(fst); stack.push(snd); },
+        Tree::Opr { fst, snd } => { stack.push(fst); stack.push(snd); },
+        Tree::Swi { fst, snd } => { stack.push(fst); stack.push(snd); },
+        Tree::Num { val } => {},
+        Tree::Var { nam } => {},
+        Tree::Era => {},
+      };
     }
+    acc
   }
 }
 
@@ -537,7 +543,6 @@ impl Book {
       }
     }
     let mut book = hvm::Book { defs: Vec::new() };
-    let mut lookup = BTreeMap::new();
     for (fid, name) in &fid_to_name {
       let ast_def = self.defs.get(name).expect("missing `@main` definition");
       let mut def = hvm::Def {
@@ -550,9 +555,8 @@ impl Book {
       };
       ast_def.build(&mut def, &name_to_fid, &mut BTreeMap::new());
       book.defs.push(def);
-      lookup.insert(name.clone(), book.defs.len() - 1);
     }
-    self.propagate_safety(&mut book, &lookup);
+    self.propagate_safety(&mut book, &name_to_fid);
     return book;
   }
 
@@ -566,25 +570,26 @@ impl Book {
   /// cloned and can generate seemingly unexpected results, such as placing eraser
   /// nodes in weird places. See HVM issue [#362](https://github.com/HigherOrderCO/HVM/issues/362)
   /// for an example.
-  fn propagate_safety(&self, compiled_book: &mut hvm::Book, lookup: &BTreeMap<String, usize>) {
+  fn propagate_safety(&self, compiled_book: &mut hvm::Book, lookup: &BTreeMap<String, u32>) {
     let rev_dependencies = self.direct_dependencies_reversed();
-    let mut visited: BTreeSet<&str> = BTreeSet::new();
     let mut stack: Vec<&str> = Vec::new();
 
     for (name, _) in self.defs.iter() {
-      let def = &compiled_book.defs[lookup[name]];
+      let def = &mut compiled_book.defs[lookup[name] as usize];
       if !def.safe {
         stack.push(&name);
+        // Temporarily set as safe so we won't need a separate "visited" set
+        def.safe = true;
       }
     }
 
     while let Some(curr) = stack.pop() {
-      if visited.contains(curr) {
+      let def = &mut compiled_book.defs[lookup[curr] as usize];
+      if !def.safe {
+        // Already visited, skip this
         continue;
       }
-      visited.insert(curr);
 
-      let def = &mut compiled_book.defs[lookup[curr]];
       def.safe = false;
 
       for &next in rev_dependencies[curr].iter() {
@@ -613,13 +618,10 @@ impl Book {
 
     let mut process = |tree: &'name Tree, name: &'name str| {
       for dependency in tree.direct_dependencies() {
-        match result.entry(dependency) {
-          Entry::Vacant(_) => panic!("global definition depends on undeclared reference"),
-          Entry::Occupied(mut entry) => {
-            // dependency => name
-            entry.get_mut().insert(name);
-          },
-        }
+        result
+          .get_mut(dependency)
+          .expect("global definition depends on undeclared reference")
+          .insert(name);
       }
     };
 
diff --git a/tests/snapshots/run__file@empty.hvm.snap b/tests/snapshots/run__file@empty.hvm.snap
index 37fe8b17..1460dd1e 100644
--- a/tests/snapshots/run__file@empty.hvm.snap
+++ b/tests/snapshots/run__file@empty.hvm.snap
@@ -4,6 +4,6 @@ expression: rust_output
 input_file: tests/programs/empty.hvm
 ---
 exit status: 101
-thread 'main' panicked at src/ast.rs:542:41:
+thread 'main' panicked at src/ast.rs:547:41:
 missing `@main` definition
 note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace

From 92773fb11d806fa74a36055efac8e65d8c8f3f98 Mon Sep 17 00:00:00 2001
From: Eduardo Sandalo Porto <edu@sandalo.dev>
Date: Mon, 10 Jun 2024 15:08:56 -0300
Subject: [PATCH 09/11] Slightly better visited definition

---
 src/ast.rs | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/ast.rs b/src/ast.rs
index 8b22f36e..29b9dc5d 100644
--- a/src/ast.rs
+++ b/src/ast.rs
@@ -577,9 +577,9 @@ impl Book {
     for (name, _) in self.defs.iter() {
       let def = &mut compiled_book.defs[lookup[name] as usize];
       if !def.safe {
-        stack.push(&name);
-        // Temporarily set as safe so we won't need a separate "visited" set
-        def.safe = true;
+        for next in rev_dependencies[name.as_str()].iter() {
+          stack.push(next);
+        }
       }
     }
 

From e016d7c049380589284929ee0675d36149d46c35 Mon Sep 17 00:00:00 2001
From: Eduardo Sandalo Porto <edu@sandalo.dev>
Date: Mon, 10 Jun 2024 15:35:53 -0300
Subject: [PATCH 10/11] Function name change

Change "dependency reversed" to "dependent"
---
 src/ast.rs | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/ast.rs b/src/ast.rs
index 29b9dc5d..25f52038 100644
--- a/src/ast.rs
+++ b/src/ast.rs
@@ -571,7 +571,7 @@ impl Book {
   /// nodes in weird places. See HVM issue [#362](https://github.com/HigherOrderCO/HVM/issues/362)
   /// for an example.
   fn propagate_safety(&self, compiled_book: &mut hvm::Book, lookup: &BTreeMap<String, u32>) {
-    let rev_dependencies = self.direct_dependencies_reversed();
+    let rev_dependencies = self.direct_dependents();
     let mut stack: Vec<&str> = Vec::new();
 
     for (name, _) in self.defs.iter() {
@@ -598,9 +598,9 @@ impl Book {
     }
   }
 
-  /// Calculates the dependencies of each definition but stores them reversed,
-  /// that is, if definition `A` requires `B`, `B: A` is in the return map.
-  /// This is used to propagate unsafe definitions to others that depend on them.
+  /// Calculates the dependents of each definition, that is, if definition `A`
+  /// requires `B`, `B: A` is in the return map. This is used to propagate unsafe
+  /// definitions to others that depend on them.
   /// 
   /// This solution has linear complexity on the number of definitions in the
   /// book and the number of direct references in each definition, but it also
@@ -610,7 +610,7 @@ impl Book {
   /// - `d` is the number of definitions in the book
   /// - `r` is the number of direct references in each definition
   /// - `t` is the number of nodes in each tree
-  fn direct_dependencies_reversed<'name>(&'name self) -> BTreeMap<&'name str, BTreeSet<&'name str>> {
+  fn direct_dependents<'name>(&'name self) -> BTreeMap<&'name str, BTreeSet<&'name str>> {
     let mut result = BTreeMap::new();
     for (name, _) in self.defs.iter() {
       result.insert(name.as_str(), BTreeSet::new());

From 23ad6dc0901d95655726528f30593a9a9507c706 Mon Sep 17 00:00:00 2001
From: Eduardo Sandalo Porto <edu@sandalo.dev>
Date: Mon, 10 Jun 2024 17:21:05 -0300
Subject: [PATCH 11/11] Fix function name change

Co-authored-by: T6 <t6@t6.fyi>
---
 src/ast.rs | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/ast.rs b/src/ast.rs
index 25f52038..56db241c 100644
--- a/src/ast.rs
+++ b/src/ast.rs
@@ -571,13 +571,13 @@ impl Book {
   /// nodes in weird places. See HVM issue [#362](https://github.com/HigherOrderCO/HVM/issues/362)
   /// for an example.
   fn propagate_safety(&self, compiled_book: &mut hvm::Book, lookup: &BTreeMap<String, u32>) {
-    let rev_dependencies = self.direct_dependents();
+    let dependents = self.direct_dependents();
     let mut stack: Vec<&str> = Vec::new();
 
     for (name, _) in self.defs.iter() {
       let def = &mut compiled_book.defs[lookup[name] as usize];
       if !def.safe {
-        for next in rev_dependencies[name.as_str()].iter() {
+        for next in dependents[name.as_str()].iter() {
           stack.push(next);
         }
       }
@@ -592,7 +592,7 @@ impl Book {
 
       def.safe = false;
 
-      for &next in rev_dependencies[curr].iter() {
+      for &next in dependents[curr].iter() {
         stack.push(next);
       }
     }