From 2d7e02deb0b74be38e4ea4be06e963b17e8baf6a Mon Sep 17 00:00:00 2001
From: Chris Fallin <chris@cfallin.org>
Date: Tue, 7 Feb 2023 17:31:52 -0800
Subject: [PATCH] Fix waffle-util build.

---
 src/bin/waffle-util.rs | 29 ++++++++++++++++-------------
 src/ir/func.rs         | 20 ++++++++++++++------
 2 files changed, 30 insertions(+), 19 deletions(-)

diff --git a/src/bin/waffle-util.rs b/src/bin/waffle-util.rs
index 52a7f9a..32a49a6 100644
--- a/src/bin/waffle-util.rs
+++ b/src/bin/waffle-util.rs
@@ -41,6 +41,19 @@ enum Command {
     },
 }
 
+fn apply_options(opts: &Options, module: &mut Module) -> Result<()> {
+    if opts.basic_opts || opts.max_ssa {
+        module.expand_all_funcs()?;
+    }
+    if opts.basic_opts {
+        module.per_func_body(|body| body.optimize());
+    }
+    if opts.max_ssa {
+        module.per_func_body(|body| body.convert_to_max_ssa());
+    }
+    Ok(())
+}
+
 fn main() -> Result<()> {
     let opts = Options::from_args();
 
@@ -50,29 +63,19 @@ fn main() -> Result<()> {
     }
     let _ = logger.try_init();
 
-    match opts.command {
+    match &opts.command {
         Command::PrintIR { wasm } => {
             let bytes = std::fs::read(wasm)?;
             debug!("Loaded {} bytes of Wasm data", bytes.len());
             let mut module = Module::from_wasm_bytes(&bytes[..])?;
-            if opts.basic_opts {
-                module.optimize();
-            }
-            if opts.max_ssa {
-                module.convert_to_max_ssa();
-            }
+            apply_options(&opts, &mut module)?;
             println!("{}", module.display());
         }
         Command::RoundTrip { input, output } => {
             let bytes = std::fs::read(input)?;
             debug!("Loaded {} bytes of Wasm data", bytes.len());
             let mut module = Module::from_wasm_bytes(&bytes[..])?;
-            if opts.basic_opts {
-                module.optimize();
-            }
-            if opts.max_ssa {
-                module.convert_to_max_ssa();
-            }
+            apply_options(&opts, &mut module)?;
             let produced = module.to_wasm_bytes()?;
             std::fs::write(output, &produced[..])?;
         }
diff --git a/src/ir/func.rs b/src/ir/func.rs
index 50458fe..eb0c1f1 100644
--- a/src/ir/func.rs
+++ b/src/ir/func.rs
@@ -34,10 +34,7 @@ impl<'a> FuncDecl<'a> {
     pub fn optimize(&mut self) {
         match self {
             FuncDecl::Body(_, body) => {
-                let cfg = crate::cfg::CFGInfo::new(body);
-                crate::passes::basic_opt::gvn(body, &cfg);
-                crate::passes::resolve_aliases::run(body);
-                crate::passes::empty_blocks::run(body);
+                body.optimize();
             }
             _ => {}
         }
@@ -46,8 +43,7 @@ impl<'a> FuncDecl<'a> {
     pub fn convert_to_max_ssa(&mut self) {
         match self {
             FuncDecl::Body(_, body) => {
-                let cfg = crate::cfg::CFGInfo::new(body);
-                crate::passes::maxssa::run(body, &cfg);
+                body.convert_to_max_ssa();
             }
             _ => {}
         }
@@ -115,6 +111,18 @@ impl FunctionBody {
         }
     }
 
+    pub fn optimize(&mut self) {
+        let cfg = crate::cfg::CFGInfo::new(self);
+        crate::passes::basic_opt::gvn(self, &cfg);
+        crate::passes::resolve_aliases::run(self);
+        crate::passes::empty_blocks::run(self);
+    }
+
+    pub fn convert_to_max_ssa(&mut self) {
+        let cfg = crate::cfg::CFGInfo::new(self);
+        crate::passes::maxssa::run(self, &cfg);
+    }
+
     pub fn add_block(&mut self) -> Block {
         let id = self.blocks.push(BlockDef::default());
         log::trace!("add_block: block {}", id);