cubical-transport-hott-lean4/native/cubical/src/composition.rs

//! # composition
//!
//! Rust implementation of `vHCompValue`, `vCompAtTerm`, `vCompNAtTerm`
//! — CCHM composition at the value level (Cubical/Eval.lean).
//!
//! `vHCompValue ℓ A φ tube base` — homogeneous composition:
//!   - φ = .top    → `vPApp tube .one`   (tube's 1-endpoint)
//!   - A = .pi     → `.vHCompFun {ℓ_c} codA φ tube base` (Π β later)
//!   - else        → `.vneu (.nhcomp {ℓ} A φ tube base)`
//!
//! `vCompAtTerm ℓ env i A φ u t` — heterogeneous composition (CTerm-input):
//!   - φ = .top            → `eval env (u.substDim i .one)`     (C1)
//!   - φ = .bot            → `eval env (.transp i {ℓ} A .bot t)` (C2)
//!   - A dim-absent from i → hcomp via `.plam i u` as tube      (C3-like)
//!   - A = .pi             → `.vCompFun {ℓ_d ℓ_c} env i domA codA φ u t`
//!   - else                → `.vneu (.ncomp {ℓ} i A φ (eval u) (eval t))`
//!
//! `vCompNAtTerm ℓ env i A clauses t` — multi-clause comp.
//!
//! ## ABI v4 — ULevel
//!
//! Each entry takes the implicit `{ℓ : ULevel}` of the Lean `vCompAt*`
//! function as `l : LeanObj` first.  When constructing CVal/CNeu carrying
//! their own implicit ULevel(s), we forward `l` (or read sub-level slots
//! from a CType.pi's [ℓ_d, ℓ_c, ...] layout).

use crate::lean_runtime::*;
use crate::tags::*;
use crate::value::*;

const LIST_NIL:  u32 = 0;
const LIST_CONS: u32 = 1;
const PROD_MK:   u32 = 0;

// ── vHCompValue ─────────────────────────────────────────────────────────────

/// Homogeneous composition at the value level.
/// Takes borrowed `l`, `a`, `phi`; owned `tube` and `base`.  Returns owned.
pub fn vhcomp_value(
    l: LeanObj,
    a: LeanObj, phi: LeanObj, tube: LeanObjMut, base: LeanObjMut,
) -> LeanObjMut {
    // (1) φ = .top → vPApp tube .one
    if ctor_tag(phi) == FACE_TOP {
        release(base as LeanObj);
        // Construct DimExpr.one.
        let one = lean_box_mut(DIM_ONE as usize);
        // vPApp consumes tube and r.
        return crate::eval::vpapp(tube, one as LeanObj);
    }

    // (2) A = .pi → .vHCompFun {ℓ_c} closure.
    // ABI v4: pi layout [ℓ_d, ℓ_c, var, A, B] — codomain B at field 4.
    if ctor_tag(a) == TY_PI {
        let lc = ctor_field(a, 1);
        let cod_a = ctor_field(a, 4);
        retain(lc); retain(cod_a); retain(phi);
        return mk_vhcompfun(lc, cod_a, phi, tube as LeanObj, base as LeanObj);
    }

    // (3) Stuck — .vneu (.nhcomp {ℓ} A φ tube base)
    retain(l); retain(a); retain(phi);
    let nhcomp = mk_nhcomp(l, a, phi, tube as LeanObj, base as LeanObj);
    mk_vneu(nhcomp as LeanObj)
}

// ── vCompAtTerm ─────────────────────────────────────────────────────────────

