[filetests]: attempt to compile run-command -> wat -> zkasm

cranelift/filetests/src/test_run_zkasm.rs

@@ -11,6 +11,7 @@ use cranelift_reader::{parse_run_command, RunCommand, TestCommand, TestFile};
 use std::borrow::Cow;
 
 use crate::zkasm_codegen;
+use crate::zkasm_runner::{self, ExecutionStatus};
 
 struct TestRunZkasm;
 
@@ -47,6 +48,7 @@ impl SubTest for TestRunZkasm {
     ) -> anyhow::Result<()> {
         let mut zkasm_functions: Vec<Vec<String>> = Vec::new();
         let mut invocations: Vec<Vec<String>> = Vec::new();
+        let mut invoke_names: Vec<String> = Vec::new();
         for (func, details) in &testfile.functions {
             zkasm_functions.push(zkasm_codegen::compile_clif_function(func));
             for comment in details.comments.iter() {
@@ -56,22 +58,30 @@ impl SubTest for TestRunZkasm {
                             todo!()
                         }
                         RunCommand::Run(invoke, compare, expected) => {
+                            invoke_names.push(zkasm_codegen::invoke_name(&invoke));
                             invocations
                                 .push(zkasm_codegen::compile_invocation(invoke, compare, expected));
                         }
                     }
                 }
             }
         }
-        let zkasm_program = zkasm_codegen::build_test_zkasm(zkasm_functions, invocations);
+        let main = zkasm_codegen::build_main(invoke_names);
+        let zkasm_program = zkasm_codegen::build_test_zkasm(zkasm_functions, invocations, main);
         println!("{}", zkasm_program);
-        // TODO: instead of printing run program, using something like this:
-        // match zkasm_runner::run_zkasm(&zkasm_program) {
-        //     // TODO: Probably here is a good place to collect info generated by assert-hostfunction
-        //     // and somehow show it
-        //     Ok(_) => Ok(()),
-        //     Err(e) => Err(e),
-        // }
+        let execution_result = zkasm_runner::run_zkasm(&zkasm_program).unwrap();
+        match execution_result.error {
+            Some(err) => panic!("Zkasm runtime error: {}", err),
+            None => (),
+        };
+        // TODO: Maybe it is better to create "Unknown error" as error in case when
+        // status is RuntimeError and current error is None, and remove this filed?
+        match execution_result.status {
+            ExecutionStatus::Success => (),
+            ExecutionStatus::RuntimeError => panic!("Unknown zkasm runtime error"),
+        };
+        // TODO: Probably here is a good place to collect info generated by assert-hostfunction
+        // and somehow show it
         Ok(())
     }
 

cranelift/filetests/src/zkasm_codegen.rs

