surrealdb-core 3.2.3

//! This module applies index mutations for a single document across different
//! index types (UNIQUE, regular, search, fulltext, Hnsw). Index keys are
//! constructed via key::index and field values are encoded using
//! key::value::Array.
//!
//! Numeric normalization in keys:
//! - Array normalizes Number values (Int/Float/Decimal) through a lexicographic numeric encoding so
//!   that byte order mirrors numeric order.
//! - Numerically equal values (e.g., 0, 0.0, 0dec) map to the same key bytes. On UNIQUE indexes,
//!   such inserts collide and produce a uniqueness error.
//!
//! Planner/executor simplification:
//! - Numeric predicates need a single probe/range in the index; per-variant fan-out is no longer
//!   required.

use std::path::PathBuf;

use anyhow::{Result, bail};
use reblessive::tree::Stk;
use surrealdb_types::ToSql;
use uuid::Uuid;

use crate::catalog::providers::TableProvider;
use crate::catalog::{
	DatabaseId, DiskAnnParams, FullTextParams, HnswParams, Index, IndexDefinition, NamespaceId,
	TableId,
};
use crate::ctx::FrozenContext;
use crate::dbs::Options;
use crate::err::Error;
use crate::expr::{Cond, Part};
use crate::idx::IndexKeyBase;
use crate::idx::ft::fulltext::{FullTextCompactionPlan, FullTextIndex};
use crate::idx::planner::iterators::{IndexCountCompactionPlan, IndexCountThingIterator};
#[cfg(diskann)]
use crate::idx::trees::diskann::index::{DiskAnnCompactionPlan, DiskAnnIndex};
use crate::idx::trees::hnsw::index::{HnswCompactionPlan, HnswIndex};
use crate::idx::trees::store::IndexStores;
use crate::key;
use crate::key::index::iu::IndexCountKey;
use crate::kvs::Transaction;
use crate::val::{Array, RecordId, Value};

pub(crate) struct IndexOperation<'a> {
	ctx: &'a FrozenContext,
	opt: &'a Options,
	ns: NamespaceId,
	db: DatabaseId,
	tb: TableId,
	ix: &'a IndexDefinition,
	ikb: IndexKeyBase,
	/// The old values (if existing)
	o: Option<Vec<Value>>,
	/// The new values (if existing)
	n: Option<Vec<Value>>,
	rid: &'a RecordId,
	/// For COUNT indexes with a WHERE condition: pre-evaluated condition results.
	/// `(old_doc_matches, new_doc_matches)` — whether the old/new document
	/// satisfies the COUNT index condition. `None` for non-COUNT indexes.
	count_cond_match: Option<(bool, bool)>,
}

impl<'a> IndexOperation<'a> {
	#[expect(clippy::too_many_arguments)]
	pub(crate) fn new(
		ctx: &'a FrozenContext,
		opt: &'a Options,
		ns: NamespaceId,
		db: DatabaseId,
		tb: TableId,
		ix: &'a IndexDefinition,
		o: Option<Vec<Value>>,
		n: Option<Vec<Value>>,
		rid: &'a RecordId,
	) -> Self {
		Self {
			ctx,
			opt,
			ns,
			db,
			tb,
			ix,
			ikb: IndexKeyBase::new(ns, db, ix.table_name.clone(), ix.index_id),
			o,
			n,
			rid,
			count_cond_match: None,
		}
	}

	pub(crate) fn with_count_cond_match(mut self, old_matches: bool, new_matches: bool) -> Self {
		self.count_cond_match = Some((old_matches, new_matches));
		self
	}

	pub(crate) async fn create_fulltext_index(
		ctx: &FrozenContext,
		ns: NamespaceId,
		db: DatabaseId,
		ix: &IndexDefinition,
	) -> Result<Option<FullTextIndex>> {
		let Index::FullText(p) = &ix.index else {
			return Ok(None);
		};
		let ikb = IndexKeyBase::new(ns, db, ix.table_name.clone(), ix.index_id);
		Ok(Some(
			FullTextIndex::new(
				ctx.get_index_stores(),
				&ctx.tx(),
				ikb,
				p,
				&ctx.config.file_allowlist,
			)
			.await?,
		))
	}

