[IR] Track inbound branches to flow nodes.

This CL updates the flow node to track the nodes which are inbound
to the current target. This allows to fix an issue with a loop,
after a continuing with a break, properly eliminating dead code.

This also allows always providing the merge target for an if, it
just knows if it's disconnected by having zero inbound branches.

Bug: tint:1718
Change-Id: Ia9b9dbc734bf1e9cd0c829093c0cb1e470efb32e
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/107800
Commit-Queue: Dan Sinclair <dsinclair@chromium.org>
Reviewed-by: Ben Clayton <bclayton@google.com>
Kokoro: Kokoro <noreply+kokoro@google.com>
diff --git a/src/tint/ir/builder.cc b/src/tint/ir/builder.cc
index 66b195b..594b279 100644
--- a/src/tint/ir/builder.cc
+++ b/src/tint/ir/builder.cc
@@ -39,19 +39,23 @@
     auto* ir_func = ir.flow_nodes.Create<Function>(ast_func);
     ir_func->start_target = CreateBlock();
     ir_func->end_target = CreateTerminator();
+
+    // Function is always branching into the start target
+    ir_func->start_target->inbound_branches.Push(ir_func);
+
     return ir_func;
 }
 
-If* Builder::CreateIf(const ast::Statement* stmt, IfFlags flags) {
+If* Builder::CreateIf(const ast::Statement* stmt) {
     auto* ir_if = ir.flow_nodes.Create<If>(stmt);
-    ir_if->false_target = CreateBlock();
     ir_if->true_target = CreateBlock();
+    ir_if->false_target = CreateBlock();
+    ir_if->merge_target = CreateBlock();
 
-    if (flags == IfFlags::kCreateMerge) {
-        ir_if->merge_target = CreateBlock();
-    } else {
-        ir_if->merge_target = nullptr;
-    }
+    // An if always branches to both the true and false block.
+    ir_if->true_target->inbound_branches.Push(ir_if);
+    ir_if->false_target->inbound_branches.Push(ir_if);
+
     return ir_if;
 }
 
@@ -61,13 +65,17 @@
     ir_loop->continuing_target = CreateBlock();
     ir_loop->merge_target = CreateBlock();
 
+    // A loop always branches to the start block.
+    ir_loop->start_target->inbound_branches.Push(ir_loop);
+
     return ir_loop;
 }
 
-void Builder::Branch(Block* from, const FlowNode* to) {
+void Builder::Branch(Block* from, FlowNode* to) {
     TINT_ASSERT(IR, from);
     TINT_ASSERT(IR, to);
     from->branch_target = to;
+    to->inbound_branches.Push(from);
 }
 
 }  // namespace tint::ir
diff --git a/src/tint/ir/builder.h b/src/tint/ir/builder.h
index 531bcce..a60e7b0 100644
--- a/src/tint/ir/builder.h
+++ b/src/tint/ir/builder.h
@@ -52,20 +52,10 @@
     /// @returns the flow node
     Function* CreateFunction(const ast::Function* func);
 
-    /// Flags used for creation of if flow nodes
-    enum class IfFlags {
-        /// Do not create a merge node, `merge_target` will be `nullptr`
-        kSkipMerge,
-        /// Create the `merge_target` block
-        kCreateMerge,
-    };
-
     /// Creates an if flow node for the given ast::IfStatement or ast::BreakIfStatement
     /// @param stmt the ast::IfStatement or ast::BreakIfStatement
-    /// @param flags the if creation flags. By default the merge block will not be created, pass
-    ///              IfFlags::kCreateMerge if creation is desired.
     /// @returns the flow node
-    If* CreateIf(const ast::Statement* stmt, IfFlags flags = IfFlags::kSkipMerge);
+    If* CreateIf(const ast::Statement* stmt);
 
     /// Creates a loop flow node for the given ast::LoopStatement
     /// @param stmt the ast::LoopStatement
@@ -75,7 +65,7 @@
     /// Branches the given block to the given flow node.
     /// @param from the block to branch from
     /// @param to the node to branch too
