blr-lang 0.1.0

//! Core is focused on closures.
//!
//! This layer prioritizes closure conversion from the curried [`crate::mantle::Expr`] to closure based [`Expr`] intermediate representation.

pub mod sexpr;

use std::{
    collections::{BTreeMap, BTreeSet, HashMap},
    panic,
};

use super::mantle::{self, Symbol};
use crate::external_type::{ExternalType, FunctionType};

#[derive(PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Hash)]
// TODO This needs to be private
// Fix Case branches to keep their var around
pub struct VarId(pub usize);

#[derive(Eq, PartialEq, Clone, Hash)]
pub struct Var {
    pub id: VarId,
    pub typ: Type,
}

impl Ord for Var {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.id.cmp(&other.id)
    }
}
impl PartialOrd for Var {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

#[derive(PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Hash)]
pub struct ItemId(pub u32);

#[derive(Clone, PartialEq)]
pub enum Expr {
    Unit,
    Var(Var),
    Int(i64),
    Float(f64),
    String(String),
    Closure(Type, ItemId, Vec<Var>),
    AppClosure(Box<Self>, Box<Self>),
    AppExternal(Type, FunctionType, ItemId, Vec<Self>),
    Local(Var, Box<Self>, Box<Self>),
    Access(Box<Self>, usize),
    Tuple(Vec<Self>, Vec<String>),
    Tag(Type, usize, Box<Self>),
    Case(Type, Box<Self>, Vec<Self>),
}

impl Expr {
    pub fn apply_closure(closure: Self, param: Self) -> Self {
        Self::AppClosure(Box::new(closure), Box::new(param))
    }

    fn local(var: Var, defn: Self, body: Self) -> Self {
        Self::Local(var, Box::new(defn), Box::new(body))
    }
    fn tuple(fields: impl IntoIterator<Item = (String, Self)>) -> Self {
        let (labels, fields) = fields.into_iter().unzip();
        Expr::Tuple(fields, labels)
    }

    fn free_vars_aux(&self, free: &mut BTreeSet<Var>) {
        match self {
            Expr::Var(var) => {
                free.insert(var.clone());
            }
            Expr::Unit | Expr::Int(_) | Expr::Float(_) | Expr::String(_) => {}
            Expr::Closure(_, _, vars) => {
                for var in vars {
                    free.insert(var.clone());
                }
            }
            Expr::AppClosure(abs, param) => {
                abs.free_vars_aux(free);
                param.free_vars_aux(free);
            }
            Expr::AppExternal(_, _, _abs, params) => {
                params.iter().for_each(|param| param.free_vars_aux(free));
            }
            Expr::Local(var, defn, body) => {
                body.free_vars_aux(free);
                defn.free_vars_aux(free);
                free.remove(var);
            }
            Expr::Access(tuple, _) => tuple.free_vars_aux(free),
            Expr::Tuple(fields, _labels) => {
                fields.iter().for_each(|field| field.free_vars_aux(free))
            }
            Expr::Tag(_, _, expr) => expr.free_vars_aux(free),
            Expr::Case(_, scrutinee, branches) => {
                scrutinee.free_vars_aux(free);
                branches.iter().for_each(|case| case.free_vars_aux(free));
            }
        }
    }

    fn free_vars(&self) -> BTreeSet<Var> {
        let mut free = BTreeSet::default();
        self.free_vars_aux(&mut free);
        free
    }

    fn access(body: Self, field: usize) -> Expr {
        Expr::Access(Box::new(body), field)
    }
    fn tag(typ: Type, index: usize, expr: Expr) -> Expr {
        Expr::Tag(typ, index, Box::new(expr))
    }
    fn case(typ: Type, scrutinee: Expr, branches: impl IntoIterator<Item = Expr>) -> Expr {
        Expr::Case(typ, Box::new(scrutinee), branches.into_iter().collect())
    }

