Compile prover functions.

Author: chriseth

executor/src/witgen/jit/block_machine_processor.rs

@@ -13,6 +13,7 @@ use crate::witgen::{
 
 use super::{
     processor::ProcessorResult,
+    prover_function_heuristics::ProverFunction,
     variable::{Cell, Variable},
     witgen_inference::{CanProcessCall, FixedEvaluator, WitgenInference},
 };
@@ -50,7 +51,7 @@ impl<'a, T: FieldElement> BlockMachineProcessor<'a, T> {
         can_process: impl CanProcessCall<T>,
         identity_id: u64,
         known_args: &BitVec,
-    ) -> Result<ProcessorResult<T>, String> {
+    ) -> Result<(ProcessorResult<T>, Vec<ProverFunction<'a, T>>), String> {
         let connection = self.machine_parts.connections[&identity_id];
         assert_eq!(connection.right.expressions.len(), known_args.len());
 
@@ -119,7 +120,7 @@ impl<'a, T: FieldElement> BlockMachineProcessor<'a, T> {
             .iter()
             .enumerate()
             .filter_map(|(i, is_input)| (!is_input).then_some(Variable::Param(i)));
-        Processor::new(
+        let result= Processor::new(
             self.fixed_data,
             self,
             identities,
@@ -129,7 +130,12 @@ impl<'a, T: FieldElement> BlockMachineProcessor<'a, T> {
         .with_block_shape_check()
         .with_block_size(self.block_size)
         .with_requested_range_constraints((0..known_args.len()).map(Variable::Param))
-        .with_prover_functions(prover_functions)
+        .with_prover_functions(
+            prover_functions
+                .iter()
+                .flat_map(|f| (0..self.block_size).map(move |row| (f.clone(), row as i32)))
+                .collect_vec()
+        )
         .generate_code(can_process, witgen)
         .map_err(|e| {
             let err_str = e.to_string_with_variable_formatter(|var| match var {
@@ -149,7 +155,8 @@ impl<'a, T: FieldElement> BlockMachineProcessor<'a, T> {
                 .take(10)
                 .format("\n  ");
             format!("Code generation failed: {shortened_error}\nRun with RUST_LOG=trace to see the code generated so far.")
-        })
+        })?;
+        Ok((result, prover_functions))
     }
 }
 

executor/src/witgen/jit/compiler.rs

@@ -6,7 +6,7 @@ use powdr_ast::{
     analyzed::{PolyID, PolynomialType},
     indent,
 };
-use powdr_jit_compiler::util_code::util_code;
+use powdr_jit_compiler::{util_code::util_code, CodeGenerator, DefinitionFetcher};
 use powdr_number::FieldElement;
 
 use crate::witgen::{
@@ -23,6 +23,7 @@ use crate::witgen::{
 
 use super::{
     effect::{Assertion, BranchCondition, Effect, ProverFunctionCall},
+    prover_function_heuristics::ProverFunction,
     symbolic_expression::{BinaryOperator, BitOperator, SymbolicExpression, UnaryOperator},
     variable::Variable,
 };
@@ -87,17 +88,28 @@ extern "C" fn call_machine<T: FieldElement, Q: QueryCallback<T>>(
 }
 
 /// Compile the given inferred effects into machine code and load it.
-pub fn compile_effects<T: FieldElement>(
+pub fn compile_effects<T: FieldElement, D: DefinitionFetcher>(
+    definitions: &D,
     column_layout: ColumnLayout,
     known_inputs: &[Variable],
     effects: &[Effect<T, Variable>],
+    prover_functions: Vec<ProverFunction<T>>,
 ) -> Result<WitgenFunction<T>, String> {
     let utils = util_code::<T>()?;
     let interface = interface_code(column_layout);
+    let mut codegen = CodeGenerator::new(definitions);
+    let prover_functions = prover_functions
+        .iter()
+        .map(|f| prover_function_code(f, &mut codegen))
+        .collect::<Result<Vec<_>, _>>()?
+        .into_iter()
+        .format("\n");
     let witgen_code = witgen_code(known_inputs, effects);
     let code = format!(
         "{utils}\n\
         //-------------------------------\n\
+        {prover_functions}\n\
+        //-------------------------------\n\
         {interface}\n\
         //-------------------------------\n\
         {witgen_code}"
@@ -522,11 +534,31 @@ fn interface_code(column_layout: ColumnLayout) -> String {
     )
 }
 
+fn prover_function_code<T: FieldElement, D: DefinitionFetcher>(
+    f: &ProverFunction<'_, T>,
+    codegen: &mut CodeGenerator<'_, T, D>,
+) -> Result<String, String> {
+    let ProverFunction::ComputeFrom(f) = f else {
+        return Err("ProvideIfUnknown functions are not supported".to_string());
+    };
+
+    let code = codegen.generate_code_for_expresson(f.computation)?;
+
+    let index = f.index;
+    Ok(format!(
+        "fn prover_function_{index}(i: u64, args: &[FieldElement]) -> FieldElement {{\n\
+            let i: ibig::IBig = i.into();\n\
+            ({code}).call(args.to_vec().into())\n\
+        }}"
+    ))
+}
+
 #[cfg(test)]
 mod tests {
 
     use std::ptr::null;
 
+    use powdr_ast::analyzed::FunctionValueDefinition;
     use pretty_assertions::assert_eq;
     use test_log::test;
 
@@ -538,18 +570,27 @@ mod tests {
 
     use super::*;
 
+    struct NoDefinitions;
+    impl DefinitionFetcher for NoDefinitions {
+        fn get_definition(&self, _: &str) -> Option<&FunctionValueDefinition> {
+            None
+        }
+    }
+
     fn compile_effects(
         column_count: usize,
         known_inputs: &[Variable],
         effects: &[Effect<GoldilocksField, Variable>],
     ) -> Result<WitgenFunction<GoldilocksField>, String> {
         super::compile_effects(
+            &NoDefinitions,
             ColumnLayout {
                 column_count,
                 first_column_id: 0,
             },
             known_inputs,
             effects,
+            vec![],
         )
     }
 

executor/src/witgen/jit/function_cache.rs

@@ -114,10 +114,13 @@ impl<'a, T: FieldElement> FunctionCache<'a, T> {
             cache_key.known_args
         );
 
-        let ProcessorResult {
-            code,
-            range_constraints,
-        } = self
+        let (
+            ProcessorResult {
+                code,
+                range_constraints,
+            },
+            prover_functions,
+        ) = self
             .processor
             .generate_code(can_process, cache_key.identity_id, &cache_key.known_args)
             .map_err(|e| {

jit-compiler/src/codegen.rs

@@ -80,6 +80,13 @@ impl<'a, T: FieldElement, Def: DefinitionFetcher> CodeGenerator<'a, T, Def> {
         Ok(self.symbol_reference(name, type_args))
     }
 
+    /// Generates code for an isolated expression. This might request code generation
+    /// for referenced symbols, this the returned code is only valid code in connection with
+    /// the code returned by `self.generated_code`.
+    pub fn generate_code_for_expresson(&mut self, e: &Expression) -> Result<String, String> {
+        self.format_expr(e, 0)
+    }
+
     /// Returns the concatenation of all successfully compiled symbols.
     pub fn generated_code(self) -> String {
         self.symbols

jit-compiler/src/lib.rs

@@ -8,13 +8,13 @@ use std::{
     sync::Arc,
 };
 
-use codegen::CodeGenerator;
 use compiler::{generate_glue_code, load_library};
 
 use itertools::Itertools;
 use powdr_ast::analyzed::Analyzed;
 use powdr_number::FieldElement;
 
+pub use codegen::{CodeGenerator, DefinitionFetcher};
 pub use compiler::call_cargo;
 
 pub struct CompiledPIL {