Add witness hints fillers

circuit/Cargo.toml

@@ -21,6 +21,7 @@ p3-koala-bear.workspace = true
 p3-matrix.workspace = true
 p3-symmetric.workspace = true
 p3-uni-stark.workspace = true
+p3-util.workspace = true
 rand.workspace = true
 
 # Other common dependencies

circuit/src/builder/circuit_builder.rs

@@ -6,10 +6,11 @@ use itertools::zip_eq;
 use p3_field::{Field, PrimeCharacteristicRing};
 
 use super::compiler::{ExpressionLowerer, NonPrimitiveLowerer, Optimizer};
-use super::{BuilderConfig, ExpressionBuilder, PublicInputTracker};
+use super::{BuilderConfig, ExpressionBuilder};
 use crate::CircuitBuilderError;
+use crate::builder::public_input_tracker::PublicInputTracker;
 use crate::circuit::Circuit;
-use crate::op::NonPrimitiveOpType;
+use crate::op::{NonPrimitiveOpType, WitnessHintFiller};
 use crate::ops::MmcsVerifyConfig;
 use crate::types::{ExprId, NonPrimitiveOpId, WitnessAllocator, WitnessId};
 
@@ -120,16 +121,28 @@ where
         self.public_tracker.count()
     }
 
-    /// Allocates a witness hint (uninitialized witness slot set during non-primitive execution).
+    /// Allocates multiple witness. Witness hints are placeholders for values that will later be provided by a
+    /// `filler`.
     #[must_use]
-    pub fn alloc_witness_hint(&mut self, label: &'static str) -> ExprId {
-        self.expr_builder.add_witness_hint(label)
+    pub fn alloc_witness_hints<W: 'static + WitnessHintFiller<F>>(
+        &mut self,
+        filler: W,
+        label: &'static str,
+    ) -> Vec<ExprId> {
+        self.expr_builder.add_witness_hints(filler, label)
     }
 
-    /// Allocates multiple witness hints.
+    /// Allocates multiple witness hints without saying how they should be filled
+    /// TODO: Remove this function
     #[must_use]
-    pub fn alloc_witness_hints(&mut self, count: usize, label: &'static str) -> Vec<ExprId> {
-        self.expr_builder.add_witness_hints(count, label)
+    pub fn alloc_witness_hints_no_filler(
+        &mut self,
+        count: usize,
+        label: &'static str,
+    ) -> Vec<ExprId> {
+        (0..count)
+            .map(|_| self.expr_builder.add_witness_hint(label))
+            .collect()
     }
 
     /// Adds a constant to the circuit (deduplicated).
@@ -394,6 +407,7 @@ where
             self.expr_builder.graph(),
             self.expr_builder.pending_connects(),
             self.public_tracker.count(),
+            self.expr_builder.hints_with_fillers(),
             self.witness_alloc,
         );
         let (primitive_ops, public_rows, expr_to_widx, public_mappings, witness_count) =
