hessboost 0.2.2

//! Exact greedy tree construction (XGBoost's `tree_method=exact`, `ColMaker`).
//!
//! For each node we scan every feature's value-sorted entries and evaluate every
//! candidate threshold the way XGBoost's `ColMaker` does: a backward
//! (descending) scan sends missing values left and runs for every feature; a
//! forward (ascending) scan sends them right and runs only for features that
//! actually have missing values and are not constant. Each scan closes with the
//! endpoint candidate that puts every present value on one side and the
//! missing mass on the other. Growth is level-wise (depth-wise): a whole level
//! is scanned per feature pass.
//!
//! Monotone and interaction constraints are honored during split search.

use super::categorical::sweep_categorical;
use super::shared::{
    BuilderConfig, InteractionState, LeafRows, finalize_leaf_values, permits, rayon_available,
    sum_rows, xgb_node_gain,
};
use super::split::{Screen, ScreenBound, SplitScorer};
use super::{
    BELOW_ALL_VALUES, BestSplit, Children, SplitLocation, SplitPos, limit_or_unbounded,
    need_replace, xgb_update,
};
use crate::K_RT_EPS_F32;
use crate::config::TrainingParams;
use crate::data::{DMatrix, FeatureType};
use crate::objective::GradPair;
use crate::tree::SplitRule;
use crate::tree::constraints::{Bounds, MonotoneConstraints};
use crate::tree::gain::{GradStats, RegParams};
use crate::tree::regtree::RegTree;
use crate::tree::reuse::{CategoricalPenalty, ReuseSet};
use crate::tree::sampler::ColumnSampler;
use rayon::prelude::*;
use std::cell::RefCell;

/// Value-sorted column index over a [`DMatrix`], built once and reused across
/// boosting rounds. Within each column, `(row, value)` pairs are sorted by
/// ascending value. Missing entries are omitted (sparsity-aware).
#[derive(Debug, Clone)]
pub struct SortedColumns {
    n_rows: usize,
    col_ptr: Vec<usize>,
    rows: Vec<u32>,
    vals: Vec<f32>,
}

impl SortedColumns {
    /// Build the value-sorted column index from a dataset.
    pub fn from_dmatrix(data: &DMatrix) -> Self {
        let csc = data.to_csc();
        let n_cols = csc.n_cols();
        // `map` keeps the exact size hint, so `collect` allocates once.
        let mut end = 0;
        let col_ptr: Vec<usize> = std::iter::once(0)
            .chain((0..n_cols).map(|c| {
                end += csc.col_len(c);
                end
            }))
            .collect();
        let nnz = col_ptr[n_cols];
        let mut rows = vec![0u32; nnz];
        let mut vals = vec![0f32; nnz];
        // Sort each column's entries by ascending value (NaN cannot appear:
        // missing entries were excluded when building the CSC). Columns are
        // independent, so large inputs sort them in parallel.
        let fill = |(c, (rows, vals)): (usize, (&mut [u32], &mut [f32]))| {
            let (crows, cvals) = csc.column(c);
            let mut order: Vec<usize> = (0..crows.len()).collect();
            // `partial_cmp` keeps `-0.0` and `0.0` in row order (stable).
            order.sort_by(|&a, &b| {
                cvals[a]
                    .partial_cmp(&cvals[b])
                    .unwrap_or(std::cmp::Ordering::Equal)
            });
            for (k, &o) in order.iter().enumerate() {
                rows[k] = crows[o];
                vals[k] = cvals[o];
            }
        };
        let mut columns: Vec<(&mut [u32], &mut [f32])> = Vec::with_capacity(n_cols);
        let (mut rest_rows, mut rest_vals) = (rows.as_mut_slice(), vals.as_mut_slice());
        for c in 0..n_cols {
            let len = col_ptr[c + 1] - col_ptr[c];
            let (r, rr) = rest_rows.split_at_mut(len);
            let (v, rv) = rest_vals.split_at_mut(len);
            columns.push((r, v));
            (rest_rows, rest_vals) = (rr, rv);
        }
        if nnz >= 65_536 && rayon::current_num_threads() > 1 {
            columns.into_par_iter().enumerate().for_each(fill);
        } else {
            columns.into_iter().enumerate().for_each(fill);
        }
        SortedColumns {
            n_rows: csc.n_rows(),
            col_ptr,
            rows,
            vals,
        }
    }

    /// Number of rows in the source matrix.
    #[inline]
    pub fn n_rows(&self) -> usize {
        self.n_rows
    }

    #[inline]
    fn column(&self, f: usize) -> (&[u32], &[f32]) {
        let (s, e) = (self.col_ptr[f], self.col_ptr[f + 1]);
        (&self.rows[s..e], &self.vals[s..e])
    }
}

