surrealdb-core 3.2.1

//! Index function system for the streaming executor.
//!
//! This module provides traits for functions that are bound to WHERE clause
//! predicates via index infrastructure. Unlike scalar functions which operate
//! purely on their arguments, index functions reference a specific predicate
//! (e.g., a MATCHES clause or a KNN operator) in the WHERE condition and need
//! access to the associated index at evaluation time.
//!
//! Each index function declares what kind of index context it requires via
//! [`IndexContextKind`]. The planner uses this to resolve the appropriate
//! [`IndexContext`] at plan time:
//!
//! - **FullText**: The `index_ref_arg_index()` argument (e.g., the `1` in `search::highlight('<b>',
//!   '</b>', 1)`) is extracted and resolved to a [`MatchContext`] via the WHERE clause's MATCHES
//!   operators.
//!
//! - **Knn**: A [`KnnContext`] is created from the KNN operator in the WHERE clause. The KNN scan
//!   operator populates it with per-row distances at execution time.
//!
//! Examples: search::highlight, search::score, search::offsets, vector::distance::knn

use std::collections::HashMap;
use std::fmt::Debug;
use std::sync::Arc;

use anyhow::Result;

use super::Signature;
use crate::catalog::Index;
use crate::exec::physical_expr::EvalContext;
use crate::exec::{BoxFut, ContextLevel, SendSyncRequirement};
use crate::expr::Kind;
use crate::expr::idiom::Idiom;
use crate::idx::IndexKeyBase;
use crate::idx::ft::MatchRef;
use crate::idx::ft::fulltext::{FullTextIndex, QueryTerms, Scorer};
use crate::kvs::index::filter_online_indexes;
use crate::val::{Number, RecordId, TableName, Value};

// =========================================================================
// IndexContextKind - what kind of index context a function requires
// =========================================================================

/// What kind of index context an [`IndexFunction`] requires.
///
/// The planner uses this to resolve the appropriate [`IndexContext`] at plan
/// time without needing to know about specific function names.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum IndexContextKind {
	/// Full-text search context (resolved from MATCHES clauses).
	FullText,
	/// KNN distance context (resolved from KNN operators).
	Knn,
}

// =========================================================================
// IndexContext - resolved context enum passed to index functions
// =========================================================================

/// Resolved index context, created at plan time and passed to index functions
/// at evaluation time.
///
/// Each variant carries the context appropriate for its index type.
/// Cloning is cheap (all variants are `Arc`-wrapped).
#[derive(Clone)]
pub enum IndexContext {
	/// Full-text search context with lazy access to FT index resources.
	FullText(Arc<MatchContext>),
	/// KNN distance context populated by the KNN scan operator.
	Knn(Arc<KnnContext>),
}

impl Debug for IndexContext {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		match self {
			Self::FullText(ctx) => f.debug_tuple("IndexContext::FullText").field(ctx).finish(),
			Self::Knn(ctx) => f.debug_tuple("IndexContext::Knn").field(ctx).finish(),
		}
	}
}

// =========================================================================
// IndexFunction trait
// =========================================================================

/// A function that is bound to a WHERE clause predicate via index infrastructure.
///
/// Index functions differ from scalar functions in that they:
/// - Are associated with a specific WHERE clause predicate (MATCHES, KNN, etc.)
/// - Need access to index infrastructure at evaluation time
/// - May have a reference argument extracted at plan time (not passed at runtime)
///
/// The planner dispatches generically based on [`index_context_kind()`] to
/// resolve the appropriate [`IndexContext`], without hardcoding function names.
pub trait IndexFunction: SendSyncRequirement + Debug {
	/// The fully qualified function name (e.g., "search::highlight", "search::score")
	fn name(&self) -> &'static str;

	/// The function signature describing arguments and return type.
	#[allow(unused)]
	fn signature(&self) -> Signature;

	/// Infer the return type given the argument types.
	///
	/// The default implementation returns the signature's return type.
	#[allow(unused)]
	fn return_type(&self, _arg_types: &[Kind]) -> Result<Kind> {
		Ok(self.signature().returns)
	}

	/// What kind of index context this function requires.
	///
	/// The planner uses this to determine how to resolve the [`IndexContext`]:
	/// - [`IndexContextKind::FullText`]: resolve via MATCHES clause match_ref
	/// - [`IndexContextKind::Knn`]: resolve via KNN operator context
	fn index_context_kind(&self) -> IndexContextKind;

	/// Which argument position contains the index reference number.
	///
	/// When `Some(idx)`, this argument is extracted at plan time by the planner
	/// and is NOT passed to `invoke_async` as a runtime argument. For full-text
	/// functions, this is the match_ref that identifies a MATCHES clause.
	///
	/// Returns `None` if no reference argument is needed (e.g., KNN functions
	/// that use a single KNN operator from the WHERE clause).
	fn index_ref_arg_index(&self) -> Option<usize>;

	/// The minimum context level required to execute this function.
	///
	/// Index functions typically need root context for transaction and
	/// index store access.
	fn required_context(&self) -> ContextLevel {
		ContextLevel::Root
	}

