vm: Make UPDATE actually update

Previously, an UPDATE statement would be the identity function. Fix this
by threading through the assignments (SET clause) and transforming the
rows matching the query.

crates/core/src/sql/compiler.rs

@@ -403,24 +403,8 @@ fn compile_update(
     } else {
         QueryExpr::new(&table.root)
     };
-    let mut cols = Vec::with_capacity(table.root.columns.len());
-
-    for field in table.root.columns.iter() {
-        let field = FieldName::named(&table.root.table_name, &field.col_name);
-        if let Some(f) = assignments.get(&field) {
-            cols.push(f.clone());
-        } else {
-            cols.push(FieldExpr::Name(field));
-        }
-    }
 
-    let insert = QueryExpr::new(&table.root).with_project(&cols, None);
-    let insert = if let Some(filter) = selection {
-        compile_where(insert, &table, filter)?
-    } else {
-        insert
-    };
-    Ok(CrudExpr::Update { insert, delete })
+    Ok(CrudExpr::Update { delete, assignments })
 }
 
 /// Compiles a `CREATE TABLE ...` clause

crates/core/src/vm.rs

@@ -11,7 +11,6 @@ use spacetimedb_lib::relation::{Header, MemTable, RelIter, RelValue, RowCount, T
 use spacetimedb_lib::table::ProductTypeMeta;
 use spacetimedb_primitives::{ColId, TableId};
 use spacetimedb_sats::{AlgebraicValue, ProductValue};
-use spacetimedb_vm::dsl::mem_table;
 use spacetimedb_vm::env::EnvDb;
 use spacetimedb_vm::errors::ErrorVm;
 use spacetimedb_vm::eval::IterRows;
@@ -308,16 +307,6 @@ impl<'db, 'tx> DbProgram<'db, 'tx> {
         }
     }
 
