reifydb-core 0.7.0

// SPDX-License-Identifier: Apache-2.0
// Copyright (c) 2026 ReifyDB

use std::{
	collections::{BTreeMap, BTreeSet, HashMap},
	fmt::Debug,
	hash::Hash,
	marker::PhantomData,
};

use reifydb_value::{Result, reifydb_assertions, value::row_number::RowNumber};
use serde::{Deserialize, Serialize, de::DeserializeOwned};

use crate::{
	encoded::key::{EncodedKey, IntoEncodedKey},
	window::{
		accumulator::WindowAccumulator,
		engine::{
			AccumulatorEvent, EmitKind, LatePolicy, MetaKey, WindowResult, meta_key_for,
			tumbling::TumblingBuckets,
		},
		span::{Slot, WindowSpan},
		state::StateCache,
		store::WindowStore,
	},
};

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(bound(
	serialize = "C: Serialize + Ord, Carry: Serialize",
	deserialize = "C: serde::de::DeserializeOwned + Ord, Carry: serde::de::DeserializeOwned"
))]
struct WindowEntry<C, Carry> {
	row_number: RowNumber,
	span: WindowSpan<C>,
	carry_out: Option<Carry>,
	has_output: bool,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(bound(
	serialize = "C: Serialize + Ord, Carry: Serialize",
	deserialize = "C: serde::de::DeserializeOwned + Ord, Carry: serde::de::DeserializeOwned"
))]
struct CarryMeta<C, Carry> {
	high_water: Option<C>,
	sealed_up_to: Option<C>,
	sealed_carry: Option<Carry>,
	windows: BTreeMap<C, WindowEntry<C, Carry>>,
}

impl<C, Carry> Default for CarryMeta<C, Carry> {
	fn default() -> Self {
		Self {
			high_water: None,
			sealed_up_to: None,
			sealed_carry: None,
			windows: BTreeMap::new(),
		}
	}
}

type MetaLoaded<G, C, Carry> = HashMap<G, CarryMeta<C, Carry>>;
type SlotResolved = Vec<Option<(RowNumber, bool)>>;

pub struct TumblingCarryEngine<G, C: Slot, Accumulator, Carry> {
	accumulators: StateCache<RowNumber, Accumulator>,
	meta: StateCache<MetaKey, CarryMeta<C, Carry>>,
	late_policy: LatePolicy,
	retention: Option<C::Duration>,
	_pd: PhantomData<G>,
}

impl<G, C, Accumulator, Carry> Default for TumblingCarryEngine<G, C, Accumulator, Carry>
where
	G: Clone + Eq + Ord + Hash + Debug + Serialize + DeserializeOwned,
	C: Slot + Hash + Serialize + DeserializeOwned,
	Accumulator: WindowAccumulator,
	Carry: Clone + Debug + Serialize + DeserializeOwned,
	for<'a> &'a G: IntoEncodedKey,
{
	fn default() -> Self {
		Self::new()
	}
}

