feat(control_tree): Phase 129-B join_k as-last lowering

src/mir/control_tree/normalized_shadow/builder.rs

@@ -14,7 +14,7 @@
 //! - Single responsibility: structure → JoinIR conversion
 
 use crate::mir::control_tree::step_tree::StepTree;
-use crate::mir::join_ir::lowering::carrier_info::{ExitMeta, JoinFragmentMeta};
+use crate::mir::join_ir::lowering::carrier_info::JoinFragmentMeta;
 use crate::mir::join_ir::JoinModule;
 use crate::mir::ValueId; // Phase 126
 use std::collections::BTreeMap; // Phase 126
@@ -56,8 +56,6 @@ impl EnvLayout {
         contract: &crate::mir::control_tree::step_tree_contract_box::StepTreeContract,
         available_inputs: &std::collections::BTreeMap<String, crate::mir::ValueId>,
     ) -> Self {
-        use std::collections::BTreeSet;
-
         // Phase 125 P2: writes from contract
         let writes: Vec<String> = contract.writes.iter().cloned().collect();
 
@@ -77,6 +75,26 @@ impl EnvLayout {
 pub struct StepTreeNormalizedShadowLowererBox;
 
 impl StepTreeNormalizedShadowLowererBox {
+    /// Phase 129-B: Expected env field count (writes + inputs)
+    pub fn expected_env_field_count(
+        step_tree: &StepTree,
+        available_inputs: &BTreeMap<String, ValueId>,
+    ) -> usize {
+        let env_layout = EnvLayout::from_contract(&step_tree.contract, available_inputs);
+        env_layout.writes.len() + env_layout.inputs.len()
+    }
+
+    /// Phase 129-B: If-as-last shape detection (no post-if)
+    pub fn expects_join_k_as_last(step_tree: &StepTree) -> bool {
+        use crate::mir::control_tree::step_tree::StepNode;
+
+        match &step_tree.root {
+            StepNode::If { .. } => true,
+            StepNode::Block(nodes) => matches!(nodes.last(), Some(StepNode::If { .. })),
+            _ => false,
+        }
+    }
+
     /// Try to lower an if-only StepTree to normalized form
     ///
     /// ## Returns
@@ -150,12 +168,19 @@ impl StepTreeNormalizedShadowLowererBox {
         step_tree: &StepTree,
         available_inputs: &BTreeMap<String, ValueId>,
     ) -> Result<Option<(JoinModule, JoinFragmentMeta)>, String> {
-        use crate::mir::join_ir::{JoinFunction, JoinFuncId, JoinInst};
+        use crate::mir::join_ir::{JoinFunction, JoinFuncId};
         use crate::mir::ValueId;
 
         // Phase 126: EnvLayout 生成（available_inputs を使用）
         let env_layout = EnvLayout::from_contract(&step_tree.contract, available_inputs);
 
+        // Phase 129-B: If-as-last join_k lowering (dev-only)
+        if let Some((module, meta)) =
+            Self::lower_if_node_as_last_stmt(step_tree, &env_layout)?
+        {
+            return Ok(Some((module, meta)));
+        }
+
         let main_func_id = JoinFuncId::new(0);
 
         // Phase 125 P2: writes 用の ValueId 生成
@@ -170,20 +195,14 @@ impl StepTreeNormalizedShadowLowererBox {
             })
             .collect();
 
-        // Phase 125 P2: inputs 用の ValueId 生成（今は空だが、P3 で available_inputs から参照）
+        // Phase 129-B: inputs 用の ValueId 生成（writes + inputs が env params のSSOT）
         let inputs_params: Vec<ValueId> = env_layout
             .inputs
             .iter()
-            .map(|name| {
-                // Phase 125 P3: available_inputs から取得
-                // 今は空なので到達しないはずだが、念のため placeholder
-                available_inputs
-                    .get(name)
-                    .copied()
-                    .unwrap_or_else(|| {
-                        // Should not reach here (inputs は available_inputs にある前提)
-                        ValueId(0) // placeholder
-                    })
+            .map(|_| {
+                let vid = ValueId(next_value_id);
+                next_value_id += 1;
+                vid
             })
             .collect();
 
@@ -196,9 +215,9 @@ impl StepTreeNormalizedShadowLowererBox {
             env.insert(name.clone(), *vid);
         }
 
-        // Phase 125 P2: 関数パラメータは writes のみ（inputs は外側から来る前提）
-        // TODO P3: inputs も params に含める必要があるか検討
-        let env_params = writes_params;
+        // Phase 129-B: 関数パラメータは writes + inputs（env params SSOT）
+        let mut env_params = writes_params;
+        env_params.extend(inputs_params);
 
         // main 関数生成
         let mut main_func = JoinFunction::new(
@@ -246,6 +265,379 @@ impl StepTreeNormalizedShadowLowererBox {
         Ok(Some((module, meta)))
     }
 
+    /// Phase 129-B: Lower if node as last statement using join_k (dev-only)
+    ///
+    /// Scope: "if-as-last" only (no post-if). If it doesn't match, return Ok(None).
+    fn lower_if_node_as_last_stmt(
+        step_tree: &StepTree,
+        env_layout: &EnvLayout,
+    ) -> Result<Option<(JoinModule, JoinFragmentMeta)>, String> {
+        use crate::mir::control_tree::step_tree::{StepNode, StepStmtKind};
+        use crate::mir::join_ir::{ConstValue, JoinFunction, JoinFuncId, JoinInst, MirLikeInst};
+        use crate::mir::join_ir::lowering::error_tags;
+        use crate::mir::ValueId;
+        use std::collections::BTreeMap;
+
+        let (prefix_nodes, if_node) = match &step_tree.root {
+            StepNode::If { .. } => (&[][..], &step_tree.root),
+            StepNode::Block(nodes) => {
+                let last = nodes.last();
+                if matches!(last, Some(StepNode::If { .. })) {
+                    (&nodes[..nodes.len() - 1], nodes.last().unwrap())
+                } else {
+                    return Ok(None);
+                }
+            }
+            _ => return Ok(None),
+        };
+
+        let if_node = match if_node {
+            StepNode::If { .. } => if_node,
+            _ => return Ok(None),
+        };
+
+        let env_fields: Vec<String> = env_layout
+            .writes
+            .iter()
+            .chain(env_layout.inputs.iter())
+            .cloned()
+            .collect();
+
+        fn alloc_value_id(next_value_id: &mut u32) -> ValueId {
+            let vid = ValueId(*next_value_id);
+            *next_value_id += 1;
+            vid
+        }
+
+        let mut next_value_id: u32 = 1;
+
+        let alloc_env_params =
+            |fields: &[String], next_value_id: &mut u32| -> Vec<ValueId> {
+                fields
+                    .iter()
+                    .map(|_| alloc_value_id(next_value_id))
+                    .collect()
+            };
+
+        let build_env_map = |fields: &[String], params: &[ValueId]| -> BTreeMap<String, ValueId> {
+            let mut env = BTreeMap::new();
+            for (name, vid) in fields.iter().zip(params.iter()) {
+                env.insert(name.clone(), *vid);
+            }
+            env
+        };
+
+        let collect_env_args = |fields: &[String], env: &BTreeMap<String, ValueId>| -> Result<Vec<ValueId>, String> {
+            let mut args = Vec::with_capacity(fields.len());
+            for name in fields {
+                let vid = env.get(name).copied().ok_or_else(|| {
+                    error_tags::freeze_with_hint(
+                        "phase129/join_k/env_missing",
+                        &format!("env missing required field '{name}'"),
+                        "ensure env layout and env map are built from the same SSOT field list",
+                    )
+                })?;
+                args.push(vid);
+            }
+            Ok(args)
+        };
+
+        // IDs (stable, dev-only)
+        let main_id = JoinFuncId::new(0);
+        let k_then_id = JoinFuncId::new(1);
+        let k_else_id = JoinFuncId::new(2);
+        let join_k_id = JoinFuncId::new(3);
+
+        // main(env)
+        let main_params = alloc_env_params(&env_fields, &mut next_value_id);
+        let mut env_main = build_env_map(&env_fields, &main_params);
+        let mut main_func = JoinFunction::new(main_id, "main".to_string(), main_params);
+
+        // Lower prefix (pre-if) statements into main
+        for n in prefix_nodes {
+            match n {
+                StepNode::Stmt { kind, .. } => match kind {
+                    StepStmtKind::Assign { target, value_ast } => {
+                        if Self::lower_assign_stmt(
+                            target,
+                            value_ast,
+                            &mut main_func.body,
+                            &mut next_value_id,
+                            &mut env_main,
+                        )
+                        .is_err()
+                        {
+                            return Ok(None);
+                        }
+                    }
+                    StepStmtKind::LocalDecl { .. } => {}
+                    _ => {
+                        return Ok(None);
+                    }
+                },
+                _ => {
+                    return Ok(None);
+                }
+            }
+        }
+
+        // Extract return variable and branch bodies.
+        fn split_branch_for_as_last(
+            branch: &StepNode,
+        ) -> Result<(&[StepNode], &crate::mir::control_tree::step_tree::AstNodeHandle), String> {
+            use crate::mir::control_tree::step_tree::StepNode;
+            use crate::mir::control_tree::step_tree::StepStmtKind;
+            use crate::mir::join_ir::lowering::error_tags;
+
+            match branch {
+                StepNode::Stmt { kind, .. } => match kind {
+                    StepStmtKind::Return { value_ast } => {
+                        let ast_handle = value_ast.as_ref().ok_or_else(|| {
+                            error_tags::freeze_with_hint(
+                                "phase129/join_k/branch_return_void",
+                                "branch return must return a variable (not void)",
+                                "use `return x` in both then/else branches",
+                            )
+                        })?;
+                        Ok((&[][..], ast_handle))
+                    }
+                    _ => Err(error_tags::freeze_with_hint(
+                        "phase129/join_k/branch_not_return",
+                        "branch must end with return",
+                        "use `return x` in both then/else branches",
+                    )),
+                },
+                StepNode::Block(nodes) => {
+                    let last = nodes.last().ok_or_else(|| {
+                        error_tags::freeze_with_hint(
+                            "phase129/join_k/branch_empty",
+                            "branch is empty",
+                            "add `return x` as the last statement of the branch",
+                        )
+                    })?;
+                    match last {
+                        StepNode::Stmt { kind, .. } => match kind {
+                            StepStmtKind::Return { value_ast } => {
+                                let ast_handle = value_ast.as_ref().ok_or_else(|| {
+                                    error_tags::freeze_with_hint(
+                                        "phase129/join_k/branch_return_void",
+                                        "branch return must return a variable (not void)",
+                                        "use `return x` in both then/else branches",
+                                    )
+                                })?;
+                                Ok((&nodes[..nodes.len() - 1], ast_handle))
+                            }
+                            _ => Err(error_tags::freeze_with_hint(
+                                "phase129/join_k/branch_not_return",
+                                "branch must end with return",
+                                "add `return x` as the last statement of the branch",
+                            )),
+                        },
+                        _ => Err(error_tags::freeze_with_hint(
+                            "phase129/join_k/branch_last_not_stmt",
+                            "branch last node must be a statement return",
+                            "ensure the branch ends with `return x`",
+                        )),
+                    }
+                }
+                _ => Err(error_tags::freeze_with_hint(
+                    "phase129/join_k/branch_node_unsupported",
+                    "unsupported branch node",
+                    "Phase 129-B supports only Block/Return branches",
+                )),
+            }
+        }
+
+        fn extract_return_var_name(ast_handle: &crate::mir::control_tree::step_tree::AstNodeHandle) -> Result<String, String> {
+            use crate::ast::ASTNode;
+            use crate::mir::join_ir::lowering::error_tags;
+            match ast_handle.0.as_ref() {
+                ASTNode::Variable { name, .. } => Ok(name.clone()),
+                _ => Err(error_tags::freeze_with_hint(
+                    "phase129/join_k/return_expr_unsupported",
+                    "branch return expression must be a variable",
+                    "use `return x` (variable) in both then/else branches",
+                )),
+            }
+        }
+
+        let (cond_ast, then_branch, else_branch) = match if_node {
+            StepNode::If {
+                cond_ast,
+                then_branch,
+                else_branch,
+                ..
+            } => (cond_ast, then_branch.as_ref(), else_branch.as_deref()),
+            _ => unreachable!(),
+        };
+
+        let else_branch = match else_branch {
+            Some(b) => b,
+            None => return Ok(None),
+        };
+
+        let (then_prefix, then_ret_ast) = match split_branch_for_as_last(then_branch) {
+            Ok(v) => v,
+            Err(_msg) => return Ok(None),
+        };
+        let (else_prefix, else_ret_ast) = match split_branch_for_as_last(else_branch) {
+            Ok(v) => v,
+            Err(_msg) => return Ok(None),
+        };
+
+        let then_ret_var = match extract_return_var_name(then_ret_ast) {
+            Ok(v) => v,
+            Err(_msg) => return Ok(None),
+        };
+        let else_ret_var = match extract_return_var_name(else_ret_ast) {
+            Ok(v) => v,
+            Err(_msg) => return Ok(None),
+        };
+
+        if then_ret_var != else_ret_var {
+            return Ok(None);
+        }
+
+        let ret_var = then_ret_var;
+        if !env_layout.writes.iter().any(|w| w == &ret_var) {
+            return Ok(None);
+        }
+
+        // join_k(env_phi): return env_phi[ret_var]
+        let join_k_params = alloc_env_params(&env_fields, &mut next_value_id);
+        let env_join_k = build_env_map(&env_fields, &join_k_params);
+        let ret_vid = env_join_k.get(&ret_var).copied().ok_or_else(|| {
+            error_tags::freeze_with_hint(
+                "phase129/join_k/ret_vid_missing",
+                "return variable not found in join_k env",
+                "ensure env layout includes the return variable in writes",
+            )
+        })?;
+        let mut join_k_func = JoinFunction::new(join_k_id, "join_k".to_string(), join_k_params);
+        join_k_func.body.push(JoinInst::Ret { value: Some(ret_vid) });
+
+        // k_then(env_in): <prefix> ; tailcall join_k(env_out)
+        let then_params = alloc_env_params(&env_fields, &mut next_value_id);
+        let mut env_then = build_env_map(&env_fields, &then_params);
+        let mut then_func = JoinFunction::new(k_then_id, "k_then".to_string(), then_params);
+        for n in then_prefix {
+            match n {
+                StepNode::Stmt { kind, .. } => match kind {
+                    StepStmtKind::Assign { target, value_ast } => {
+                        if Self::lower_assign_stmt(
+                            target,
+                            value_ast,
+                            &mut then_func.body,
+                            &mut next_value_id,
+                            &mut env_then,
+                        )
+                        .is_err()
+                        {
+                            return Ok(None);
+                        }
+                    }
+                    StepStmtKind::LocalDecl { .. } => {}
+                    _ => {
+                        return Ok(None);
+                    }
+                },
+                _ => {
+                    return Ok(None);
+                }
+            }
+        }
+        let then_args = collect_env_args(&env_fields, &env_then)?;
+        then_func.body.push(JoinInst::Call {
+            func: join_k_id,
+            args: then_args,
+            k_next: None,
+            dst: None,
+        });
+
+        // k_else(env_in): <prefix> ; tailcall join_k(env_out)
+        let else_params = alloc_env_params(&env_fields, &mut next_value_id);
+        let mut env_else = build_env_map(&env_fields, &else_params);
+        let mut else_func = JoinFunction::new(k_else_id, "k_else".to_string(), else_params);
+        for n in else_prefix {
+            match n {
+                StepNode::Stmt { kind, .. } => match kind {
+                    StepStmtKind::Assign { target, value_ast } => {
+                        if Self::lower_assign_stmt(
+                            target,
+                            value_ast,
+                            &mut else_func.body,
+                            &mut next_value_id,
+                            &mut env_else,
+                        )
+                        .is_err()
+                        {
+                            return Ok(None);
+                        }
+                    }
+                    StepStmtKind::LocalDecl { .. } => {}
+                    _ => {
+                        return Ok(None);
+                    }
+                },
+                _ => {
+                    return Ok(None);
+                }
+            }
+        }
+        let else_args = collect_env_args(&env_fields, &env_else)?;
+        else_func.body.push(JoinInst::Call {
+            func: join_k_id,
+            args: else_args,
+            k_next: None,
+            dst: None,
+        });
+
+        // main: cond compare + conditional jump to k_then, else to k_else
+        let (lhs_var, op, rhs_literal) = Self::parse_minimal_compare(&cond_ast.0)?;
+        let lhs_vid = env_main.get(&lhs_var).copied().ok_or_else(|| {
+            error_tags::freeze_with_hint(
+                "phase129/join_k/cond_lhs_missing",
+                &format!("condition lhs var '{lhs_var}' not found in env"),
+                "ensure the if condition uses a variable from writes or captured inputs",
+            )
+        })?;
+        let rhs_vid = alloc_value_id(&mut next_value_id);
+        main_func.body.push(JoinInst::Compute(MirLikeInst::Const {
+            dst: rhs_vid,
+            value: ConstValue::Integer(rhs_literal),
+        }));
+        let cond_vid = alloc_value_id(&mut next_value_id);
+        main_func.body.push(JoinInst::Compute(MirLikeInst::Compare {
+            dst: cond_vid,
+            op,
+            lhs: lhs_vid,
+            rhs: rhs_vid,
+        }));
+
+        let main_args = collect_env_args(&env_fields, &env_main)?;
+        main_func.body.push(JoinInst::Jump {
+            cont: k_then_id.as_cont(),
+            args: main_args.clone(),
+            cond: Some(cond_vid),
+        });
+        main_func.body.push(JoinInst::Jump {
+            cont: k_else_id.as_cont(),
+            args: main_args,
+            cond: None,
+        });
+
+        // Build module
+        let mut module = JoinModule::new();
+        module.add_function(main_func);
+        module.add_function(then_func);
+        module.add_function(else_func);
+        module.add_function(join_k_func);
+        module.entry = Some(main_id);
+        module.mark_normalized();
+
+        Ok(Some((module, JoinFragmentMeta::empty())))
+    }
+
     /// Phase 123-128 P1-P3: Lower node from StepTree
     ///
     /// ## Support (Phase 123-128)
@@ -263,10 +655,8 @@ impl StepTreeNormalizedShadowLowererBox {
         env: &mut std::collections::BTreeMap<String, crate::mir::ValueId>,
         contract: &crate::mir::control_tree::step_tree_contract_box::StepTreeContract,
     ) -> Result<(), String> {
-        use crate::ast::{ASTNode, LiteralValue};
         use crate::mir::control_tree::step_tree::{StepNode, StepStmtKind};
         use crate::mir::join_ir::JoinInst;
-        use crate::mir::ValueId;
 
         match node {
             StepNode::Block(nodes) => {
@@ -402,9 +792,8 @@ impl StepTreeNormalizedShadowLowererBox {
         env: &mut std::collections::BTreeMap<String, crate::mir::ValueId>,
         contract: &crate::mir::control_tree::step_tree_contract_box::StepTreeContract,
     ) -> Result<(), String> {
-        use crate::ast::{ASTNode, BinaryOperator};
         use crate::mir::control_tree::step_tree::StepNode;
-        use crate::mir::join_ir::{CompareOp, ConstValue, JoinInst, MirLikeInst};
+        use crate::mir::join_ir::{ConstValue, JoinInst, MirLikeInst};
         use crate::mir::ValueId;
 
         if let StepNode::If {
@@ -995,7 +1384,6 @@ mod tests {
         // Phase 124 P3: Test Return(Variable) when variable is in env (writes)
         use crate::ast::{ASTNode, Span};
         use crate::mir::control_tree::step_tree::AstNodeHandle;
-        use std::collections::BTreeSet;
 
         // Create StepTree with "local x; return x"
         let mut tree = make_if_only_tree();