Lower branch-yield result joins into Core

crates/edict-cli/src/main.rs

@@ -1508,15 +1508,22 @@ fn core_node_review(node: &CoreNode) -> Value {
             "body": core_block_review(body),
         }),
         CoreNode::Branch {
+            binding,
             predicate,
             then_block,
             else_block,
-        } => json!({
-            "kind": "branch",
-            "predicate": core_predicate_review(predicate),
-            "then": core_block_review(then_block),
-            "else": core_block_review(else_block),
-        }),
+        } => {
+            let mut review = json!({
+                "kind": "branch",
+                "predicate": core_predicate_review(predicate),
+                "then": core_block_review(then_block),
+                "else": core_block_review(else_block),
+            });
+            if let Some(binding) = binding {
+                review["binding"] = local_ref_review(binding);
+            }
+            review
+        }
     }
 }
 

crates/edict-syntax/src/canonical.rs

@@ -1095,15 +1095,22 @@ fn core_node_value(node: &CoreNode) -> Result<CanonicalValue, CanonicalError> {
             ("body", core_block_value(body)?),
         ])),
         CoreNode::Branch {
+            binding,
             predicate,
             then_block,
             else_block,
-        } => Ok(map([
-            ("kind", text("branch")),
-            ("predicate", core_predicate_value(predicate)?),
-            ("then", core_block_value(then_block)?),
-            ("else", core_block_value(else_block)?),
-        ])),
+        } => {
+            let mut fields = vec![
+                ("kind", text("branch")),
+                ("predicate", core_predicate_value(predicate)?),
+                ("then", core_block_value(then_block)?),
+                ("else", core_block_value(else_block)?),
+            ];
+            if let Some(binding) = binding {
+                fields.push(("binding", local_ref_value(binding)));
+            }
+            Ok(map(fields))
+        }
     }
 }
 

crates/edict-syntax/src/compiler.rs