    fn rename(&mut self, subst: &HashMap<Var, Var>) {
        match self {
            Expr::Var(var) => {
                if let Some(new_var) = subst.get(var) {
                    *var = new_var.clone();
                }
            }
            Expr::Unit | Expr::Int(_) | Expr::Float(_) | Expr::String(_) => {}
            Expr::Closure(_, _, vars) => {
                for var in vars.iter_mut() {
                    if let Some(new_var) = subst.get(var) {
                        *var = new_var.clone();
                    }
                }
            }
            Expr::AppClosure(abs, param) => {
                abs.rename(subst);
                param.rename(subst);
            }
            Expr::AppExternal(_, _, _abs, params) => {
                params.iter_mut().for_each(|param| param.rename(subst));
            }
            Expr::Local(_, defn, body) => {
                defn.rename(subst);
                body.rename(subst);
            }
            Expr::Access(body, _) => body.rename(subst),
            Expr::Tuple(fields, _labels) => fields.iter_mut().for_each(|field| field.rename(subst)),
            Expr::Tag(_, _, expr) => expr.rename(subst),
            Expr::Case(_, scrutinee, branches) => {
                scrutinee.rename(subst);
                branches.iter_mut().for_each(|case| case.rename(subst));
            }
        }
    }

    pub fn type_of(&self) -> Type {
        match self {
            Expr::Var(var) => var.typ.clone(),
            Expr::Unit => Type::Unit,
            Expr::Int(_) => Type::Int,
            Expr::Float(_) => Type::Float,
            Expr::String(_) => Type::String,
            Expr::Closure(ty, _, _) => ty.clone(),
            Expr::AppClosure(closure, arg) => {
                let Type::Closure(arg_ty, ret_ty) = closure.type_of() else {
                    panic!("Apply was applied to non closure type");
                };
                if *arg_ty != arg.type_of() {
                    panic!("Closure was applied to argument of wrong type");
                }
                *ret_ty
            }
            Expr::AppExternal(typ, _, _, _) => typ.clone(),
            Expr::Local(_, _, body) => body.type_of(),
            Expr::Access(tuple, field) => match tuple.type_of() {
                Type::Tuple(fields, _labels) => fields[*field].clone(),
                _ => panic!("Access was applied to non closure env node"),
            },
            Expr::Tuple(fields, labels) => Type::labeled_tuple(
                fields
                    .iter()
                    .zip(labels)
                    .map(|(field, name)| (name.clone(), field.type_of())),
            ),
            Expr::Tag(_, _, _expr) => todo!(),
            Expr::Case(_, _, _) => todo!(),
        }
    }
}

#[derive()]
pub enum Item {
    /// A top level native function with implementation code within the current module.
    /// The function follows a partial application convention taking as its first parameter a
    /// closure env, an arbitrary second parmaeter, and returning a closure struct or final result.
    Local(LocalItem),
    /// A top level native function with its implementation code outside the current module.
    /// The function follows a partial application convention taking as its first parameter a
    /// closure env, an arbitrary second parmaeter, and returning a closure struct or final result.
    Remote(RemoteItem),
    /// A top level function with its implementation code outside the current module and outside
    /// this compiler. The function follows WASM WIT conventions and its implementation must be
    /// found during linking.
    External(ExternalItem),
}

#[derive(PartialEq)]
pub struct LocalItem {
    pub name: String,
    pub params: Vec<Var>,
    pub ret_typ: Type,
    pub body: Expr,
}

#[derive(PartialEq)]
pub struct RemoteItem {
    pub symbol: Symbol,
    pub param_typ: Type,
    pub ret_typ: Type,
}
#[derive()]
pub struct ExternalItem {
    pub symbol: Symbol,
    pub param_typs: Vec<Type>,
    pub ret_typs: Vec<Type>,
    pub external_typ: FunctionType,
}

#[derive(Default)]
pub struct VarSupply {
    next: usize,
    cache: HashMap<mantle::VarId, VarId>,
}

impl VarSupply {
    fn supply_for(&mut self, var: mantle::VarId) -> VarId {
        self.cache
            .entry(var)
            .or_insert_with(|| {
                let id = self.next;
                self.next += 1;
                VarId(id)
            })
            .to_owned()
    }

