Resolve imports and alternatives outside of the ast traversal

6 files changed, 294 insertions, 130 deletions

diff --git a/dhall/src/ctxt.rs b/dhall/src/ctxt.rs
index 3809bc9..aad1a1b 100644
--- a/dhall/src/ctxt.rs
+++ b/dhall/src/ctxt.rs
@@ -3,29 +3,19 @@ use once_cell::sync::OnceCell;
 use std::marker::PhantomData;
 use std::ops::{Deref, Index};
 
-use crate::semantics::{Import, ImportLocation};
+use crate::semantics::{Import, ImportLocation, ImportNode};
 use crate::syntax::Span;
 use crate::Typed;
 
-#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
-pub struct ImportId<'cx>(usize, PhantomData<&'cx ()>);
-#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
-pub struct ImportResultId<'cx>(usize, PhantomData<&'cx ()>);
-
-/// What's stored for each `ImportId`. Allows getting and setting a result for this import.
-pub struct StoredImport<'cx> {
-    cx: Ctxt<'cx>,
-    pub base_location: ImportLocation,
-    pub import: Import,
-    pub span: Span,
-    result: OnceCell<ImportResultId<'cx>>,
-}
+/////////////////////////////////////////////////////////////////////////////////////////////////////
+// Ctxt
 
 /// Implementation detail. Made public for the `Index` instances.
 #[derive(Default)]
 pub struct CtxtS<'cx> {
     imports: FrozenVec<Box<StoredImport<'cx>>>,
-    import_results: FrozenVec<Box<Typed<'cx>>>,
+    import_alternatives: FrozenVec<Box<StoredImportAlternative<'cx>>>,
+    import_results: FrozenVec<Box<StoredImportResult<'cx>>>,
 }
 
 /// Context for the dhall compiler. Stores various global maps.
@@ -40,33 +30,46 @@ impl Ctxt<'_> {
         f(cx)
     }
 }
