diff --git a/compiler/build/src/program.rs b/compiler/build/src/program.rs index df8d1eb578..1a8bc8f33c 100644 --- a/compiler/build/src/program.rs +++ b/compiler/build/src/program.rs @@ -227,8 +227,14 @@ pub fn build( // Populate Procs further and get the low-level Expr from the canonical Expr let main_body = Expr::new(&mut mono_env, loc_expr.value, &mut procs); + let mut headers = { + let num_headers = match &procs.pending_specializations { + Some(map) => map.len(), + None => 0, + }; - let mut headers = Vec::with_capacity(procs.pending_specializations.len()); + Vec::with_capacity(num_headers) + }; let mut layout_cache = LayoutCache::default(); let mut procs = roc_mono::expr::specialize_all(&mut mono_env, procs, &mut layout_cache); diff --git a/compiler/gen/tests/helpers/eval.rs b/compiler/gen/tests/helpers/eval.rs index a29049dcd4..2997e0bab2 100644 --- a/compiler/gen/tests/helpers/eval.rs +++ b/compiler/gen/tests/helpers/eval.rs @@ -75,8 +75,14 @@ macro_rules! assert_llvm_evals_to { }; let main_body = Expr::new(&mut mono_env, loc_expr.value, &mut procs); + let mut headers = { + let num_headers = match &procs.pending_specializations { + Some(map) => map.len(), + None => 0 + }; - let mut headers = Vec::with_capacity(procs.pending_specializations.len()); + Vec::with_capacity(num_headers) + }; let mut layout_cache = roc_mono::layout::LayoutCache::default(); let mut procs = roc_mono::expr::specialize_all(&mut mono_env, procs, &mut layout_cache); @@ -246,7 +252,14 @@ macro_rules! assert_opt_evals_to { }; let main_body = Expr::new(&mut mono_env, loc_expr.value, &mut procs); - let mut headers = Vec::with_capacity(procs.pending_specializations.len()); + let mut headers = { + let num_headers = match &procs.pending_specializations { + Some(map) => map.len(), + None => 0 + }; + + Vec::with_capacity(num_headers) + }; let mut layout_cache = roc_mono::layout::LayoutCache::default(); let mut procs = roc_mono::expr::specialize_all(&mut mono_env, procs, &mut layout_cache); diff --git a/compiler/mono/src/expr.rs b/compiler/mono/src/expr.rs index 52f50fd5d5..c2a7765d80 100644 --- a/compiler/mono/src/expr.rs +++ b/compiler/mono/src/expr.rs @@ -38,7 +38,8 @@ pub struct Proc<'a> { pub struct Procs<'a> { pub partial_procs: MutMap>, pub module_thunks: MutSet, - pub pending_specializations: MutMap, PendingSpecialization<'a>>>, + pub pending_specializations: + Option, PendingSpecialization<'a>>>>, pub specialized: MutMap<(Symbol, Layout<'a>), Proc<'a>>, pub runtime_errors: MutSet, } @@ -81,8 +82,7 @@ impl<'a> Procs<'a> { ret_var: Variable, layout_cache: &mut LayoutCache<'a>, ) -> Layout<'a> { - let (pattern_vars, pattern_symbols, body) = - patterns_to_when(env, loc_args, ret_var, loc_body); + let (_, pattern_symbols, body) = patterns_to_when(env, loc_args, ret_var, loc_body); // an anonymous closure. These will always be specialized already // by the surrounding context, so we can add pending specializations @@ -104,20 +104,19 @@ impl<'a> Procs<'a> { layout } +} - fn add_pending_specialization( - &mut self, - symbol: Symbol, - layout: Layout<'a>, - pending: PendingSpecialization<'a>, - ) { - let all_pending = self - .pending_specializations - .entry(symbol) - .or_insert_with(|| HashMap::with_capacity_and_hasher(1, default_hasher())); +fn add_pending<'a>( + pending_specializations: &mut MutMap, PendingSpecialization<'a>>>, + symbol: Symbol, + layout: Layout<'a>, + pending: PendingSpecialization<'a>, +) { + let all_pending = pending_specializations + .entry(symbol) + .or_insert_with(|| HashMap::with_capacity_and_hasher(1, default_hasher())); - all_pending.insert(layout, pending); - } + all_pending.insert(layout, pending); } #[derive(Default)] @@ -1319,19 +1318,40 @@ fn call_by_name<'a>( fn_var, }; - // TODO should pending_procs hold a Rc? - let partial_proc = procs - .partial_procs - .get(&proc_name) - .unwrap_or_else(|| panic!("Could not find partial_proc for {:?}", proc_name)) - .clone(); - - match specialize(env, procs, proc_name, layout_cache, pending, partial_proc) { - Ok(proc) => { - procs.specialized.insert((proc_name, layout.clone()), proc); + // When requested (that is, when procs.pending_specializations is `Some`), + // store a pending specialization rather than specializing immediately. + // + // We do this so that we can do specialization in two passes: first, + // build the mono_expr with all the specialized calls in place (but + // no specializations performed yet), and then second, *after* + // de-duplicating requested specializations (since multiple modules + // which could be getting monomorphized in parallel might request + // the same specialization independently), we work through the + // queue of pending specializations to complete each specialization + // exactly once. + match &mut procs.pending_specializations { + Some(pending_specializations) => { + // register the pending specialization, so this gets code genned later + add_pending(pending_specializations, proc_name, layout.clone(), pending); } - Err(_) => { - procs.runtime_errors.insert(proc_name); + None => { + // TODO should pending_procs hold a Rc? + let partial_proc = procs + .partial_procs + .get(&proc_name) + .unwrap_or_else(|| { + panic!("Could not find partial_proc for {:?}", proc_name) + }) + .clone(); + + match specialize(env, procs, proc_name, layout_cache, pending, partial_proc) { + Ok(proc) => { + procs.specialized.insert((proc_name, layout.clone()), proc); + } + Err(_) => { + procs.runtime_errors.insert(proc_name); + } + } } } @@ -1354,61 +1374,39 @@ pub fn specialize_all<'a>( mut procs: Procs<'a>, layout_cache: &mut LayoutCache<'a>, ) -> Procs<'a> { - let mut is_finished = procs.pending_specializations.is_empty(); + let mut pending_specializations = procs.pending_specializations.unwrap_or(MutMap::default()); - // TODO replace this synchronous loop with a work-stealing queue which - // processes each entry in pending_specializations in parallel, one - // module at a time (because the &mut env will need exclusive access to - // that module's IdentIds; the only reason Env is &mut in specialize is - // that we need to generate unique symbols and register them in them module's - // IdentIds). - while !is_finished { - let Procs { - partial_procs, - module_thunks, - mut pending_specializations, - specialized, - runtime_errors, - } = procs; + // When calling from_can, pending_specializations should be unavailable. + // This must be a single pass, and we must not add any more entries to it! + procs.pending_specializations = None; - procs = Procs { - partial_procs, - module_thunks, - pending_specializations: MutMap::default(), - specialized, - runtime_errors, - }; + for (name, mut by_layout) in pending_specializations.drain() { + // Use the function's symbol's home module as the home module + // when doing canonicalization. This will be important to determine + // whether or not it's safe to defer specialization. + env.home = name.module_id(); - for (name, mut by_layout) in pending_specializations.drain() { - // Use the function's symbol's home module as the home module - // when doing canonicalization. This will be important to determine - // whether or not it's safe to defer specialization. - env.home = name.module_id(); + for (layout, pending) in by_layout.drain() { + // If we've already seen this (Symbol, Layout) combination before, + // don't try to specialize it again. If we do, we'll loop forever! + if !procs.specialized.contains_key(&(name, layout.clone())) { + // TODO should pending_procs hold a Rc? + let partial_proc = procs + .partial_procs + .get(&name) + .unwrap_or_else(|| panic!("Could not find partial_proc for {:?}", name)) + .clone(); - for (layout, pending) in by_layout.drain() { - // If we've already seen this (Symbol, Layout) combination before, - // don't try to specialize it again. If we do, we'll loop forever! - if !procs.specialized.contains_key(&(name, layout.clone())) { - // TODO should pending_procs hold a Rc? - let partial_proc = procs - .partial_procs - .get(&name) - .unwrap_or_else(|| panic!("Could not find partial_proc for {:?}", name)) - .clone(); - - match specialize(env, &mut procs, name, layout_cache, pending, partial_proc) { - Ok(proc) => { - procs.specialized.insert((name, layout), proc); - } - Err(_) => { - procs.runtime_errors.insert(name); - } + match specialize(env, &mut procs, name, layout_cache, pending, partial_proc) { + Ok(proc) => { + procs.specialized.insert((name, layout), proc); + } + Err(_) => { + procs.runtime_errors.insert(name); } } } } - - is_finished = procs.pending_specializations.is_empty(); } procs