    pub fn supply(&mut self) -> VarId {
        let id = self.next;
        self.next += 1;
        VarId(id)
    }
}

#[derive(Debug, Default)]
struct ItemSupply {
    next: u32,
    cache: BTreeMap<mantle::ItemId, ItemId>,
}

impl ItemSupply {
    fn supply_for(&mut self, item_id: mantle::ItemId) -> ItemId {
        self.cache
            .entry(item_id)
            .or_insert_with(|| {
                let item_id = self.next;
                self.next += 1;
                ItemId(item_id)
            })
            .to_owned()
    }
    fn supply(&mut self) -> ItemId {
        let item_id = self.next;
        self.next += 1;
        ItemId(item_id)
    }
}

#[derive(Eq, PartialEq, Clone, Hash)]
pub enum Type {
    Unit,
    Int,
    Float,
    String,
    Closure(Box<Self>, Box<Self>),
    Abstraction(Vec<Self>, Box<Self>),
    Tuple(Vec<Self>, Option<Vec<String>>),
    Variant(Vec<Self>, Option<Vec<String>>),
    DataFrame,
}

impl Type {
    pub fn closure(arg: Self, ret: Self) -> Self {
        Self::Closure(Box::new(arg), Box::new(ret))
    }

    // Helper function to create a tuple type from a closure env
    pub fn closure_env(closure_ty: Type, vars: impl IntoIterator<Item = Self>) -> Type {
        Self::tuple([closure_ty].into_iter().chain(vars))
    }
    pub fn labeled_tuple(fields: impl IntoIterator<Item = (String, Self)>) -> Self {
        let (labels, fields) = fields.into_iter().unzip();
        Self::Tuple(fields, Some(labels))
    }
    pub fn tuple(fields: impl IntoIterator<Item = Self>) -> Self {
        Self::Tuple(fields.into_iter().collect(), None)
    }
    pub fn labeled_variant(fields: impl IntoIterator<Item = (String, Self)>) -> Self {
        let (labels, cases) = fields.into_iter().unzip();
        Self::Variant(cases, Some(labels))
    }
    pub fn variant(cases: impl IntoIterator<Item = Self>) -> Self {
        Self::Variant(cases.into_iter().collect(), None)
    }

    pub fn abs(parameters: impl IntoIterator<Item = Self>, ret: Self) -> Self {
        Self::Abstraction(parameters.into_iter().collect(), Box::new(ret))
    }
}

fn lower_typ(ty: mantle::Type) -> Type {
    match ty {
        mantle::Type::Unit => Type::Unit,
        mantle::Type::Int => Type::Int,
        mantle::Type::Float => Type::Float,
        mantle::Type::String => Type::String,
        mantle::Type::DataFrame => Type::DataFrame,
        mantle::Type::Abs(arg, ret) => Type::closure(lower_typ(*arg), lower_typ(*ret)),
        mantle::Type::Var(_) | mantle::Type::TypAbs(_, _) => Type::Unit,
        // HACK:
        // We need better dead code elimination in mantle::simplify
        // This hack ignores the fact we allow polymorphic types to exist in never accessed fields
        // of tuples.
        //panic!(
        //    "ICE: Type function or variable appeared in closure conversion. This should've been handled by monomorphization."
        //),
        mantle::Type::Prod(row) => match row {
            mantle::Row::Open(_type_var) => panic!("Cannot lower open row type"),
            mantle::Row::Closed(items) => Type::labeled_tuple(
                items
                    .iter()
                    .map(|(name, typ)| (name.clone(), lower_typ(typ.clone()))),
            ),
        },
        mantle::Type::Sum(row) => match row {
            mantle::Row::Open(_type_var) => panic!("ICE: Cannot lower open row type"),
            mantle::Row::Closed(items) => Type::labeled_variant(
                items
                    .iter()
                    .map(|(name, typ)| (name.clone(), lower_typ(typ.clone()))),
            ),
        },
    }
}

struct ClosureConvert<'a> {
    var_supply: VarSupply,
    item_supply: &'a mut ItemSupply,
    item_kinds: &'a BTreeMap<ItemId, ItemKind>,
    items: BTreeMap<ItemId, Item>,
}