/// Heterogeneous composition taking CTerm inputs (so C1 can substitute at
/// the term level).
/// All arguments borrowed; returns owned CVal.
pub fn vcomp_at_term(
    l: LeanObj,
    env: LeanObj, i: LeanObj, a: LeanObj, phi: LeanObj,
    u: LeanObj, t: LeanObj,
) -> LeanObjMut {
    // (1) φ = .top → C1: eval env (u.substDim i .one)
    if ctor_tag(phi) == FACE_TOP {
        let one = lean_box_mut(DIM_ONE as usize);
        let u_at_1 = crate::subst::cterm_subst_dim(i, one as LeanObj, u);
        let result = crate::eval::eval(env, u_at_1 as LeanObj);
        release(u_at_1 as LeanObj);
        return result;
    }

    // (2) φ = .bot → C2: eval env (.transp i {ℓ} A .bot t)
    if ctor_tag(phi) == FACE_BOT {
        // ABI v4: TERM_TRANSP layout [i, ℓ, A, φ, t] (5 fields).
        retain(l); retain(i); retain(a);
        let bot = lean_box_mut(FACE_BOT as usize);
        let transp_term = alloc_ctor(TERM_TRANSP, 5);
        ctor_set_field(transp_term, 0, i);
        ctor_set_field(transp_term, 1, l);
        ctor_set_field(transp_term, 2, a);
        ctor_set_field(transp_term, 3, bot as LeanObj);
        retain(t);
        ctor_set_field(transp_term, 4, t);
        let result = crate::eval::eval(env, transp_term as LeanObj);
        release(transp_term as LeanObj);
        return result;
    }

    // (3) A dim-absent from i → hcomp on A with `.plam i u` as tube
    if crate::dim_absent::ctype_absent(i, a) {
        // Build `.plam i u`.
        retain(i); retain(u);
        let plam_u = alloc_ctor(TERM_PLAM, 2);
        ctor_set_field(plam_u, 0, i);
        ctor_set_field(plam_u, 1, u);
        let tube_val = crate::eval::eval(env, plam_u as LeanObj);
        release(plam_u as LeanObj);
        let base_val = crate::eval::eval(env, t);
        return vhcomp_value(l, a, phi, tube_val, base_val);
    }

    // (4) A = .pi → .vCompFun {ℓ_d ℓ_c} env i domA codA φ u t closure.
    // ABI v4: pi layout [ℓ_d, ℓ_c, var, A, B].
    if ctor_tag(a) == TY_PI {
        let ld = ctor_field(a, 0);
        let lc = ctor_field(a, 1);
        let dom_a = ctor_field(a, 3);
        let cod_a = ctor_field(a, 4);
        retain(ld); retain(lc);
        retain(env); retain(i);
        retain(dom_a); retain(cod_a); retain(phi);
        retain(u); retain(t);
        return mk_vcompfun(ld, lc, env, i, dom_a, cod_a, phi, u, t);
    }

    // (5) Stuck — .vneu (.ncomp {ℓ} i A φ (eval u) (eval t))
    let u_val = crate::eval::eval(env, u);
    let t_val = crate::eval::eval(env, t);
    retain(l); retain(i); retain(a); retain(phi);
    let ncomp = mk_ncomp(l, i, a, phi, u_val as LeanObj, t_val as LeanObj);
    mk_vneu(ncomp as LeanObj)
}

// ── vCompNAtTerm ────────────────────────────────────────────────────────────

/// Find the first clause whose face is `.top`.  Returns the matching body
/// CTerm (borrowed) or null pointer if none found.
fn find_top_clause(clauses: LeanObj) -> LeanObj {
    let mut cur = clauses;
    loop {
        match ctor_tag(cur) {
            LIST_NIL => return core::ptr::null(),
            LIST_CONS => {
                let head = ctor_field(cur, 0);
                let face = ctor_field(head, 0);
                if ctor_tag(face) == FACE_TOP {
                    return ctor_field(head, 1);  // body
                }
                cur = ctor_field(cur, 1);
            }
            _ => return core::ptr::null(),
        }
    }
}

/// Build a new clause list containing only the non-`.bot` clauses of the
/// input.  Fields retained.  Returns the length as well.
fn filter_live_clauses(clauses: LeanObj) -> (LeanObjMut, u32) {
    // Collect non-bot clauses via a simple recursive traversal.
    build_live(clauses)
}

/// Recursive helper; returns `(list, length)`.
fn build_live(cur: LeanObj) -> (LeanObjMut, u32) {
    match ctor_tag(cur) {
        LIST_NIL => (lean_box_mut(LIST_NIL as usize), 0),
        LIST_CONS => {
            let head = ctor_field(cur, 0);
            let face = ctor_field(head, 0);
            let tail = ctor_field(cur, 1);
            let (tail_list, tail_len) = build_live(tail);
            if ctor_tag(face) == FACE_BOT {
                // Drop this clause.
                (tail_list, tail_len)
            } else {
                // Retain head pair fields and rebuild.
                let body = ctor_field(head, 1);
                retain(face); retain(body);
                let new_pair = alloc_ctor(PROD_MK, 2);
                ctor_set_field(new_pair, 0, face);
                ctor_set_field(new_pair, 1, body);
                let new_cons = alloc_ctor(LIST_CONS, 2);
                ctor_set_field(new_cons, 0, new_pair as LeanObj);
                ctor_set_field(new_cons, 1, tail_list as LeanObj);
                (new_cons, tail_len + 1)
            }
        }
        _ => (lean_box_mut(LIST_NIL as usize), 0),
    }
}

