From ac75d66992e80174f5724040ce40c94694372bac Mon Sep 17 00:00:00 2001
From: Eduard-Mihai Burtescu <eddyb@vectorware.com>
Date: Fri, 26 Sep 2025 17:51:48 +0300
Subject: [PATCH] cf: fully structurize functions that never return, as well.

---
 src/cf/structurize.rs  | 50 +++++++++++++++++++++---------------------
 src/passes/legalize.rs |  5 +++--
 2 files changed, 28 insertions(+), 27 deletions(-)

diff --git a/src/cf/structurize.rs b/src/cf/structurize.rs
index f8e08c0d..7d358120 100644
--- a/src/cf/structurize.rs
+++ b/src/cf/structurize.rs
@@ -4,7 +4,7 @@
 
 use crate::cf::SelectionKind;
 use crate::cf::unstructured::{
-    ControlFlowGraph, ControlInst, ControlInstKind, IncomingEdgeCount, LoopFinder, TraversalState,
+    ControlInst, ControlInstKind, IncomingEdgeCount, LoopFinder, TraversalState,
 };
 use crate::transform::{InnerInPlaceTransform as _, Transformed, Transformer};
 use crate::{
@@ -55,6 +55,8 @@ pub struct Structurizer<'a> {
     /// Scrutinee value for [`SelectionKind::BoolCond`], for the "else" case.
     const_false: Const,
 
+    func_ret_types: &'a [Type],
+
     func_def_body: &'a mut FuncDefBody,
 
     // FIXME(eddyb) this feels a bit inefficient (are many-exit loops rare?).
@@ -546,7 +548,12 @@ struct ClaimedRegion {
 }
 
 impl<'a> Structurizer<'a> {
-    pub fn new(cx: &'a Context, func_def_body: &'a mut FuncDefBody) -> Self {
+    pub fn new(cx: &'a Context, func_decl: &'a mut crate::FuncDecl) -> Self {
+        // TODO(eddyb) find a way to change this so it doesn't need to panic.
+        let crate::DeclDef::Present(func_def_body) = &mut func_decl.def else {
+            unreachable!();
+        };
+
         // FIXME(eddyb) SPIR-T should have native booleans itself.
         let wk = &spv::spec::Spec::get().well_known;
         let type_bool = cx.intern(TypeKind::SpvInst {
@@ -614,6 +621,13 @@ impl<'a> Structurizer<'a> {
             const_true,
             const_false,
 
+            func_ret_types: {
+                let is_void = match &cx[func_decl.ret_type].kind {
+                    TypeKind::SpvInst { spv_inst, .. } => spv_inst.opcode == wk.OpTypeVoid,
+                    _ => false,
+                };
+                if is_void { &[][..] } else { std::slice::from_ref(&func_decl.ret_type) }
+            },
             func_def_body,
 
             loop_header_to_exit_targets,
@@ -661,30 +675,16 @@ impl<'a> Structurizer<'a> {
         };
 
         match func_body_deferred_edges {
-            // FIXME(eddyb) also support structured return when the whole body
-            // is divergent, by generating undef constants (needs access to the
-            // whole `FuncDecl`, not just `FuncDefBody`, to get the right types).
+            // FIXME(eddyb) is there a way to do this without returning `undef`s?
             DeferredEdgeBundleSet::Unreachable => {
-                // HACK(eddyb) replace the CFG with one that only contains an
-                // `Unreachable` terminator for the body, comparable to what
-                // `rebuild_cfg_from_unclaimed_region_deferred_edges` would do
-                // in the general case (but special-cased because this is very
-                // close to being structurizable, just needs a bit of plumbing).
-                let mut control_inst_on_exit_from = EntityOrientedDenseMap::new();
-                control_inst_on_exit_from.insert(
-                    self.func_def_body.body,
-                    ControlInst {
-                        attrs: AttrSet::default(),
-                        kind: ControlInstKind::Unreachable,
-                        inputs: [].into_iter().collect(),
-                        targets: [].into_iter().collect(),
-                        target_inputs: FxIndexMap::default(),
-                    },
-                );
-                self.func_def_body.unstructured_cfg = Some(ControlFlowGraph {
-                    control_inst_on_exit_from,
-                    loop_merge_to_loop_header: Default::default(),
-                });
+                let undef_outputs = self
+                    .func_ret_types
+                    .iter()
+                    .map(|&ty| Value::Const(self.const_undef(ty)))
+                    .collect();
+                let body_def = self.func_def_body.at_mut_body().def();
+                body_def.outputs = undef_outputs;
+                self.func_def_body.unstructured_cfg = None;
             }
 
             // Structured return, the function is fully structurized.
diff --git a/src/passes/legalize.rs b/src/passes/legalize.rs
index 3f4cf5ba..885400ac 100644
--- a/src/passes/legalize.rs
+++ b/src/passes/legalize.rs
@@ -21,8 +21,9 @@ pub fn structurize_func_cfgs(module: &mut Module) {
     }
 
     for &func in &collector.seen_funcs {
-        if let DeclDef::Present(func_def_body) = &mut module.funcs[func].def {
-            cf::structurize::Structurizer::new(cx, func_def_body).structurize_func();
+        let func_decl = &mut module.funcs[func];
+        if let DeclDef::Present(_) = func_decl.def {
+            cf::structurize::Structurizer::new(cx, func_decl).structurize_func();
         }
     }
 }