spacetimedb_vm/
eval.rs

1use crate::errors::ErrorVm;
2use crate::expr::{Code, ColumnOp, Expr, JoinExpr, ProjectExpr, SourceSet};
3use crate::program::{ProgramVm, Sources};
4use crate::rel_ops::RelOps;
5use crate::relation::RelValue;
6use spacetimedb_sats::ProductValue;
7use spacetimedb_table::table::RowRef;
8
9pub type IterRows<'a> = dyn RelOps<'a> + 'a;
10
11/// Utility to simplify the creation of a boxed iterator.
12pub fn box_iter<'a, T: Iterator<Item = RowRef<'a>> + 'a>(iter: T) -> Box<dyn Iterator<Item = RowRef<'a>> + 'a> {
13    Box::new(iter)
14}
15
16pub fn build_select<'a>(base: impl RelOps<'a> + 'a, cmp: &'a ColumnOp) -> Box<IterRows<'a>> {
17    Box::new(base.select(move |row| cmp.eval_bool(row)))
18}
19
20pub fn build_project<'a>(base: impl RelOps<'a> + 'a, proj: &'a ProjectExpr) -> Box<IterRows<'a>> {
21    Box::new(base.project(&proj.cols, move |cols, row| {
22        RelValue::Projection(row.project_owned(cols))
23    }))
24}
25
26pub fn join_inner<'a>(lhs: impl RelOps<'a> + 'a, rhs: impl RelOps<'a> + 'a, q: &'a JoinExpr) -> Box<IterRows<'a>> {
27    let col_lhs = q.col_lhs.idx();
28    let col_rhs = q.col_rhs.idx();
29    let key_lhs = move |row: &RelValue<'_>| row.read_column(col_lhs).unwrap().into_owned();
30    let key_rhs = move |row: &RelValue<'_>| row.read_column(col_rhs).unwrap().into_owned();
31    let pred = move |l: &RelValue<'_>, r: &RelValue<'_>| l.read_column(col_lhs) == r.read_column(col_rhs);
32
33    if q.inner.is_some() {
34        Box::new(lhs.join_inner(rhs, key_lhs, key_rhs, pred, move |l, r| l.extend(r)))
35    } else {
36        Box::new(lhs.join_inner(rhs, key_lhs, key_rhs, pred, move |l, _| l))
37    }
38}
39
40/// Execute the code
41pub fn eval<const N: usize, P: ProgramVm>(p: &mut P, code: Code, sources: Sources<'_, N>) -> Code {
42    match code {
43        c @ (Code::Value(_) | Code::Halt(_) | Code::Table(_)) => c,
44        Code::Block(lines) => {
45            let mut result = Vec::with_capacity(lines.len());
46            for x in lines {
47                match eval(p, x, sources) {
48                    Code::Pass(None) => {}
49                    r => result.push(r),
50                };
51            }
52
53            match result.len() {
54                0 => Code::Pass(None),
55                1 => result.pop().unwrap(),
56                _ => Code::Block(result),
57            }
58        }
59        Code::Crud(q) => p.eval_query(q, sources).unwrap_or_else(|err| Code::Halt(err.into())),
60        Code::Pass(x) => Code::Pass(x),
61    }
62}
63
64fn to_vec(of: Vec<Expr>) -> Code {
65    let mut new = Vec::with_capacity(of.len());
66    for ast in of {
67        let code = match ast {
68            Expr::Block(x) => to_vec(x),
69            Expr::Crud(x) => Code::Crud(*x),
70            x => Code::Halt(ErrorVm::Unsupported(format!("{x:?}")).into()),
71        };
72        new.push(code);
73    }
74    Code::Block(new)
75}
76
77/// Optimize, compile & run the [Expr]
78pub fn run_ast<const N: usize, P: ProgramVm>(
79    p: &mut P,
80    ast: Expr,
81    mut sources: SourceSet<Vec<ProductValue>, N>,
82) -> Code {
83    let code = match ast {
84        Expr::Block(x) => to_vec(x),
85        Expr::Crud(x) => Code::Crud(*x),
86        Expr::Value(x) => Code::Value(x),
87        Expr::Halt(err) => Code::Halt(err),
88        Expr::Ident(x) => Code::Halt(ErrorVm::Unsupported(format!("Ident {x}")).into()),
89    };
90    eval(p, code, &mut sources)
91}
92
93/// Used internally for testing SQL JOINS.
94#[doc(hidden)]
95pub mod test_helpers {
96    use crate::relation::MemTable;
97    use core::hash::BuildHasher as _;
98    use spacetimedb_data_structures::map::DefaultHashBuilder;
99    use spacetimedb_primitives::TableId;
100    use spacetimedb_sats::{product, AlgebraicType, AlgebraicValue, ProductType, ProductValue};
101    use spacetimedb_schema::relation::{Column, FieldName, Header};
102    use std::sync::Arc;
103
104    pub fn mem_table_without_table_name(mem: &MemTable) -> (&[Column], &[ProductValue]) {
105        (&mem.head.fields, &mem.data)
106    }
107
108    pub fn header_for_mem_table(table_id: TableId, fields: ProductType) -> Header {
109        let hash = DefaultHashBuilder::default().hash_one(&fields);
110        let table_name = format!("mem#{hash:x}").into();
111
112        let cols = Vec::from(fields.elements)
113            .into_iter()
114            .enumerate()
115            .map(|(pos, f)| Column::new(FieldName::new(table_id, pos.into()), f.algebraic_type))
116            .collect();
117
118        Header::new(table_id, table_name, cols, Vec::new())
119    }
120
121    pub fn mem_table_one_u64(table_id: TableId) -> MemTable {
122        let ty = ProductType::from([AlgebraicType::U64]);
123        mem_table(table_id, ty, product![1u64])
124    }
125
126    pub fn mem_table<T: Into<ProductValue>>(
127        table_id: TableId,
128        ty: impl Into<ProductType>,
129        iter: impl IntoIterator<Item = T>,
130    ) -> MemTable {
131        let head = header_for_mem_table(table_id, ty.into());
132        MemTable::from_iter(Arc::new(head), iter.into_iter().map(Into::into))
133    }
134
135    pub fn scalar(of: impl Into<AlgebraicValue>) -> AlgebraicValue {
136        of.into()
137    }
138
139    pub struct GameData {
140        pub location: MemTable,
141        pub inv: MemTable,
142        pub player: MemTable,
143        pub location_ty: ProductType,
144        pub inv_ty: ProductType,
145        pub player_ty: ProductType,
146    }
147
148    pub fn create_game_data() -> GameData {
149        let inv_ty = ProductType::from([("inventory_id", AlgebraicType::U64), ("name", AlgebraicType::String)]);
150        let row = product!(1u64, "health");
151        let inv = mem_table(0.into(), inv_ty.clone(), [row]);
152
153        let player_ty = ProductType::from([("entity_id", AlgebraicType::U64), ("inventory_id", AlgebraicType::U64)]);
154        let row1 = product!(100u64, 1u64);
155        let row2 = product!(200u64, 1u64);
156        let row3 = product!(300u64, 1u64);
157        let player = mem_table(1.into(), player_ty.clone(), [row1, row2, row3]);
158
159        let location_ty = ProductType::from([
160            ("entity_id", AlgebraicType::U64),
161            ("x", AlgebraicType::F32),
162            ("z", AlgebraicType::F32),
163        ]);
164        let row1 = product!(100u64, 0.0f32, 32.0f32);
165        let row2 = product!(100u64, 1.0f32, 31.0f32);
166        let location = mem_table(2.into(), location_ty.clone(), [row1, row2]);
167
168        GameData {
169            location,
170            inv,
171            player,
172            inv_ty,
173            player_ty,
174            location_ty,
175        }
176    }
177}
178
179#[cfg(test)]
180pub mod tests {
181    #![allow(clippy::disallowed_macros)]
182
183    use super::test_helpers::*;
184    use super::*;
185    use crate::expr::{CrudExpr, Query, QueryExpr, SourceExpr, SourceSet};
186    use crate::iterators::RelIter;
187    use crate::relation::MemTable;
188    use spacetimedb_lib::operator::{OpCmp, OpLogic};
189    use spacetimedb_primitives::ColId;
190    use spacetimedb_sats::{product, AlgebraicType, ProductType};
191    use spacetimedb_schema::def::error::RelationError;
192    use spacetimedb_schema::relation::{FieldName, Header};
193
194    /// From an original source of `result`s, applies `queries` and returns a final set of results.
195    fn build_query<'a, const N: usize>(
196        mut result: Box<IterRows<'a>>,
197        queries: &'a [Query],
198        sources: Sources<'_, N>,
199    ) -> Box<IterRows<'a>> {
200        for q in queries {
201            result = match q {
202                Query::IndexScan(_) | Query::IndexJoin(_) => panic!("unsupported on memory tables"),
203                Query::Select(cmp) => build_select(result, cmp),
204                Query::Project(proj) => build_project(result, proj),
205                Query::JoinInner(q) => {
206                    let rhs = build_source_expr_query(sources, &q.rhs.source);
207                    let rhs = build_query(rhs, &q.rhs.query, sources);
208                    join_inner(result, rhs, q)
209                }
210            };
211        }
212        result
213    }
214
215    fn build_source_expr_query<'a, const N: usize>(sources: Sources<'_, N>, source: &SourceExpr) -> Box<IterRows<'a>> {
216        let source_id = source.source_id().unwrap();
217        let table = sources.take(source_id).unwrap();
218        Box::new(RelIter::new(table.into_iter().map(RelValue::Projection)))
219    }
220
221    /// A default program that run in-memory without a database
222    struct Program;
223
224    impl ProgramVm for Program {
225        fn eval_query<const N: usize>(&mut self, query: CrudExpr, sources: Sources<'_, N>) -> Result<Code, ErrorVm> {
226            match query {
227                CrudExpr::Query(query) => {
228                    let result = build_source_expr_query(sources, &query.source);
229                    let rows = build_query(result, &query.query, sources).collect_vec(|row| row.into_product_value());
230
231                    let head = query.head().clone();
232
233                    Ok(Code::Table(MemTable::new(head, query.source.table_access(), rows)))
234                }
235                _ => todo!(),
236            }
237        }
238    }
239
240    fn run_query<const N: usize>(ast: Expr, sources: SourceSet<Vec<ProductValue>, N>) -> MemTable {
241        match run_ast(&mut Program, ast, sources) {
242            Code::Table(x) => x,
243            x => panic!("Unexpected result on query: {x}"),
244        }
245    }
246
247    fn get_field_pos(table: &MemTable, pos: usize) -> FieldName {
248        *table.head.fields.get(pos).map(|x| &x.field).unwrap()
249    }
250
251    #[test]
252    fn test_select() {
253        let input = mem_table_one_u64(0.into());
254        let field = get_field_pos(&input, 0);
255        let mut sources = SourceSet::<_, 1>::empty();
256        let source_expr = sources.add_mem_table(input);
257
258        let q = QueryExpr::new(source_expr)
259            .with_select_cmp(OpCmp::Eq, field, scalar(1u64))
260            .unwrap();
261
262        let head = q.head().clone();
263
264        let result = run_query(q.into(), sources);
265        let row = product![1u64];
266        assert_eq!(result, MemTable::from_iter(head, [row]), "Query");
267    }
268
269    #[test]
270    fn test_project() {
271        let p = &mut Program;
272        let table = mem_table_one_u64(0.into());
273
274        let mut sources = SourceSet::<_, 1>::empty();
275        let source_expr = sources.add_mem_table(table.clone());
276
277        let source = QueryExpr::new(source_expr);
278        let field = get_field_pos(&table, 0);
279        let q = source.clone().with_project([field.into()].into(), None).unwrap();
280        let head = q.head().clone();
281
282        let result = run_ast(p, q.into(), sources);
283        let row = product![1u64];
284        assert_eq!(result, Code::Table(MemTable::from_iter(head.clone(), [row])), "Project");
285    }
286
287    #[test]
288    fn test_project_out_of_bounds() {
289        let table = mem_table_one_u64(0.into());
290
291        let mut sources = SourceSet::<_, 1>::empty();
292        let source_expr = sources.add_mem_table(table.clone());
293
294        let source = QueryExpr::new(source_expr);
295        // This field is out of bounds of `table`'s header, so `run_ast` will panic.
296        let field = FieldName::new(table.head.table_id, 1.into());
297        assert!(matches!(
298            source.with_project([field.into()].into(), None).unwrap_err(),
299            RelationError::FieldNotFound(_, f) if f == field,
300        ));
301    }
302
303    #[test]
304    fn test_join_inner() {
305        let table_id = 0.into();
306        let table = mem_table_one_u64(table_id);
307        let col: ColId = 0.into();
308        let field = table.head.fields[col.idx()].clone();
309
310        let mut sources = SourceSet::<_, 2>::empty();
311        let source_expr = sources.add_mem_table(table.clone());
312        let second_source_expr = sources.add_mem_table(table);
313
314        let q = QueryExpr::new(source_expr).with_join_inner(second_source_expr, col, col, false);
315        let result = run_query(q.into(), sources);
316
317        // The expected result.
318        let head = Header::new(table_id, "".into(), [field.clone(), field].into(), Vec::new());
319        let input = MemTable::from_iter(head.into(), [product!(1u64, 1u64)]);
320
321        println!("{}", &result.head);
322        println!("{}", &input.head);
323
324        assert_eq!(
325            mem_table_without_table_name(&result),
326            mem_table_without_table_name(&input),
327            "Project"
328        );
329    }
330
331    #[test]
332    fn test_semijoin() {
333        let table_id = 0.into();
334        let table = mem_table_one_u64(table_id);
335        let col = 0.into();
336
337        let mut sources = SourceSet::<_, 2>::empty();
338        let source_expr = sources.add_mem_table(table.clone());
339        let second_source_expr = sources.add_mem_table(table);
340
341        let q = QueryExpr::new(source_expr).with_join_inner(second_source_expr, col, col, true);
342        let result = run_query(q.into(), sources);
343
344        // The expected result.
345        let inv = ProductType::from([(None, AlgebraicType::U64)]);
346        let input = mem_table(table_id, inv, [product![1u64]]);
347
348        println!("{}", &result.head);
349        println!("{}", &input.head);
350
351        assert_eq!(
352            mem_table_without_table_name(&result),
353            mem_table_without_table_name(&input),
354            "Semijoin should not be projected",
355        );
356    }
357
358    #[test]
359    fn test_query_logic() {
360        let inv = ProductType::from([("id", AlgebraicType::U64), ("name", AlgebraicType::String)]);
361
362        let row = product![1u64, "health"];
363
364        let input = mem_table(0.into(), inv, vec![row]);
365        let inv = input.clone();
366
367        let mut sources = SourceSet::<_, 1>::empty();
368        let source_expr = sources.add_mem_table(input.clone());
369
370        let q = QueryExpr::new(source_expr.clone())
371            .with_select_cmp(OpLogic::And, scalar(true), scalar(true))
372            .unwrap();
373
374        let result = run_query(q.into(), sources);
375
376        assert_eq!(result, inv.clone(), "Query And");
377
378        let mut sources = SourceSet::<_, 1>::empty();
379        let source_expr = sources.add_mem_table(input);
380
381        let q = QueryExpr::new(source_expr)
382            .with_select_cmp(OpLogic::Or, scalar(true), scalar(false))
383            .unwrap();
384
385        let result = run_query(q.into(), sources);
386
387        assert_eq!(result, inv, "Query Or");
388    }
389
390    #[test]
391    /// Inventory
392    /// | id: u64 | name : String |
393    fn test_query_inner_join() {
394        let inv = ProductType::from([("id", AlgebraicType::U64), ("name", AlgebraicType::String)]);
395
396        let row = product![1u64, "health"];
397
398        let table_id = 0.into();
399        let input = mem_table(table_id, inv, [row]);
400        let col = 0.into();
401
402        let mut sources = SourceSet::<_, 2>::empty();
403        let source_expr = sources.add_mem_table(input.clone());
404        let second_source_expr = sources.add_mem_table(input);
405
406        let q = QueryExpr::new(source_expr).with_join_inner(second_source_expr, col, col, false);
407
408        let result = run_query(q.into(), sources);
409
410        //The expected result
411        let inv = ProductType::from([
412            (None, AlgebraicType::U64),
413            (Some("id"), AlgebraicType::U64),
414            (Some("name"), AlgebraicType::String),
415        ]);
416        let row = product![1u64, "health", 1u64, "health"];
417        let input = mem_table(table_id, inv, vec![row]);
418        assert_eq!(result.data, input.data, "Project");
419    }
420
421    #[test]
422    /// Inventory
423    /// | id: u64 | name : String |
424    fn test_query_semijoin() {
425        let inv = ProductType::from([("id", AlgebraicType::U64), ("name", AlgebraicType::String)]);
426
427        let row = product![1u64, "health"];
428
429        let table_id = 0.into();
430        let input = mem_table(table_id, inv, [row]);
431        let col = 0.into();
432
433        let mut sources = SourceSet::<_, 2>::empty();
434        let source_expr = sources.add_mem_table(input.clone());
435        let second_source_expr = sources.add_mem_table(input);
436
437        let q = QueryExpr::new(source_expr).with_join_inner(second_source_expr, col, col, true);
438
439        let result = run_query(q.into(), sources);
440
441        // The expected result.
442        let inv = ProductType::from([(None, AlgebraicType::U64), (Some("name"), AlgebraicType::String)]);
443        let row = product![1u64, "health"];
444        let input = mem_table(table_id, inv, vec![row]);
445        assert_eq!(result.data, input.data, "Semijoin should not project");
446    }
447
448    #[test]
449    /// Inventory
450    /// | inventory_id: u64 | name : String |
451    /// Player
452    /// | entity_id: u64 | inventory_id : u64 |
453    /// Location
454    /// | entity_id: u64 | x : f32 | z : f32 |
455    fn test_query_game() {
456        // See table above.
457        let data = create_game_data();
458        let inv @ [inv_inventory_id, _] = [0, 1].map(|c| c.into());
459        let inv_head = data.inv.head.clone();
460        let inv_expr = |col: ColId| inv_head.fields[col.idx()].field.into();
461        let [location_entity_id, location_x, location_z] = [0, 1, 2].map(|c| c.into());
462        let [player_entity_id, player_inventory_id] = [0, 1].map(|c| c.into());
463        let loc_head = data.location.head.clone();
464        let loc_field = |col: ColId| loc_head.fields[col.idx()].field;
465        let inv_table_id = data.inv.head.table_id;
466        let player_table_id = data.player.head.table_id;
467
468        let mut sources = SourceSet::<_, 2>::empty();
469        let player_source_expr = sources.add_mem_table(data.player.clone());
470        let location_source_expr = sources.add_mem_table(data.location.clone());
471
472        // SELECT
473        // Player.*
474        //     FROM
475        // Player
476        // JOIN Location
477        // ON Location.entity_id = Player.entity_id
478        // WHERE x > 0 AND x <= 32 AND z > 0 AND z <= 32
479        let q = QueryExpr::new(player_source_expr)
480            .with_join_inner(location_source_expr, player_entity_id, location_entity_id, true)
481            .with_select_cmp(OpCmp::Gt, loc_field(location_x), scalar(0.0f32))
482            .unwrap()
483            .with_select_cmp(OpCmp::LtEq, loc_field(location_x), scalar(32.0f32))
484            .unwrap()
485            .with_select_cmp(OpCmp::Gt, loc_field(location_z), scalar(0.0f32))
486            .unwrap()
487            .with_select_cmp(OpCmp::LtEq, loc_field(location_z), scalar(32.0f32))
488            .unwrap();
489
490        let result = run_query(q.into(), sources);
491
492        let ty = ProductType::from([("entity_id", AlgebraicType::U64), ("inventory_id", AlgebraicType::U64)]);
493        let row1 = product!(100u64, 1u64);
494        let input = mem_table(player_table_id, ty, [row1]);
495
496        assert_eq!(
497            mem_table_without_table_name(&result),
498            mem_table_without_table_name(&input),
499            "Player"
500        );
501
502        let mut sources = SourceSet::<_, 3>::empty();
503        let player_source_expr = sources.add_mem_table(data.player);
504        let location_source_expr = sources.add_mem_table(data.location);
505        let inventory_source_expr = sources.add_mem_table(data.inv);
506
507        // SELECT
508        // Inventory.*
509        //     FROM
510        // Inventory
511        // JOIN Player
512        // ON Inventory.inventory_id = Player.inventory_id
513        // JOIN Location
514        // ON Player.entity_id = Location.entity_id
515        // WHERE x > 0 AND x <= 32 AND z > 0 AND z <= 32
516        let q = QueryExpr::new(inventory_source_expr)
517            // NOTE: The way this query is set up, the first join must be an inner join, not a semijoin,
518            // so that the second join has access to the `Player.entity_id` field.
519            // This necessitates a trailing `project` to get just `Inventory.*`.
520            .with_join_inner(player_source_expr, inv_inventory_id, player_inventory_id, false)
521            .with_join_inner(
522                location_source_expr,
523                (inv_head.fields.len() + player_entity_id.idx()).into(),
524                location_entity_id,
525                true,
526            )
527            .with_select_cmp(OpCmp::Gt, loc_field(location_x), scalar(0.0f32))
528            .unwrap()
529            .with_select_cmp(OpCmp::LtEq, loc_field(location_x), scalar(32.0f32))
530            .unwrap()
531            .with_select_cmp(OpCmp::Gt, loc_field(location_z), scalar(0.0f32))
532            .unwrap()
533            .with_select_cmp(OpCmp::LtEq, loc_field(location_z), scalar(32.0f32))
534            .unwrap()
535            .with_project(inv.map(inv_expr).into(), Some(inv_table_id))
536            .unwrap();
537
538        let result = run_query(q.into(), sources);
539
540        let ty = ProductType::from([("inventory_id", AlgebraicType::U64), ("name", AlgebraicType::String)]);
541        let row1 = product!(1u64, "health");
542        let input = mem_table(inv_table_id, ty, [row1]);
543
544        assert_eq!(
545            mem_table_without_table_name(&result),
546            mem_table_without_table_name(&input),
547            "Inventory"
548        );
549    }
550}
spacetimedb_vm/eval.rs

spacetimedb_vm/
eval.rs
