diff --git a/src/compile.rs b/src/compile.rs index 64b88b3e..dca119e7 100644 --- a/src/compile.rs +++ b/src/compile.rs @@ -1,6 +1,6 @@ #![cfg(feature = "std")] -use alloc::collections::BTreeSet; +use alloc::collections::{BTreeMap, BTreeSet}; use crate::prelude::*; @@ -12,6 +12,12 @@ use crate::{ use core::{fmt::Write, hash::Hasher}; use std::hash::DefaultHasher; +struct DefInfo<'a> { + rust_name: String, + def: &'a Def>, + refs: BTreeSet<&'a str>, +} + /// Compiles a [`Host`] to Rust, returning a file to replace `gen.rs`. pub fn compile_host(host: &Host) -> String { _compile_host(host).unwrap() @@ -20,11 +26,16 @@ pub fn compile_host(host: &Host) -> String { fn _compile_host(host: &Host) -> Result { let mut code = String::default(); - let defs = host - .defs - .iter() - .filter_map(|(name, def)| Some((name, def.downcast_ref::>()?))) - .map(|(raw_name, def)| (raw_name, sanitize_name(raw_name), def)); + let mut def_infos: BTreeMap<&str, DefInfo<'_>> = BTreeMap::new(); + for (hvmc_name, def) in &host.defs { + if let Some(def) = def.downcast_ref::>() { + def_infos.insert(hvmc_name, DefInfo { + rust_name: sanitize_name(hvmc_name), + refs: refs(host, &def.data.0.instr), + def, + }); + } + } writeln!(code, "#![allow(non_camel_case_types, unused_imports, unused_variables)]")?; writeln!( @@ -36,39 +47,42 @@ fn _compile_host(host: &Host) -> Result { writeln!(code, "pub fn insert_into_host(host: &mut Host) {{")?; // insert empty defs - for (hvmc_name, rust_name, def) in defs.clone() { + for (hvmc_name, DefInfo { rust_name, def, .. }) in &def_infos { let labs = compile_lab_set(&def.labs)?; writeln!( code, - r##" host.insert_def(r#"{hvmc_name}"#, unsafe {{ HostedDef::::new({labs}) }});"## + r##" host.insert_def(r#"{hvmc_name}"#, unsafe {{ HostedDef::::new({labs}) }});"##, )?; } writeln!(code)?; // hoist all unique refs present in the right hand side of some def - for hvmc_name in defs.clone().flat_map(|(_, _, def)| refs(host, &def.data.0.instr)).collect::>() { - let rust_name = sanitize_name(hvmc_name); + for hvmc_name in def_infos.values().flat_map(|info| &info.refs).collect::>() { + let rust_name = &def_infos[hvmc_name].rust_name; writeln!(code, r##" let def_{rust_name} = Port::new_ref(&host.defs[r#"{hvmc_name}"#]);"##)?; } writeln!(code)?; - // initialize each def with the refs it makes use of - for (hvmc_name, rust_name, def) in defs.clone() { - let refs = - refs(host, &def.data.0.instr).iter().map(|r| format!("def_{}", sanitize_name(r))).collect::>().join(", "); + // initialize defs that have refs + for (hvmc_name, DefInfo { rust_name, refs, .. }) in &def_infos { + if refs.is_empty() { + continue; + } + + let fields = refs.iter().map(|r| format!("def_{}", sanitize_name(r))).collect::>().join(", "); writeln!( code, - r##" host.get_mut::>(r#"{hvmc_name}"#).data.0 = Def_{rust_name} {{ {refs} }};"## + r##" host.get_mut::>(r#"{hvmc_name}"#).data.0 = Def_{rust_name} {{ {fields} }};"## )?; } writeln!(code, "}}")?; writeln!(code)?; - for (_, rust_name, def) in defs.clone() { - compile_struct(&mut code, host, &rust_name, def)?; + for DefInfo { rust_name, def, .. } in def_infos.values() { + compile_struct(&mut code, host, rust_name, def)?; } Ok(code)