Fix replace *rel = if/case with comprehension branches crashes

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<String, FuncId>,
 
+    // 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<ast::Span>,
+
     // 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<ast::Span>) {
+        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

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}"
+    );
+}