Added support for multiple outputs in test cases

examples/ZoKrates/mpc/array_tests/2pc_array_ret.zok

@@ -0,0 +1,2 @@
+def main(private<1> u32 a, private<2> u32 b) -> u32[2]:
+    return [a, b]

examples/ZoKrates/mpc/array_tests/2pc_array_sum.zok

@@ -0,0 +1,3 @@
+def main(private<1> u32 a, private<2> u32 b) -> u32:
+    u32[2] c = [a, b]
+    return c[0] + c[1] 

examples/circ.rs

@@ -61,7 +61,7 @@ fn main() {
         Mode::Opt => opt(cs, vec![Opt::ConstantFold]),
         Mode::Mpc(_) => opt(
             cs,
-            vec![], //Opt::Sha, Opt::ConstantFold, Opt::Mem, Opt::ConstantFold],
+            vec![Opt::Sha, Opt::ConstantFold, Opt::Mem, Opt::ConstantFold],
         ),
         Mode::Proof => opt(
             cs,

scripts/build_mpc_zokrates_test.zsh

@@ -56,14 +56,16 @@ mpc_test 2 ./examples/ZoKrates/mpc/ite_tests/2pc_ite_ret_bool.zok
 mpc_test 2 ./examples/ZoKrates/mpc/ite_tests/2pc_ite_ret_int.zok
 
 
-# build mpc misc tests
-mpc_test 2 ./examples/ZoKrates/mpc/2pc_millionaire.zok
-
-
 # build mpc const tests
 mpc_test 2 ./examples/ZoKrates/mpc/const_tests/2pc_const_arith.zok
 mpc_test 2 ./examples/ZoKrates/mpc/const_tests/2pc_const_bool.zok
 
+
+# build mpc array tests
+mpc_test 2 ./examples/ZoKrates/mpc/array_tests/2pc_array_sum.zok
+mpc_test 2 ./examples/ZoKrates/mpc/array_tests/2pc_array_ret.zok
+
+
 # build mpc misc tests
 mpc_test 2 ./examples/ZoKrates/mpc/2pc_millionaire.zok
 mpc_test 2 ./examples/ZoKrates/mpc/2pc_conv.zok

scripts/test_aby.py

@@ -428,6 +428,23 @@ ite_tests = [
     ],
 ]
 
+arr_tests = [
+    [
+        "Array sum test",
+        3,
+        "./third_party/ABY/build/bin/2pc_array_sum_test",
+        {"a": 2, "b": 0},
+        {"a": 0, "b": 1},
+    ], 
+    [
+        "Array ret test",
+        "2\n1",
+        "./third_party/ABY/build/bin/2pc_array_ret_test",
+        {"a": 2, "b": 0},
+        {"a": 0, "b": 1},
+    ], 
+]
+
 
 misc_tests = [
     [
@@ -468,9 +485,6 @@ def flatten_args(args: dict) -> list:
         flat_args.append(str(v))
     return flat_args
 
-def build_expected(expected) -> str:
-    return "output: "+str(expected)
-
 def build_server_cmd(exec: str, args: dict) -> List[str]:
     return [exec, "-r", "0", "-i"] + flatten_args(args)
 
@@ -496,8 +510,6 @@ def run_test(desc: str, expected: str, server_cmd: List[str], client_cmd: List[s
         server_out = server_out.decode('utf-8').strip()
         client_out = client_out.decode('utf-8').strip()
 
-        assert server_out.startswith("output: "), "server output did not start with \"output:\", but instead with: "+server_out
-        assert client_out.startswith("output: "), "client output did not start with \"output:\", but instead with: "+client_out
         assert server_out == client_out, "server out != client out\nserver_out: "+server_out+"\nclient_out: "+client_out
         assert server_out == expected, "output != expected\nserver_out: "+server_out+"\nexpected: "+expected
         return True
@@ -520,15 +532,16 @@ def main():
         nary_boolean_tests + \
         const_tests + \
         ite_tests + \
+        arr_tests + \
         misc_tests
-    # tests = misc_tests
+    # tests = arr_tests
 
     failed_test_descs = []
     num_retries = 3
     for test in tests:
         assert len(test) == 5, "test configurations are wrong for test: "+test[0]
         desc = test[0]
-        expected = build_expected(test[1])
+        expected = str(test[1])
         path = test[2]
         server_cmd = build_server_cmd(path, test[3])
         client_cmd = build_client_cmd(path, test[4])

src/target/aby/assignment/mod.rs

@@ -1,6 +1,6 @@
 //! Machinery for assigning operations to sharing schemes
 
-use crate::ir::term::{Computation, Op, BvNaryOp, PostOrderIter, TermMap};
+use crate::ir::term::{BvNaryOp, Computation, Op, PostOrderIter, TermMap};
 
 pub mod ilp;
 
@@ -51,7 +51,7 @@ pub fn some_arith_sharing(c: &Computation) -> SharingMap {
                     BvNaryOp::Add => (term.clone(), ShareType::Arithmetic),
                     BvNaryOp::Mul => (term.clone(), ShareType::Arithmetic),
                     _ => (term.clone(), ShareType::Boolean),
-                }
+                },
                 _ => (term.clone(), ShareType::Boolean),
             })
         })

src/target/aby/mod.rs

