feat: separate contract/program compilation from IO (#967)

* feat: separate contract/program compilation from IO

* chore: merge some imports

crates/nargo/src/cli/compile_cmd.rs

@@ -1,8 +1,8 @@
-use acvm::acir::circuit::Circuit;
-use acvm::ProofSystemCompiler;
-use noirc_driver::CompileOptions;
-use noirc_driver::Driver;
-use noirc_frontend::node_interner::FuncId;
+use acvm::{acir::circuit::Circuit, ProofSystemCompiler};
+use iter_extended::{try_btree_map, try_vecmap};
+use noirc_driver::{CompileOptions, CompiledProgram, Driver};
+use noirc_frontend::{hir::def_map::Contract, node_interner::FuncId};
+use std::collections::BTreeMap;
 use std::path::{Path, PathBuf};
 
 use clap::Args;
@@ -26,27 +26,43 @@ pub(crate) struct CompileCommand {
     compile_options: CompileOptions,
 }
 
-pub(crate) fn run(mut args: CompileCommand, config: NargoConfig) -> Result<(), CliError> {
+struct CompiledContract {
+    /// The name of the contract.
+    name: String,
+    /// Each of the contract's functions are compiled into a separate `CompiledProgram`
+    /// stored in this `BTreeMap`.
+    functions: BTreeMap<String, CompiledProgram>,
+}
+
+pub(crate) fn run(args: CompileCommand, config: NargoConfig) -> Result<(), CliError> {
     let driver = check_crate(&config.program_dir, &args.compile_options)?;
 
     let mut circuit_dir = config.program_dir;
     circuit_dir.push(TARGET_DIR);
 
     // If contracts is set we're compiling every function in a 'contract' rather than just 'main'.
     if args.contracts {
-        let circuit_name = args.circuit_name.clone();
-
-        for contract in driver.get_all_contracts() {
-            for function in contract.functions {
-                let name = driver.function_name(function);
-                args.circuit_name = format!("{}-{}-{name}", circuit_name, &contract.name);
-                compile_and_save_program(&driver, function, &args, &circuit_dir)?;
-            }
+        let compiled_contracts = try_vecmap(driver.get_all_contracts(), |contract| {
+            compile_contract(&driver, contract, &args.compile_options)
+        })?;
+
+        // Flatten each contract into a list of its functions, each being assigned a unique name.
+        let compiled_programs = compiled_contracts.into_iter().flat_map(|contract| {
+            let contract_id = format!("{}-{}", args.circuit_name, &contract.name);
+            contract.functions.into_iter().map(move |(function, program)| {
+                let program_name = format!("{}-{}", contract_id, function);
+                (program_name, program)
+            })
+        });
+
+        for (circuit_name, compiled_program) in compiled_programs {
+            save_and_preprocess_program(&compiled_program, &circuit_name, &circuit_dir)?
         }
         Ok(())
     } else {
         let main = driver.main_function().map_err(|_| CliError::CompilationError)?;
-        compile_and_save_program(&driver, main, &args, &circuit_dir)
+        let program = compile_program(&driver, main, &args.compile_options, &args.circuit_name)?;
+        save_and_preprocess_program(&program, &args.circuit_name, &circuit_dir)
     }
 }
 
@@ -57,20 +73,48 @@ fn setup_driver(program_dir: &Path) -> Result<Driver, CliError> {
     Ok(driver)
 }
 
-/// Compile and save a program to disk with the given main function.
-fn compile_and_save_program(
+fn check_crate(program_dir: &Path, options: &CompileOptions) -> Result<Driver, CliError> {
+    let mut driver = setup_driver(program_dir)?;
+    driver.check_crate(options).map_err(|_| CliError::CompilationError)?;
+    Ok(driver)
+}
+
+/// Compiles all of the functions associated with a Noir contract.
+fn compile_contract(
+    driver: &Driver,
+    contract: Contract,
+    compile_options: &CompileOptions,
+) -> Result<CompiledContract, CliError> {
+    let functions = try_btree_map(&contract.functions, |function| {
+        let function_name = driver.function_name(*function).to_owned();
+        let program_id = format!("{}-{}", contract.name, function_name);
+
+        compile_program(driver, *function, compile_options, &program_id)
+            .map(|program| (function_name, program))
+    })?;
+
+    Ok(CompiledContract { name: contract.name, functions })
+}
+
+fn compile_program(
     driver: &Driver,
     main: FuncId,
-    args: &CompileCommand,
+    compile_options: &CompileOptions,
+    program_id: &str,
+) -> Result<CompiledProgram, CliError> {
+    driver
+        .compile_no_check(compile_options, main)
+        .map_err(|_| CliError::Generic(format!("'{}' failed to compile", program_id)))
+}
+
+/// Save a program to disk along with proving and verification keys.
+fn save_and_preprocess_program(
+    compiled_program: &CompiledProgram,
+    circuit_name: &str,
     circuit_dir: &Path,
 ) -> Result<(), CliError> {
-    let compiled_program = driver
-        .compile_no_check(&args.compile_options, main)
-        .map_err(|_| CliError::Generic(format!("'{}' failed to compile", args.circuit_name)))?;
-
-    save_program_to_file(&compiled_program, &args.circuit_name, circuit_dir);
-
-    preprocess_with_path(&args.circuit_name, circuit_dir, &compiled_program.circuit)?;
+    save_program_to_file(compiled_program, circuit_name, circuit_dir);
+    preprocess_with_path(circuit_name, circuit_dir, &compiled_program.circuit)?;
     Ok(())
 }
 
@@ -82,12 +126,6 @@ pub(crate) fn compile_circuit(
     driver.compile_main(compile_options).map_err(|_| CliError::CompilationError)
 }
 
-fn check_crate(program_dir: &Path, options: &CompileOptions) -> Result<Driver, CliError> {
-    let mut driver = setup_driver(program_dir)?;
-    driver.check_crate(options).map_err(|_| CliError::CompilationError)?;
-    Ok(driver)
-}
-
 fn preprocess_with_path<P: AsRef<Path>>(
     key_name: &str,
     preprocess_dir: P,