diff --git a/compiler/rustc_mir_transform/src/uninhabited_enum_branching.rs b/compiler/rustc_mir_transform/src/uninhabited_enum_branching.rs
index 77bc209539b34..cda9ba9dcc87f 100644
--- a/compiler/rustc_mir_transform/src/uninhabited_enum_branching.rs
+++ b/compiler/rustc_mir_transform/src/uninhabited_enum_branching.rs
@@ -3,8 +3,7 @@
 use crate::MirPass;
 use rustc_data_structures::stable_set::FxHashSet;
 use rustc_middle::mir::{
-    BasicBlock, BasicBlockData, Body, Local, Operand, Rvalue, StatementKind, SwitchTargets,
-    TerminatorKind,
+    BasicBlockData, Body, Local, Operand, Rvalue, StatementKind, SwitchTargets, TerminatorKind,
 };
 use rustc_middle::ty::layout::TyAndLayout;
 use rustc_middle::ty::{Ty, TyCtxt};
@@ -56,7 +55,10 @@ fn variant_discriminants<'tcx>(
     match &layout.variants {
         Variants::Single { index } => {
             let mut res = FxHashSet::default();
-            res.insert(index.as_u32() as u128);
+            res.insert(
+                ty.discriminant_for_variant(tcx, *index)
+                    .map_or(index.as_u32() as u128, |discr| discr.val),
+            );
             res
         }
         Variants::Multiple { variants, .. } => variants
@@ -75,16 +77,9 @@ impl<'tcx> MirPass<'tcx> for UninhabitedEnumBranching {
     }
 
     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
-        if body.source.promoted.is_some() {
-            return;
-        }
-
         trace!("UninhabitedEnumBranching starting for {:?}", body.source);
 
-        let basic_block_count = body.basic_blocks().len();
-
-        for bb in 0..basic_block_count {
-            let bb = BasicBlock::from_usize(bb);
+        for bb in body.basic_blocks().indices() {
             trace!("processing block {:?}", bb);
 
             let Some(discriminant_ty) = get_switched_on_type(&body.basic_blocks()[bb], tcx, body) else {