impl<G, C, Accumulator, Carry> TumblingCarryEngine<G, C, Accumulator, Carry>
where
	G: Clone + Eq + Ord + Hash + Debug + Serialize + DeserializeOwned,
	C: Slot + Hash + Serialize + DeserializeOwned,
	Accumulator: WindowAccumulator,
	Carry: Clone + Debug + Serialize + DeserializeOwned,
	for<'a> &'a G: IntoEncodedKey,
{
	pub fn new() -> Self {
		Self::with_late_policy_and_retention(LatePolicy::Drop, None)
	}

	pub fn with_late_policy_and_retention(late_policy: LatePolicy, retention: Option<C::Duration>) -> Self {
		Self {
			accumulators: StateCache::<RowNumber, Accumulator>::new(8),
			meta: StateCache::<MetaKey, CarryMeta<C, Carry>>::new_internal(64),
			late_policy,
			retention,
			_pd: PhantomData,
		}
	}

	#[allow(clippy::too_many_arguments)]
	pub fn apply<S, K, NA, BO, CF, Output>(
		&mut self,
		store: &mut S,
		buckets: TumblingBuckets<G, C, Accumulator::Contribution>,
		row_key: K,
		new_accumulator: NA,
		build_output: BO,
		carry_forward: CF,
	) -> Result<Vec<WindowResult<G, C, Output>>>
	where
		S: WindowStore,
		K: Fn(&G, C) -> EncodedKey,
		NA: Fn() -> Accumulator,
		BO: Fn(&G, WindowSpan<C>, &Accumulator::Output, Option<&Carry>) -> Option<Output>,
		CF: Fn(&Accumulator::Output, Option<&Carry>) -> Option<Carry>,
	{
		if buckets.is_empty() {
			return Ok(Vec::new());
		}
		let late_policy = self.late_policy;
		let retention = self.retention;
		let mut meta_loaded = self.warm_and_load_meta(store, &buckets)?;
		let slot_resolved = self.resolve_survivor_rows(store, &buckets, &meta_loaded, &row_key)?;

		let mut earliest_affected: HashMap<G, C> = HashMap::new();
		for (((group, span), events), slot_pre) in buckets.into_iter().zip(slot_resolved) {
			let entry = meta_loaded.entry(group.clone()).or_default();
			if matches!(entry.sealed_up_to, Some(s) if span.start <= s) {
				continue;
			}
			let drop_adds = matches!(late_policy, LatePolicy::Drop)
				&& matches!(entry.high_water, Some(hw) if span.start < hw);

			let row_number = match entry.windows.get(&span.start).map(|w| w.row_number) {
				Some(rn) => rn,
				None => match slot_pre {
					Some((rn, _)) => rn,
					None => continue,
				},
			};

			let mut accumulator: Accumulator =
				self.accumulators.get(store, &row_number)?.unwrap_or_else(&new_accumulator);
			let mut changed = false;
			for event in events {
				match event {
					AccumulatorEvent::Add(c) => {
						if drop_adds {
							continue;
						}
						accumulator.add(&c);
						changed = true;
					}
					AccumulatorEvent::Remove(c) => {
						if accumulator.is_empty() {
							continue;
						}
						accumulator.remove(&c);
						changed = true;
					}
				}
			}
			if !changed {
				continue;
			}
			self.accumulators.put(store, &row_number, accumulator)?;

			entry.windows.entry(span.start).or_insert_with(|| WindowEntry {
				row_number,
				span,
				carry_out: None,
				has_output: false,
			});
			if entry.high_water.is_none_or(|hw| span.start > hw) {
				entry.high_water = Some(span.start);
			}

			let e = earliest_affected.entry(group).or_insert(span.start);
			if span.start < *e {
				*e = span.start;
			}
		}

		let mut results: Vec<WindowResult<G, C, Output>> = Vec::new();
		for (group, start) in earliest_affected {
			let meta = meta_loaded.get_mut(&group).expect("affected group has meta");

			let mut prev_carry: Option<Carry> = match meta.windows.range(..start).next_back() {
				Some((_, w)) => w.carry_out.clone(),
				None => meta.sealed_carry.clone(),
			};

			let coords: Vec<C> = meta.windows.range(start..).map(|(c, _)| *c).collect();
			let mut emptied: Vec<C> = Vec::new();
			for coord in coords {
				let (row_number, span, had_output) = {
					let w = meta.windows.get(&coord).expect("window entry present");
					(w.row_number, w.span, w.has_output)
				};
				let value = self.accumulators.get(store, &row_number)?.and_then(|a| a.finalize());
				match value.as_ref().and_then(|v| build_output(&group, span, v, prev_carry.as_ref())) {
					Some(out) => {
						let new_carry = value
							.as_ref()
							.and_then(|v| carry_forward(v, prev_carry.as_ref()));
						let kind = if had_output {
							EmitKind::Update
						} else {
							EmitKind::Insert
						};
						results.push(WindowResult {
							row_number,
							group: group.clone(),
							span,
							value: out,
							prior: None,
							kind,
						});
						let w = meta.windows.get_mut(&coord).expect("window entry present");
						w.carry_out = new_carry.clone();
						w.has_output = true;
						if new_carry.is_some() {
							prev_carry = new_carry;
						}
					}
					None => emptied.push(coord),
				}
			}
			for coord in emptied {
				meta.windows.remove(&coord);
			}

			if let (Some(retention), Some(hw)) = (retention, meta.high_water) {
				loop {
					let Some((&first, w)) = meta.windows.iter().next() else {
						break;
					};
					if hw - first <= retention {
						break;
					}
					let carry_out = w.carry_out.clone();
					let row_number = w.row_number;
					meta.windows.remove(&first);
					meta.sealed_up_to = Some(first);
					meta.sealed_carry = carry_out;
					self.accumulators.remove(store, &row_number)?;
				}
			}
		}

		self.persist_meta(store, meta_loaded)?;
		Ok(results)
	}

	pub fn flush<S: WindowStore>(&mut self, store: &mut S) -> Result<()> {
		self.accumulators.flush(store)?;
		self.meta.flush(store)?;
		Ok(())
	}

	fn warm_and_load_meta<S: WindowStore>(
		&mut self,
		store: &mut S,
		buckets: &TumblingBuckets<G, C, Accumulator::Contribution>,
	) -> Result<MetaLoaded<G, C, Carry>> {
		let meta_keys: Vec<MetaKey> = buckets
			.keys()
			.map(|(group, _)| group)
			.collect::<BTreeSet<_>>()
			.into_iter()
			.map(meta_key_for)
			.collect();
		self.meta.warm(store, &meta_keys)?;

		let mut meta_loaded: MetaLoaded<G, C, Carry> = HashMap::new();
		for (group, _) in buckets.keys() {
			if !meta_loaded.contains_key(group) {
				let m = self.meta.get(store, &meta_key_for(group))?.unwrap_or_default();
				meta_loaded.insert(group.clone(), m);
			}
		}
		Ok(meta_loaded)
	}

	fn resolve_survivor_rows<S, K>(
		&mut self,
		store: &mut S,
		buckets: &TumblingBuckets<G, C, Accumulator::Contribution>,
		meta_loaded: &MetaLoaded<G, C, Carry>,
		row_key: &K,
	) -> Result<SlotResolved>
	where
		S: WindowStore,
		K: Fn(&G, C) -> EncodedKey,
	{
		let mut survivor_keys: Vec<EncodedKey> = Vec::new();
		let mut slot_survives: Vec<bool> = Vec::with_capacity(buckets.len());
		for (group, span) in buckets.keys() {
			let meta = meta_loaded.get(group);
			let initial_high_water = meta.and_then(|m| m.high_water);
			let sealed = matches!(meta.and_then(|m| m.sealed_up_to), Some(s) if span.start <= s);
			let survives = !sealed
				&& (matches!(self.late_policy, LatePolicy::Process)
					|| initial_high_water.is_none_or(|hw| span.start >= hw));
			slot_survives.push(survives);
			if survives {
				survivor_keys.push(row_key(group, span.start));
			}
		}
		let resolved_rows = store.get_or_create_row_numbers(&survivor_keys)?;
		reifydb_assertions! {
			let survivors = survivor_keys.len();
			let resolved = resolved_rows.len();
			assert!(
				resolved == survivors,
				"get_or_create_row_numbers must return exactly one row per survivor key; a short batch would \
				 leave a surviving slot with no resolved row, so the slot_resolved zip below pairs it with None \
				 and apply silently re-creates a fresh row instead of reusing the existing window state, \
				 double-counting it (survivor_keys={survivors}, resolved_rows={resolved})"
			);
		}
		let accumulator_keys: Vec<RowNumber> = resolved_rows.iter().map(|(rn, _)| *rn).collect();
		self.accumulators.warm(store, &accumulator_keys)?;
		let mut resolved_rows = resolved_rows.into_iter();
		Ok(slot_survives
			.into_iter()
			.map(|survives| {
				if survives {
					resolved_rows.next()
				} else {
					None
				}
			})
			.collect())
	}

	fn persist_meta<S: WindowStore>(&mut self, store: &mut S, meta_loaded: MetaLoaded<G, C, Carry>) -> Result<()> {
		for (group, meta) in meta_loaded {
			self.meta.set(store, &meta_key_for(&group), &meta)?;
		}
		Ok(())
	}
}

