computing cases

1 files changed, 283 insertions, 0 deletions

diff --git a/src/checker_set.rs b/src/checker_set.rs
new file mode 100644
index 0000000..a14ac9a
--- /dev/null
+++ b/src/checker_set.rs
@@ -0,0 +1,283 @@
+use crate::ast::*;
+use crate::checker_state::*;
+
+use std::iter::zip;
+use tracing::instrument;
+
+impl CheckerState {
+    #[instrument(skip(self), level = "debug", fields(%set))]
+    pub fn check_set(&self, set: &Set) -> Result<Set, CheckerError> {
+        match set {
+            Set::BuiltIn(_) => Ok(set.clone()),
+            Set::Record(fields) => {
+                let fields = fields
+                    .iter()
+                    .map(|RecordField { name, set }| {
+                        let set = self.check_set(set)?;
+                        Ok(RecordField {
+                            name: name.clone(),
+                            set,
+                        })
+                    })
+                    .collect::<Result<Vec<_>, _>>()?;
+                Ok(Set::Record(fields))
+            }
+            Set::Variant(fields) => {
+                let fields = fields
+                    .iter()
+                    .map(|VariantField { name, set }| {
+                        let set = self.check_set(set)?;
+                        Ok(VariantField {
+                            name: name.clone(),
+                            set,
+                        })
+                    })
+                    .collect::<Result<Vec<_>, _>>()?;
+                Ok(Set::Variant(fields))
+            }
+            Set::ClaimedSet(_) => Err(CheckerError::Unimplemented("instances as sets".to_string())),
+            Set::Var(v) => {
+                let deref = self.lookup_set(v)?;
+                Ok(deref.clone())
+            }
+        }
+    }
+
+    fn _check_literal_set_helper(&self, claimed: &Set, should_be: Set) -> Result<(), CheckerError> {
+        if !self.set_equal(claimed, &should_be) {
+            Err(CheckerError::WrongSetForElement(claimed.clone(), should_be))
+        } else {
+            Ok(())
+        }
+    }
+
+    #[instrument(skip(self), level = "debug", fields(%element, %set))]
+    pub fn check_element(&self, element: &Element, set: &Set) -> Result<Element, CheckerError> {
+        match element {
+            Element::Literal(lit) => {
+                // we may infer the type from the element
+                match lit {
+                    Literal::Int(_) => {
+                        self._check_literal_set_helper(set, Set::BuiltIn(BuiltIn::Int))?;
+                    }
+                    Literal::Nat(_) => {
+                        self._check_literal_set_helper(set, Set::BuiltIn(BuiltIn::Nat))?;
+                    }
+                    Literal::Str(_) => {
+                        self._check_literal_set_helper(set, Set::BuiltIn(BuiltIn::Str))?;
+                    }
+                    Literal::Bool(_) => {
+                        self._check_literal_set_helper(set, Set::BuiltIn(BuiltIn::Bool))?;
+                    }
+                    Literal::Float(_) => {
+                        self._check_literal_set_helper(set, Set::BuiltIn(BuiltIn::Float))?;
+                    }
+                }
+                Ok(element.clone())
+            }
+            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) {
+                    return Err(CheckerError::WrongSetForElement(
+                        set.clone(),
+                        lookup.set.clone(),
+                    ));
+                }
+                Ok(lookup.element.clone())
+            }
+            Element::Record(assignations) => {
+                let rej = |reason| CheckerError::ElementDoesNotBelong {
+                    element: element.clone(),
+                    claimed: set.clone(),
+                    reason,
+                };
+
+                // make sure we are filling a record
+                let fields = if let Set::Record(fields) = set {
+                    Ok(fields)
+                } else {
+                    Err(rej("element is a record instance".to_string()))
+                }?;
+
+                let (set_fnames, set_fsets): (Vec<String>, Vec<Set>) = fields
+                    .iter()
+                    .map(|RecordField { name, set }| (name.clone(), set.clone()))
+                    .unzip();
+                let mut set_fnames_sorted = set_fnames.clone();
+                set_fnames_sorted.sort();
+
+                let (element_fnames, element_felements): (Vec<String>, Vec<&Element>) =
+                    assignations
+                        .iter()
+                        .map(|ElemAssign { name, element }| (name.clone(), element))
+                        .unzip();
+
+                let mut element_fnames_sorted = element_fnames.clone();
+                element_fnames_sorted.sort();
+
+                if set_fnames_sorted != element_fnames_sorted {
+                    return Err(rej(format!(
+                        "expected [{}] but found [{}]",
+                        set_fnames.join(", "),
+                        element_fnames.join(", "),
+                    )));
+                }
+
+                // 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))
+                    .collect::<Result<Vec<_>, _>>()?;
+                // rebuild
+                let assignations = zip(element_fnames, sub_els)
+                    .map(|(name, element)| ElemAssign { name, element })
+                    .collect();
+                // resign?
+                Ok(Element::Record(assignations))
+            }
+            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)?;
+
+                // enforce the correct typing of the claimed result
+                if !self.set_equal(set, field_set) {
+                    return Err(CheckerError::WrongSetForElement(
+                        set.clone(),
+                        field_set.clone(),
+                    ));
+                }
+
+                // 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();
+
+                Ok(sub_element)
+            }
+            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)?;
+
+                // enforce the correct typing of the claimed result
+                if !self.set_equal(set, owner_set) {
+                    return Err(CheckerError::WrongSetForElement(
+                        set.clone(),
+                        owner_set.clone(),
+                    ));
+                }
+
+                // enforce the correct typing of the element
+                let element = self.check_element(inner, field_set)?;
+
+                Ok(Element::Inject {
+                    element: Box::new(element),
+                    field: field.clone(),
+                })
+            }
+            Element::Case { arms, scrutinee } => {
+                // TODO: do we allow mapping out of bottom?
+                if arms.is_empty() {
+                    return Err(CheckerError::Unimplemented(
+                        "mapping out of bottom types".to_string(),
+                    ));
+                }
+
+                // 1. Syntactic checks
+                // -------------------
+                // arms agree on the set to which the scrutinee should belong
+                let arm_owners = arms
+                    .iter()
+                    .map(|ca| self.lookup_variant_field(&ca.tag).map(|sf| &sf.owner_set))
+                    .collect::<Result<Vec<_>, _>>()?;
+                let owner = arm_owners[0]; // safe because of the above decision about bottom
+                if !arm_owners.iter().all(|o| self.set_equal(owner, o)) {
+                    return Err(CheckerError::IncosistentCaseScrutineeSet(element.clone()));
+                }
+
+                // all cases are handled
+                let Set::Variant(fields) = owner else {
+                    panic!(
+                        "invariant violation: looking up the owner of a variant field resulted in a non-variant set",
+                    )
+                };
+                let mut required_field_names_sorted: Vec<String> =
+                    fields.iter().map(|vf| vf.name.clone()).collect();
+                required_field_names_sorted.sort();
+                let mut covered_field_names_sorted: Vec<String> =
+                    arms.iter().map(|ca| ca.tag.clone()).collect();
+                covered_field_names_sorted.sort();
+                if required_field_names_sorted != covered_field_names_sorted {
+                    return Err(CheckerError::IncompleteCaseAnalysis {
+                        found: covered_field_names_sorted,
+                        required: required_field_names_sorted,
+                    });
+                }
+
+                // 2. semantic checks
+                // ------------------
+
+                // 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)?;
+
+                // 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"
+                    );
+                };
+
+                // 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 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)?;
+
+                Ok(computed)
+            }
+        }
+    }
+}