@@ -427,7 +441,6 @@ mod tests {
 
     use p3_baby_bear::BabyBear;
     use p3_field::PrimeCharacteristicRing;
-    use proptest::prelude::*;
 
     use super::*;
 
@@ -700,6 +713,66 @@ mod tests {
         assert_eq!(circuit.primitive_ops.len(), 2);
     }
 
+    #[derive(Debug, Clone)]
+    struct ConstantHint<const C: usize> {
+        inputs: Vec<ExprId>,
+    }
+
+    impl<const C: usize> ConstantHint<C> {
+        pub fn new(input: ExprId) -> Self {
+            Self {
+                inputs: vec![input],
+            }
+        }
+    }
+
+    impl<F: Field, const C: usize> WitnessHintFiller<F> for ConstantHint<C> {
+        fn inputs(&self) -> &[ExprId] {
+            &self.inputs
+        }
+
+        fn n_outputs(&self) -> usize {
+            1
+        }
+
+        fn compute_outputs(&self, _inputs_val: Vec<F>) -> Result<Vec<F>, crate::CircuitError> {
+            Ok(vec![F::from_usize(C)])
+        }
+    }
+    #[test]
+    fn test_build_with_witness_hint() {
+        let mut builder = CircuitBuilder::<BabyBear>::new();
+        let a = builder.add_const(BabyBear::ZERO);
+        let mock_filler = ConstantHint::<1>::new(a);
+        let b = builder.alloc_witness_hints(mock_filler, "a");
+        assert_eq!(b.len(), 1);
+        let circuit = builder
+            .build()
+            .expect("Circuit with operations should build");
+
+        assert_eq!(circuit.witness_count, 2);
+        assert_eq!(circuit.primitive_ops.len(), 2);
+
+        match &circuit.primitive_ops[1] {
+            crate::op::Op::Unconstrained {
+                inputs, outputs, ..
+            } => {
+                assert_eq!(*inputs, vec![WitnessId(0)]);
+                assert_eq!(*outputs, vec![WitnessId(1)]);
+            }
+            _ => panic!("Expected Unconstrained at index 0"),
+        }
+    }
+}
+
+#[cfg(test)]
+mod proptests {
+    use p3_baby_bear::BabyBear;
+    use p3_field::PrimeCharacteristicRing;
+    use proptest::prelude::*;
+
+    use super::*;
+
     // Strategy for generating valid field elements
     fn field_element() -> impl Strategy<Value = BabyBear> {
         any::<u64>().prop_map(BabyBear::from_u64)

circuit/src/builder/compiler/expression_lowerer.rs

@@ -1,3 +1,5 @@
+use alloc::boxed::Box;
+use alloc::string::ToString;
 use alloc::vec;
 use alloc::vec::Vec;
 
@@ -8,6 +10,7 @@ use crate::Op;
 use crate::builder::CircuitBuilderError;
 use crate::builder::compiler::get_witness_id;
 use crate::expr::{Expr, ExpressionGraph};
+use crate::op::WitnessHintFiller;
 use crate::types::{ExprId, WitnessAllocator, WitnessId};
 
 /// Sparse disjoint-set "find" with path compression over a HashMap (iterative).
@@ -70,10 +73,15 @@ pub struct ExpressionLowerer<'a, F> {
     /// Number of public inputs
     public_input_count: usize,
 
+    /// The hint witnesses with their respective filler
+    hints_with_fillers: &'a [HintsWithFiller<F>],
+
     /// Witness allocator
     witness_alloc: WitnessAllocator,
 }
 
+pub type HintsWithFiller<F> = (Vec<ExprId>, Box<dyn WitnessHintFiller<F>>);
+
 impl<'a, F> ExpressionLowerer<'a, F>
 where
     F: Clone + PrimeCharacteristicRing + PartialEq + Eq + core::hash::Hash,
@@ -83,12 +91,14 @@ where
         graph: &'a ExpressionGraph<F>,
         pending_connects: &'a [(ExprId, ExprId)],
         public_input_count: usize,
+        hints_with_fillers: &'a [HintsWithFiller<F>],
         witness_alloc: WitnessAllocator,
     ) -> Self {
         Self {
             graph,
             pending_connects,
             public_input_count,
+            hints_with_fillers,
             witness_alloc,
         }
     }
@@ -162,7 +172,7 @@ where
             }
         };
 
-        // Pass B: emit public inputs
+        // Pass B: emit public inputs and process witness hints
         for (expr_idx, expr) in self.graph.nodes().iter().enumerate() {
             if let Expr::Public(pos) = expr {
                 let id = ExprId(expr_idx as u32);
@@ -179,18 +189,57 @@ where
             }
         }
 
