From b30b3f438e8ed6953f2fec9c610619ac4fb17553 Mon Sep 17 00:00:00 2001
From: Tom French <15848336+TomAFrench@users.noreply.github.com>
Date: Thu, 10 Aug 2023 16:40:05 +0100
Subject: [PATCH] fix: Optimize contracts built by `nargo info` (#2259)

fix: optimize contracts built by `nargo info`
---
 crates/nargo_cli/src/cli/compile_cmd.rs | 53 +++++++++++++++----------
 crates/nargo_cli/src/cli/info_cmd.rs    |  9 ++++-
 2 files changed, 40 insertions(+), 22 deletions(-)

diff --git a/crates/nargo_cli/src/cli/compile_cmd.rs b/crates/nargo_cli/src/cli/compile_cmd.rs
index 4430b9bb851..fad9783e654 100644
--- a/crates/nargo_cli/src/cli/compile_cmd.rs
+++ b/crates/nargo_cli/src/cli/compile_cmd.rs
@@ -7,7 +7,8 @@ use nargo::prepare_package;
 use nargo::{artifacts::contract::PreprocessedContract, NargoError};
 use nargo_toml::{find_package_manifest, resolve_workspace_from_toml};
 use noirc_driver::{
-    compile_contracts, compile_main, CompileOptions, CompiledProgram, ErrorsAndWarnings, Warnings,
+    compile_contracts, compile_main, CompileOptions, CompiledContract, CompiledProgram,
+    ContractFunction, ErrorsAndWarnings, Warnings,
 };
 use noirc_frontend::graph::CrateName;
 use noirc_frontend::hir::Context;
@@ -62,31 +63,31 @@ pub(crate) fn run<B: Backend>(
         if package.is_contract() {
             let result = compile_contracts(&mut context, crate_id, &args.compile_options);
             let contracts = report_errors(result, &context, args.compile_options.deny_warnings)?;
+            let optimized_contracts =
+                try_vecmap(contracts, |contract| optimize_contract(backend, contract))?;
 
             // TODO(#1389): I wonder if it is incorrect for nargo-core to know anything about contracts.
             // As can be seen here, It seems like a leaky abstraction where ContractFunctions (essentially CompiledPrograms)
             // are compiled via nargo-core and then the PreprocessedContract is constructed here.
             // This is due to EACH function needing it's own CRS, PKey, and VKey from the backend.
             let preprocessed_contracts: Result<Vec<PreprocessedContract>, CliError<B>> =
-                try_vecmap(contracts, |contract| {
-                    let preprocessed_contract_functions =
-                        try_vecmap(contract.functions, |mut func| {
-                            func.bytecode = optimize_circuit(backend, func.bytecode)?.0;
-                            common_reference_string = update_common_reference_string(
-                                backend,
-                                &common_reference_string,
-                                &func.bytecode,
-                            )
-                            .map_err(CliError::CommonReferenceStringError)?;
-
-                            preprocess_contract_function(
-                                backend,
-                                args.include_keys,
-                                &common_reference_string,
-                                func,
-                            )
-                            .map_err(CliError::ProofSystemCompilerError)
-                        })?;
+                try_vecmap(optimized_contracts, |contract| {
+                    let preprocessed_contract_functions = try_vecmap(contract.functions, |func| {
+                        common_reference_string = update_common_reference_string(
+                            backend,
+                            &common_reference_string,
+                            &func.bytecode,
+                        )
+                        .map_err(CliError::CommonReferenceStringError)?;
+
+                        preprocess_contract_function(
+                            backend,
+                            args.include_keys,
+                            &common_reference_string,
+                            func,
+                        )
+                        .map_err(CliError::ProofSystemCompilerError)
+                    })?;
 
                     Ok(PreprocessedContract {
                         name: contract.name,
@@ -162,6 +163,18 @@ pub(super) fn optimize_circuit<B: Backend>(
     Ok(result)
 }
 
+pub(super) fn optimize_contract<B: Backend>(
+    backend: &B,
+    contract: CompiledContract,
+) -> Result<CompiledContract, CliError<B>> {
+    let functions = try_vecmap(contract.functions, |func| {
+        let optimized_bytecode = optimize_circuit(backend, func.bytecode)?.0;
+        Ok::<_, CliError<B>>(ContractFunction { bytecode: optimized_bytecode, ..func })
+    })?;
+
+    Ok(CompiledContract { functions, ..contract })
+}
+
 /// Helper function for reporting any errors in a Result<(T, Warnings), ErrorsAndWarnings>
 /// structure that is commonly used as a return result in this file.
 pub(crate) fn report_errors<T>(
diff --git a/crates/nargo_cli/src/cli/info_cmd.rs b/crates/nargo_cli/src/cli/info_cmd.rs
index 2518c402605..d972aeb0511 100644
--- a/crates/nargo_cli/src/cli/info_cmd.rs
+++ b/crates/nargo_cli/src/cli/info_cmd.rs
@@ -8,7 +8,10 @@ use noirc_frontend::graph::CrateName;
 
 use crate::{cli::compile_cmd::compile_package, errors::CliError};
 
-use super::{compile_cmd::report_errors, NargoConfig};
+use super::{
+    compile_cmd::{optimize_contract, report_errors},
+    NargoConfig,
+};
 
 /// Provides detailed information on a circuit
 ///
@@ -76,8 +79,10 @@ fn count_opcodes_and_gates_in_contracts<B: Backend>(
     let (mut context, crate_id) = prepare_package(package);
     let result = compile_contracts(&mut context, crate_id, compile_options);
     let contracts = report_errors(result, &context, compile_options.deny_warnings)?;
+    let optimized_contracts =
+        try_vecmap(contracts, |contract| optimize_contract(backend, contract))?;
 
-    for contract in contracts {
+    for contract in optimized_contracts {
         let function_info: Vec<(String, usize, u32)> = try_vecmap(contract.functions, |function| {
             let num_opcodes = function.bytecode.opcodes.len();
             let exact_circuit_size = backend.get_exact_circuit_size(&function.bytecode)?;