asm: add num_registers_available() (#11714)

This change adds a way to inspect how many registers an instruction can
address. We assume for now that legacy instructions can always emit the
REX prefix and thus access 16 registers; thus, the only encoding that
allows access to 32 registers is EVEX.

This also refactors some of the top-level generation logic to be less
verbose.

Author: abrown

cranelift/assembler-x64/meta/src/generate.rs

@@ -12,11 +12,8 @@ use cranelift_srcgen::{Formatter, fmtln};
 pub fn rust_assembler(f: &mut Formatter, insts: &[dsl::Inst]) {
     // Generate "all instructions" enum.
     generate_inst_enum(f, insts);
+    generate_inst_impls(f, insts);
     generate_inst_display_impl(f, insts);
-    generate_inst_encode_impl(f, insts);
-    generate_inst_visit_impl(f, insts);
-    generate_inst_is_available_impl(f, insts);
-    generate_inst_features_impl(f, insts);
 
     // Generate per-instruction structs.
     f.empty_line();
@@ -46,22 +43,14 @@ fn generate_inst_enum(f: &mut Formatter, insts: &[dsl::Inst]) {
     });
 }
 
-/// `#[derive(...)]`
-fn generate_derive(f: &mut Formatter) {
-    fmtln!(f, "#[derive(Copy, Clone, Debug)]");
-    fmtln!(
-        f,
-        "#[cfg_attr(any(test, feature = \"fuzz\"), derive(arbitrary::Arbitrary))]"
-    );
-}
-
-/// Adds a custom bound to the `Arbitrary` implementation which ensures that
-/// the associated registers are all `Arbitrary` as well.
-fn generate_derive_arbitrary_bounds(f: &mut Formatter) {
-    fmtln!(
-        f,
-        "#[cfg_attr(any(test, feature = \"fuzz\"), arbitrary(bound = \"R: crate::fuzz::RegistersArbitrary\"))]"
-    );
+/// Helper for emitting `match self { ... }` blocks over all instructions. For each instruction in
+/// `insts`, this generate a separate match arm containing `invoke`.
+fn match_variants(f: &mut Formatter, insts: &[dsl::Inst], invoke: &str) {
+    f.add_block("match self", |f| {
+        for inst in insts.iter().map(|i| i.name()) {
+            fmtln!(f, "Self::{inst}(i) => i.{invoke},");
+        }
+    });
 }
 
 /// `impl std::fmt::Display for Inst { ... }`
@@ -70,81 +59,52 @@ fn generate_inst_display_impl(f: &mut Formatter, insts: &[dsl::Inst]) {
         f.add_block(
             "fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result",
             |f| {
-                f.add_block("match self", |f| {
-                    for inst in insts {
-                        let variant_name = inst.name();
-                        fmtln!(f, "Self::{variant_name}(i) => i.fmt(f),");
-                    }
-                });
+                match_variants(f, insts, "fmt(f)");
             },
         );
     });
 }
 
