Revert "fix(universaldb): bound rocksdb range scans"

This reverts commit 7835a88.

Author: NathanFlurry

engine/packages/universaldb/src/driver/rocksdb/transaction_task.rs

@@ -12,7 +12,7 @@ use crate::{
 	error::DatabaseError,
 	key_selector::KeySelector,
 	options::{ConflictRangeType, MutationType},
-	tx_ops::{Operation, range_begin_contains, range_end_contains},
+	tx_ops::Operation,
 	value::{KeyValue, Slice, Values},
 	versionstamp::{
 		generate_versionstamp, substitute_raw_versionstamp, substitute_versionstamp_if_incomplete,
@@ -438,8 +438,17 @@ impl TransactionTask {
 		let txn = self.create_transaction();
 		let read_opts = ReadOptions::default();
 
+		// Resolve the begin selector
+		let resolved_begin =
+			self.resolve_key_selector_for_range(&txn, &begin, begin_or_equal, begin_offset)?;
+
+		// Resolve the end selector
+		let resolved_end =
+			self.resolve_key_selector_for_range(&txn, &end, end_or_equal, end_offset)?;
+
+		// Now execute the range query with resolved keys
 		let iter = txn.iterator_opt(
-			rocksdb::IteratorMode::From(&begin, rocksdb::Direction::Forward),
+			rocksdb::IteratorMode::From(&resolved_begin, rocksdb::Direction::Forward),
 			read_opts,
 		);
 
@@ -448,12 +457,8 @@ impl TransactionTask {
 
 		for item in iter {
 			let (k, v) = item.context("failed to iterate rocksdb for get range")?;
-
-			if !range_begin_contains(k.as_ref(), &begin, begin_or_equal, begin_offset) {
-				continue;
-			}
-
-			if !range_end_contains(k.as_ref(), &end, end_or_equal, end_offset) {
+			// Check if we've reached the end key
+			if k.as_ref() >= resolved_end.as_slice() {
 				break;
 			}
 
@@ -472,6 +477,64 @@ impl TransactionTask {
 		Ok(Values::new(results))
 	}
 
+	fn resolve_key_selector_for_range(
+		&self,
+		txn: &RocksDbTransaction<OptimisticTransactionDB>,
+		key: &[u8],
+		or_equal: bool,
+		offset: i32,
+	) -> Result<Vec<u8>> {
+		// Based on PostgreSQL's interpretation:
+		// (false, 1) => first_greater_or_equal
+		// (true, 1) => first_greater_than
+		// (false, 0) => last_less_than
+		// (true, 0) => last_less_or_equal
+
+		let read_opts = ReadOptions::default();
+
+		match (or_equal, offset) {
+			(false, 1) => {
+				// first_greater_or_equal: find first key >= search_key
+				let iter = txn.iterator_opt(
+					rocksdb::IteratorMode::From(key, rocksdb::Direction::Forward),
+					read_opts,
+				);
+				for item in iter {
+					let (k, _v) = item.context(
+						"failed to iterate rocksdb for range selector first_greater_or_equal",
+					)?;
+					return Ok(k.to_vec());
+				}
+				// If no key found, return a key that will make the range empty
+				Ok(vec![0xff; 255])
+			}
+			(true, 1) => {
+				// first_greater_than: find first key > search_key
+				let iter = txn.iterator_opt(
+					rocksdb::IteratorMode::From(key, rocksdb::Direction::Forward),
+					read_opts,
+				);
+				for item in iter {
+					let (k, _v) = item.context(
+						"failed to iterate rocksdb for range selector first_greater_than",
+					)?;
+					// Skip if it's the exact key
+					if k.as_ref() == key {
+						continue;
+					}
+					return Ok(k.to_vec());
+				}
+				// If no key found, return a key that will make the range empty
+				Ok(vec![0xff; 255])
+			}
+			_ => {
+				// For other cases, just use the key as-is for now
+				// This is a simplification - full implementation would handle all cases
+				Ok(key.to_vec())
+			}
+		}
+	}
+
 	async fn handle_get_estimated_range_size(&mut self, begin: &[u8], end: &[u8]) -> Result<i64> {
 		let range = rocksdb::Range::new(begin, end);
 

engine/packages/universaldb/src/tx_ops.rs

@@ -342,6 +342,9 @@ impl TransactionOperations {
 			return Ok(db_values);
 		}
 
+		let begin = opt.begin.key();
+		let end = opt.end.key();
+
 		// Start with database results in a map
 		let mut result_map = BTreeMap::new();
 		for kv in db_values.into_iter() {
@@ -354,7 +357,7 @@ impl TransactionOperations {
 		for op in &*self.operations() {
 			match op {
 				Operation::Set { key, value } => {
-					if range_contains(key.as_slice(), opt) {
+					if key.as_slice() >= begin && key.as_slice() < end {
 						result_map.insert(key.clone(), value.clone());
 					}
 				}
@@ -379,7 +382,7 @@ impl TransactionOperations {
 					param,
 					op_type,
 				} => {
-					if range_contains(key.as_slice(), opt) {
+					if key.as_slice() >= begin && key.as_slice() < end {
 						// Get current value for this key (from result_map or empty if not exists)
 						let current_value = result_map.get(key);
 						let current_slice = current_value.map(|v| &**v);
@@ -420,32 +423,3 @@ impl TransactionOperations {
 			.push((begin.to_vec(), end.to_vec(), conflict_type));
 	}
 }
-
-fn range_contains(key: &[u8], opt: &RangeOption<'_>) -> bool {
-	range_begin_contains(
-		key,
-		opt.begin.key(),
-		opt.begin.or_equal(),
-		opt.begin.offset(),
-	) && range_end_contains(key, opt.end.key(), opt.end.or_equal(), opt.end.offset())
-}
-
-pub(crate) fn range_begin_contains(key: &[u8], begin: &[u8], or_equal: bool, offset: i32) -> bool {
-	match (or_equal, offset) {
-		(false, 1) => key >= begin,
-		(true, 1) => key > begin,
-		(false, 0) => key > begin,
-		(true, 0) => key >= begin,
-		_ => key >= begin,
-	}
-}
-
-pub(crate) fn range_end_contains(key: &[u8], end: &[u8], or_equal: bool, offset: i32) -> bool {
-	match (or_equal, offset) {
-		(false, 1) => key < end,
-		(true, 1) => key <= end,
-		(false, 0) => key < end,
-		(true, 0) => key <= end,
-		_ => key < end,
-	}
-}

engine/packages/universaldb/tests/integration.rs

@@ -705,41 +705,6 @@ async fn test_range_options(db: &Database) {
 	assert_eq!(results[2].key(), test_subspace.pack(&("range_e",)));
 	assert_eq!(results[2].value(), b"val_e");
 
-	// Test 5: local writes outside the range should not be merged into the result
-	let results = db
-		.run(|tx| async move {
-			let test_subspace = Subspace::from("test");
-			let key_b = test_subspace.pack(&("range_b",));
-			let key_d = test_subspace.pack(&("range_d",));
-			let key_z = test_subspace.pack(&("range_z",));
-
-			tx.set(&key_z, b"val_z");
-
-			let range = RangeOption {
-				begin: KeySelector::first_greater_or_equal(Cow::Owned(key_b)),
-				end: KeySelector::first_greater_or_equal(Cow::Owned(key_d)),
-				limit: None,
-				reverse: false,
-				mode: StreamingMode::WantAll,
-				target_bytes: 0,
-				..RangeOption::default()
-			};
-
-			let vals = tx.get_range(&range, 1, Serializable).await?;
-			Ok(vals.into_vec())
-		})
-		.await
-		.unwrap();
-
-	assert_eq!(
-		results.len(),
-		2,
-		"Expected local key outside [b, d) to be excluded"
-	);
-	assert!(!results
-		.iter()
-		.any(|r| r.key() == test_subspace.pack(&("range_z",))));
-
 	// Clear test data
 	db.run(|tx| async move {
 		let test_subspace = Subspace::from("test");