Merge pull request #17 from arnoox/add-opt-switch

Enhance optimization handling by introducing HERKOS_OPTIMIZE flag and updating transpile options

Author: arnoox

.github/workflows/ci.yml

@@ -7,6 +7,7 @@ on:
 
 env:
   CARGO_TERM_COLOR: always
+  FORCE_JAVASCRIPT_ACTIONS_TO_NODE24: true
 
 jobs:
   check:
@@ -29,8 +30,15 @@ jobs:
       - name: Build
         run: cargo build --all-targets
 
-      - name: Test
+      - name: Test (optimized)
         run: cargo test
+        env:
+          HERKOS_OPTIMIZE: "1"
+
+      - name: Test (unoptimized)
+        run: cargo test
+        env:
+          HERKOS_OPTIMIZE: "0"
 
       - name: Clippy
         run: cargo clippy --all-targets -- -D warnings

crates/herkos-tests/build.rs

@@ -5,18 +5,26 @@ use std::fs;
 use std::path::{Path, PathBuf};
 
 fn main() -> Result<()> {
-    // Rerun build script if this file or any data files change
+    // Rerun build script if this file, any data files, or the optimize flag changes
     println!("cargo:rerun-if-changed=build.rs");
     println!("cargo:rerun-if-changed=data");
+    println!("cargo:rerun-if-env-changed=HERKOS_OPTIMIZE");
 
     // Use cargo's OUT_DIR for generated files (parallel-build safe)
     let out_dir = PathBuf::from(env::var("OUT_DIR").context("OUT_DIR not set")?);
+
+    // HERKOS_OPTIMIZE=1 enables IR optimizations; any other value (or unset) disables them.
+    let optimize = env::var("HERKOS_OPTIMIZE").is_ok_and(|v| v == "1");
     eprintln!(
-        "Generating WASM and transpiled modules to: {}",
-        out_dir.display()
+        "Generating WASM and transpiled modules to: {} (optimize={})",
+        out_dir.display(),
+        optimize
     );
 
-    let options = TranspileOptions::default();
+    let options = TranspileOptions {
+        optimize,
+        ..TranspileOptions::default()
+    };
     let mut module_names = Vec::new();
 
     // 1. WAT test cases from data/wat/*.wat

crates/herkos/src/lib.rs

@@ -15,7 +15,7 @@ use backend::SafeBackend;
 use codegen::CodeGenerator;
 use ir::builder::build_module_info;
 use ir::{lower_phis, LoweredModuleInfo};
-use optimizer::optimize_ir;
+use optimizer::{optimize_ir, optimize_lowered_ir};
 use parser::parse_wasm;
 
 /// Configuration options for transpilation
@@ -25,13 +25,16 @@ pub struct TranspileOptions {
     pub mode: String,
     /// Maximum memory pages (used when Wasm module declares no maximum)
     pub max_pages: usize,
+    /// Enable optimizations
+    pub optimize: bool,
 }
 
 impl Default for TranspileOptions {
     fn default() -> Self {
         Self {
             mode: "safe".to_string(),
             max_pages: 256,
+            optimize: false,
         }
     }
 }
@@ -65,16 +68,17 @@ pub fn transpile(wasm_bytes: &[u8], options: &TranspileOptions) -> Result<String
     let module_info =
         build_module_info(&parsed, options).context("failed to build module metadata")?;
 
+    // Optimize the pure SSA IR.
+    let module_info = optimize_ir(module_info, options.optimize)?;
+
     // SSA destruction: lower phi nodes to predecessor assignments.
-    // Must run before the optimizer so that all block-elimination passes
-    // operate on phi-free IR.
-    let module_info = lower_phis::lower(module_info);
+    let lowered_module_info = lower_phis::lower(module_info);
 
-    // Optimize the IR
-    let module_info = optimize_ir(module_info)?;
+    // Optimize the lowered IR
+    let lowered_module_info = optimize_lowered_ir(lowered_module_info, options.optimize)?;
 
     // Generate Rust source code
-    let rust_code = generate_rust_code(&module_info)?;
+    let rust_code = generate_rust_code(&lowered_module_info)?;
 
     Ok(rust_code)
 }

crates/herkos/src/main.rs

@@ -14,6 +14,10 @@ struct Cli {
     /// Output Rust source file
     #[arg(long, short)]
     output: Option<PathBuf>,
+
+    /// Enable IR optimizations
+    #[arg(long, short = 'O')]
+    optimize: bool,
 }
 
 fn main() -> Result<()> {
@@ -29,6 +33,7 @@ fn main() -> Result<()> {
     let options = TranspileOptions {
         mode: "safe".to_string(),
         max_pages: 256,
+        optimize: cli.optimize,
     };
 
     // Transpile using library function

crates/herkos/src/optimizer/mod.rs

@@ -6,24 +6,38 @@
 //! Each optimization is a self-contained sub-module. The top-level
 //! [`optimize_ir`] function runs all passes in order.
 
-use crate::ir::LoweredModuleInfo;
+use crate::ir::{LoweredModuleInfo, ModuleInfo};
 use anyhow::Result;
 
 // ── Passes ───────────────────────────────────────────────────────────────────
 mod dead_blocks;
 
-/// Optimizes the IR representation by running all passes in order.
+/// Optimizes the pure SSA IR before phi lowering.
 ///
-/// Expects a [`LoweredModuleInfo`] — i.e. phi nodes have already been lowered
-/// by [`crate::ir::lower_phis::lower`] before calling this function.
-pub fn optimize_ir(module_info: LoweredModuleInfo) -> Result<LoweredModuleInfo> {
+/// Passes here operate on [`ModuleInfo`] with phi nodes still intact.
+/// Control-flow based passes (e.g. dead block elimination) belong here
+/// because reachability is identical in SSA and lowered form.
+pub fn optimize_ir(module_info: ModuleInfo, do_opt: bool) -> Result<ModuleInfo> {
     let mut module_info = module_info;
-    for func in &mut module_info.ir_functions {
-        dead_blocks::eliminate(func)?;
+    if do_opt {
+        for func in &mut module_info.ir_functions {
+            dead_blocks::eliminate(func)?;
+        }
     }
     Ok(module_info)
 }
 
+/// Optimizes the lowered IR after phi nodes have been eliminated.
+///
+/// Passes here operate on [`LoweredModuleInfo`] where all `IrInstr::Phi`
+/// nodes have been replaced by `IrInstr::Assign` in predecessor blocks.
+pub fn optimize_lowered_ir(
+    module_info: LoweredModuleInfo,
+    _do_opt: bool,
+) -> Result<LoweredModuleInfo> {
+    Ok(module_info)
+}
+
 // ── optimize_ir integration tests ─────────────────────────────────────────────
 
 #[cfg(test)]
@@ -74,7 +88,7 @@ mod tests {
             ..Default::default()
         };
 
-        let result = super::optimize_ir(crate::ir::lower_phis::lower(module)).unwrap();
+        let result = super::optimize_ir(module, true).unwrap();
         assert_eq!(
             result.ir_functions[0].blocks.len(),
             1,