cljrs_runtime/tiered/

lower.rs

//! Rust-native IR lowering orchestration.
//!
//! Calls the Rust `cljrs_ir::lower` pipeline directly (no Clojure interpreter
//! round-trip).  Macro expansion still runs through the interpreter since
//! macros are user-defined Clojure functions.

use std::sync::Arc;

use cljrs_ir::lower::CoreShadows;
use cljrs_ir::{IrFunction, Repr};
use cljrs_reader::Form;
use cljrs_value::TypeHint;

use crate::builtins::form::resolve_auto_forms;
use crate::env::env::{Env, GlobalEnv};

/// Map per-parameter primitive type hints onto representation seeds for type
/// inference, positional with the function's fixed parameters.  `^long`/`^int`
/// → [`Repr::Long`], `^double`/`^float` → [`Repr::Double`], `^boolean` →
/// [`Repr::Bool`].  Array hints and unrecognized/absent hints seed [`Repr::Boxed`]
/// (no scalar unboxing).  Returns an empty vec when no hint is usable, so the
/// caller leaves `IrFunction::seed_reprs` empty.
pub fn seed_reprs_from_hints(param_hints: &[Option<TypeHint>]) -> Vec<Repr> {
    if !param_hints.iter().any(|h| {
        matches!(
            h,
            Some(
                TypeHint::Long
                    | TypeHint::Int
                    | TypeHint::Double
                    | TypeHint::Float
                    | TypeHint::Bool
                    | TypeHint::LongArray
                    | TypeHint::DoubleArray
            )
        )
    }) {
        return Vec::new();
    }
    param_hints
        .iter()
        .map(|h| match h {
            Some(TypeHint::Long | TypeHint::Int) => Repr::Long,
            Some(TypeHint::Double | TypeHint::Float) => Repr::Double,
            Some(TypeHint::Bool) => Repr::Bool,
            // `^longs`/`^doubles` enable unboxed element access via aget/aset.
            Some(TypeHint::LongArray) => Repr::LongArray,
            Some(TypeHint::DoubleArray) => Repr::DoubleArray,
            // Other array hints / absent / non-primitive hints stay boxed (they
            // still work through the boxed array bridge).
            _ => Repr::Boxed,
        })
        .collect()
}

/// What `ns` binds that `clojure.core` also publishes, for the lowerer's
/// call-position core check (see [`cljrs_ir::lower::CoreShadows`]).
///
/// Must be sampled on the mutator thread: it reads the namespace tables, which
/// the background lowering worker may not touch.
pub fn core_shadows_for(globals: &GlobalEnv, ns: &str) -> CoreShadows {
    CoreShadows::new(globals.core_shadowed_names(ns))
}

// ── Error type ──────────────────────────────────────────────────────────────

#[derive(Debug)]
pub enum LowerError {
    /// The Rust lowering function failed.
    LowerFailed(String),
}

impl std::fmt::Display for LowerError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            LowerError::LowerFailed(msg) => write!(f, "lowering failed: {msg}"),
        }
    }
}

// ── Public entry point ──────────────────────────────────────────────────────

/// Lower a function arity's body to IR using the native Rust compiler pipeline.
///
/// `destructure_params` carries the original destructuring patterns for any
/// parameters the interpreter replaced with gensym placeholders (each paired
/// with its index into `params`); `destructure_rest` is the rest parameter's
/// pattern, if it is itself destructured.  Both are expanded into explicit
/// bindings in the IR prologue.
#[allow(clippy::too_many_arguments)]
pub fn lower_arity(
    name: Option<&str>,
    params: &[Arc<str>],
    rest_param: Option<&Arc<str>>,
    destructure_params: &[(usize, Form)],
    destructure_rest: Option<&Form>,
    body: &[Form],
    ns: &Arc<str>,
    env: &mut Env,
    is_async: bool,
) -> Result<IrFunction, LowerError> {
    lower_arity_inner(
        name,
        params,
        rest_param,
        destructure_params,
        destructure_rest,
        body,
        ns,
        env,
        false,
        is_async,
    )
    .map(|(ir, _)| ir)
}

