blr-lang 0.1.0

use im::HashMap;
use std::collections::HashSet;
use std::ops::ControlFlow;
use tracing::trace;

use crate::compiler::mantle::Branch;

use super::{Expr, Kind, Module, TypApp, Type, Var, VarId};

#[allow(dead_code)]
pub enum Param {
    Ty(Kind),
    Val(Var),
}

pub trait IRExt {
    fn is_trivial(&self) -> bool;

    fn is_value(&self) -> bool;

    fn split_funs(self) -> (Vec<Param>, Expr);

    fn size(&self) -> usize;
}
impl IRExt for Expr {
    fn is_trivial(&self) -> bool {
        matches!(self, Expr::Variable(_) | Expr::Integer(_))
    }

    fn is_value(&self) -> bool {
        match self {
            Expr::Unit
            | Expr::Variable(_)
            | Expr::Integer(_)
            | Expr::Float(_)
            | Expr::String(_)
            | Expr::Abstraction(_, _)
            | Expr::Field(_, _) => true,
            Expr::TypAbs(_, expr) => expr.is_value(),
            Expr::TypApp(expr, _) => expr.is_value(),
            Expr::Local(_, defn, body) => defn.is_value() && body.is_value(),
            Expr::Application(_, _) => false,
            Expr::Tag(_, _, expr) => expr.is_value(),
            Expr::Case(_, _, _) => false,
            Expr::Item(_, _, _) => false,
            Expr::Tuple(_) => false,
        }
    }

    fn split_funs(self) -> (Vec<Param>, Expr) {
        fn split_funs(expr: Expr, params: &mut Vec<Param>) -> Expr {
            match expr {
                Expr::TypAbs(kind, expr) => {
                    params.push(Param::Ty(kind));
                    split_funs(*expr, params)
                }
                Expr::Abstraction(var, expr) => {
                    params.push(Param::Val(var));
                    split_funs(*expr, params)
                }
                expr => expr,
            }
        }
        let mut params = vec![];
        let body = split_funs(self, &mut params);
        (params, body)
    }

    fn size(&self) -> usize {
        fn size_app(expr: &Expr, params: usize) -> usize {
            match expr {
                Expr::Application(fun, param) => param.size() + size_app(fun, params + 1),
                Expr::TypApp(expr, _) => size_app(expr, params),
                expr => expr.size() + 10 * (1 + params),
            }
        }
        match self {
            Expr::Unit
            | Expr::Variable(_)
            | Expr::Integer(_)
            | Expr::Float(_)
            | Expr::String(_) => 0,
            Expr::Abstraction(_, body) => 10 + body.size(),
            Expr::Application(fun, param) => param.size() + size_app(fun, 1),
            Expr::TypAbs(_, expr) => expr.size(),
            Expr::TypApp(expr, _) => expr.size(),
            Expr::Local(var, defn, body) => {
                defn.size() + body.size() + (if var.typ.is_stack_alloc() { 0 } else { 10 })
            }
            Expr::Tuple(expr) => expr
                .iter()
                .fold(10, |size, (_name, expr)| size + expr.size()),
            // TODO Use index
            Expr::Field(product, _index) => product.size(),
            Expr::Tag(_, _, body) => body.size(),
            Expr::Case(_, scrutinee, branches) => {
                branches.iter().fold(scrutinee.size() + 20, |size, branch| {
                    size + branch.body.size()
                })
            }
            Expr::Item(_, _, _) => 0,
        }
    }
}

trait TypeExt {
    fn is_stack_alloc(&self) -> bool;
}
impl TypeExt for Type {
    fn is_stack_alloc(&self) -> bool {
        // Our only stack allocated type is Int
        matches!(self, Type::Int)
    }
}

