zoisite 0.1.0

Zoisite is a programming language designed specifically for competitive programming.
Documentation
use crate::resolve_context::ResolveContext;
use itertools::Itertools;
use std::collections::{BTreeSet, HashMap};
use std::iter::once;

#[derive(PartialEq, Eq, Hash, Clone, Debug)]
pub enum Type {
    TyVar(usize),
    Unit,
    Int,
    Float,
    Bool,
    Str,
    Char,
    Array(Box<Type>),
    Tuple(Vec<Type>),
    Option(Box<Type>),
    Invalid,
}

impl Type {
    pub fn inner_ty(&self) -> Option<Type> {
        match self {
            Type::Str => Some(Type::Char),
            Type::Array(inner_ty) => Some(*inner_ty.clone()),
            Type::Option(inner_ty) => Some(*inner_ty.clone()),
            _ => None,
        }
    }
    pub fn contains_ty_var(&self) -> bool {
        match self {
            Type::TyVar(_) => true,
            Type::Unit => false,
            Type::Int => false,
            Type::Float => false,
            Type::Bool => false,
            Type::Str => false,
            Type::Char => false,
            Type::Array(inner_ty) => inner_ty.contains_ty_var(),
            Type::Tuple(inner_tys) => inner_tys.iter().any(Self::contains_ty_var),
            Type::Option(inner_ty) => inner_ty.contains_ty_var(),
            Type::Invalid => false,
        }
    }
    pub fn wrap_in_array(self) -> Type {
        Type::Array(Box::new(self))
    }
    pub fn wrap_in_option(self) -> Type {
        Type::Option(Box::new(self))
    }
    pub fn substitute_with(&self, subst: &HashMap<usize, Type>) -> Type {
        match self {
            Type::TyVar(id) => subst.get(id).map_or(self.clone(), |ty| ty.substitute_with(subst)),
            Type::Array(inner_ty) => inner_ty.substitute_with(subst).wrap_in_array(),
            Type::Tuple(inner_ty) => Type::Tuple(inner_ty.iter().map(|ty| ty.substitute_with(subst)).collect()),
            Type::Option(inner_ty) => inner_ty.substitute_with(subst).wrap_in_option(),
            _ => self.clone(),
        }
    }
    pub fn list_ty_var(&self) -> BTreeSet<usize> {
        match self {
            Type::TyVar(ty_var_id) => BTreeSet::from([*ty_var_id]),
            Type::Array(inner_ty) => inner_ty.list_ty_var(),
            Type::Tuple(inner_tys) => inner_tys.iter().flat_map(|inner_ty| inner_ty.list_ty_var()).collect(),
            Type::Option(inner_ty) => inner_ty.list_ty_var(),
            _ => BTreeSet::new(),
        }
    }
    pub fn mangle(&self) -> String {
        match self {
            Type::TyVar(_) => "a".to_owned(),
            Type::Unit => "u".to_owned(),
            Type::Int => "i".to_owned(),
            Type::Float => "f".to_owned(),
            Type::Bool => "b".to_owned(),
            Type::Str => "s".to_owned(),
            Type::Char => "c".to_owned(),
            Type::Array(inner_ty) => inner_ty.mangle() + "[]",
            Type::Tuple(inner_tys) => inner_tys.iter().map(Self::mangle).join(""),
            Type::Option(inner_ty) => inner_ty.mangle() + "?",
            Type::Invalid => "x".to_owned(),
        }
    }
}

#[derive(Debug, Clone, Hash, Eq, PartialEq)]
pub struct FuncType {
    pub params_ty: Vec<Type>,
    pub return_ty: Type,
}

impl FuncType {
    pub fn new(params_ty: Vec<Type>, return_ty: Type) -> Self {
        Self {
            params_ty,
            return_ty,
        }
    }
    pub fn ty_vars(&self) -> BTreeSet<usize> {
        self.params_ty.iter().chain(once(&self.return_ty)).flat_map(|ty| ty.list_ty_var()).collect()
    }
    pub fn instantiate(&self, resolve_ctx: &mut ResolveContext) -> Self {
        let ty_vars_subst = self.ty_vars().into_iter().map(|ty_var_id| (ty_var_id, resolve_ctx.new_ty_var())).collect();
        let params_ty = self.params_ty.iter().map(|ty| ty.substitute_with(&ty_vars_subst)).collect();
        let return_ty = self.return_ty.substitute_with(&ty_vars_subst);
        Self::new(params_ty, return_ty)
    }
    pub fn substitute_with(&self, subst: &HashMap<usize, Type>) -> Self {
        Self {
            params_ty: self.params_ty.iter().map(|ty| ty.substitute_with(subst)).collect(),
            return_ty: self.return_ty.substitute_with(subst),
        }
    }
    pub fn mangle(&self) -> String {
        format!("{}:{}", self.params_ty.iter().map(Type::mangle).join(""), self.return_ty.mangle())
    }
}