kevy-rt 6.4.0

//! zset/set algebra `*STORE` orchestrator, step-1→step-2
//! transition (same two-hop shape as [`crate::exec_rename`]).
//!
//! Step 1 (built by `build_zalgebra_store`): scored/set gathers fan
//! out to the source keys' shards, folding into
//! [`Agg::ZStoreGather`]. When the last gather lands, `fold` routes
//! the agg here: the origin computes the combination via the shared
//! `kevy_store` pure algebra, re-arms the slot as a plain
//! `Agg::SumInt`, and ships `Op::ZStoreResult` / `Op::SetStoreResult`
//! to `dst`'s owning shard (or executes locally). Step 2's
//! `Part::Int(cardinality)` folds through the SumInt into the `:n`
//! reply — no dedicated finalize needed.

use std::collections::HashMap;

/// One source key's scored members in request order.
type ScoredInput = Vec<(Vec<u8>, f64)>;

use crate::Commands;
use crate::message::{Agg, Gathered, Inbound, Op, SmallReply, ZCombine};
use crate::shard::Shard;

impl<C: Commands> Shard<C> {
    pub(crate) fn finalize_zstore_agg(&mut self, conn_id: u64, seq: u64, agg: Agg) {
        let Agg::ZStoreGather { combine, weights, aggregate, dst, keys, got } = agg else {
            return;
        };
        // Rebuild inputs in request-key order; WRONGTYPE on any source
        // aborts with the standard error (Redis behavior).
        let (zset_inputs, wrongtype) = collect_scored(&keys, &got);
        if wrongtype {
            self.fill_zstore_slot(
                conn_id,
                seq,
                b"-WRONGTYPE Operation against a key holding the wrong kind of value\r\n".to_vec(),
            );
            return;
        }
        let dst_shard = self.shard_of(&dst);
        let op = build_store_op(combine, zset_inputs, weights, aggregate, dst);
        self.ship_store_op(conn_id, seq, dst_shard, op);
    }

    /// Step 2 of a `*STORE` orchestrator: re-arm the pending slot so one
    /// `Part::Int` (the stored cardinality) folds into the `:n` reply, then
    /// run the store Op on the DESTINATION's owning shard. Shared with the
    /// geo `*STORE` orchestrator ([`crate::exec_geostore`]) — both compute
    /// elsewhere and materialize here.
    pub(crate) fn ship_store_op(&mut self, conn_id: u64, seq: u64, dst_shard: usize, op: Op) {
        if let Some(c) = self.conns.get_mut(&conn_id) {
            let idx = (seq - c.next_emit) as usize;
            if let Some(slot) = c.pending.get_mut(idx) {
                slot.remaining = 1;
                slot.agg = Agg::SumInt(0);
            }
        }
        if dst_shard == self.id {
            let part = self.exec_op(op);
            self.fold(conn_id, seq, part);
        } else {
            self.send_to(dst_shard, Inbound::Request { origin: self.id, conn: conn_id, seq, op });
        }
    }

    /// Re-arm the slot for a continuation phase and fan the new
    /// argv to every shard (stateless two-phase — see exec.rs fold).
    pub(crate) fn start_extension_phase(&mut self, conn_id: u64, seq: u64, argv: Vec<Vec<u8>>) {
        // One buffer, N refcount bumps — the phase-2 argv is as shared as
        // the phase-1 one, and for the same reason.
        let argv: std::sync::Arc<[Vec<u8>]> = argv.into();
        if let Some(c) = self.conns.get_mut(&conn_id) {
            let idx = (seq - c.next_emit) as usize;
            if let Some(slot) = c.pending.get_mut(idx) {
                slot.remaining = self.nshards as u32;
                slot.agg = Agg::ExtensionGather { argv: argv.clone(), chunks: Vec::new() };
            }
        }
        let targets: Vec<(usize, Op)> =
            (0..self.nshards).map(|s| (s, Op::Extension { argv: argv.clone() })).collect();
        self.dispatch_targets(conn_id, seq, targets);
    }

    /// Complete an extension fan-out slot with the reduced reply.
    pub(crate) fn fill_extension_slot(&mut self, conn_id: u64, seq: u64, reply: Vec<u8>) {
        self.fill_zstore_slot(conn_id, seq, reply);
    }

    /// Complete the pending slot with a pre-encoded reply (parse /
    /// WRONGTYPE aborts).
    pub(crate) fn fill_zstore_slot(&mut self, conn_id: u64, seq: u64, reply: Vec<u8>) {
        if let Some(c) = self.conns.get_mut(&conn_id) {
            let idx = (seq - c.next_emit) as usize;
            if let Some(slot) = c.pending.get_mut(idx) {
                slot.remaining = 1;
                slot.agg = Agg::First(None);
            }
        }
        self.fold(conn_id, seq, crate::message::Part::Reply(SmallReply::from_vec(reply)));
    }
}

/// Inputs in request order: `Scored` as-is, set `Members` at 1.0,
/// missing = empty. Second return = a WRONGTYPE was gathered.
fn collect_scored(keys: &[Vec<u8>], got: &HashMap<Vec<u8>, Gathered>) -> (Vec<ScoredInput>, bool) {
    let mut inputs = Vec::with_capacity(keys.len());
    for k in keys {
        match got.get(k) {
            Some(Gathered::Scored(p)) => inputs.push(p.clone()),
            Some(Gathered::Members(m)) => {
                inputs.push(m.iter().map(|v| (v.clone(), 1.0)).collect());
            }
            Some(Gathered::WrongType) => return (inputs, true),
            _ => inputs.push(Vec::new()),
        }
    }
    (inputs, false)
}

/// Combine the gathered inputs per `combine` and build the step-2 store
/// Op. Extracted verbatim from [`Shard::finalize_zstore_agg`] (single
/// call site, `inline(always)`) purely for the 50-LOC fn rule.
#[inline(always)]
fn build_store_op(
    combine: ZCombine,
    zset_inputs: Vec<Vec<(Vec<u8>, f64)>>,
    weights: Option<Vec<f64>>,
    aggregate: kevy_store::ZAggregate,
    dst: Vec<u8>,
) -> Op {
    match combine {
        ZCombine::ZInter | ZCombine::ZUnion | ZCombine::ZDiff => {
            let pairs = match combine {
                ZCombine::ZInter => kevy_store::zinter(&zset_inputs, weights.as_deref(), aggregate),
                ZCombine::ZUnion => kevy_store::zunion(&zset_inputs, weights.as_deref(), aggregate),
                _ => kevy_store::zdiff(&zset_inputs),
            };
            Op::ZStoreResult { dst, pairs }
        }
        ZCombine::SInter | ZCombine::SUnion | ZCombine::SDiff => {
            // Set forms ride the Members gather: inputs carry score
            // 1.0 placeholders only when sources were sets; strip.
            let sets: Vec<Vec<Vec<u8>>> = zset_inputs
                .into_iter()
                .map(|inp| inp.into_iter().map(|(m, _)| m).collect())
                .collect();
            let members = match combine {
                ZCombine::SInter => crate::reduce::set_intersect(&sets),
                ZCombine::SUnion => crate::reduce::set_union(&sets),
                _ => crate::reduce::set_diff(&sets),
            };
            Op::SetStoreResult { dst, members }
        }
    }
}