Add SelectionKind::Switch, using one scalar::Const for each case'…

…s constant.

src/lib.rs

@@ -895,13 +895,24 @@ pub enum ControlNodeKind {
     },
 }
 
+// FIXME(eddyb) consider interning this, perhaps in a similar vein to `DataInstForm`.
 #[derive(Clone)]
 pub enum SelectionKind {
     /// Two-case selection based on boolean condition, i.e. `if`-`else`, with
     /// the two cases being "then" and "else" (in that order).
     BoolCond,
 
-    SpvInst(spv::Inst),
+    /// `N+1`-case selection based on comparing an integer scrutinee against
+    /// `N` constants, i.e. `switch`, with the last case being the "default"
+    /// (making it the only case without a matching entry in `case_consts`).
+    Switch {
+        // FIXME(eddyb) avoid some of the `scalar::Const` overhead here, as there
+        // is only a single type and we shouldn't need to store more bits per case,
+        // than the actual width of the integer type.
+        // FIXME(eddyb) consider storing this more like sorted compressed keyset,
+        // as there can be no duplicates, and in many cases it may be contiguous.
+        case_consts: Vec<scalar::Const>,
+    },
 }
 
 /// Entity handle for a [`DataInstDef`](crate::DataInstDef) (an SSA instruction).

src/print/mod.rs

@@ -2953,7 +2953,7 @@ impl Print for FuncAt<'_, ControlNode> {
                     (
                         pretty::join_comma_sep(
                             "(",
-                            input_decls_and_uses.clone().zip(initial_inputs).map(
+                            input_decls_and_uses.clone().zip_eq(initial_inputs).map(
                                 |((input_decl, input_use), initial)| {
                                     pretty::Fragment::new([
                                         input_decl.print(printer).insert_name_before_def(
@@ -3483,7 +3483,7 @@ impl SelectionKind {
         mut cases: impl ExactSizeIterator<Item = pretty::Fragment>,
     ) -> pretty::Fragment {
         let kw = |kw| kw_style.apply(kw).into();
-        match *self {
+        match self {
             SelectionKind::BoolCond => {
                 assert_eq!(cases.len(), 2);
                 let [then_case, else_case] = [cases.next().unwrap(), cases.next().unwrap()];
@@ -3500,27 +3500,36 @@ impl SelectionKind {
                     "}".into(),
                 ])
             }
-            SelectionKind::SpvInst(spv::Inst { opcode, ref imms }) => {
-                let header = printer.pretty_spv_inst(
-                    kw_style,
-                    opcode,
-                    imms,
-                    [Some(scrutinee.print(printer))]
-                        .into_iter()
-                        .chain((0..cases.len()).map(|_| None)),
-                );
+            SelectionKind::Switch { case_consts } => {
+                assert_eq!(cases.len(), case_consts.len() + 1);
+
+                let case_patterns = case_consts
+                    .iter()
+                    .map(|&ct| {
+                        let int_to_string = (ct.int_as_u128().map(|x| x.to_string()))
+                            .or_else(|| ct.int_as_i128().map(|x| x.to_string()));
+                        match int_to_string {
+                            Some(v) => printer.numeric_literal_style().apply(v).into(),
+                            None => {
+                                let ct: Const = printer.cx.intern(ct);
+                                ct.print(printer)
+                            }
+                        }
+                    })
+                    .chain(["_".into()]);
 
                 pretty::Fragment::new([
-                    header,
+                    kw("switch"),
+                    " ".into(),
+                    scrutinee.print(printer),
                     " {".into(),
                     pretty::Node::IndentedBlock(
-                        cases
-                            .map(|case| {
+                        case_patterns
+                            .zip_eq(cases)
+                            .map(|(case_pattern, case)| {
                                 pretty::Fragment::new([
                                     pretty::Node::ForceLineSeparation.into(),
-                                    // FIXME(eddyb) this should pull information out
-                                    // of the instruction to be more precise.
-                                    kw("case"),
+                                    case_pattern,
                                     " => {".into(),
                                     pretty::Node::IndentedBlock(vec![case]).into(),
                                     "}".into(),

src/spv/canonical.rs

@@ -165,7 +165,11 @@ def_mappable_ops! {
 }
 
 impl scalar::Const {
-    fn try_decode_from_spv_imms(ty: scalar::Type, imms: &[spv::Imm]) -> Option<scalar::Const> {
+    // HACK(eddyb) this is not private so `spv::lower` can use it for `OpSwitch`.
+    pub(super) fn try_decode_from_spv_imms(
+        ty: scalar::Type,
+        imms: &[spv::Imm],
+    ) -> Option<scalar::Const> {
         // FIXME(eddyb) don't hardcode the 128-bit limitation,
         // but query `scalar::Const` somehow instead.
         if ty.bit_width() > 128 {
@@ -198,7 +202,8 @@ impl scalar::Const {
         }
     }
 
-    fn encode_as_spv_imms(&self) -> impl Iterator<Item = spv::Imm> {
+    // HACK(eddyb) this is not private so `spv::lift` can use it for `OpSwitch`.
+    pub(super) fn encode_as_spv_imms(&self) -> impl Iterator<Item = spv::Imm> {
         let wk = &spec::Spec::get().well_known;
 
         let ty = self.ty();

src/spv/lift.rs

@@ -1309,6 +1309,14 @@ impl LazyInst<'_, '_> {
                 ids: [merge_label_id, continue_label_id].into_iter().collect(),
             },
             Self::Terminator { parent_func, terminator } => {
+                let mut ids: SmallVec<[_; 4]> = terminator
+                    .inputs
+                    .iter()
+                    .map(|&v| value_to_id(parent_func, v))
+                    .chain(terminator.targets.iter().map(|&target| parent_func.label_ids[&target]))
+                    .collect();
+
+                // FIXME(eddyb) move some of this to `spv::canonical`.
                 let inst = match &*terminator.kind {
                     cfg::ControlInstKind::Unreachable => wk.OpUnreachable.into(),
                     cfg::ControlInstKind::Return => {
@@ -1327,23 +1335,21 @@ impl LazyInst<'_, '_> {
                     cfg::ControlInstKind::SelectBranch(SelectionKind::BoolCond) => {
                         wk.OpBranchConditional.into()
                     }
-                    cfg::ControlInstKind::SelectBranch(SelectionKind::SpvInst(inst)) => {
-                        inst.clone()
+                    cfg::ControlInstKind::SelectBranch(SelectionKind::Switch { case_consts }) => {
+                        // HACK(eddyb) move the default case from last back to first.
+                        let default_target = ids.pop().unwrap();
+                        ids.insert(1, default_target);
+
+                        spv::Inst {
+                            opcode: wk.OpSwitch,
+                            imms: case_consts
+                                .iter()
+                                .flat_map(|ct| ct.encode_as_spv_imms())
+                                .collect(),
+                        }
                     }
                 };
-                spv::InstWithIds {
-                    without_ids: inst,
-                    result_type_id: None,
-                    result_id: None,
-                    ids: terminator
-                        .inputs
-                        .iter()
-                        .map(|&v| value_to_id(parent_func, v))
-                        .chain(
-                            terminator.targets.iter().map(|&target| parent_func.label_ids[&target]),
-                        )
-                        .collect(),
-                }
+                spv::InstWithIds { without_ids: inst, result_type_id: None, result_id: None, ids }
             }
             Self::OpFunctionEnd => spv::InstWithIds {
                 without_ids: wk.OpFunctionEnd.into(),