surrealdb-core 3.2.1

//! Index scan operator for B-tree index access.
//!
//! This operator retrieves records using B-tree index structures (Idx and Uniq),
//! supporting equality lookups, range scans, and union operations.

use std::collections::HashSet;
use std::sync::Arc;

use surrealdb_types::ToSql;

use super::common::{fetch_and_filter_records_batch, resolve_version_stamp};
use super::pipeline::{build_field_state, eval_limit_expr};
use super::resolved::ResolvedTableContext;
use crate::err::Error;
use crate::exec::index::access_path::{BTreeAccess, IndexRef};
use crate::exec::index::iterator::btree::{CompoundEqualIterator, CompoundRangeIterator};
use crate::exec::index::iterator::{
	IndexEqualIterator, IndexRangeIterator, UniqueEqualIterator, UniqueRangeIterator,
};
use crate::exec::permission::{
	PhysicalPermission, convert_permission_to_physical_runtime, should_check_perms,
	validate_record_user_access,
};
use crate::exec::{
	AccessMode, ContextLevel, ControlFlowExt, ExecOperator, ExecutionContext, FlowResult,
	OperatorMetrics, PhysicalExpr, ValueBatch, ValueBatchStream, monitor_stream,
};
use crate::expr::ControlFlow;
use crate::iam::Action;
use crate::idx::planner::ScanDirection;
use crate::kvs::CachePolicy;

/// Index scan operator for B-tree indexes (Idx and Uniq).
///
/// Retrieves records using an index access path, then fetches the full
/// record data and applies permission filtering.
///
/// When `limit` and/or `start` are provided (pushed down from the planner),
/// the operator stops iteration early once the limit is reached, avoiding
/// unnecessary index and record reads.
#[derive(Debug)]
pub struct IndexScan {
	/// Reference to the index definition
	pub index_ref: IndexRef,
	/// How to access the index
	pub access: BTreeAccess,
	/// Scan direction (forward or backward)
	pub direction: ScanDirection,
	/// Table name for record fetching
	pub table_name: crate::val::TableName,
	/// Pushed-down LIMIT expression (evaluated at execution time).
	pub(crate) limit: Option<Arc<dyn PhysicalExpr>>,
	/// Pushed-down START expression (evaluated at execution time).
	pub(crate) start: Option<Arc<dyn PhysicalExpr>>,
	/// Optional VERSION timestamp for time-travel queries.
	pub(crate) version: Option<Arc<dyn PhysicalExpr>>,
	/// Plan-time resolved table context. When present, `execute()` skips
	/// runtime table def + permission lookup.
	pub(crate) resolved: Option<ResolvedTableContext>,
	/// Projection-aware field set for computed-field materialization.
	/// Outer `None` = sub-operator mode (parent handles fields).
	/// `Some(None)` = all fields, `Some(Some(set))` = specific fields.
	pub(crate) needed_fields: Option<Option<HashSet<String>>>,
	/// Full WHERE predicate as a streaming physical expression. Unlike
	/// [`TableScan`](super::TableScan), index seeks only narrow candidates;
	/// this predicate is evaluated on fetched rows to enforce any conditions
	/// beyond what the indexer proved (matching `DynamicScan`'s unified
	/// outer pipeline behaviour).
	pub(crate) where_predicate: Option<Arc<dyn PhysicalExpr>>,
	/// Per-batch size ceiling when LIMIT wasn't pushed (due to a residual
	/// filter preventing direct pushdown).  The scan reads batches of at
	/// most this many entries, enabling faster early termination from the
	/// downstream Limit operator.  Does NOT cap total entries — the loop
	/// continues until either the range is exhausted or the consumer
	/// drops the stream.
	pub(crate) batch_ceiling: Option<Arc<dyn PhysicalExpr>>,
	/// Per-operator runtime metrics for EXPLAIN ANALYZE.
	pub(crate) metrics: Arc<OperatorMetrics>,
}

