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::{BTreeMap, BTreeSet};
 
 // Types
 // -----
@@ -488,6 +489,25 @@ impl Tree {
       },
     }
   }
+
+  pub fn direct_dependencies<'name>(&'name self) -> BTreeSet<&'name str> {
+    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
+  }
 }
 
 impl Net {
@@ -536,6 +556,82 @@ impl Book {
       ast_def.build(&mut def, &name_to_fid, &mut BTreeMap::new());
       book.defs.push(def);
     }
+    self.propagate_safety(&mut book, &name_to_fid);
     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, u32>) {
+    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 dependents[name.as_str()].iter() {
+          stack.push(next);
+        }
+      }
+    }
+
+    while let Some(curr) = stack.pop() {
+      let def = &mut compiled_book.defs[lookup[curr] as usize];
+      if !def.safe {
+        // Already visited, skip this
+        continue;
+      }
+
+      def.safe = false;
+
+      for &next in dependents[curr].iter() {
+        stack.push(next);
+      }
+    }
+  }
+
+  /// 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
+  /// 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_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());
+    }
+
+    let mut process = |tree: &'name Tree, name: &'name str| {
+      for dependency in tree.direct_dependencies() {
+        result
+          .get_mut(dependency)
+          .expect("global definition depends on undeclared reference")
+          .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:547: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.