diff --git a/crates/knot-compiler/src/codegen.rs b/crates/knot-compiler/src/codegen.rs index 3724f454..f54afc88 100644 --- a/crates/knot-compiler/src/codegen.rs +++ b/crates/knot-compiler/src/codegen.rs @@ -189,6 +189,15 @@ pub struct Codegen { // Trampolines for user functions used as values: fn_name -> trampoline_func_id user_fn_trampolines: HashMap, + // Do-blocks sitting in the value position of a `set`/`replace`, keyed by + // span. These are relational comprehensions even when they bind from a + // source (`x <- *rel`), which `is_io_do_block` would otherwise read as IO. + // A do-block directly under the `=` is handled by `compile_set_value_expr`, + // but one nested inside an `if`/`case` branch is reached through the + // generic `compile_expr` path, which has no way to know it is producing the + // relation being written — hence the span set. + relational_do_spans: HashSet, + // Resolved monad types for desugared do-blocks (from type inference) monad_info: MonadInfo, /// Static dispatch type per trait-method occurrence, from inference. @@ -920,6 +929,7 @@ impl Codegen { fetch_route_entries: HashMap::new(), type_aliases: HashMap::new(), user_fn_trampolines: HashMap::new(), + relational_do_spans: HashSet::new(), monad_info: HashMap::new(), trait_call_targets: HashMap::new(), trait_method_traits: HashMap::new(), @@ -4674,7 +4684,11 @@ impl Codegen { } ast::ExprKind::Do(stmts) => { - if self.is_io_do_block(stmts) { + if self.relational_do_spans.contains(&expr.span) { + // Produces the relation written by an enclosing + // set/replace, even if it binds from a source. + self.compile_do(builder, stmts, env, db) + } else if self.is_io_do_block(stmts) { self.compile_io_do(builder, stmts, env, db) } else if self.in_io_eager && !stmts.iter().any(|s| matches!(&s.node, ast::StmtKind::Bind { .. })) @@ -7923,10 +7937,49 @@ impl Codegen { ast::ExprKind::Refine(inner) => { self.compile_set_value_expr(builder, inner, env, db) } + // `if`/`case` in set-value position: each branch is itself a + // set value, so a do-block in a branch is a relational + // comprehension too. The branches are compiled by the generic + // `compile_expr` path, so record their do-block spans first and + // let the `Do` arm of `compile_expr` consult the set. + ast::ExprKind::If { .. } | ast::ExprKind::Case { .. } => { + Self::collect_relational_do_spans(value, &mut self.relational_do_spans); + self.compile_expr(builder, value, env, db) + } _ => self.compile_expr(builder, value, env, db), } } + /// Record the spans of do-blocks that a set/replace value produces its + /// relation from. Result position extends through type/unit wrappers and + /// through `if`/`case` branches (which may nest arbitrarily), so the walk + /// mirrors `compile_set_value_expr`'s own recursion. + fn collect_relational_do_spans(value: &ast::Expr, spans: &mut HashSet) { + match &value.node { + ast::ExprKind::Do(_) => { + spans.insert(value.span); + } + ast::ExprKind::UnitLit { value: inner, .. } + | ast::ExprKind::TimeUnitLit { value: inner, .. } + | ast::ExprKind::Annot { expr: inner, .. } + | ast::ExprKind::Refine(inner) => Self::collect_relational_do_spans(inner, spans), + ast::ExprKind::If { + then_branch, + else_branch, + .. + } => { + Self::collect_relational_do_spans(then_branch, spans); + Self::collect_relational_do_spans(else_branch, spans); + } + ast::ExprKind::Case { arms, .. } => { + for arm in arms { + Self::collect_relational_do_spans(&arm.body, spans); + } + } + _ => {} + } + } + fn is_io_do_block(&self, stmts: &[ast::Stmt]) -> bool { // Do-blocks with groupBy always need relational iteration (compile_do), // even if they contain IO-like expressions, because groupBy requires diff --git a/crates/knot-compiler/tests/regress_codegen_fixes.rs b/crates/knot-compiler/tests/regress_codegen_fixes.rs index cc3bfff3..1283c4d1 100644 --- a/crates/knot-compiler/tests/regress_codegen_fixes.rs +++ b/crates/knot-compiler/tests/regress_codegen_fixes.rs @@ -10,6 +10,9 @@ //! unexecuted IO thunks (`expected Relation in len, got IO` at runtime). //! 5. Trampoline curry chains emitting unsorted env-record keys for //! functions with >= 12 parameters ("10" sorts before "2"). +//! 6. Relation comprehensions in an `if`/`case` branch of a `set`/`replace` +//! value compiling as IO thunks rather than relations +//! (`source_write expects a Relation, got IO` at runtime). //! //! Each test compiles a small Knot program with the real `knot` binary into //! its own scratch directory (so `knot.db` lands there) and asserts on the @@ -862,3 +865,111 @@ main = do "a whole-relation bind must keep binding the relation, got:\n{stdout}" ); } + +// ── Finding 6: comprehensions in an if/case branch of a set value ── + +#[test] +fn set_value_comprehension_inside_if_branch() { + // A do-block directly under `replace *rel =` is a relational + // comprehension, but one inside an `if` branch was reached through the + // generic expression path, where `is_io_do_block` sees the `x <- *other` + // bind and compiles it as an IO thunk. The thunk was then handed to + // `knot_source_write`, aborting with + // `source_write expects a Relation, got IO`. + let (stdout, stderr, ok) = compile_and_run( + "set_if_branch_comprehension", + r#"*other : [{name: Text}] +*items : [{name: Text}] + +main = do + replace *other = [{name: "a"}, {name: "b"}] + replace *items = + if True + then do + x <- *other + where x.name == "a" + yield x + else [] + rows <- *items + println ("then: " ++ show rows) + + -- the same comprehension in the else branch, one level of nesting down + replace *items = + if False + then [] + else if True + then do + x <- *other + yield x + else [] + all <- *items + println ("else: " ++ show (count all)) + + -- a branch that is not a comprehension still writes its own rows + replace *items = + if False + then do + x <- *other + yield x + else [{name: "z"}] + lit <- *items + println ("lit: " ++ show lit) +"#, + ); + assert!(ok, "program failed:\nstdout: {stdout}\nstderr: {stderr}"); + // Each branch iterates the source's ROWS, exactly as the same do-block + // written directly under the `=` would. + assert!( + stdout.contains("then: [{name: a}]"), + "an if-branch comprehension must write the filtered rows, got:\n{stdout}" + ); + assert!( + stdout.contains("else: 2"), + "a nested else-branch comprehension must write every row, got:\n{stdout}" + ); + assert!( + stdout.contains("lit: [{name: z}]"), + "a non-comprehension branch must still write its own value, got:\n{stdout}" + ); +} + +#[test] +fn set_value_comprehension_inside_case_arm() { + // Same defect via `case`: the arm bodies are set values too. + let (stdout, stderr, ok) = compile_and_run( + "set_case_arm_comprehension", + r#"data Mode = Copy | Clear + +*other : [{name: Text}] +*items : [{name: Text}] + +main = do + replace *other = [{name: "a"}, {name: "b"}] + replace *items = + case Copy of + Copy -> do + x <- *other + yield x + Clear -> [] + copied <- *items + println ("copy: " ++ show (count copied)) + replace *items = + case Clear of + Copy -> do + x <- *other + yield x + Clear -> [] + cleared <- *items + println ("clear: " ++ show (count cleared)) +"#, + ); + assert!(ok, "program failed:\nstdout: {stdout}\nstderr: {stderr}"); + assert!( + stdout.contains("copy: 2"), + "a case-arm comprehension must write the source's rows, got:\n{stdout}" + ); + assert!( + stdout.contains("clear: 0"), + "the non-comprehension arm must still clear the relation, got:\n{stdout}" + ); +}