From 51c195410f57143232226ec40de8e79a241d80d8 Mon Sep 17 00:00:00 2001
From: Viktar Makouski <mcviktor@tut.by>
Date: Tue, 27 Feb 2024 20:23:43 +0300
Subject: [PATCH 1/9] [filetests]: attempt to compile run-command -> wat ->
 zkasm. WIP

---
 cranelift/filetests/src/zkasm_codegen.rs | 82 ++++++++++++++++++------
 1 file changed, 62 insertions(+), 20 deletions(-)

diff --git a/cranelift/filetests/src/zkasm_codegen.rs b/cranelift/filetests/src/zkasm_codegen.rs
index 9c4cb1bc3e56..bf7d3c60ca73 100644
--- a/cranelift/filetests/src/zkasm_codegen.rs
+++ b/cranelift/filetests/src/zkasm_codegen.rs
@@ -159,7 +159,7 @@ fn fix_relocs(
             if let FinalizedRelocTarget::ExternalName(ExternalName::User(name)) = reloc.target {
                 let name = &params.user_named_funcs()[name];
                 if name.index == 0 {
-                    b"  B :ASSERT".to_vec()
+                    b"  $${assert_eq(A, B, label)}".to_vec()
                 } else {
                     format!("  zkPC + 2 => RR\n  :JMP(function_{})", name.index)
                         .as_bytes()
@@ -296,28 +296,70 @@ start:
     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()
+}
+
+
 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"];
-
-    let args = invoke.args;
-    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
+    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());
+    // TODO: remove this debug prints
+    println!("{}", wat);
+    println!("Translated wasm:\n============================================");
+    let wasm_module = wat::parse_str(wat).unwrap();
+    println!("{}\n================================ End of compiled wasm", generate_zkasm(&wasm_module));
+    // TODO: we should not use generate_zkasm itself, but a bit changed version.
+    generate_zkasm(&wasm_module).split("\n").map(|s| s.to_string()).collect()
+    // TODO: I think here is a good place to replace function_<number> in generated zkasm
+    // to name of invoked function, and replace default label in assert helper to something
+    // meaningful
 }

From 7e91b252d7261e7e7222efb05df8f6694e685099 Mon Sep 17 00:00:00 2001
From: Viktar Makouski <mcviktor@tut.by>
Date: Thu, 29 Feb 2024 17:41:36 +0300
Subject: [PATCH 2/9] [filetests]: generate labels

---
 cranelift/filetests/src/test_run_zkasm.rs |  18 ++--
 cranelift/filetests/src/zkasm_codegen.rs  | 105 +++++++++++++++-------
 2 files changed, 81 insertions(+), 42 deletions(-)

diff --git a/cranelift/filetests/src/test_run_zkasm.rs b/cranelift/filetests/src/test_run_zkasm.rs
index ec7d145778ee..81f864324f62 100644
--- a/cranelift/filetests/src/test_run_zkasm.rs
+++ b/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;
 
 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,20 @@ 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));
+                                .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),
-        // }
+        zkasm_runner::run_zkasm(&zkasm_program).unwrap();
+        // TODO: Probably here is a good place to collect info generated by assert-hostfunction
+        // and somehow show it
         Ok(())
     }
 
diff --git a/cranelift/filetests/src/zkasm_codegen.rs b/cranelift/filetests/src/zkasm_codegen.rs
index bf7d3c60ca73..de94250cee44 100644
--- a/cranelift/filetests/src/zkasm_codegen.rs
+++ b/cranelift/filetests/src/zkasm_codegen.rs
@@ -258,37 +258,52 @@ 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()
+}
+
+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.
-pub fn build_test_zkasm(functions: Vec<Vec<String>>, invocations: Vec<Vec<String>>) -> String {
+pub fn invoke_name(invoke: &Invocation) -> String {
+    let mut res = invoke.func.clone();
+    for arg in &invoke.args {
+        res.push_str(&format!("_{}", arg));
+    }
+    res
+}
+
+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:
   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);
     }
@@ -296,12 +311,16 @@ start:
     program.join("\n")
 }
 