// Subsitute the payload type in the haystack expression using DeBrujin index substitution
pub fn subst_typ(haystack: Expr, payload: Type) -> Expr {
    match haystack {
        Expr::Variable(var) => Expr::Variable(var.map_typ(|ty| ty.subst_typ(payload))),
        Expr::Unit => Expr::Unit,
        Expr::Integer(i) => Expr::Integer(i),
        Expr::Float(f) => Expr::Float(f),
        Expr::String(s) => Expr::String(s),
        Expr::Abstraction(var, expr) => Expr::abs(
            var.map_typ(|ty| ty.subst_typ(payload.clone())),
            subst_typ(*expr, payload),
        ),
        Expr::Application(fun, param) => {
            Expr::app(subst_typ(*fun, payload.clone()), subst_typ(*param, payload))
        }
        Expr::TypAbs(kind, expr) => Expr::ty_abs(kind, subst_typ(*expr, payload)),
        Expr::TypApp(expr, ty) => {
            Expr::ty_app(subst_typ(*expr, payload.clone()), ty.subst_typ(payload))
        }
        Expr::Local(var, defn, body) => Expr::local(
            var.map_typ(|ty| ty.subst_typ(payload.clone())),
            subst_typ(*defn, payload.clone()),
            subst_typ(*body, payload),
        ),
        Expr::Tuple(expr) => Expr::tuple(
            expr.into_iter()
                .map(|(name, expr)| (name, subst_typ(expr, payload.clone()))),
        ),
        Expr::Field(product, index) => Expr::field(subst_typ(*product, payload), index),
        Expr::Tag(_, _, _) => todo!(),
        Expr::Case(typ, expr, branches) => Expr::case(
            typ.subst_typ(payload.clone()),
            subst_typ(*expr, payload.clone()),
            branches.into_iter().map(|branch| Branch {
                param: branch.param.map_typ(|typ| typ.subst_typ(payload.clone())),
                body: subst_typ(branch.body, payload.clone()),
            }),
        ),
        Expr::Item(typ, item, symbol) => Expr::Item(typ.subst_typ(payload), item, symbol),
    }
}

#[derive(PartialEq, Eq, Clone, Copy, Debug)]
enum Occurrence {
    /// Our variable never occurs
    Dead,
    /// Our variable appears once.
    Once,
    /// Our variable appears once inside a Fun node.
    OnceInFun,
    /// Our variable appears many times.
    Many,
}

#[derive(Default, Debug, Clone)]
struct Occurrences {
    vars: HashMap<VarId, Occurrence>,
}
impl Occurrences {
    fn with_var_once(var: VarId) -> Self {
        let mut vars = HashMap::default();
        vars.insert(var, Occurrence::Once);
        Self { vars }
    }

    fn in_fun(self, free: &HashSet<VarId>) -> Self {
        Self {
            vars: self
                .vars
                .into_iter()
                .map(|(id, occ)| {
                    (
                        id,
                        match occ {
                            // Only mark the free variables of a function as OnceInFun.
                            // This prevents marking bindings fully encapsulated by the function as OnceInFun.
                            // In term:
                            // (fun [V0] (let [V1 3] V1))
                            // V1 should be considered Once not OnceInFun.
                            Occurrence::Once if free.contains(&id) => Occurrence::OnceInFun,
                            occ => occ,
                        },
                    )
                })
                .collect(),
        }
    }

    fn merge(mut self, other: Self) -> Self {
        for (var, occ) in other.vars {
            self.vars
                .entry(var)
                .and_modify(|self_occ| {
                    *self_occ = match (*self_occ, occ) {
                        (Occurrence::Dead, occ) | (occ, Occurrence::Dead) => occ,
                        (Occurrence::Many, _) | (_, Occurrence::Many) => Occurrence::Many,
                        (Occurrence::Once, Occurrence::Once) => Occurrence::Many,
                        (Occurrence::Once, _) | (Occurrence::OnceInFun, _) => Occurrence::Many,
                    };
                })
                .or_insert(occ);
        }
        self
    }

    fn lookup_var(&self, var: &Var) -> Occurrence {
        self.vars.get(&var.id).copied().unwrap_or(Occurrence::Dead)
    }

    fn mark_dead(&mut self, id: &VarId) {
        self.vars.remove(id);
    }
}

