minor touchups

Author: sai-deng

circuit/src/builder/compiler/non_primitive_lowerer.rs

@@ -210,7 +210,10 @@ mod tests {
         // Test MmcsVerify with mock config (simplest case: 1 leaf + 1 index + 1 root)
         let mock_config = MmcsVerifyConfig::mock_config();
         assert_eq!(mock_config.ext_field_digest_elems, 1);
-        assert_eq!(mock_config.input_size(), 2 * mock_config.ext_field_digest_elems + 1);
+        assert_eq!(
+            mock_config.input_size(),
+            2 * mock_config.ext_field_digest_elems + 1
+        );
 
         let mut config = BuilderConfig::new();
         config.enable_mmcs(&mock_config);

circuit/src/builder/config.rs

@@ -2,41 +2,50 @@ use hashbrown::HashMap;
 
 use crate::op::{NonPrimitiveOpConfig, NonPrimitiveOpType};
 
+/// Configuration for the circuit builder.
 #[derive(Debug, Clone, Default)]
 pub struct BuilderConfig {
+    /// Enabled non-primitive operation types with their respective configuration.
     enabled_ops: HashMap<NonPrimitiveOpType, NonPrimitiveOpConfig>,
 }
 
 impl BuilderConfig {
+    /// Creates a new builder configuration.
     pub fn new() -> Self {
         Self {
             enabled_ops: HashMap::new(),
         }
     }
 
+    /// Enables a non-primitive operation type with its configuration.
     pub fn enable_op(&mut self, op: NonPrimitiveOpType, cfg: NonPrimitiveOpConfig) {
         self.enabled_ops.insert(op, cfg);
     }
 
+    /// Enables MMCS verification operations with the given configuration.
     pub fn enable_mmcs(&mut self, mmcs_config: &crate::ops::MmcsVerifyConfig) {
         self.enable_op(
             NonPrimitiveOpType::MmcsVerify,
             NonPrimitiveOpConfig::MmcsVerifyConfig(mmcs_config.clone()),
         );
     }
 
+    /// Enables FRI verification operations.
     pub fn enable_fri(&mut self) {
-        // TODO when available
+        // TODO: Add FRI ops when available.
     }
 
+    /// Checks whether an operation type is enabled.
     pub fn is_op_enabled(&self, op: &NonPrimitiveOpType) -> bool {
         self.enabled_ops.contains_key(op)
     }
 
+    /// Gets the configuration for an operation type, if enabled.
     pub fn get_op_config(&self, op: &NonPrimitiveOpType) -> Option<&NonPrimitiveOpConfig> {
         self.enabled_ops.get(op)
     }
 
+    /// Consumes the config and returns the enabled operations map.
     pub fn into_enabled_ops(self) -> HashMap<NonPrimitiveOpType, NonPrimitiveOpConfig> {
         self.enabled_ops
     }

circuit/src/lib.rs

@@ -23,7 +23,7 @@ pub use builder::{CircuitBuilder, CircuitBuilderError};
 pub use circuit::{Circuit, CircuitField};
 pub use errors::CircuitError;
 pub use expr::{Expr, ExpressionGraph};
-pub use op::{NonPrimitiveOp, NonPrimitiveOpPrivateData, NonPrimitiveOpType, Op};
+pub use op::{NonPrimitiveOpPrivateData, NonPrimitiveOpType, Op};
 pub use ops::{FriOps, MmcsOps};
 pub use tables::{CircuitRunner, MmcsPrivateData, MmcsTrace, Traces};
 pub use types::{ExprId, NonPrimitiveOpId, WitnessAllocator, WitnessId};

circuit/src/tables/runner.rs

@@ -71,17 +71,42 @@ impl<F: CircuitField> CircuitRunner<F> {
             });
         }
 
-        // Validate that the private data matches the operation type (if any)
+        // Validate that the private data matches the operation type
         if let Op::NonPrimitiveOpWithExecutor { executor, .. } =
             &self.circuit.non_primitive_ops[op_id.0 as usize]
-            && let (
-                crate::op::NonPrimitiveOpType::MmcsVerify,
-                NonPrimitiveOpPrivateData::MmcsVerify(_),
-            ) = (executor.op_type(), &private_data)
         {
-            // ok
+            match (executor.op_type(), &private_data) {
+                (
+                    crate::op::NonPrimitiveOpType::MmcsVerify,
+                    NonPrimitiveOpPrivateData::MmcsVerify(_),
+                ) => {
+                    // ok
+                }
+                (op_ty, _) => {
+                    // Other ops currently don't expect private data
+                    return Err(CircuitError::IncorrectNonPrimitiveOpPrivateData {
+                        op: op_ty.clone(),
+                        operation_index: op_id,
+                        expected: format!("no private data"),
+                        got: format!("{private_data:?}"),
+                    });
+                }
+            }
+        }
+
+        // Disallow double-setting private data
+        if self.non_primitive_op_private_data[op_id.0 as usize].is_some() {
+            if let Op::NonPrimitiveOpWithExecutor { executor, .. } =
+                &self.circuit.non_primitive_ops[op_id.0 as usize]
+            {
+                return Err(CircuitError::IncorrectNonPrimitiveOpPrivateData {
+                    op: executor.op_type().clone(),
+                    operation_index: op_id,
+                    expected: format!("private data not previously set"),
+                    got: format!("already set"),
+                });
+            }
         }
-        // Other ops currently don't expect private data
 
         // Store private data for this operation
         self.non_primitive_op_private_data[op_id.0 as usize] = Some(private_data);