Make global store of imports and import results

6 files changed, 217 insertions, 62 deletions

diff --git a/Cargo.lock b/Cargo.lock
index fdbb3ed..b89c7cd 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -300,6 +300,7 @@ dependencies = [
  "annotate-snippets",
  "anyhow",
  "colored-diff",
+ "elsa",
  "fs_extra",
  "hex",
  "itertools 0.9.0",
@@ -374,6 +375,15 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e78d4f1cc4ae33bbfc157ed5d5a5ef3bc29227303d595861deb238fcec4e9457"
 
 [[package]]
+name = "elsa"
+version = "1.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0daf12e2857c9485cb2386b18a9ecd5a17c7cee2fd10afb87d436b10ced678e7"
+dependencies = [
+ "stable_deref_trait",
+]
+
+[[package]]
 name = "encoding_rs"
 version = "0.8.24"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1457,6 +1467,12 @@ dependencies = [
 ]
 
 [[package]]
+name = "stable_deref_trait"
+version = "1.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3"
+
+[[package]]
 name = "static_assertions"
 version = "1.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
diff --git a/dhall/Cargo.toml b/dhall/Cargo.toml
index 1b4f9a3..1e9cfbc 100644
--- a/dhall/Cargo.toml
+++ b/dhall/Cargo.toml
@@ -21,6 +21,7 @@ path = "tests/spec.rs"
 
 [dependencies]
 annotate-snippets = "0.9.0"
+elsa = "1.3.2"
 hex = "0.4.2"
 itertools = "0.9.0"
 lazy_static = "1.4.0"
diff --git a/dhall/src/ctxt.rs b/dhall/src/ctxt.rs
new file mode 100644
index 0000000..1d97232
--- /dev/null
+++ b/dhall/src/ctxt.rs
@@ -0,0 +1,109 @@
+use elsa::vec::FrozenVec;
+use once_cell::sync::OnceCell;
+
+use crate::semantics::{Import, ImportLocation};
+use crate::syntax::Span;
+use crate::Typed;
+
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
+pub struct ImportId(usize);
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
+pub struct ImportResultId(usize);
+
+struct StoredImport {
+    base_location: ImportLocation,
+    import: Import,
+    span: Span,
+    result: OnceCell<ImportResultId>,
+}
+
+#[derive(Default)]
+struct CtxtS {
+    imports: FrozenVec<Box<StoredImport>>,
+    import_results: FrozenVec<Box<Typed>>,
+}
+
+/// Context for the dhall compiler. Stores various global maps.
+#[derive(Copy, Clone)]
+pub struct Ctxt<'cx>(&'cx CtxtS);
+
+impl Ctxt<'_> {
+    pub fn with_new<T>(f: impl for<'cx> FnOnce(Ctxt<'cx>) -> T) -> T {
+        let cx = CtxtS::default();
+        let cx = Ctxt(&cx);
+        f(cx)
+    }
+}
+impl<'cx> Ctxt<'cx> {
+    fn get_stored_import(self, import: ImportId) -> &'cx StoredImport {
+        self.0.imports.get(import.0).unwrap()
+    }
+    pub fn get_import_result(self, id: ImportResultId) -> &'cx Typed {
+        &self.0.import_results.get(id.0).unwrap()
+    }
+
+    /// Store an import and the location relative to which it must be resolved.
+    pub fn push_import(
+        self,
+        base_location: ImportLocation,
+        import: Import,
+        span: Span,
+    ) -> ImportId {
+        let stored = StoredImport {
+            base_location,
+            import,
+            span,
+            result: OnceCell::new(),
+        };
+        let id = self.0.imports.len();
+        self.0.imports.push(Box::new(stored));
+        ImportId(id)
+    }
+    /// Store the result of fetching an import.
+    pub fn push_import_result(self, res: Typed) -> ImportResultId {
+        let id = self.0.import_results.len();
+        self.0.import_results.push(Box::new(res));
+        ImportResultId(id)
+    }
+
+    pub fn get_import(self, import: ImportId) -> &'cx Import {
+        &self.get_stored_import(import).import
+    }
+    pub fn get_import_base_location(
+        self,
+        import: ImportId,
+    ) -> &'cx ImportLocation {
+        &self.get_stored_import(import).base_location
+    }
+    pub fn get_import_span(self, import: ImportId) -> Span {
+        self.get_stored_import(import).span.clone()
+    }
+    /// Get the result of fetching this import. Returns `None` if the result has not yet been
+    /// fetched.
+    pub fn get_result_of_import(self, import: ImportId) -> Option<&'cx Typed> {
+        let res = self.get_resultid_of_import(import)?;
+        Some(self.get_import_result(res))
+    }
+    /// Get the id of the result of fetching this import. Returns `None` if the result has not yet
+    /// been fetched.
+    pub fn get_resultid_of_import(
+        self,
+        import: ImportId,
+    ) -> Option<ImportResultId> {
+        self.get_stored_import(import).result.get().copied()
+    }
+    /// Store the result of fetching this import.
+    pub fn set_result_of_import(
+        self,
+        import: ImportId,
+        res: Typed,
+    ) -> ImportResultId {
+        let res = self.push_import_result(res);
+        self.set_resultid_of_import(import, res);
+        res
+    }
+    /// Store the result of fetching this import.
+    pub fn set_resultid_of_import(self, import: ImportId, res: ImportResultId) {
+        let _ = self.get_stored_import(import).result.set(res);
+    }
+}
diff --git a/dhall/src/lib.rs b/dhall/src/lib.rs
index 6747eff..d079e92 100644
--- a/dhall/src/lib.rs
+++ b/dhall/src/lib.rs
@@ -9,6 +9,7 @@
 )]
 
 pub mod builtins;
