feat/misc: add const_fold to all possible immediate values / test utils

crates/assembler/src/parser.rs

@@ -258,33 +258,38 @@ impl ParseInstruction for Instruction {
                         if tokens.len() < 8 {
                             return None;
                         }
-                        match (
-                            &tokens[1],
-                            &tokens[2],
-                            &tokens[3],
-                            &tokens[4],
-                            &tokens[5],
-                            &tokens[6],
-                            &tokens[7],
-                        ) {
-                            (
-                                Token::LeftBracket(_),
-                                Token::Register(_, _),
-                                Token::BinaryOp(_, _),
-                                Token::ImmediateValue(_, _),
-                                Token::RightBracket(_),
-                                Token::Comma(_),
-                                Token::Register(_, _)
-                            ) => {
-                                operands.push(tokens[2].clone());
-                                operands.push(tokens[4].clone());
-                                operands.push(tokens[7].clone());
-                            }
-                            _ => {
-                                return None;
+                        let (value, advance_token_num) = inline_and_fold_constant(tokens, const_map, 4);
+                        if let Some(value) = value {
+                            match (
+                                &tokens[1],
+                                &tokens[2],
+                                &tokens[3],
+                                // Fourth operand is folded to an immediate value
+                                &tokens[5],
+                                &tokens[6],
+                                &tokens[7],
+                            ) {
+                                (
+                                    Token::LeftBracket(_),
+                                    Token::Register(_, _),
+                                    Token::BinaryOp(_, _),
+                                    // Fourth operand is folded to an immediate value
+                                    Token::RightBracket(_),
+                                    Token::Comma(_),
+                                    Token::Register(_, _)
+                                ) => {
+                                    operands.push(tokens[2].clone());
+                                    operands.push(Token::ImmediateValue(value, 0));
+                                    operands.push(tokens[7].clone());
+                                }
+                                _ => {
+                                    return None;
+                                }
                             }
+                            next_token_num = advance_token_num + 3;
+                        } else {
+                            return None;
                         }
-                        next_token_num = 8;
                     }
                     Opcode::Add32 | Opcode::Sub32 | Opcode::Mul32 
                     | Opcode::Div32 | Opcode::Or32 | Opcode::And32 
@@ -351,56 +356,78 @@ impl ParseInstruction for Instruction {
                         if tokens.len() < 6 {
                             return None;
                         }
-                        match (
-                            &tokens[1],
-                            &tokens[2],
-                            &tokens[3],
-                            &tokens[4],
-                            &tokens[5],
-                        ) {
-                            (
-                                Token::Register(_, _),
-                                Token::Comma(_),
-                                Token::ImmediateValue(_, _),
-                                Token::Comma(_),
-                                Token::Identifier(_, _)
-                            ) => {
-                                opcode = FromPrimitive::from_u8((opcode as u8) + 1).expect("Invalid opcode conversion"); 
-                                operands.push(tokens[1].clone());
-                                operands.push(tokens[3].clone());
-                                operands.push(tokens[5].clone());
-                            }
-                            (
-                                Token::Register(_, _),
-                                Token::Comma(_),
-                                Token::Register(_, _),
-                                Token::Comma(_),
-                                Token::Identifier(_, _)
-                            ) => {
-                                opcode = FromPrimitive::from_u8((opcode as u8) + 2).expect("Invalid opcode conversion"); 
-                                operands.push(tokens[1].clone());
-                                operands.push(tokens[3].clone());
-                                operands.push(tokens[5].clone());
+                        let (value, advance_token_num) = inline_and_fold_constant(tokens, const_map, 3);
+                        if let Some(value) = value {
+                            match (
+                                &tokens[1],
+                                &tokens[2],
+                                // Third operand is folded to an immediate value
+                                &tokens[4],
+                                &tokens[5],
+                            ) {
+                                (
+                                    Token::Register(_, _),
+                                    Token::Comma(_),
+                                    // Third operand is folded to an immediate value
+                                    Token::Comma(_),
+                                    Token::Identifier(_, _)
+                                ) => {
+                                    opcode = FromPrimitive::from_u8((opcode as u8) + 1).expect("Invalid opcode conversion"); 
+                                    operands.push(tokens[1].clone());
+                                    operands.push(Token::ImmediateValue(value, 0));
+                                    operands.push(tokens[5].clone());
+                                }
+                                _ => {
+                                    return None;
+                                }
                             }
-                            _ => {
-                                return None;
+                            next_token_num = advance_token_num + 2;
+                        } else {
+                            match (
+                                &tokens[1],
+                                &tokens[2],
+                                &tokens[3],
+                                &tokens[4],
+                                &tokens[5],
+                            ) {
+                                (
+                                    Token::Register(_, _),
+                                    Token::Comma(_),
+                                    Token::Register(_, _),
+                                    Token::Comma(_),
+                                    Token::Identifier(_, _)
+                                ) => {
+                                    opcode = FromPrimitive::from_u8((opcode as u8) + 2).expect("Invalid opcode conversion"); 
+                                    operands.push(tokens[1].clone());
+                                    operands.push(tokens[3].clone());
+                                    operands.push(tokens[5].clone());
+                                }
+                                _ => {
+                                    return None;
+                                }
                             }
+                            next_token_num = 6;
                         }
-                        next_token_num = 6;
                     }
                     Opcode::Ja => {
                         if tokens.len() < 2 {
                             return None;
                         }
-                        match &tokens[1] {
-                            Token::Identifier(_, _) | Token::ImmediateValue(_, _) => {
-                                operands.push(tokens[1].clone());
-                            }
-                            _ => {
-                                return None;
+                        let (value, advance_token_num) = inline_and_fold_constant(tokens, const_map, 1);
+                        if let Some(value) = value {
+                            operands.push(Token::ImmediateValue(value, 0));
+                            next_token_num = advance_token_num;
+                        } else {
+                            match &tokens[1] {
+                                Token::Identifier(_, _) => {
+                                    operands.push(tokens[1].clone());
+                                }
+                                _ => {
+                                    return None;
+                                }
                             }
+                            next_token_num = 2;
                         }
-                        next_token_num = 2;
                     }
                     Opcode::Call => {
                         if tokens.len() < 2 {

tests/test_memo.rs

@@ -1,132 +1,42 @@
 mod utils;
 
-use utils::EnvVarGuard;
-
-use std::process::Command;
-use std::fs;
+use utils::{TestEnv, init_project, verify_project_structure, run_light_build, verify_so_files, run_tests, update_assembly_file};
 
 #[test]
 fn test_memo_project_e2e() {
-    let current_dir = std::env::current_dir().expect("Failed to get current directory");
-    let sbpf_path = current_dir.join("target").join("debug").join("sbpf");
-    let _guard = EnvVarGuard::new("SBPF_BIN", sbpf_path.to_string_lossy());
-    let sbpf_bin = std::env::var("SBPF_BIN").expect("SBPF_BIN not set");
-    let temp_dir = std::env::temp_dir().join("sbpf_test_memo_e2e");
-    if temp_dir.exists() {
-        fs::remove_dir_all(&temp_dir).expect("Failed to remove existing test directory");
-    }
-    fs::create_dir_all(&temp_dir).expect("Failed to create test directory");
-    println!("Using test directory: {:?}", temp_dir);
-    // Step 1: Initialize the memo project using the built binary
-    println!("Step 1: Initializing memo project...");
-    let init_output = Command::new(&sbpf_bin)
-        .current_dir(&temp_dir)
-        .arg("init")
-        .arg("memo")
-        .output()
-        .expect("Failed to execute target/debug/sbpf init memo");
-    if !init_output.status.success() {
-        let stderr = String::from_utf8_lossy(&init_output.stderr);
-        let stdout = String::from_utf8_lossy(&init_output.stdout);
-        panic!("target/debug/sbpf init memo failed:\nSTDOUT: {}\nSTDERR: {}", stdout, stderr);
-    }
-    println!("✅ Project initialized successfully");
-    println!("STDOUT: {}", String::from_utf8_lossy(&init_output.stdout));
-    // Step 2: Set up memo directory path
-    let memo_dir = temp_dir.join("memo");
-    println!("Using memo directory: {:?}", memo_dir);
-    assert!(memo_dir.join("src").exists(), "src directory should exist");
-    assert!(memo_dir.join("deploy").exists(), "deploy directory should exist");
-    assert!(memo_dir.join("src/memo").exists(), "src/memo directory should exist");
-    assert!(memo_dir.join("src/memo/memo.s").exists(), "src/memo/memo.s should exist");
-    assert!(memo_dir.join("src/lib.rs").exists(), "src/lib.rs should exist");
-    assert!(memo_dir.join("Cargo.toml").exists(), "Cargo.toml should exist");
-    // Step 3: Run light-build using the built binary
-    println!("Step 3: Running light-build...");
-    let light_build_output = Command::new(&sbpf_bin)
-        .current_dir(&memo_dir)
-        .arg("light-build")
-        .output()
-        .expect("Failed to execute target/debug/sbpf light-build");
-    if !light_build_output.status.success() {
-        let stderr = String::from_utf8_lossy(&light_build_output.stderr);
-        let stdout = String::from_utf8_lossy(&light_build_output.stdout);
-        panic!("target/debug/sbpf light-build failed:\nSTDOUT: {}\nSTDERR: {}", stdout, stderr);
-    }
-    println!("✅ Light-build completed successfully");
-    println!("STDOUT: {}", String::from_utf8_lossy(&light_build_output.stdout));
-    // Verify that .so files were created in the deploy directory
-    let deploy_dir = memo_dir.join("deploy");
-    assert!(deploy_dir.exists(), "deploy directory should exist");
-    let so_files: Vec<_> = fs::read_dir(&deploy_dir)
-        .expect("Failed to read deploy directory")
-        .filter_map(|entry| entry.ok())
-        .filter(|entry| {
-            entry.path()
-                .extension()
-                .and_then(|ext| ext.to_str())
-                .map(|ext| ext == "so")
-                .unwrap_or(false)
-        })
-        .collect();
-    assert!(!so_files.is_empty(), "At least one .so file should be created in deploy directory");
-    println!("Found {} .so file(s) in deploy directory", so_files.len());
-    // Step 4: Run tests using the built binary
-    println!("Step 4: Running tests...");
-    let test_output = Command::new(&sbpf_bin)
-        .current_dir(&memo_dir)
-        .arg("test")
-        .output()
-        .expect("Failed to execute target/debug/sbpf test");
-    if !test_output.status.success() {
-        let stderr = String::from_utf8_lossy(&test_output.stderr);
-        let stdout = String::from_utf8_lossy(&test_output.stdout);
-        panic!("target/debug/sbpf test failed:\nSTDOUT: {}\nSTDERR: {}", stdout, stderr);
-    }
-    println!("✅ Tests completed successfully");
-    println!("STDOUT: {}", String::from_utf8_lossy(&test_output.stdout));
-    // Step 5: Clean up
-    fs::remove_dir_all(&temp_dir).expect("Failed to remove test directory");
+    let env = TestEnv::new("memo");
+    
+    // Step 1: Initialize the memo project
+    init_project(&env, "memo");
+    
+    // Step 2: Verify project structure
+    verify_project_structure(&env, "memo");
+    
+    // Step 3: Run light-build
+    run_light_build(&env);
+    
+    // Step 4: Verify .so files were created
+    verify_so_files(&env);
+    
+    // Step 5: Run tests
+    run_tests(&env);
+    
+    // Step 6: Clean up
+    env.cleanup();
     println!("🎉 All tests passed! Memo project E2E test completed successfully.");
 }
 
 #[test]
 fn test_memo_project_e2e_second() {
-    let current_dir = std::env::current_dir().expect("Failed to get current directory");
-    let sbpf_path = current_dir.join("target").join("debug").join("sbpf");
-    let _guard = EnvVarGuard::new("SBPF_BIN", sbpf_path.to_string_lossy());
-    let sbpf_bin = std::env::var("SBPF_BIN").expect("SBPF_BIN not set");
-    let temp_dir = std::env::temp_dir().join("sbpf_test_memo_e2e_second");
-    if temp_dir.exists() {
-        fs::remove_dir_all(&temp_dir).expect("Failed to remove existing test directory");
-    }
-    fs::create_dir_all(&temp_dir).expect("Failed to create test directory");
-    println!("Using system temp directory for second test: {:?}", temp_dir);
-    // Step 1: Initialize the memo project using the built binary
-    println!("Step 1: Initializing second memo project...");
-    let init_output = Command::new(&sbpf_bin)
-        .current_dir(&temp_dir)
-        .arg("init")
-        .arg("memo2")
-        .output()
-        .expect("Failed to execute target/debug/sbpf init memo2");
-    if !init_output.status.success() {
-        let stderr = String::from_utf8_lossy(&init_output.stderr);
-        let stdout = String::from_utf8_lossy(&init_output.stdout);
-        panic!("target/debug/sbpf init memo2 failed:\nSTDOUT: {}\nSTDERR: {}", stdout, stderr);
-    }
-    println!("✅ Second project initialized successfully");
-    println!("STDOUT: {}", String::from_utf8_lossy(&init_output.stdout));
-    // Step 2: Set up memo2 directory path
-    let memo_dir = temp_dir.join("memo2");
-    println!("Using memo2 directory: {:?}", memo_dir);
-    assert!(memo_dir.join("src").exists(), "src directory should exist");
-    assert!(memo_dir.join("deploy").exists(), "deploy directory should exist");
-    assert!(memo_dir.join("src/memo2").exists(), "src/memo2 directory should exist");
-    assert!(memo_dir.join("src/memo2/memo2.s").exists(), "src/memo2/memo2.s should exist");
-    assert!(memo_dir.join("src/lib.rs").exists(), "src/lib.rs should exist");
-    assert!(memo_dir.join("Cargo.toml").exists(), "Cargo.toml should exist");
-    // Update the memo2.s content to the specified content
+    let env = TestEnv::new("memo2");
+    
+    // Step 1: Initialize the memo2 project
+    init_project(&env, "memo2");
+    
+    // Step 2: Verify project structure
+    verify_project_structure(&env, "memo2");
+    
+    // Step 3: Update the memo2.s content to the specified content
     let new_memo_content = r#".equ NUM_ACCOUNTS, 0x00
 .equ DATA_LEN, 0x08
 .equ DATA, 0x10
@@ -138,53 +48,18 @@ entrypoint:
   add64 r1, DATA
   call sol_log_
   exit"#;
-    fs::write(memo_dir.join("src/memo2/memo2.s"), new_memo_content).expect("Failed to write new memo2.s content");
-    println!("✅ Updated memo2.s with specified content");
-    // Step 3: Run light-build using the built binary
-    println!("Step 3: Running light-build on second project...");
-    let light_build_output = Command::new(&sbpf_bin)
-        .current_dir(&memo_dir)
-        .arg("light-build")
-        .output()
-        .expect("Failed to execute target/debug/sbpf light-build");
-    if !light_build_output.status.success() {
-        let stderr = String::from_utf8_lossy(&light_build_output.stderr);
-        let stdout = String::from_utf8_lossy(&light_build_output.stdout);
-        panic!("target/debug/sbpf light-build failed:\nSTDOUT: {}\nSTDERR: {}", stdout, stderr);
-    }
-    println!("✅ Second light-build completed successfully");
-    println!("STDOUT: {}", String::from_utf8_lossy(&light_build_output.stdout));
-    // Verify that .so files were created in the deploy directory
-    let deploy_dir = memo_dir.join("deploy");
-    assert!(deploy_dir.exists(), "deploy directory should exist");
-    let so_files: Vec<_> = fs::read_dir(&deploy_dir)
-        .expect("Failed to read deploy directory")
-        .filter_map(|entry| entry.ok())
-        .filter(|entry| {
-            entry.path()
-                .extension()
-                .and_then(|ext| ext.to_str())
-                .map(|ext| ext == "so")
-                .unwrap_or(false)
-        })
-        .collect();
-    assert!(!so_files.is_empty(), "At least one .so file should be created in deploy directory");
-    println!("Found {} .so file(s) in deploy directory", so_files.len());
-    // Step 4: Run tests using the built binary
-    println!("Step 4: Running tests on second project...");
-    let test_output = Command::new(&sbpf_bin)
-        .current_dir(&memo_dir)
-        .arg("test")
-        .output()
-        .expect("Failed to execute target/debug/sbpf test");
-    if !test_output.status.success() {
-        let stderr = String::from_utf8_lossy(&test_output.stderr);
-        let stdout = String::from_utf8_lossy(&test_output.stdout);
-        panic!("target/debug/sbpf test failed:\nSTDOUT: {}\nSTDERR: {}", stdout, stderr);
-    }
-    println!("✅ Second tests completed successfully");
-    println!("STDOUT: {}", String::from_utf8_lossy(&test_output.stdout));
-    // Step 5: Clean up
-    fs::remove_dir_all(&temp_dir).expect("Failed to remove test directory");
+    update_assembly_file(&env, "memo2", new_memo_content);
+    
+    // Step 4: Run light-build
+    run_light_build(&env);
+    
+    // Step 5: Verify .so files were created
+    verify_so_files(&env);
+    
+    // Step 6: Run tests
+    run_tests(&env);
+    
+    // Step 7: Clean up
+    env.cleanup();
     println!("🎉 Second test passed! Memo2 project E2E test completed successfully.");
 }

tests/utils.rs

@@ -22,3 +22,146 @@ impl Drop for EnvVarGuard {
         }
     }
 }
+
+use std::process::{Command, Output};
+use std::fs;
+use std::path::PathBuf;
+
+/// Test environment setup for SBPF tests
+pub struct TestEnv {
+    pub sbpf_bin: String,
+    pub temp_dir: PathBuf,
+    pub project_dir: PathBuf,
+    pub _guard: EnvVarGuard,
+}
+
+impl TestEnv {
+    /// Set up a test environment with SBPF binary and temporary directory
+    pub fn new(project_name: &str) -> Self {
+        let current_dir = std::env::current_dir().expect("Failed to get current directory");
+        let sbpf_path = current_dir.join("target").join("debug").join("sbpf");
+        let guard = EnvVarGuard::new("SBPF_BIN", sbpf_path.to_string_lossy());
+        let sbpf_bin = std::env::var("SBPF_BIN").expect("SBPF_BIN not set");
+        
+        let temp_dir = std::env::temp_dir().join(format!("sbpf_test_{}_e2e", project_name));
+        if temp_dir.exists() {
+            fs::remove_dir_all(&temp_dir).expect("Failed to remove existing test directory");
+        }
+        fs::create_dir_all(&temp_dir).expect("Failed to create test directory");
+        
+        let project_dir = temp_dir.join(project_name);
+        
+        println!("Using test directory: {:?}", temp_dir);
+        println!("Using project directory: {:?}", project_dir);
+        
+        Self {
+            sbpf_bin,
+            temp_dir,
+            project_dir,
+            _guard: guard,
+        }
+    }
+    
+    /// Clean up the test environment
+    pub fn cleanup(self) {
+        fs::remove_dir_all(&self.temp_dir).expect("Failed to remove test directory");
+    }
+}
+
+/// Run a command and return the output, panicking on failure
+pub fn run_command(cmd: &mut Command, operation_name: &str) -> Output {
+    let output = cmd.output().expect(&format!("Failed to execute {}", operation_name));
+    
+    if !output.status.success() {
+        let stderr = String::from_utf8_lossy(&output.stderr);
+        let stdout = String::from_utf8_lossy(&output.stdout);
+        panic!("{} failed:\nSTDOUT: {}\nSTDERR: {}", operation_name, stdout, stderr);
+    }
+    
+    println!("✅ {} completed successfully", operation_name);
+    println!("STDOUT: {}", String::from_utf8_lossy(&output.stdout));
+    
+    output
+}
+
+/// Initialize a new SBPF project
+pub fn init_project(env: &TestEnv, project_name: &str) {
+    println!("Step 1: Initializing {} project...", project_name);
+    
+    let init_output = run_command(
+        Command::new(&env.sbpf_bin)
+            .current_dir(&env.temp_dir)
+            .arg("init")
+            .arg(project_name),
+        &format!("target/debug/sbpf init {}", project_name)
+    );
+    
+    println!("✅ Project initialized successfully");
+    println!("STDOUT: {}", String::from_utf8_lossy(&init_output.stdout));
+}
+
+/// Verify that the project structure is correct
+pub fn verify_project_structure(env: &TestEnv, project_name: &str) {
+    let project_dir = &env.project_dir;
+    
+    assert!(project_dir.join("src").exists(), "src directory should exist");
+    assert!(project_dir.join("deploy").exists(), "deploy directory should exist");
+    assert!(project_dir.join(format!("src/{}", project_name)).exists(), 
+            "src/{} directory should exist", project_name);
+    assert!(project_dir.join(format!("src/{}/{}.s", project_name, project_name)).exists(), 
+            "src/{}/{}.s should exist", project_name, project_name);
+    assert!(project_dir.join("src/lib.rs").exists(), "src/lib.rs should exist");
+    assert!(project_dir.join("Cargo.toml").exists(), "Cargo.toml should exist");
+}
+
+/// Run light-build on the project
+pub fn run_light_build(env: &TestEnv) {
+    println!("Step 3: Running light-build...");
+    
+    run_command(
+        Command::new(&env.sbpf_bin)
+            .current_dir(&env.project_dir)
+            .arg("light-build"),
+        "target/debug/sbpf light-build"
+    );
+}
+
+/// Verify that .so files were created in the deploy directory
+pub fn verify_so_files(env: &TestEnv) {
+    let deploy_dir = env.project_dir.join("deploy");
+    assert!(deploy_dir.exists(), "deploy directory should exist");
+    
+    let so_files: Vec<_> = fs::read_dir(&deploy_dir)
+        .expect("Failed to read deploy directory")
+        .filter_map(|entry| entry.ok())
+        .filter(|entry| {
+            entry.path()
+                .extension()
+                .and_then(|ext| ext.to_str())
+                .map(|ext| ext == "so")
+                .unwrap_or(false)
+        })
+        .collect();
+    
+    assert!(!so_files.is_empty(), "At least one .so file should be created in deploy directory");
+    println!("Found {} .so file(s) in deploy directory", so_files.len());
+}
+
+/// Run tests on the project
+pub fn run_tests(env: &TestEnv) {
+    println!("Step 4: Running tests...");
+    
+    run_command(
+        Command::new(&env.sbpf_bin)
+            .current_dir(&env.project_dir)
+            .arg("test"),
+        "target/debug/sbpf test"
+    );
+}
+
+/// Update the project's assembly file content
+pub fn update_assembly_file(env: &TestEnv, project_name: &str, content: &str) {
+    let assembly_path = env.project_dir.join(format!("src/{}/{}.s", project_name, project_name));
+    fs::write(&assembly_path, content).expect(&format!("Failed to write new {}.s content", project_name));
+    println!("✅ Updated {}.s with specified content", project_name);
+}