fn occurrence_analysis(expr: &Expr) -> (HashSet<VarId>, Occurrences) {
    match expr {
        Expr::Variable(var) => {
            let mut free = HashSet::default();
            free.insert(var.id);
            (free, Occurrences::with_var_once(var.id))
        }
        Expr::Unit | Expr::Integer(_) | Expr::Float(_) | Expr::String(_) => {
            (HashSet::default(), Occurrences::default())
        }
        Expr::Abstraction(var, expr) => {
            let (mut free, occs) = occurrence_analysis(expr);
            free.remove(&var.id);
            let occs = occs.in_fun(&free);
            (free, occs)
        }
        Expr::Application(fun, param) => {
            let (mut fun_free, fun_occs) = occurrence_analysis(fun);
            let (param_free, param_occurences) = occurrence_analysis(param);
            fun_free.extend(param_free);
            (fun_free, fun_occs.merge(param_occurences))
        }
        Expr::TypAbs(_, expr) => occurrence_analysis(expr),
        Expr::TypApp(expr, _) => occurrence_analysis(expr),
        Expr::Local(var, defn, body) => {
            let (mut free, occs) = occurrence_analysis(body);
            let (defn_free, defn_occs) = occurrence_analysis(defn);
            free.extend(defn_free);
            free.remove(&var.id);
            (free, defn_occs.merge(occs))
        }
        Expr::Tuple(expr) => expr.iter().fold(
            (HashSet::default(), Occurrences::default()),
            |(mut free, occrs), (_name, expr)| {
                let (f, o) = occurrence_analysis(expr);
                free.extend(f);
                (free, occrs.merge(o))
            },
        ),
        Expr::Field(expr, _) => occurrence_analysis(expr),
        Expr::Tag(_, _, expr) => occurrence_analysis(expr),
        Expr::Case(_, expr, branches) => {
            branches
                .iter()
                .fold(occurrence_analysis(expr), |(mut free, occrs), branch| {
                    let (mut f, o) = occurrence_analysis(&branch.body);
                    f.remove(&branch.param.id);
                    let o = o.in_fun(&f);
                    free.extend(f);
                    (free, occrs.merge(o))
                })
        }
        Expr::Item(_, _, _) => (HashSet::default(), Occurrences::default()),
    }
}

#[derive(Debug, Clone)]
enum SubstLocalDefn {
    // Delay simplifying the definition of the local until we have its surrounding context.
    // We only suspend simplifying local definitions that are used once so that we avoid exponential
    // runtimes.
    Suspend(Expr, Subst),
    // The local has already been simplified and can simply be inlined.
    Done(Expr),
}
type Subst = HashMap<VarId, SubstLocalDefn>;

#[derive(Debug, Clone)]
enum Definition {
    Unknown,
    BoundTo(Expr, Occurrence),
}
type InScope = HashMap<VarId, Definition>;

#[derive(Debug, Clone)]
struct Simplifier {
    // state to perform simplification
    occs: Occurrences,
    // Track which variables can be substituted with their definition.
    subst: Subst,

    // stats, used to determine if we simplified
    saturated_fun_count: usize,
    saturated_ty_fun_count: usize,
    locals_inlined: usize,
    field_access_inlined: usize,

    // configurable flags for simplification
    inline_size_threshold: usize,
}

impl Default for Simplifier {
    fn default() -> Self {
        Self {
            inline_size_threshold: 60,
            occs: Default::default(),
            subst: Default::default(),
            saturated_fun_count: Default::default(),
            saturated_ty_fun_count: Default::default(),
            locals_inlined: Default::default(),
            field_access_inlined: Default::default(),
        }
    }
}

// ContextEntry stores top-down state.
// This state is reconstructed during rebuild.
//
// TypApp and TypAbs together implement a direct monomorphization algorithm.
// Any direct type application of a type abstraction is simplified away as there is by definition
// only a single instantion of that type abstraction.
#[derive(Debug, Clone)]
enum ContextEntry {
    App {
        parameter: Expr,
    },
    TypApp {
        type_parameter: TypApp,
    },
    TypAbs {
        kind: Kind,
    },
    Local {
        var: Var,
        occ: Occurrence,
        body: Expr,
    },
    Field {
        index: usize,
    },
}

type Context = Vec<(ContextEntry, Subst)>;

