Add environment, name generator, qualified types, type schemes

numbat/src/typechecker/environment.rs

@@ -0,0 +1,32 @@
+use super::substitutions::{ApplySubstitution, Substitution, SubstitutionError};
+use super::type_scheme::TypeScheme;
+
+use std::collections::HashMap;
+
+type Identifier = String; // TODO ?
+
+#[derive(Clone)]
+pub struct Environment(HashMap<Identifier, TypeScheme>);
+
+impl Environment {
+    pub fn new() -> Environment {
+        Environment(HashMap::new())
+    }
+
+    pub fn add(&mut self, v: Identifier, t: TypeScheme) {
+        self.0.insert(v, t);
+    }
+
+    pub(crate) fn lookup(&self, v: &Identifier) -> Option<&TypeScheme> {
+        self.0.get(v)
+    }
+}
+
+impl ApplySubstitution for Environment {
+    fn apply(&mut self, substitution: &Substitution) -> Result<(), SubstitutionError> {
+        for (_, t) in &mut self.0 {
+            t.apply(substitution)?;
+        }
+        Ok(())
+    }
+}

numbat/src/typechecker/mod.rs

@@ -2,9 +2,13 @@
 mod tests;
 
 mod constraints;
+mod environment;
 mod error;
 mod incompatible_dimensions;
+mod name_generator;
+mod qualified_type;
 mod substitutions;
+mod type_scheme;
 
 use std::collections::{HashMap, HashSet};
 

numbat/src/typechecker/name_generator.rs

@@ -0,0 +1,18 @@
+use crate::type_variable::TypeVariable;
+
+#[derive(Clone)]
+pub struct NameGenerator {
+    counter: u64,
+}
+
+impl NameGenerator {
+    pub fn new() -> NameGenerator {
+        NameGenerator { counter: 0 }
+    }
+
+    pub fn fresh_type_variable(&mut self) -> TypeVariable {
+        let name = format!("T{}", self.counter);
+        self.counter += 1;
+        TypeVariable::new(name)
+    }
+}

numbat/src/typechecker/qualified_type.rs

@@ -0,0 +1,117 @@
+use crate::{type_variable::TypeVariable, Type};
+
+use super::{
+    substitutions::{ApplySubstitution, Substitution, SubstitutionError},
+    type_scheme::TypeScheme,
+};
+
+/// A predicate on type variables.
+#[derive(Clone, Debug, PartialEq)]
+pub enum Bound {
+    // TODO: At the moment, this can only be IsDim(Type::TVar(_)). Maybe we can find a better representation.
+    IsDim(Type),
+}
+
+#[derive(Clone, Debug, PartialEq)]
+pub struct Bounds(Vec<Bound>);
+
+impl Bounds {
+    pub fn none() -> Bounds {
+        Bounds(vec![])
+    }
+
+    pub fn iter(&self) -> std::slice::Iter<Bound> {
+        self.0.iter()
+    }
+
+    pub fn iter_mut(&mut self) -> std::slice::IterMut<Bound> {
+        self.0.iter_mut()
+    }
+
+    pub fn union(self, other: Bounds) -> Bounds {
+        let mut bounds = self;
+        bounds.0.extend(other.0);
+        bounds
+    }
+}
+
+impl FromIterator<Bound> for Bounds {
+    fn from_iter<T: IntoIterator<Item = Bound>>(iter: T) -> Self {
+        Bounds(iter.into_iter().collect())
+    }
+}
+
+/// A qualified type is a type with a (potentially empty) set of bounds
+/// on the type variables.
+///
+/// For example, the type of the square-function (D -> D^2), needs an
+/// additional `D: Dim` bound, as arbitrary types (like Bool) can not
+/// be squared.
+#[derive(Clone, Debug, PartialEq)]
+pub struct QualifiedType {
+    pub inner: Type,
+    pub bounds: Bounds,
+}
+
+impl QualifiedType {
+    pub fn new(inner: Type, bounds: Bounds) -> QualifiedType {
+        QualifiedType { inner, bounds }
+    }
+
+    pub fn pretty_print(&self) -> String {
+        let bounds = self
+            .bounds
+            .iter()
+            .map(|b| match b {
+                Bound::IsDim(t) => format!("Dim({t})", t = t),
+            })
+            .collect::<Vec<String>>()
+            .join(", ");
+        if bounds.is_empty() {
+            self.inner.to_string()
+        } else {
+            format!("{} where {}", self.inner.to_string(), bounds)
+        }
+    }
+
+    pub fn quantify(&self, variables: &[TypeVariable]) -> TypeScheme {
+        let mut qt = self.clone();
+        for (i, v) in variables.iter().enumerate() {
+            qt.apply(&Substitution::single(
+                v.clone(),
+                Type::TVar(TypeVariable::Quantified(i)),
+            ))
+            .unwrap();
+        }
+
+        TypeScheme::new(variables.len(), qt)
+    }
+
+    pub fn type_variables(&self) -> Vec<TypeVariable> {
+        self.inner.type_variables()
+    }
+
+    pub(crate) fn instantiate(&self, type_variables: &[TypeVariable]) -> QualifiedType {
+        QualifiedType {
+            inner: self.inner.instantiate(type_variables),
+            bounds: self
+                .bounds
+                .iter()
+                .map(|b| match b {
+                    Bound::IsDim(t) => Bound::IsDim(t.instantiate(type_variables)),
+                })
+                .collect(),
+        }
+    }
+}
+
+impl ApplySubstitution for QualifiedType {
+    fn apply(&mut self, substitution: &Substitution) -> Result<(), SubstitutionError> {
+        self.inner.apply(substitution)?;
+
+        for Bound::IsDim(v) in self.bounds.iter_mut() {
+            v.apply(substitution)?;
+        }
+        Ok(())
+    }
+}