-/// `impl Inst { fn encode... }`
-fn generate_inst_encode_impl(f: &mut Formatter, insts: &[dsl::Inst]) {
+fn generate_inst_impls(f: &mut Formatter, insts: &[dsl::Inst]) {
     f.add_block("impl<R: Registers> Inst<R>", |f| {
         f.add_block("pub fn encode(&self, b: &mut impl CodeSink)", |f| {
-            f.add_block("match self", |f| {
-                for inst in insts {
-                    let variant_name = inst.name();
-                    fmtln!(f, "Self::{variant_name}(i) => i.encode(b),");
-                }
-            });
+            match_variants(f, insts, "encode(b)");
         });
-    });
-}
-
-/// `impl Inst { fn visit... }`
-fn generate_inst_visit_impl(f: &mut Formatter, insts: &[dsl::Inst]) {
-    fmtln!(f, "impl<R: Registers> Inst<R> {{");
-    f.indent(|f| {
-        fmtln!(
-            f,
-            "pub fn visit(&mut self, v: &mut impl RegisterVisitor<R>) {{"
+        f.add_block(
+            "pub fn visit(&mut self, v: &mut impl RegisterVisitor<R>)",
+            |f| {
+                match_variants(f, insts, "visit(v)");
+            },
         );
-        f.indent(|f| {
-            fmtln!(f, "match self {{");
-            f.indent_push();
-            for inst in insts {
-                let variant_name = inst.name();
-                fmtln!(f, "Self::{variant_name}(i) => i.visit(v),");
-            }
-            f.indent_pop();
-            fmtln!(f, "}}");
-        });
-        fmtln!(f, "}}");
-    });
-    fmtln!(f, "}}");
-}
-
-/// `impl Inst { fn is_available... }`
-fn generate_inst_is_available_impl(f: &mut Formatter, insts: &[dsl::Inst]) {
-    f.add_block("impl<R: Registers> Inst<R>", |f| {
         f.add_block(
             "pub fn is_available(&self, f: &impl AvailableFeatures) -> bool",
             |f| {
-                f.add_block("match self", |f| {
-                    for inst in insts {
-                        let variant_name = inst.name();
-                        fmtln!(f, "Self::{variant_name}(i) => i.is_available(f),");
-                    }
-                });
+                match_variants(f, insts, "is_available(f)");
             },
         );
-    });
-}
-
-/// `impl Inst { fn features... }`
-fn generate_inst_features_impl(f: &mut Formatter, insts: &[dsl::Inst]) {
-    f.add_block("impl<R: Registers> Inst<R>", |f| {
         f.add_block("pub fn features(&self) -> &'static Features", |f| {
-            f.add_block("match self", |f| {
-                for inst in insts {
-                    let variant_name = inst.name();
-                    fmtln!(f, "Self::{variant_name}(i) => i.features(),");
-                }
-            });
+            match_variants(f, insts, "features()");
+        });
+        f.add_block("pub fn num_registers_available(&self) -> usize", |f| {
+            match_variants(f, insts, "num_registers_available()");
         });
     });
 }
+
+/// `#[derive(...)]`
+fn generate_derive(f: &mut Formatter) {
+    fmtln!(f, "#[derive(Copy, Clone, Debug)]");
+    fmtln!(
+        f,
+        "#[cfg_attr(any(test, feature = \"fuzz\"), derive(arbitrary::Arbitrary))]"
+    );
+}
+
+/// Adds a custom bound to the `Arbitrary` implementation which ensures that
+/// the associated registers are all `Arbitrary` as well.
+fn generate_derive_arbitrary_bounds(f: &mut Formatter) {
+    fmtln!(
+        f,
+        "#[cfg_attr(any(test, feature = \"fuzz\"), arbitrary(bound = \"R: crate::fuzz::RegistersArbitrary\"))]"
+    );
+}

cranelift/assembler-x64/meta/src/generate/inst.rs

@@ -68,6 +68,8 @@ impl dsl::Inst {
             self.generate_is_available_function(f);
             f.empty_line();
             self.generate_features_function(f);
+            f.empty_line();
+            self.generate_num_registers_function(f);
         });
     }
 
@@ -244,6 +246,19 @@ impl dsl::Inst {
         });
     }
 
+    /// `fn num_registers_available(&self) -> usize { ... }`
+    fn generate_num_registers_function(&self, f: &mut Formatter) {
+        fmtln!(f, "#[must_use]");
+        f.add_block("pub fn num_registers_available(&self) -> usize", |f| {
+            use dsl::Encoding::*;
+            let n = match self.encoding {
+                Rex(_) | Vex(_) => 16,
+                Evex(_) => 32,
+            };
+            fmtln!(f, "{n}")
+        });
+    }
+
     /// `impl Display for <inst> { ... }`
     pub fn generate_display_impl(&self, f: &mut Formatter) {
         use crate::dsl::Customization::*;