#[derive(Debug, Clone, Copy)]
enum ItemKind {
    Local,
    External,
    Remote,
}

impl<'a> ClosureConvert<'a> {
    fn make_closure(
        &mut self,
        item_id: ItemId,
        var: Var,
        body: mantle::Expr,
        env: im::HashMap<mantle::Var, Var>,
    ) -> (Expr, LocalItem) {
        let ret = lower_typ(body.type_of());
        let mut body = self.convert(body, env);
        let mut free_vars = body.free_vars();
        free_vars.remove(&var);

        let vars: Vec<Var> = free_vars.iter().cloned().collect();
        let closure_ty = Type::closure(var.typ.clone(), ret.clone());
        let env_var = Var {
            id: self.var_supply.supply(),
            typ: Type::closure_env(closure_ty.clone(), vars.iter().map(|var| var.typ.clone())),
        };
        let subst = free_vars
            .into_iter()
            .enumerate()
            .map(|(i, var)| {
                let id = self.var_supply.supply();
                let new_var = Var {
                    id,
                    typ: var.typ.clone(),
                };
                body = Expr::local(
                    new_var.clone(),
                    Expr::access(Expr::Var(env_var.clone()), i + 1),
                    body.clone(),
                );
                (var, new_var)
            })
            .collect::<HashMap<_, _>>();

        body.rename(&subst);

        let params = vec![env_var, var];
        let item = LocalItem {
            name: format!("anon_local{}", item_id.0),
            params,
            ret_typ: ret,
            body,
        };
        (Expr::Closure(closure_ty, item_id, vars), item)
    }

