From 2bceaf0503094a3657747df24d568a9c31d8cb2a Mon Sep 17 00:00:00 2001 From: Folkert Date: Tue, 13 Oct 2020 16:10:38 +0200 Subject: [PATCH] implement LayoutCache on top of ena --- Cargo.lock | 1 + compiler/mono/Cargo.toml | 1 + compiler/mono/src/ir.rs | 3 + compiler/mono/src/layout.rs | 109 +++++++++++++++++++++++++++++++----- 4 files changed, 101 insertions(+), 13 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 7a80914d80..88b35d558d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2510,6 +2510,7 @@ dependencies = [ "roc_solve", "roc_types", "roc_unify", + "ven_ena", "ven_pretty", ] diff --git a/compiler/mono/Cargo.toml b/compiler/mono/Cargo.toml index 079e0f0da0..0a812aa98b 100644 --- a/compiler/mono/Cargo.toml +++ b/compiler/mono/Cargo.toml @@ -17,6 +17,7 @@ roc_solve = { path = "../solve" } roc_problem = { path = "../problem" } ven_pretty = { path = "../../vendor/pretty" } bumpalo = { version = "3.2", features = ["collections"] } +ven_ena = { path = "../../vendor/ena" } [dev-dependencies] roc_constrain = { path = "../constrain" } diff --git a/compiler/mono/src/ir.rs b/compiler/mono/src/ir.rs index 80628ced52..7a143d6d31 100644 --- a/compiler/mono/src/ir.rs +++ b/compiler/mono/src/ir.rs @@ -1490,6 +1490,7 @@ fn specialize_solved_type<'a>( use roc_types::subs::VarStore; let snapshot = env.subs.snapshot(); + let cache_snapshot = layout_cache.snapshot(); let mut free_vars = FreeVars::default(); let mut var_store = VarStore::new_from_subs(env.subs); @@ -1514,10 +1515,12 @@ fn specialize_solved_type<'a>( .from_var(&env.arena, fn_var, env.subs) .unwrap_or_else(|err| panic!("TODO handle invalid function {:?}", err)); env.subs.rollback_to(snapshot); + layout_cache.rollback_to(cache_snapshot); Ok((proc, layout)) } Err(error) => { env.subs.rollback_to(snapshot); + layout_cache.rollback_to(cache_snapshot); Err(error) } } diff --git a/compiler/mono/src/layout.rs b/compiler/mono/src/layout.rs index 2208121a9e..8cf3a1af2e 100644 --- a/compiler/mono/src/layout.rs +++ b/compiler/mono/src/layout.rs @@ -226,9 +226,51 @@ impl<'a> Layout<'a> { } /// Avoid recomputing Layout from Variable multiple times. +/// We use `ena` for easy snapshots and rollbacks of the cache. +/// During specialization, a type variable `a` can be specialized to different layouts, +/// e.g. `identity : a -> a` could be specialized to `Bool -> Bool` or `Str -> Str`. +/// Therefore in general it's invalid to store a map from variables to layouts +/// But if we're careful when to invalidate certain keys, we still get some benefit #[derive(Default, Debug)] pub struct LayoutCache<'a> { - layouts: MutMap, LayoutProblem>>, + layouts: ven_ena::unify::UnificationTable>>, +} + +#[derive(Debug, Clone)] +pub enum CachedLayout<'a> { + Cached(Layout<'a>), + NotCached, + Problem(LayoutProblem), +} + +/// Must wrap so we can define a specific UnifyKey instance +/// PhantomData so we can store the 'a lifetime, which is needed to implement the UnifyKey trait, +/// specifically so we can use `type Value = CachedLayout<'a>` +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub struct CachedVariable<'a>(Variable, std::marker::PhantomData<&'a ()>); + +impl<'a> CachedVariable<'a> { + fn new(var: Variable) -> Self { + CachedVariable(var, std::marker::PhantomData) + } +} + +// use ven_ena::unify::{InPlace, Snapshot, UnificationTable, UnifyKey}; + +impl<'a> ven_ena::unify::UnifyKey for CachedVariable<'a> { + type Value = CachedLayout<'a>; + + fn index(&self) -> u32 { + self.0.index() + } + + fn from_index(index: u32) -> Self { + CachedVariable(Variable::from_index(index), std::marker::PhantomData) + } + + fn tag() -> &'static str { + "CachedVariable" + } } impl<'a> LayoutCache<'a> { @@ -244,19 +286,60 @@ impl<'a> LayoutCache<'a> { // Store things according to the root Variable, to avoid duplicate work. let var = subs.get_root_key_without_compacting(var); - let mut env = Env { - arena, - subs, - seen: MutSet::default(), - }; + let cached_var = CachedVariable::new(var); - /* - self.layouts - .entry(var) - .or_insert_with(|| Layout::from_var(&mut env, var)) - .clone() - */ - Layout::from_var(&mut env, var) + self.expand_to_fit(cached_var); + + use CachedLayout::*; + match self.layouts.probe_value(cached_var) { + Cached(result) => Ok(result), + Problem(problem) => Err(problem), + NotCached => { + let mut env = Env { + arena, + subs, + seen: MutSet::default(), + }; + + let result = Layout::from_var(&mut env, var); + + let cached_layout = match &result { + Ok(layout) => Cached(layout.clone()), + Err(problem) => Problem(problem.clone()), + }; + + self.layouts + .update_value(cached_var, |existing| existing.value = cached_layout); + + result + } + } + } + + fn expand_to_fit(&mut self, var: CachedVariable<'a>) { + use ven_ena::unify::UnifyKey; + + let required = (var.index() as isize) - (self.layouts.len() as isize) + 1; + if required > 0 { + self.layouts.reserve(required as usize); + + for _ in 0..required { + self.layouts.new_key(CachedLayout::NotCached); + } + } + } + + pub fn snapshot( + &mut self, + ) -> ven_ena::unify::Snapshot>> { + self.layouts.snapshot() + } + + pub fn rollback_to( + &mut self, + snapshot: ven_ena::unify::Snapshot>>, + ) { + self.layouts.rollback_to(snapshot) } }