molgfx-render 0.3.2

//! Cost-aware deterministic selection for compact ligand pose batches.
//!
//! Planning is `O(batches log batches)` only when scene revisions change.
//! Stable frames return before this module runs. Realtime selection keeps
//! complete candidates under one analytic-instance ceiling and a fixed draw
//! activation ceiling; quality selection is exact.

use std::cmp::Ordering;

/// Maximum beauty-pass analytic instances generated by compact pose batches.
///
/// This power-of-two ceiling leaves enough frame budget for the receptor
/// surface on the reference adapter while preserving complete candidates.
pub(super) const REALTIME_POSE_PRIMITIVES: u64 = 8 * 1_024;
/// Maximum beauty-pass indirect draws generated by compact pose batches.
pub(super) const REALTIME_POSE_DRAW_GROUPS: u32 = 128;
/// Opaque pose instances below this limit are repeated in the shadow pass.
pub(super) const REALTIME_POSE_SHADOW_PRIMITIVES: u64 = 4 * 1_024;
/// Beauty plus the bounded shadow repeat can rasterize at most this many rows.
#[cfg(test)]
pub(super) const REALTIME_POSE_RASTER_PRIMITIVES: u64 =
    REALTIME_POSE_PRIMITIVES + REALTIME_POSE_SHADOW_PRIMITIVES;
/// Every shadow draw is a subset of the bounded beauty draw groups.
#[cfg(test)]
pub(super) const REALTIME_POSE_RASTER_DRAWS: u32 = REALTIME_POSE_DRAW_GROUPS * 2;

const DRAW_ACTIVATION_WEIGHT: u64 = REALTIME_POSE_PRIMITIVES / REALTIME_POSE_DRAW_GROUPS as u64;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(super) struct PoseBatchSample {
    pub(super) plan_index: usize,
    pub(super) stable_key: u64,
    pub(super) opaque_source: u32,
    pub(super) translucent_source: u32,
    pub(super) primitive_cost: u32,
    pub(super) shape_count: u32,
    pub(super) selected_opaque: u32,
    pub(super) selected_translucent: u32,
    allowed_groups: u32,
}

impl PoseBatchSample {
    pub(super) const fn new(
        plan_index: usize,
        stable_key: u64,
        opaque_source: u32,
        translucent_source: u32,
        primitive_cost: u32,
        shape_count: u32,
    ) -> Self {
        Self {
            plan_index,
            stable_key,
            opaque_source,
            translucent_source,
            primitive_cost,
            shape_count,
            selected_opaque: 0,
            selected_translucent: 0,
            allowed_groups: 0,
        }
    }

    #[cfg(test)]
    pub(super) const fn selected(&self) -> u32 {
        self.selected_opaque + self.selected_translucent
    }

    #[cfg(test)]
    pub(super) const fn draw_groups(&self) -> u32 {
        opacity_classes(self.selected_opaque, self.selected_translucent) * self.shape_count
    }

    fn source(&self) -> u64 {
        u64::from(self.opaque_source) + u64::from(self.translucent_source)
    }
}

/// Selects whole candidates with a deterministic throughput-oriented policy.
///
/// Realtime ranks batches by amortized instance and draw-activation cost. This
/// prevents thousands of one-pose batches from consuming all draw slots ahead
/// of dense batches, while cheaper topologies still win equal activation
/// efficiency. Opacity classes split one batch allocation; they never receive
/// independent fairness shares.
pub(super) fn select_pose_batches(samples: &mut [PoseBatchSample], quality: bool) {
    if quality {
        for sample in samples {
            sample.selected_opaque = sample.opaque_source;
            sample.selected_translucent = sample.translucent_source;
        }
        return;
    }

    for sample in samples.iter_mut() {
        sample.selected_opaque = 0;
        sample.selected_translucent = 0;
        sample.allowed_groups = 0;
    }
    samples.sort_unstable_by(compare_efficiency);
    let mut remaining_groups = REALTIME_POSE_DRAW_GROUPS;
    for sample in samples.iter_mut() {
        if sample.source() == 0 {
            continue;
        }
        if sample.primitive_cost == 0 || sample.shape_count == 0 {
            sample.selected_opaque = sample.opaque_source;
            sample.selected_translucent = sample.translucent_source;
            continue;
        }
        if u64::from(sample.primitive_cost) > REALTIME_POSE_PRIMITIVES {
            continue;
        }
        let full_groups =
            opacity_classes(sample.opaque_source, sample.translucent_source) * sample.shape_count;
        sample.allowed_groups = if full_groups <= remaining_groups {
            full_groups
        } else if sample.shape_count <= remaining_groups {
            sample.shape_count
        } else {
            0
        };
        remaining_groups -= sample.allowed_groups;
    }

    samples.sort_unstable_by(compare_candidate_cost);
    let mut remaining_instances = REALTIME_POSE_PRIMITIVES;
    for sample in samples {
        if sample.allowed_groups == 0 || sample.primitive_cost == 0 {
            continue;
        }
        let available = remaining_instances / u64::from(sample.primitive_cost);
        let bounded = available.min(sample.source());
        let candidate_limit = crate::fallback(u32::try_from(bounded), u32::MAX);
        let (opaque, translucent) = select_classes(sample, candidate_limit, sample.allowed_groups);
        let selected = opaque + translucent;
        sample.selected_opaque = opaque;
        sample.selected_translucent = translucent;
        remaining_instances -= u64::from(selected) * u64::from(sample.primitive_cost);
    }
}

