Untangle ImportEnv::handle_import

2 files changed, 37 insertions, 29 deletions

diff --git a/dhall/src/semantics/resolve/env.rs b/dhall/src/semantics/resolve/env.rs
index 6346a6d..25d2277 100644
--- a/dhall/src/semantics/resolve/env.rs
+++ b/dhall/src/semantics/resolve/env.rs
@@ -69,9 +69,20 @@ impl ImportEnv {
         ImportEnv::default()
     }
 
-    pub fn handle_import(
+    pub fn get_from_cache<'a>(
+        &'a self,
+        location: &ImportLocation,
+    ) -> Option<&'a TypedHir> {
+        self.cache.get(location)
+    }
+
+    pub fn set_cache(&mut self, location: ImportLocation, expr: TypedHir) {
+        self.cache.insert(location, expr);
+    }
+
+    pub fn with_cycle_detection(
         &mut self,
-        mut location: ImportLocation,
+        location: ImportLocation,
         do_resolve: impl FnOnce(&mut Self) -> Result<TypedHir, Error>,
     ) -> Result<TypedHir, Error> {
         if self.stack.contains(&location) {
@@ -79,25 +90,13 @@ impl ImportEnv {
                 ImportError::ImportCycle(self.stack.clone(), location).into()
             );
         }
-        Ok(match self.cache.get(&location) {
-            Some(expr) => expr.clone(),
-            None => {
-                let expr = {
-                    // Push the current location on the stack
-                    self.stack.push(location);
-                    // Resolve the import recursively
-                    // WARNING: do not propagate errors here or the stack will get messed up.
-                    let result = do_resolve(self);
-                    // Remove location from the stack.
-                    location = self.stack.pop().unwrap();
-                    result
-                }?;
-
-                // Add the resolved import to the cache
-                self.cache.insert(location, expr.clone());
-
-                expr
-            }
-        })
+        // Push the current location on the stack
+        self.stack.push(location);
+        // Resolve the import recursively
+        // WARNING: do not propagate errors here or the stack will get messed up.
+        let result = do_resolve(self);
+        // Remove location from the stack.
+        self.stack.pop().unwrap();
+        result
     }
 }
diff --git a/dhall/src/semantics/resolve/resolve.rs b/dhall/src/semantics/resolve/resolve.rs
index 2b401dc..f6d664e 100644
--- a/dhall/src/semantics/resolve/resolve.rs
+++ b/dhall/src/semantics/resolve/resolve.rs
@@ -224,12 +224,10 @@ fn resolve_one_import(
     env: &mut ImportEnv,
     cache: &Cache,
     import: &Import,
-    location: &ImportLocation,
+    location: ImportLocation,
     span: Span,
 ) -> Result<TypedHir, Error> {
-    let do_sanity_check = import.mode != ImportMode::Location;
-    let location = location.chain(&import.location, do_sanity_check)?;
-    env.handle_import(location.clone(), |env| match import.mode {
+    match import.mode {
         ImportMode::Code => {
             let (hir, ty) = cache.caching_import(
                 import,
@@ -278,7 +276,7 @@ fn resolve_one_import(
             let ty = hir.typecheck_noenv()?.ty().clone();
             Ok((hir, ty))
         }
-    })
+    }
 }
 
 /// Desugar a `with` expression.
@@ -396,12 +394,23 @@ fn resolve_with_env(
     cache: &Cache,
     parsed: Parsed,
 ) -> Result<Resolved, Error> {
-    let Parsed(expr, location) = parsed;
+    let Parsed(expr, base_location) = parsed;
     let resolved = traverse_resolve_expr(
         &mut NameEnv::new(),
         &expr,
         &mut |import, span| {
-            resolve_one_import(env, cache, &import, &location, span)
+            let do_sanity_check = import.mode != ImportMode::Location;
+            let location =
+                base_location.chain(&import.location, do_sanity_check)?;
+            if let Some(expr) = env.get_from_cache(&location) {
+                return Ok(expr.clone());
+            }
+            let expr = env.with_cycle_detection(location.clone(), |env| {
+                resolve_one_import(env, cache, &import, location.clone(), span)
+            })?;
+            // Add the resolved import to the cache
+            env.set_cache(location, expr.clone());
+            Ok(expr)
         },
     )?;
     Ok(Resolved(resolved))