Feature/optimise step selector builder

src/plonkish/compiler/step_selector.rs

@@ -188,6 +188,68 @@ impl StepSelectorBuilder for TwoStepsSelectorBuilder {
     }
 }
 
+#[derive(Debug, Default, Clone)]
+pub struct LogNSelectorBuilder {}
+
+impl StepSelectorBuilder for LogNSelectorBuilder {
+    fn build<F: Field>(&self, unit: &mut CompilationUnit<F>) {
+        let mut selector: StepSelector<F> = StepSelector {
+            selector_expr: HashMap::new(),
+            selector_expr_not: HashMap::new(),
+            selector_assignment: HashMap::new(),
+            columns: Vec::new(),
+        };
+
+        let n_step_types = unit.step_types.len() as u64;
+        let n_cols = (n_step_types as f64 + 1.0).log2().ceil() as u64;
+
+        let mut annotation;
+        for index in 0..n_cols {
+            annotation = format!("'binary selector column {}'", index);
+
+            let column = Column::advice(annotation.clone(), 0);
+            selector.columns.push(column.clone());
+        }
+
+        let mut step_value = 1;
+        for step in unit.step_types.values() {
+            let mut combined_expr = PolyExpr::Const(F::ZERO);
+            let mut combined_expr_not = PolyExpr::Const(F::ZERO);
+            let mut assignments = Vec::new();
+
+            for i in 0..n_cols {
+                let bit = (step_value >> i) & 1; // Extract the i-th bit of step_value
+                let column = &selector.columns[i as usize];
+
+                if bit == 1 {
+                    combined_expr = combined_expr * column.query(0, format!("Column {}", i));
+                    combined_expr_not = combined_expr_not
+                        * (PolyExpr::Const(F::ONE) - column.query(0, format!("Column {}", i)));
+                    assignments.push((column.query(0, format!("Column {}", i)), F::ONE));
+                } else {
+                    combined_expr = combined_expr
+                        * (PolyExpr::Const(F::ONE) - column.query(0, format!("Column {}", i)));
+                    combined_expr_not =
+                        combined_expr_not * column.query(0, format!("Column {}", i));
+                }
+            }
+
+            selector
+                .selector_expr
+                .insert(step.uuid(), combined_expr.clone());
+            selector
+                .selector_expr_not
+                .insert(step.uuid(), combined_expr_not.clone());
+            selector
+                .selector_assignment
+                .insert(step.uuid(), assignments);
+            step_value += 1;
+        }
+        unit.columns.extend_from_slice(&selector.columns);
+        unit.selector = selector;
+    }
+}
+
 fn other_step_type<F>(unit: &CompilationUnit<F>, uuid: UUID) -> Option<Rc<StepType<F>>> {
     for step_type in unit.step_types.values() {
         if step_type.uuid() != uuid {
@@ -197,3 +259,137 @@ fn other_step_type<F>(unit: &CompilationUnit<F>, uuid: UUID) -> Option<Rc<StepTy
 
     None
 }
+
+#[cfg(test)]
+mod tests {
+    use halo2curves::bn256::Fr;
+    use uuid::Uuid;
+
+    use super::*;
+
+    fn mock_compilation_unit<F>() -> CompilationUnit<F> {
+        CompilationUnit::default()
+    }
+
+    #[test]
+    fn test_log_n_selector_builder_3_step_types() {
+        // Setup
+        let builder = LogNSelectorBuilder {};
+        let mut unit = mock_compilation_unit::<Fr>();
+
+        // Assuming you have 3 step types for this test
+        unit.step_types.insert(
+            Uuid::now_v1(&[1, 2, 3, 4, 5, 6]).as_u128(),
+            Rc::new(StepType::new(
+                Uuid::now_v1(&[1, 2, 3, 4, 5, 6]).as_u128(),
+                "StepType1".to_owned(),
+            )),
+        );
+        unit.step_types.insert(
+            Uuid::now_v1(&[1, 2, 3, 4, 5, 6]).as_u128(),
+            Rc::new(StepType::new(
+                Uuid::now_v1(&[1, 2, 3, 4, 5, 6]).as_u128(),
+                "StepType2".to_owned(),
+            )),
+        );
+        unit.step_types.insert(
+            Uuid::now_v1(&[1, 2, 3, 4, 5, 6]).as_u128(),
+            Rc::new(StepType::new(
+                Uuid::now_v1(&[1, 2, 3, 4, 5, 6]).as_u128(),
+                "StepType3".to_owned(),
+            )),
+        );
+
+        builder.build(&mut unit);
+
+        // Assert number of columns
+        assert_eq!(unit.columns.len(), 2);
+        assert_eq!(unit.selector.columns.len(), 2);
+
+        // TODO: Assert expressions
+
+        // TODO: Assert assignments
+
+    }
+
+    #[test]
+    fn test_log_n_selector_builder_4_step_types() {
+        // Setup
+        let builder = LogNSelectorBuilder {};
+        let mut unit = mock_compilation_unit::<Fr>();
+
+        // Assuming you have 4 step types for this test
+        // Assuming you have 3 step types for this test
+        unit.step_types.insert(
+            Uuid::now_v1(&[1, 2, 3, 4, 5, 6]).as_u128(),
+            Rc::new(StepType::new(
+                Uuid::now_v1(&[1, 2, 3, 4, 5, 6]).as_u128(),
+                "StepType1".to_owned(),
+            )),
+        );
+        unit.step_types.insert(
+            Uuid::now_v1(&[1, 2, 3, 4, 5, 6]).as_u128(),
+            Rc::new(StepType::new(
+                Uuid::now_v1(&[1, 2, 3, 4, 5, 6]).as_u128(),
+                "StepType2".to_owned(),
+            )),
+        );
+        unit.step_types.insert(
+            Uuid::now_v1(&[1, 2, 3, 4, 5, 6]).as_u128(),
+            Rc::new(StepType::new(
+                Uuid::now_v1(&[1, 2, 3, 4, 5, 6]).as_u128(),
+                "StepType3".to_owned(),
+            )),
+        );
+        unit.step_types.insert(
+            Uuid::now_v1(&[1, 2, 3, 4, 5, 6]).as_u128(),
+            Rc::new(StepType::new(
+                Uuid::now_v1(&[1, 2, 3, 4, 5, 6]).as_u128(),
+                "StepType43".to_owned(),
+            )),
+        );
+
+        builder.build(&mut unit);
+
+        // Assert number of columns
+        assert_eq!(unit.columns.len(), 3);
+        assert_eq!(unit.selector.columns.len(), 3);
+
+        // TODO: Assert expressions
+
+        // TODO: Assert assignments
+
+    }
+
+    #[test]
+    fn test_log_n_selector_builder_10_step_types() {
+        // Setup
+        let builder = LogNSelectorBuilder {};
+        let mut unit = mock_compilation_unit::<Fr>();
+
+        // Assuming you have 10 step types for this test
+        let n_step_types = 10;
+
+        for i in 0..n_step_types {
+            unit.step_types.insert(
+                Uuid::now_v1(&[1, 2, 3, 4, 5, 6]).as_u128(),
+                Rc::new(StepType::new(
+                    Uuid::now_v1(&[1, 2, 3, 4, 5, 6]).as_u128(),
+                    format!("StepType{}", i),
+                )),
+            );
+        }
+
+        builder.build(&mut unit);
+
+        // Assert number of columns
+        let n_cols = (n_step_types as f64).log2().ceil() as usize;
+        assert_eq!(unit.columns.len(), n_cols);
+        assert_eq!(unit.selector.columns.len(), n_cols);
+
+        // TODO: Assert expressions
+
+        // TODO: Assert assignments
+
+    }
+}