    fn convert_native_item(
        &mut self,
        item_id: ItemId,
        item: mantle::NativeItem,
        env: im::HashMap<mantle::Var, Var>,
    ) -> Expr {
        let mantle::Expr::Abstraction(fun_var, body) = item.expr else {
            panic!("local item should be an abstraction: {}", item.symbol.field)
        };
        let var = Var {
            id: self.var_supply.supply_for(fun_var.id),
            typ: lower_typ(fun_var.typ.clone()),
        };
        let (expr, mut local_item) =
            self.make_closure(item_id, var.clone(), *body, env.update(fun_var, var));
        local_item.name = item.symbol.field;
        // HACK? For local main item only ignore the closure parameter if it doesn't close over any
        // variables.
        if local_item.name == "main"
            && let Type::Tuple(fields, _) = &local_item.params[0].typ
            && fields.len() == 1
        {
            local_item.params = vec![local_item.params.pop().unwrap()];
        }

        self.items.insert(item_id, Item::Local(local_item));
        expr
    }
    // Construct a local item that is capable of calling the external item.
    // This is in effect currying the external item so we can partially apply the item.
    fn convert_external_item(&mut self, item_id: ItemId, item: mantle::ExternalItem) {
        let ret_typ = convert_external_type(*item.external_type.ret.clone());
        let mut param_typs: Vec<_> = item
            .external_type
            .parameter_typs
            .iter()
            .map(|param| convert_external_type(param.clone()))
            .collect();
        let external_item = self.item_supply.supply();
        self.items.insert(
            external_item,
            Item::External(ExternalItem {
                param_typs: param_typs.to_vec(),
                ret_typs: vec![ret_typ.clone()],
                symbol: item.symbol.clone(),
                external_typ: item.external_type.clone(),
            }),
        );
        // Build chain curried chain of local item closures
        let item_ids: Vec<_> = [item_id]
            .into_iter()
            .chain(param_typs.iter().skip(1).map(|_| self.item_supply.supply()))
            .collect();

        let names: Vec<_> = item_ids
            .iter()
            .enumerate()
            .map(|(i, _)| {
                if i == 0 {
                    item.symbol.field.clone()
                } else {
                    format!("{}[{}]", item.symbol.field.clone(), i)
                }
            })
            .collect();
        for (i, ((first, item_id), name)) in param_typs
            .iter()
            .zip(item_ids.iter())
            .zip(&names)
            .enumerate()
            .take(param_typs.len() - 1)
        {
            let mut vars: Vec<_> = param_typs
                .iter()
                .take(i)
                .map(|ty| Var {
                    id: self.var_supply.supply(),
                    typ: ty.clone(),
                })
                .collect();
            let closure_ty = param_typs
                .iter()
                .skip(i)
                .rfold(ret_typ.clone(), |ret, param| {
                    Type::closure(param.clone(), ret)
                });
            let env_var = Var {
                id: self.var_supply.supply(),
                typ: Type::closure_env(closure_ty, param_typs.iter().take(i).cloned()),
            };
            let first_var = Var {
                id: self.var_supply.supply(),
                typ: first.clone(),
            };
            vars.push(first_var.clone());

            let ret_closure_typ = param_typs
                .iter()
                .skip(i + 1)
                .rfold(ret_typ.clone(), |ret, param| {
                    Type::closure(param.clone(), ret)
                });

            let body = vars.iter().take(vars.len() - 1).enumerate().fold(
                Expr::Closure(ret_closure_typ.clone(), item_ids[i + 1], vars.clone()),
                |body, (i, var)| {
                    Expr::local(
                        var.clone(),
                        Expr::access(Expr::Var(env_var.clone()), i + 1),
                        body,
                    )
                },
            );

            self.items.insert(
                *item_id,
                Item::Local(LocalItem {
                    name: name.clone(),
                    params: vec![env_var, first_var],
                    ret_typ: Type::closure_env(
                        ret_closure_typ,
                        param_typs.iter().take(i + 1).cloned(),
                    ),
                    body,
                }),
            );
        }

        let last = param_typs.pop().unwrap();
        let last_var = Var {
            id: self.var_supply.supply(),
            typ: last.clone(),
        };
        let closure_typ = Type::closure(last.clone(), ret_typ.clone());
        let env_var = Var {
            id: self.var_supply.supply(),
            typ: Type::closure_env(closure_typ.clone(), param_typs.clone()),
        };
        let body = Expr::AppExternal(
            ret_typ.clone(),
            item.external_type,
            external_item,
            param_typs
                .iter()
                .enumerate()
                .map(|(i, _param)| Expr::access(Expr::Var(env_var.clone()), i + 1))
                .chain([Expr::Var(last_var.clone())])
                .collect(),
        );
        self.items.insert(
            // use last item id
            item_ids[item_ids.len() - 1],
            Item::Local(LocalItem {
                name: names[names.len() - 1].clone(),
                params: vec![env_var, last_var],
                ret_typ,
                body,
            }),
        );
    }
    fn convert(&mut self, expr: mantle::Expr, env: im::HashMap<mantle::Var, Var>) -> Expr {
        match expr {
            mantle::Expr::Unit => Expr::Unit,
            mantle::Expr::Integer(i) => Expr::Int(i),
            mantle::Expr::Float(f) => Expr::Float(f),
            mantle::Expr::String(s) => Expr::String(s),
            mantle::Expr::Variable(var) => Expr::Var(
                env.get(&var)
                    .expect(&format!("var should be defined: {var:?}"))
                    .clone(),
            ),
            mantle::Expr::TypAbs(_, _) | mantle::Expr::TypApp(_, _) => {
                panic!("ICE: Generics appeared after monomorphizing")
            }
            mantle::Expr::Local(var, defn, body) => {
                let defn = self.convert(*defn, env.clone());
                let v = Var {
                    id: self.var_supply.supply_for(var.id),
                    typ: lower_typ(var.typ.clone()),
                };
                let body = self.convert(*body, env.update(var, v.clone()));
                Expr::local(v, defn, body)
            }
            mantle::Expr::Abstraction(fun_var, body) => {
                let var = Var {
                    id: self.var_supply.supply_for(fun_var.id),
                    typ: lower_typ(fun_var.typ.clone()),
                };
                let item_id = self.item_supply.supply();
                let (expr, item) =
                    self.make_closure(item_id, var.clone(), *body, env.update(fun_var, var));
                self.items.insert(item_id, Item::Local(item));
                expr
            }
            mantle::Expr::Application(fun, arg) => {
                let closure = self.convert(*fun, env.clone());
                let arg = self.convert(*arg, env);
                Expr::apply_closure(closure, arg)
            }
            mantle::Expr::Tuple(fields) => Expr::tuple(
                fields
                    .into_iter()
                    .map(|(name, field)| (name, self.convert(field, env.clone()))),
            ),
            mantle::Expr::Field(body, index) => Expr::access(self.convert(*body, env), index),
            mantle::Expr::Tag(typ, idx, expr) => Expr::tag(
                lower_typ(typ.clone()),
                idx,
                self.convert(*expr, env.clone()),
            ),
            mantle::Expr::Case(typ, scrutinee, branches) => Expr::case(
                lower_typ(typ.clone()),
                self.convert(*scrutinee, env.clone()),
                branches.into_iter().map(|branch| {
                    //Branches are just abstractions
                    self.convert(mantle::Expr::abs(branch.param, branch.body), env.clone())
                }),
            ),
            mantle::Expr::Item(typ, item_id, symbol) => {
                let ty = lower_typ(typ.clone());
                let Type::Closure(param_ty, ret_ty) = ty.clone() else {
                    panic!("item should have closure type")
                };
                let item_id = self.item_supply.supply_for(item_id);
                match self.item_kind(item_id) {
                    ItemKind::Local => {}
                    ItemKind::External => {}
                    ItemKind::Remote => {
                        self.items.insert(
                            item_id,
                            Item::Remote(RemoteItem {
                                param_typ: *param_ty,
                                ret_typ: *ret_ty,
                                symbol,
                            }),
                        );
                    }
                };
                Expr::Closure(ty, item_id, vec![])
            }
        }
    }
    fn item_kind(&self, item_id: ItemId) -> ItemKind {
        self.item_kinds
            .get(&item_id)
            .cloned()
            .unwrap_or(ItemKind::Remote)
    }
}

