path: root/src/checker.rs

use crate::ast::*;
use tracing::{debug, instrument};

use derive_more::Display;
use std::collections::HashMap;
use std::fmt;
use std::iter::zip;

impl Programme {
    pub fn check(&self) -> Result<(), CheckError> {
        let mut state = CheckState::default();
        state.check(self)
    }
}

#[derive(Display)]
pub enum CheckError {
    #[display("Unbound: {_0}")]
    Unbound(String),
    #[display("Rebinding: {_0}")]
    Rebinding(String),
    #[display("The following functionality is unimplemented: {_0}")]
    Unimplemented(String),
    #[display("Element claimed to belong to {_0} but actually belongs to {_1}")]
    WrongSetForElement(Set, Set),
    #[display("Element {element} does belong to set {claimed}: {reason}")]
    ElementDoesNotBelong {
        element: Element,
        claimed: Set,
        reason: String,
    },
}

#[derive(Display)]
#[display("{value} @ {belongs_to}")]
struct SetField {
    value: Set,
    belongs_to: Set,
}

#[derive(Display)]
#[display("{element} : {set}")]
struct CheckedElement {
    element: Element,
    set: Set,
}

#[derive(Default)]
struct CheckState {
    wf_sets: HashMap<String, Set>,
    wf_elements: HashMap<String, CheckedElement>,
    wf_signatures: HashMap<String, Signature>,
    wf_instances: HashMap<String, Instance>,
    record_fields: HashMap<String, SetField>,
    variant_fields: HashMap<String, SetField>,
}

impl fmt::Display for CheckState {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fn section<K, V>(f: &mut fmt::Formatter<'_>, name: &str, map: &HashMap<K, V>) -> fmt::Result
        where
            K: fmt::Display + Ord,
            V: fmt::Display,
        {
            write!(f, " {name} = {{")?;
            let mut entries: Vec<_> = map.iter().collect();
            entries.sort_by(|a, b| a.0.cmp(b.0));
            for (k, v) in entries {
                write!(f, "{k} ~> {v}, ")?;
            }
            write!(f, "}},")
        }

        write!(f, "CheckState {{")?;
        section(f, "sets", &self.wf_sets)?;
        section(f, "elements", &self.wf_elements)?;
        section(f, "record_fields", &self.record_fields)?;
        section(f, "variant_fields", &self.variant_fields)?;
        section(f, "signatures", &self.wf_signatures)?;
        section(f, "instances", &self.wf_instances)?;
        write!(f, " }}")?;
        Ok(())
    }
}

impl CheckState {
    #[instrument(skip(self), level = "debug")]
    fn assert_unbound_set(&self, name: &String) -> Result<(), CheckError> {
        if self.wf_sets.contains_key(name) {
            Err(CheckError::Rebinding(name.clone()))
        } else {
            Ok(())
        }
    }

    #[instrument(skip(self), level = "debug")]
    fn assert_unbound_element(&self, name: &String) -> Result<(), CheckError> {
        if self.wf_elements.contains_key(name) {
            Err(CheckError::Rebinding(name.clone()))
        } else {
            Ok(())
        }
    }

    #[instrument(skip(self), level = "debug", fields(%name, %set_ref, %belongs_to))]
    fn assert_correct_owner(
        &self,
        name: &String,
        set_ref: &SetField,
        belongs_to: &Set,
    ) -> Result<(), CheckError> {
        let SetField {
            value: _,
            belongs_to: owner,
        } = set_ref;
        if !self.set_equal(owner, belongs_to) {
            Err(CheckError::Rebinding(name.clone()))
        } else {
            Ok(())
        }
    }

    #[instrument(skip(self), level = "debug", fields(%name, %set, %belongs_to))]
    fn add_record_field(
        &mut self,
        name: &String,
        set: &Set,
        belongs_to: &Set,
    ) -> Result<(), CheckError> {
        if let Some(set_ref) = self.record_fields.get(name) {
            self.assert_correct_owner(name, set_ref, belongs_to)?;
        };
        self.record_fields.insert(
            name.clone(),
            SetField {
                value: set.clone(),
                belongs_to: belongs_to.clone(),
            },
        );
        Ok(())
    }