@@ -8,11 +8,11 @@ pub mod trans;
 /// The ABY Circuit consists of three Vec<String>: setup, circ, and closer
 /// *setup* holds code for initializing the ABY party, sharing scheme, and input values
 /// *circs* holds the lowered code from the IR to ABY Circuits
-/// *closer* holds the code for executing the ABY Circuits and printing the output value
+/// *output* holds the code for printing the output value
 pub struct ABY {
     setup: Vec<String>,
     circs: Vec<String>,
-    closer: Vec<String>,
+    output: Vec<String>,
 }
 
 impl ABY {
@@ -21,7 +21,7 @@ impl ABY {
         ABY {
             setup: Vec::new(),
             circs: Vec::new(),
-            closer: Vec::new(),
+            output: Vec::new(),
         }
     }
 }

src/target/aby/output.rs

@@ -96,7 +96,7 @@ fn write_h_file(filename: &String) {
 }
 
 /// Using the cpp_template.txt, write the .cpp file for the new test case
-fn write_circ_file(filename: &String, circ: String) {
+fn write_circ_file(filename: &String, circ: String, output: &String) {
     let template = fs::read_to_string("third_party/ABY_templates/cpp_template.txt")
         .expect("Unable to read file");
     let path = format!(
@@ -108,7 +108,8 @@ fn write_circ_file(filename: &String, circ: String) {
         path.clone(),
         template
             .replace("{fn}", &*filename)
-            .replace("{circ}", &circ),
+            .replace("{circ}", &circ)
+            .replace("{output}", &output),
     )
     .expect("Failed to write to cmake file");
 }
@@ -121,6 +122,6 @@ pub fn write_aby_exec(aby: ABY, path_buf: PathBuf) {
     write_test_cmake_file(&filename);
     write_test_file(&filename);
     write_h_file(&filename);
-    let circ_str = aby.setup.join("\n\t") + &aby.circs.join("\n\t") + &aby.closer.join("\n\t");
-    write_circ_file(&filename, circ_str);
+    let circ_str = aby.setup.join("\n\t") + &aby.circs.join("\n\t");
+    write_circ_file(&filename, circ_str, &aby.output.join("\n\t"));
 }

src/target/aby/trans.rs

@@ -5,10 +5,8 @@
 //! [Link to comment in EzPC Compiler](https://github.com/mpc-msri/EzPC/blob/da94a982709123c8186d27c9c93e27f243d85f0e/EzPC/EzPC/codegen.ml)
 
 use crate::ir::term::*;
-use crate::target::aby::assignment::ilp::assign;
-use crate::target::aby::assignment::{
-    some_arith_sharing, ShareType, SharingMap,
-};
+// use crate::target::aby::assignment::ilp::assign;
+use crate::target::aby::assignment::{some_arith_sharing, ShareType, SharingMap};
 use crate::target::aby::*;
 
 const NO_ROLE: u8 = u8::MAX;
@@ -28,6 +26,7 @@ struct ToABY {
     inputs: TermMap<Option<PartyId>>,
     cache: TermMap<EmbeddedTerm>,
     s_map: SharingMap,
+    output_counter: u8,
 }
 
 impl ToABY {
@@ -38,6 +37,7 @@ impl ToABY {
             inputs: TermMap::new(),
             cache: TermMap::new(),
             s_map: s_map,
+            output_counter: 0,
         }
     }
 
@@ -500,16 +500,31 @@ impl ToABY {
     /// the circuit to a share      
     ///
     /// Return a String of the resulting Circuit
-    fn add_output_gate(&mut self, circ: String) -> String {
-        format!("\ts_out = bcirc->PutOUTGate({}, ALL);\n", circ)
+    fn add_output_gate(&mut self, t: Term, circ: String) -> String {
+        format!(
+            "\ts_out_{} = {}->PutOUTGate({}, ALL);\n",
+            self.output_counter,
+            self.get_sharetype_circ(t),
+            circ
+        )
     }
 
     /// Given a term `t`, lower `t` to ABY Circuits
     fn lower(&mut self, t: Term) {
-        println!("{:?}", t);
-        let mut output_circ = self.embed(t);
-        output_circ = self.add_output_gate(output_circ);
+        let mut output_circ = self.embed(t.clone());
+
+        output_circ = self.add_output_gate(t, output_circ);
         self.aby.circs.push(output_circ);
+        self.aby
+            .setup
+            .push(format!("share *s_out_{};", self.output_counter));
+        self.aby.output.push(format!(
+            "uint32_t output_{} = s_out_{}->get_clear_value<uint32_t>();\n 
+            \tstd::cout << output_{} << std::endl;",
+            self.output_counter, self.output_counter, self.output_counter
+        ));
+
+        self.output_counter += 1;
     }
 }
 

third_party/ABY_templates/cpp_template.txt

@@ -15,15 +15,14 @@ int32_t test_{fn}_circuit(std::map<std::string, std::string> params, e_role role
 	Circuit* bcirc = sharings[S_BOOL]->GetCircuitBuildRoutine();
 	Circuit* ycirc = sharings[S_YAO]->GetCircuitBuildRoutine();
 
-	share *s_out;
-
 	// compiled circuit
 	{circ}
 
-	// output
 	party->ExecCircuit();
-	uint32_t output = s_out->get_clear_value<uint32_t>();
-	std::cout << "output: " << output << std::endl;
+
+	// output
+	{output}
+
 	delete party;
 	return 0;
 }