/// Exact greedy tree builder.
pub struct ExactTreeBuilder<'a> {
    config: BuilderConfig<'a>,
    /// Opt-in reuse penalties (`toad_penalty_*`): the ensemble's used features
    /// and thresholds, extended by every split this builder commits. `None`
    /// on the default path.
    reuse: Option<RefCell<ReuseSet>>,
}

impl<'a> ExactTreeBuilder<'a> {
    /// Create a builder bound to a training configuration.
    pub fn new(params: &'a TrainingParams) -> Self {
        ExactTreeBuilder {
            config: BuilderConfig::new(params),
            reuse: None,
        }
    }

    /// Penalize candidates by the reuse penalties of `set` (the ensemble's
    /// used features and thresholds; `None` keeps the default gain). Splits
    /// the builder commits extend its own copy, so later levels (and later
    /// trees grown by this builder) reuse them for free.
    #[must_use]
    pub(crate) fn with_reuse(mut self, set: Option<&ReuseSet>) -> Self {
        self.reuse = set.cloned().map(RefCell::new);
        self
    }

    /// Grow a single tree.
    ///
    /// * `cols`: value-sorted column index over the *full* dataset.
    /// * `data`: the dataset (for routing rows after a split).
    /// * `gpair`: per-row gradient/Hessian (length = dataset rows).
    /// * `row_subset`: the sampled rows to train this tree on.
    /// * `sampler`: per-tree column sampler. The exact builder draws one subset
    ///   per level (shared across that level's nodes).
    pub fn build(
        &self,
        cols: &SortedColumns,
        data: &DMatrix,
        gpair: &[GradPair],
        row_subset: &[u32],
        sampler: &mut ColumnSampler,
    ) -> RegTree {
        self.build_inner(cols, data, gpair, row_subset, sampler, false)
            .0
    }

    /// [`Self::build`], also returning the sampled rows each leaf holds
    /// (ascending), so the training margins can be updated per leaf instead
    /// of routing every row through the tree again.
    pub(crate) fn build_with_leaf_rows(
        &self,
        cols: &SortedColumns,
        data: &DMatrix,
        gpair: &[GradPair],
        row_subset: &[u32],
        sampler: &mut ColumnSampler,
    ) -> (RegTree, Vec<LeafRows>) {
        self.build_inner(cols, data, gpair, row_subset, sampler, true)
    }