	/// Async invocation with index context.
	///
	/// # Arguments
	/// * `ctx` - The evaluation context with access to current row and parameters
	/// * `index_ctx` - The resolved index context (FullText or Knn)
	/// * `args` - The evaluated function arguments, WITHOUT any plan-time extracted arguments
	///
	/// # Returns
	/// The computed value
	fn invoke_async<'a>(
		&'a self,
		ctx: &'a EvalContext<'_>,
		index_ctx: &'a IndexContext,
		args: Vec<Value>,
	) -> BoxFut<'a, Result<Value>>;
}

// =========================================================================
// MatchContext - resolved context for a single MATCHES clause
// =========================================================================

/// Resolved context for a single MATCHES clause, created at plan time.
///
/// This captures the field path and query string from a `WHERE field @N@ 'query'`
/// expression and provides lazy access to the associated full-text index
/// infrastructure. The expensive FullTextIndex/QueryTerms/Scorer are initialized
/// only on first use and then cached for all subsequent rows.
pub struct MatchContext {
	/// The field path from the left side of the MATCHES operator.
	pub idiom: Idiom,
	/// The search query string from the right side of the MATCHES operator.
	pub query: String,
	/// The table name for index lookup.
	pub table: TableName,
	/// Lazily initialized full-text index resources.
	ft_cache: tokio::sync::OnceCell<(FullTextIndex, QueryTerms, Option<Scorer>)>,
}

impl MatchContext {
	/// Create a new MatchContext from resolved MATCHES clause info.
	pub fn new(idiom: Idiom, query: String, table: TableName) -> Self {
		Self {
			idiom,
			query,
			table,
			ft_cache: tokio::sync::OnceCell::new(),
		}
	}

	/// Get or lazily initialize the full-text index resources.
	///
	/// On first call, this looks up the full-text index definition for the
	/// table/idiom, opens the FullTextIndex, extracts QueryTerms, and
	/// optionally creates a Scorer. Subsequent calls return the cached result.
	pub async fn ft_resources(
		&self,
		ctx: &EvalContext<'_>,
	) -> Result<&(FullTextIndex, QueryTerms, Option<Scorer>)> {
		self.ft_cache
			.get_or_try_init(|| async {
				use crate::catalog::providers::TableProvider;

				let frozen = ctx.exec_ctx.ctx();
				let root = ctx.exec_ctx.root();
				let opt = root
					.options
					.as_ref()
					.ok_or_else(|| anyhow::anyhow!("IndexFunction requires Options context"))?;
				let tx = ctx.txn();

				// Get namespace and database IDs from the execution context
				let db_ctx = ctx.exec_ctx.database().map_err(|e| {
					anyhow::anyhow!("IndexFunction requires database context: {}", e)
				})?;
				let ns_id = db_ctx.ns_ctx.ns.namespace_id;
				let db_id = db_ctx.db.database_id;

				// Find the full-text index for this table and idiom
				let indexes = tx
					.all_tb_indexes(ns_id, db_id, &self.table, ctx.exec_ctx.version_stamp())
					.await?;
				let indexes = if ctx.exec_ctx.version_stamp().is_none() {
					// MATCHES/scoring must only open indexes that durable build
					// state has published as queryable.
					filter_online_indexes(tx.as_ref(), ns_id, db_id, indexes).await?
				} else {
					indexes
				};
				let index_def = indexes
					.iter()
					.find(|idx| {
						matches!(&idx.index, Index::FullText(_))
							&& idx.cols.iter().any(|col| col.0 == self.idiom.0)
					})
					.ok_or_else(|| {
						anyhow::anyhow!(
							"No full-text index found for field {:?} on table {}",
							self.idiom,
							self.table
						)
					})?;

				let ft_params = match &index_def.index {
					Index::FullText(params) => params,
					_ => unreachable!("Already checked for FullText above"),
				};

				let ikb = IndexKeyBase::new(ns_id, db_id, self.table.clone(), index_def.index_id);

				// Open the full-text index
				let fti = FullTextIndex::new(
					frozen.get_index_stores(),
					tx.as_ref(),
					ikb,
					ft_params,
					&frozen.config.file_allowlist,
				)
				.await?;

				// Extract query terms
				let query_terms = {
					let mut stack = reblessive::TreeStack::new();
					stack
						.enter(|stk| {
							fti.extract_querying_terms(stk, frozen, opt, self.query.clone())
						})
						.finish()
						.await?
				};

				// Create scorer if BM25 is configured
				let scorer = fti.new_scorer(frozen).await?;

				Ok((fti, query_terms, scorer))
			})
			.await
	}
}

impl Debug for MatchContext {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		f.debug_struct("MatchContext")
			.field("idiom", &self.idiom)
			.field("query", &self.query)
			.field("table", &self.table)
			.field("initialized", &self.ft_cache.initialized())
			.finish()
	}
}

// =========================================================================
// KnnContext - distance context populated by KNN scan operators
// =========================================================================