impl IndexScan {
	#[allow(clippy::too_many_arguments)]
	pub(crate) fn new(
		index_ref: IndexRef,
		access: BTreeAccess,
		direction: ScanDirection,
		table_name: crate::val::TableName,
		limit: Option<Arc<dyn PhysicalExpr>>,
		start: Option<Arc<dyn PhysicalExpr>>,
		version: Option<Arc<dyn PhysicalExpr>>,
		needed_fields: Option<Option<HashSet<String>>>,
		where_predicate: Option<Arc<dyn PhysicalExpr>>,
	) -> Self {
		Self {
			index_ref,
			access,
			direction,
			table_name,
			limit,
			start,
			version,
			resolved: None,
			needed_fields,
			where_predicate,
			batch_ceiling: None,
			metrics: Arc::new(OperatorMetrics::new()),
		}
	}

	/// Set the plan-time resolved table context.
	pub(crate) fn with_resolved(mut self, resolved: ResolvedTableContext) -> Self {
		self.resolved = Some(resolved);
		self
	}

	/// Set a per-batch ceiling for downstream LIMIT awareness.
	///
	/// When the planner knows there is a downstream LIMIT but cannot push
	/// it to the scan (residual filter), it passes the user's LIMIT here.
	/// This makes each batch small so the downstream Limit operator can
	/// terminate the stream quickly instead of waiting for a full 1000-entry
	/// batch.
	pub(crate) fn with_batch_ceiling(mut self, ceiling: Option<Arc<dyn PhysicalExpr>>) -> Self {
		self.batch_ceiling = ceiling;
		self
	}
}
impl ExecOperator for IndexScan {
	fn name(&self) -> &'static str {
		"IndexScan"
	}

	fn attrs(&self) -> Vec<(String, String)> {
		let access_str = match &self.access {
			BTreeAccess::Equality(v) => format!("= {}", v.to_sql()),
			BTreeAccess::Range {
				from,
				to,
			} => {
				let from_str = match from {
					Some(r) => format!(
						"{}{}",
						if r.inclusive {
							">="
						} else {
							">"
						},
						r.value.to_sql()
					),
					None => String::new(),
				};
				let to_str = match to {
					Some(r) => format!(
						"{}{}",
						if r.inclusive {
							"<="
						} else {
							"<"
						},
						r.value.to_sql()
					),
					None => String::new(),
				};
				format!("{from_str} {to_str}").trim().to_string()
			}
			BTreeAccess::Compound {
				prefix,
				range,
			} => {
				let prefix_str = prefix.iter().map(|v| v.to_sql()).collect::<Vec<_>>().join(", ");
				if let Some((op, val)) = range {
					let val_sql = val.to_sql();
					format!("[{prefix_str}] {op:?} {val_sql}")
				} else {
					format!("[{prefix_str}]")
				}
			}
			// FullText and KNN should use dedicated operators
			BTreeAccess::FullText {
				..
			}
			| BTreeAccess::Knn {
				..
			} => {
				unreachable!("IndexScan does not support FullText or KNN access")
			}
		};
		let mut attrs = vec![
			("index".to_string(), self.index_ref.name.to_string()),
			("access".to_string(), access_str),
			("direction".to_string(), format!("{:?}", self.direction)),
		];
		if let Some(ref limit) = self.limit {
			attrs.push(("limit".to_string(), limit.to_sql()));
		}
		if let Some(ref start) = self.start {
			attrs.push(("offset".to_string(), start.to_sql()));
		}
		attrs
	}

	fn required_context(&self) -> ContextLevel {
		ContextLevel::Database
	}

	fn access_mode(&self) -> AccessMode {
		let mut mode = AccessMode::ReadOnly;
		if let Some(ref limit) = self.limit {
			mode = mode.combine(limit.access_mode());
		}
		if let Some(ref start) = self.start {
			mode = mode.combine(start.access_mode());
		}
		if let Some(ref pred) = self.where_predicate {
			mode = mode.combine(pred.access_mode());
		}
		mode
	}

	fn metrics(&self) -> Option<&OperatorMetrics> {
		Some(&self.metrics)
	}

	fn output_ordering(&self) -> crate::exec::OutputOrdering {
		use crate::exec::operators::SortDirection;
		use crate::exec::ordering::SortProperty;

		let dir = match self.direction {
			ScanDirection::Forward => SortDirection::Asc,
			ScanDirection::Backward => SortDirection::Desc,
		};

		// For compound access with an equality prefix, the prefix columns all
		// have the same value within the scan and do not define ordering.
		// Skip them so that the effective ordering starts from the first
		// non-equality column. This allows `satisfies()` to match ORDER BY
		// on the column after the prefix (e.g. `ORDER BY modified DESC`
		// with `idx(IsVisible, modified)` and `WHERE IsVisible = true`).
		//
		// For single-column Equality access (`WHERE col = val`), ALL index
		// columns are constant, so we skip them all.  The effective ordering
		// is then just the implicit record-id tail for non-unique indexes.
		let skip_cols = match &self.access {
			BTreeAccess::Compound {
				prefix,
				..
			} => prefix.len(),
			BTreeAccess::Equality(_) => self.index_ref.definition().cols.len(),
			_ => 0,
		};

		let ix_def = self.index_ref.definition();
		let mut cols: Vec<SortProperty> = ix_def
			.cols
			.iter()
			.skip(skip_cols)
			.filter_map(|idiom| {
				crate::exec::field_path::FieldPath::try_from(idiom).ok().map(|path| SortProperty {
					path,
					direction: dir,
					collate: false,
					numeric: false,
				})
			})
			.collect();

		// For non-unique indexes (Idx), the record ID is stored in the BTree
		// key after the field values.  This means entries are implicitly
		// sorted by record ID after the declared index columns.  Expose
		// this so that ORDER BY (col DESC, id DESC) is recognised as
		// satisfied by a backward index scan.
		//
		// When all index columns are skipped (e.g., single-column Equality),
		// the effective ordering is *only* by record ID.  We still append
		// the `id` property so `ORDER BY id` can be satisfied.
		if !self.index_ref.is_unique() {
			// Only append if the index actually has columns (guards against
			// degenerate case of zero-column index definitions).
			if !ix_def.cols.is_empty() {
				cols.push(SortProperty {
					path: crate::exec::field_path::FieldPath::field("id"),
					direction: dir,
					collate: false,
					numeric: false,
				});
			}
		}

		if cols.is_empty() {
			crate::exec::OutputOrdering::Unordered
		} else {
			crate::exec::OutputOrdering::Sorted(cols)
		}
	}

	fn constant_output_fields(&self) -> Vec<crate::exec::field_path::FieldPath> {
		use crate::exec::index::access_path::BTreeAccess;

		let ix_def = self.index_ref.definition();
		match &self.access {
			// All index columns have the same value for Equality scans
			BTreeAccess::Equality(_) => ix_def
				.cols
				.iter()
				.filter_map(|idiom| crate::exec::field_path::FieldPath::try_from(idiom).ok())
				.collect(),
			// Compound prefix columns are all equality-pinned
			BTreeAccess::Compound {
				prefix,
				..
			} => ix_def
				.cols
				.iter()
				.take(prefix.len())
				.filter_map(|idiom| crate::exec::field_path::FieldPath::try_from(idiom).ok())
				.collect(),
			_ => vec![],
		}
	}

	fn execute(&self, ctx: &ExecutionContext) -> FlowResult<ValueBatchStream> {
		let db_ctx = ctx.database()?.clone();

		// Validate record user has access to this namespace/database
		validate_record_user_access(&db_ctx)?;

		// Check if we need to enforce permissions
		let check_perms = should_check_perms(&db_ctx, Action::View)?;

		// Clone for the async block
		let index_ref = self.index_ref.clone();
		let access = self.access.clone();
		let direction = self.direction;
		let table_name = self.table_name.clone();
		let limit_expr = self.limit.clone();
		let start_expr = self.start.clone();
		let ceiling_expr = self.batch_ceiling.clone();
		let version_expr = self.version.clone();
		let resolved = self.resolved.clone();
		let needed_fields = self.needed_fields.clone();
		let where_predicate = self.where_predicate.clone();
		let ctx = ctx.clone();

		let stream = async_stream::try_stream! {
			let db_ctx = ctx.database()?;
			let txn = ctx.txn();
			let ns = Arc::clone(&db_ctx.ns_ctx.ns);
			let db = Arc::clone(&db_ctx.db);
			let ns_id = ns.namespace_id;
			let db_id = db.database_id;

			// Evaluate pushed-down LIMIT and START expressions
			let limit_val: Option<usize> = match &limit_expr {
				Some(expr) => Some(eval_limit_expr(&**expr, &ctx).await?),
				None => None,
			};
			let start_val: usize = match &start_expr {
				Some(expr) => eval_limit_expr(&**expr, &ctx).await?,
				None => 0,
			};

			// Evaluate batch ceiling expression (downstream LIMIT hint for
			// residual-filter queries).  Each batch reads at most this many
			// index entries so the downstream Limit operator can stop the
			// stream quickly instead of waiting for a full 1000-entry batch.
			// We use a 4x multiplier to account for rows rejected by the
			// residual filter — this keeps the batch large enough to avoid
			// excessive small-batch round-trips while still being far smaller
			// than the default 1000-entry INDEX_BATCH_SIZE.
			let batch_max: u32 = match &ceiling_expr {
				Some(expr) => {
					let c = eval_limit_expr(&**expr, &ctx).await?;
					c.saturating_add(start_val).saturating_mul(4).clamp(1, 1000) as u32
				}
				None => u32::MAX, // next_batch caps at INDEX_BATCH_SIZE internally
			};

			// Resolve VERSION timestamp; see [`resolve_version_stamp`] for
			// why we prefer the stamp already set by the enclosing
			// `VersionScope` over re-evaluating `version_expr` here.
			let version: Option<u64> = resolve_version_stamp(&ctx, version_expr.as_ref()).await?;

			// Resolve table permissions: plan-time fast path or runtime fallback
			let select_permission = if let Some(ref res) = resolved {
				res.select_permission(check_perms)
			} else if check_perms {
				let table_def = db_ctx
					.get_table_def(&table_name, version)
					.await
					.context("Failed to get table")?;

				if let Some(def) = &table_def {
					convert_permission_to_physical_runtime(&def.permissions.select, ctx.ctx())
						.await
						.context("Failed to convert permission")?
				} else {
					Err(ControlFlow::Err(anyhow::Error::new(Error::TbNotFound {
						name: table_name.clone(),
					})))?
				}
			} else {
				PhysicalPermission::Allow
			};

			// Early exit if denied
			if matches!(select_permission, PhysicalPermission::Deny) {
				return;
			}

			if limit_val == Some(0) {
				return;
			}

			// Resolve field state for computed fields and field-level
			// permissions. When needed_fields is None (sub-operator mode),
			// the parent operator handles field processing.
			let field_state = match &needed_fields {
				Some(nf) => {
					if let Some(ref res) = resolved {
						res.field_state_for_projection(nf.as_ref())
					} else {
						build_field_state(
							&ctx, &table_name, check_perms, nf.as_ref(),
						).await?
					}
				}
				None => super::pipeline::FieldState::empty(),
			};

			// Table-level permissions are already handled by
			// fetch_and_filter_records_batch, so the pipeline uses Allow
			// to avoid double-checking.
			let mut pipeline = super::pipeline::ScanPipeline::new(
				PhysicalPermission::Allow,
				where_predicate,
				field_state,
				check_perms,
				limit_val,
				start_val,
			);

			// Create the appropriate iterator based on access type and index uniqueness
			let is_unique = index_ref.is_unique();
			let ix = index_ref.definition();

			// Collect record IDs from index and batch-fetch full records
			match (&access, is_unique) {
				// Unique equality — single record for non-NULL values,
				// but NONE/NULL entries use prefix scanning and can
				// match multiple records requiring pagination.
				(BTreeAccess::Equality(value), true) => {
					let mut iter = UniqueEqualIterator::new(ns_id, db_id, ix, value)
						.context("Failed to create iterator")?;

					loop {
						if ctx.cancellation().is_cancelled() {
							Err(ControlFlow::Err(anyhow::anyhow!(
								crate::err::Error::QueryCancelled
							)))?;
						}
						let rids = iter.next_batch(&txn).await
							.context("Failed to iterate index")?;
						if rids.is_empty() {
							break;
						}

						let mut values = fetch_and_filter_records_batch(
							&ctx, &txn, ns_id, db_id, &rids, &select_permission, check_perms, version,
							CachePolicy::ReadOnly,
						).await?;

						let cont = pipeline.process_batch(&mut values, &ctx).await?;

						if !values.is_empty() {
							yield ValueBatch { values };
						}
						if !cont {
							break;
						}
					}
				}

				// Non-unique equality - multiple records possible
				(BTreeAccess::Equality(value), false) => {
					let reverse = matches!(direction, ScanDirection::Backward);
					let mut iter = IndexEqualIterator::with_direction(ns_id, db_id, ix, value, reverse)
						.context("Failed to create iterator")?;

					loop {
						if ctx.cancellation().is_cancelled() {
							Err(ControlFlow::Err(anyhow::anyhow!(
								crate::err::Error::QueryCancelled
							)))?;
						}
						let rids = iter.next_batch(&txn).await
							.context("Failed to iterate index")?;
						if rids.is_empty() {
							break;
						}

						let mut values = fetch_and_filter_records_batch(
							&ctx, &txn, ns_id, db_id, &rids, &select_permission, check_perms, version,
							CachePolicy::ReadOnly,
						).await?;

						let cont = pipeline.process_batch(&mut values, &ctx).await?;

						if !values.is_empty() {
							yield ValueBatch { values };
						}
						if !cont {
							break;
						}
					}
				}

				// Range scan (unique or non-unique).
				//
				// Both branches share the same batch-fetch-yield loop; they
				// differ only in iterator construction.  We keep them as two
				// explicit `loop` blocks rather than abstracting over the
				// iterator type because `async_stream` closures cannot
				// easily hold trait objects or generics.
			 (BTreeAccess::Range { from, to }, true) => {
					let mut iter = UniqueRangeIterator::new(ns_id, db_id, ix, from.as_ref(), to.as_ref(), direction).context("Failed to create iterator")?;

					loop {
						if ctx.cancellation().is_cancelled() {
							Err(ControlFlow::Err(anyhow::anyhow!(
								Error::QueryCancelled
							)))?;
						}
						let rids = iter.next_batch(&txn).await
							.context("Failed to iterate index")?;
						if rids.is_empty() { break; }

						let mut values = fetch_and_filter_records_batch(
							&ctx, &txn, ns_id, db_id, &rids, &select_permission, check_perms, version,
							CachePolicy::ReadOnly,
						).await?;

						let cont = pipeline.process_batch(&mut values, &ctx).await?;

						if !values.is_empty() {
							yield ValueBatch { values };
						}
						if !cont {
							break;
						}
					}
				}

				(BTreeAccess::Range { from, to }, false) => {
					let mut iter = IndexRangeIterator::new(ns_id, db_id, ix, from.as_ref(), to.as_ref(), direction).context("Failed to create iterator")?;

					loop {
						if ctx.cancellation().is_cancelled() {
							Err(ControlFlow::Err(anyhow::anyhow!(
								Error::QueryCancelled
							)))?
						}
						let rids = iter.next_batch(&txn).await
							.context("Failed to iterate index")?;
						if rids.is_empty() { break; }

						let mut values = fetch_and_filter_records_batch(
							&ctx, &txn, ns_id, db_id, &rids, &select_permission, check_perms, version,
							CachePolicy::ReadOnly,
						).await?;

						let cont = pipeline.process_batch(&mut values, &ctx).await?;

						if !values.is_empty() {
							yield ValueBatch { values };
						}
						if !cont {
							break;
						}
					}
				}

				// Compound index access — equality prefix only (no range)
				(BTreeAccess::Compound { prefix, range: None }, _) => {

					let mut iter = CompoundEqualIterator::new(ns_id, db_id, ix, prefix, None, direction).context("Failed to create compound iterator")?;

					// Compute the maximum number of index entries we need.
					// When a LIMIT + START is pushed down AND permissions
					// won't filter rows, we can cap the scan at limit+start
					// index entries. With conditional permissions, rows may
					// be denied after fetch, so we must not cap — let the
					// pipeline's limit/start tracking terminate the loop.
					let can_cap = !matches!(select_permission, PhysicalPermission::Conditional(_));
					let mut remaining: u32 = match (limit_val, can_cap) {
						(Some(l), true) => l.saturating_add(start_val).min(u32::MAX as usize) as u32,
						_ => u32::MAX,
					};

					// Fetch the first batch of record IDs sequentially.
					// Use batch_max to keep batches small when a downstream
					// LIMIT exists but wasn't pushed (residual filter).
					let mut rids = iter.next_batch(&txn, remaining.min(batch_max)).await
						.context("Failed to iterate compound index")?;

					while !rids.is_empty() {
						if ctx.cancellation().is_cancelled() {
							Err(ControlFlow::Err(anyhow::anyhow!(
								crate::err::Error::QueryCancelled
							)))?;
						}
						remaining = remaining.saturating_sub(rids.len() as u32);

						// Overlap: fetch records for the current batch while
						// scanning the next batch of index entries concurrently.
						// This halves serial latency on TiKV.
						let (values_result, next_rids_result) = if remaining > 0 {
							let fetch_fut = fetch_and_filter_records_batch(
								&ctx, &txn, ns_id, db_id, &rids, &select_permission, check_perms, version,
								CachePolicy::ReadOnly,
							);
							let scan_fut = iter.next_batch(&txn, remaining.min(batch_max));
							let (v, n) = futures::join!(fetch_fut, scan_fut);
							(v, Some(n))
						} else {
							// No more entries needed; skip the prefetch.
							let v = fetch_and_filter_records_batch(
								&ctx, &txn, ns_id, db_id, &rids, &select_permission, check_perms, version,
								CachePolicy::ReadOnly,
							).await;
							(v, None)
						};

						let mut values = values_result?;
						let cont = pipeline.process_batch(&mut values, &ctx).await?;

						if !values.is_empty() {
							yield ValueBatch { values };
						}
						if !cont || remaining == 0 {
							break;
						}

						rids = match next_rids_result {
							Some(r) => r.context("Failed to iterate compound index")?,
							// Iterator exhausted before remaining reached 0 — stop cleanly.
							None => break,
						};
					}
				}

				// Compound index access — equality prefix with range on next column
				(BTreeAccess::Compound { prefix, range: Some(range) }, _) => {
					let mut iter = CompoundRangeIterator::new(ns_id, db_id, ix, prefix, range, direction).context("Failed to create compound range iterator")?;

					// Same cap logic as the equality-only compound branch:
					// only cap when permissions won't filter rows post-fetch.
					let can_cap = !matches!(select_permission, PhysicalPermission::Conditional(_));
					let mut remaining: u32 = match (limit_val, can_cap) {
						(Some(l), true) => l.saturating_add(start_val).min(u32::MAX as usize) as u32,
						_ => u32::MAX,
					};

					// Fetch the first batch of record IDs sequentially.
					// Use batch_max to keep batches small when a downstream
					// LIMIT exists but wasn't pushed (residual filter).
					let mut rids = iter.next_batch(&txn, remaining.min(batch_max)).await
						.context("Failed to iterate compound index")?;

					while !rids.is_empty() {
						if ctx.cancellation().is_cancelled() {
							Err(ControlFlow::Err(anyhow::anyhow!(
								crate::err::Error::QueryCancelled
							)))?;
						}
						remaining = remaining.saturating_sub(rids.len() as u32);

						// Overlap: fetch records for the current batch while
						// scanning the next batch of index entries concurrently.
						let (values_result, next_rids_result) = if remaining > 0 {
							let fetch_fut = fetch_and_filter_records_batch(
								&ctx, &txn, ns_id, db_id, &rids, &select_permission, check_perms, version,
								CachePolicy::ReadOnly,
							);
							let scan_fut = iter.next_batch(&txn, remaining.min(batch_max));
							let (v, n) = futures::join!(fetch_fut, scan_fut);
							(v, Some(n))
						} else {
							let v = fetch_and_filter_records_batch(
								&ctx, &txn, ns_id, db_id, &rids, &select_permission, check_perms, version,
								CachePolicy::ReadOnly,
							).await;
							(v, None)
						};

						let mut values = values_result?;
						let cont = pipeline.process_batch(&mut values, &ctx).await?;

						if !values.is_empty() {
							yield ValueBatch { values };
						}
						if !cont || remaining == 0 {
							break;
						}

						rids = match next_rids_result {
							Some(r) => r.context("Failed to iterate compound index")?,
							// Iterator exhausted before remaining reached 0 — stop cleanly.
							None => break,
						};
					}
				}

				// FullText and KNN should use dedicated operators
				(BTreeAccess::FullText { .. }, _) | (BTreeAccess::Knn { .. }, _) => {
					Err(ControlFlow::Err(anyhow::anyhow!(
						"IndexScan does not support FullText or KNN access - use dedicated operators"
					)))?
				}
			}
		};

		Ok(monitor_stream(Box::pin(stream), "IndexScan", &self.metrics))
	}
}