    fn build_inner(
        &self,
        cols: &SortedColumns,
        data: &DMatrix,
        gpair: &[GradPair],
        row_subset: &[u32],
        sampler: &mut ColumnSampler,
        capture_rows: bool,
    ) -> (RegTree, Vec<LeafRows>) {
        let n_rows = cols.n_rows();

        // node_of_row[r] = current node id for row r, or -1 if r is not sampled.
        let mut node_of_row = vec![-1i32; n_rows];
        for &r in row_subset {
            node_of_row[r as usize] = 0;
        }
        let root = sum_rows(gpair, row_subset);

        let mut tree = RegTree::with_root(root.hess as f32);
        let mut node_stats: Vec<GradStats> = vec![root];
        // Per-node monotone weight bounds (default `±∞` when unconstrained).
        let mut node_bounds: Vec<Bounds> = vec![Bounds::default()];
        let mut node_allowed: Vec<Option<InteractionState>> = vec![None];

        let ftypes = data.feature_types();

        let depth_limit = limit_or_unbounded(self.config.params.max_depth);
        let mut active: Vec<usize> = vec![0];
        let mut depth = 0;

        while depth < depth_limit && !active.is_empty() {
            // One column subset for the whole level (bylevel ∘ bynode).
            let feature_subset = sampler.sample(depth);
            let k = active.len();
            // slot_of_node maps an active node id to its dense slot index.
            let mut slot_of_node = vec![usize::MAX; tree.num_nodes()];
            for (slot, &nid) in active.iter().enumerate() {
                slot_of_node[nid] = slot;
            }

            let mut best = vec![BestSplit::none(); k];

            // XGBoost's `root_gain`: the node's own structure score as `f32`.
            let mut root_gain = vec![0f32; k];
            for (slot, &nid) in active.iter().enumerate() {
                root_gain[slot] =
                    xgb_node_gain(node_stats[nid], &self.config.reg, node_bounds[nid]);
            }

            let row_slot: Vec<u32> = node_of_row
                .iter()
                .map(|&nid| match usize::try_from(nid) {
                    Ok(nid) if slot_of_node[nid] != usize::MAX => slot_of_node[nid] as u32,
                    _ => u32::MAX,
                })
                .collect();
            let reuse_guard = self.reuse.as_ref().map(RefCell::borrow);
            let level = Level {
                reg: &self.config.reg,
                cons: &self.config.cons,
                cols,
                ftypes,
                gpair,
                n_rows,
                row_slot: &row_slot,
                node_allowed: &node_allowed,
                node_stats: &node_stats,
                node_bounds: &node_bounds,
                root_gain: &root_gain,
                active: &active,
                reuse: reuse_guard.as_deref(),
            };
            if feature_subset.len() > 1 && n_rows >= PARALLEL_LEVEL_ROWS && rayon_available() {
                // Each task scans a run of consecutive features in order,
                // so later features are screened against the run's best so
                // far; a scan only reads the level state. The runs' winners
                // per node are merged in feature order with XGBoost's tie
                // rule, which is the sequential search's outcome.
                let per_task = feature_subset
                    .len()
                    .div_ceil(rayon::current_num_threads())
                    .max(1);
                let per_run: Vec<Vec<BestSplit>> = feature_subset
                    .par_chunks(per_task)
                    .map(|run| {
                        let mut scratch = Scratch::new(k);
                        let mut local = vec![BestSplit::none(); k];
                        for &f in run {
                            level.scan_feature(f, &mut local, &mut scratch);
                        }
                        local
                    })
                    .collect();
                for local in per_run {
                    for (best, candidate) in best.iter_mut().zip(local) {
                        // A run's entry replaced `BestSplit::none()` only
                        // with a positive loss change.
                        if candidate.loss_chg > 0.0
                            && need_replace(
                                best.loss_chg as f32,
                                best.feature,
                                candidate.loss_chg as f32,
                                candidate.feature,
                            )
                        {
                            *best = candidate;
                        }
                    }
                }
            } else {
                let mut scratch = Scratch::new(k);
                for &f in feature_subset.iter() {
                    level.scan_feature(f, &mut best, &mut scratch);
                }
            }
            drop(reuse_guard);

            let mut next_active = Vec::new();

            for &nid in &active {
                let slot = slot_of_node[nid];
                let b = &best[slot];
                if !b.valid(self.config.params.gamma, self.config.reg.min_child_weight) {
                    continue; // stays a leaf; value finalized below
                }

                // Monotone child bounds derived from the (bounded) child weights.
                let dir = self.config.cons.dir(b.feature as usize);
                let (lb_bounds, rb_bounds) = b.child_bounds(node_bounds[nid], dir);

                // Children carry XGBoost's bounded `f32` weight, so the
                // `leaf_value` field of nodes that later split records the value
                // they had as a leaf at expansion time. Leaves are overwritten
                // by the finalize pass below.
                let rule = match &b.location {
                    SplitLocation::Numeric(pos) => {
                        SplitRule::numeric(b.feature, value_threshold(*pos), b.default_left)
                    }
                    SplitLocation::Categories(categories) => {
                        SplitRule::categorical(b.feature, categories, b.default_left)
                    }
                };
                let (left_id, right_id) = b.expand(&mut tree, nid, rule);
                if let Some(reuse) = &self.reuse {
                    let mut reuse = reuse.borrow_mut();
                    match &b.location {
                        SplitLocation::Numeric(pos) => {
                            reuse.record_numeric(b.feature, value_threshold(*pos));
                        }
                        SplitLocation::Categories(categories) => {
                            reuse.record_categorical(b.feature, categories);
                        }
                    }
                }
                debug_assert_eq!(left_id, node_stats.len());
                node_stats.push(b.left);
                node_stats.push(b.right);
                node_bounds.push(lb_bounds);
                node_bounds.push(rb_bounds);
                let allowed = self
                    .config
                    .next_allowed(node_allowed[nid].as_ref(), b.feature);
                node_allowed.push(allowed.clone());
                node_allowed.push(allowed);
                next_active.push(left_id);
                next_active.push(right_id);
            }

            // Route each sampled row of a node split at this level into its
            // child (rows of every other node sit at a leaf).
            if !next_active.is_empty() {
                let tree = &tree;
                let route = |(r, nid): (usize, &mut i32)| {
                    if *nid < 0 {
                        return;
                    }
                    let node = tree.node(*nid as usize);
                    if node.is_leaf() {
                        return;
                    }
                    let value = data.get(r, node.split_feature as usize);
                    *nid = tree.child(*nid as usize, value) as i32;
                };
                // Rows route independently.
                if n_rows >= PARALLEL_LEVEL_ROWS && rayon_available() {
                    node_of_row
                        .par_iter_mut()
                        .with_min_len(PARALLEL_LEVEL_ROWS)
                        .enumerate()
                        .for_each(route);
                } else {
                    node_of_row.iter_mut().enumerate().for_each(route);
                }
            }

            active = next_active;
            depth += 1;
        }

        // Finalize every leaf's weight (respecting each leaf's monotone bounds).
        finalize_leaf_values(&mut tree, &node_stats, &node_bounds, &self.config.reg);
        if !capture_rows {
            return (tree, Vec::new());
        }
        // Every sampled row now sits at its leaf.
        let mut slot_of_leaf = vec![usize::MAX; tree.num_nodes()];
        let mut leaves: Vec<LeafRows> = Vec::new();
        for (r, &nid) in node_of_row.iter().enumerate() {
            let Ok(nid) = usize::try_from(nid) else {
                continue;
            };
            let slot = &mut slot_of_leaf[nid];
            if *slot == usize::MAX {
                *slot = leaves.len();
                leaves.push(LeafRows {
                    node: nid,
                    rows: Vec::new(),
                });
            }
            leaves[*slot].rows.push(r as u32);
        }
        (tree, leaves)
    }
}

