From 9da822f59923a9953894c43afd1ddbeffa871dbf Mon Sep 17 00:00:00 2001
From: Tom French <15848336+TomAFrench@users.noreply.github.com>
Date: Tue, 12 Sep 2023 00:15:52 +0100
Subject: [PATCH] fix(wasm): Apply transformation map to circuit debug
 information in `noir_wasm` (#2635)

---
 Cargo.lock                   |  2 ++
 compiler/wasm/Cargo.toml     |  2 ++
 compiler/wasm/src/compile.rs | 52 ++++++++++++------------------------
 3 files changed, 21 insertions(+), 35 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 872fdf6c79f..043387bde89 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2285,7 +2285,9 @@ dependencies = [
  "fm",
  "getrandom",
  "gloo-utils",
+ "iter-extended",
  "log",
+ "nargo",
  "noirc_driver",
  "noirc_frontend",
  "serde",
diff --git a/compiler/wasm/Cargo.toml b/compiler/wasm/Cargo.toml
index efc50ef0476..1049dc92f47 100644
--- a/compiler/wasm/Cargo.toml
+++ b/compiler/wasm/Cargo.toml
@@ -13,6 +13,8 @@ crate-type = ["cdylib"]
 [dependencies]
 acvm.workspace = true
 fm.workspace = true
+iter-extended.workspace = true
+nargo.workspace = true
 noirc_driver.workspace = true
 noirc_frontend.workspace = true
 wasm-bindgen.workspace = true
diff --git a/compiler/wasm/src/compile.rs b/compiler/wasm/src/compile.rs
index a4e7f93ca79..f610ff7e79f 100644
--- a/compiler/wasm/src/compile.rs
+++ b/compiler/wasm/src/compile.rs
@@ -1,10 +1,9 @@
-use acvm::acir::circuit::Circuit;
 use fm::FileManager;
 use gloo_utils::format::JsValueSerdeExt;
+use iter_extended::try_vecmap;
 use log::debug;
 use noirc_driver::{
-    add_dep, check_crate, compile_contracts, compile_no_check, prepare_crate, prepare_dependency,
-    CompileOptions, CompiledContract,
+    add_dep, compile_contracts, compile_main, prepare_crate, prepare_dependency, CompileOptions,
 };
 use noirc_frontend::{graph::CrateGraph, hir::Context};
 use serde::{Deserialize, Serialize};
@@ -111,7 +110,10 @@ pub fn compile(args: JsValue) -> JsValue {
         add_noir_lib(&mut context, dependency.as_str());
     }
 
-    check_crate(&mut context, crate_id, false).expect("Crate check failed");
+    // For now we default to plonk width = 3, though we can add it as a parameter
+    let np_language = acvm::Language::PLONKCSat { width: 3 };
+    #[allow(deprecated)]
+    let is_opcode_supported = acvm::pwg::default_is_opcode_supported(np_language);
 
     if options.contracts {
         let compiled_contracts =
@@ -119,41 +121,21 @@ pub fn compile(args: JsValue) -> JsValue {
                 .expect("Contract compilation failed")
                 .0;
 
-        let optimized_contracts: Vec<CompiledContract> =
-            compiled_contracts.into_iter().map(optimize_contract).collect();
+        let optimized_contracts = try_vecmap(compiled_contracts, |contract| {
+            nargo::ops::optimize_contract(contract, np_language, &is_opcode_supported)
+        })
+        .expect("Contract optimization failed");
 
         <JsValue as JsValueSerdeExt>::from_serde(&optimized_contracts).unwrap()
     } else {
-        let main = context.get_main_function(&crate_id).expect("Could not find main function!");
-        let mut compiled_program =
-            compile_no_check(&context, &options.compile_options, main).expect("Compilation failed");
+        let compiled_program = compile_main(&mut context, crate_id, &options.compile_options)
+            .expect("Compilation failed")
+            .0;
 
-        compiled_program.circuit = optimize_circuit(compiled_program.circuit);
-
-        <JsValue as JsValueSerdeExt>::from_serde(&compiled_program).unwrap()
-    }
-}
+        let optimized_program =
+            nargo::ops::optimize_program(compiled_program, np_language, &is_opcode_supported)
+                .expect("Program optimization failed");
 
-fn optimize_contract(contract: CompiledContract) -> CompiledContract {
-    CompiledContract {
-        name: contract.name,
-        functions: contract
-            .functions
-            .into_iter()
-            .map(|mut func| {
-                func.bytecode = optimize_circuit(func.bytecode);
-                func
-            })
-            .collect(),
+        <JsValue as JsValueSerdeExt>::from_serde(&optimized_program).unwrap()
     }
 }
-
-fn optimize_circuit(circuit: Circuit) -> Circuit {
-    // For now we default to plonk width = 3, though we can add it as a parameter
-    let language = acvm::Language::PLONKCSat { width: 3 };
-    #[allow(deprecated)]
-    let opcode_supported = acvm::pwg::default_is_opcode_supported(language);
-    acvm::compiler::compile(circuit, language, opcode_supported)
-        .expect("Circuit optimization failed")
-        .0
-}