core: add NDArrayType & OptionType

nac3core/src/toplevel/numpy.rs

@@ -2,7 +2,7 @@ use crate::{
     toplevel::helper::PrimDef,
     typecheck::{
         type_inferencer::PrimitiveStore,
-        typedef::{Type, TypeEnum, Unifier, VarMap},
+        typedef::{into_var_map, Object, Type, TypeEnum, TypeVar, TypeVarId, Unifier, VarMap},
     },
 };
 use itertools::Itertools;
@@ -79,3 +79,83 @@ pub fn unpack_ndarray_params(
     let ndims = *params.get(&store.ndarray_ndims_tvar.id).unwrap();
     NDArrayParams { dtype, ndims }
 }
+
+fn into_subst_map<I>(iterator: I) -> VarMap
+where
+    I: IntoIterator<Item = (TypeVarId, Option<Type>)>,
+{
+    let mut subst_map = VarMap::new();
+    for (tvar_id, maybe_ty) in iterator.into_iter() {
+        if let Some(ty) = maybe_ty {
+            subst_map.insert(tvar_id, ty);
+        }
+    }
+    subst_map
+}
+
+/// `Option<T>`
+pub struct OptionType {
+    pub ty: Type,
+
+    /// The [`TypeVar`] of the contained type with the `Option`.
+    /// This is the `T` in `Option<T>`.
+    pub type_tvar: TypeVar,
+}
+
+impl OptionType {
+    fn make(&self, unifier: &mut Unifier, typ: Type) -> Type {
+        self.subst(unifier, self.ty, Some(typ))
+    }
+
+    fn subst(&self, unifier: &mut Unifier, option: Type, typ: Option<Type>) -> Type {
+        // Check the object is indeed an option.
+        let ty_enum = &*unifier.get_ty(option);
+        ty_enum.into_object().with_id(PrimDef::NDArray.id());
+
+        let subst_map = into_subst_map([(self.type_tvar.id, typ)]);
+        unifier.subst(option, &subst_map).unwrap_or(option)
+    }
+
+    fn unpack_type_immutable(&self, unifier: &Unifier, option: Type) -> Type {
+        let ty_enum = &*unifier.get_ty_immutable(option);
+        ty_enum.into_object().with_id(PrimDef::Option.id()).get_param(self.type_tvar.id)
+    }
+}
+
+/// `NDArray<dtype, ndims>`
+pub struct NDArrayType {
+    pub ty: Type,
+    pub dtype_tvar: TypeVar,
+    pub ndims_tvar: TypeVar,
+}
+
+impl NDArrayType {
+    fn make(&self, unifier: &mut Unifier, dtype: Type, ndims: Type) -> Type {
+        self.subst(unifier, self.ty, Some(dtype), Some(ndims))
+    }
+
+    fn subst(
+        &self,
+        unifier: &mut Unifier,
+        ndarray: Type,
+        dtype: Option<Type>,
+        ndims: Option<Type>,
+    ) -> Type {
+        // Check the object is indeed an ndarray.
+        let ty_enum = &*unifier.get_ty(ndarray);
+        ty_enum.into_object().with_id(PrimDef::NDArray.id());
+
+        let subst_map = into_subst_map([(self.dtype_tvar.id, dtype), (self.ndims_tvar.id, ndims)]);
+        unifier.subst(ndarray, &subst_map).unwrap_or(ndarray)
+    }
+
+    fn unpack_ndims_immutable(&self, unifier: &Unifier, ndarray: Type) -> Type {
+        let ty_enum = &*unifier.get_ty_immutable(ndarray);
+        ty_enum.into_object().with_id(PrimDef::NDArray.id()).get_param(self.ndims_tvar.id)
+    }
+
+    fn unpack_dtype_immutable(&self, unifier: &Unifier, ndarray: Type) -> Type {
+        let ty_enum = &*unifier.get_ty_immutable(ndarray);
+        ty_enum.into_object().with_id(PrimDef::NDArray.id()).get_param(self.dtype_tvar.id)
+    }
+}

nac3core/src/typecheck/typedef/mod.rs

@@ -82,14 +82,14 @@ pub struct Call {
     pub loc: Option<Location>,
 }
 
-#[derive(Clone)]
+#[derive(Debug, Clone)]
 pub struct FuncArg {
     pub name: StrRef,
     pub ty: Type,
     pub default_value: Option<SymbolValue>,
 }
 
-#[derive(Clone)]
+#[derive(Debug, Clone)]
 pub struct FunSignature {
     pub args: Vec<FuncArg>,
     pub ret: Type,
@@ -165,8 +165,27 @@ impl RecordField {
     }
 }
 
+#[derive(Debug, Clone)]
+pub struct Object<'a> {
+    pub id: DefinitionId,
+    pub fields: &'a Mapping<StrRef, (Type, bool)>,
+    pub params: &'a VarMap,
+}
+
+impl<'a> Object<'a> {
+    /// debug assert `Object::id` == `expected_id`.
+    pub fn with_id(self, expected_id: DefinitionId) -> Self {
+        debug_assert_eq!(self.id, expected_id);
+        self
+    }
+
+    pub fn get_param(&self, tvar_id: TypeVarId) -> Type {
+        *self.params.get(&tvar_id).unwrap()
+    }
+}
+
 /// Category of variable and value types.
-#[derive(Clone)]
+#[derive(Debug, Clone)]
 pub enum TypeEnum {
     TRigidVar {
         id: TypeVarId,
@@ -246,6 +265,24 @@ impl TypeEnum {
             TypeEnum::TFunc { .. } => "TFunc",
         }
     }
+
+    /// Attempt to cast this [`TypeEnum`] into a [`TObj`], return None otherwise.
+    pub fn try_into_object(&self) -> Option<Object<'_>> {
+        match self {
+            TypeEnum::TObj { obj_id, params, fields } => {
+                Some(Object { id: *obj_id, params, fields })
+            }
+            _ => None,
+        }
+    }
+
+    /// Cast this [`TypeEnum`] into a [`TObj`], panic otherwise.
+    pub fn into_object<'a>(&'a self) -> Object<'a> {
+        match self.try_into_object() {
+            Some(obj) => obj,
+            None => panic!("Expecting TObj but got {self:?}"),
+        }
+    }
 }
 
 pub type SharedUnifier = Arc<Mutex<(UnificationTable<TypeEnum>, u32, Vec<Call>)>>;