-        // Pass C: emit arithmetic ops in creation order; tie outputs to class slot if connected
+        // Pass C: collect witness hints and emit `Unconstrained` operatios
+        for (expr_idx, _) in self
+            .graph
+            .nodes()
+            .iter()
+            .enumerate()
+            .filter(|(_, expr)| matches!(expr, Expr::Witness))
+        {
+            let expr_id = ExprId(expr_idx as u32);
+            let out_widx = alloc_witness_id_for_expr(expr_idx);
+            expr_to_widx.insert(expr_id, out_widx);
+        }
+        for (witness_hints, filler) in self.hints_with_fillers.iter() {
+            let inputs = filler
+                .inputs()
+                .iter()
+                .map(|expr_id| {
+                    expr_to_widx
+                        .get(expr_id)
+                        .ok_or(CircuitBuilderError::MissingExprMapping {
+                            expr_id: *expr_id,
+                            context: "Unconstrained op".to_string(),
+                        })
+                        .copied()
+                })
+                .collect::<Result<Vec<WitnessId>, _>>()?;
+            let outputs = witness_hints
+                .iter()
+                .map(|expr_id| {
+                    expr_to_widx
+                        .get(expr_id)
+                        .ok_or(CircuitBuilderError::MissingExprMapping {
+                            expr_id: *expr_id,
+                            context: "Unconstrained op".to_string(),
+                        })
+                        .copied()
+                })
+                .collect::<Result<Vec<WitnessId>, _>>()?;
+            primitive_ops.push(Op::Unconstrained {
+                inputs,
+                outputs,
+                filler: filler.clone(),
+            })
+        }
+
+        // Pass D: emit arithmetic ops in creation order; tie outputs to class slot if connected
         for (expr_idx, expr) in self.graph.nodes().iter().enumerate() {
             let expr_id = ExprId(expr_idx as u32);
             match expr {
-                Expr::Const(_) | Expr::Public(_) => { /* handled above */ }
-                Expr::Witness => {
-                    // Allocate a fresh witness slot (non-primitive op)
-                    // Allows non-primitive operations to set values during execution that
-                    // are not part of the central Witness bus.
-                    let out_widx = alloc_witness_id_for_expr(expr_idx);
-                    expr_to_widx.insert(expr_id, out_widx);
-                }
+                Expr::Const(_) | Expr::Public(_) | Expr::Witness => { /* handled above */ }
+
                 Expr::Add { lhs, rhs } => {
                     let out_widx = alloc_witness_id_for_expr(expr_idx);
                     let a_widx = get_witness_id(
@@ -394,9 +443,10 @@ mod tests {
         let quot = graph.add_expr(Expr::Div { lhs: diff, rhs: p2 });
 
         let connects = vec![];
+        let hints_with_fillers = vec![];
         let alloc = WitnessAllocator::new();
 
-        let lowerer = ExpressionLowerer::new(&graph, &connects, 3, alloc);
+        let lowerer = ExpressionLowerer::new(&graph, &connects, 3, &hints_with_fillers, alloc);
         let (prims, public_rows, expr_map, public_map, witness_count) = lowerer.lower().unwrap();
 
         // Verify Primitives
@@ -563,9 +613,10 @@ mod tests {
         // Group B: p1 ~ p2 ~ p3 (transitive)
         // Group C: sum ~ p4 (operation result shared)
         let connects = vec![(c_42, p0), (p1, p2), (p2, p3), (sum, p4)];
+        let hints_with_fillers = vec![];
         let alloc = WitnessAllocator::new();
 
-        let lowerer = ExpressionLowerer::new(&graph, &connects, 5, alloc);
+        let lowerer = ExpressionLowerer::new(&graph, &connects, 5, &hints_with_fillers, alloc);
         let (prims, public_rows, expr_map, public_map, witness_count) = lowerer.lower().unwrap();
 
         // Verify Primitives
@@ -703,8 +754,9 @@ mod tests {
         });
 
         let connects = vec![];
+        let hints_with_fillers = vec![];
         let alloc = WitnessAllocator::new();
-        let lowerer = ExpressionLowerer::new(&graph, &connects, 0, alloc);
+        let lowerer = ExpressionLowerer::new(&graph, &connects, 0, &hints_with_fillers, alloc);
         let result = lowerer.lower();
 
         assert!(result.is_err());
@@ -724,8 +776,9 @@ mod tests {
         });
 
         let connects = vec![];
+        let hints_with_fillers = vec![];
         let alloc = WitnessAllocator::new();
-        let lowerer = ExpressionLowerer::new(&graph, &connects, 0, alloc);
+        let lowerer = ExpressionLowerer::new(&graph, &connects, 0, &hints_with_fillers, alloc);
         let result = lowerer.lower();
 
         assert!(result.is_err());

circuit/src/builder/compiler/mod.rs

@@ -4,7 +4,7 @@ mod expression_lowerer;
 mod non_primitive_lowerer;
 mod optimizer;
 
-pub use expression_lowerer::ExpressionLowerer;
+pub use expression_lowerer::{ExpressionLowerer, HintsWithFiller};
 use hashbrown::HashMap;
 pub use non_primitive_lowerer::NonPrimitiveLowerer;
 pub use optimizer::Optimizer;

circuit/src/builder/errors.rs

@@ -23,6 +23,10 @@ pub enum CircuitBuilderError {
         got: usize,
     },
 
+    /// Unconstrained op received an unexpected number of input expressions.
+    #[error("Expects exactly {expected} witness expressions, got {got}")]
+    UnconstrainedeOpArity { expected: String, got: usize },
+
     /// Non-primitive operation rejected by the active policy/profile.
     #[error("Operation {op:?} is not allowed by the current profile")]
     OpNotAllowed { op: NonPrimitiveOpType },