@@ -181,6 +181,8 @@ fn fix_relocs(
             if let FinalizedRelocTarget::ExternalName(ExternalName::User(name)) = reloc.target {
                 let name = &params.user_named_funcs()[name];
                 if name.index == 0 {
+                    // Codegen line after migrating to new assert:
+                    // b"  $${assert_eq(A, B, label)}".to_vec()
                     b"  B :ASSERT".to_vec()
                 } else {
                     format!("  zkPC + 2 => RR\n  :JMP(function_{})", name.index)
@@ -281,68 +283,158 @@ pub fn compile_clif_function(func: &Function) -> Vec<String> {
     // TODO: I believe it can be done more beautiful way
     let mut funcname = func.name.to_string();
     funcname.remove(0);
-    funcname.push(':');
-    let mut res = vec![funcname];
+    let mut res = vec![format!("{}:", funcname)];
     res.append(&mut lines);
+    res.into_iter()
+        .map(|s| s.replace("label", &format!("label_{}", funcname)))
+        .collect()
+}
+
+/// Builds main for test program
+pub fn build_main(invoke_names: Vec<String>) -> Vec<String> {
+    let mut res = vec![
+        "main:".to_string(),
+        "  SP - 1 => SP".to_string(),
+        "  RR :MSTORE(SP)".to_string(),
+    ];
+    for name in invoke_names {
+        res.push("  zkPC + 2 => RR".to_string());
+        res.push(format!("  :JMP({})", name));
+    }
+    res.push("  $ => RR :MLOAD(SP)".to_string());
+    res.push("  SP + 1 => SP".to_string());
+    res.push("  :JMP(RR)".to_string());
     res
 }
 
-// TODO: this function should be much rewrited,
-// now it works for one very basic case:
-// Simple progam which don't contain globals or some other speciefic preamble\postamble
-// Program don't need helper functions (for example 2-exp.zkasm)
-// How to fix it? Use generate_preamble and provide correct inputs for it.
-/// Builds zkASM used in filetests.
-pub fn build_test_zkasm(functions: Vec<Vec<String>>, invocations: Vec<Vec<String>>) -> String {
+/// Generate invoke name in format <function_name>_<arg>_..._<arg>
+pub fn invoke_name(invoke: &Invocation) -> String {
+    let mut res = invoke.func.clone();
+    for arg in &invoke.args {
+        res.push_str(&format!("_{}", arg));
+    }
+    res
+}
+
+/// Assembles all parts of zkasm test program together
+pub fn build_test_zkasm(
+    functions: Vec<Vec<String>>,
+    invocations: Vec<Vec<String>>,
+    main: Vec<String>,
+) -> String {
     // TODO: use generate_preamble to get preamble
     let preamble = "\
 start:
+  0xffff => SP
   zkPC + 2 => RR
     :JMP(main)
     :JMP(finalizeExecution)";
-    let mut main = vec![
-        "main:".to_string(),
-        "  SP - 1 => SP".to_string(),
-        "  RR :MSTORE(SP)".to_string(),
-    ];
-    for invocation in invocations {
-        main.extend(invocation);
-    }
-    main.push("  SP - 1 => SP".to_string());
-    main.push("  :JMP(RR)".to_string());
     let mut postamble = generate_postamble();
     let mut program = vec![preamble.to_string()];
-    program.append(&mut main);
+    program.extend(main);
+    for inv in invocations {
+        program.extend(inv);
+    }
     for foo in functions {
         program.extend(foo);
     }
     program.append(&mut postamble);
     program.join("\n")
 }
 
+fn runcommand_to_wasm(
+    invoke: Invocation,
+    _compare: Comparison,
+    expected: Vec<DataValue>,
+) -> String {
+    // TODO: support different amounts of outputs
+    let res_bitness = match expected[0] {
+        DataValue::I32(_) => "i32",
+        DataValue::I64(_) => "i64",
+        _ => unimplemented!(),
+    };
+    let func_name = invoke.func;
+    let expected_result = expected[0].clone();
+    let mut arg_types = String::new();
+    let mut args_pushing = String::new();
+    for arg in &invoke.args {
+        let arg_type = match arg {
+            DataValue::I32(_) => "i32",
+            DataValue::I64(_) => "i64",
+            _ => unimplemented!(),
+        };
+        arg_types.push_str(arg_type);
+        arg_types.push_str(" ");
+
+        args_pushing.push_str(&format!("{arg_type}.const {arg}\n        "));
+    }
+    if arg_types.len() > 0 {
+        arg_types.pop();
+    }
+    // TODO: remove line with 8 whitespaces in the end of args_pushing
+    let wat_code = format!(
+        r#"(module
+    (import "env" "assert_eq" (func $assert_eq (param {res_bitness} {res_bitness})))
+    (import "env" "{func_name}" (func ${func_name} (param {arg_types}) (result {res_bitness})))
+    (func $main
+        {args_pushing}
+        call ${func_name}
+        {res_bitness}.const {expected_result}
+        call $assert_eq
+    )
+    (start $main)
+)"#,
+        args_pushing = args_pushing,
+        arg_types = arg_types,
+        res_bitness = res_bitness,
+        func_name = func_name,
+        expected_result = expected_result,
+    );
+    wat_code.to_string()
+}
+
 /// Compiles a invocation.
 pub fn compile_invocation(
     invoke: Invocation,
     compare: Comparison,
     expected: Vec<DataValue>,
 ) -> Vec<String> {
-    // Here I assume that each "function" in zkasm gets it's arguments from first N registers
-    // and put result in A.
-    // TODO: should be more robust way to do it, we need somehow define inputs and outputs
-    let mut res: Vec<String> = Default::default();
-    let registers = vec!["A", "B", "C", "D", "E"];
+    // TODO: don't do this clones
+    let cmp = if compare == Comparison::Equals {
+        Comparison::Equals
+    } else {
+        Comparison::NotEquals
+    };
+    let inv = Invocation {
+        func: invoke.func.clone(),
+        args: invoke.args.clone(),
+    };
+    let wat = runcommand_to_wasm(inv, cmp, expected.clone());
+    let wasm_module = wat::parse_str(wat).unwrap();
+
+    let settings = ZkasmSettings::default();
 
-    let args = invoke.args;
+    // TODO: we should not use generate_zkasm itself, but a bit changed version?
+    let generated: Vec<String> = generate_zkasm(&settings, &wasm_module)
+        .split("\n")
+        .map(|s| s.to_string())
+        .collect();
+    let new_label = invoke_name(&invoke);
     let funcname = invoke.func;
 
-    // TODO: here we should pay attention to type of DataValue (I64 or I32)
-    for (idx, arg) in args.iter().enumerate() {
-        res.push(format!("  {} => {}", arg, registers[idx]))
-    }
-    res.push(format!("    :JMP({})", funcname));
-    // TODO: handle functions with multiple outputs
-    res.push(format!("  {} => B", expected[0]));
-    // TODO: replace with call to host function
-    res.push(format!("  CALL AWESOME ASSERT ({})", compare));
-    res
+    // TODO: will this always function_2?
+    let start_index = generated.iter().position(|r| r == "function_2:").unwrap();
+    let end_index = generated.iter().rposition(|r| r == "  :JMP(RR)").unwrap();
+    let mut generated_function = generated[start_index..=end_index].to_vec();
+
+    generated_function[0] = format!("{}:", new_label);
+    // TODO: will this always function_1?
+    let generated_replaced: Vec<String> = generated_function
+        .iter()
+        .map(|s| {
+            s.replace("label", &new_label)
+                .replace("function_1", &funcname)
+        })
+        .collect();
+    generated_replaced
 }

cranelift/filetests/src/zkasm_runner.rs

@@ -28,9 +28,9 @@ pub struct ExecutionResult {
     /// Path to the main zkAsm file that was executed.
     path: String,
     /// Status of the execution.
-    status: ExecutionStatus,
+    pub status: ExecutionStatus,
     /// Error message in case the execution failed.
-    error: Option<String>,
+    pub error: Option<String>,
     /// Profiling information about this execution.
     /// Only populated for the successful executions.
     counters: Option<Counters>,
@@ -76,7 +76,7 @@ pub fn run_zkasm_path(input_path: &Path) -> anyhow::Result<Vec<ExecutionResult>>
     let mut output_file = NamedTempFile::new()?;
     let common_args = [
         "--prefix",
-        "../../tests/zkasm",
+        "../tests/zkasm",
         "test",
         input_path.to_str().unwrap(),
         output_file.path().to_str().unwrap(),