    #[instrument(skip(self), level = "debug", fields(%name, %set, %belongs_to))]
    fn add_variant_field(
        &mut self,
        name: &String,
        set: &Set,
        belongs_to: &Set,
    ) -> Result<(), CheckError> {
        if let Some(set_ref) = self.variant_fields.get(name) {
            self.assert_correct_owner(name, set_ref, belongs_to)?;
        };
        self.variant_fields.insert(
            name.clone(),
            SetField {
                value: set.clone(),
                belongs_to: belongs_to.clone(),
            },
        );
        Ok(())
    }

    #[instrument(skip(self), level = "debug", fields(%name, %set))]
    fn add_set(&mut self, name: &String, set: Set) -> Result<(), CheckError> {
        self.assert_unbound_set(name)?;
        match &set {
            Set::Record(fields) => {
                for RecordField {
                    name: rfn,
                    set: rset,
                } in fields
                {
                    self.add_record_field(rfn, rset, &set)?;
                }
            }
            Set::Variant(fields) => {
                for VariantField {
                    name: vfn,
                    set: vset,
                } in fields
                {
                    self.add_variant_field(vfn, vset, &set)?;
                }
            }
            _ => (),
        };
        self.wf_sets.insert(name.clone(), set);
        Ok(())
    }

    fn add_element(&mut self, name: &String, element: Element, set: Set) -> Result<(), CheckError> {
        self.assert_unbound_element(name)?;
        self.wf_elements
            .insert(name.clone(), CheckedElement { element, set });
        Ok(())
    }
}

// The invariant we're maintaining is that everything is fully evaluated before
// we commit it to be stored in the state.
impl CheckState {
    // Because of our invariant we don't actually need to do anything
    // non-trivial here.
    #[instrument(skip(self), level = "debug", fields(%set_a, %set_b))]
    fn set_equal(&self, set_a: &Set, set_b: &Set) -> bool {
        set_a == set_b
    }

    #[instrument(skip(self, prog), level = "debug")]
    pub fn check(&mut self, prog: &Programme) -> Result<(), CheckError> {
        let Programme(decls) = prog;

        for decl in decls {
            debug!(%self);
            debug!(%decl);
            match decl {
                Decl::Set { name, set } => {
                    let set = self.check_set(set)?;
                    self.add_set(name, set)
                }

                Decl::Element { name, element, set } => {
                    let set = self.check_set(set)?;
                    let element = self.check_element(element, &set)?;
                    self.add_element(name, element, set)
                }
                Decl::Signature { .. } => {
                    return Err(CheckError::Unimplemented("signatures".to_string()));
                }
                Decl::Instance { .. } => {
                    return Err(CheckError::Unimplemented("instances".to_string()));
                }
            }?;
        }
        debug!(%self, "END");
        Ok(())
    }

    #[instrument(skip(self), level = "debug", fields(%set))]
    fn check_set(&self, set: &Set) -> Result<Set, CheckError> {
        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(CheckError::Unimplemented("instances as sets".to_string())),
            Set::Var(v) => {
                if let Some(deref) = self.wf_sets.get(v) {
                    Ok(deref.clone())
                } else {
                    Err(CheckError::Unbound(v.clone()))
                }
            }
        }
    }

    fn _check_literal_set_helper(&self, claimed: &Set, should_be: Set) -> Result<(), CheckError> {
        if !self.set_equal(claimed, &should_be) {
            Err(CheckError::WrongSetForElement(claimed.clone(), should_be))
        } else {
            Ok(())
        }
    }

    #[instrument(skip(self), level = "debug", fields(%element, %set))]
    fn check_element(&self, element: &Element, set: &Set) -> Result<Element, CheckError> {
        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) => {
                if let Some(CheckedElement {
                    element: found_element,
                    set: found_set,
                }) = self.wf_elements.get(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, found_set) {
                        return Err(CheckError::WrongSetForElement(
                            set.clone(),
                            found_set.clone(),
                        ));
                    }
                    Ok(found_element.clone())
                } else {
                    Err(CheckError::Unbound(v.clone()))
                }
            }
            Element::Record(assignations) => {
                let rej = |reason| CheckError::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 { .. } => {
                Err(CheckError::Unimplemented("element project".to_string()))
            }
            Element::Inject { .. } => Err(CheckError::Unimplemented("element inject".to_string())),
            Element::App(_, _) => Err(CheckError::Unimplemented("element app".to_string())),
            Element::Case { .. } => Err(CheckError::Unimplemented("element case".to_string())),
        }
    }
}