/// The threshold of an exact split position, which search records in value
/// space ([`SplitPos::Value`]).
fn value_threshold(pos: SplitPos) -> f32 {
    match pos {
        SplitPos::Value(threshold) => threshold,
        SplitPos::Bin(_) | SplitPos::BelowBins => {
            unreachable!("exact search records value-space thresholds")
        }
    }
}

/// Datasets with at least this many rows scan a level's features in
/// parallel.
const PARALLEL_LEVEL_ROWS: usize = 4096;

/// One active node's running scan state for the current feature, with the
/// constants its per-row step needs gathered in one place.
#[derive(Debug, Clone, Copy)]
struct SlotScan {
    /// Statistics of the rows scanned so far.
    acc: GradStats,
    /// Value of the last row scanned.
    last_val: f32,
    /// The node's statistics.
    total: GradStats,
    /// The node's [`SplitScorer::screen`], `None` where it may not screen
    /// (no division-free bound, or reuse penalties).
    node_screen: Option<Screen>,
    /// `node_screen` bound to the incumbent while it comes from this or an
    /// earlier feature (a candidate then needs a strictly larger loss change
    /// to replace it), else `None`.
    screen: Option<ScreenBound>,
    /// The node may split on the current feature.
    allowed: bool,
}

impl SlotScan {
    /// Whether `children` certainly cannot replace the incumbent. Only for
    /// candidates whose children both have non-negative Hessians (which the
    /// scan's `min_child_weight` tests establish first).
    #[inline(always)]
    fn rules_out(&self, children: &Children) -> bool {
        self.screen
            .is_some_and(|screen| screen.rules_out(children.left, children.right))
    }

    /// Take the incumbent from `best` for a scan of feature `f`.
    #[inline]
    fn sync(&mut self, best: &BestSplit, f: u32) {
        self.screen = self
            .node_screen
            .filter(|_| best.feature <= f)
            .map(|screen| screen.bound(self.total.hess, best.loss_chg));
    }
}

/// Per-node scan state, reused across features and scan directions.
struct Scratch {
    slots: Vec<SlotScan>,
}

impl Scratch {
    fn new(k: usize) -> Self {
        Scratch {
            slots: Vec::with_capacity(k),
        }
    }
}

/// The read-only state of one level's split search.
struct Level<'a> {
    reg: &'a RegParams,
    cons: &'a MonotoneConstraints,
    cols: &'a SortedColumns,
    ftypes: &'a [FeatureType],
    gpair: &'a [GradPair],
    n_rows: usize,
    /// Each row's slot among the active nodes, or `u32::MAX` (unsampled, or
    /// at a leaf).
    row_slot: &'a [u32],
    node_allowed: &'a [Option<InteractionState>],
    node_stats: &'a [GradStats],
    node_bounds: &'a [Bounds],
    /// XGBoost's `root_gain` of each active slot.
    root_gain: &'a [f32],
    active: &'a [usize],
    reuse: Option<&'a ReuseSet>,
}

