From 02cacc08f27634bffd1ae0f4dfafb7ba2a5bd829 Mon Sep 17 00:00:00 2001
From: Eduard-Mihai Burtescu <eddyb@vectorware.com>
Date: Tue, 9 Sep 2025 13:47:07 +0300
Subject: [PATCH] spv/lift: elide empty targets of an `if`-`else`/`switch`, as
 allowed by SPIR-V.

---
 src/spv/lift.rs | 62 ++++++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 56 insertions(+), 6 deletions(-)

diff --git a/src/spv/lift.rs b/src/spv/lift.rs
index 5bbb655..c0c8348 100644
--- a/src/spv/lift.rs
+++ b/src/spv/lift.rs
@@ -863,22 +863,72 @@ impl<'a> FuncLifting<'a> {
         // HACK(eddyb) this takes advantage of `blocks` being an `IndexMap`,
         // to iterate at the same time as mutating other entries.
         for block_idx in (0..blocks.len()).rev() {
-            let BlockLifting { terminator: original_terminator, .. } = &blocks[block_idx];
+            // HACK(eddyb) elide empty cases of an `if`-`else`/`switch`, as
+            // SPIR-V allows their targets to just be the whole merge block
+            // (the same one that `OpSelectionMerge` describes).
+            let block = &blocks[block_idx];
+            if let (cfg::ControlInstKind::SelectBranch(_), Some(Merge::Selection(merge_point))) =
+                (&*block.terminator.kind, block.terminator.merge)
+            {
+                for target_idx in 0..block.terminator.targets.len() {
+                    let block = &blocks[block_idx];
+                    let target = block.terminator.targets[target_idx];
+                    if !block
+                        .terminator
+                        .target_phi_values
+                        .get(&target)
+                        .copied()
+                        .unwrap_or_default()
+                        .is_empty()
+                    {
+                        continue;
+                    }
+
+                    let target_is_trivial_branch = {
+                        let BlockLifting {
+                            phis,
+                            insts,
+                            terminator:
+                                Terminator { attrs, kind, inputs, targets, target_phi_values, merge },
+                        } = &blocks[&target];
+
+                        (phis.is_empty()
+                            && insts.iter().all(|insts| insts.is_empty())
+                            && *attrs == AttrSet::default()
+                            && matches!(**kind, cfg::ControlInstKind::Branch)
+                            && inputs.is_empty()
+                            && targets.len() == 1
+                            && target_phi_values.is_empty()
+                            && merge.is_none())
+                        .then(|| targets[0])
+                    };
+                    if let Some(target_of_target) = target_is_trivial_branch {
+                        // FIXME(eddyb) what does it mean for this to not be true?
+                        // (can it even happen?)
+                        if target_of_target == merge_point {
+                            blocks[block_idx].terminator.targets[target_idx] = target_of_target;
+                            *use_counts.get_mut(&target).unwrap() -= 1;
+                            *use_counts.get_mut(&target_of_target).unwrap() += 1;
+                        }
+                    }
+                }
+            }
 
+            let block = &blocks[block_idx];
             let is_trivial_branch = {
                 let Terminator { attrs, kind, inputs, targets, target_phi_values, merge } =
-                    original_terminator;
+                    &block.terminator;
 
-                *attrs == AttrSet::default()
+                (*attrs == AttrSet::default()
                     && matches!(**kind, cfg::ControlInstKind::Branch)
                     && inputs.is_empty()
                     && targets.len() == 1
                     && target_phi_values.is_empty()
-                    && merge.is_none()
+                    && merge.is_none())
+                .then(|| targets[0])
             };
 
-            if is_trivial_branch {
-                let target = original_terminator.targets[0];
+            if let Some(target) = is_trivial_branch {
                 let target_use_count = use_counts.get_mut(&target).unwrap();
 
                 if *target_use_count == 1 {