chore(vm): Pass database address to the compiler (#637)

The database address will be needed to optimize joins.
In particular it will be used to look up table size metrics.

crates/core/src/sql/compiler.rs

@@ -24,7 +24,7 @@ pub fn compile_sql(db: &RelationalDB, tx: &MutTxId, sql_text: &str) -> Result<Ve
     let mut results = Vec::with_capacity(ast.len());
 
     for sql in ast {
-        results.push(compile_statement(sql).map_err(|error| DBError::Plan {
+        results.push(compile_statement(db, sql).map_err(|error| DBError::Plan {
             sql: sql_text.to_string(),
             error,
         })?);
@@ -248,7 +248,7 @@ fn compile_drop(name: String, kind: DbType, table_access: StAccess) -> Result<Cr
 }
 
 /// Compiles a `SQL` clause
-fn compile_statement(statement: SqlAst) -> Result<CrudExpr, PlanError> {
+fn compile_statement(db: &RelationalDB, statement: SqlAst) -> Result<CrudExpr, PlanError> {
     let q = match statement {
         SqlAst::Select {
             from,
@@ -270,7 +270,7 @@ fn compile_statement(statement: SqlAst) -> Result<CrudExpr, PlanError> {
         } => compile_drop(name, kind, table_access)?,
     };
 
-    Ok(q.optimize())
+    Ok(q.optimize(Some(db.address())))
 }
 
 #[cfg(test)]
@@ -1052,7 +1052,7 @@ mod tests {
 
         // Optimize the query plan for the incremental update.
         let expr = query::to_mem_table(expr, &insert);
-        let expr = expr.optimize();
+        let expr = expr.optimize(Some(db.address()));
 
         let QueryExpr {
             source:

crates/core/src/subscription/subscription.rs

@@ -591,7 +591,7 @@ impl<'a> IncrementalJoin<'a> {
         let mut inserts = {
             // Replan query after replacing left table with virtual table,
             // since join order may need to be reversed.
-            let lhs_virt = query::to_mem_table(self.expr.clone(), &self.lhs.inserts()).optimize();
+            let lhs_virt = query::to_mem_table(self.expr.clone(), &self.lhs.inserts()).optimize(Some(db.address()));
             let rhs_virt = self.to_mem_table_rhs(self.rhs.inserts());
 
             // {A+ join B}
@@ -617,7 +617,7 @@ impl<'a> IncrementalJoin<'a> {
         let mut deletes = {
             // Replan query after replacing left table with virtual table,
             // since join order may need to be reversed.
-            let lhs_virt = query::to_mem_table(self.expr.clone(), &self.lhs.deletes()).optimize();
+            let lhs_virt = query::to_mem_table(self.expr.clone(), &self.lhs.deletes()).optimize(Some(db.address()));
             let rhs_virt = self.to_mem_table_rhs(self.rhs.deletes());
 
             // {A- join B}

crates/core/src/vm.rs

@@ -3,6 +3,7 @@ use std::collections::HashMap;
 use std::ops::RangeBounds;
 
 use itertools::Itertools;
+use spacetimedb_lib::Address;
 use tracing::debug;
 
 use crate::db::cursor::{CatalogCursor, IndexCursor, TableCursor};
@@ -395,6 +396,10 @@ impl<'db, 'tx> DbProgram<'db, 'tx> {
 }
 
 impl ProgramVm for DbProgram<'_, '_> {
+    fn address(&self) -> Option<Address> {
+        Some(self.db.address())
+    }
+
     fn env(&self) -> &EnvDb {
         &self.env
     }

crates/vm/src/eval.rs

@@ -77,7 +77,7 @@ fn build_typed<P: ProgramVm>(p: &mut P, node: Expr) -> ExprOpt {
             }
         }
         Expr::Crud(q) => {
-            let q = q.optimize();
+            let q = q.optimize(p.address());
             match q {
                 CrudExpr::Query(q) => {
                     let source = build_query_opt(q);

crates/vm/src/expr.rs

@@ -9,7 +9,7 @@ use crate::errors::{ErrorKind, ErrorLang, ErrorType, ErrorVm};
 use crate::functions::{FunDef, Param};
 use crate::operator::{Op, OpCmp, OpLogic, OpQuery};
 use crate::types::Ty;
-use spacetimedb_lib::Identity;
+use spacetimedb_lib::{Address, Identity};
 use spacetimedb_primitives::*;
 use spacetimedb_sats::algebraic_type::AlgebraicType;
 use spacetimedb_sats::algebraic_value::AlgebraicValue;
@@ -476,9 +476,9 @@ pub enum CrudExpr {
 }
 
 impl CrudExpr {
-    pub fn optimize(self) -> Self {
+    pub fn optimize(self, db: Option<Address>) -> Self {
         match self {
-            CrudExpr::Query(x) => CrudExpr::Query(x.optimize()),
+            CrudExpr::Query(x) => CrudExpr::Query(x.optimize(db)),
             _ => self,
         }
     }
@@ -1129,7 +1129,7 @@ impl QueryExpr {
     //
     // Ex. SELECT Left.* FROM Left JOIN Right ON Left.id = Right.id ...
     // where `Left` has an index defined on `id`.
-    fn try_index_join(mut query: QueryExpr) -> QueryExpr {
+    fn try_index_join(mut query: QueryExpr, _db: Option<Address>) -> QueryExpr {
         // We expect 2 and only 2 operations - a join followed by a wildcard projection.
         if query.query.len() != 2 {
             return query;
@@ -1325,7 +1325,7 @@ impl QueryExpr {
         })
     }
 
-    pub fn optimize(self) -> Self {
+    pub fn optimize(self, db: Option<Address>) -> Self {
         let mut q = Self {
             source: self.source.clone(),
             query: Vec::with_capacity(self.query.len()),
@@ -1350,11 +1350,11 @@ impl QueryExpr {
                 Query::Select(op) => {
                     q = Self::optimize_select(q, op, &tables);
                 }
-                Query::JoinInner(join) => q = q.with_join_inner(join.rhs.optimize(), join.col_lhs, join.col_rhs),
+                Query::JoinInner(join) => q = q.with_join_inner(join.rhs.optimize(db), join.col_lhs, join.col_rhs),
                 _ => q.query.push(query),
             };
         }
-        Self::try_index_join(q)
+        Self::try_index_join(q, db)
     }
 }
 

crates/vm/src/program.rs

@@ -2,6 +2,7 @@
 //!
 //! It carries an [EnvDb] with the functions, idents, types.
 use spacetimedb_lib::identity::AuthCtx;
+use spacetimedb_lib::Address;
 use spacetimedb_sats::relation::{MemTable, RelIter, Relation, Table};
 use std::collections::HashMap;
 
@@ -66,6 +67,7 @@ pub trait ProgramVm {
         env.functions.ops = ops
     }
 
+    fn address(&self) -> Option<Address>;
     fn env(&self) -> &EnvDb;
     fn env_mut(&mut self) -> &mut EnvDb;
     fn ctx(&self) -> &dyn ProgramVm;
@@ -151,6 +153,10 @@ impl Program {
 }
 
 impl ProgramVm for Program {
+    fn address(&self) -> Option<Address> {
+        None
+    }
+
     fn env(&self) -> &EnvDb {
         &self.env
     }