impl Level<'_> {
    /// Offer every candidate split of feature `f` to the active nodes' `best`
    /// entries, in `ColMaker` order.
    fn scan_feature(&self, f: u32, best: &mut [BestSplit], scratch: &mut Scratch) {
        let (crows, cvals) = self.cols.column(f as usize);
        let dir = self.cons.dir(f as usize);
        let gpair = self.gpair;
        let reg = self.reg;
        // Scoring context of active node `nid` (dense slot `slot`) for `f`.
        let scorer = |slot: usize, nid: usize| SplitScorer {
            reg,
            root_gain: self.root_gain[slot],
            bounds: self.node_bounds[nid],
            dir,
        };
        // Row `r`'s node and its dense slot, when that node is active
        // at this level and may split on `f`.
        let active_slot = |r: usize| {
            let slot = self.row_slot[r];
            if slot == u32::MAX {
                return None;
            }
            let slot = slot as usize;
            let nid = self.active[slot];
            permits(self.node_allowed[nid].as_ref(), f).then_some((nid, slot))
        };

        // Categorical features use a set-membership split instead of a
        // numeric threshold, over every category of the column in
        // ascending order (the histogram builder's category bins).
        if self.ftypes[f as usize] == FeatureType::Categorical {
            let k = self.active.len();
            // The column is sorted, so equal categories are adjacent.
            let mut categories: Vec<u32> = Vec::new();
            let mut cat_stats: Vec<Vec<GradStats>> = vec![Vec::new(); k];
            for (&rr, &val) in crows.iter().zip(cvals) {
                let cat = val as u32;
                if categories.last() != Some(&cat) {
                    categories.push(cat);
                }
                let Some((_, slot)) = active_slot(rr as usize) else {
                    continue;
                };
                let stats = &mut cat_stats[slot];
                stats.resize(categories.len(), GradStats::default());
                stats[categories.len() - 1].add(GradStats::from_pair(gpair[rr as usize]));
            }
            for (slot, &nid) in self.active.iter().enumerate() {
                if !permits(self.node_allowed[nid].as_ref(), f) {
                    continue;
                }
                let stats = &cat_stats[slot];
                let cats: Vec<(u32, GradStats)> = categories
                    .iter()
                    .enumerate()
                    .map(|(i, &c)| (c, stats.get(i).copied().unwrap_or_default()))
                    .collect();
                sweep_categorical(
                    &mut best[slot],
                    &cats,
                    self.node_stats[nid],
                    &scorer(slot, nid),
                    f,
                    self.reuse.map(|r| r as &dyn CategoricalPenalty),
                );
            }
            return;
        }

        // `NeedForwardSearch`: only a column with missing values that is
        // not constant scans forward (missing right); every column
        // scans backward (missing left).
        let indicator = !cvals.is_empty() && cvals[0] == cvals[cvals.len() - 1];
        let slots = &mut scratch.slots;
        slots.clear();
        slots.extend(self.active.iter().enumerate().map(|(slot, &nid)| {
            let mut s = SlotScan {
                acc: GradStats::default(),
                last_val: 0.0,
                total: self.node_stats[nid],
                node_screen: if self.reuse.is_none() {
                    scorer(slot, nid).screen()
                } else {
                    None
                },
                screen: None,
                allowed: permits(self.node_allowed[nid].as_ref(), f),
            };
            s.sync(&best[slot], f);
            s
        }));
        let (row_slot, mcw) = (self.row_slot, reg.min_child_weight);
        let mut scan = |d_step: i8, slots: &mut [SlotScan]| {
            // Loop invariants copied into the closure's own locals: read
            // through its captures, they were reloaded past every store of
            // the per-row step instead of staying in registers.
            let (row_slot, mcw, gpair) = (row_slot, mcw, gpair);
            for s in slots.iter_mut() {
                s.acc = GradStats::default();
            }
            // The per-row step on slot `$slot`'s state `$s`.
            macro_rules! visit {
                ($s:expr, $slot:expr, $r:expr, $val:expr, $step:expr) => {{
                    let (s, slot, r, val, d_step): (&mut SlotScan, usize, usize, f32, i8) =
                        (&mut $s, $slot, $r, $val, $step);
                    if s.allowed {
                        let e = s.acc;
                        // `UpdateEnumeration`: the first rows with positive
                        // Hessian only seed the running statistics.
                        if e.hess != 0.0 && val != s.last_val && e.hess >= mcw {
                            let c = s.total.sub(e);
                            if c.hess >= mcw {
                                let children = if d_step < 0 {
                                    Children::new(true, c, e)
                                } else {
                                    Children::new(false, e, c)
                                };
                                // ColMaker's midpoint `(fvalue + last) * 0.5f`
                                // overflows to `±inf` for two same-sign
                                // values near `±f32::MAX`. Only then fall
                                // back to the halved form, which is finite
                                // and still lies between the two values, so
                                // the partition is unchanged. Trees must
                                // stay finite.
                                let last = s.last_val;
                                let thr = || {
                                    let mut mid = f32::midpoint(val, last);
                                    if !mid.is_finite() {
                                        mid = val * 0.5 + last * 0.5;
                                    }
                                    if mid == val { last } else { mid }
                                };
                                if !s.rules_out(&children) {
                                    let nid = self.active[slot];
                                    self.offer(
                                        &mut best[slot],
                                        &scorer(slot, nid),
                                        f,
                                        thr,
                                        children,
                                    );
                                    s.sync(&best[slot], f);
                                }
                            }
                        }
                        s.acc.add(GradStats::from_pair(gpair[r]));
                        s.last_val = val;
                    }
                }};
            }
            // The column's `(row, value)` pairs in scan order.
            macro_rules! each_row {
                (|$step:ident, $r:ident, $val:ident| $body:block) => {
                    // Each direction's loop sees its step as a constant.
                    if d_step > 0 {
                        let $step: i8 = 1;
                        for (&$r, &$val) in crows.iter().zip(cvals) $body
                    } else {
                        let $step: i8 = -1;
                        for (&$r, &$val) in crows.iter().zip(cvals).rev() $body
                    }
                };
            }
            if let [only] = slots {
                // A single active node (the root level): its state lives in
                // registers instead of a store-to-load chain through memory.
                let mut state = *only;
                each_row!(|step, r, val| {
                    let r = r as usize;
                    if row_slot[r] != u32::MAX {
                        visit!(state, 0, r, val, step);
                    }
                });
                *only = state;
            } else {
                each_row!(|step, r, val| {
                    let r = r as usize;
                    let slot = row_slot[r];
                    if slot != u32::MAX {
                        let slot = slot as usize;
                        visit!(slots[slot], slot, r, val, step);
                    }
                });
            }
            // Endpoint: every present value on the scanned side, the
            // missing mass on the other.
            for (slot, &nid) in self.active.iter().enumerate() {
                let s = &mut slots[slot];
                let e = s.acc;
                let c = s.total.sub(e);
                if e.hess >= reg.min_child_weight && c.hess >= reg.min_child_weight {
                    let last = s.last_val;
                    let gap = last.abs() + K_RT_EPS_F32;
                    let thr = if d_step > 0 { last + gap } else { last - gap };
                    // ColMaker's `last_fvalue ± delta` overflows to `±inf`
                    // for `|last|` near `f32::MAX`; the tree must stay
                    // finite. Backward (missing left) needs every present
                    // `v >= thr`, which `BELOW_ALL_VALUES` satisfies.
                    // Forward (missing right) needs `v < thr`: `f32::MAX`
                    // works unless `last` is itself `f32::MAX`, in which
                    // case no finite threshold represents the partition
                    // and the candidate is skipped.
                    let thr = if thr.is_finite() {
                        thr
                    } else if d_step < 0 {
                        BELOW_ALL_VALUES
                    } else if last < f32::MAX {
                        f32::MAX
                    } else {
                        continue;
                    };
                    let children = if d_step < 0 {
                        Children::new(true, c, e)
                    } else {
                        Children::new(false, e, c)
                    };
                    if !s.rules_out(&children) {
                        self.offer(&mut best[slot], &scorer(slot, nid), f, || thr, children);
                        s.sync(&best[slot], f);
                    }
                }
            }
        };
        if cvals.len() < self.n_rows && !indicator {
            scan(1, slots);
        }
        scan(-1, slots);
    }

    /// Evaluate one candidate partition exactly as XGBoost's `ColMaker` does
    /// (`CalcSplitGain − root_gain` in `f32`, `SplitEntry::Update` tie rule)
    /// and record it in `best` when it wins. The scan skips candidates that
    /// [`SlotScan::rules_out`] before offering them.
    #[inline]
    fn offer(
        &self,
        best: &mut BestSplit,
        scorer: &SplitScorer,
        feature: u32,
        threshold: impl FnOnce() -> f32,
        children: Children,
    ) {
        if let Some(mut score) = scorer.loss_chg(children.left, children.right) {
            let threshold = threshold();
            if let Some(reuse) = self.reuse {
                score.loss_chg -= reuse.numeric_penalty(feature, threshold);
            }
            xgb_update(best, feature, SplitPos::Value(threshold), children, score);
        }
    }
}

