pulley: add instructions for pushing/popping list of registers

bytecodealliance · Aug 9, 2024 · 5d70d66 · 5d70d66
1 parent 8d5c235
commit 5d70d66
Show file tree

Hide file tree

Showing 10 changed files with 243 additions and 11 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/pulley/Cargo.toml b/pulley/Cargo.toml
@@ -13,6 +13,7 @@ workspace = true
 
 [dependencies]
 arbitrary = { workspace = true, optional = true }
+cranelift-bitset = { path = "../cranelift/bitset" }
 log = { workspace = true }
 sptr = { workspace = true }
 
@@ -21,7 +22,7 @@ env_logger = { workspace = true }
 
 [features]
 std = []
-arbitrary = ["dep:arbitrary", "arbitrary/derive", "std"]
+arbitrary = ["dep:arbitrary", "arbitrary/derive", "std", "cranelift-bitset/arbitrary"]
 encode = []
 decode = []
 disas = ["decode"]

diff --git a/pulley/fuzz/src/interp.rs b/pulley/fuzz/src/interp.rs
@@ -82,10 +82,16 @@ fn op_is_safe_for_fuzzing(op: &Op) -> bool {
         Op::Store64(_) => false,
         Op::Store32SOffset8(_) => false,
         Op::Store64Offset8(_) => false,
-        Op::XPush32(_) | Op::XPush64(_) => true,
-        Op::XPop32(_) | Op::XPop64(_) => true,
-        Op::PushFrame(_) => false,
-        Op::PopFrame(_) => false,
+        Op::XPush32(_)
+        | Op::XPush64(_)
+        | Op::XPush32Many(_)
+        | Op::XPush64Many(_)
+        | Op::XPop32(_)
+        | Op::XPop64(_)
+        | Op::XPop32Many(_)
+        | Op::XPop64Many(_)
+        | Op::PushFrame(_)
+        | Op::PopFrame(_) => false,
         Op::BitcastIntFromFloat32(op::BitcastIntFromFloat32 { .. })
         | Op::BitcastIntFromFloat64(op::BitcastIntFromFloat64 { .. }) => true,
         Op::BitcastFloatFromInt32(_) => true,

diff --git a/pulley/src/decode.rs b/pulley/src/decode.rs
@@ -1,6 +1,8 @@
 //! Decoding support for pulley bytecode.
 
 use alloc::vec::Vec;
+use cranelift_bitset::scalar::ScalarBitSetStorage;
+use cranelift_bitset::ScalarBitSet;
 
 use crate::imms::*;
 use crate::opcode::*;
@@ -394,6 +396,24 @@ impl<R: Reg> Decode for BinaryOperands<R> {
     }
 }
 
+impl<S: Decode + ScalarBitSetStorage> Decode for ScalarBitSet<S> {
+    fn decode<T>(bytecode: &mut T) -> Result<Self, T::Error>
+    where
+        T: BytecodeStream,
+    {
+        S::decode(bytecode).map(ScalarBitSet::from)
+    }
+}
+
+impl<R: Reg + Decode> Decode for RegSet<R> {
+    fn decode<T>(bytecode: &mut T) -> Result<Self, T::Error>
+    where
+        T: BytecodeStream,
+    {
+        ScalarBitSet::decode(bytecode).map(Self::from)
+    }
+}
+
 /// A Pulley bytecode decoder.
 ///
 /// Does not materialize bytecode instructions, instead all decoding methods are

diff --git a/pulley/src/disas.rs b/pulley/src/disas.rs
@@ -133,13 +133,26 @@ impl Disas for PcRelOffset {
     }
 }
 
