surrealdb_expr/expr/
mock.rs

1use surrealdb_types::{SqlFormat, ToSql};
2
3use crate::val::range::{IntegerRangeIter, TypedRange};
4use crate::val::{RecordId, RecordIdKey, TableName};
5
6pub struct IntoIter {
7	table: TableName,
8	key: IntoIterKey,
9}
10
11pub enum IntoIterKey {
12	Count(i64),
13	Range(IntegerRangeIter),
14}
15
16impl Iterator for IntoIter {
17	type Item = RecordId;
18
19	fn next(&mut self) -> Option<RecordId> {
20		match self.key {
21			IntoIterKey::Count(ref mut c) => {
22				if *c == 0 {
23					None
24				} else {
25					*c -= 1;
26					Some(RecordId {
27						table: self.table.clone(),
28						key: RecordIdKey::rand(),
29					})
30				}
31			}
32			IntoIterKey::Range(ref mut r) => {
33				let k = r.next()?;
34				Some(RecordId {
35					table: self.table.clone(),
36					key: RecordIdKey::from(k),
37				})
38			}
39		}
40	}
41
42	fn size_hint(&self) -> (usize, Option<usize>) {
43		const MAX: usize = const {
44			match usize::BITS {
45				64 => i64::MAX as usize,
46				32 => usize::MAX,
47				_ => panic!("unsupported pointer width"),
48			}
49		};
50
51		match &self.key {
52			IntoIterKey::Count(x) => {
53				let x = *x;
54				let lower = x.min(MAX as i64) as usize;
55				let upper = if x < MAX as i64 {
56					Some(x as usize)
57				} else {
58					None
59				};
60				(lower, upper)
61			}
62			IntoIterKey::Range(r) => r.size_hint(),
63		}
64	}
65}
66
67#[derive(Clone, Debug, Eq, PartialEq, Hash)]
68pub enum Mock {
69	Count(TableName, i64),
70	Range(TableName, TypedRange<i64>),
71}
72
73impl Mock {
74	pub fn table(&self) -> &TableName {
75		match self {
76			Mock::Count(t, _) => t,
77			Mock::Range(t, _) => t,
78		}
79	}
80}
81
82impl IntoIterator for Mock {
83	type Item = RecordId;
84	type IntoIter = IntoIter;
85	fn into_iter(self) -> Self::IntoIter {
86		match self {
87			Mock::Count(t, k) => IntoIter {
88				table: t,
89				key: IntoIterKey::Count(k.max(0)),
90			},
91			Mock::Range(t, r) => IntoIter {
92				table: t,
93				key: IntoIterKey::Range(r.iter()),
94			},
95		}
96	}
97}
98
99impl ToSql for Mock {
100	fn fmt_sql(&self, f: &mut String, sql_fmt: SqlFormat) {
101		let mock: crate::sql::Mock = self.clone().into();
102		mock.fmt_sql(f, sql_fmt);
103	}
104}
surrealdb_expr/expr/mock.rs

surrealdb_expr/expr/
mock.rs