/// Utility: the full row index `0..n_rows` as `u32` (no subsampling).
pub fn all_rows(n_rows: usize) -> Vec<u32> {
    (0..n_rows as u32).collect()
}

#[cfg(test)]
mod tests {
    use super::super::test_support::{gp, grow_exact, monotone_v_shape_data, non_decreasing};
    use super::*;
    use crate::config::TrainingParams;
    use crate::model::Iterations;
    use crate::model::ModelFormat;
    use crate::objective::{Objective, RegLoss};

    /// A clean separable problem: feature 0 perfectly separates the sign of the
    /// gradient at threshold 0.5, so the root should split there.
    #[test]
    fn splits_on_separating_feature() {
        // 4 rows, 1 feature. values 0,0,1,1. gradients push low->+, high->-.
        let x = vec![0.0f32, 0.0, 1.0, 1.0];
        let data = DMatrix::from_dense(&x, 4, 1).unwrap();
        // squared-error-like gradients: left group wants negative weight, right positive
        let gpair = vec![gp(1.0, 1.0), gp(1.0, 1.0), gp(-1.0, 1.0), gp(-1.0, 1.0)];

        let params = TrainingParams::builder()
            .max_depth(1)
            .lambda(0.0)
            .min_child_weight(0.0)
            .gamma(0.0)
            .build()
            .unwrap();
        let tree = grow_exact(&params, &data, &gpair);

        assert_eq!(
            tree.num_nodes(),
            3,
            "root should have split into two leaves"
        );
        let root = tree.node(0);
        assert_eq!(root.split_feature, 0);
        assert!((root.split_cond - 0.5).abs() < 1e-6);
        // left leaf: G=2,H=2 -> w=-1 ; right leaf: G=-2,H=2 -> w=+1
        assert!((tree.predict_row(&data, 0) - (-1.0)).abs() < 1e-6);
        assert!((tree.predict_row(&data, 2) - 1.0).abs() < 1e-6);
    }