@@ -927,6 +927,7 @@ impl<'a> TypeChecker<'a> {
                 span,
             } => self.check_let_stmt(
                 intent,
+                output_shape,
                 LetStatement {
                     name,
                     ty: ty.as_ref(),
@@ -1065,6 +1066,7 @@ impl<'a> TypeChecker<'a> {
         };
         state.accumulated_steps = then_steps.max(state.accumulated_steps);
         state.nodes.push(CoreNode::Branch {
+            binding: None,
             predicate,
             then_block,
             else_block,
@@ -1639,6 +1641,7 @@ impl<'a> TypeChecker<'a> {
     fn check_let_stmt(
         &mut self,
         intent: &ResolvedIntent,
+        output_shape: &TypeShape,
         stmt: LetStatement<'_>,
         env: &mut BTreeMap<String, (LocalRef, TypeShape)>,
         locals: &mut Vec<LocalRef>,
@@ -1647,13 +1650,14 @@ impl<'a> TypeChecker<'a> {
         if let Some(handler) = stmt.handler {
             self.check_effectful_let(intent, stmt, handler, env, locals, state);
         } else {
-            self.check_pure_let(intent, stmt, env, locals, state);
+            self.check_pure_let(intent, output_shape, stmt, env, locals, state);
         }
     }
 
     fn check_pure_let(
         &mut self,
         intent: &ResolvedIntent,
+        output_shape: &TypeShape,
         stmt: LetStatement<'_>,
         env: &mut BTreeMap<String, (LocalRef, TypeShape)>,
         locals: &mut Vec<LocalRef>,
@@ -1662,6 +1666,27 @@ impl<'a> TypeChecker<'a> {
         if !self.check_known_effect_profiles(intent, stmt.value) {
             return;
         }
+        if let Expr::IfYield {
+            pred,
+            then_block,
+            else_block,
+            span,
+        } = stmt.value
+        {
+            self.check_branch_yield_let(
+                intent,
+                output_shape,
+                &stmt,
+                pred,
+                then_block,
+                else_block,
+                *span,
+                env,
+                locals,
+                state,
+            );
+            return;
+        }
         let annotation_shape = match stmt.ty {
             Some(annotation) => match self.type_ref_shape(annotation, stmt.span, None) {
                 Some(shape) => Some(shape),
@@ -1686,6 +1711,173 @@ impl<'a> TypeChecker<'a> {
         env.insert(stmt.name.to_owned(), (local, binding_shape));
     }
 
+    #[allow(clippy::too_many_arguments)]
+    fn check_branch_yield_let(
+        &mut self,
+        intent: &ResolvedIntent,
+        output_shape: &TypeShape,
+        stmt: &LetStatement<'_>,
+        pred: &Expr,
+        then_source: &YieldBlock,
+        else_source: &YieldBlock,
+        span: Span,
+        env: &mut BTreeMap<String, (LocalRef, TypeShape)>,
+        locals: &mut Vec<LocalRef>,
+        state: &mut BodyState,
+    ) {
+        let Some(predicate) = self.check_predicate(pred, env) else {
+            return;
+        };
+        let annotation_shape = match stmt.ty {
+            Some(annotation) => match self.type_ref_shape(annotation, stmt.span, None) {
+                Some(shape) => Some(shape),
+                None => return,
+            },
+            None => None,
+        };
+        let baseline_steps = state.accumulated_steps;
+        let (then_block, then_shape, else_block, else_shape) =
+            if annotation_shape.is_none() && is_bare_integer_literal(&then_source.value) {
+                let Some((else_block, else_shape)) =
+                    self.check_yield_block(intent, output_shape, else_source, env, state, None)
+                else {
+                    return;
+                };
+                let else_steps = state.accumulated_steps;
+                state.accumulated_steps = baseline_steps;
+                let Some((then_block, then_shape)) = self.check_yield_block(
+                    intent,
+                    output_shape,
+                    then_source,
+                    env,
+                    state,
+                    Some(&else_shape),
+                ) else {
+                    return;
+                };
+                state.accumulated_steps = else_steps.max(state.accumulated_steps);
+                (then_block, then_shape, else_block, else_shape)
+            } else {
+                let Some((then_block, then_shape)) = self.check_yield_block(
+                    intent,
+                    output_shape,
+                    then_source,
+                    env,
+                    state,
+                    annotation_shape.as_ref(),
+                ) else {
+                    return;
+                };
+                let then_steps = state.accumulated_steps;
+                state.accumulated_steps = baseline_steps;
+                let branch_expectation = annotation_shape.as_ref().unwrap_or(&then_shape);
+                let Some((else_block, else_shape)) = self.check_yield_block(
+                    intent,
+                    output_shape,
+                    else_source,
+                    env,
+                    state,
+                    Some(branch_expectation),
+                ) else {
+                    return;
+                };
+                state.accumulated_steps = then_steps.max(state.accumulated_steps);
+                (then_block, then_shape, else_block, else_shape)
+            };
+        let binding_shape = if let Some(annotation_shape) = annotation_shape {
+            if !compatible(&annotation_shape, &then_shape)
+                || !compatible(&annotation_shape, &else_shape)
+            {
+                self.errors.push(error(
+                    CompilerStage::TypeCheck,
+                    CompilerErrorKind::TypeMismatch,
+                    "branch-yield results do not match the annotated type",
+                    span,
+                ));
+                return;
+            }
+            annotation_shape
+        } else if compatible(&then_shape, &else_shape) {
+            then_shape
+        } else if compatible(&else_shape, &then_shape) {
+            else_shape
+        } else {
+            self.errors.push(error(
+                CompilerStage::TypeCheck,
+                CompilerErrorKind::TypeMismatch,
+                "branch-yield results do not have a compatible bounded type",
+                span,
+            ));
+            return;
+        };
+
+        let binding = next_local(&mut state.local_index, binding_shape.coord.clone());
+        state.nodes.push(CoreNode::Branch {
+            binding: Some(binding.clone()),
+            predicate,
+            then_block,
+            else_block,
+        });
+        locals.push(binding.clone());
+        env.insert(stmt.name.to_owned(), (binding, binding_shape));
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    fn check_yield_block(
+        &mut self,
+        intent: &ResolvedIntent,
+        output_shape: &TypeShape,
+        block: &YieldBlock,
+        env: &BTreeMap<String, (LocalRef, TypeShape)>,
+        state: &mut BodyState,
+        expected: Option<&TypeShape>,
+    ) -> Option<(CoreBlock, TypeShape)> {
+        let error_count = self.errors.len();
+        let mut nested_env = env.clone();
+        let mut nested_locals = Vec::new();
+        let mut nested_state = BodyState {
+            local_index: state.local_index,
+            obstruction_index: state.obstruction_index,
+            step_factor: state.step_factor,
+            accumulated_steps: state.accumulated_steps,
+            ..BodyState::default()
+        };
+        for stmt in &block.stmts {
+            if let Stmt::Return { span, .. } = stmt {
+                self.errors.push(error(
+                    CompilerStage::TypeCheck,
+                    CompilerErrorKind::UnsupportedSourceShape,
+                    "return is not legal inside a branch-yield block",
+                    *span,
+                ));
+                continue;
+            }
+            self.check_body_stmt(
+                intent,
+                output_shape,
+                stmt,
+                &mut nested_env,
+                &mut nested_locals,
+                &mut nested_state,
+            );
+        }
+        let value = self.check_expr_with_expected(&block.value, &nested_env, expected)?;
+        state.local_index = nested_state.local_index;
+        state.obstruction_index = nested_state.obstruction_index;
+        state.accumulated_steps = nested_state.accumulated_steps;
+        if self.errors.len() != error_count {
+            return None;
+        }
+        Some((
+            CoreBlock {
+                locals: nested_locals,
+                nodes: nested_state.nodes,
+                result: value.expr,
+            },
+            value.ty,
+        ))
+    }
+
     fn pure_let_binding_shape(
         &mut self,
         stmt: &LetStatement<'_>,
@@ -2801,6 +2993,14 @@ fn next_local(index: &mut usize, ty: String) -> LocalRef {
 }
 
 fn compatible(expected: &TypeShape, actual: &TypeShape) -> bool {
+    if let (TypeKind::Record(expected), TypeKind::Record(actual)) = (&expected.kind, &actual.kind) {
+        return expected.len() == actual.len()
+            && expected.iter().all(|(name, expected_ty)| {
+                actual
+                    .get(name)
+                    .is_some_and(|actual_ty| compatible(expected_ty, actual_ty))
+            });
+    }
     if expected.coord == actual.coord {
         return true;
     }
@@ -2815,14 +3015,6 @@ fn compatible(expected: &TypeShape, actual: &TypeShape) -> bool {
                 canonical: actual_canonical,
             },
         ) => actual_max <= expected_max && actual_canonical == expected_canonical,
-        (TypeKind::Record(expected), TypeKind::Record(actual)) => {
-            expected.len() == actual.len()
-                && expected.iter().all(|(name, expected_ty)| {
-                    actual
-                        .get(name)
-                        .is_some_and(|actual_ty| compatible(expected_ty, actual_ty))
-                })
-        }
         (TypeKind::Bool, TypeKind::Bool) => true,
         (TypeKind::Int { width: expected }, TypeKind::Int { width: actual }) => expected == actual,
         (TypeKind::Bytes { max: expected }, TypeKind::Bytes { max: actual }) => actual <= expected,

crates/edict-syntax/src/core_ir.rs

@@ -313,6 +313,9 @@ pub enum CoreNode {
         body: CoreBlock,
     },
     Branch {
+        /// When present, the selected block result becomes this local.
+        /// Statement-only branches leave the binding absent.
+        binding: Option<LocalRef>,
         predicate: CorePredicate,
         then_block: CoreBlock,
         else_block: CoreBlock,