#[cfg(test)]
mod tests {
	use std::{collections::HashMap, ops::Bound};

	use postcard::{from_bytes, to_allocvec};

	use super::*;
	use crate::{encoded::key::EncodedKeyRange, window::accumulator::invertible::RetainedAccumulator};

	// In-memory store that allocates a distinct row number per key (the state.rs
	// mock collapses every key onto row 1, which would alias all window
	// accumulators and defeat a storage-bound test).
	#[derive(Default)]
	struct CountingStore {
		data: HashMap<Vec<u8>, Vec<u8>>,
		internal: HashMap<Vec<u8>, Vec<u8>>,
		rows: HashMap<Vec<u8>, RowNumber>,
		next_row: u64,
	}

	impl WindowStore for CountingStore {
		fn state_get<V: DeserializeOwned>(&mut self, key: &EncodedKey) -> Result<Option<V>> {
			Ok(self.data.get(key.as_bytes()).map(|b| from_bytes(b).expect("decode")))
		}
		fn state_get_many_visit<V: DeserializeOwned>(
			&mut self,
			keys: &[EncodedKey],
			visit: &mut dyn FnMut(EncodedKey, V) -> Result<()>,
		) -> Result<()> {
			for key in keys {
				if let Some(b) = self.data.get(key.as_bytes()) {
					visit(key.clone(), from_bytes(b).expect("decode"))?;
				}
			}
			Ok(())
		}
		fn state_set<V: Serialize>(&mut self, key: &EncodedKey, value: &V) -> Result<()> {
			self.data.insert(key.as_bytes().to_vec(), to_allocvec(value).expect("encode"));
			Ok(())
		}
		fn state_remove(&mut self, key: &EncodedKey) -> Result<()> {
			self.data.remove(key.as_bytes());
			Ok(())
		}
		fn state_drop(&mut self, key: &EncodedKey) -> Result<()> {
			self.data.remove(key.as_bytes());
			Ok(())
		}
		fn internal_get<V: DeserializeOwned>(&mut self, key: &EncodedKey) -> Result<Option<V>> {
			Ok(self.internal.get(key.as_bytes()).map(|b| from_bytes(b).expect("decode")))
		}
		fn internal_get_many_visit<V: DeserializeOwned>(
			&mut self,
			keys: &[EncodedKey],
			visit: &mut dyn FnMut(EncodedKey, V) -> Result<()>,
		) -> Result<()> {
			for key in keys {
				if let Some(b) = self.internal.get(key.as_bytes()) {
					visit(key.clone(), from_bytes(b).expect("decode"))?;
				}
			}
			Ok(())
		}
		fn internal_set<V: Serialize>(&mut self, key: &EncodedKey, value: &V) -> Result<()> {
			self.internal.insert(key.as_bytes().to_vec(), to_allocvec(value).expect("encode"));
			Ok(())
		}
		fn internal_remove(&mut self, key: &EncodedKey) -> Result<()> {
			self.internal.remove(key.as_bytes());
			Ok(())
		}
		fn internal_drop(&mut self, key: &EncodedKey) -> Result<()> {
			self.internal.remove(key.as_bytes());
			Ok(())
		}
		fn internal_range_visit<V: DeserializeOwned>(
			&mut self,
			range: EncodedKeyRange,
			visit: &mut dyn FnMut(EncodedKey, V) -> Result<()>,
		) -> Result<()> {
			let after_start = |k: &[u8]| match &range.start {
				Bound::Included(s) => k >= s.as_bytes(),
				Bound::Excluded(s) => k > s.as_bytes(),
				Bound::Unbounded => true,
			};
			let before_end = |k: &[u8]| match &range.end {
				Bound::Included(e) => k <= e.as_bytes(),
				Bound::Excluded(e) => k < e.as_bytes(),
				Bound::Unbounded => true,
			};
			let mut matched: Vec<(Vec<u8>, Vec<u8>)> = self
				.internal
				.iter()
				.filter(|(k, _)| after_start(k) && before_end(k))
				.map(|(k, v)| (k.clone(), v.clone()))
				.collect();
			matched.sort_by(|a, b| a.0.cmp(&b.0));
			for (k, b) in matched {
				visit(EncodedKey::new(k), from_bytes(&b).expect("decode"))?;
			}
			Ok(())
		}
		fn get_or_create_row_number(&mut self, key: &EncodedKey) -> Result<(RowNumber, bool)> {
			if let Some(rn) = self.rows.get(key.as_bytes()) {
				return Ok((*rn, false));
			}
			self.next_row += 1;
			let rn = RowNumber(self.next_row);
			self.rows.insert(key.as_bytes().to_vec(), rn);
			Ok((rn, true))
		}
		fn get_or_create_row_numbers(&mut self, keys: &[EncodedKey]) -> Result<Vec<(RowNumber, bool)>> {
			keys.iter().map(|k| self.get_or_create_row_number(k)).collect()
		}
		fn allocate_row_numbers(&mut self, count: u64) -> Result<RowNumber> {
			let start = self.next_row + 1;
			self.next_row += count;
			Ok(RowNumber(start))
		}
		fn clock_now_nanos(&self) -> u64 {
			0
		}
	}