    #[test]
    fn no_split_when_gain_below_gamma() {
        let x = vec![0.0f32, 1.0];
        let data = DMatrix::from_dense(&x, 2, 1).unwrap();
        let gpair = vec![gp(1.0, 1.0), gp(-1.0, 1.0)];
        let params = TrainingParams::builder()
            .max_depth(3)
            .gamma(1e9) // impossibly high min split loss
            .build()
            .unwrap();
        let tree = grow_exact(&params, &data, &gpair);
        assert_eq!(tree.num_nodes(), 1, "no split should be taken");
    }

    #[test]
    fn missing_values_pick_a_direction() {
        // 3 rows, feature 0 missing for row 2. Non-missing rows separate cleanly.
        let x = vec![0.0f32, 1.0, f32::NAN];
        let data = DMatrix::from_dense(&x, 3, 1).unwrap();
        // row2 (missing) shares the sign of the high group.
        let gpair = vec![gp(1.0, 1.0), gp(-1.0, 1.0), gp(-1.0, 1.0)];
        let params = TrainingParams::builder()
            .max_depth(1)
            .lambda(0.0)
            .min_child_weight(0.0)
            .gamma(0.0)
            .build()
            .unwrap();
        let tree = grow_exact(&params, &data, &gpair);
        assert_eq!(tree.num_nodes(), 3);
        // The missing row should be routed with the negative-gradient group
        // (right, positive weight). default_left should therefore be false.
        assert!(!tree.node(0).default_left);
        assert!(tree.predict_row(&data, 2) > 0.0);
    }

    #[test]
    fn monotone_increasing_is_enforced() {
        use crate::config::Monotone;
        let (data, gpair) = monotone_v_shape_data();
        let params = TrainingParams::builder()
            .max_depth(4)
            .min_child_weight(0.0)
            .gamma(0.0)
            .lambda(1.0)
            .monotone_constraints(vec![Monotone::Increasing])
            .build()
            .unwrap();
        let tree = grow_exact(&params, &data, &gpair);
        // Predictions must be non-decreasing in x under the increasing constraint.
        assert!(non_decreasing(&tree, &data), "monotonicity violated");

        // Sanity: the unconstrained fit on the same data is *not* monotone, so
        // the constraint is doing real work.
        let unconstrained = TrainingParams::builder()
            .max_depth(4)
            .min_child_weight(0.0)
            .gamma(0.0)
            .lambda(1.0)
            .build()
            .unwrap();
        let free = grow_exact(&unconstrained, &data, &gpair);
        assert!(
            !non_decreasing(&free, &data),
            "unconstrained fit should be non-monotone"
        );
    }