-    void Branch(Block* from, const FlowNode* to);
+    void Branch(Block* from, FlowNode* to);
 
     /// The IR module.
     Module ir;
diff --git a/src/tint/ir/builder_impl.cc b/src/tint/ir/builder_impl.cc
index 5f7553f..9c5caff 100644
--- a/src/tint/ir/builder_impl.cc
+++ b/src/tint/ir/builder_impl.cc
@@ -50,23 +50,41 @@
     BuilderImpl* impl_;
 };
 
+bool IsBranched(const Block* b) {
+    return b->branch_target != nullptr;
+}
+
+bool IsConnected(const FlowNode* b) {
+    // Function is always connected as it's the start.
+    if (b->Is<ir::Function>()) {
+        return true;
+    }
+
+    for (auto* parent : b->inbound_branches) {
+        if (IsConnected(parent)) {
+            return true;
+        }
+    }
+    // Getting here means all the incoming branches are disconnected.
+    return false;
+}
+
 }  // namespace
 
 BuilderImpl::BuilderImpl(const Program* program) : builder_(program) {}
 
 BuilderImpl::~BuilderImpl() = default;
 
-void BuilderImpl::BranchTo(const FlowNode* node) {
+void BuilderImpl::BranchTo(FlowNode* node) {
     TINT_ASSERT(IR, current_flow_block_);
-    TINT_ASSERT(IR, !current_flow_block_->branch_target);
+    TINT_ASSERT(IR, !IsBranched(current_flow_block_));
 
     builder_.Branch(current_flow_block_, node);
-    current_flow_block_->branch_target = node;
     current_flow_block_ = nullptr;
 }
 
-void BuilderImpl::BranchToIfNeeded(const FlowNode* node) {
-    if (!current_flow_block_ || current_flow_block_->branch_target) {
+void BuilderImpl::BranchToIfNeeded(FlowNode* node) {
+    if (!current_flow_block_ || IsBranched(current_flow_block_)) {
         return;
     }
     BranchTo(node);
@@ -168,7 +186,7 @@
 
         // If the current flow block has a branch target then the rest of the statements in this
         // block are dead code. Skip them.
-        if (!current_flow_block_ || current_flow_block_->branch_target) {
+        if (!current_flow_block_ || IsBranched(current_flow_block_)) {
             break;
         }
     }
@@ -239,20 +257,19 @@
     // If both branches went somewhere, then they both returned, continued or broke. So,
     // there is no need for the if merge-block and there is nothing to branch to the merge
     // block anyway.
-    if (if_node->true_target->branch_target && if_node->false_target->branch_target) {
+    if (IsBranched(if_node->true_target) && IsBranched(if_node->false_target)) {
         return true;
     }
 
-    if_node->merge_target = builder_.CreateBlock();
     current_flow_block_ = if_node->merge_target;
 
     // If the true branch did not execute control flow, then go to the merge target
-    if (!if_node->true_target->branch_target) {
-        if_node->true_target->branch_target = if_node->merge_target;
+    if (!IsBranched(if_node->true_target)) {
+        builder_.Branch(if_node->true_target, if_node->merge_target);
     }
     // If the false branch did not execute control flow, then go to the merge target
-    if (!if_node->false_target->branch_target) {
-        if_node->false_target->branch_target = if_node->merge_target;
+    if (!IsBranched(if_node->false_target)) {
+        builder_.Branch(if_node->false_target, if_node->merge_target);
     }
 
     return true;
@@ -287,7 +304,12 @@
         BranchToIfNeeded(loop_node->start_target);
     }
 
+    // The loop merge can get disconnected if the loop returns directly, or the continuing target
+    // branches, eventually, to the merge, but nothing branched to the continuing target.
     current_flow_block_ = loop_node->merge_target;
+    if (!IsConnected(loop_node->merge_target)) {
+        current_flow_block_ = nullptr;
+    }
     return true;
 }
 
@@ -337,7 +359,7 @@
 }
 
 bool BuilderImpl::EmitBreakIf(const ast::BreakIfStatement* stmt) {
-    auto* if_node = builder_.CreateIf(stmt, Builder::IfFlags::kCreateMerge);
+    auto* if_node = builder_.CreateIf(stmt);
 
     // TODO(dsinclair): Emit the condition expression into the current block
 
diff --git a/src/tint/ir/builder_impl.h b/src/tint/ir/builder_impl.h
index 8daeebc..1ca2482 100644
--- a/src/tint/ir/builder_impl.h
+++ b/src/tint/ir/builder_impl.h
@@ -133,8 +133,8 @@
   private:
     enum class ControlFlags { kNone, kExcludeSwitch };
 
-    void BranchTo(const ir::FlowNode* node);
-    void BranchToIfNeeded(const ir::FlowNode* node);
+    void BranchTo(ir::FlowNode* node);
+    void BranchToIfNeeded(ir::FlowNode* node);
 
     FlowNode* FindEnclosingControl(ControlFlags flags);
 
diff --git a/src/tint/ir/builder_impl_test.cc b/src/tint/ir/builder_impl_test.cc
index fa37ac5..eb932da 100644
--- a/src/tint/ir/builder_impl_test.cc
+++ b/src/tint/ir/builder_impl_test.cc
@@ -36,6 +36,9 @@
     EXPECT_NE(f->start_target, nullptr);
     EXPECT_NE(f->end_target, nullptr);
 
+    EXPECT_EQ(1u, f->start_target->inbound_branches.Length());
+    EXPECT_EQ(1u, f->end_target->inbound_branches.Length());
+
     EXPECT_EQ(f->start_target->branch_target, f->end_target);
 }
 
@@ -82,6 +85,13 @@
     ASSERT_EQ(1u, m.functions.Length());
     auto* func = m.functions[0];
 
+    EXPECT_EQ(1u, flow->inbound_branches.Length());
+    EXPECT_EQ(1u, flow->true_target->inbound_branches.Length());
+    EXPECT_EQ(1u, flow->false_target->inbound_branches.Length());
+    EXPECT_EQ(2u, flow->merge_target->inbound_branches.Length());
+    EXPECT_EQ(1u, func->start_target->inbound_branches.Length());
+    EXPECT_EQ(1u, func->end_target->inbound_branches.Length());
+
     EXPECT_EQ(func->start_target->branch_target, flow);
     EXPECT_EQ(flow->true_target->branch_target, flow->merge_target);
     EXPECT_EQ(flow->false_target->branch_target, flow->merge_target);
@@ -116,6 +126,13 @@
     ASSERT_EQ(1u, m.functions.Length());
     auto* func = m.functions[0];
 
+    EXPECT_EQ(1u, flow->inbound_branches.Length());
+    EXPECT_EQ(1u, flow->true_target->inbound_branches.Length());
+    EXPECT_EQ(1u, flow->false_target->inbound_branches.Length());
+    EXPECT_EQ(1u, flow->merge_target->inbound_branches.Length());
+    EXPECT_EQ(1u, func->start_target->inbound_branches.Length());
+    EXPECT_EQ(2u, func->end_target->inbound_branches.Length());
+
     EXPECT_EQ(func->start_target->branch_target, flow);
     EXPECT_EQ(flow->true_target->branch_target, func->end_target);
     EXPECT_EQ(flow->false_target->branch_target, flow->merge_target);
@@ -150,6 +167,13 @@
     ASSERT_EQ(1u, m.functions.Length());
     auto* func = m.functions[0];
 
+    EXPECT_EQ(1u, flow->inbound_branches.Length());
+    EXPECT_EQ(1u, flow->true_target->inbound_branches.Length());
+    EXPECT_EQ(1u, flow->false_target->inbound_branches.Length());
+    EXPECT_EQ(1u, flow->merge_target->inbound_branches.Length());
+    EXPECT_EQ(1u, func->start_target->inbound_branches.Length());
+    EXPECT_EQ(2u, func->end_target->inbound_branches.Length());
+
     EXPECT_EQ(func->start_target->branch_target, flow);
     EXPECT_EQ(flow->true_target->branch_target, flow->merge_target);
     EXPECT_EQ(flow->false_target->branch_target, func->end_target);
@@ -179,14 +203,21 @@
     auto* flow = ir_if->As<ir::If>();
     ASSERT_NE(flow->true_target, nullptr);
     ASSERT_NE(flow->false_target, nullptr);
+    ASSERT_NE(flow->merge_target, nullptr);
 
     ASSERT_EQ(1u, m.functions.Length());
     auto* func = m.functions[0];
 
+    EXPECT_EQ(1u, flow->inbound_branches.Length());
+    EXPECT_EQ(1u, flow->true_target->inbound_branches.Length());
+    EXPECT_EQ(1u, flow->false_target->inbound_branches.Length());
+    EXPECT_EQ(0u, flow->merge_target->inbound_branches.Length());
+    EXPECT_EQ(1u, func->start_target->inbound_branches.Length());
+    EXPECT_EQ(2u, func->end_target->inbound_branches.Length());
+
     EXPECT_EQ(func->start_target->branch_target, flow);
     EXPECT_EQ(flow->true_target->branch_target, func->end_target);
     EXPECT_EQ(flow->false_target->branch_target, func->end_target);
-    EXPECT_EQ(flow->merge_target, nullptr);
 }
 
 TEST_F(IRBuilderImplTest, Loop_WithBreak) {
@@ -214,6 +245,13 @@
     ASSERT_EQ(1u, m.functions.Length());
     auto* func = m.functions[0];
 
+    EXPECT_EQ(1u, flow->inbound_branches.Length());
+    EXPECT_EQ(2u, flow->start_target->inbound_branches.Length());
+    EXPECT_EQ(0u, flow->continuing_target->inbound_branches.Length());
+    EXPECT_EQ(1u, flow->merge_target->inbound_branches.Length());
+    EXPECT_EQ(1u, func->start_target->inbound_branches.Length());
+    EXPECT_EQ(1u, func->end_target->inbound_branches.Length());
+
     EXPECT_EQ(func->start_target->branch_target, flow);
     EXPECT_EQ(flow->start_target->branch_target, flow->merge_target);
     EXPECT_EQ(flow->continuing_target->branch_target, flow->start_target);
@@ -260,6 +298,17 @@
     ASSERT_EQ(1u, m.functions.Length());
     auto* func = m.functions[0];
 
+    EXPECT_EQ(1u, loop_flow->inbound_branches.Length());
+    EXPECT_EQ(2u, loop_flow->start_target->inbound_branches.Length());
+    EXPECT_EQ(1u, loop_flow->continuing_target->inbound_branches.Length());
+    EXPECT_EQ(1u, loop_flow->merge_target->inbound_branches.Length());
+    EXPECT_EQ(1u, if_flow->inbound_branches.Length());
+    EXPECT_EQ(1u, if_flow->true_target->inbound_branches.Length());
+    EXPECT_EQ(1u, if_flow->false_target->inbound_branches.Length());
+    EXPECT_EQ(1u, if_flow->merge_target->inbound_branches.Length());
+    EXPECT_EQ(1u, func->start_target->inbound_branches.Length());
+    EXPECT_EQ(1u, func->end_target->inbound_branches.Length());
+
     EXPECT_EQ(func->start_target->branch_target, loop_flow);
     EXPECT_EQ(loop_flow->start_target->branch_target, if_flow);
     EXPECT_EQ(if_flow->true_target->branch_target, loop_flow->merge_target);
@@ -308,6 +357,17 @@
     ASSERT_EQ(1u, m.functions.Length());
     auto* func = m.functions[0];
 
+    EXPECT_EQ(1u, loop_flow->inbound_branches.Length());
+    EXPECT_EQ(2u, loop_flow->start_target->inbound_branches.Length());
+    EXPECT_EQ(1u, loop_flow->continuing_target->inbound_branches.Length());
+    EXPECT_EQ(1u, loop_flow->merge_target->inbound_branches.Length());
+    EXPECT_EQ(1u, break_if_flow->inbound_branches.Length());
+    EXPECT_EQ(1u, break_if_flow->true_target->inbound_branches.Length());
+    EXPECT_EQ(1u, break_if_flow->false_target->inbound_branches.Length());
+    EXPECT_EQ(1u, break_if_flow->merge_target->inbound_branches.Length());
+    EXPECT_EQ(1u, func->start_target->inbound_branches.Length());
+    EXPECT_EQ(1u, func->end_target->inbound_branches.Length());
+
     EXPECT_EQ(func->start_target->branch_target, loop_flow);
     EXPECT_EQ(loop_flow->start_target->branch_target, loop_flow->continuing_target);
     EXPECT_EQ(loop_flow->continuing_target->branch_target, break_if_flow);
@@ -325,7 +385,7 @@
     //   [if false] -> if merge
     //   [if merge] -> loop continuing
     //   [loop continuing] -> loop start
-    //   [loop merge] -> func end
+    //   [loop merge] -> nullptr
     //
     auto* ast_if = If(true, Block(Return()));
     auto* ast_loop = Loop(Block(ast_if, Continue()));
@@ -357,14 +417,147 @@
     ASSERT_EQ(1u, m.functions.Length());
     auto* func = m.functions[0];
 
+    EXPECT_EQ(1u, loop_flow->inbound_branches.Length());
+    EXPECT_EQ(2u, loop_flow->start_target->inbound_branches.Length());
+    EXPECT_EQ(1u, loop_flow->continuing_target->inbound_branches.Length());
+    EXPECT_EQ(0u, loop_flow->merge_target->inbound_branches.Length());
+    EXPECT_EQ(1u, if_flow->inbound_branches.Length());
+    EXPECT_EQ(1u, if_flow->true_target->inbound_branches.Length());
+    EXPECT_EQ(1u, if_flow->false_target->inbound_branches.Length());
+    EXPECT_EQ(1u, if_flow->merge_target->inbound_branches.Length());
+    EXPECT_EQ(1u, func->start_target->inbound_branches.Length());
+    EXPECT_EQ(1u, func->end_target->inbound_branches.Length());
+
     EXPECT_EQ(loop_flow->start_target->branch_target, if_flow);
     EXPECT_EQ(if_flow->true_target->branch_target, func->end_target);
     EXPECT_EQ(if_flow->false_target->branch_target, if_flow->merge_target);
-
     EXPECT_EQ(loop_flow->continuing_target->branch_target, loop_flow->start_target);
 
     EXPECT_EQ(func->start_target->branch_target, ir_loop);
-    EXPECT_EQ(loop_flow->merge_target->branch_target, func->end_target);
+    EXPECT_EQ(loop_flow->merge_target->branch_target, nullptr);
+}
+
+TEST_F(IRBuilderImplTest, Loop_WithOnlyReturn) {
+    // {
+    //   loop {
+    //     return;
+    //     continue;
+    //   }
+    //   if true { return; }
+    // }
+    //
+    // func -> start -> loop -> loop start -> return -> func end
+    //
+    //   [loop continuing] -> loop start
+    //   [loop merge] -> nullptr
+    //
+    // Note, the continue; is here is a dead call, so we won't emit a branch to the continuing block
+    // so the inbound_branches will be zero for continuing.
+    //
+    // The `if` after the `loop` is also eliminated as there is no control-flow path reaching the
+    // block.
+    auto* ast_loop = Loop(Block(Return(), Continue()));
+    WrapInFunction(ast_loop, If(true, Block(Return())));
+    auto& b = Build();
+
+    auto r = b.Build();
+    ASSERT_TRUE(r) << b.error();
+    auto m = r.Move();
+
+    auto* ir_loop = b.FlowNodeForAstNode(ast_loop);
+    ASSERT_NE(ir_loop, nullptr);
+    EXPECT_TRUE(ir_loop->Is<ir::Loop>());
+
+    auto* loop_flow = ir_loop->As<ir::Loop>();
+    ASSERT_NE(loop_flow->start_target, nullptr);
+    ASSERT_NE(loop_flow->continuing_target, nullptr);
+    ASSERT_NE(loop_flow->merge_target, nullptr);
+
+    ASSERT_EQ(1u, m.functions.Length());
+    auto* func = m.functions[0];
+
+    EXPECT_EQ(1u, loop_flow->inbound_branches.Length());
+    EXPECT_EQ(2u, loop_flow->start_target->inbound_branches.Length());
+    EXPECT_EQ(0u, loop_flow->continuing_target->inbound_branches.Length());
+    EXPECT_EQ(0u, loop_flow->merge_target->inbound_branches.Length());
+    EXPECT_EQ(1u, func->start_target->inbound_branches.Length());
+    EXPECT_EQ(1u, func->end_target->inbound_branches.Length());
+
+    EXPECT_EQ(loop_flow->start_target->branch_target, func->end_target);
+    EXPECT_EQ(loop_flow->continuing_target->branch_target, loop_flow->start_target);
+
+    EXPECT_EQ(func->start_target->branch_target, ir_loop);
+}
+
+TEST_F(IRBuilderImplTest, Loop_WithOnlyReturn_ContinuingBreakIf) {
+    // {
+    //   loop {
+    //     return;
+    //     continuing {
+    //       break if true;
+    //     }
+    //   }
+    //   if (true) { return; }
+    // }
+    //
+    // func -> start -> loop -> loop start -> return -> func end
+    //
+    //   [loop continuing] -> break if true
+    //                     -> break if false
+    //   [break if true] -> loop merge
+    //   [break if false] -> if merge
+    //   [break if merge] -> loop start
+    //   [loop merge] -> nullptr
+    //
+    // In this case, the continuing block is dead code, but we don't really know that when parsing
+    // so we end up with a branch into the loop merge target. The loop merge can tell it's dead code
+    // so we can drop the if ater the loop.
+    auto* ast_break_if = BreakIf(true);
+    auto* ast_loop = Loop(Block(Return()), Block(ast_break_if));
+    auto* ast_if = If(true, Block(Return()));
+    WrapInFunction(Block(ast_loop, ast_if));
+    auto& b = Build();
+
+    auto r = b.Build();
+    ASSERT_TRUE(r) << b.error();
+    auto m = r.Move();
+
+    auto* ir_loop = b.FlowNodeForAstNode(ast_loop);
+    ASSERT_NE(ir_loop, nullptr);
+    EXPECT_TRUE(ir_loop->Is<ir::Loop>());
+
+    auto* loop_flow = ir_loop->As<ir::Loop>();
+    ASSERT_NE(loop_flow->start_target, nullptr);
+    ASSERT_NE(loop_flow->continuing_target, nullptr);
+    ASSERT_NE(loop_flow->merge_target, nullptr);
+
+    auto* ir_if = b.FlowNodeForAstNode(ast_if);
+    EXPECT_EQ(ir_if, nullptr);
+
+    auto* ir_break_if = b.FlowNodeForAstNode(ast_break_if);
+    ASSERT_NE(ir_break_if, nullptr);
+    EXPECT_TRUE(ir_break_if->Is<ir::If>());
+
+    auto* break_if_flow = ir_break_if->As<ir::If>();
+    ASSERT_NE(break_if_flow->true_target, nullptr);
+    ASSERT_NE(break_if_flow->false_target, nullptr);
+    ASSERT_NE(break_if_flow->merge_target, nullptr);
+
+    ASSERT_EQ(1u, m.functions.Length());
+    auto* func = m.functions[0];
+
+    EXPECT_EQ(1u, loop_flow->inbound_branches.Length());
+    EXPECT_EQ(2u, loop_flow->start_target->inbound_branches.Length());
+    EXPECT_EQ(0u, loop_flow->continuing_target->inbound_branches.Length());
+    EXPECT_EQ(1u, loop_flow->merge_target->inbound_branches.Length());
+    EXPECT_EQ(1u, func->start_target->inbound_branches.Length());
+    // This is 1 because only the loop branch happens. The subsequent if return is dead code.
+    EXPECT_EQ(1u, func->end_target->inbound_branches.Length());
+
+    EXPECT_EQ(loop_flow->start_target->branch_target, func->end_target);
+    EXPECT_EQ(loop_flow->continuing_target->branch_target, break_if_flow);
+
+    EXPECT_EQ(func->start_target->branch_target, ir_loop);
 }
 
 TEST_F(IRBuilderImplTest, Loop_WithIf_BothBranchesBreak) {
@@ -402,10 +595,22 @@
     auto* if_flow = ir_if->As<ir::If>();
     ASSERT_NE(if_flow->true_target, nullptr);
     ASSERT_NE(if_flow->false_target, nullptr);
+    ASSERT_NE(if_flow->merge_target, nullptr);
 
     ASSERT_EQ(1u, m.functions.Length());
     auto* func = m.functions[0];
 
+    EXPECT_EQ(1u, loop_flow->inbound_branches.Length());
+    EXPECT_EQ(2u, loop_flow->start_target->inbound_branches.Length());
+    EXPECT_EQ(0u, loop_flow->continuing_target->inbound_branches.Length());
+    EXPECT_EQ(2u, loop_flow->merge_target->inbound_branches.Length());
+    EXPECT_EQ(1u, if_flow->inbound_branches.Length());
+    EXPECT_EQ(1u, if_flow->true_target->inbound_branches.Length());
+    EXPECT_EQ(1u, if_flow->false_target->inbound_branches.Length());
+    EXPECT_EQ(0u, if_flow->merge_target->inbound_branches.Length());
+    EXPECT_EQ(1u, func->start_target->inbound_branches.Length());
+    EXPECT_EQ(1u, func->end_target->inbound_branches.Length());
+
     // Note, the `continue` is dead code because both if branches go out of loop, so it just gets
     // dropped.
 
@@ -413,7 +618,6 @@
     EXPECT_EQ(loop_flow->start_target->branch_target, if_flow);
     EXPECT_EQ(if_flow->true_target->branch_target, loop_flow->merge_target);
     EXPECT_EQ(if_flow->false_target->branch_target, loop_flow->merge_target);
-    EXPECT_EQ(if_flow->merge_target, nullptr);
     EXPECT_EQ(loop_flow->continuing_target->branch_target, loop_flow->start_target);
     EXPECT_EQ(loop_flow->merge_target->branch_target, func->end_target);
 }
@@ -551,6 +755,41 @@
     ASSERT_EQ(1u, m.functions.Length());
     auto* func = m.functions[0];
 
+    EXPECT_EQ(1u, loop_flow_a->inbound_branches.Length());
+    EXPECT_EQ(2u, loop_flow_a->start_target->inbound_branches.Length());
+    EXPECT_EQ(1u, loop_flow_a->continuing_target->inbound_branches.Length());
+    EXPECT_EQ(1u, loop_flow_a->merge_target->inbound_branches.Length());
+    EXPECT_EQ(1u, loop_flow_b->inbound_branches.Length());
+    EXPECT_EQ(2u, loop_flow_b->start_target->inbound_branches.Length());
+    EXPECT_EQ(2u, loop_flow_b->continuing_target->inbound_branches.Length());
+    EXPECT_EQ(1u, loop_flow_b->merge_target->inbound_branches.Length());
+    EXPECT_EQ(1u, loop_flow_c->inbound_branches.Length());
+    EXPECT_EQ(2u, loop_flow_c->start_target->inbound_branches.Length());
+    EXPECT_EQ(0u, loop_flow_c->continuing_target->inbound_branches.Length());
+    EXPECT_EQ(1u, loop_flow_c->merge_target->inbound_branches.Length());
+    EXPECT_EQ(1u, loop_flow_d->inbound_branches.Length());
+    EXPECT_EQ(2u, loop_flow_d->start_target->inbound_branches.Length());
+    EXPECT_EQ(1u, loop_flow_d->continuing_target->inbound_branches.Length());
+    EXPECT_EQ(1u, loop_flow_d->merge_target->inbound_branches.Length());
+    EXPECT_EQ(1u, if_flow_a->inbound_branches.Length());
+    EXPECT_EQ(1u, if_flow_a->true_target->inbound_branches.Length());
+    EXPECT_EQ(1u, if_flow_a->false_target->inbound_branches.Length());
+    EXPECT_EQ(1u, if_flow_a->merge_target->inbound_branches.Length());
+    EXPECT_EQ(1u, if_flow_b->inbound_branches.Length());
+    EXPECT_EQ(1u, if_flow_b->true_target->inbound_branches.Length());
+    EXPECT_EQ(1u, if_flow_b->false_target->inbound_branches.Length());
+    EXPECT_EQ(1u, if_flow_b->merge_target->inbound_branches.Length());
+    EXPECT_EQ(1u, if_flow_c->inbound_branches.Length());
+    EXPECT_EQ(1u, if_flow_c->true_target->inbound_branches.Length());
+    EXPECT_EQ(1u, if_flow_c->false_target->inbound_branches.Length());
+    EXPECT_EQ(1u, if_flow_c->merge_target->inbound_branches.Length());
+    EXPECT_EQ(1u, if_flow_d->inbound_branches.Length());
+    EXPECT_EQ(1u, if_flow_d->true_target->inbound_branches.Length());
+    EXPECT_EQ(1u, if_flow_d->false_target->inbound_branches.Length());
+    EXPECT_EQ(1u, if_flow_d->merge_target->inbound_branches.Length());
+    EXPECT_EQ(1u, func->start_target->inbound_branches.Length());
+    EXPECT_EQ(1u, func->end_target->inbound_branches.Length());
+
     EXPECT_EQ(func->start_target->branch_target, loop_flow_a);
     EXPECT_EQ(loop_flow_a->start_target->branch_target, loop_flow_b);
     EXPECT_EQ(loop_flow_b->start_target->branch_target, if_flow_a);
diff --git a/src/tint/ir/flow_node.h b/src/tint/ir/flow_node.h
index 0158af4..2b5ec39 100644
--- a/src/tint/ir/flow_node.h
+++ b/src/tint/ir/flow_node.h
@@ -16,6 +16,7 @@
 #define SRC_TINT_IR_FLOW_NODE_H_
 
 #include "src/tint/castable.h"
+#include "src/tint/utils/vector.h"
 
 namespace tint::ir {
 
@@ -24,6 +25,13 @@
   public:
     ~FlowNode() override;
 
+    /// The list of flow nodes which branch into this node. This list maybe empty for several
+    /// reasons:
+    ///   - Node is a start node
+    ///   - Node is a merge target outside control flow (if that returns in both branches)
+    ///   - Node is a continue target outside control flow (loop that returns)
+    utils::Vector<FlowNode*, 2> inbound_branches;
+
   protected:
     /// Constructor
     FlowNode();
diff --git a/src/tint/ir/if.h b/src/tint/ir/if.h
index 84967b0..2d28aa1 100644
--- a/src/tint/ir/if.h
+++ b/src/tint/ir/if.h
@@ -25,8 +25,7 @@
 
 namespace tint::ir {
 
-/// A flow node representing an if statement. The node always contains a true and a false block. It
-/// may contain a merge block where the true/false blocks will merge too unless they return.
+/// A flow node representing an if statement.
 class If : public Castable<If, FlowNode> {
   public:
     /// Constructor
@@ -41,7 +40,8 @@
     Block* true_target = nullptr;
     /// The false branch block
     Block* false_target = nullptr;
-    /// An optional block where the true/false blocks will branch too if needed.
+    /// An block to reconvert the true/false barnches. The block always exists, but there maybe no
+    /// branches into it. (e.g. if both branches `return`)
     Block* merge_target = nullptr;
 };