+pub mod ctxt;
 pub mod error;
 pub mod operations;
 pub mod semantics;
@@ -26,6 +27,8 @@ use crate::semantics::resolve::ImportLocation;
 use crate::semantics::{typecheck, typecheck_with, Hir, Nir, Tir, Type};
 use crate::syntax::Expr;
 
+pub use ctxt::{Ctxt, ImportId, ImportResultId};
+
 #[derive(Debug, Clone)]
 pub struct Parsed(Expr, ImportLocation);
 
diff --git a/dhall/src/semantics/resolve/env.rs b/dhall/src/semantics/resolve/env.rs
index 29dd16b..82f21e8 100644
--- a/dhall/src/semantics/resolve/env.rs
+++ b/dhall/src/semantics/resolve/env.rs
@@ -3,7 +3,7 @@ use std::collections::HashMap;
 use crate::error::{Error, ImportError};
 use crate::semantics::{AlphaVar, Cache, ImportLocation, VarEnv};
 use crate::syntax::{Hash, Label, V};
-use crate::Typed;
+use crate::{Ctxt, ImportResultId, Typed};
 
 /// Environment for resolving names.
 #[derive(Debug, Clone, Default)]
@@ -11,14 +11,13 @@ pub struct NameEnv {
     names: Vec<Label>,
 }
 
-pub type MemCache = HashMap<ImportLocation, Typed>;
 pub type CyclesStack = Vec<ImportLocation>;
 
 /// Environment for resolving imports
-#[derive(Debug)]
-pub struct ImportEnv {
-    disk_cache: Option<Cache>, // Missing if it failed to initialize
-    mem_cache: MemCache,
+pub struct ImportEnv<'cx> {
+    cx: Ctxt<'cx>,
+    disk_cache: Option<Cache>, // `None` if it failed to initialize
+    mem_cache: HashMap<ImportLocation, ImportResultId>,
     stack: CyclesStack,
 }
 
@@ -66,26 +65,28 @@ impl NameEnv {
     }
 }
 
