added ITE optimization

examples/C/mpc/playground.c

@@ -1,7 +1,7 @@
 int main(__attribute__((private(0))) int a, __attribute__((private(1))) int b) { 
-    int c[1];
     if (a < b) {
-        c[0] = 1;
+        return 1;
+    } else {
+        return 2;
     }
-    return c[0];
 }

examples/circ.rs

@@ -254,6 +254,7 @@ fn main() {
                     // The linear scan pass produces more tuples, that must be eliminated
                     Opt::Tuple,
                     Opt::ConstantFold(Box::new(ignore.clone())),
+                    Opt::Ite,
                     // Inline Function Calls
                     // Opt::Link,
                     // Opt::Tuple,

scripts/build_mpc_c_test.zsh

@@ -36,7 +36,7 @@ function mpc_test_bool  {
     RUST_BACKTRACE=1 measure_time $BIN --parties $parties $cpath mpc --cost-model "hycc" --selection-scheme "b"
 }
 
-#  mpc_test 2 ./examples/C/mpc/playground.c
+# mpc_test 2 ./examples/C/mpc/playground.c
 
 # build mpc arithmetic tests
 mpc_test 2 ./examples/C/mpc/unit_tests/arithmetic_tests/2pc_add.c

src/ir/opt/ite.rs

@@ -0,0 +1,48 @@
+//! Rewrite ITE terms
+
+use crate::ir::opt::cfold::fold;
+use crate::ir::opt::visit::RewritePass;
+use crate::ir::term::*;
+
+/// ITE cache.
+#[derive(Default)]
+struct IteRewriter;
+
+impl RewritePass for IteRewriter {
+    fn visit<F: Fn() -> Vec<Term>>(
+        &mut self,
+        _computation: &mut Computation,
+        orig: &Term,
+        rewritten_children: F,
+    ) -> Option<Term> {
+        let cs = rewritten_children();
+        match &orig.op {
+            Op::Ite => {
+                let sel = cs[0].clone();
+                let t = cs[1].clone();
+                let f = cs[2].clone();
+                match f.op {
+                    Op::Ite => {
+                        let child_sel = f.cs[0].clone();
+                        let child_t = f.cs[1].clone();
+                        if sel
+                            == term![AND; term![Op::Not; child_sel.clone()], term![Op::Not; child_sel.clone()]]
+                        {
+                            Some(term![Op::Ite; child_sel, child_t, t])
+                        } else {
+                            None
+                        }
+                    }
+                    _ => None,
+                }
+            }
+            _ => None,
+        }
+    }
+}
+
+/// Binarize (expand) n-ary terms.
+pub fn rewrite_ites(c: &mut Computation) {
+    let mut pass = IteRewriter;
+    pass.traverse(c);
+}

src/ir/opt/mod.rs

@@ -3,6 +3,7 @@ pub mod binarize;
 pub mod cfold;
 pub mod flat;
 pub mod inline;
+pub mod ite;
 pub mod link;
 pub mod mem;
 pub mod scalarize_vars;
@@ -38,6 +39,8 @@ pub enum Opt {
     FlattenAssertions,
     /// Find outputs like `(= variable term)`, and substitute out `variable`
     Inline,
+    /// Ite peephole optimizations
+    Ite,
     /// Link function calls
     Link,
     /// Eliminate tuples
@@ -108,6 +111,9 @@ pub fn opt<I: IntoIterator<Item = Opt>>(mut fs: Functions, optimizations: I) ->
                 Opt::Binarize => {
                     binarize::binarize(comp);
                 }
+                Opt::Ite => {
+                    ite::rewrite_ites(comp);
+                }
                 Opt::Inline => {
                     let public_inputs = comp
                         .metadata