enum Arg<'a> {
    Val(&'a Expr),
    #[allow(dead_code)]
    Ty(&'a Type),
}

impl Simplifier {
    fn new(occs: Occurrences) -> Self {
        Self {
            occs,
            ..Default::default()
        }
    }

    fn some_benefit(&self, expr: &Expr, in_scope: &InScope, ctx: &Context) -> bool {
        let (fun_params, _) = expr.clone().split_funs();
        let provided = ctx
            .iter()
            .rev()
            .map_while(|entry| match entry {
                (ContextEntry::App { parameter: param }, _) => Some(Arg::Val(param)),
                (ContextEntry::TypApp { type_parameter: ty }, _) => match ty {
                    TypApp::Ty(ty) => Some(Arg::Ty(ty)),
                    TypApp::Row(_row) => todo!(),
                },
                _ => None,
            })
            .collect::<Vec<_>>();

        // We have enough parameters to saturate our parameters.
        // We know this is a local function so there is benefit to inline
        // If we saturate all params to our function and then apply more params to the body there is value
        // in inlining.
        if provided.len() >= fun_params.len() {
            return true;
        }

        // If we have a non trivial parameter in context, there's some benefit.
        provided
            .iter()
            .take(fun_params.len())
            .any(|param| match param {
                Arg::Val(param) => {
                    !param.is_trivial()
                        || match param {
                            Expr::Variable(var) => {
                                matches!(in_scope.get(&var.id), Some(Definition::BoundTo(_, _)))
                            }
                            _ => false,
                        }
                }
                Arg::Ty(_) => false,
            })
    }

    fn rebuild(&mut self, mut expr: Expr, in_scope: InScope, mut ctx: Context) -> Expr {
        while let Some((entry, subst)) = ctx.pop() {
            self.subst = subst;
            match entry {
                ContextEntry::App { parameter: param } => {
                    if let Expr::Abstraction(var, body) = expr {
                        self.saturated_fun_count += 1;
                        return self.simplify(Expr::local(var, param, *body), in_scope, ctx);
                    } else {
                        let param = self.simplify(param, in_scope.clone(), vec![]);
                        expr = Expr::app(expr, param);
                    }
                }
                ContextEntry::TypApp { type_parameter: ty } => {
                    expr = if let Expr::TypAbs(_, body) = expr {
                        self.saturated_ty_fun_count += 1;
                        match ty {
                            TypApp::Ty(ty) => subst_typ(*body, ty),
                            TypApp::Row(_row) => todo!(),
                        }
                    } else {
                        Expr::ty_app(expr, ty)
                    }
                }
                ContextEntry::TypAbs { kind } => {
                    expr = Expr::ty_abs(kind, expr);
                }
                ContextEntry::Local { var, occ, body } => {
                    if expr.is_trivial() {
                        self.locals_inlined += 1;
                        self.subst.insert(var.id, SubstLocalDefn::Done(expr));
                        return self.simplify(body, in_scope, ctx);
                    } else {
                        let body = self.simplify(
                            body,
                            in_scope.update(var.id, Definition::BoundTo(expr.clone(), occ)),
                            vec![],
                        );
                        // We might have inlined all occurrences of var while simplifying body.
                        // If our binding is now dead, remove it.
                        expr = if let Occurrence::Dead = self.occs.lookup_var(&var) {
                            self.locals_inlined += 1;
                            body
                        } else {
                            let expr_typ = match &expr {
                                Expr::Abstraction(param, _) => {
                                    let mut typ = param.typ.clone();
                                    while let Type::TypAbs(_, inner) = typ {
                                        typ = *inner;
                                    }
                                    typ
                                }
                                _ => expr.type_of(),
                            };
                            let var = var.map_typ(|_| expr_typ);
                            Expr::local(var, expr, body)
                        };
                    }
                }
                ContextEntry::Field { index: idx } => match expr {
                    // Direct field inline
                    Expr::Tuple(mut fields) => {
                        self.field_access_inlined += 1;
                        return self.simplify(fields.swap_remove(idx).1, in_scope, ctx);
                    }
                    // Inline a field access through variable reference.
                    // Hitting this case means that the tuple was not already inlined,
                    // likely because it occurs many times.
                    //
                    // We still want to inline in cases where the field accessed is a polymorphic function,
                    // as it enables direct monomorphization. We could do this more explicitly in
                    // the monomorph pass later but doing so here means we can utilize direct
                    // monomorphization and potentially do futher simplification passses that remove
                    // the tuple entirely. (RowEvidence branch polymorphic functions is a common
                    // case).
                    //
                    // A successful inlining will result in the variable reference count decrementing, we rely on the
                    // next simplification pass to do any further work on this specific variable.
                    Expr::Variable(ref var) => match in_scope.get(&var.id) {
                        Some(Definition::BoundTo(Expr::Tuple(fields), _occ)) => {
                            if matches!(&fields[idx], (_, Expr::TypAbs(_, _))) {
                                self.field_access_inlined += 1;
                                return self.simplify(fields[idx].1.clone(), in_scope, ctx);
                            } else {
                                expr = Expr::field(expr, idx);
                            }
                        }
                        _ => {
                            expr = Expr::field(expr, idx);
                        }
                    },
                    _ => {
                        expr = Expr::field(expr, idx);
                    }
                },
            }
        }
        expr
    }

    fn simplify(&mut self, mut expr: Expr, in_scope: InScope, mut ctx: Context) -> Expr {
        loop {
            trace!(?expr, ?ctx, "simplify iteration");
            expr = match expr {
                Expr::Application(fun, param) => {
                    ctx.push((ContextEntry::App { parameter: *param }, self.subst.clone()));
                    *fun
                }
                Expr::TypApp(ty_fun, ty_app) => {
                    ctx.push((
                        ContextEntry::TypApp {
                            type_parameter: ty_app,
                        },
                        self.subst.clone(),
                    ));
                    *ty_fun
                }
                Expr::Unit => break self.rebuild(Expr::Unit, in_scope, ctx),
                Expr::Integer(i) => break self.rebuild(Expr::Integer(i), in_scope, ctx),
                Expr::Float(f) => break self.rebuild(Expr::Float(f), in_scope, ctx),
                Expr::String(s) => break self.rebuild(Expr::String(s), in_scope, ctx),
                Expr::TypAbs(kind, body) => {
                    ctx.push((ContextEntry::TypAbs { kind }, self.subst.clone()));
                    *body
                }
                Expr::Abstraction(var, body) => {
                    let body =
                        self.simplify(*body, in_scope.update(var.id, Definition::Unknown), vec![]);
                    break self.rebuild(Expr::abs(var, body), in_scope, ctx);
                }
                Expr::Local(var, defn, body) => self.simplify_local(var, *defn, *body, &mut ctx),
                Expr::Variable(var) => match self.simplify_var(var, in_scope.clone(), &ctx) {
                    ControlFlow::Continue(expr) => expr,
                    ControlFlow::Break(var) => {
                        break self.rebuild(Expr::Variable(var), in_scope, ctx);
                    }
                },
                Expr::Tuple(fields) => {
                    let fields = fields
                        .into_iter()
                        .map(|(label, field)| {
                            (label, self.simplify(field, in_scope.clone(), vec![]))
                        })
                        .collect::<Vec<_>>();
                    break self.rebuild(Expr::tuple(fields), in_scope, ctx);
                }
                Expr::Field(tup, index) => {
                    ctx.push((ContextEntry::Field { index }, self.subst.clone()));
                    *tup
                }
                Expr::Tag(typ, tag, body) => {
                    let body = self.simplify(*body, in_scope.clone(), vec![]);
                    break self.rebuild(Expr::tag(typ, tag, body), in_scope, ctx);
                }
                Expr::Case(typ, expr, branches) => {
                    let expr = self.simplify(*expr, in_scope.clone(), vec![]);
                    let branches: Vec<_> = branches
                        .into_iter()
                        .map(|branch| {
                            let id = branch.param.id;
                            Expr::branch(
                                branch.param,
                                self.simplify(
                                    branch.body,
                                    in_scope.update(id, Definition::Unknown),
                                    vec![],
                                ),
                            )
                        })
                        .collect();
                    break self.rebuild(Expr::case(typ, expr, branches), in_scope, ctx);
                }
                Expr::Item(ty, id, symbol) => {
                    break self.rebuild(Expr::Item(ty, id, symbol), in_scope, ctx);
                }
            }
        }
    }

    fn simplify_local(&mut self, var: Var, defn: Expr, body: Expr, ctx: &mut Context) -> Expr {
        match self.occs.lookup_var(&var) {
            // Throwaway dead bindings.
            Occurrence::Dead => {
                self.locals_inlined += 1;
                body
            }
            // Inline and remove locals used once.
            Occurrence::Once => {
                self.locals_inlined += 1;
                let subst = self.subst.clone();
                self.subst
                    .insert(var.id, SubstLocalDefn::Suspend(defn, subst));
                body
            }
            occ => {
                ctx.push((ContextEntry::Local { var, occ, body }, self.subst.clone()));
                defn
            }
        }
    }

    fn simplify_var(
        &mut self,
        var: Var,
        in_scope: InScope,
        ctx: &Context,
    ) -> ControlFlow<Var, Expr> {
        match self.subst.remove(&var.id) {
            Some(SubstLocalDefn::Suspend(payload, subst)) => {
                self.subst = subst;
                ControlFlow::Continue(payload)
            }
            Some(SubstLocalDefn::Done(payload)) => {
                self.subst = Subst::default();
                ControlFlow::Continue(payload)
            }
            None => self.callsite_inline(var, in_scope, ctx),
        }
    }

    fn callsite_inline(
        &mut self,
        var: Var,
        in_scope: InScope,
        ctx: &Context,
    ) -> ControlFlow<Var, Expr> {
        let id = var.id;
        in_scope
            .get(&var.id)
            .map(|bind| match bind {
                Definition::BoundTo(definition, occ)
                    if self.should_inline(definition, *occ, &in_scope, ctx) =>
                {
                    self.subst = Subst::default();
                    // If we inlined a OnceInFun occurrence, mark its binding as dead.
                    // This prevents us from having to run more simplification passes to clean up dead
                    // bindings
                    if let Occurrence::OnceInFun = occ {
                        self.occs.mark_dead(&var.id);
                    }
                    ControlFlow::Continue(definition.clone())
                }
                _ => ControlFlow::Break(var),
            })
            .unwrap_or_else(|| {
                panic!("ICE: Unbound variable encountered in simplification: {id:?}, {ctx:?}")
            })
    }

    fn should_inline(
        &self,
        expr: &Expr,
        occ: Occurrence,
        in_scope: &InScope,
        ctx: &Context,
    ) -> bool {
        match occ {
            Occurrence::Dead | Occurrence::Once => panic!(
                "ICE: should_inline encountered unexpected dead or once occurrence. This should've been handled prior"
            ),
            Occurrence::OnceInFun => expr.is_value() && self.some_benefit(expr, in_scope, ctx),
            Occurrence::Many => {
                let small_enough = expr.size() <= self.inline_size_threshold;
                expr.is_value() && small_enough && self.some_benefit(expr, in_scope, ctx)
            }
        }
    }

    fn did_no_work(&self) -> bool {
        self.saturated_fun_count == 0
            && self.saturated_ty_fun_count == 0
            && self.locals_inlined == 0
            && self.field_access_inlined == 0
    }
}

pub fn simplify_module(module: Module) -> Module {
    module.map(simplify)
}
pub fn simplify(mut expr: Expr) -> Expr {
    for _ in 0..2 {
        let (_, occs) = occurrence_analysis(&expr);
        trace!(?occs, "occurrence_analysis");
        let mut simplifier = Simplifier::new(occs);
        expr = simplifier.simplify(expr, InScope::default(), vec![]);
        trace!(
            simplifier.saturated_fun_count,
            simplifier.saturated_ty_fun_count,
            simplifier.locals_inlined,
            simplifier.field_access_inlined,
            "simplifier work"
        );
        if simplifier.did_no_work() {
            break;
        }
    }
    expr
}

#[cfg(test)]
mod tests {
    use super::*;

    use std::collections::BTreeSet;

    use test_log::test;

    use crate::compiler::mantle::{self, lower_module};
    use crate::compiler::sexpr::{FromSExpr as _, parse_one};

    use super::super::crust::{
        self, ItemSource,
        builder::{ExprBuilder, make_vars},
        type_infer_with_items,
    };

    fn test_simplify(expr: crust::Expr<crust::Var>) -> Expr {
        let out = type_infer_with_items(
            ItemSource::default(),
            crust::Module {
                items: std::collections::BTreeMap::from_iter([(
                    crust::ItemId(0),
                    crust::Item::Native(crust::NativeItem {
                        symbol: crust::Symbol {
                            module: "".to_string(),
                            field: "main".to_string(),
                        },
                        abstraction: expr,
                        typ: crust::TypeScheme {
                            unbound_rows: Default::default(),
                            unbound_tys: BTreeSet::from_iter([
                                crust::TypeVar(0),
                                crust::TypeVar(1),
                            ]),
                            evidence: Default::default(),
                            typ: crust::Type::abstraction(
                                crust::Type::Var(crust::TypeVar(0)),
                                crust::Type::Var(crust::TypeVar(1)),
                            ),
                        },
                    }),
                )]),
            },
        )
        .expect("Type checking failed");
        let module = lower_module(&crust::ItemSource::default(), out);
        match super::simplify_module(module)
            .items
            .into_values()
            .next()
            .unwrap()
        {
            mantle::Item::Native(native_item) => native_item.expr,
            mantle::Item::External(_external_item) => unreachable!(),
        }
    }

    #[test]
    #[ignore]
    fn simple_removes_unused_vars() {
        let b = ExprBuilder::default();
        let [x, y, z] = make_vars();
        let ast = b.fun(x, b.locals([(y, b.int(1)), (z, b.int(2))], b.var(x)));

        let expr = test_simplify(ast);

        let expect = expect_test::expect![[r#"
        (ty_fun [Type]
          (fun [V0]
            V0))"#]];
        expect.assert_debug_eq(&expr);
    }

    #[test]
    fn simple_inline_once() {
        let b = ExprBuilder::default();
        let [u, y, z] = make_vars();
        let ast = b.fun(
            u,
            b.app(b.fun(y, b.app(b.var(y), b.int(157))), b.fun(z, b.var(z))),
        );

        let expr = test_simplify(ast);

        let expect = expect_test::expect![[r#"
            (typ_abs (type) (abs (var (var_id 0) (typ_id 0)) (i 157)))
        "#]];
        expect.assert_debug_eq(&expr);
    }

    #[test]
    #[ignore]
    fn simple_many_trivial_expression_is_inlined() {
        let b = ExprBuilder::default();
        let ast = b.app(
            b.make_funs(|[x, f]| {
                b.app(
                    b.app(b.var(f), b.var(x)),
                    b.apps(b.var(f), [b.var(x), b.var(x)]),
                )
            }),
            b.int(3005),
        );

        let expr = test_simplify(ast);

        let expect = expect_test::expect![[r#"
        (fun [V1]
          (V1 3005 (V1 3005 3005)))"#]];
        expect.assert_debug_eq(&expr);
    }

    #[test]
    #[ignore]
    fn simple_once_nontrivial_expression_is_inlined() {
        let [x, y, f] = make_vars();
        let b = ExprBuilder::default();
        let ast = b.fun(
            f,
            b.app(
                b.fun(x, b.app(b.var(f), b.var(x))),
                b.fun(y, b.app(b.var(y), b.int(3005))),
            ),
        );

        let expr = test_simplify(ast);

        let expect = expect_test::expect![[r#"
        (ty_fun [Type Type]
          (fun [V0]
            (V0 (fun [V2] (V2 3005)))))"#]];
        expect.assert_debug_eq(&expr);
    }

    #[test]
    fn simple_many_nontrivial_expression_is_not_inlined() {
        let b = ExprBuilder::default();
        let [u, x, y, f] = make_vars();
        let ast = b.fun(
            u,
            b.app(
                b.fun(x, b.fun(f, b.app(b.app(b.var(f), b.var(x)), b.var(x)))),
                b.fun(y, b.app(b.var(y), b.int(3005))),
            ),
        );

        let expr = test_simplify(ast);

        let expect = expect_test::expect![[r#"
            (typ_abs
              (type)
              (typ_abs
                (type)
                (typ_abs
                  (type)
                  (abs
                    (var (var_id 0) (typ_id 2))
                    (let
                      (var (var_id 1) (fn (int) (typ_id 1)))
                      (abs
                        (var (var_id 3) (fn (int) (typ_id 1)))
                        (app (var (var_id 3) (fn (int) (typ_id 1))) (i 3005)))
                      (abs
                        (var
                          (var_id 2)
                          (fn
                            (fn (fn (int) (typ_id 1)) (typ_id 1))
                            (fn (fn (fn (int) (typ_id 1)) (typ_id 1)) (typ_id 0))))
                        (app
                          (var
                            (var_id 2)
                            (fn
                              (fn (fn (int) (typ_id 1)) (typ_id 1))
                              (fn (fn (fn (int) (typ_id 1)) (typ_id 1)) (typ_id 0))))
                          (var (var_id 1) (fn (fn (int) (typ_id 1)) (typ_id 1)))
                          (var (var_id 1) (fn (fn (int) (typ_id 1)) (typ_id 1))))))))))
        "#]];
        expect.assert_debug_eq(&expr);
    }

    #[test]
    #[ignore]
    fn simple_onceinfun_uninteresting_context_is_not_inlined() {
        let b = ExprBuilder::default();
        let [x, y, f] = make_vars();
        let ast = b.app(
            b.funs([f, x], b.var(f)),
            b.fun(y, b.app(b.var(y), b.int(3005))),
        );

        let expr = test_simplify(ast);

        let expect = expect_test::expect![[r#"
            (ty_fun [Type Type]
              (let[(V0 (fun [V2] (V2 3005)))] (fun [V1] V0)))"#]];
        expect.assert_debug_eq(&expr);
    }

    #[test]
    #[ignore]
    fn simple_onceinfun_interesting_param_is_inlined() {
        let b = ExprBuilder::default();
        let [x, y, w, f, g, h] = make_vars();
        // An interesting parameter is any expression that isn't trivial.
        // We use a function for that purpose here.
        let interesting_param = b.fun(g, b.var(g));
        let ast = b.app(
            b.funs([f, x], b.app(b.var(f), interesting_param)),
            b.funs([y, w], b.app(b.var(y), b.fun(h, b.var(h)))),
        );

        let expr = test_simplify(ast);

        let expect = expect_test::expect![[r#"
        (ty_fun [Type Type Type]
          (fun [V1, V4, V5]
            V5))"#]];
        expect.assert_debug_eq(&expr);
    }

    #[test]
    #[ignore]
    fn simple_big_expr() {
        let build = ExprBuilder::default();
        let [a, b, c, d, e, f, g, h, i, j, k, l, m] = make_vars();
        let ast = build.locals(
            [
                (
                    a,
                    build.locals(
                        [(
                            b,
                            build.locals(
                                [(
                                    c,
                                    build.locals(
                                        [(d, build.locals([(e, build.int(1))], build.var(e)))],
                                        build.var(d),
                                    ),
                                )],
                                build.var(c),
                            ),
                        )],
                        build.var(b),
                    ),
                ),
                (i, build.int(2)),
                (g, build.int(3)),
                (h, build.int(4)),
                (f, build.funs([j, k, l, m], build.var(j))),
            ],
            build.apps(
                build.var(f),
                [build.var(a), build.var(i), build.var(g), build.var(h)],
            ),
        );

        let expr = test_simplify(ast);

        let expect = expect_test::expect!["1"];
        expect.assert_debug_eq(&expr);
    }

    #[test]
    fn variable_field_inline_typabs() {
        // Tests that a field access through a variable to a typ_abs will be inlined in order to get
        // direct monomorphization.
        //
        // This is essential to making row evidense work for sums types.
        // Row evidence is compiled to a tuple of functions, the branch function is polymorphic.
        // Without this logic we would need more complex logic in the monomorph pass to do inlining
        // there.
        //
        // With the logic as part of the simplify pass we can also find cases where the entire row
        // evidence tuple can be inlined and removed.
        let input = r#"
(let
    (var
        (var_id 0)
        (prd
            (closed
                a
                (typ_abs (type) (fn (typ_id 0) (typ_id 0)))
                b
                (int))))
    (tuple
        a
        (typ_abs (type) (abs (var (var_id 1) (typ_id 0)) (var (var_id 1) (typ_id 0))))
        b
        (i 42))
    (app
        (ty_app
            (field
                (var
                    (var_id 0)
                    (prd
                        (closed
                            a
                            (typ_abs (type) (fn (typ_id 0) (typ_id 0)))
                            b
                            (int))))
                0)
            (ty (int)))
        (field
            (var
                (var_id 0)
                (prd
                    (closed
                        a
                        (typ_abs (type) (fn (typ_id 0) (typ_id 0)))
                        b
                        (int))))
            1)))"#;

        let expr = Expr::from_sexp(&parse_one(input).unwrap(), &mut ()).unwrap();
        let result = simplify(expr);

        let expect = expect_test::expect![[r#"
            (i 42)
        "#]];
        expect.assert_debug_eq(&result);
    }
}