-impl ImportEnv {
-    pub fn new() -> Self {
+impl<'cx> ImportEnv<'cx> {
+    pub fn new(cx: Ctxt<'cx>) -> Self {
         ImportEnv {
+            cx,
             disk_cache: Cache::new().ok(),
             mem_cache: Default::default(),
             stack: Default::default(),
         }
     }
 
+    pub fn cx(&self) -> Ctxt<'cx> {
+        self.cx
+    }
+
     pub fn get_from_mem_cache(
-        &mut self,
+        &self,
         location: &ImportLocation,
-    ) -> Option<Typed> {
-        Some(self.mem_cache.get(location)?.clone())
+    ) -> Option<ImportResultId> {
+        Some(*self.mem_cache.get(location)?)
     }
 
-    pub fn get_from_disk_cache(
-        &mut self,
-        hash: &Option<Hash>,
-    ) -> Option<Typed> {
+    pub fn get_from_disk_cache(&self, hash: &Option<Hash>) -> Option<Typed> {
         let hash = hash.as_ref()?;
         let expr = self.disk_cache.as_ref()?.get(hash).ok()?;
         Some(expr)
@@ -94,12 +95,12 @@ impl ImportEnv {
     pub fn write_to_mem_cache(
         &mut self,
         location: ImportLocation,
-        expr: Typed,
+        result: ImportResultId,
     ) {
-        self.mem_cache.insert(location, expr);
+        self.mem_cache.insert(location, result);
     }
 
-    pub fn write_to_disk_cache(&mut self, hash: &Option<Hash>, expr: &Typed) {
+    pub fn write_to_disk_cache(&self, hash: &Option<Hash>, expr: &Typed) {
         if let Some(disk_cache) = self.disk_cache.as_ref() {
             if let Some(hash) = hash {
                 let _ = disk_cache.insert(hash, &expr);
diff --git a/dhall/src/semantics/resolve/resolve.rs b/dhall/src/semantics/resolve/resolve.rs
index 16987de..da17c0a 100644
--- a/dhall/src/semantics/resolve/resolve.rs
+++ b/dhall/src/semantics/resolve/resolve.rs
@@ -15,7 +15,7 @@ use crate::syntax::{
     Expr, ExprKind, FilePath, FilePrefix, Hash, ImportMode, ImportTarget, Span,
     UnspannedExpr, URL,
 };
-use crate::{Parsed, Resolved, Typed};
+use crate::{Ctxt, ImportId, Parsed, Resolved, Typed};
 
 // TODO: evaluate import headers
 pub type Import = syntax::Import<()>;
@@ -227,7 +227,11 @@ impl ImportLocation {
     }
 
     /// Fetches the expression corresponding to this location.
-    fn fetch(&self, env: &mut ImportEnv, span: Span) -> Result<Typed, Error> {
+    fn fetch<'cx>(
+        &self,
+        env: &mut ImportEnv<'cx>,
+        span: Span,
+    ) -> Result<Typed, Error> {
         let (hir, ty) = match self.mode {
             ImportMode::Code => {
                 let parsed = self.kind.fetch_dhall()?;
@@ -386,8 +390,63 @@ fn traverse_resolve_expr(
     })
 }
 
-fn resolve_with_env(
-    env: &mut ImportEnv,
+/// Fetch the import and store the result in the global context.
+fn fetch_import<'cx>(
+    env: &mut ImportEnv<'cx>,
+    import_id: ImportId,
+) -> Result<(), Error> {
+    let base_location = env.cx().get_import_base_location(import_id);
+    let import = env.cx().get_import(import_id);
+    let span = env.cx().get_import_span(import_id);
+    let location = base_location.chain(&import)?;
+
+    // If the import is in the in-memory cache, or the hash is in the on-disk cache, return
+    // the cached contents.
+    if let Some(res_id) = env.get_from_mem_cache(&location) {
+        env.cx().set_resultid_of_import(import_id, res_id);
+        // The same location may be used with different or no hashes. Thus we need to check
+        // the hashes every time.
+        let typed = env.cx().get_import_result(res_id);
+        check_hash(import, typed, span)?;
+        env.write_to_disk_cache(&import.hash, typed);
+        return Ok(());
+    }
+    if let Some(typed) = env.get_from_disk_cache(&import.hash) {
+        // No need to check the hash, it was checked before reading the file. We also don't
+        // write to the in-memory cache, because the location might be completely unrelated
+        // to the cached file (e.g. `missing sha256:...` is valid).
+        // This actually means that importing many times a same hashed import will take
+        // longer than importing many times a same non-hashed import.
+        env.cx().set_result_of_import(import_id, typed);
+        return Ok(());
+    }
+
+    // Resolve this import, making sure that recursive imports don't cycle back to the
+    // current one.
+    let res = env.with_cycle_detection(location.clone(), |env| {
+        location.fetch(env, span.clone())
+    });
+    let typed = match res {
+        Ok(typed) => typed,
+        Err(e) => mkerr(
+            ErrorBuilder::new("error")
+                .span_err(span.clone(), e.to_string())
+                .format(),
+        )?,
+    };
+
+    // Add the resolved import to the caches
+    let import = env.cx().get_import(import_id);
+    check_hash(import, &typed, span)?;
+    env.write_to_disk_cache(&import.hash, &typed);
+    let res_id = env.cx().set_result_of_import(import_id, typed);
+    env.write_to_mem_cache(location, res_id);
+
+    Ok(())
+}
+
+fn resolve_with_env<'cx>(
+    env: &mut ImportEnv<'cx>,
     parsed: Parsed,
 ) -> Result<Resolved, Error> {
     let Parsed(expr, base_location) = parsed;
@@ -395,52 +454,18 @@ fn resolve_with_env(
         &mut NameEnv::new(),
         &expr,
         &mut |import, span| {
-            let location = base_location.chain(&import)?;
-
-            // If the import is in the in-memory cache, or the hash is in the on-disk cache, return
-            // the cached contents.
-            if let Some(typed) = env.get_from_mem_cache(&location) {
-                // The same location may be used with different or no hashes. Thus we need to check
-                // the hashes every time.
-                check_hash(&import, &typed, span)?;
-                env.write_to_disk_cache(&import.hash, &typed);
-                return Ok(typed);
-            }
-            if let Some(typed) = env.get_from_disk_cache(&import.hash) {
-                // No need to check the hash, it was checked before reading the file. We also don't
-                // write to the in-memory cache, because the location might be completely unrelated
-                // to the cached file (e.g. `missing sha256:...` is valid).
-                // This actually means that importing many times a same hashed import will take
-                // longer than importing many times a same non-hashed import.
-                return Ok(typed);
-            }
-
-            // Resolve this import, making sure that recursive imports don't cycle back to the
-            // current one.
-            let res = env.with_cycle_detection(location.clone(), |env| {
-                location.fetch(env, span.clone())
-            });
-            let typed = match res {
-                Ok(typed) => typed,
-                Err(e) => mkerr(
-                    ErrorBuilder::new("error")
-                        .span_err(span.clone(), e.to_string())
-                        .format(),
-                )?,
-            };
-
-            // Add the resolved import to the caches
-            check_hash(&import, &typed, span)?;
-            env.write_to_disk_cache(&import.hash, &typed);
-            env.write_to_mem_cache(location, typed.clone());
-            Ok(typed)
+            let import_id =
+                env.cx().push_import(base_location.clone(), import, span);
+            fetch_import(env, import_id)?;
+            // TODO: store import id in Hir
+            Ok(env.cx().get_result_of_import(import_id).unwrap().clone())
         },
     )?;
     Ok(Resolved(resolved))
 }
 
 pub fn resolve(parsed: Parsed) -> Result<Resolved, Error> {
-    resolve_with_env(&mut ImportEnv::new(), parsed)
+    Ctxt::with_new(|cx| resolve_with_env(&mut ImportEnv::new(cx), parsed))
 }
 
 pub fn skip_resolve_expr(expr: &Expr) -> Result<Hir, Error> {