Find the unified merge block that all literal match arms jump to

Author: cburgdorf

crates/codegen/src/yul/emitter.rs

@@ -5,7 +5,7 @@ use hir::hir_def::{
 };
 use mir::{
     BasicBlockId, CallOrigin, LoopInfo, MirFunction, Terminator, ValueId, ValueOrigin,
-    ir::{MatchArmPattern, SwitchOrigin, SwitchValue},
+    ir::{MatchArmPattern, SwitchOrigin, SwitchTarget, SwitchValue},
     lower_module,
 };
 use rustc_hash::FxHashMap;
@@ -207,6 +207,10 @@ impl<'db> YulEmitter<'db> {
                         "  default ",
                         vec![YulDoc::line(format!("{temp} := {default_expr}"))],
                     ));
+                    if let Some(merge_block) = self.match_merge_block(targets, default)? {
+                        let next_docs = self.emit_block_with_ctx(merge_block, loop_ctx, state)?;
+                        docs.extend(next_docs);
+                    }
                     Ok(docs)
                 }
                 SwitchOrigin::None => {
@@ -263,6 +267,45 @@ impl<'db> YulEmitter<'db> {
         }
     }
 
+    /// Finds the unified merge block that all literal match arms jump to, if any.
+    fn match_merge_block(
+        &self,
+        targets: &[SwitchTarget],
+        default: BasicBlockId,
+    ) -> Result<Option<BasicBlockId>, YulError> {
+        let mut merge = None;
+        for block_id in targets
+            .iter()
+            .map(|target| target.block)
+            .chain(std::iter::once(default))
+        {
+            let block = self
+                .mir_func
+                .body
+                .blocks
+                .get(block_id.index())
+                .ok_or_else(|| YulError::Unsupported("invalid block in match lowering".into()))?;
+            match block.terminator {
+                Terminator::Goto { target } => match merge {
+                    Some(existing) if existing != target => {
+                        return Err(YulError::Unsupported(
+                            "match arms must converge to a single merge block".into(),
+                        ));
+                    }
+                    None => merge = Some(target),
+                    _ => {}
+                },
+                Terminator::Unreachable => {}
+                _ => {
+                    return Err(YulError::Unsupported(
+                        "match arms must jump to a merge block".into(),
+                    ));
+                }
+            }
+        }
+        Ok(merge)
+    }
+
     /// Looks up metadata about the loop that starts at `header`, if it exists.
     fn loop_info(&self, header: BasicBlockId) -> Option<LoopInfo> {
         self.mir_func.body.loop_headers.get(&header).copied()

crates/codegen/tests/fixtures/match_literal.snap

@@ -13,6 +13,7 @@ input_file: tests/fixtures/match_literal.fe
       default {
         v0 := 2
       }
+    ret := v0
   }
 }
 {
@@ -28,6 +29,7 @@ input_file: tests/fixtures/match_literal.fe
       default {
         v0 := 3
       }
+    ret := v0
   }
 }
 {
@@ -43,6 +45,7 @@ input_file: tests/fixtures/match_literal.fe
       default {
         v0 := 3
       }
+    ret := v0
   }
 }
 {
@@ -58,6 +61,7 @@ input_file: tests/fixtures/match_literal.fe
       default {
         v0 := 3
       }
+    ret := v0
   }
 }
 {
@@ -73,6 +77,7 @@ input_file: tests/fixtures/match_literal.fe
       default {
         v0 := 3
       }
+    ret := v0
   }
 }
 {
@@ -88,6 +93,7 @@ input_file: tests/fixtures/match_literal.fe
       default {
         v0 := 3
       }
+    ret := v0
   }
 }
 {
@@ -103,6 +109,7 @@ input_file: tests/fixtures/match_literal.fe
       default {
         v0 := 3
       }
+    ret := v0
   }
 }
 {
@@ -118,6 +125,7 @@ input_file: tests/fixtures/match_literal.fe
       default {
         v0 := 3
       }
+    ret := v0
   }
 }
 {
@@ -133,6 +141,7 @@ input_file: tests/fixtures/match_literal.fe
       default {
         v0 := 3
       }
+    ret := v0
   }
 }
 {
@@ -148,6 +157,7 @@ input_file: tests/fixtures/match_literal.fe
       default {
         v0 := 3
       }
+    ret := v0
   }
 }
 {
@@ -163,6 +173,7 @@ input_file: tests/fixtures/match_literal.fe
       default {
         v0 := 3
       }
+    ret := v0
   }
 }
 {
@@ -178,6 +189,7 @@ input_file: tests/fixtures/match_literal.fe
       default {
         v0 := 3
       }
+    ret := v0
   }
 }
 {
@@ -193,5 +205,6 @@ input_file: tests/fixtures/match_literal.fe
       default {
         v0 := 3
       }
+    ret := v0
   }
 }