Plumb source locations through classic CLVM evaluation

src/classic/clvm_tools/debug.rs

@@ -119,13 +119,12 @@ where
     let mut map_result: Vec<A::NodePtr> = Vec::new();
 
     for (k, v) in constants_lookup.iter() {
-        let v_export = allocator.export(v);
         let vloc = allocator.loc(v);
-        let run_result = run_program
-            .run_program(allocator.allocator(), v_export, NodePtr::NIL, None)
-            .map_err(|e| allocator.map_err(vloc.clone(), e))?;
+        let nil = allocator.new_atom(vloc.clone(), &[])?;
+        let run_result = allocator.run_clvm(run_program.clone(), v, &nil)?;
+        let run_result = allocator.export(&run_result);
 
-        let sha256 = sha256tree(allocator.allocator(), run_result.1).hex();
+        let sha256 = sha256tree(allocator.allocator(), run_result).hex();
         let sha_atom = allocator.new_atom(vloc.clone(), sha256.as_bytes())?;
         let name_atom = allocator.new_atom(vloc.clone(), &k.clone())?;
 

src/classic/clvm_tools/stages/stage_2/abstraction.rs

@@ -7,6 +7,10 @@ use clvm_rs::error::EvalErr;
 
 use crate::classic::clvm::__type_compatibility__::bi_zero;
 use crate::classic::clvm_tools::binutils::disassemble;
+use crate::classic::clvm_tools::stages::stage_0::TRunProgram;
+use crate::compiler::clvm;
+use crate::compiler::prims;
+use crate::compiler::runtypes::RunFailure;
 use crate::compiler::sexp::SExp as ModernSExp;
 use crate::compiler::srcloc::Srcloc;
 use crate::util::{number_from_u8, u8_from_number};
@@ -70,6 +74,17 @@ pub trait ClassicAllocator {
     ) -> Result<Self::NodePtr, ClError>;
     fn import(&mut self, loc: Srcloc, node: NodePtr) -> Result<Self::NodePtr, ClError>;
     fn export(&self, node: &Self::NodePtr) -> NodePtr;
+    /// Run CLVM using the representation best suited to this allocator.
+    ///
+    /// The supplied runner remains the fallback dialect for non-core
+    /// operators. Located allocators may evaluate core CLVM without first
+    /// exporting the program to clvmr.
+    fn run_clvm(
+        &mut self,
+        runner: Rc<dyn TRunProgram>,
+        program: &Self::NodePtr,
+        args: &Self::NodePtr,
+    ) -> Result<Self::NodePtr, ClError>;
 }
 
 thread_local! {
@@ -127,6 +142,18 @@ impl ClassicAllocator for Allocator {
     fn export(&self, node: &Self::NodePtr) -> NodePtr {
         *node
     }
+    fn run_clvm(
+        &mut self,
+        runner: Rc<dyn TRunProgram>,
+        program: &Self::NodePtr,
+        args: &Self::NodePtr,
+    ) -> Result<Self::NodePtr, ClError> {
+        let loc = self.loc(program);
+        runner
+            .run_program(self, *program, *args, None)
+            .map(|result| result.1)
+            .map_err(|err| ClError(loc, err))
+    }
 }
 
 /// A stage-2 node backed by the compiler's location-aware S-expression.
@@ -304,4 +331,90 @@ impl ClassicAllocator for SExpClassicAllocator {
     fn export(&self, node: &Self::NodePtr) -> NodePtr {
         node.raw
     }
+
+    fn run_clvm(
+        &mut self,
+        runner: Rc<dyn TRunProgram>,
+        program: &Self::NodePtr,
+        args: &Self::NodePtr,
+    ) -> Result<Self::NodePtr, ClError> {
+        let result = clvm::run(
+            &mut self.allocator,
+            runner,
+            prims::prim_map(),
+            program.sexp.clone(),
+            args.sexp.clone(),
+            None,
+            None,
+        )
+        .map_err(|err| match err {
+            RunFailure::RunErr(loc, message) => {
+                ClError(loc, EvalErr::InternalError(NodePtr::NIL, message))
+            }
+            RunFailure::RunExn(loc, value) => ClError(
+                loc,
+                EvalErr::InternalError(NodePtr::NIL, format!("exception: {value}")),
+            ),
+        })?;
+
+        self.from_sexp(result)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{ClassicAllocator, SExpClassicAllocator};
+    use crate::classic::clvm_tools::stages::stage_0::{DefaultProgramRunner, TRunProgram};
+    use crate::compiler::sexp::{enlist, SExp};
+    use crate::compiler::srcloc::Srcloc;
+    use num_bigint::ToBigInt;
+    use std::rc::Rc;
+
+    fn quote(loc: Srcloc, value: Rc<SExp>) -> Rc<SExp> {
+        Rc::new(SExp::Cons(
+            loc.clone(),
+            Rc::new(SExp::Integer(loc, 1_i32.to_bigint().unwrap())),
+            value,
+        ))
+    }
+
+    #[test]
+    fn located_allocator_runs_core_clvm_without_losing_result_locations() {
+        let file = Rc::new("*located-runner-test*".to_string());
+        let program_loc = Srcloc::new(file.clone(), 1, 1);
+        let operator_loc = Srcloc::new(file.clone(), 1, 2);
+        let value_loc = Srcloc::new(file.clone(), 1, 5);
+        let nil_loc = Srcloc::new(file.clone(), 1, 10);
+        let args_loc = Srcloc::new(file, 1, 12);
+
+        let value = Rc::new(SExp::Atom(value_loc.clone(), b"value".to_vec()));
+        let nil = Rc::new(SExp::Nil(nil_loc.clone()));
+        let program = Rc::new(enlist(
+            program_loc,
+            &[
+                Rc::new(SExp::Integer(
+                    operator_loc.clone(),
+                    4_i32.to_bigint().unwrap(),
+                )),
+                quote(value_loc.clone(), value),
+                quote(nil_loc.clone(), nil),
+            ],
+        ));
+        let expected_result_loc = program.loc();
+
+        let mut allocator = SExpClassicAllocator::new();
+        let program = allocator.from_sexp(program).unwrap();
+        let args = allocator.from_sexp(Rc::new(SExp::Nil(args_loc))).unwrap();
+        let runner: Rc<dyn TRunProgram> = Rc::new(DefaultProgramRunner::new());
+        let result = allocator.run_clvm(runner, &program, &args).unwrap();
+
+        match result.sexp.as_ref() {
+            SExp::Cons(loc, first, rest) => {
+                assert_eq!(*loc, expected_result_loc);
+                assert_eq!(first.loc(), value_loc);
+                assert_eq!(rest.loc(), nil_loc);
+            }
+            result => panic!("expected cons result, got {result}"),
+        }
+    }
 }

src/classic/clvm_tools/stages/stage_2/module.rs

@@ -676,32 +676,19 @@ where
                                 produce_extra_info
                             )?;
 
-                        let compiled_export = allocator.export(&compiled);
-                        let loc = allocator.loc(&compiled);
-                        let compilation_result =
-                            run_program.run_program(
-                                allocator.allocator(),
-                                compiled_export,
-                                NodePtr::NIL,
-                                None
-                            ).map_err(|e| {
-                                ClError(loc.clone(), e)
-                            })?;
-
-                        let result =
-                            run_program.run_program(
-                                allocator.allocator(),
-                                compilation_result.1,
-                                NodePtr::NIL,
-                                None
-                            ).map_err(|e| {
-                                ClError(loc.clone(), e)
-                            })?;
-
-                        let result_imp = allocator.import(loc, result.1)?;
+                        let compilation_result = allocator.run_clvm(
+                            run_program.clone(),
+                            &compiled,
+                            &nil_import,
+                        )?;
+                        let result = allocator.run_clvm(
+                            run_program.clone(),
+                            &compilation_result,
+                            &nil_import,
+                        )?;
                         delayed_constants.remove(name);
                         result_collection.constants.insert(
-                            name.to_vec(), quote(allocator, &result_imp)?
+                            name.to_vec(), quote(allocator, &result)?
                         );
                     }
 

src/classic/clvm_tools/stages/stage_2/optimize.rs

@@ -123,11 +123,9 @@ where
         );
     }
     if sc_r && nn_r {
-        let r_export = allocator.export(r);
-        let res = runner
-            .run_program(allocator.allocator(), r_export, NodePtr::NIL, None)
-            .map_err(|e| ClError(allocator.loc(r), e))?;
-        let r1 = allocator.import(allocator.loc(r), res.1)?;
+        let loc = allocator.loc(r);
+        let nil = allocator.new_atom(loc, &[])?;
+        let r1 = allocator.run_clvm(runner, r, &nil)?;
         if DIAG_OPTIMIZATIONS {
             println!(
                 "CONSTANT_OPTIMIZER {} TO {}",