    #[test]
    fn categorical_splits_on_non_ordinal_pattern() {
        use crate::data::FeatureType;
        // 4 categories with a NON-ordinal target: {0,2} vs {1,3}. A numeric
        // threshold cannot separate them; a set-membership split can.
        let mut x = Vec::new();
        let mut gpair = Vec::new();
        for _ in 0..10 {
            for c in 0u32..4 {
                x.push(c as f32);
                // Residual around 0.5: even cats want negative weight, odd positive.
                let g = if c % 2 == 0 { 0.5 } else { -0.5 };
                gpair.push(gp(g, 1.0));
            }
        }
        let n = x.len();
        let data = DMatrix::from_dense(&x, n, 1)
            .unwrap()
            .with_feature_types(&[FeatureType::Categorical])
            .unwrap();
        let params = TrainingParams::builder()
            .max_depth(1)
            .min_child_weight(0.0)
            .gamma(0.0)
            .lambda(1.0)
            .build()
            .unwrap();
        let tree = grow_exact(&params, &data, &gpair);

        assert_eq!(tree.num_nodes(), 3, "root should split");
        assert!(tree.node(0).is_categorical, "split should be categorical");

        // Prediction for a bare category value.
        let pred = |c: f32| tree.leaf_id_with(|_| Some(c));
        // Even categories share a leaf; odd categories share the other leaf.
        assert_eq!(pred(0.0), pred(2.0));
        assert_eq!(pred(1.0), pred(3.0));
        assert_ne!(pred(0.0), pred(1.0), "the two groups must be separated");

        // Even cats (positive grad) want negative weight; odd cats positive.
        let val = |c: f32| tree.node(pred(c)).leaf_value;
        assert!(val(0.0) < 0.0 && val(2.0) < 0.0);
        assert!(val(1.0) > 0.0 && val(3.0) > 0.0);
    }

    /// Train a small exact regressor on one feature and check that every split
    /// threshold is finite, prediction works (debug builds assert finiteness
    /// in the compact forest), and the model survives a native JSON round trip.
    /// Returns the predictions.
    fn train_exact_finite(x: &[f32], y: &[f32]) -> Vec<f32> {
        use crate::config::TreeMethod;
        use crate::{model::BoostedModel, training::train};
        let n = x.len();
        let data = crate::test_support::labeled_dense(x, n, 1, y);
        let params = TrainingParams::builder()
            .objective(Objective::SquaredError(RegLoss::default()))
            .tree_method(TreeMethod::Exact)
            .max_depth(2)
            .eta(0.5)
            .build()
            .unwrap();
        let model = train(&params, &data, 3).unwrap();
        let mut n_splits = 0;
        for tree in model.trees() {
            for node in tree.nodes().iter().filter(|n| !n.is_leaf()) {
                n_splits += 1;
                assert!(
                    node.split_cond.is_finite(),
                    "non-finite split_cond {}",
                    node.split_cond
                );
            }
        }
        assert!(n_splits > 0, "the model should have split");
        let pred = model.predict(&data, Iterations::Best).unwrap();
        let back =
            BoostedModel::decode(model.encode(ModelFormat::Json).unwrap(), ModelFormat::Json)
                .unwrap();
        assert_eq!(back.predict(&data, Iterations::Best).unwrap(), pred);
        pred.into_vec()
    }

    #[test]
    fn backward_endpoint_near_neg_max_stays_finite() {
        // Constant `-f32::MAX` column plus missing rows: only the backward
        // endpoint separates them, and ColMaker's `last - (|last| + eps)`
        // is `-inf` there. The fallback `f32::MIN` keeps every present value
        // on the right (`v >= thr`) with missing on the left.
        let x = [
            -f32::MAX,
            -f32::MAX,
            -f32::MAX,
            -f32::MAX,
            f32::NAN,
            f32::NAN,
        ];
        let y = [0.0, 0.0, 0.0, 0.0, 10.0, 10.0];
        let pred = train_exact_finite(&x, &y);
        assert!(
            pred[0] < pred[4],
            "present rows must be separated from missing"
        );
        assert_eq!(pred[0], pred[3]);
        assert_eq!(pred[4], pred[5]);
    }

    #[test]
    fn forward_endpoint_near_pos_max_stays_finite() {
        // `3e38` rows, zeros and missing rows: the forward scan's endpoint
        // (`last + (|last| + eps)`) overflows to `+inf`; the fallback
        // `f32::MAX` still sends every present value left (`v < thr`).
        let x = [3e38f32, 3e38, 0.0, 0.0, f32::NAN, f32::NAN];
        let y = [0.0, 0.0, 0.0, 0.0, 10.0, 10.0];
        let pred = train_exact_finite(&x, &y);
        assert!(
            pred[0] < pred[4],
            "present rows must be separated from missing"
        );
        assert_eq!(pred[0], pred[2]);
        assert_eq!(pred[4], pred[5]);
    }

    #[test]
    fn midpoint_near_pos_max_stays_finite() {
        // Two same-sign values near `f32::MAX`: `(2e38 + 3e38) * 0.5` is
        // `+inf`; the halved form `2.5e38` lies between them.
        let x = [2e38f32, 2e38, 3e38, 3e38];
        let y = [0.0, 0.0, 10.0, 10.0];
        let pred = train_exact_finite(&x, &y);
        assert!(pred[0] < pred[2], "the two value groups must be separated");
        assert_eq!(pred[0], pred[1]);
        assert_eq!(pred[2], pred[3]);
    }
}