-fn runcommand_to_wasm(invoke: Invocation, _compare: Comparison, expected: Vec<DataValue>) -> String {
+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!()
+        _ => unimplemented!(),
     };
     let func_name = invoke.func;
     let expected_result = expected[0].clone();
@@ -311,7 +330,7 @@ fn runcommand_to_wasm(invoke: Invocation, _compare: Comparison, expected: Vec<Da
         let arg_type = match arg {
             DataValue::I32(_) => "i32",
             DataValue::I64(_) => "i64",
-            _ => unimplemented!()
+            _ => unimplemented!(),
         };
         arg_types.push_str(arg_type);
         arg_types.push_str(" ");
@@ -343,23 +362,43 @@ fn runcommand_to_wasm(invoke: Invocation, _compare: Comparison, expected: Vec<Da
     wat_code.to_string()
 }
 
-
 pub fn compile_invocation(
     invoke: Invocation,
     compare: Comparison,
     expected: Vec<DataValue>,
 ) -> Vec<String> {
-    let cmp =  if compare == Comparison::Equals {Comparison::Equals} else {Comparison::NotEquals};
-    let inv = Invocation {func: invoke.func.clone(), args: invoke.args.clone()};
+    // 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());
-    // TODO: remove this debug prints
-    println!("{}", wat);
-    println!("Translated wasm:\n============================================");
     let wasm_module = wat::parse_str(wat).unwrap();
-    println!("{}\n================================ End of compiled wasm", generate_zkasm(&wasm_module));
+
     // TODO: we should not use generate_zkasm itself, but a bit changed version.
-    generate_zkasm(&wasm_module).split("\n").map(|s| s.to_string()).collect()
-    // TODO: I think here is a good place to replace function_<number> in generated zkasm
-    // to name of invoked function, and replace default label in assert helper to something
-    // meaningful
+    let generated: Vec<String> = generate_zkasm(&wasm_module)
+        .split("\n")
+        .map(|s| s.to_string())
+        .collect();
+    let new_label = invoke_name(&invoke);
+    let funcname = invoke.func;
+
+    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);
+    let generated_replaced: Vec<String> = generated_function
+        .iter()
+        .map(|s| {
+            s.replace("label", &new_label)
+                .replace("function_1", &funcname)
+        })
+        .collect();
+    generated_replaced
 }

From 92b60d6df8c3ea5fdd9f9933e2d332782387e40b Mon Sep 17 00:00:00 2001
From: Viktar Makouski <mcviktor@tut.by>
Date: Fri, 1 Mar 2024 13:26:56 +0300
Subject: [PATCH 3/9] refactoring

---
 cranelift/filetests/src/test_run_zkasm.rs | 8 ++++++--
 cranelift/filetests/src/zkasm_codegen.rs  | 1 +
 cranelift/filetests/src/zkasm_runner.rs   | 6 +++---
 3 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/cranelift/filetests/src/test_run_zkasm.rs b/cranelift/filetests/src/test_run_zkasm.rs
index 81f864324f62..c38616d9a44e 100644
--- a/cranelift/filetests/src/test_run_zkasm.rs
+++ b/cranelift/filetests/src/test_run_zkasm.rs
@@ -11,7 +11,7 @@ use cranelift_reader::{parse_run_command, RunCommand, TestCommand, TestFile};
 use std::borrow::Cow;
 
 use crate::zkasm_codegen;
-use crate::zkasm_runner;
+use crate::zkasm_runner::{self, ExecutionResult};
 
 struct TestRunZkasm;
 