/// Like [`lower_arity`], but also runs the region-optimization pass.
#[allow(clippy::too_many_arguments)]
pub fn lower_and_optimize_arity(
    name: Option<&str>,
    params: &[Arc<str>],
    rest_param: Option<&Arc<str>>,
    destructure_params: &[(usize, Form)],
    destructure_rest: Option<&Form>,
    body: &[Form],
    ns: &Arc<str>,
    env: &mut Env,
    is_async: bool,
) -> Result<IrFunction, LowerError> {
    lower_and_optimize_arity_tracked(
        name,
        params,
        rest_param,
        destructure_params,
        destructure_rest,
        body,
        ns,
        env,
        is_async,
    )
    .map(|(ir, _)| ir)
}

/// Like [`lower_and_optimize_arity`], but also returns the `(ns, name)` set
/// of cross-defn externals (see [`crate::tiered::defn_registry`]) the optimizer
/// consulted — the caller must register those as invalidation dependencies.
#[allow(clippy::too_many_arguments)]
pub fn lower_and_optimize_arity_tracked(
    name: Option<&str>,
    params: &[Arc<str>],
    rest_param: Option<&Arc<str>>,
    destructure_params: &[(usize, Form)],
    destructure_rest: Option<&Form>,
    body: &[Form],
    ns: &Arc<str>,
    env: &mut Env,
    is_async: bool,
) -> Result<(IrFunction, Vec<(Arc<str>, Arc<str>)>), LowerError> {
    lower_arity_inner(
        name,
        params,
        rest_param,
        destructure_params,
        destructure_rest,
        body,
        ns,
        env,
        true,
        is_async,
    )
}

#[allow(clippy::too_many_arguments)]
fn lower_arity_inner(
    name: Option<&str>,
    params: &[Arc<str>],
    rest_param: Option<&Arc<str>>,
    destructure_params: &[(usize, Form)],
    destructure_rest: Option<&Form>,
    body: &[Form],
    ns: &Arc<str>,
    env: &mut Env,
    do_optimize: bool,
    is_async: bool,
) -> Result<(IrFunction, Vec<(Arc<str>, Arc<str>)>), LowerError> {
    // Temporarily set current_ns to the function's defining namespace so that
    // ::kw resolution in macroexpand_body uses the correct namespace.
    let prev_ns = std::mem::replace(&mut env.current_ns, ns.clone());
    let expanded_body = macroexpand_body(body, env);
    // Auto-resolved identifiers are qualified here, the last boundary holding
    // an Env; the lowerer refuses any that survive. An unresolvable one is left
    // in place for the lowerer to refuse, so the tree-walker reports it.
    let resolved: Result<Vec<Form>, _> = expanded_body
        .iter()
        .map(|f| resolve_auto_forms(f, env))
        .collect();
    let expanded_body = resolved.unwrap_or(expanded_body);
    env.current_ns = prev_ns;
    let shadows = core_shadows_for(&env.globals, ns);
    lower_expanded_arity(
        name,
        params,
        rest_param,
        destructure_params,
        destructure_rest,
        &expanded_body,
        ns,
        &shadows,
        env.globals.id(),
        None,
        do_optimize,
        is_async,
    )
}

/// Macro-expand a function body on the calling thread.
///
/// Macros are user-defined Clojure functions, so expansion must run through
/// the interpreter with a live `Env` — this is the only part of lowering that
/// cannot move to the background worker.  Forms that fail to expand are kept
/// unexpanded (lowering will reject them if they matter).
pub fn macroexpand_body(body: &[Form], env: &mut Env) -> Vec<Form> {
    // Guard against re-entrant lowering during macro expansion.
    use crate::tiered::apply::IR_LOWERING_ACTIVE;
    let was_active = IR_LOWERING_ACTIVE.get();
    IR_LOWERING_ACTIVE.set(true);

    let expanded_body: Vec<Form> = body
        .iter()
        .map(|f| crate::interp::macros::macroexpand_all(f, env).unwrap_or_else(|_| f.clone()))
        .collect();

    IR_LOWERING_ACTIVE.with(|c| c.set(was_active));
    expanded_body
}