-impl<R: Disas> Disas for BinaryOperands<R> {
-    fn disas(&self, position: usize, disas: &mut String) {
-        self.dst.disas(position, disas);
-        write!(disas, ", ").unwrap();
-        self.src1.disas(position, disas);
+fn disas_list<T: Disas>(position: usize, disas: &mut String, iter: impl IntoIterator<Item = T>) {
+    let mut iter = iter.into_iter();
+    let Some(first) = iter.next() else { return };
+    first.disas(position, disas);
+
+    for item in iter {
         write!(disas, ", ").unwrap();
-        self.src2.disas(position, disas);
+        item.disas(position, disas);
+    }
+}
+
+impl<R: Reg + Disas> Disas for BinaryOperands<R> {
+    fn disas(&self, position: usize, disas: &mut String) {
+        disas_list(position, disas, [self.dst, self.src1, self.src2])
+    }
+}
+
+impl<R: Reg + Disas> Disas for RegSet<R> {
+    fn disas(&self, position: usize, disas: &mut String) {
+        disas_list(position, disas, *self)
     }
 }
 

diff --git a/pulley/src/encode.rs b/pulley/src/encode.rs
@@ -136,6 +136,15 @@ impl<R: Reg> Encode for BinaryOperands<R> {
     }
 }
 
+impl<R: Reg + Encode> Encode for RegSet<R> {
+    fn encode<E>(&self, sink: &mut E)
+    where
+        E: Extend<u8>,
+    {
+        self.to_bitset().0.encode(sink);
+    }
+}
+
 macro_rules! impl_encoders {
     (
         $(

diff --git a/pulley/src/interp.rs b/pulley/src/interp.rs
@@ -961,23 +961,51 @@ impl OpVisitor for InterpreterVisitor<'_> {
         Continuation::Continue
     }
 
+    fn xpush32_many(&mut self, srcs: RegSet<XReg>) -> Self::Return {
+        for src in srcs {
+            self.xpush32(src);
+        }
+        Continuation::Continue
+    }
+
     fn xpush64(&mut self, src: XReg) -> Self::Return {
         self.state.push(self.state[src].get_u64());
         Continuation::Continue
     }
 
+    fn xpush64_many(&mut self, srcs: RegSet<XReg>) -> Self::Return {
+        for src in srcs {
+            self.xpush64(src);
+        }
+        Continuation::Continue
+    }
+
     fn xpop32(&mut self, dst: XReg) -> Self::Return {
         let val = self.state.pop();
         self.state[dst].set_u32(val);
         Continuation::Continue
     }
 
+    fn xpop32_many(&mut self, dsts: RegSet<XReg>) -> Self::Return {
+        for dst in dsts.into_iter().rev() {
+            self.xpop32(dst);
+        }
+        Continuation::Continue
+    }
+
     fn xpop64(&mut self, dst: XReg) -> Self::Return {
         let val = self.state.pop();
         self.state[dst].set_u64(val);
         Continuation::Continue
     }
 
+    fn xpop64_many(&mut self, dsts: RegSet<XReg>) -> Self::Return {
+        for dst in dsts.into_iter().rev() {
+            self.xpop64(dst);
+        }
+        Continuation::Continue
+    }
+
     /// `push lr; push fp; fp = sp`
     fn push_frame(&mut self) -> Self::Return {
         self.state.push(self.state[SReg::lr].get_ptr::<u8>());

diff --git a/pulley/src/lib.rs b/pulley/src/lib.rs
@@ -128,13 +128,21 @@ macro_rules! for_each_op {
 
             /// `*sp = low32(src); sp += 4`
             xpush32 = XPush32 { src: XReg };
+            /// `for src in srcs { xpush32 src }`
+            xpush32_many = XPush32Many { srcs: RegSet<XReg> };
             /// `*sp = src; sp += 8`
             xpush64 = XPush64 { src: XReg };
+            /// `for src in srcs { xpush64 src }`
+            xpush64_many = XPush64Many { srcs: RegSet<XReg> };
 
             /// `*dst = *sp; sp -= 4`
             xpop32 = XPop32 { dst: XReg };
+            /// `for dst in dsts.rev() { xpop32 dst }`
+            xpop32_many = XPop32Many { dsts: RegSet<XReg> };
             /// `*dst = *sp; sp -= 8`
             xpop64 = XPop64 { dst: XReg };
+            /// `for dst in dsts.rev() { xpop64 dst }`
+            xpop64_many = XPop64Many { dsts: RegSet<XReg> };
 
             /// `low32(dst) = bitcast low32(src) as i32`
             bitcast_int_from_float_32 = BitcastIntFromFloat32 { dst: XReg, src: FReg };

diff --git a/pulley/src/regs.rs b/pulley/src/regs.rs
@@ -1,8 +1,11 @@
 //! Pulley registers.
 
 use core::hash::Hash;
+use core::marker::PhantomData;
 use core::{fmt, ops::Range};
 
+use cranelift_bitset::ScalarBitSet;
+
 /// Trait for common register operations.
 pub trait Reg: Sized + Copy + Eq + Ord + Hash + Into<AnyReg> + fmt::Debug + fmt::Display {
     /// Range of valid register indices.
@@ -185,6 +188,114 @@ impl<R: Reg> BinaryOperands<R> {
     }
 }
 
+/// A set of registers, packed into a 32-bit bitset.
+pub struct RegSet<R> {
+    bitset: ScalarBitSet<u32>,
+    phantom: PhantomData<R>,
+}
+
+impl<R: Reg> RegSet<R> {
+    /// Create a `RegSet` from a `ScalarBitSet`.
+    pub fn from_bitset(bitset: ScalarBitSet<u32>) -> Self {
+        Self {
+            bitset,
+            phantom: PhantomData,
+        }
+    }
+
+    /// Convert a `RegSet` into a `ScalarBitSet`.
+    pub fn to_bitset(self) -> ScalarBitSet<u32> {
+        self.bitset
+    }
+}
+
+impl<R: Reg> From<ScalarBitSet<u32>> for RegSet<R> {
+    fn from(bitset: ScalarBitSet<u32>) -> Self {
+        Self {
+            bitset,
+            phantom: PhantomData,
+        }
+    }
+}
+
+impl<R: Reg> Into<ScalarBitSet<u32>> for RegSet<R> {
+    fn into(self) -> ScalarBitSet<u32> {
+        self.bitset
+    }
+}
+
+impl<R: Reg> IntoIterator for RegSet<R> {
+    type Item = R;
+    type IntoIter = core::iter::FilterMap<cranelift_bitset::scalar::Iter<u32>, fn(u8) -> Option<R>>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.bitset.into_iter().filter_map(R::new)
+    }
+}
+
+impl<R: Reg> FromIterator<R> for RegSet<R> {
+    fn from_iter<I: IntoIterator<Item = R>>(iter: I) -> Self {
+        let mut set = ScalarBitSet::new();
+        for reg in iter {
+            set.insert(reg.to_u8());
+        }
+        RegSet::from(set)
+    }
+}
+
+impl<R: Reg> Default for RegSet<R> {
+    fn default() -> Self {
+        Self {
+            bitset: Default::default(),
+            phantom: Default::default(),
+        }
+    }
+}
+
+impl<R: Reg> Copy for RegSet<R> {}
+impl<R: Reg> Clone for RegSet<R> {
+    fn clone(&self) -> Self {
+        *self
+    }
+}
+
+impl<R: Reg> PartialEq for RegSet<R> {
+    fn eq(&self, other: &Self) -> bool {
+        self.bitset == other.bitset
+    }
+}
+impl<R: Reg> Eq for RegSet<R> {}
+
+impl<R: Reg> PartialOrd for RegSet<R> {
+    fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
+        self.bitset.partial_cmp(&other.bitset)
+    }
+}
+impl<R: Reg> Ord for RegSet<R> {
+    fn cmp(&self, other: &Self) -> core::cmp::Ordering {
+        self.bitset.cmp(&other.bitset)
+    }
+}
+
+impl<R: Reg> Hash for RegSet<R> {
+    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
+        self.bitset.hash(state);
+    }
+}
+
+impl<R: Reg> fmt::Debug for RegSet<R> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_set().entries(self.into_iter()).finish()
+    }
+}
+
+#[cfg(feature = "arbitrary")]
+impl<'a, R: Reg> arbitrary::Arbitrary<'a> for RegSet<R> {
+    fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
+        ScalarBitSet::arbitrary(u).map(Self::from)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

diff --git a/pulley/tests/all/disas.rs b/pulley/tests/all/disas.rs
@@ -50,3 +50,38 @@ fn simple() {
         "#,
     );
 }
