surrealdb_expr/val/

set.rs

use revision::revisioned;
use storekey::{BorrowDecode, Encode};
use surrealdb_collections::{VecSet, VecSetIntoIter};
use surrealdb_types::{SqlFormat, ToSql, write_sql};

use crate::expr::Expr;
use crate::val::{IndexFormat, Value};

/// Internal Set type that stores unique values
///
/// Sets use [`VecSet`] internally to maintain uniqueness and sorted order.
///
/// - **Rev 1** — `u16 revision || VecSet<Value>` (length-prefixed). Byte-identical to the legacy
///   on-disk encoding.
/// - **Rev 2** — optimised envelope (`u16 revision || u32_le payload_length`), inner `VecSet`
///   written via the indexed-set prologue past `OFFSET_TABLE_MIN_LEN = 8`. Walker descent stays
///   zero-allocation through the Wire-repr fast path (skip + borrow).
#[revisioned(revision(1), revision(2, optimised))]
#[derive(Clone, Debug, Default, Eq, Ord, PartialEq, PartialOrd, Hash, Encode, BorrowDecode)]
#[storekey(format = "()")]
#[storekey(format = "IndexFormat")]
pub struct Set(#[revision(indexed_set)] pub VecSet<Value>);

impl Set {
	/// Create a new empty set
	pub fn new() -> Self {
		Set(VecSet::new())
	}

	/// Get the number of elements in the set
	pub fn len(&self) -> usize {
		self.0.len()
	}

	/// Check if the set is empty
	pub fn is_empty(&self) -> bool {
		self.0.is_empty()
	}

	/// Get the first value in the set
	pub fn first(&self) -> Option<&Value> {
		self.0.first()
	}

	/// Get the last value in the set
	pub fn last(&self) -> Option<&Value> {
		self.0.last()
	}

	/// Get the nth value in the set
	pub fn nth(&self, index: usize) -> Option<&Value> {
		self.0.iter().nth(index)
	}

	/// Get a mutable reference to the first value in the set
	pub fn first_mut(&mut self) -> Option<&mut Value> {
		self.0.first_mut()
	}

	/// Get a mutable reference to the last value in the set
	pub fn last_mut(&mut self) -> Option<&mut Value> {
		self.0.last_mut()
	}

	/// Get a mutable reference to the nth value in the set
	pub fn nth_mut(&mut self, index: usize) -> Option<&mut Value> {
		self.0.get_mut(index)
	}

	/// Get an iterator over the values in the set
	pub fn iter(&self) -> impl Iterator<Item = &Value> {
		self.0.iter()
	}

	/// Insert a value into the set
	/// Returns true if the value was newly inserted
	pub fn insert(&mut self, value: Value) -> bool {
		self.0.insert(value)
	}

	/// Check if the set contains a value
	pub fn contains(&self, value: &Value) -> bool {
		self.0.contains(value)
	}

	/// Remove a value from the set
	/// Returns true if the value was present
	pub fn remove(&mut self, value: &Value) -> bool {
		self.0.remove(value)
	}

	/// Convert into a literal expression
	pub fn into_literal(self) -> Vec<Expr> {
		self.0.into_iter().map(Value::into_literal).collect()
	}

	/// Return the union of this set with another (A ∪ B)
	pub fn union(&self, other: &Set) -> Set {
		Set(self.0.union(&other.0))
	}

	/// Return the intersection of this set with another (A ∩ B)
	pub fn intersection(&self, other: &Set) -> Set {
		Set(self.0.intersection(&other.0))
	}

	/// Return the symmetric difference (A △ B) - elements in either but not both
	pub fn symmetric_difference(&self, other: &Set) -> Set {
		Set(self.0.symmetric_difference(&other.0))
	}

	/// Return the relative complement (A \ B) - elements in self but not in other
	pub fn complement(&self, other: &Set) -> Set {
		Set(self.0.difference(&other.0))
	}

	/// Rebuild the set so equal values collapse to one entry.
	///
	/// Used after in-place element transformations
	/// (for example a `.*` field VALUE clause)
	pub fn normalize(self) -> Self {
		Set::from(Vec::from(self))
	}

	/// Flatten nested sets and arrays into a single set
	pub fn flatten(self) -> Set {
		let mut out = Set::new();
		for v in self {
			match v {
				Value::Array(arr) => {
					for item in arr.0 {
						out.insert(item);
					}
				}
				Value::Set(set) => {
					for item in set.0 {
						out.insert(item);
					}
				}
				_ => {
					out.insert(v);
				}
			}
		}
		out
	}
}

impl<T> From<Vec<T>> for Set
where
	Value: From<T>,
{
	fn from(v: Vec<T>) -> Self {
		Set(v.into_iter().map(Value::from).collect())
	}
}

impl From<std::collections::BTreeSet<Value>> for Set {
	fn from(set: std::collections::BTreeSet<Value>) -> Self {
		Set(set.into())
	}
}

impl From<VecSet<Value>> for Set {
	fn from(set: VecSet<Value>) -> Self {
		Set(set)
	}
}

impl From<Set> for Vec<Value> {
	fn from(s: Set) -> Self {
		s.0.into_iter().collect()
	}
}

impl TryFrom<Set> for surrealdb_types::Set {
	type Error = anyhow::Error;

	fn try_from(s: Set) -> Result<Self, Self::Error> {
		Ok(surrealdb_types::Set::from(
			s.0.into_iter().map(surrealdb_types::Value::try_from).collect::<Result<Vec<_>, _>>()?,
		))
	}
}

impl From<surrealdb_types::Set> for Set {
	fn from(s: surrealdb_types::Set) -> Self {
		Set(s.into_iter().map(Value::from).collect())
	}
}

impl FromIterator<Value> for Set {
	fn from_iter<I: IntoIterator<Item = Value>>(iter: I) -> Self {
		Set(iter.into_iter().collect())
	}
}

impl IntoIterator for Set {
	type Item = Value;
	type IntoIter = VecSetIntoIter<Value>;
	fn into_iter(self) -> Self::IntoIter {
		self.0.into_iter()
	}
}

impl ToSql for Set {
	fn fmt_sql(&self, f: &mut String, sql_fmt: SqlFormat) {
		if self.is_empty() {
			return f.push_str("{,}");
		}

		// Format as Python-style set literal: `{,}`, `{val,}`, `{val, val, val}`
		f.push('{');
		let len = self.len();
		for (i, v) in self.iter().enumerate() {
			write_sql!(f, sql_fmt, "{}", v);
			// If this is not the last element, add a comma.
			// If this is the first element, add a comma.
			if len == 1 {
				f.push(',');
			} else if i < len - 1 {
				f.push_str(", ");
			}
		}
		f.push('}');
	}
}