	pub(crate) async fn compute(
		&mut self,
		stk: &mut Stk,
		require_compaction: &mut bool,
	) -> Result<()> {
		// Index operation dispatching
		match &self.ix.index {
			Index::Uniq => self.index_unique().await,
			Index::Idx => self.index_non_unique().await,
			Index::FullText(p) => self.index_fulltext(stk, p, require_compaction).await,
			Index::Hnsw(p) => self.index_hnsw(p, require_compaction).await,
			Index::DiskAnn(p) => self.index_diskann(p, require_compaction).await,
			Index::Count(c) => self.index_count(stk, c.as_ref(), require_compaction).await,
		}
	}

	/// Build the KV key for a unique index. The Array encodes values in
	/// a canonical, lexicographically ordered byte form which normalizes numeric
	/// types (Int/Float/Decimal). This means equal numeric values like 0, 0.0 and
	/// 0dec map to the same index key and therefore conflict on UNIQUE indexes.
	fn get_unique_index_key(&self, v: &'a Array) -> Result<key::index::Index<'_>> {
		Ok(key::index::Index::new(self.ns, self.db, &self.ix.table_name, self.ix.index_id, v, None))
	}

	/// Build the KV key for a non-unique index. The record id is appended
	/// to the encoded field values so multiple records can share the same field
	/// bytes; numeric values inside fd are normalized via Array.
	fn get_non_unique_index_key(&self, v: &'a Array) -> Result<key::index::Index<'_>> {
		Ok(key::index::Index::new(
			self.ns,
			self.db,
			&self.ix.table_name,
			self.ix.index_id,
			v,
			Some(&self.rid.key),
		))
	}

	async fn index_unique(&mut self) -> Result<()> {
		let txn = self.ctx.tx();
		// Delete the old index data
		if let Some(o) = self.o.take() {
			let i = Indexable::new(o, self.ix);
			for o in i {
				if o.is_any_none_or_null() {
					// NONE/NULL tuples use the non-unique key format (with
					// record ID suffix) so multiple such entries can coexist.
					let key = self.get_non_unique_index_key(&o)?;
					match txn.delc(&key, Some(self.rid)).await {
						Err(e)
							if matches!(
								e.downcast_ref::<Error>(),
								Some(Error::Kvs(crate::kvs::Error::TransactionConditionNotMet))
							) => {}
						Err(e) => return Err(e),
						Ok(()) => {}
					}
				} else {
					let key = self.get_unique_index_key(&o)?;
					match txn.delc(&key, Some(self.rid)).await {
						Err(e)
							if matches!(
								e.downcast_ref::<Error>(),
								Some(Error::Kvs(crate::kvs::Error::TransactionConditionNotMet))
							) => {}
						Err(e) => return Err(e),
						Ok(()) => {}
					}
				}
			}
		}
		// Create the new index data
		if let Some(n) = self.n.take() {
			let i = Indexable::new(n, self.ix);
			for n in i {
				if n.is_any_none_or_null() {
					// NONE/NULL tuples are stored with the non-unique key
					// format so they remain visible to index scans. No
					// uniqueness check — NULL != NULL per SQL convention.
					let key = self.get_non_unique_index_key(&n)?;
					txn.set(&key, self.rid).await?;
				} else {
					let key = self.get_unique_index_key(&n)?;
					if txn.putc(&key, self.rid, None).await.is_err() {
						let key = self.get_unique_index_key(&n)?;
						let rid: RecordId =
							txn.get(&key, None).await?.expect("record should exist");
						return self.err_index_exists(rid, n);
					}
				}
			}
		}
		Ok(())
	}

	async fn index_non_unique(&mut self) -> Result<()> {
		// Lock the transaction
		let txn = self.ctx.tx();
		// Delete the old index data
		if let Some(o) = self.o.take() {
			let i = Indexable::new(o, self.ix);
			for o in i {
				let key = self.get_non_unique_index_key(&o)?;
				match txn.delc(&key, Some(self.rid)).await {
					Err(e) => {
						if matches!(
							e.downcast_ref::<Error>(),
							Some(Error::Kvs(crate::kvs::Error::TransactionConditionNotMet))
						) {
							Ok(())
						} else {
							Err(e)
						}
					}
					Ok(v) => Ok(v),
				}?
			}
		}
		// Create the new index data
		if let Some(n) = self.n.take() {
			let i = Indexable::new(n, self.ix);
			for n in i {
				let key = self.get_non_unique_index_key(&n)?;
				txn.set(&key, self.rid).await?;
			}
		}
		Ok(())
	}

	async fn index_count(
		&mut self,
		_stk: &mut Stk,
		cond: Option<&Cond>,
		require_compaction: &mut bool,
	) -> Result<()> {
		let mut relative_count: i8 = 0;
		if let Some(_c) = cond {
			let (old_matches, new_matches) = self.count_cond_match.unwrap_or((false, false));
			if self.o.is_some() && old_matches {
				relative_count -= 1;
			}
			if self.n.is_some() && new_matches {
				relative_count += 1;
			}
		} else {
			if self.o.is_some() {
				relative_count -= 1;
			}
			if self.n.is_some() {
				relative_count += 1;
			}
		}
		if relative_count == 0 {
			return Ok(());
		}
		let key = IndexCountKey::new(
			self.ns,
			self.db,
			&self.ix.table_name,
			self.ix.index_id,
			Some((self.ctx.node_id(), uuid::Uuid::now_v7())),
			relative_count > 0,
			relative_count.unsigned_abs() as u64,
		);
		self.ctx.tx().put(&key, &()).await?;
		*require_compaction = true;
		Ok(())
	}

	/// Creates the read-phase plan for full-text compaction.
	///
	/// The caller owns the transaction split so this can run in a read-only
	/// transaction and be applied later with a short write transaction.
	pub(crate) async fn prepare_fulltext_compaction(
		ixs: &IndexStores,
		ikb: &IndexKeyBase,
		tx: &Transaction,
		p: &FullTextParams,
		allow_list: &[PathBuf],
	) -> Result<FullTextCompactionPlan> {
		let ft = FullTextIndex::new(ixs, tx, ikb.clone(), p, allow_list).await?;
		ft.prepare_compaction(tx).await
	}

	/// Applies a prepared full-text compaction plan.
	///
	/// Returns `false` when there is no work or another compactor advanced the
	/// generation first.
	pub(crate) async fn apply_fulltext_compaction(
		ixs: &IndexStores,
		ikb: &IndexKeyBase,
		tx: &Transaction,
		p: &FullTextParams,
		allow_list: &[PathBuf],
		plan: FullTextCompactionPlan,
	) -> Result<bool> {
		let ft = FullTextIndex::new(ixs, tx, ikb.clone(), p, allow_list).await?;
		ft.apply_compaction(tx, plan).await
	}

	/// Creates the read-phase plan for HNSW pending compaction.
	pub(crate) async fn prepare_hnsw_compaction(
		ctx: &FrozenContext,
		ikb: &IndexKeyBase,
	) -> Result<HnswCompactionPlan> {
		HnswIndex::prepare_compaction(ctx, ikb).await
	}

	/// Applies a prepared HNSW pending compaction plan.
	///
	/// Returns `false` when there is no work, another compactor advanced the
	/// generation first, or a captured pending key changed before the write.
	pub(crate) async fn apply_hnsw_compaction(
		ctx: &FrozenContext,
		ixs: &IndexStores,
		ikb: &IndexKeyBase,
		ix: &IndexDefinition,
		p: &HnswParams,
		plan: HnswCompactionPlan,
	) -> Result<bool> {
		let tx = ctx.tx();
		if let Some(tb) = tx.get_tb(ikb.ns(), ikb.db(), ikb.table(), None).await? {
			let hnsw = ixs.get_index_hnsw(ikb.ns(), ikb.db(), ctx, tb.table_id, ix, p).await?;
			return hnsw.apply_compaction(ctx, plan).await;
		}
		Ok(false)
	}

	#[cfg(diskann)]
	/// Creates the read-phase plan for DiskANN pending compaction.
	pub(crate) async fn prepare_diskann_compaction(
		ctx: &FrozenContext,
		ikb: &IndexKeyBase,
	) -> Result<DiskAnnCompactionPlan> {
		DiskAnnIndex::prepare_compaction(ctx, ikb).await
	}

	#[cfg(diskann)]
	/// Applies a prepared DiskANN pending compaction plan.
	pub(crate) async fn apply_diskann_compaction(
		ctx: &FrozenContext,
		ixs: &IndexStores,
		ikb: &IndexKeyBase,
		ix: &IndexDefinition,
		p: &DiskAnnParams,
		plan: DiskAnnCompactionPlan,
	) -> Result<bool> {
		let tx = ctx.tx();
		if let Some(tb) = tx.get_tb(ikb.ns(), ikb.db(), ikb.table(), None).await? {
			let diskann = ixs.get_index_diskann(ikb.ns(), ikb.db(), tb.table_id, ix, p).await?;
			return diskann.apply_compaction(ctx, plan).await;
		}
		Ok(false)
	}

	/// Creates the read-phase plan for count-index compaction.
	pub(crate) async fn prepare_count_compaction(
		ikb: &IndexKeyBase,
		tx: &Transaction,
	) -> Result<IndexCountCompactionPlan> {
		IndexCountThingIterator::new(ikb.ns(), ikb.db(), ikb.table(), ikb.index())?
			.prepare_compaction(ikb, tx)
			.await
	}

	/// Applies a prepared count-index compaction plan.
	///
	/// Returns `false` when there is no work or another compactor advanced the
	/// generation first.
	pub(crate) async fn apply_count_compaction(
		ikb: &IndexKeyBase,
		tx: &Transaction,
		plan: IndexCountCompactionPlan,
	) -> Result<bool> {
		IndexCountThingIterator::apply_compaction(ikb, tx, plan).await
	}

	/// Construct a consistent uniqueness violation error message.
	/// Formats the conflicting value as a single value or array depending on
	/// the number of indexed fields.
	fn err_index_exists(&self, rid: RecordId, mut n: Array) -> Result<()> {
		bail!(Error::IndexExists {
			record: rid,
			index: self.ix.name.to_string(),
			value: match n.0.len() {
				1 => n.0.remove(0).to_sql(),
				_ => n.to_sql(),
			},
		})
	}

	async fn index_fulltext(
		&mut self,
		stk: &mut Stk,
		p: &FullTextParams,
		require_compaction: &mut bool,
	) -> Result<()> {
		// Build a FullText instance
		let fti = FullTextIndex::new(
			self.ctx.get_index_stores(),
			&self.ctx.tx(),
			self.ikb.clone(),
			p,
			&self.ctx.config.file_allowlist,
		)
		.await?;
		self.compute_fulltext_with_index(stk, &fti, require_compaction).await
	}

	pub(crate) async fn compute_fulltext_with_index(
		&mut self,
		stk: &mut Stk,
		fti: &FullTextIndex,
		require_compaction: &mut bool,
	) -> Result<()> {
		let mut rc = false;
		// Delete the old index data
		let doc_id = if let Some(o) = self.o.take() {
			fti.remove_content(stk, self.ctx, self.opt, self.rid, o, &mut rc).await?
		} else {
			None
		};
		// Create the new index data
		if let Some(n) = self.n.take() {
			fti.index_content(stk, self.ctx, self.opt, self.rid, n, &mut rc).await?;
		} else {
			// It is a deletion, we can remove the doc
			if let Some(doc_id) = doc_id {
				fti.remove_doc(self.ctx, doc_id).await?;
			}
		}
		// Do we need to trigger the compaction?
		if rc {
			*require_compaction = true;
		}
		Ok(())
	}

	pub(crate) async fn trigger_compaction(&self) -> Result<()> {
		IndexOperation::compaction_trigger(&self.ikb, &self.ctx.tx(), self.ctx.node_id()).await
	}

	/// Triggers index compaction.
	///
	/// This method adds an entry to the index compaction queue by creating an
	/// `Ic` key for the specified index. The index compaction thread will
	/// later process this entry and perform the actual compaction via
	/// [`Datastore::index_compaction`].
	///
	/// Compaction helps optimize index performance after many mutations.
	/// For full-text indexes it consolidates term frequency and document
	/// length data; for HNSW indexes it processes pending vector operations;
	/// for count indexes it reconciles count tracking entries.
	pub(crate) async fn compaction_trigger(
		ikb: &IndexKeyBase,
		tx: &Transaction,
		nid: Uuid,
	) -> Result<()> {
		let ic = ikb.new_ic_key(nid);
		tx.put(&ic, &()).await?;
		Ok(())
	}

	async fn index_hnsw(&mut self, p: &HnswParams, require_compaction: &mut bool) -> Result<()> {
		let hnsw = self
			.ctx
			.get_index_stores()
			.get_index_hnsw(self.ns, self.db, self.ctx, self.tb, self.ix, p)
			.await?;
		let old_values = self.o.take();
		let new_values = self.n.take();
		if old_values.is_some() || new_values.is_some() {
			hnsw.index(self.ctx, &self.rid.key, old_values, new_values).await?;
			*require_compaction = true;
		}
		Ok(())
	}

	async fn index_diskann(
		&mut self,
		p: &DiskAnnParams,
		require_compaction: &mut bool,
	) -> Result<()> {
		#[cfg(not(diskann))]
		{
			let _ = (p, require_compaction);
			bail!("DISKANN indexes require a 64-bit, non-WASM platform")
		}
		#[cfg(diskann)]
		{
			let diskann = self
				.ctx
				.get_index_stores()
				.get_index_diskann(self.ns, self.db, self.tb, self.ix, p)
				.await?;
			let old_values = self.o.take();
			let new_values = self.n.take();
			if old_values.is_some() || new_values.is_some() {
				diskann.index(self.ctx, &self.rid.key, old_values, new_values).await?;
				*require_compaction = true;
			}
			Ok(())
		}
	}
}

/// Extract from the given document, the values required by the index and put
/// then in an array. Eg. IF the index is composed of the columns `name` and
/// `instrument` Given this doc: { "id": 1, "instrument":"piano", "name":"Tobie"
/// } It will return: ["Tobie", "piano"]
struct Indexable(Vec<(Value, bool)>);

impl Indexable {
	fn new(vals: Vec<Value>, ix: &IndexDefinition) -> Self {
		let mut source = Vec::with_capacity(vals.len());
		for (v, i) in vals.into_iter().zip(ix.cols.iter()) {
			let f = matches!(i.0.last(), Some(&Part::Flatten));
			source.push((v, f));
		}
		Self(source)
	}
}

impl IntoIterator for Indexable {
	type Item = Array;
	type IntoIter = Combinator;

	fn into_iter(self) -> Self::IntoIter {
		Combinator::new(self.0)
	}
}

struct Combinator {
	iterators: Vec<Box<dyn ValuesIterator>>,
	has_next: bool,
}

impl Combinator {
	fn new(source: Vec<(Value, bool)>) -> Self {
		let mut iterators: Vec<Box<dyn ValuesIterator>> = Vec::new();
		// We create an iterator for each idiom
		for (v, f) in source {
			if !f {
				// Iterator for not flattened values
				if let Value::Array(v) = v {
					iterators.push(Box::new(MultiValuesIterator::new(v.0)));
					continue;
				}
			}
			iterators.push(Box::new(SingleValueIterator(v)));
		}
		Self {
			iterators,
			has_next: true,
		}
	}
}

impl Iterator for Combinator {
	type Item = Array;

	fn next(&mut self) -> Option<Self::Item> {
		if !self.has_next {
			return None;
		}
		let mut o = Vec::with_capacity(self.iterators.len());
		// Create the combination and advance to the next
		self.has_next = false;
		for i in &mut self.iterators {
			o.push(i.current().clone());
			if !self.has_next {
				// We advance only one iterator per iteration
				if i.next() {
					self.has_next = true;
				}
			}
		}
		let o = Array::from(o);
		Some(o)
	}
}

trait ValuesIterator: Send {
	fn next(&mut self) -> bool;
	fn current(&self) -> &Value;
}

struct MultiValuesIterator {
	vals: Vec<Value>,
	done: bool,
	current: usize,
	end: usize,
}

impl MultiValuesIterator {
	fn new(vals: Vec<Value>) -> Self {
		let len = vals.len();
		if len == 0 {
			Self {
				vals,
				done: true,
				current: 0,
				end: 0,
			}
		} else {
			Self {
				vals,
				done: false,
				current: 0,
				end: len - 1,
			}
		}
	}
}

impl ValuesIterator for MultiValuesIterator {
	fn next(&mut self) -> bool {
		if self.done {
			return false;
		}
		if self.current == self.end {
			self.done = true;
			return false;
		}
		self.current += 1;
		true
	}

	fn current(&self) -> &Value {
		self.vals.get(self.current).unwrap_or(&Value::Null)
	}
}

struct SingleValueIterator(Value);

impl ValuesIterator for SingleValueIterator {
	fn next(&mut self) -> bool {
		false
	}

	fn current(&self) -> &Value {
		&self.0
	}
}