numbat/src/typechecker/type_scheme.rs

@@ -0,0 +1,42 @@
+use super::name_generator::NameGenerator;
+use super::qualified_type::{Bounds, QualifiedType};
+use super::substitutions::{ApplySubstitution, Substitution, SubstitutionError};
+use crate::typed_ast::Type;
+
+/// A type *scheme* is a type with a set of universally quantified
+/// type variables.
+///
+/// The way this is represented is by using type variables of kind
+/// `Quantified(i)` in the type tree, where `i` is a unique index
+/// for each universally quantified type variable.
+///
+/// Type schemes can be created by calling .quantify() on a qualified
+/// type.
+#[derive(Clone, Debug, PartialEq)]
+pub struct TypeScheme(usize, QualifiedType);
+
+impl TypeScheme {
+    pub fn new(num_quantified: usize, qt: QualifiedType) -> TypeScheme {
+        TypeScheme(num_quantified, qt)
+    }
+
+    pub fn instantiate(&self, name_generator: &mut NameGenerator) -> QualifiedType {
+        // Replace $gen_i by fresh type variables
+
+        let new_type_variables = (0..self.0)
+            .map(|_| name_generator.fresh_type_variable())
+            .collect::<Vec<_>>();
+
+        self.1.instantiate(&new_type_variables)
+    }
+
+    pub fn from_type(type_: Type) -> TypeScheme {
+        TypeScheme(0, QualifiedType::new(type_, Bounds::none()))
+    }
+}
+
+impl ApplySubstitution for TypeScheme {
+    fn apply(&mut self, substitution: &Substitution) -> Result<(), SubstitutionError> {
+        self.1.apply(substitution)
+    }
+}

numbat/src/typed_ast.rs

@@ -174,6 +174,14 @@ impl Type {
             (t1, t2) => t1 == t2,
         }
     }
+
+    pub(crate) fn type_variables(&self) -> Vec<TypeVariable> {
+        todo!()
+    }
+
+    pub(crate) fn instantiate(&self, type_variables: &[TypeVariable]) -> Type {
+        todo!()
+    }
 }
 
 #[derive(Debug, Clone, PartialEq)]