registers in memory

pipeline/src/lib.rs

@@ -148,6 +148,42 @@ pub fn serde_data_to_query_callback<T: FieldElement, S: serde::Serialize + Send
     }
 }
 
+pub fn dict_data_to_query_callback<T: FieldElement>(
+    dict: BTreeMap<u32, Vec<T>>,
+) -> impl QueryCallback<T> {
+    move |query: &str| -> Result<Option<T>, String> {
+        let (id, data) = parse_query(query)?;
+        match id {
+            "None" => Ok(None),
+            "DataIdentifier" => {
+                let [index, cb_channel] = data[..] else {
+                    panic!()
+                };
+                let cb_channel = cb_channel
+                    .parse::<u32>()
+                    .map_err(|e| format!("Error parsing callback data channel: {e})"))?;
+
+                if !dict.contains_key(&cb_channel) {
+                    return Err("Callback channel mismatch".to_string());
+                }
+
+                let index = index
+                    .parse::<usize>()
+                    .map_err(|e| format!("Error parsing index: {e})"))?;
+
+                let bytes = dict.get(&cb_channel).unwrap();
+
+                // query index 0 means the length
+                Ok(Some(match index {
+                    0 => (bytes.len() as u64).into(),
+                    index => bytes[index - 1],
+                }))
+            }
+            _ => Err(format!("Unsupported query: {query}")),
+        }
+    }
+}
+
 pub fn inputs_to_query_callback<T: FieldElement>(inputs: Vec<T>) -> impl QueryCallback<T> {
     move |query: &str| -> Result<Option<T>, String> {
         let (id, data) = parse_query(query)?;

pipeline/src/pipeline.rs

@@ -29,9 +29,11 @@ use powdr_number::{write_polys_csv_file, write_polys_file, CsvRenderMode, FieldE
 use powdr_schemas::SerializedAnalyzed;
 
 use crate::{
-    handle_simple_queries_callback, inputs_to_query_callback, serde_data_to_query_callback,
+    dict_data_to_query_callback, handle_simple_queries_callback, inputs_to_query_callback,
+    serde_data_to_query_callback,
     util::{try_read_poly_set, FixedPolySet, WitnessPolySet},
 };
+use std::collections::BTreeMap;
 
 type Columns<T> = Vec<(String, Vec<T>)>;
 
@@ -268,6 +270,10 @@ impl<T: FieldElement> Pipeline<T> {
         self.add_query_callback(Arc::new(inputs_to_query_callback(inputs)))
     }
 
+    pub fn with_prover_dict_inputs(self, inputs: BTreeMap<u32, Vec<T>>) -> Self {
+        self.add_query_callback(Arc::new(dict_data_to_query_callback(inputs)))
+    }
+
     pub fn with_backend(mut self, backend: BackendType, options: Option<BackendOptions>) -> Self {
         self.arguments.backend = Some(backend);
         self.arguments.backend_options = options.unwrap_or_default();