// TODO This is lossy, we need to know the core type independently
fn convert_external_type(typ: ExternalType) -> Type {
    match typ {
        ExternalType::Unit => Type::Unit,
        ExternalType::Int => Type::Int,
        ExternalType::Float => Type::Float,
        ExternalType::String => Type::String,
        ExternalType::Resource(_) => Type::DataFrame,
        ExternalType::Fun(FunctionType {
            parameter_names: _,
            parameter_typs,
            ret,
        }) => Type::abs(
            parameter_typs.into_iter().map(convert_external_type),
            convert_external_type(*ret),
        ),
        ExternalType::Record(_, fields) => Type::labeled_tuple(
            fields
                .into_iter()
                .map(|(label, field_typ)| (label, convert_external_type(field_typ))),
        ),
    }
}

#[derive()]
pub struct Module {
    pub closure_items: BTreeMap<ItemId, Item>,
}

pub fn closure_convert(module: mantle::Module) -> Module {
    let mut item_supply = ItemSupply::default();
    let item_kinds = module
        .items
        .iter()
        .map(|(item_id, item)| match item {
            mantle::Item::Native(_) => (item_supply.supply_for(*item_id), ItemKind::Local),
            // NOTE the item_id
            mantle::Item::External(_) => (item_supply.supply_for(*item_id), ItemKind::External),
        })
        .collect();

    let closure_items = module
        .items
        .into_iter()
        .flat_map(|(item_id, item)| {
            let item_id = item_supply.supply_for(item_id);
            let var_supply = VarSupply::default();
            let mut conversion = ClosureConvert {
                var_supply,
                item_kinds: &item_kinds,
                item_supply: &mut item_supply,
                items: BTreeMap::default(),
            };
            match item {
                mantle::Item::Native(native_item) => {
                    let env = im::HashMap::default();
                    conversion.convert_native_item(item_id, native_item, env);
                }
                mantle::Item::External(external_item) => {
                    conversion.convert_external_item(item_id, external_item);
                }
            };
            conversion.items
        })
        .collect();

    Module { closure_items }
}