-impl<'cx> Ctxt<'cx> {
-    /// 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<'cx> {
-        let stored = StoredImport {
-            cx: self,
-            base_location,
-            import,
-            span,
-            result: OnceCell::new(),
-        };
-        let id = self.0.imports.len();
-        self.0.imports.push(Box::new(stored));
-        ImportId(id, PhantomData)
+impl<'cx> Deref for Ctxt<'cx> {
+    type Target = &'cx CtxtS<'cx>;
+    fn deref(&self) -> &&'cx CtxtS<'cx> {
+        &self.0
     }
-    /// Store the result of fetching an import.
-    pub fn push_import_result(self, res: Typed<'cx>) -> ImportResultId<'cx> {
-        let id = self.0.import_results.len();
-        self.0.import_results.push(Box::new(res));
-        ImportResultId(id, PhantomData)
+}
+impl<'a, 'cx, T> Index<&'a T> for CtxtS<'cx>
+where
+    Self: Index<T>,
+    T: Copy,
+{
+    type Output = <Self as Index<T>>::Output;
+    fn index(&self, id: &'a T) -> &Self::Output {
+        &self[*id]
+    }
+}
+
+/// Empty impl, because `FrozenVec` does not implement `Debug` and I can't be bothered to do it
+/// myself.
+impl<'cx> std::fmt::Debug for Ctxt<'cx> {
+    fn fmt(&self, _: &mut std::fmt::Formatter) -> std::fmt::Result {
+        Ok(())
     }
 }
 
+/////////////////////////////////////////////////////////////////////////////////////////////////////
+// Imports
+
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
+pub struct ImportId<'cx>(usize, PhantomData<&'cx ()>);
+
+/// What's stored for each `ImportId`. Allows getting and setting a result for this import.
+pub struct StoredImport<'cx> {
+    cx: Ctxt<'cx>,
+    pub base_location: ImportLocation,
+    pub import: Import,
+    pub span: Span,
+    result: OnceCell<ImportResultId<'cx>>,
+}
+
 impl<'cx> StoredImport<'cx> {
     /// Get the id of the result of fetching this import. Returns `None` if the result has not yet
     /// been fetched.
@@ -79,58 +82,130 @@ impl<'cx> StoredImport<'cx> {
     }
     /// Get the result of fetching this import. Returns `None` if the result has not yet been
     /// fetched.
-    pub fn get_result(&self) -> Option<&'cx Typed<'cx>> {
+    pub fn get_result(&self) -> Option<&'cx StoredImportResult<'cx>> {
         let res = self.get_resultid()?;
         Some(&self.cx[res])
     }
     /// Get the result of fetching this import. Panicx if the result has not yet been
     /// fetched.
-    pub fn unwrap_result(&self) -> &'cx Typed<'cx> {
+    pub fn unwrap_result(&self) -> &'cx StoredImportResult<'cx> {
         self.get_result()
             .expect("imports should all have been resolved at this stage")
     }
     /// Store the result of fetching this import.
-    pub fn set_result(&self, res: Typed<'cx>) -> ImportResultId<'cx> {
+    pub fn set_result(
+        &self,
+        res: StoredImportResult<'cx>,
+    ) -> ImportResultId<'cx> {
         let res = self.cx.push_import_result(res);
         self.set_resultid(res);
         res
     }
 }
-
-impl<'cx> Deref for Ctxt<'cx> {
-    type Target = &'cx CtxtS<'cx>;
-    fn deref(&self) -> &&'cx CtxtS<'cx> {
-        &self.0
+impl<'cx> Ctxt<'cx> {
+    /// 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<'cx> {
+        let stored = StoredImport {
+            cx: self,
+            base_location,
+            import,
+            span,
+            result: OnceCell::new(),
+        };
+        let id = self.0.imports.len();
+        self.0.imports.push(Box::new(stored));
+        ImportId(id, PhantomData)
     }
 }
-
 impl<'cx> Index<ImportId<'cx>> for CtxtS<'cx> {
     type Output = StoredImport<'cx>;
     fn index(&self, id: ImportId<'cx>) -> &StoredImport<'cx> {
         &self.imports[id.0]
     }
 }
-impl<'cx> Index<ImportResultId<'cx>> for CtxtS<'cx> {
-    type Output = Typed<'cx>;
-    fn index(&self, id: ImportResultId<'cx>) -> &Typed<'cx> {
-        &self.import_results[id.0]
+
+/////////////////////////////////////////////////////////////////////////////////////////////////////
+// Import alternatives
+
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
+pub struct ImportAlternativeId<'cx>(usize, PhantomData<&'cx ()>);
+
+/// What's stored for each `ImportAlternativeId`.
+pub struct StoredImportAlternative<'cx> {
+    pub left_imports: Box<[ImportNode<'cx>]>,
+    pub right_imports: Box<[ImportNode<'cx>]>,
+    /// `true` for left, `false` for right.
+    selected: OnceCell<bool>,
+}
+
+impl<'cx> StoredImportAlternative<'cx> {
+    /// Get which alternative got selected. `true` for left, `false` for right.
+    pub fn get_selected(&self) -> Option<bool> {
+        self.selected.get().copied()
+    }
+    /// Get which alternative got selected. `true` for left, `false` for right.
+    pub fn unwrap_selected(&self) -> bool {
+        self.get_selected()
+            .expect("imports should all have been resolved at this stage")
+    }
+    /// Set which alternative got selected. `true` for left, `false` for right.
+    pub fn set_selected(&self, selected: bool) {
+        let _ = self.selected.set(selected);
     }
 }
-impl<'a, 'cx, T> Index<&'a T> for CtxtS<'cx>
-where
-    Self: Index<T>,
-    T: Copy,
-{
-    type Output = <Self as Index<T>>::Output;
-    fn index(&self, id: &'a T) -> &Self::Output {
-        &self[*id]
+impl<'cx> Ctxt<'cx> {
+    pub fn push_import_alternative(
+        self,
+        left_imports: Box<[ImportNode<'cx>]>,
+        right_imports: Box<[ImportNode<'cx>]>,
+    ) -> ImportAlternativeId<'cx> {
+        let stored = StoredImportAlternative {
+            left_imports,
+            right_imports,
+            selected: OnceCell::new(),
+        };
+        let id = self.0.import_alternatives.len();
+        self.0.import_alternatives.push(Box::new(stored));
+        ImportAlternativeId(id, PhantomData)
+    }
+}
+impl<'cx> Index<ImportAlternativeId<'cx>> for CtxtS<'cx> {
+    type Output = StoredImportAlternative<'cx>;
+    fn index(
+        &self,
+        id: ImportAlternativeId<'cx>,
+    ) -> &StoredImportAlternative<'cx> {
+        &self.import_alternatives[id.0]
     }
 }
 
-/// Empty impl, because `FrozenVec` does not implement `Debug` and I can't be bothered to do it
-/// myself.
-impl<'cx> std::fmt::Debug for Ctxt<'cx> {
-    fn fmt(&self, _: &mut std::fmt::Formatter) -> std::fmt::Result {
-        Ok(())
+/////////////////////////////////////////////////////////////////////////////////////////////////////
+// Import results
+
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
+pub struct ImportResultId<'cx>(usize, PhantomData<&'cx ()>);
+
+type StoredImportResult<'cx> = Typed<'cx>;
+
+impl<'cx> Ctxt<'cx> {
+    /// Store the result of fetching an import.
+    pub fn push_import_result(
+        self,
+        res: StoredImportResult<'cx>,
+    ) -> ImportResultId<'cx> {
+        let id = self.0.import_results.len();
+        self.0.import_results.push(Box::new(res));
+        ImportResultId(id, PhantomData)
+    }
+}
+impl<'cx> Index<ImportResultId<'cx>> for CtxtS<'cx> {
+    type Output = StoredImportResult<'cx>;
+    fn index(&self, id: ImportResultId<'cx>) -> &StoredImportResult<'cx> {
+        &self.import_results[id.0]
     }
 }
diff --git a/dhall/src/lib.rs b/dhall/src/lib.rs
index c27b633..03d3931 100644
--- a/dhall/src/lib.rs
+++ b/dhall/src/lib.rs
@@ -26,7 +26,7 @@ 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};
+pub use ctxt::*;
 
 #[derive(Debug, Clone)]
 pub struct Parsed(Expr, ImportLocation);
diff --git a/dhall/src/semantics/nze/normalize.rs b/dhall/src/semantics/nze/normalize.rs
index 0a09a80..59710d1 100644
--- a/dhall/src/semantics/nze/normalize.rs
+++ b/dhall/src/semantics/nze/normalize.rs
@@ -160,6 +160,14 @@ pub fn normalize_hir<'cx>(env: &NzEnv<'cx>, hir: &Hir<'cx>) -> NirKind<'cx> {
             let typed = env.cx()[import].unwrap_result();
             normalize_hir(env, &typed.hir)
         }
+        HirKind::ImportAlternative(alt, left, right) => {
+            let hir = if env.cx()[alt].unwrap_selected() {
+                left
+            } else {
+                right
+            };
+            normalize_hir(env, hir)
+        }
         HirKind::Expr(ExprKind::Lam(binder, annot, body)) => {
             let annot = annot.eval(env);
             NirKind::LamClosure {
diff --git a/dhall/src/semantics/resolve/hir.rs b/dhall/src/semantics/resolve/hir.rs
index 05a8550..8baad10 100644
--- a/dhall/src/semantics/resolve/hir.rs
+++ b/dhall/src/semantics/resolve/hir.rs
@@ -1,7 +1,7 @@
 use crate::error::TypeError;
 use crate::semantics::{type_with, typecheck, NameEnv, Nir, NzEnv, Tir, TyEnv};
 use crate::syntax::{Expr, ExprKind, Span, V};
-use crate::{Ctxt, ImportId, ToExprOptions};
+use crate::{Ctxt, ImportAlternativeId, ImportId, ToExprOptions};
 
 /// Stores an alpha-normalized variable.
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
@@ -15,8 +15,10 @@ pub enum HirKind<'cx> {
     Var(AlphaVar),
     /// A variable that couldn't be resolved. Detected during resolution, but causes an error during typeck.
     MissingVar(V),
-    /// An import. It must have been resolved by the time we get to typechecking/normalization.
+    /// An import. It must have been resolved after resolution.
     Import(ImportId<'cx>),
+    /// An import alternative. It must have been decided after resolution.
+    ImportAlternative(ImportAlternativeId<'cx>, Hir<'cx>, Hir<'cx>),
     // Forbidden ExprKind variants: Var, Import, Completion
     Expr(ExprKind<Hir<'cx>>),
 }
@@ -111,6 +113,14 @@ fn hir_to_expr<'cx>(
             let typed = cx[import].unwrap_result();
             return hir_to_expr(cx, &typed.hir, opts, &mut NameEnv::new());
         }
+        HirKind::ImportAlternative(alt, left, right) => {
+            let hir = if cx[alt].unwrap_selected() {
+                left
+            } else {
+                right
+            };
+            return hir_to_expr(cx, hir, opts, env);
+        }
         HirKind::Expr(e) => {
             let e = e.map_ref_maybe_binder(|l, hir| {
                 if let Some(l) = l {
diff --git a/dhall/src/semantics/resolve/resolve.rs b/dhall/src/semantics/resolve/resolve.rs
index c4cd518..116e1a5 100644
--- a/dhall/src/semantics/resolve/resolve.rs
+++ b/dhall/src/semantics/resolve/resolve.rs
@@ -15,7 +15,10 @@ use crate::syntax::{
     Expr, ExprKind, FilePath, FilePrefix, Hash, ImportMode, ImportTarget, Span,
     UnspannedExpr, URL,
 };
-use crate::{Ctxt, ImportId, ImportResultId, Parsed, Resolved, Typed};
+use crate::{
+    Ctxt, ImportAlternativeId, ImportId, ImportResultId, Parsed, Resolved,
+    Typed,
+};
 
 // TODO: evaluate import headers
 pub type Import = syntax::Import<()>;
@@ -252,7 +255,7 @@ impl ImportLocation {
         let typed = match self.mode {
             ImportMode::Code => {
                 let parsed = self.kind.fetch_dhall()?;
-                let typed = resolve_with_env(env, parsed)?.typecheck(cx)?;
+                let typed = parsed.resolve_with_env(env)?.typecheck(cx)?;
                 Typed {
                     // TODO: manage to keep the Nir around. Will need fixing variables.
                     hir: typed.normalize(cx).to_hir(),
@@ -366,56 +369,6 @@ fn desugar(expr: &Expr) -> Cow<'_, Expr> {
     }
 }
 
-/// Traverse the expression, handling import alternatives and passing
-/// found imports to the provided function. Also resolving names.
-fn traverse_resolve_expr<'cx>(
-    name_env: &mut NameEnv,
-    expr: &Expr,
-    f: &mut impl FnMut(Import, Span) -> Result<ImportId<'cx>, Error>,
-) -> Result<Hir<'cx>, Error> {
-    let expr = desugar(expr);
-    Ok(match expr.kind() {
-        ExprKind::Var(var) => match name_env.unlabel_var(&var) {
-            Some(v) => Hir::new(HirKind::Var(v), expr.span()),
-            None => Hir::new(HirKind::MissingVar(var.clone()), expr.span()),
-        },
-        ExprKind::Op(OpKind::BinOp(BinOp::ImportAlt, l, r)) => {
-            match traverse_resolve_expr(name_env, l, f) {
-                Ok(l) => l,
-                Err(_) => {
-                    match traverse_resolve_expr(name_env, r, f) {
-                        Ok(r) => r,
-                        // TODO: keep track of the other error too
-                        Err(e) => return Err(e),
-                    }
-                }
-            }
-        }
-        kind => {
-            let kind = kind.traverse_ref_maybe_binder(|l, e| {
-                if let Some(l) = l {
-                    name_env.insert_mut(l);
-                }
-                let hir = traverse_resolve_expr(name_env, e, f)?;
-                if l.is_some() {
-                    name_env.remove_mut();
-                }
-                Ok::<_, Error>(hir)
-            })?;
-            let kind = match kind {
-                ExprKind::Import(import) => {
-                    // TODO: evaluate import headers
-                    let import = import.traverse_ref(|_| Ok::<_, Error>(()))?;
-                    let import_id = f(import, expr.span())?;
-                    HirKind::Import(import_id)
-                }
-                kind => HirKind::Expr(kind),
-            };
-            Hir::new(kind, expr.span())
-        }
-    })
-}
-
 /// Fetch the import and store the result in the global context.
 fn fetch_import<'cx>(
     env: &mut ImportEnv<'cx>,
@@ -469,22 +422,123 @@ fn fetch_import<'cx>(
     Ok(res_id)
 }
 
+/// Part of a tree of imports.
+#[derive(Debug, Clone, Copy)]
+pub enum ImportNode<'cx> {
+    Import(ImportId<'cx>),
+    Alternative(ImportAlternativeId<'cx>),
+}
+
+/// Traverse the expression and replace each import and import alternative by an id into the global
+/// context. The ids are also accumulated into `nodes` so that we can resolve them afterwards.
+fn traverse_accumulate<'cx>(
+    env: &mut ImportEnv<'cx>,
+    name_env: &mut NameEnv,
+    nodes: &mut Vec<ImportNode<'cx>>,
+    base_location: &ImportLocation,
+    expr: &Expr,
+) -> Hir<'cx> {
+    let cx = env.cx();
+    let expr = desugar(expr);
+    let kind = match expr.kind() {
+        ExprKind::Var(var) => match name_env.unlabel_var(&var) {
+            Some(v) => HirKind::Var(v),
+            None => HirKind::MissingVar(var.clone()),
+        },
+        ExprKind::Op(OpKind::BinOp(BinOp::ImportAlt, l, r)) => {
+            let mut imports_l = Vec::new();
+            let l = traverse_accumulate(
+                env,
+                name_env,
+                &mut imports_l,
+                base_location,
+                l,
+            );
+            let mut imports_r = Vec::new();
+            let r = traverse_accumulate(
+                env,
+                name_env,
+                &mut imports_r,
+                base_location,
+                r,
+            );
+            let alt =
+                cx.push_import_alternative(imports_l.into(), imports_r.into());
+            nodes.push(ImportNode::Alternative(alt));
+            HirKind::ImportAlternative(alt, l, r)
+        }
+        kind => {
+            let kind = kind.map_ref_maybe_binder(|l, e| {
+                if let Some(l) = l {
+                    name_env.insert_mut(l);
+                }
+                let hir =
+                    traverse_accumulate(env, name_env, nodes, base_location, e);
+                if l.is_some() {
+                    name_env.remove_mut();
+                }
+                hir
+            });
+            match kind {
+                ExprKind::Import(import) => {
+                    // TODO: evaluate import headers
+                    let import = import.map_ref(|_| ());
+                    let import_id = cx.push_import(
+                        base_location.clone(),
+                        import,
+                        expr.span(),
+                    );
+                    nodes.push(ImportNode::Import(import_id));
+                    HirKind::Import(import_id)
+                }
+                kind => HirKind::Expr(kind),
+            }
+        }
+    };
+    Hir::new(kind, expr.span())
+}
+
+/// Take a list of nodes and recursively resolve them.
+fn resolve_nodes<'cx>(
+    env: &mut ImportEnv<'cx>,
+    nodes: &[ImportNode<'cx>],
+) -> Result<(), Error> {
+    for &node in nodes {
+        match node {
+            ImportNode::Import(import) => {
+                let res_id = fetch_import(env, import)?;
+                env.cx()[import].set_resultid(res_id);
+            }
+            ImportNode::Alternative(alt) => {
+                let alt = &env.cx()[alt];
+                if resolve_nodes(env, &alt.left_imports).is_ok() {
+                    alt.set_selected(true);
+                } else {
+                    resolve_nodes(env, &alt.right_imports)?;
+                    alt.set_selected(false);
+                }
+            }
+        }
+    }
+    Ok(())
+}
+
 fn resolve_with_env<'cx>(
     env: &mut ImportEnv<'cx>,
     parsed: Parsed,
 ) -> Result<Resolved<'cx>, Error> {
     let Parsed(expr, base_location) = parsed;
-    let resolved = traverse_resolve_expr(
+    let mut nodes = Vec::new();
+    // First we collect all imports.
+    let resolved = traverse_accumulate(
+        env,
         &mut NameEnv::new(),
+        &mut nodes,
+        &base_location,
         &expr,
-        &mut |import, span| {
-            let import_id =
-                env.cx().push_import(base_location.clone(), import, span);
-            let res_id = fetch_import(env, import_id)?;
-            env.cx()[import_id].set_resultid(res_id);
-            Ok(import_id)
-        },
-    )?;
+    );
+    // Then we resolve them and choose sides for the alternatives.
+    resolve_nodes(env, &nodes)?;
     Ok(Resolved(resolved))
 }
 
@@ -494,10 +548,10 @@ pub fn resolve<'cx>(
     cx: Ctxt<'cx>,
     parsed: Parsed,
 ) -> Result<Resolved<'cx>, Error> {
-    resolve_with_env(&mut ImportEnv::new(cx), parsed)
+    parsed.resolve_with_env(&mut ImportEnv::new(cx))
 }
 
-/// Resolves names and errors if we find any imports.
+/// Resolves names, and errors if we find any imports.
 pub fn skip_resolve<'cx>(
     cx: Ctxt<'cx>,
     parsed: Parsed,
@@ -506,6 +560,15 @@ pub fn skip_resolve<'cx>(
     Ok(resolve(cx, parsed)?)
 }
 
+impl Parsed {
+    fn resolve_with_env<'cx>(
+        self,
+        env: &mut ImportEnv<'cx>,
+    ) -> Result<Resolved<'cx>, Error> {
+        resolve_with_env(env, self)
+    }
+}
+
 pub trait Canonicalize {
     fn canonicalize(&self) -> Self;
 }
diff --git a/dhall/src/semantics/tck/typecheck.rs b/dhall/src/semantics/tck/typecheck.rs
index e14bcc6..23c2bd2 100644
--- a/dhall/src/semantics/tck/typecheck.rs
+++ b/dhall/src/semantics/tck/typecheck.rs
@@ -193,6 +193,14 @@ pub fn type_with<'cx, 'hir>(
             let typed = env.cx()[import].unwrap_result();
             Tir::from_hir(hir, typed.ty.clone())
         }
+        HirKind::ImportAlternative(alt, left, right) => {
+            let hir = if env.cx()[alt].unwrap_selected() {
+                left
+            } else {
+                right
+            };
+            return type_with(env, hir, annot);
+        }
         HirKind::Expr(ExprKind::Var(_)) => {
             unreachable!("Hir should contain no unresolved variables")
         }