query rewriter (#2031)

crates/core/src/sql/parser.rs

@@ -3,14 +3,13 @@ use crate::db::relational_db::RelationalDB;
 use crate::sql::ast::SchemaViewer;
 use spacetimedb_expr::check::parse_and_type_sub;
 use spacetimedb_expr::errors::TypingError;
-use spacetimedb_expr::expr::RelExpr;
-use spacetimedb_expr::ty::TyCtx;
+use spacetimedb_expr::expr::Project;
 use spacetimedb_lib::db::raw_def::v9::RawRowLevelSecurityDefV9;
 use spacetimedb_lib::identity::AuthCtx;
 use spacetimedb_schema::schema::RowLevelSecuritySchema;
 
 pub struct RowLevelExpr {
-    pub sql: RelExpr,
+    pub sql: Project,
     pub def: RowLevelSecuritySchema,
 }
 
@@ -21,12 +20,11 @@ impl RowLevelExpr {
         auth_ctx: &AuthCtx,
         rls: &RawRowLevelSecurityDefV9,
     ) -> Result<Self, TypingError> {
-        let mut ctx = TyCtx::default();
-        let sql = parse_and_type_sub(&mut ctx, &rls.sql, &SchemaViewer::new(stdb, tx, auth_ctx))?;
+        let sql = parse_and_type_sub(&rls.sql, &SchemaViewer::new(stdb, tx, auth_ctx))?;
 
         Ok(Self {
             def: RowLevelSecuritySchema {
-                table_id: sql.table_id(&mut ctx)?,
+                table_id: sql.table_id().unwrap(),
                 sql: rls.sql.clone(),
             },
             sql,

crates/core/src/subscription/module_subscription_actor.rs

@@ -16,7 +16,6 @@ use crate::worker_metrics::WORKER_METRICS;
 use parking_lot::RwLock;
 use spacetimedb_client_api_messages::websocket::FormatSwitch;
 use spacetimedb_expr::check::compile_sql_sub;
-use spacetimedb_expr::ty::TyCtx;
 use spacetimedb_lib::identity::AuthCtx;
 use spacetimedb_lib::Identity;
 use spacetimedb_vm::errors::ErrorVm;
@@ -88,11 +87,7 @@ impl ModuleSubscriptions {
             } else {
                 // NOTE: The following ensures compliance with the 1.0 sql api.
                 // Come 1.0, it will have replaced the current compilation stack.
-                compile_sql_sub(
-                    &mut TyCtx::default(),
-                    sql,
-                    &SchemaViewer::new(&self.relational_db, &*tx, &auth),
-                )?;
+                compile_sql_sub(sql, &SchemaViewer::new(&self.relational_db, &*tx, &auth))?;
 
                 let mut compiled = compile_read_only_query(&self.relational_db, &auth, &tx, sql)?;
                 // Note that no error path is needed here.

crates/expr/src/check.rs

@@ -1,187 +1,150 @@
+use std::collections::HashMap;
+use std::ops::{Deref, DerefMut};
 use std::sync::Arc;
 
-use crate::statement::Statement;
-use crate::ty::TyId;
-use spacetimedb_schema::schema::{ColumnSchema, TableSchema};
+use crate::expr::{Expr, Project};
+use crate::{expr::LeftDeepJoin, statement::Statement};
+use spacetimedb_lib::AlgebraicType;
+use spacetimedb_schema::schema::TableSchema;
+use spacetimedb_sql_parser::ast::BinOp;
 use spacetimedb_sql_parser::{
-    ast::{
-        self,
-        sub::{SqlAst, SqlSelect},
-        SqlFrom, SqlIdent, SqlJoin,
-    },
+    ast::{sub::SqlSelect, SqlFrom, SqlIdent, SqlJoin},
     parser::sub::parse_subscription,
 };
 
 use super::{
-    assert_eq_types,
     errors::{DuplicateName, TypingError, Unresolved, Unsupported},
-    expr::{Expr, Let, RelExpr},
-    ty::{Symbol, TyCtx, TyEnv},
+    expr::RelExpr,
     type_expr, type_proj, type_select, StatementCtx, StatementSource,
 };
 
 /// The result of type checking and name resolution
 pub type TypingResult<T> = core::result::Result<T, TypingError>;
 
+/// A view of the database schema
 pub trait SchemaView {
     fn schema(&self, name: &str) -> Option<Arc<TableSchema>>;
 }
 
+#[derive(Default)]
+pub struct Relvars(HashMap<Box<str>, Arc<TableSchema>>);
+
+impl Deref for Relvars {
+    type Target = HashMap<Box<str>, Arc<TableSchema>>;
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl DerefMut for Relvars {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
 pub trait TypeChecker {
     type Ast;
     type Set;
 
-    fn type_ast(ctx: &mut TyCtx, ast: Self::Ast, tx: &impl SchemaView) -> TypingResult<RelExpr>;
+    fn type_ast(ast: Self::Ast, tx: &impl SchemaView) -> TypingResult<Project>;
 
-    fn type_set(ctx: &mut TyCtx, ast: Self::Set, tx: &impl SchemaView) -> TypingResult<RelExpr>;
+    fn type_set(ast: Self::Set, vars: &mut Relvars, tx: &impl SchemaView) -> TypingResult<Project>;
 
-    fn type_from(
-        ctx: &mut TyCtx,
-        from: SqlFrom<Self::Ast>,
-        tx: &impl SchemaView,
-    ) -> TypingResult<(RelExpr, Option<Symbol>)> {
+    fn type_from(from: SqlFrom, vars: &mut Relvars, tx: &impl SchemaView) -> TypingResult<RelExpr> {
         match from {
-            SqlFrom::Expr(expr, None) => Self::type_rel(ctx, expr, tx),
-            SqlFrom::Expr(expr, Some(SqlIdent(alias))) => {
-                let (expr, _) = Self::type_rel(ctx, expr, tx)?;
-                let symbol = ctx.gen_symbol(alias);
-                Ok((expr, Some(symbol)))
+            SqlFrom::Expr(SqlIdent(name), SqlIdent(alias)) => {
+                let schema = Self::type_relvar(tx, &name)?;
+                vars.insert(alias.clone(), schema.clone());
+                Ok(RelExpr::RelVar(schema, alias))
             }
-            SqlFrom::Join(r, SqlIdent(alias), joins) => {
-                // The type environment with which to type the join expressions
-                let mut env = TyEnv::default();
-                // The lowered inputs to the join operator
-                let mut inputs = Vec::new();
-                // The join expressions or predicates
-                let mut exprs = Vec::new();
-                // The types of the join variables or aliases
-                let mut types = Vec::new();
-
-                let input = Self::type_rel(ctx, r, tx)?.0;
-                let ty = input.ty_id();
-                let name = ctx.gen_symbol(alias);
-
-                env.add(name, ty);
-                inputs.push(input);
-                types.push((name, ty));
+            SqlFrom::Join(SqlIdent(name), SqlIdent(alias), joins) => {
+                let schema = Self::type_relvar(tx, &name)?;
+                vars.insert(alias.clone(), schema.clone());
+                let mut join = RelExpr::RelVar(schema, alias);
 
                 for SqlJoin {
-                    expr,
+                    var: SqlIdent(name),
                     alias: SqlIdent(alias),
                     on,
                 } in joins
                 {
-                    let input = Self::type_rel(ctx, expr, tx)?.0;
-                    let ty = input.ty_id();
-                    let name = ctx.gen_symbol(&alias);
-
-                    // New join variable is now in scope
-                    if env.add(name, ty).is_some() {
+                    // Check for duplicate aliases
+                    if vars.contains_key(&alias) {
                         return Err(DuplicateName(alias.into_string()).into());
                     }
 
-                    inputs.push(input);
-                    types.push((name, ty));
+                    let rhs = Self::type_relvar(tx, &name)?;
+                    let lhs = Box::new(join);
+                    let var = alias;
+
+                    vars.insert(var.clone(), rhs.clone());
 
-                    // Type check join expression with current type environment
                     if let Some(on) = on {
-                        exprs.push(type_expr(ctx, &env, on, Some(TyId::BOOL))?);
+                        if let Expr::BinOp(BinOp::Eq, a, b) = type_expr(vars, on, Some(&AlgebraicType::Bool))? {
+                            if let (Expr::Field(a), Expr::Field(b)) = (*a, *b) {
+                                join = RelExpr::EqJoin(LeftDeepJoin { lhs, rhs, var }, a, b);
+                                continue;
+                            }
+                        }
+                        unreachable!("Unreachability guaranteed by parser")
                     }
+
+                    join = RelExpr::LeftDeepJoin(LeftDeepJoin { lhs, rhs, var });
                 }
 
-                let ty = ctx.add_row_type(types.clone());
-                let input = RelExpr::Join(inputs.into(), ty);
-                let vars = types
-                    .into_iter()
-                    .enumerate()
-                    .map(|(i, (name, ty))| (name, Expr::Field(Box::new(Expr::Input(input.ty_id())), i, ty)))
-                    .collect();
-                Ok((RelExpr::select(input, Let { vars, exprs }), None))
+                Ok(join)
             }
         }
     }
 
-    fn type_rel(
-        ctx: &mut TyCtx,
-        expr: ast::RelExpr<Self::Ast>,
-        tx: &impl SchemaView,
-    ) -> TypingResult<(RelExpr, Option<Symbol>)> {
-        match expr {
-            ast::RelExpr::Var(SqlIdent(var)) => {
-                let schema = tx
-                    .schema(&var)
-                    .ok_or_else(|| Unresolved::table(&var))
-                    .map_err(TypingError::from)?;
-                let mut types = Vec::new();
-                for ColumnSchema { col_name, col_type, .. } in schema.columns() {
-                    let id = ctx.add_algebraic_type(col_type);
-                    let name = ctx.gen_symbol(col_name);
-                    types.push((name, id));
-                }
-                let id = ctx.add_var_type(schema.table_id, types);
-                let symbol = ctx.gen_symbol(var);
-                Ok((RelExpr::RelVar(schema, id), Some(symbol)))
-            }
-            ast::RelExpr::Ast(ast) => Ok((Self::type_ast(ctx, *ast, tx)?, None)),
-        }
+    fn type_relvar(tx: &impl SchemaView, name: &str) -> TypingResult<Arc<TableSchema>> {
+        tx.schema(name)
+            .ok_or_else(|| Unresolved::table(name))
+            .map_err(TypingError::from)
     }
 }
 
 /// Type checker for subscriptions
 struct SubChecker;
 
 impl TypeChecker for SubChecker {
-    type Ast = SqlAst;
-    type Set = SqlAst;
+    type Ast = SqlSelect;
+    type Set = SqlSelect;
 
-    fn type_ast(ctx: &mut TyCtx, ast: Self::Ast, tx: &impl SchemaView) -> TypingResult<RelExpr> {
-        Self::type_set(ctx, ast, tx)
+    fn type_ast(ast: Self::Ast, tx: &impl SchemaView) -> TypingResult<Project> {
+        Self::type_set(ast, &mut Relvars::default(), tx)
     }
 
-    fn type_set(ctx: &mut TyCtx, ast: Self::Set, tx: &impl SchemaView) -> TypingResult<RelExpr> {
+    fn type_set(ast: Self::Set, vars: &mut Relvars, tx: &impl SchemaView) -> TypingResult<Project> {
         match ast {
-            SqlAst::Union(a, b) => {
-                let a = Self::type_ast(ctx, *a, tx)?;
-                let b = Self::type_ast(ctx, *b, tx)?;
-                assert_eq_types(ctx, a.ty_id(), b.ty_id())?;
-                Ok(RelExpr::Union(Box::new(a), Box::new(b)))
-            }
-            SqlAst::Minus(a, b) => {
-                let a = Self::type_ast(ctx, *a, tx)?;
-                let b = Self::type_ast(ctx, *b, tx)?;
-                assert_eq_types(ctx, a.ty_id(), b.ty_id())?;
-                Ok(RelExpr::Minus(Box::new(a), Box::new(b)))
-            }
-            SqlAst::Select(SqlSelect {
+            SqlSelect {
                 project,
                 from,
                 filter: None,
-            }) => {
-                let (input, alias) = Self::type_from(ctx, from, tx)?;
-                type_proj(ctx, input, alias, project)
+            } => {
+                let input = Self::type_from(from, vars, tx)?;
+                type_proj(input, project, vars)
             }
-            SqlAst::Select(SqlSelect {
+            SqlSelect {
                 project,
                 from,
                 filter: Some(expr),
-            }) => {
-                let (from, alias) = Self::type_from(ctx, from, tx)?;
-                let input = type_select(ctx, from, alias, expr)?;
-                type_proj(ctx, input, alias, project)
+            } => {
+                let input = Self::type_from(from, vars, tx)?;
+                type_proj(type_select(input, expr, vars)?, project, vars)
             }
         }
     }
 }
 
 /// Parse and type check a subscription query
-pub fn parse_and_type_sub(ctx: &mut TyCtx, sql: &str, tx: &impl SchemaView) -> TypingResult<RelExpr> {
-    let expr = SubChecker::type_ast(ctx, parse_subscription(sql)?, tx)?;
-    expect_table_type(ctx, expr)
+pub fn parse_and_type_sub(sql: &str, tx: &impl SchemaView) -> TypingResult<Project> {
+    expect_table_type(SubChecker::type_ast(parse_subscription(sql)?, tx)?)
 }
 
 /// Parse and type check a *subscription* query into a `StatementCtx`
-pub fn compile_sql_sub<'a>(ctx: &mut TyCtx, sql: &'a str, tx: &impl SchemaView) -> TypingResult<StatementCtx<'a>> {
-    let expr = parse_and_type_sub(ctx, sql, tx)?;
+pub fn compile_sql_sub<'a>(sql: &'a str, tx: &impl SchemaView) -> TypingResult<StatementCtx<'a>> {
+    let expr = parse_and_type_sub(sql, tx)?;
     Ok(StatementCtx {
         statement: Statement::Select(expr),
         sql,
@@ -190,8 +153,10 @@ pub fn compile_sql_sub<'a>(ctx: &mut TyCtx, sql: &'a str, tx: &impl SchemaView)
 }
 
 /// Returns an error if the input type is not a table type or relvar
-fn expect_table_type(ctx: &TyCtx, expr: RelExpr) -> TypingResult<RelExpr> {
-    let _ = expr.ty(ctx)?.expect_relvar().map_err(|_| Unsupported::ReturnType)?;
+fn expect_table_type(expr: Project) -> TypingResult<Project> {
+    if let Project::Fields(..) = expr {
+        return Err(Unsupported::ReturnType.into());
+    }
     Ok(expr)
 }
 
@@ -233,7 +198,6 @@ pub mod test_utils {
 #[cfg(test)]
 mod tests {
     use crate::check::test_utils::{build_module_def, SchemaViewer};
-    use crate::ty::TyCtx;
     use spacetimedb_lib::{AlgebraicType, ProductType};
     use spacetimedb_schema::def::ModuleDef;
 
@@ -274,14 +238,11 @@ mod tests {
             "select * from t where t.u32 = 1 or t.str = ''",
             "select * from s where s.bytes = 0xABCD or bytes = X'ABCD'",
             "select * from s as r where r.bytes = 0xABCD or bytes = X'ABCD'",
-            "select * from (select t.* from t join s)",
-            "select * from (select t.* from t join s join s as r where t.u32 = s.u32 and s.u32 = r.u32)",
-            "select * from (select t.* from t join s on t.u32 = s.u32 where t.f32 = 0.1)",
-            "select * from (select t.* from t join (select s.u32 from s) s on t.u32 = s.u32)",
-            "select * from (select t.* from t join (select u32 as a from s) s on t.u32 = s.a)",
-            "select * from (select * from t union all select * from t)",
+            "select t.* from t join s",
+            "select t.* from t join s join s as r where t.u32 = s.u32 and s.u32 = r.u32",
+            "select t.* from t join s on t.u32 = s.u32 where t.f32 = 0.1",
         ] {
-            let result = parse_and_type_sub(&mut TyCtx::default(), sql, &tx);
+            let result = parse_and_type_sub(sql, &tx);
             assert!(result.is_ok());
         }
     }
@@ -308,21 +269,13 @@ mod tests {
             // Subscriptions must be typed to a single table
             "select * from t join s",
             // Self join requires aliases
-            "select * from (select t.* from t join t)",
+            "select t.* from t join t",
             // Product values are not comparable
-            "select * from (select t.* from t join s on t.arr = s.arr)",
-            // Subscriptions must be typed to a single table
-            "select * from (select s.* from t join (select u32 from s) s on t.u32 = s.u32)",
-            // Field u32 has been renamed
-            "select * from (select t.* from t join (select u32 as a from s) s on t.u32 = s.u32)",
-            // Field bytes is no longer in scope
-            "select * from (select t.* from t join (select u32 from s) s on s.bytes = 0xABCD)",
+            "select t.* from t join s on t.arr = s.arr",
             // Alias r is not in scope when it is referenced
-            "select * from (select t.* from t join s on t.u32 = r.u32 join s as r)",
-            // Union arguments are of different types
-            "select * from (select * from t union all select * from s)",
+            "select t.* from t join s on t.u32 = r.u32 join s as r",
         ] {
-            let result = parse_and_type_sub(&mut TyCtx::default(), sql, &tx);
+            let result = parse_and_type_sub(sql, &tx);
             assert!(result.is_err());
         }
     }