pub(super) const fn realtime_shadow_instances(opaque_instances: u64) -> u64 {
    if opaque_instances > 0 && opaque_instances <= REALTIME_POSE_SHADOW_PRIMITIVES {
        opaque_instances
    } else {
        0
    }
}

fn select_classes(
    sample: &PoseBatchSample,
    candidate_limit: u32,
    remaining_groups: u32,
) -> (u32, u32) {
    if candidate_limit == 0 || remaining_groups < sample.shape_count {
        return (0, 0);
    }
    let balanced = proportional_split(sample, candidate_limit);
    if opacity_classes(balanced.0, balanced.1) * sample.shape_count <= remaining_groups {
        return balanced;
    }
    let opaque = sample.opaque_source.min(candidate_limit);
    let translucent = sample.translucent_source.min(candidate_limit);
    if opaque >= translucent {
        (opaque, 0)
    } else {
        (0, translucent)
    }
}

fn proportional_split(sample: &PoseBatchSample, selected: u32) -> (u32, u32) {
    let source = sample.source();
    if selected == 0 || source == 0 {
        return (0, 0);
    }
    let selected_u32 = selected;
    let selected = u64::from(selected_u32);
    let opaque_product = selected * u64::from(sample.opaque_source);
    let translucent_product = selected * u64::from(sample.translucent_source);
    let Ok(mut opaque) = u32::try_from(opaque_product / source) else {
        return (0, 0);
    };
    let Ok(mut translucent) = u32::try_from(translucent_product / source) else {
        return (0, 0);
    };
    if opaque + translucent < selected_u32 {
        let opaque_remainder = opaque_product % source;
        let translucent_remainder = translucent_product % source;
        if opaque_remainder >= translucent_remainder {
            opaque += 1;
        } else {
            translucent += 1;
        }
    }
    (opaque, translucent)
}

fn compare_efficiency(left: &PoseBatchSample, right: &PoseBatchSample) -> Ordering {
    let (left_cost, left_capacity) = amortized_cost(left);
    let (right_cost, right_capacity) = amortized_cost(right);
    (u128::from(left_cost) * u128::from(right_capacity))
        .cmp(&(u128::from(right_cost) * u128::from(left_capacity)))
        .then_with(|| left.primitive_cost.cmp(&right.primitive_cost))
        .then_with(|| right.source().cmp(&left.source()))
        .then_with(|| left.stable_key.cmp(&right.stable_key))
        .then_with(|| left.opaque_source.cmp(&right.opaque_source))
        .then_with(|| left.translucent_source.cmp(&right.translucent_source))
        .then_with(|| left.plan_index.cmp(&right.plan_index))
}

fn compare_candidate_cost(left: &PoseBatchSample, right: &PoseBatchSample) -> Ordering {
    left.primitive_cost
        .cmp(&right.primitive_cost)
        .then_with(|| left.stable_key.cmp(&right.stable_key))
        .then_with(|| left.plan_index.cmp(&right.plan_index))
}

fn amortized_cost(sample: &PoseBatchSample) -> (u64, u64) {
    if sample.source() == 0 {
        return (u64::MAX, 1);
    }
    if sample.primitive_cost == 0 || sample.shape_count == 0 {
        return (0, sample.source());
    }
    let capacity = sample
        .source()
        .min(REALTIME_POSE_PRIMITIVES / u64::from(sample.primitive_cost));
    let selected = crate::fallback(u32::try_from(capacity), u32::MAX);
    let (opaque, translucent) = proportional_split(sample, selected);
    let groups = opacity_classes(opaque, translucent) * sample.shape_count;
    let instance_cost = capacity * u64::from(sample.primitive_cost);
    let activation_cost = u64::from(groups) * DRAW_ACTIVATION_WEIGHT;
    (instance_cost + activation_cost, capacity.max(1))
}

const fn opacity_classes(opaque: u32, translucent: u32) -> u32 {
    let opaque_class = if opaque > 0 { 1 } else { 0 };
    let translucent_class = if translucent > 0 { 1 } else { 0 };
    opaque_class + translucent_class
}

pub(super) fn distributed_source_index(ordinal: u32, source: u32, selected: u32) -> u32 {
    if selected >= source {
        return ordinal;
    }
    if selected <= 1 {
        return source / 2;
    }
    let source_span = source - 1;
    let sample_span = selected - 1;
    let quotient = source_span / sample_span;
    let remainder = source_span % sample_span;
    ordinal * quotient + ordinal * remainder / sample_span
}

#[cfg(test)]
#[path = "ligand_pose_sampling_tests.rs"]
mod tests;