Support compiling prover functions. (#2478)

executor/src/witgen/jit/compiler.rs

@@ -547,13 +547,20 @@ fn prover_function_code<T: FieldElement, D: DefinitionFetcher>(
             codegen.generate_code_for_expression(code)?
         ),
         ProverFunctionComputation::ProvideIfUnknown(code) => {
+            assert!(!f.compute_multi);
             format!("({}).call()", codegen.generate_code_for_expression(code)?)
         }
     };
-
+    let code = if f.compute_multi {
+        format!("({code}).as_slice().try_into().unwrap()")
+    } else {
+        assert_eq!(f.target.len(), 1);
+        format!("[{code}]")
+    };
+    let length = f.target.len();
     let index = f.index;
     Ok(format!(
-        "fn prover_function_{index}(i: u64, args: &[FieldElement]) -> FieldElement {{\n\
+        "fn prover_function_{index}(i: u64, args: &[FieldElement]) -> [FieldElement; {length}] {{\n\
             let i: ibig::IBig = i.into();\n\
             {code}
         }}"

executor/src/witgen/jit/function_cache.rs

@@ -6,10 +6,7 @@ use powdr_number::{FieldElement, KnownField};
 
 use crate::witgen::{
     data_structures::finalizable_data::{ColumnLayout, CompactDataRef},
-    jit::{
-        effect::{format_code, Effect},
-        processor::ProcessorResult,
-    },
+    jit::{effect::format_code, processor::ProcessorResult},
     machines::{
         profiling::{record_end, record_start},
         LookupCell, MachineParts,
@@ -170,22 +167,6 @@ impl<'a, T: FieldElement> FunctionCache<'a, T> {
             );
         }
 
-        // TODO remove this once code generation for prover functions is working.
-        if code
-            .iter()
-            .flat_map(|e| -> Box<dyn Iterator<Item = &Effect<_, _>>> {
-                if let Effect::Branch(_, first, second) = e {
-                    Box::new(first.iter().chain(second))
-                } else {
-                    Box::new(std::iter::once(e))
-                }
-            })
-            .any(|e| matches!(e, Effect::ProverFunctionCall { .. }))
-        {
-            log::debug!("Inferred code contains call to prover function, which is not yet implemented. Using runtime solving instead.");
-            return None;
-        }
-
         log::trace!("Generated code ({} steps)", code.len());
         let known_inputs = cache_key
             .known_args

executor/src/witgen/jit/prover_function_heuristics.rs

@@ -16,7 +16,6 @@ pub trait TrySymbolByName: Copy {
     fn try_symbol_by_name<'a>(&'a self, name: &str) -> Option<&'a Symbol>;
 }
 
-#[allow(unused)]
 #[derive(Clone)]
 pub struct ProverFunction<'a, T> {
     pub index: usize,
@@ -29,7 +28,6 @@ pub struct ProverFunction<'a, T> {
     pub computation: ProverFunctionComputation<'a>,
 }
 
-#[allow(unused)]
 #[derive(Clone)]
 pub enum ProverFunctionComputation<'a> {
     /// The expression `f` in `query |i| std::prover::provide_if_unknown(Y, i, f)`,

jit-compiler/src/compiler.rs

@@ -106,14 +106,10 @@ pub fn call_cargo(code: &str, opt_level: Option<u32>) -> Result<PathInTempDir, S
         .output()
         .unwrap();
     if !out.status.success() {
-        if log::log_enabled!(log::Level::Debug) {
-            let stderr = from_utf8(&out.stderr).unwrap_or("UTF-8 error in error message.");
-            return Err(format!(
+        let stderr = from_utf8(&out.stderr).unwrap_or("UTF-8 error in error message.");
+        return Err(format!(
                 "Rust compiler error when JIT-compiling. Will use interpreter instead. Error message:\n{stderr}."
             ));
-        } else {
-            return Err("Rust compiler error when JIT-compiling. Will use interpreter instead. Set log level to DEBUG for reason.".to_string());
-        }
     }
     #[allow(clippy::print_stdout)]
     if output_asm {

jit-compiler/src/includes/types.rs

@@ -27,6 +27,10 @@ impl<T> PilVec<T> {
     fn len(&self) -> usize {
         self.0.len()
     }
+
+    fn as_slice(&self) -> &[T] {
+        self.0.as_ref()
+    }
 }
 impl<T> From<Vec<T>> for PilVec<T> {
     fn from(v: Vec<T>) -> Self {