/// Evaluate every clause's body CTerm under `env`, producing a new
/// `List (FaceFormula × CVal)`.
fn eval_clause_bodies(env: LeanObj, clauses: LeanObjMut) -> LeanObjMut {
    eval_clauses_rec(env, clauses as LeanObj)
}

fn eval_clauses_rec(env: LeanObj, cur: LeanObj) -> LeanObjMut {
    match ctor_tag(cur) {
        LIST_NIL => lean_box_mut(LIST_NIL as usize),
        LIST_CONS => {
            let head = ctor_field(cur, 0);
            let face = ctor_field(head, 0);
            let body = ctor_field(head, 1);
            let tail = ctor_field(cur, 1);

            retain(face);
            let body_val = crate::eval::eval(env, body);
            let new_pair = alloc_ctor(PROD_MK, 2);
            ctor_set_field(new_pair, 0, face);
            ctor_set_field(new_pair, 1, body_val as LeanObj);

            let new_tail = eval_clauses_rec(env, tail);
            let new_cons = alloc_ctor(LIST_CONS, 2);
            ctor_set_field(new_cons, 0, new_pair as LeanObj);
            ctor_set_field(new_cons, 1, new_tail as LeanObj);
            new_cons
        }
        _ => lean_box_mut(LIST_NIL as usize),
    }
}

/// Multi-clause heterogeneous composition.
pub fn vcompn_at_term(
    l: LeanObj,
    env: LeanObj, i: LeanObj, a: LeanObj,
    clauses: LeanObj, t: LeanObj,
) -> LeanObjMut {
    // (1) Any `.top` clause → fire it at i := 1.
    let top_body = find_top_clause(clauses);
    if !top_body.is_null() {
        let one = lean_box_mut(DIM_ONE as usize);
        let u_at_1 = crate::subst::cterm_subst_dim(i, one as LeanObj, top_body);
        let result = crate::eval::eval(env, u_at_1 as LeanObj);
        release(u_at_1 as LeanObj);
        return result;
    }

    // (2) Strip `.bot` clauses.
    let (live, len) = filter_live_clauses(clauses);

    match len {
        0 => {
            // No clauses — plain transport of base via .transp i {ℓ} A .bot t.
            release(live as LeanObj);
            retain(l); retain(i); retain(a);
            let bot = lean_box_mut(FACE_BOT as usize);
            // ABI v4: TERM_TRANSP has 5 fields [i, ℓ, A, φ, t].
            let transp_term = alloc_ctor(TERM_TRANSP, 5);
            ctor_set_field(transp_term, 0, i);
            ctor_set_field(transp_term, 1, l);
            ctor_set_field(transp_term, 2, a);
            ctor_set_field(transp_term, 3, bot as LeanObj);
            retain(t);
            ctor_set_field(transp_term, 4, t);
            let result = crate::eval::eval(env, transp_term as LeanObj);
            release(transp_term as LeanObj);
            result
        }
        1 => {
            // Single clause — delegate to vCompAtTerm.
            let head = ctor_field(live as LeanObj, 0);
            let face = ctor_field(head, 0);
            let body = ctor_field(head, 1);
            let result = vcomp_at_term(l, env, i, a, face, body, t);
            release(live as LeanObj);
            result
        }
        _ => {
            // Stuck — produce ncompN with evaluated clause bodies.
            let evaled_clauses = eval_clause_bodies(env, live);
            let t_val = crate::eval::eval(env, t);
            retain(l); retain(env); retain(i); retain(a);
            let ncompn = mk_ncompn(
                l, env, i, a, evaled_clauses as LeanObj, t_val as LeanObj);
            mk_vneu(ncompn as LeanObj)
        }
    }
}