Handle restrictions on alloc instruction

2 files changed, 81 insertions, 0 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 1fcd9d0..79a5184 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1307,6 +1307,32 @@ impl yaxpeax_arch::LengthedInstruction for InstructionBundle {
 }
 impl yaxpeax_arch::Instruction for InstructionBundle {
     fn well_defined(&self) -> bool {
+        // Alloc has some special rules that need to be checked.
+        let validate_alloc = |insn: &Instruction| {
+            let sof = insn.operands[2].as_unsigned_imm();
+            let sol = insn.operands[3].as_unsigned_imm();
+            let sor = insn.operands[4].as_unsigned_imm();
+
+            // Frame size cannot be bigger than 96 registers.
+            sof <= 95 &&
+            // Rotation size cannot be larger than frame.
+            (sor << 3) <= sof &&
+            // Local size cannot be larger than frame.
+            sol <= sof &&
+            // Alloc cannot be predicated.
+            insn.predicate() == 0
+        };
+
+        if self.instructions[0].opcode() == Opcode::Alloc && !validate_alloc(&self.instructions[0]) {
+            return false;
+        }
+
+        // If alloc is in slot 1, there must be a stop before it.
+        let stop_before_slot1 = BUNDLE_TAGS[self.bundle_tag as usize].map_or(false, |(_, stops)| stops & 0b100 != 0);
+        if self.instructions[1].opcode() == Opcode::Alloc && (!stop_before_slot1 || !validate_alloc(&self.instructions[1])) {
+            return false;
+        }
+
         true
     }
 }
diff --git a/tests/test.rs b/tests/test.rs
index 3dd0ef2..3734e20 100644
--- a/tests/test.rs
+++ b/tests/test.rs
@@ -2332,3 +2332,58 @@ fn test_indirection() {
     let inst = decoder.decode(&mut reader).unwrap();
     assert_eq!(format!("{}", inst), expected);
 }
+
+// alloc has special restrictions that need to be checked.
+#[test]
+fn test_alloc_restrictions() {
+    use yaxpeax_arch::Instruction;
+    let decoder = InstDecoder::default();
+
+    // size of frame cannot be greater than 96.
+    let expected = "[MMI] alloc r34=ar.pfs,127,5,0; break.m 0x0; break.i 0x0";
+    let data = [0x08, 0x10, 0xfd, 0x0b, 0x80, 0x05, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00];
+    let mut reader = U8Reader::new(&data[..]);
+    let inst = decoder.decode(&mut reader).unwrap();
+    assert_eq!(format!("{}", inst), expected);
+    assert!(!inst.well_defined());
+
+    // size of rotation cannot be greater than size of frame.
+    let expected = "[MMI] alloc r34=ar.pfs,5,5,1; break.m 0x0; break.i 0x0";
+    let data = [0x08, 0x10, 0x15, 0x0a, 0x81, 0x05, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00];
+    let mut reader = U8Reader::new(&data[..]);
+    let inst = decoder.decode(&mut reader).unwrap();
+    assert_eq!(format!("{}", inst), expected);
+    assert!(!inst.well_defined());
+
+    // size of locals cannot be greater than size of frame.
+    let expected = "[MMI] alloc r34=ar.pfs,5,6,0; break.m 0x0; break.i 0x0";
+    let data = [0x08, 0x10, 0x15, 0x0c, 0x80, 0x05, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00];
+    let mut reader = U8Reader::new(&data[..]);
+    let inst = decoder.decode(&mut reader).unwrap();
+    assert_eq!(format!("{}", inst), expected);
+    assert!(!inst.well_defined());
+
+    // alloc cannot be predicated.
+    let expected = "[MMI] (p02) alloc r34=ar.pfs,5,5,0; break.m 0x0; break.i 0x0";
+    let data = [0x48, 0x10, 0x15, 0x0a, 0x80, 0x05, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00];
+    let mut reader = U8Reader::new(&data[..]);
+    let inst = decoder.decode(&mut reader).unwrap();
+    assert_eq!(format!("{}", inst), expected);
+    assert!(!inst.well_defined());
+
+    // alloc cannot be in slot 1 of a bundle type without a stop before it.
+    let expected = "[MMI] break.m 0x0; alloc r34=ar.pfs,5,5,0; break.i 0x0";
+    let data = [0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x20, 0x2a, 0x14, 0x00, 0x0b, 0x00, 0x00, 0x00, 0x00, 0x00];
+    let mut reader = U8Reader::new(&data[..]);
+    let inst = decoder.decode(&mut reader).unwrap();
+    assert_eq!(format!("{}", inst), expected);
+    assert!(!inst.well_defined());
+
+    // but alloc can be in slot 1 of a bundle type with a stop before it.
+    let expected = "[MMI] break.m 0x0;; alloc r34=ar.pfs,5,5,0; break.i 0x0";
+    let data = [0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x20, 0x2a, 0x14, 0x00, 0x0b, 0x00, 0x00, 0x00, 0x00, 0x00];
+    let mut reader = U8Reader::new(&data[..]);
+    let inst = decoder.decode(&mut reader).unwrap();
+    assert_eq!(format!("{}", inst), expected);
+    assert!(inst.well_defined());
+}