/// Lower an already macro-expanded arity body to (optionally optimized) IR.
///
/// Env-free and callable from the background lowering worker (Phase 10.7):
/// everything below operates on plain `Form`/IR data.  `shadows` is the one
/// piece that has to be sampled from the environment first — see
/// [`core_shadows_for`].  `globals_id` is
/// `GlobalEnv::id`, scoping the cross-defn registry lookups to one runtime.
///
/// `arity_id` selects the externals protocol:
/// - `Some(id)` (background worker): externals are fetched via
///   `defn_registry::snapshot_externals`, which atomically records the
///   dependent edges while holding the registry lock — required off the
///   mutator thread, where a rebind can interleave with lowering.
/// - `None` (synchronous mutator-thread callers): legacy `externals_for`;
///   the caller records dependents from the returned `used` set, which is
///   race-free on the single mutator thread.
#[allow(clippy::too_many_arguments)]
pub fn lower_expanded_arity(
    name: Option<&str>,
    params: &[Arc<str>],
    rest_param: Option<&Arc<str>>,
    destructure_params: &[(usize, Form)],
    destructure_rest: Option<&Form>,
    expanded_body: &[Form],
    ns: &Arc<str>,
    shadows: &CoreShadows,
    globals_id: u64,
    arity_id: Option<u64>,
    do_optimize: bool,
    is_async: bool,
) -> Result<(IrFunction, Vec<(Arc<str>, Arc<str>)>), LowerError> {
    tracing::debug!(
        target: "lower",
        "lowering {:?}/{:?} optimize? {}",
        ns,
        name,
        do_optimize
    );

    // Build the flat params list (includes rest param as last element if present).
    let mut all_params: Vec<Arc<str>> = params.to_vec();
    if let Some(rest) = rest_param {
        all_params.push(rest.clone());
    }

    // Build the combined destructuring list, indexed into `all_params`.  Fixed
    // params keep their recorded index; the rest param, if destructured, sits at
    // the final position (`params.len()`).
    let mut destructures: Vec<(usize, Form)> = destructure_params.to_vec();
    if let Some(rest_pat) = destructure_rest {
        destructures.push((params.len(), rest_pat.clone()));
    }

    let kwargs_rest_index = destructure_rest
        .filter(|pat| pat.is_kwargs_rest_pattern())
        .map(|_| params.len());
    let ir = cljrs_ir::lower::lower_fn_body_shadowed_kwargs(
        name,
        ns,
        &all_params,
        &destructures,
        kwargs_rest_index,
        expanded_body,
        is_async,
        shadows,
    )
    .map_err(|e| LowerError::LowerFailed(format!("{e:?}")))?;

    if !do_optimize {
        return Ok((ir, Vec::new()));
    }

    // Make previously-lowered defns this function references visible to
    // escape analysis and stage-4 cross-function region promotion (the
    // script/REPL counterpart of AOT's whole-program lowering).
    let referenced = referenced_globals(&ir);
    let externals = match arity_id {
        Some(id) => crate::tiered::defn_registry::snapshot_externals(globals_id, id, &referenced),
        None => crate::tiered::defn_registry::externals_for(globals_id, &referenced),
    };
    let (ir, used) = cljrs_ir::lower::optimize_with_externals(ir, &externals);
    Ok((ir, used.into_iter().collect()))
}

/// Collect every `(ns, name)` pair the IR tree loads as a global — the
/// candidate set of cross-defn externals.
fn referenced_globals(ir: &IrFunction) -> std::collections::HashSet<(Arc<str>, Arc<str>)> {
    use cljrs_ir::Inst;
    let mut out = std::collections::HashSet::new();
    fn walk(f: &IrFunction, out: &mut std::collections::HashSet<(Arc<str>, Arc<str>)>) {
        for block in &f.blocks {
            for inst in block.phis.iter().chain(block.insts.iter()) {
                if let Inst::LoadGlobal(_, ns, name) | Inst::LoadVar(_, ns, name) = inst {
                    out.insert((ns.clone(), name.clone()));
                }
            }
        }
        for sub in &f.subfunctions {
            walk(sub, out);
        }
    }
    walk(ir, &mut out);
    out
}