+
+#[test]
+fn push_pop_many() {
+    assert_disas(
+        &[
+            // Prologue.
+            Op::PushFrame(PushFrame {}),
+            Op::XPush32Many(XPush32Many {
+                srcs: RegSet::from_iter([XReg::x0, XReg::x1, XReg::x2, XReg::x3, XReg::x4]),
+            }),
+            // Function body.
+            Op::Xadd32(Xadd32 {
+                operands: BinaryOperands {
+                    dst: XReg::x0,
+                    src1: XReg::x0,
+                    src2: XReg::x1,
+                },
+            }),
+            // Epilogue.
+            Op::XPop32Many(XPop32Many {
+                dsts: RegSet::from_iter([XReg::x0, XReg::x1, XReg::x2, XReg::x3, XReg::x4]),
+            }),
+            Op::PopFrame(PopFrame {}),
+            Op::Ret(Ret {}),
+        ],
+        r#"
+       0: 2a                              push_frame
+       1: 2d 1f 00 00 00                  xpush32_many x0, x1, x2, x3, x4
+       6: 12 00 04                        xadd32 x0, x0, x1
+       9: 31 1f 00 00 00                  xpop32_many x0, x1, x2, x3, x4
+       e: 2b                              pop_frame
+       f: 00                              ret
+        "#,
+    );
+}