	type Engine = TumblingCarryEngine<String, u64, RetainedAccumulator<u64, f64>, f64>;

	const WINDOW: u64 = 60;

	// Feed one event into window `ws` for group "BTC" as its own batch, so the
	// high-water mark advances one window per call.
	fn feed(engine: &mut Engine, store: &mut CountingStore, ws: u64, price: f64) {
		let mut buckets: TumblingBuckets<String, u64, (u64, f64)> = BTreeMap::new();
		let span = WindowSpan::for_slot(ws, WINDOW);
		buckets.insert(("BTC".to_string(), span), vec![AccumulatorEvent::Add((ws, price))]);
		let _: Vec<WindowResult<String, u64, f64>> = engine
			.apply(
				store,
				buckets,
				|g: &String, w: u64| EncodedKey::builder().str(g).u64(w).build(),
				RetainedAccumulator::<u64, f64>::default,
				|_g: &String, _s: WindowSpan<u64>, v: &BTreeMap<u64, f64>, _p: Option<&f64>| {
					(!v.is_empty()).then(|| v.values().sum::<f64>())
				},
				|v: &BTreeMap<u64, f64>, _p: Option<&f64>| v.last_key_value().map(|(_, val)| *val),
			)
			.expect("apply");
	}

	#[test]
	fn retention_seals_old_windows_and_reclaims_accumulator_rows() {
		// With a 2-window retention horizon, only the windows within `hw - 120`
		// stay live; every older window must seal to the O(1) carry scalar and
		// have its accumulator row removed from the store. After 60 windows the
		// live accumulator-row count must stay bounded by the horizon, not grow
		// with the number of windows seen.
		let mut store = CountingStore::default();
		let mut engine = Engine::with_late_policy_and_retention(LatePolicy::Drop, Some(2 * WINDOW));
		for i in 0..60u64 {
			feed(&mut engine, &mut store, i * WINDOW, i as f64);
		}
		engine.flush(&mut store).expect("flush");
		assert!(
			store.data.len() <= 4,
			"sealed windows must reclaim their accumulator rows; found {} live rows after 60 windows",
			store.data.len()
		);
	}

	#[test]
	fn without_retention_every_window_accumulator_is_retained() {
		// The contrast that proves the bound above is the sealing, not some
		// other cap: with no retention configured the engine keeps every window's
		// accumulator forever (the pre-sealing behavior).
		let mut store = CountingStore::default();
		let mut engine = Engine::new();
		for i in 0..60u64 {
			feed(&mut engine, &mut store, i * WINDOW, i as f64);
		}
		engine.flush(&mut store).expect("flush");
		assert_eq!(
			store.data.len(),
			60,
			"with no retention the carry engine retains every window's accumulator row"
		);
	}
}