add distinction between hypothetical and concrete elements to the set checker, it now enforces that all arms in case are well typed!

4 files changed, 153 insertions, 73 deletions

diff --git a/src/checker.rs b/src/checker.rs
index 3ebcb4d..1760c0e 100644
--- a/src/checker.rs
+++ b/src/checker.rs
@@ -1,5 +1,5 @@
 use crate::ast::*;
-use crate::checker_state::{CheckerError, CheckerState};
+use crate::checker_state::{CheckerError, CheckerState, ElementValue};
 
 use tracing::{debug, instrument};
 
@@ -25,7 +25,12 @@ impl CheckerState {
 
                 Decl::Element { name, element, set } => {
                     let set = self.check_set(set)?;
-                    let element = self.check_element(element, &set)?;
+                    let element = self.check_element(element.clone().into(), &set)?;
+                    if matches!(element, ElementValue::Hypothetical(_)) {
+                        panic!(
+                            "invariant violation: from concrete values at the top level we returned a hypothetical value"
+                        );
+                    }
                     self.add_element(name.clone(), element, set)
                 }
                 Decl::Signature { .. } => {
diff --git a/src/checker_set.rs b/src/checker_set.rs
index 24327cf..5410dfe 100644
--- a/src/checker_set.rs
+++ b/src/checker_set.rs
@@ -55,9 +55,27 @@ impl CheckerState {
     }
 
     #[instrument(skip(self), level = "debug", fields(%element, %set))]
-    pub fn check_element(&self, element: &Element, set: &Set) -> Result<Element, CheckerError> {
+    pub fn check_element(
+        &self,
+        element: ElementValue,
+        set: &Set,
+    ) -> Result<ElementValue, CheckerError> {
         match element {
-            Element::Literal(lit) => {
+            // This arm ensures that if we ever in the position of obtaining a
+            // hypothetical from a call to check_element, in the context of
+            // check_element, we can safely ignore its payload. I'll point this
+            // out later as (*)
+            ElementValue::Hypothetical(h_set) => {
+                if !self.set_equal(set, &h_set) {
+                    Err(CheckerError::WrongSetForElement {
+                        claimed: set.clone(),
+                        real: h_set,
+                    })
+                } else {
+                    Ok(ElementValue::Hypothetical(h_set))
+                }
+            }
+            ElementValue::Concrete(ref element @ Element::Literal(ref lit)) => {
                 // we may infer the type from the element
                 match lit {
                     Literal::Int(_) => {
@@ -76,10 +94,10 @@ impl CheckerState {
                         self._check_literal_set_helper(set, Set::BuiltIn(BuiltIn::Float))?;
                     }
                 }
-                Ok(element.clone())
+                Ok(element.clone().into())
             }
-            Element::Var(v) => {
-                let lookup = self.lookup_element(v)?;
+            ElementValue::Concrete(Element::Var(v)) => {
+                let lookup = self.lookup_element(&v)?;
                 // we have previously done the work to discover the type of
                 // this element, so what we're claiming now must match!
                 if !self.set_equal(set, &lookup.set) {
@@ -88,11 +106,11 @@ impl CheckerState {
                         real: lookup.set.clone(),
                     });
                 }
-                Ok(lookup.element.clone())
+                Ok(lookup.value.clone())
             }
-            Element::Record(assignations) => {
+            ElementValue::Concrete(ref concrete @ Element::Record(ref assignations)) => {
                 let rej = |reason| CheckerError::ElementDoesNotBelong {
-                    element: element.clone(),
+                    element: concrete.clone(),
                     claimed: set.clone(),
                     reason,
                 };
@@ -130,25 +148,32 @@ impl CheckerState {
 
                 // recurse, sets have already been completely expanded
                 let sub_els = zip(element_felements, set_fsets)
-                    .map(|(e_f, e_s)| self.check_element(e_f, &e_s))
+                    .map(|(e_f, e_s)| self.check_element(e_f.clone().into(), &e_s))
                     .collect::<Result<Vec<_>, _>>()?;
-                // rebuild
+                // rebuild, hypotheticals are contagious
                 let assignations = zip(element_fnames, sub_els)
-                    .map(|(name, element)| ElemAssign { name, element })
+                    .map(|(name, element)| match element {
+                        ElementValue::Concrete(element) => Some(ElemAssign { name, element }),
+                        ElementValue::Hypothetical(_) => None,
+                    })
                     .collect();
                 // resign?
-                Ok(Element::Record(assignations))
+                Ok(if let Some(assignations) = assignations {
+                    Element::Record(assignations).into()
+                } else {
+                    ElementValue::Hypothetical(set.clone())
+                })
             }
-            Element::Project {
+            ElementValue::Concrete(Element::Project {
                 element: inner,
                 field,
-            } => {
+            }) => {
                 // globally unique projections mean we know what the sets going
                 // in and out must be
                 let SetField {
                     field_set,
                     owner_set,
-                } = self.lookup_record_field(field)?;
+                } = self.lookup_record_field(&field)?;
 
                 // enforce the correct typing of the claimed result
                 if !self.set_equal(set, field_set) {
@@ -159,32 +184,41 @@ impl CheckerState {
                 }
 
                 // enforce the correct typing of the element
-                let inner = self.check_element(inner, owner_set)?;
-
-                // Unfortunately we still have to do something nasty here to obtain the data
-                let Element::Record(assignations) = inner else {
-                    panic!("invariant violation: check_element returned non-record for record set");
-                };
-                let sub_element = assignations
-                    .into_iter()
-                    .find(|a| a.name == *field)
-                    .expect("invariant violation: record missing field that was type-checked")
-                    .element
-                    .clone();
+                let inner = self.check_element((*inner).into(), owner_set)?;
+                match inner {
+                    ElementValue::Concrete(inner) => {
+                        // Unfortunately we still have to do something nasty here to obtain the data
+                        let Element::Record(assignations) = inner else {
+                            panic!(
+                                "invariant violation: check_element returned non-record for record set"
+                            );
+                        };
+                        let sub_element = assignations
+                            .into_iter()
+                            .find(|a| a.name == field)
+                            .expect(
+                                "invariant violation: record missing field that was type-checked",
+                            )
+                            .element
+                            .clone();
 
-                Ok(sub_element)
+                        Ok(sub_element.into())
+                    }
+                    // correct by (*)
+                    ElementValue::Hypothetical(_) => Ok(ElementValue::Hypothetical(set.clone())),
+                }
             }
-            Element::Inject {
+            ElementValue::Concrete(Element::Inject {
                 element: inner,
                 field,
-            } => {
+            }) => {
                 // globally unique injections mean that we know what the sets
                 // going in and out must be, but compared to projections their
                 // roles are here interchanged
                 let SetField {
                     field_set,
                     owner_set,
-                } = self.lookup_variant_field(field)?;
+                } = self.lookup_variant_field(&field)?;
 
                 // enforce the correct typing of the claimed result
                 if !self.set_equal(set, owner_set) {
@@ -195,14 +229,24 @@ impl CheckerState {
                 }
 
                 // enforce the correct typing of the element
-                let element = self.check_element(inner, field_set)?;
+                let element = self.check_element((*inner).into(), field_set)?;
 
-                Ok(Element::Inject {
-                    element: Box::new(element),
-                    field: field.clone(),
-                })
+                match element {
+                    ElementValue::Concrete(element) => Ok(Element::Inject {
+                        element: Box::new(element),
+                        field: field.clone(),
+                    }
+                    .into()),
+                    // correct by (*)
+                    ElementValue::Hypothetical(_) => Ok(ElementValue::Hypothetical(set.clone())),
+                }
             }
-            Element::Case { arms, scrutinee } => {
+            ElementValue::Concrete(
+                ref element @ Element::Case {
+                    ref arms,
+                    ref scrutinee,
+                },
+            ) => {
                 // TODO: do we allow mapping out of bottom?
                 if arms.is_empty() {
                     return Err(CheckerError::Unimplemented(
@@ -247,39 +291,52 @@ impl CheckerState {
                 // scrutinee must be of the same set that all the arms are
                 // implying, in particular this implies that the following holds
                 // `inner : self.lookup_variant_field(field).field_set`
-                let scrutinee = self.check_element(scrutinee, owner)?;
+                let scrutinee = self.check_element((*scrutinee.clone()).into(), owner)?;
 
-                // which variant are we?
-                let Element::Inject {
-                    field,
-                    element: inner,
-                } = scrutinee
-                else {
-                    panic!(
-                        "invariant violation: we believe element is of a variant set but it's not an injection"
-                    );
+                // which variant are we, if any
+                let matching: Option<(String, Element)> = match scrutinee {
+                    ElementValue::Hypothetical(_) => None,
+                    ElementValue::Concrete(Element::Inject {
+                        field,
+                        element: inner,
+                    }) => Some((field, *inner)),
+                    _ => {
+                        panic!(
+                            "invariant violation: we believe element is of a variant set but it's not an injection"
+                        );
+                    }
                 };
 
-                // TODO: we would like to check that each arm is correct, but
-                // there's no easy way to do this? we can insert hypothetical
-                // elements of the correct type into the checkerstate, but if
-                // the body exacts non-trivial computation we won't be to pass
-                // further checks. In the future would could build first class
-                // support for hypothetical elements and do proper bi-di
-                // checking, but for now we only check the branch that matters.
+                let mut computed_output = None;
+                for arm in arms {
+                    let SetField { field_set, .. } = self.lookup_variant_field(&arm.tag)?;
 
-                let CaseArm { tag, bound, body } = arms.iter().find(|ca| ca.tag == field).expect("invariant violation: we know that all cases are covered and that the element is of the valid type");
-                let SetField { field_set, .. } = self.lookup_variant_field(tag)?;
-                // TODO: if we were worried about overhead we'd have a separate
-                // locals stack, though truly if we were worried about overhead
-                // we'd not have NNN instances of clone elsewhere in the
-                // codebase and we wouldn't be eagerly evaluating all
-                // expressions fully.
-                let mut new_context = self.clone();
-                new_context.add_element(bound.clone(), *inner.clone(), field_set.clone())?;
-                let computed = new_context.check_element(body, set)?;
+                    // TODO: if we were worried about overhead we'd have a separate
+                    // locals stack, though truly if we were worried about overhead
+                    // we'd not have NNN instances of clone elsewhere in the
+                    // codebase and we wouldn't be eagerly evaluating all
+                    // expressions fully.
 
-                Ok(computed)
+                    let mut new_context = self.clone();
+                    let binding_name = arm.bound.clone();
+                    let binding_set = field_set.clone();
+
+                    if let Some((tag, inner)) = &matching
+                        && *tag == arm.tag
+                    {
+                        new_context.add_element(binding_name, inner.clone().into(), binding_set)?;
+                        let output = new_context.check_element(arm.body.clone().into(), set)?;
+                        computed_output = Some(output);
+                    } else {
+                        new_context.add_element(
+                            binding_name,
+                            ElementValue::Hypothetical(field_set.clone()),
+                            binding_set,
+                        )?;
+                        new_context.check_element(arm.body.clone().into(), set)?;
+                    };
+                }
+                Ok(computed_output.unwrap_or(ElementValue::Hypothetical(set.clone())))
             }
         }
     }
diff --git a/src/checker_state.rs b/src/checker_state.rs
index 30f83be..7590457 100644
--- a/src/checker_state.rs
+++ b/src/checker_state.rs
@@ -41,9 +41,22 @@ pub struct SetField {
 }
 
 #[derive(Display, Clone)]
-#[display("{element} : {set}")]
+pub enum ElementValue {
+    Concrete(Element),
+    #[display("_ : {_0}")]
+    Hypothetical(Set),
+}
+
+impl From<Element> for ElementValue {
+    fn from(e: Element) -> ElementValue {
+        ElementValue::Concrete(e)
+    }
+}
+
+#[derive(Display, Clone)]
+#[display("{value} : {set}")]
 pub struct CheckedElement {
-    pub element: Element,
+    pub value: ElementValue,
     pub set: Set,
 }
 
@@ -198,12 +211,17 @@ impl CheckerState {
     pub fn add_element(
         &mut self,
         name: String,
-        element: Element,
+        element: ElementValue,
         set: Set,
     ) -> Result<(), CheckerError> {
         self.assert_unbound_element(&name)?;
-        self.wf_elements
-            .insert(name, CheckedElement { element, set });
+        self.wf_elements.insert(
+            name,
+            CheckedElement {
+                value: element,
+                set,
+            },
+        );
         Ok(())
     }
 
diff --git a/src/main.rs b/src/main.rs
index 984a438..b0bc33c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -34,7 +34,7 @@ let element the_nat : Nat = z .y .n
 
 let element injected : W = z. z
 
-let element compute : Nat = case injected of [ z. myz => myz .y .n | f. myf => 2 ]
+let element compute : Nat = case injected of [ z. myz => myz .y .n | f. myf => myf ]
 
 // let signature Graph = theory {
 //     .Node :: Set,