-    fn insert_query(&mut self, table: &Table, query: QueryCode) -> Result<Code, ErrorVm> {
-        let result = self._eval_query(query)?;
-        match result {
-            Code::Table(result) => {
-                self._execute_insert(table, result.data.into_iter().map(|row| row.data).collect_vec())
-            }
-            _ => Ok(result),
-        }
-    }
-
     fn create_table(
         &mut self,
         table_name: &str,
@@ -402,7 +391,10 @@ impl ProgramVm for DbProgram<'_, '_> {
         match query {
             CrudCode::Query(query) => self._eval_query(query),
             CrudCode::Insert { table, rows } => self._execute_insert(&table, rows),
-            CrudCode::Update { mut insert, delete } => {
+            CrudCode::Update {
+                delete,
+                mut assignments,
+            } => {
                 let table = delete.table.clone();
                 let result = self._eval_query(delete)?;
 
@@ -415,11 +407,38 @@ impl ProgramVm for DbProgram<'_, '_> {
                     deleted.data.clone().into_iter().map(|row| row.data).collect_vec(),
                 )?;
 
-                let to_insert = mem_table(table.head().clone(), deleted.data.into_iter().map(|row| row.data));
-                insert.table = Table::MemTable(to_insert);
+                // Replace the columns in the matched rows with the assigned
+                // values. No typechecking is performed here, nor that all
+                // assignments are consumed.
+                let exprs: Vec<Option<FieldExpr>> = table
+                    .head()
+                    .fields
+                    .iter()
+                    .map(|col| assignments.remove(&col.field))
+                    .collect();
+                let insert_rows = deleted
+                    .data
+                    .into_iter()
+                    .map(|row| {
+                        let elements = row
+                            .data
+                            .elements
+                            .into_iter()
+                            .zip(&exprs)
+                            .map(|(val, expr)| {
+                                if let Some(FieldExpr::Value(assigned)) = expr {
+                                    assigned.clone()
+                                } else {
+                                    val
+                                }
+                            })
+                            .collect();
+
+                        ProductValue { elements }
+                    })
+                    .collect_vec();
 
-                let result = self.insert_query(&table, insert)?;
-                Ok(result)
+                self._execute_insert(&table, insert_rows)
             }
             CrudCode::Delete { query } => {
                 let result = self.delete_query(query)?;

crates/vm/src/eval.rs

@@ -101,11 +101,10 @@ fn build_typed<P: ProgramVm>(p: &mut P, node: Expr) -> ExprOpt {
                 }
                 ExprOpt::Crud(Box::new(CrudExprOpt::Insert { source, rows }))
             }
-            CrudExpr::Update { insert, delete } => {
-                let insert = build_query_opt(insert);
+            CrudExpr::Update { delete, assignments } => {
                 let delete = build_query_opt(delete);
 
-                ExprOpt::Crud(Box::new(CrudExprOpt::Update { insert, delete }))
+                ExprOpt::Crud(Box::new(CrudExprOpt::Update { delete, assignments }))
             }
             CrudExpr::Delete { query } => {
                 let query = build_query_opt(query);
@@ -270,10 +269,9 @@ fn compile<P: ProgramVm>(p: &mut P, node: ExprOpt) -> Result<Code, ErrorVm> {
                     };
                     Code::Crud(q)
                 }
-                CrudExprOpt::Update { insert, delete } => {
-                    let insert = compile_query(insert);
+                CrudExprOpt::Update { delete, assignments } => {
                     let delete = compile_query(delete);
-                    Code::Crud(CrudCode::Update { insert, delete })
+                    Code::Crud(CrudCode::Update { delete, assignments })
                 }
                 CrudExprOpt::Delete { query } => {
                     let query = compile_query(query);

crates/vm/src/expr.rs

@@ -400,16 +400,16 @@ pub enum Crud {
     Drop(DbType),
 }
 
-#[derive(Debug, Clone, Eq, PartialEq, PartialOrd, Ord)]
+#[derive(Debug, Clone, Eq, PartialEq)]
 pub enum CrudExpr {
     Query(QueryExpr),
     Insert {
         source: SourceExpr,
         rows: Vec<Vec<FieldExpr>>,
     },
     Update {
-        insert: QueryExpr,
         delete: QueryExpr,
+        assignments: HashMap<FieldName, FieldExpr>,
     },
     Delete {
         query: QueryExpr,
@@ -1081,8 +1081,8 @@ pub enum CrudExprOpt {
         rows: Vec<ProductValue>,
     },
     Update {
-        insert: QueryExprOpt,
         delete: QueryExprOpt,
+        assignments: HashMap<FieldName, FieldExpr>,
     },
     Delete {
         query: QueryExprOpt,
@@ -1343,8 +1343,8 @@ pub enum CrudCode {
         rows: Vec<ProductValue>,
     },
     Update {
-        insert: QueryCode,
         delete: QueryCode,
+        assignments: HashMap<FieldName, FieldExpr>,
     },
     Delete {
         query: QueryCode,
@@ -1576,11 +1576,13 @@ mod tests {
 
     #[test]
     fn test_auth_crud_code_update() {
-        let mut qc = query_codes().into_iter();
-        let insert = qc.next().unwrap();
-        let delete = qc.next().unwrap();
-        let crud = CrudCode::Update { insert, delete };
-        assert_owner_required(crud);
+        for qc in query_codes() {
+            let crud = CrudCode::Update {
+                delete: qc,
+                assignments: Default::default(),
+            };
+            assert_owner_required(crud);
+        }
     }
 
     #[test]

crates/vm/src/typecheck.rs

@@ -136,7 +136,7 @@ pub(crate) fn check_types(env: &mut EnvTy, ast: &ExprOpt) -> Result<Ty, ErrorTyp
             let q = &**q;
             match q {
                 CrudExprOpt::Insert { source, .. } => Ok(ty_source(source)),
-                CrudExprOpt::Update { insert, .. } => Ok(ty_source(&insert.source)),
+                CrudExprOpt::Update { delete, .. } => Ok(ty_source(&delete.source)),
                 CrudExprOpt::Delete { query } => Ok(ty_source(&query.source)),
                 CrudExprOpt::CreateTable { columns, .. } => Ok(AlgebraicType::Product(columns.columns.clone()).into()),
                 CrudExprOpt::Drop { .. } => {