Reorg builtins

jit-compiler/Cargo.toml

@@ -16,6 +16,7 @@ log = "0.4.18"
 mktemp = "0.5.0"
 itertools = "0.13"
 libloading = "0.8"
+lazy_static = "1.4.0"
 
 [dev-dependencies]
 powdr-pil-analyzer.workspace = true

jit-compiler/src/codegen.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::{collections::HashMap, sync::OnceLock};
 
 use itertools::Itertools;
 use powdr_ast::{
@@ -76,7 +76,7 @@ impl<'a, T: FieldElement> CodeGenerator<'a, T> {
 
     fn generate_code(&mut self, symbol: &str) -> Result<String, String> {
         if let Some(code) = try_generate_builtin::<T>(symbol) {
-            return Ok(code);
+            return Ok(code.clone());
         }
 
         let Some((_, Some(FunctionValueDefinition::Expression(value)))) =
@@ -297,7 +297,7 @@ impl<'a, T: FieldElement> CodeGenerator<'a, T> {
     /// This is either the escaped name of the symbol or a deref operator
     /// applied to it.
     fn symbol_reference(&self, symbol: &str) -> String {
-        let needs_deref = if is_builtin(symbol) {
+        let needs_deref = if is_builtin::<T>(symbol) {
             false
         } else {
             let (_, def) = self.analyzed.definitions.get(symbol).as_ref().unwrap();
@@ -354,25 +354,43 @@ fn map_type(ty: &Type) -> String {
     }
 }
 
-fn is_builtin(symbol: &str) -> bool {
-    matches!(
-        symbol,
-        "std::check::panic" | "std::field::modulus" | "std::convert::fe"
-    )
+fn get_builtins<T: FieldElement>() -> &'static HashMap<String, String> {
+    static BUILTINS: OnceLock<HashMap<String, String>> = OnceLock::new();
+    BUILTINS.get_or_init(|| {
+        [
+            (
+                "std::array::len",
+                "<T>(a: Vec<T>) -> ibig::IBig { ibig::IBig::from(a.len()) }".to_string(),
+            ),
+            (
+                "std::check::panic",
+                "(s: &str) -> ! { panic!(\"{s}\"); }".to_string(),
+            ),
+            (
+                "std::field::modulus",
+                format!(
+                    "() -> ibig::IBig {{ {} }}",
+                    format_unsigned_integer(&T::modulus().to_arbitrary_integer())
+                ),
+            ),
+        ]
+        .into_iter()
+        .map(|(name, code)| {
+            (
+                name.to_string(),
+                format!("fn {}{code}", escape_symbol(name)),
+            )
+        })
+        .collect()
+    })
 }
 
-fn try_generate_builtin<T: FieldElement>(symbol: &str) -> Option<String> {
-    let code = match symbol {
-        "std::array::len" => "<T>(a: Vec<T>) -> ibig::IBig { ibig::IBig::from(a.len()) }".to_string(),
-        "std::check::panic" => "(s: &str) -> ! { panic!(\"{s}\"); }".to_string(),
-        "std::field::modulus" => {
-            format!("() -> ibig::IBig {{ {} }}", format_unsigned_integer(&T::modulus().to_arbitrary_integer()))
-        }
-        "std::convert::fe" => "(n: ibig::IBig) -> FieldElement {\n    <FieldElement as PrimeField>::BigInt::try_from(n.to_biguint().unwrap()).unwrap().into()\n}"
-            .to_string(),
-        _ => return None,
-    };
-    Some(format!("fn {}{code}", escape_symbol(symbol)))
+fn is_builtin<T: FieldElement>(symbol: &str) -> bool {
+    get_builtins::<T>().contains_key(symbol)
+}
+
+fn try_generate_builtin<T: FieldElement>(symbol: &str) -> Option<&String> {
+    get_builtins::<T>().get(symbol)
 }
 
 #[cfg(test)]

jit-compiler/tests/execution.rs

@@ -6,7 +6,13 @@ use powdr_pil_analyzer::analyze_string;
 
 fn compile(input: &str, symbol: &str) -> LoadedFunction {
     let analyzed = analyze_string::<GoldilocksField>(input).unwrap();
-    powdr_jit_compiler::compile(&analyzed, &[symbol]).unwrap()[symbol].clone()
+    powdr_jit_compiler::compile(&analyzed, &[symbol])
+        .map_err(|e| {
+            eprintln!("Error jit-compiling:\n{e}");
+            e
+        })
+        .unwrap()[symbol]
+        .clone()
 }
 
 #[test]