/// KNN distance context, populated by KnnScan at execution time.
///
/// Created at plan time and shared (via `Arc`) between the KNN scan operator
/// and the `vector::distance::knn()` index function. The scan operator writes
/// per-row distances after the HNSW search completes; the function reads them
/// during projection evaluation.
///
/// This is the KNN equivalent of [`MatchContext`] -- same lifecycle pattern
/// (plan-time creation, shared via `Arc`, deferred population at execution time).
pub struct KnnContext {
	/// Per-row distances keyed by RecordId, populated by the KNN scan operator.
	///
	/// Uses `tokio::sync::RwLock` so lock acquisition is async and cannot
	/// block the tokio runtime.
	distances: tokio::sync::RwLock<HashMap<RecordId, Number>>,
}

impl KnnContext {
	/// Create a new empty KnnContext.
	pub fn new() -> Self {
		Self {
			distances: tokio::sync::RwLock::new(HashMap::new()),
		}
	}

	/// Record the distance for a record. Called by KnnScan after HNSW search.
	pub async fn insert(&self, rid: RecordId, dist: Number) {
		self.distances.write().await.insert(rid, dist);
	}

	/// Look up the distance for a record. Called by vector::distance::knn().
	pub async fn get(&self, rid: &RecordId) -> Option<Number> {
		self.distances.read().await.get(rid).copied()
	}
}

impl Default for KnnContext {
	fn default() -> Self {
		Self::new()
	}
}

impl Debug for KnnContext {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		match self.distances.try_read() {
			Ok(guard) => f.debug_struct("KnnContext").field("entries", &guard.len()).finish(),
			Err(_) => f.debug_struct("KnnContext").field("entries", &"<locked>").finish(),
		}
	}
}

// =========================================================================
// MatchesContext - planning-time map of all MATCHES clauses
// =========================================================================

/// Information about a single MATCHES clause extracted from the WHERE condition.
#[derive(Debug, Clone)]
pub struct MatchInfo {
	/// The field path from the left side of the MATCHES operator.
	pub idiom: Idiom,
	/// The search query string from the right side of the MATCHES operator.
	pub query: String,
}

/// Planning-time context mapping match_ref numbers to MATCHES clause info.
///
/// Built by analyzing the WHERE clause AST during query planning. Each entry
/// maps a match_ref number (e.g., `1` from `@1@`) to the idiom and query
/// string of the corresponding MATCHES operator.
#[derive(Debug, Clone)]
pub struct MatchesContext {
	matches: HashMap<MatchRef, MatchInfo>,
	/// The table name from the FROM clause, set during planning.
	table: Option<TableName>,
}

impl MatchesContext {
	/// Create a new empty MatchesContext.
	pub fn new() -> Self {
		Self {
			matches: HashMap::new(),
			table: None,
		}
	}

	/// Set the table name for index lookup.
	pub fn set_table(&mut self, table: TableName) {
		self.table = Some(table);
	}

	/// Get the table name.
	pub fn table(&self) -> Option<&TableName> {
		self.table.as_ref()
	}

	/// Insert a MATCHES clause entry.
	pub fn insert(&mut self, match_ref: MatchRef, info: MatchInfo) {
		self.matches.insert(match_ref, info);
	}

	/// Look up a MATCHES clause by its match_ref number.
	pub fn get(&self, match_ref: MatchRef) -> Option<&MatchInfo> {
		self.matches.get(&match_ref)
	}

	/// Get the first available MatchInfo (for when there's only one MATCHES clause).
	pub fn first(&self) -> Option<(MatchRef, &MatchInfo)> {
		self.matches.iter().next().map(|(&k, v)| (k, v))
	}

	/// Check if the context has no MATCHES entries.
	pub fn is_empty(&self) -> bool {
		self.matches.is_empty()
	}

	/// Create a MatchContext for a given match_ref, resolving against this context.
	///
	/// If the match_ref is found, creates a MatchContext with the resolved
	/// idiom, query, and table name. If not found and there's exactly one
	/// entry, falls back to that entry (common case: single MATCHES clause).
	pub fn resolve(&self, match_ref: MatchRef, table: TableName) -> Result<Arc<MatchContext>> {
		let info = self.get(match_ref).or_else(|| {
			// Fall back to the single entry if there's only one
			if self.matches.len() == 1 {
				self.first().map(|(_, info)| info)
			} else {
				None
			}
		});

		match info {
			Some(info) => {
				Ok(Arc::new(MatchContext::new(info.idiom.clone(), info.query.clone(), table)))
			}
			None => {
				// If there are no MATCHES clauses at all, provide a clear error
				if self.matches.is_empty() {
					Err(anyhow::anyhow!("no MATCHES clause found in WHERE condition"))
				} else {
					Err(anyhow::anyhow!(
						"no MATCHES clause found for match_ref {} (available: {:?})",
						match_ref,
						self.matches.keys().collect::<Vec<_>>()
					))
				}
			}
		}
	}
}

impl Default for MatchesContext {
	fn default() -> Self {
		Self::new()
	}
}