@@ -69,7 +69,11 @@ impl SubTest for TestRunZkasm {
         let main = zkasm_codegen::build_main(invoke_names);
         let zkasm_program = zkasm_codegen::build_test_zkasm(zkasm_functions, invocations, main);
         println!("{}", zkasm_program);
-        zkasm_runner::run_zkasm(&zkasm_program).unwrap();
+        let execution_result = zkasm_runner::run_zkasm(&zkasm_program).unwrap();
+        match execution_result.error {
+            Some(err) => println!("{}", err),
+            None => (),
+        };
         // TODO: Probably here is a good place to collect info generated by assert-hostfunction
         // and somehow show it
         Ok(())
diff --git a/cranelift/filetests/src/zkasm_codegen.rs b/cranelift/filetests/src/zkasm_codegen.rs
index de94250cee44..0a6f99c0d0f3 100644
--- a/cranelift/filetests/src/zkasm_codegen.rs
+++ b/cranelift/filetests/src/zkasm_codegen.rs
@@ -295,6 +295,7 @@ pub fn build_test_zkasm(
     // TODO: use generate_preamble to get preamble
     let preamble = "\
 start:
+  0xffff => SP
   zkPC + 2 => RR
     :JMP(main)
     :JMP(finalizeExecution)";
diff --git a/cranelift/filetests/src/zkasm_runner.rs b/cranelift/filetests/src/zkasm_runner.rs
index df14d5d389d5..dbbb9f542d58 100644
--- a/cranelift/filetests/src/zkasm_runner.rs
+++ b/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(),

From dd102c6b1c21176bb97f735342266ea1d1b461f3 Mon Sep 17 00:00:00 2001
From: Viktar Makouski <mcviktor@tut.by>
Date: Fri, 1 Mar 2024 14:08:43 +0300
Subject: [PATCH 4/9] post-merge work

---
 cranelift/filetests/src/test_run_zkasm.rs | 10 +++-
 cranelift/filetests/src/zkasm_codegen.rs  | 63 ++++-------------------
 2 files changed, 18 insertions(+), 55 deletions(-)

diff --git a/cranelift/filetests/src/test_run_zkasm.rs b/cranelift/filetests/src/test_run_zkasm.rs
index c38616d9a44e..126916456e1b 100644
--- a/cranelift/filetests/src/test_run_zkasm.rs
+++ b/cranelift/filetests/src/test_run_zkasm.rs
@@ -11,7 +11,7 @@ use cranelift_reader::{parse_run_command, RunCommand, TestCommand, TestFile};
 use std::borrow::Cow;
 
 use crate::zkasm_codegen;
-use crate::zkasm_runner::{self, ExecutionResult};
+use crate::zkasm_runner::{self, ExecutionResult, ExecutionStatus};
 
 struct TestRunZkasm;
 
@@ -71,9 +71,15 @@ impl SubTest for TestRunZkasm {
         println!("{}", zkasm_program);
         let execution_result = zkasm_runner::run_zkasm(&zkasm_program).unwrap();
         match execution_result.error {
-            Some(err) => println!("{}", err),
+            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(())
diff --git a/cranelift/filetests/src/zkasm_codegen.rs b/cranelift/filetests/src/zkasm_codegen.rs
index 9149c5ba1dac..0f749e67083a 100644
--- a/cranelift/filetests/src/zkasm_codegen.rs
+++ b/cranelift/filetests/src/zkasm_codegen.rs
@@ -286,6 +286,7 @@ pub fn compile_clif_function(func: &Function) -> Vec<String> {
     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(),
@@ -302,6 +303,7 @@ pub fn build_main(invoke_names: Vec<String>) -> Vec<String> {
     res
 }
 
+/// Generate invoke name in format <function_name>_<arg>_..._<expected_result>
 pub fn invoke_name(invoke: &Invocation) -> String {
     let mut res = invoke.func.clone();
     for arg in &invoke.args {
@@ -310,6 +312,7 @@ pub fn invoke_name(invoke: &Invocation) -> String {
     res
 }
 
+/// Assembles all parts of zkasm test program together
 pub fn build_test_zkasm(
     functions: Vec<Vec<String>>,
     invocations: Vec<Vec<String>>,
@@ -386,57 +389,6 @@ fn runcommand_to_wasm(
     wat_code.to_string()
 }
 
-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,
@@ -456,19 +408,24 @@ pub fn compile_invocation(
     let wat = runcommand_to_wasm(inv, cmp, expected.clone());
     let wasm_module = wat::parse_str(wat).unwrap();
 
-    // TODO: we should not use generate_zkasm itself, but a bit changed version.
-    let generated: Vec<String> = generate_zkasm(&wasm_module)
+    let settings = ZkasmSettings::default();
+
+    // 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: 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| {

From c6fe1a66cef65a0d31679d99d4346ad6a0f764a3 Mon Sep 17 00:00:00 2001
From: Viktar Makouski <mcviktor@tut.by>
Date: Mon, 4 Mar 2024 15:56:10 +0300
Subject: [PATCH 5/9] small fixes

---
 cranelift/filetests/src/test_run_zkasm.rs |  8 ++++----
 cranelift/filetests/src/zkasm_codegen.rs  | 13 +++++++------
 2 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/cranelift/filetests/src/test_run_zkasm.rs b/cranelift/filetests/src/test_run_zkasm.rs
index 126916456e1b..94620c84aa19 100644
--- a/cranelift/filetests/src/test_run_zkasm.rs
+++ b/cranelift/filetests/src/test_run_zkasm.rs
@@ -11,7 +11,7 @@ use cranelift_reader::{parse_run_command, RunCommand, TestCommand, TestFile};
 use std::borrow::Cow;
 
 use crate::zkasm_codegen;
-use crate::zkasm_runner::{self, ExecutionResult, ExecutionStatus};
+use crate::zkasm_runner::{self, ExecutionStatus};
 
 struct TestRunZkasm;
 
@@ -60,7 +60,7 @@ impl SubTest for TestRunZkasm {
                         RunCommand::Run(invoke, compare, expected) => {
                             invoke_names.push(zkasm_codegen::invoke_name(&invoke));
                             invocations
-                                .push(zkasm_codegen::compile_invocation(invoke, compare, expected));   
+                                .push(zkasm_codegen::compile_invocation(invoke, compare, expected));
                         }
                     }
                 }
@@ -76,9 +76,9 @@ impl SubTest for TestRunZkasm {
         };
         // 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 { 
+        match execution_result.status {
             ExecutionStatus::Success => (),
-            ExecutionStatus::RuntimeError => panic!("Unknown zkasm runtime error")
+            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
diff --git a/cranelift/filetests/src/zkasm_codegen.rs b/cranelift/filetests/src/zkasm_codegen.rs
index 0f749e67083a..4cf375c4c58f 100644
--- a/cranelift/filetests/src/zkasm_codegen.rs
+++ b/cranelift/filetests/src/zkasm_codegen.rs
@@ -283,7 +283,9 @@ pub fn compile_clif_function(func: &Function) -> Vec<String> {
     funcname.remove(0);
     let mut res = vec![format!("{}:", funcname)];
     res.append(&mut lines);
-    res.into_iter().map(|s| s.replace("label", &format!("label_{}", funcname))).collect()
+    res.into_iter()
+        .map(|s| s.replace("label", &format!("label_{}", funcname)))
+        .collect()
 }
 
 /// Builds main for test program
@@ -303,7 +305,7 @@ pub fn build_main(invoke_names: Vec<String>) -> Vec<String> {
     res
 }
 
-/// Generate invoke name in format <function_name>_<arg>_..._<expected_result>
+/// 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 {
@@ -411,19 +413,18 @@ pub fn compile_invocation(
     let settings = ZkasmSettings::default();
 
     // TODO: we should not use generate_zkasm itself, but a bit changed version?
-    let generated: Vec<String> = generate_zkasm(&settings,&wasm_module)
+    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: will this always function_2? 
+    // 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

From edae5e7a9f0cc9f91b4172fa8918bccd15c2d0b6 Mon Sep 17 00:00:00 2001
From: Viktar Makouski <mcviktor@tut.by>
Date: Mon, 4 Mar 2024 16:51:32 +0300
Subject: [PATCH 6/9] revert assert update

---
 cranelift/filetests/src/zkasm_codegen.rs | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/cranelift/filetests/src/zkasm_codegen.rs b/cranelift/filetests/src/zkasm_codegen.rs
index 4cf375c4c58f..f0ff5c836ef1 100644
--- a/cranelift/filetests/src/zkasm_codegen.rs
+++ b/cranelift/filetests/src/zkasm_codegen.rs
@@ -181,7 +181,9 @@ fn fix_relocs(
             if let FinalizedRelocTarget::ExternalName(ExternalName::User(name)) = reloc.target {
                 let name = &params.user_named_funcs()[name];
                 if name.index == 0 {
-                    b"  $${assert_eq(A, B, label)}".to_vec()
+                    // 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)
                         .as_bytes()

From 5a2ecf8cbedd6cc860ad067e34c7dcf5c8bf7b3e Mon Sep 17 00:00:00 2001
From: Viktar Makouski <mcviktor@tut.by>
Date: Wed, 6 Mar 2024 14:49:48 +0300
Subject: [PATCH 7/9] postmerge fixes

---
 cranelift/filetests/src/zkasm_codegen.rs | 10 +++-------
 1 file changed, 3 insertions(+), 7 deletions(-)

diff --git a/cranelift/filetests/src/zkasm_codegen.rs b/cranelift/filetests/src/zkasm_codegen.rs
index c7d410d0eab3..9f88565d355c 100644
--- a/cranelift/filetests/src/zkasm_codegen.rs
+++ b/cranelift/filetests/src/zkasm_codegen.rs
@@ -349,11 +349,7 @@ start:
     program.join("\n")
 }
 
-fn runcommand_to_wasm(
-    invoke: Invocation,
-    _compare: Comparison,
-    expected: Vec<DataValue>,
-) -> String {
+fn runcommand_to_wat(invoke: Invocation, _compare: Comparison, expected: Vec<DataValue>) -> String {
     // TODO: support different amounts of outputs
     let res_bitness = match expected[0] {
         DataValue::I32(_) => "i32",
@@ -389,7 +385,7 @@ fn runcommand_to_wasm(
         {res_bitness}.const {expected_result}
         call $assert_eq
     )
-    (start $main)
+    (export "main" (func $main))
 )"#,
         args_pushing = args_pushing,
         arg_types = arg_types,
@@ -416,7 +412,7 @@ pub fn compile_invocation(
         func: invoke.func.clone(),
         args: invoke.args.clone(),
     };
-    let wat = runcommand_to_wasm(inv, cmp, expected.clone());
+    let wat = runcommand_to_wat(inv, cmp, expected.clone());
     let wasm_module = wat::parse_str(wat).unwrap();
 
     let settings = ZkasmSettings::default();

From a91348b19ba019233ade45ca186f9dd06cb76000 Mon Sep 17 00:00:00 2001
From: Viktar Makouski <mcviktor@tut.by>
Date: Thu, 7 Mar 2024 14:27:08 +0300
Subject: [PATCH 8/9] adding comments & review fixes

---
 Cargo.lock                                | 43 ++++++++++++++++++++--
 cranelift/filetests/Cargo.toml            |  1 +
 cranelift/filetests/src/test_run_zkasm.rs |  7 ++--
 cranelift/filetests/src/zkasm_codegen.rs  | 44 +++++++++++++++--------
 4 files changed, 74 insertions(+), 21 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 632c593e3071..f2111fe8c784 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -474,6 +474,15 @@ version = "0.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2da6da31387c7e4ef160ffab6d5e7f00c42626fe39aea70a7b0f1773f7dd6c1b"
 
+[[package]]
+name = "cmake"
+version = "0.1.50"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a31c789563b815f77f4250caee12365734369f942439b7defd71e18a48197130"
+dependencies = [
+ "cc",
+]
+
 [[package]]
 name = "codespan-reporting"
 version = "0.11.1"
@@ -669,6 +678,7 @@ dependencies = [
  "tempfile",
  "thiserror",
  "toml",
+ "wabt",
  "walkdir",
  "wasmparser",
  "wasmtime",
@@ -1328,6 +1338,12 @@ dependencies = [
  "stable_deref_trait",
 ]
 
+[[package]]
+name = "glob"
+version = "0.2.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8be18de09a56b60ed0edf84bc9df007e30040691af7acd1c41874faac5895bfb"
+
 [[package]]
 name = "glob"
 version = "0.3.0"
@@ -2974,6 +2990,29 @@ version = "0.9.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "49874b5167b65d7193b8aba1567f5c7d93d001cafc34600cee003eda787e483f"
 
+[[package]]
+name = "wabt"
+version = "0.10.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "00bef93d5e6c81a293bccf107cf43aa47239382f455ba14869d36695d8963b9c"
+dependencies = [
+ "serde",
+ "serde_derive",
+ "serde_json",
+ "wabt-sys",
+]
+
+[[package]]
+name = "wabt-sys"
+version = "0.8.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1a4e043159f63e16986e713e9b5e1c06043df4848565bf672e27c523864c7791"
+dependencies = [
+ "cc",
+ "cmake",
+ "glob 0.2.11",
+]
+
 [[package]]
 name = "wait-timeout"
 version = "0.2.0"
@@ -3989,7 +4028,7 @@ dependencies = [
 name = "winch-test-macros"
 version = "0.0.0"
 dependencies = [
- "glob",
+ "glob 0.3.0",
  "proc-macro2",
  "quote",
  "syn 2.0.32",
@@ -4003,7 +4042,7 @@ dependencies = [
  "capstone",
  "clap",
  "cranelift-codegen",
- "glob",
+ "glob 0.3.0",
  "serde",
  "similar",
  "target-lexicon",
diff --git a/cranelift/filetests/Cargo.toml b/cranelift/filetests/Cargo.toml
index 4a85138ac583..56acb7fcadee 100644
--- a/cranelift/filetests/Cargo.toml
+++ b/cranelift/filetests/Cargo.toml
@@ -47,3 +47,4 @@ regex = { workspace = true }
 wasmtime = { workspace = true, features = ["cranelift", "runtime"] }
 walkdir = { workspace = true }
 tempfile = { workspace = true }
+wabt = "0.10.0"
diff --git a/cranelift/filetests/src/test_run_zkasm.rs b/cranelift/filetests/src/test_run_zkasm.rs
index 94620c84aa19..223c2eb07df5 100644
--- a/cranelift/filetests/src/test_run_zkasm.rs
+++ b/cranelift/filetests/src/test_run_zkasm.rs
@@ -66,7 +66,7 @@ impl SubTest for TestRunZkasm {
                 }
             }
         }
-        let main = zkasm_codegen::build_main(invoke_names);
+        let main = zkasm_codegen::build_test_main(invoke_names);
         let zkasm_program = zkasm_codegen::build_test_zkasm(zkasm_functions, invocations, main);
         println!("{}", zkasm_program);
         let execution_result = zkasm_runner::run_zkasm(&zkasm_program).unwrap();
@@ -76,10 +76,7 @@ impl SubTest for TestRunZkasm {
         };
         // 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"),
-        };
+        debug_assert!(matches!(execution_result.status, ExecutionStatus::Success));
         // TODO: Probably here is a good place to collect info generated by assert-hostfunction
         // and somehow show it
         Ok(())
diff --git a/cranelift/filetests/src/zkasm_codegen.rs b/cranelift/filetests/src/zkasm_codegen.rs
index 9f88565d355c..af3d8bcd46fc 100644
--- a/cranelift/filetests/src/zkasm_codegen.rs
+++ b/cranelift/filetests/src/zkasm_codegen.rs
@@ -15,6 +15,8 @@ use cranelift_reader::Comparison;
 use cranelift_reader::Invocation;
 use cranelift_wasm::{translate_module, ZkasmEnvironment};
 
+use wabt::{wasm2wat, wat2wasm};
+
 /// ISA specific settings for zkASM codegen.
 #[derive(Default, Debug)]
 pub struct ZkasmSettings {
@@ -188,7 +190,7 @@ 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:
+                    // TODO(#246): Codegen line after migrating to new assert:
                     // b"  $${assert_eq(A, B, label)}".to_vec()
                     b"  B :ASSERT".to_vec()
                 } else {
@@ -267,8 +269,8 @@ fn optimize_labels(code: &[&str], func_index: usize) -> Vec<String> {
     lines
 }
 
-// TODO: fix same label names in different functions
-/// Compiles a clif function.
+/// Compiles a clif function into zkasm, to be used in test program construction.
+/// Result is Vec<String> where each String represents line of zkasm code
 pub fn compile_clif_function(func: &Function) -> Vec<String> {
     let flag_builder = settings::builder();
     let isa_builder = zkasm::isa_builder("zkasm-unknown-unknown".parse().unwrap());
@@ -287,18 +289,23 @@ pub fn compile_clif_function(func: &Function) -> Vec<String> {
     );
     let code = std::str::from_utf8(&code_buffer).unwrap();
     let mut lines: Vec<String> = code.lines().map(|s| s.to_string()).collect();
+
+    // Below we replacing default function name to same name as will be used in invocations
+    // to call it
     // TODO: I believe it can be done more beautiful way
     let mut funcname = func.name.to_string();
     funcname.remove(0);
     let mut res = vec![format!("{}:", funcname)];
     res.append(&mut lines);
+    // Here we adding function-speciefic prefix to internal labels of the function
+    // to make labels different in different functions.
     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> {
+pub fn build_test_main(invoke_names: Vec<String>) -> Vec<String> {
     let mut res = vec![
         "main:".to_string(),
         "  SP - 1 => SP".to_string(),
@@ -315,6 +322,7 @@ pub fn build_main(invoke_names: Vec<String>) -> Vec<String> {
 }
 
 /// Generate invoke name in format <function_name>_<arg>_..._<arg>
+/// This name can be used as label in zkasm, or as tag for helpers
 pub fn invoke_name(invoke: &Invocation) -> String {
     let mut res = invoke.func.clone();
     for arg in &invoke.args {
@@ -324,6 +332,8 @@ pub fn invoke_name(invoke: &Invocation) -> String {
 }
 
 /// Assembles all parts of zkasm test program together
+/// To get this parts you can use [compile_invocation], [compile_clif_function]
+/// and [build_test_main]
 pub fn build_test_zkasm(
     functions: Vec<Vec<String>>,
     invocations: Vec<Vec<String>>,
@@ -349,6 +359,9 @@ start:
     program.join("\n")
 }
 
+/// This function serves to generate wat as intermediate step
+/// for [compile_invocation].
+/// It returns wat generated from invocation as single string
 fn runcommand_to_wat(invoke: Invocation, _compare: Comparison, expected: Vec<DataValue>) -> String {
     // TODO: support different amounts of outputs
     let res_bitness = match expected[0] {
@@ -369,10 +382,7 @@ fn runcommand_to_wat(invoke: Invocation, _compare: Comparison, expected: Vec<Dat
         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();
+        args_pushing.push_str(&format!("{arg_type}.const {arg}\n"));
     }
     // TODO: remove line with 8 whitespaces in the end of args_pushing
     let wat_code = format!(
@@ -385,18 +395,20 @@ fn runcommand_to_wat(invoke: Invocation, _compare: Comparison, expected: Vec<Dat
         {res_bitness}.const {expected_result}
         call $assert_eq
     )
-    (export "main" (func $main))
-)"#,
+    (export "main" (func $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()
+    // This wasm2wat and wat2wasm used to automatically format wat code.
+    wasm2wat(wat2wasm(wat_code).unwrap()).unwrap()
 }
 
-/// Compiles a invocation.
+/// This function compiles clif run command (aka invocation) into zkasm.
+/// To properly allocate registers, it use wat as intermediate step.
+/// Result is Vec<Strings> where String represents line of zkasm code
 pub fn compile_invocation(
     invoke: Invocation,
     compare: Comparison,
@@ -425,13 +437,17 @@ pub fn compile_invocation(
     let new_label = invoke_name(&invoke);
     let funcname = invoke.func;
 
-    // TODO: will this always function_2?
+    // Next two lines used to cut preamble and postamble generated by [generate_zkasm]
+    // We need only function body, since it will be only part of zkasm test program, not a single program
+    // TODO: do not rely on label name ("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();
 
+    // Here we replace call to function from default generated label by [generate_zkasm] to
+    // to proper name. Also, here we replace default tag for assert-helper to something more verbose
+    // TODO: Do not rely on function name ("function_1")
     generated_function[0] = format!("{}:", new_label);
-    // TODO: will this always function_1?
     let generated_replaced: Vec<String> = generated_function
         .iter()
         .map(|s| {

From a7c59e10f2709eea64fdd5fcfdb7bdd478f4cea2 Mon Sep 17 00:00:00 2001
From: Viktar Makouski <mcviktor@tut.by>
Date: Wed, 13 Mar 2024 16:11:50 +0300
Subject: [PATCH 9/9] revert wabt usage

---
 Cargo.lock                               | 43 ++----------------------
 cranelift/filetests/Cargo.toml           |  1 -
 cranelift/filetests/src/zkasm_codegen.rs |  6 ++--
 3 files changed, 4 insertions(+), 46 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index f2111fe8c784..632c593e3071 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -474,15 +474,6 @@ version = "0.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2da6da31387c7e4ef160ffab6d5e7f00c42626fe39aea70a7b0f1773f7dd6c1b"
 
-[[package]]
-name = "cmake"
-version = "0.1.50"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a31c789563b815f77f4250caee12365734369f942439b7defd71e18a48197130"
-dependencies = [
- "cc",
-]
-
 [[package]]
 name = "codespan-reporting"
 version = "0.11.1"
@@ -678,7 +669,6 @@ dependencies = [
  "tempfile",
  "thiserror",
  "toml",
- "wabt",
  "walkdir",
  "wasmparser",
  "wasmtime",
@@ -1338,12 +1328,6 @@ dependencies = [
  "stable_deref_trait",
 ]
 
-[[package]]
-name = "glob"
-version = "0.2.11"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8be18de09a56b60ed0edf84bc9df007e30040691af7acd1c41874faac5895bfb"
-
 [[package]]
 name = "glob"
 version = "0.3.0"
@@ -2990,29 +2974,6 @@ version = "0.9.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "49874b5167b65d7193b8aba1567f5c7d93d001cafc34600cee003eda787e483f"
 
-[[package]]
-name = "wabt"
-version = "0.10.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "00bef93d5e6c81a293bccf107cf43aa47239382f455ba14869d36695d8963b9c"
-dependencies = [
- "serde",
- "serde_derive",
- "serde_json",
- "wabt-sys",
-]
-
-[[package]]
-name = "wabt-sys"
-version = "0.8.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1a4e043159f63e16986e713e9b5e1c06043df4848565bf672e27c523864c7791"
-dependencies = [
- "cc",
- "cmake",
- "glob 0.2.11",
-]
-
 [[package]]
 name = "wait-timeout"
 version = "0.2.0"
@@ -4028,7 +3989,7 @@ dependencies = [
 name = "winch-test-macros"
 version = "0.0.0"
 dependencies = [
- "glob 0.3.0",
+ "glob",
  "proc-macro2",
  "quote",
  "syn 2.0.32",
@@ -4042,7 +4003,7 @@ dependencies = [
  "capstone",
  "clap",
  "cranelift-codegen",
- "glob 0.3.0",
+ "glob",
  "serde",
  "similar",
  "target-lexicon",
diff --git a/cranelift/filetests/Cargo.toml b/cranelift/filetests/Cargo.toml
index 56acb7fcadee..4a85138ac583 100644
--- a/cranelift/filetests/Cargo.toml
+++ b/cranelift/filetests/Cargo.toml
@@ -47,4 +47,3 @@ regex = { workspace = true }
 wasmtime = { workspace = true, features = ["cranelift", "runtime"] }
 walkdir = { workspace = true }
 tempfile = { workspace = true }
-wabt = "0.10.0"
diff --git a/cranelift/filetests/src/zkasm_codegen.rs b/cranelift/filetests/src/zkasm_codegen.rs
index af3d8bcd46fc..26d4d18e33af 100644
--- a/cranelift/filetests/src/zkasm_codegen.rs
+++ b/cranelift/filetests/src/zkasm_codegen.rs
@@ -15,8 +15,6 @@ use cranelift_reader::Comparison;
 use cranelift_reader::Invocation;
 use cranelift_wasm::{translate_module, ZkasmEnvironment};
 
-use wabt::{wasm2wat, wat2wasm};
-
 /// ISA specific settings for zkASM codegen.
 #[derive(Default, Debug)]
 pub struct ZkasmSettings {
@@ -402,8 +400,8 @@ fn runcommand_to_wat(invoke: Invocation, _compare: Comparison, expected: Vec<Dat
         func_name = func_name,
         expected_result = expected_result,
     );
-    // This wasm2wat and wat2wasm used to automatically format wat code.
-    wasm2wat(wat2wasm(wat_code).unwrap()).unwrap()
+    // TODO: format wat code using something like wasm2wat(wat2wasm(wat_code)).
+    wat_code
 }
 
 /// This function compiles clif run command (aka invocation) into zkasm.