Propagate unsafety to definitions that depend on unsafe definitions

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"

src/ast.rs

@@ -1,7 +1,8 @@
 use TSPL::{new_parser, Parser};
 use highlight_error::highlight_error;
 use crate::hvm;
-use std::{collections::BTreeMap, fmt::{Debug, Display}};
+use std::fmt::{Debug, Display};
+use std::collections::{btree_map::Entry, BTreeMap, BTreeSet};
 
 // Types
 // -----
@@ -488,6 +489,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 {
@@ -523,6 +537,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 +550,86 @@ 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);
     return book;
   }
+
+  /// Propagate unsafe definitions to those that reference them.
+  ///
+  /// 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();
+    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 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
+  /// traverses each definition's trees entirely once.
+  ///
+  /// 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() {
+      result.insert(name.as_str(), BTreeSet::new());
+    }
+
+    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);
+          },
+        }
+      }
+    };
+
+    for (name, net) in self.defs.iter() {
+      process(&net.root, name);
+      for (_, r1, r2) in net.rbag.iter() {
+        process(r1, name);
+        process(r2, name);
+      }
+    }
+    result
+  }
 }

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

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");

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

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.