adapt OperandVisitor to protected mode too

1 files changed, 83 insertions, 2 deletions

diff --git a/src/protected_mode/mod.rs b/src/protected_mode/mod.rs
index 2403a8a..b0a3f62 100644
--- a/src/protected_mode/mod.rs
+++ b/src/protected_mode/mod.rs
@@ -450,14 +450,14 @@ impl SaeMode {
     /// a human-friendly label for this `SaeMode`:
     ///
     /// ```
-    /// use yaxpeax_x86::long_mode::SaeMode;
+    /// use yaxpeax_x86::protected_mode::SaeMode;
     ///
     /// assert_eq!(SaeMode::RoundNearest.label(), "{rne-sae}");
     /// assert_eq!(SaeMode::RoundDown.label(), "{rd-sae}");
     /// assert_eq!(SaeMode::RoundUp.label(), "{ru-sae}");
     /// assert_eq!(SaeMode::RoundZero.label(), "{rz-sae}");
     /// ```
-    pub fn label(&self) -> &'static str {
+    pub const fn label(&self) -> &'static str {
         match self {
             SaeMode::RoundNearest => "{rne-sae}",
             SaeMode::RoundDown => "{rd-sae}",
@@ -477,6 +477,42 @@ impl SaeMode {
         SAE_MODES[idx]
     }
 }
+
+pub trait OperandVisitor {
+    type Ok;
+    type Error;
+
+    fn visit_reg(&mut self, reg: RegSpec) -> Result<Self::Ok, Self::Error>;
+    fn visit_deref(&mut self, reg: RegSpec) -> Result<Self::Ok, Self::Error>;
+    fn visit_disp(&mut self, reg: RegSpec, disp: i32) -> Result<Self::Ok, Self::Error>;
+    fn visit_reg_scale(&mut self, reg: RegSpec, scale: u8) -> Result<Self::Ok, Self::Error>;
+    fn visit_index_base_scale(&mut self, base: RegSpec, index: RegSpec, scale: u8) -> Result<Self::Ok, Self::Error>;
+    fn visit_index_base_scale_disp(&mut self, base: RegSpec, index: RegSpec, scale: u8, disp: i32) -> Result<Self::Ok, Self::Error>;
+    fn visit_reg_scale_disp(&mut self, reg: RegSpec, scale: u8, disp: i32) -> Result<Self::Ok, Self::Error>;
+    fn visit_i8(&mut self, imm: i8) -> Result<Self::Ok, Self::Error>;
+    fn visit_u8(&mut self, imm: u8) -> Result<Self::Ok, Self::Error>;
+    fn visit_i16(&mut self, imm: i16) -> Result<Self::Ok, Self::Error>;
+    fn visit_u16(&mut self, imm: u16) -> Result<Self::Ok, Self::Error>;
+    fn visit_i32(&mut self, imm: i32) -> Result<Self::Ok, Self::Error>;
+    fn visit_u32(&mut self, imm: u32) -> Result<Self::Ok, Self::Error>;
+    fn visit_abs_u16(&mut self, imm: u16) -> Result<Self::Ok, Self::Error>;
+    fn visit_abs_u32(&mut self, imm: u32) -> Result<Self::Ok, Self::Error>;
+    fn visit_reg_mask_merge(&mut self, base: RegSpec, mask: RegSpec, merge_mode: MergeMode) -> Result<Self::Ok, Self::Error>;
+    fn visit_reg_mask_merge_sae(&mut self, base: RegSpec, mask: RegSpec, merge_mode: MergeMode, sae_mode: SaeMode) -> Result<Self::Ok, Self::Error>;
+    fn visit_reg_mask_merge_sae_noround(&mut self, base: RegSpec, mask: RegSpec, merge_mode: MergeMode) -> Result<Self::Ok, Self::Error>;
+    fn visit_reg_disp_masked(&mut self, base: RegSpec, disp: i32, mask_reg: RegSpec) -> Result<Self::Ok, Self::Error>;
+    fn visit_reg_deref_masked(&mut self, base: RegSpec, mask_reg: RegSpec) -> Result<Self::Ok, Self::Error>;
+    fn visit_reg_scale_masked(&mut self, base: RegSpec, scale: u8, mask_reg: RegSpec) -> Result<Self::Ok, Self::Error>;
+    fn visit_reg_scale_disp_masked(&mut self, base: RegSpec, scale: u8, disp: i32, mask_reg: RegSpec) -> Result<Self::Ok, Self::Error>;
+    fn visit_index_base_masked(&mut self, base: RegSpec, index: RegSpec, mask_reg: RegSpec) -> Result<Self::Ok, Self::Error>;
+    fn visit_index_base_disp_masked(&mut self, base: RegSpec, index: RegSpec, disp: i32, mask_reg: RegSpec) -> Result<Self::Ok, Self::Error>;
+    fn visit_index_base_scale_masked(&mut self, base: RegSpec, index: RegSpec, scale: u8, mask_reg: RegSpec) -> Result<Self::Ok, Self::Error>;
+    fn visit_index_base_scale_disp_masked(&mut self, base: RegSpec, index: RegSpec, scale: u8, disp: i32, mask_reg: RegSpec) -> Result<Self::Ok, Self::Error>;
+    fn visit_absolute_far_address(&mut self, segment: u16, address: u32) -> Result<Self::Ok, Self::Error>;
+
+    fn visit_other(&mut self) -> Result<Self::Ok, Self::Error>;
+}
+
 impl Operand {
     fn from_spec(inst: &Instruction, spec: OperandSpec) -> Operand {
         match spec {
@@ -700,6 +736,51 @@ impl Operand {
             }
         }
     }
+
+    /// provided for parity with [`Instruction::visit_operand`]. this has little utility other than
+    /// to reuse an `OperandVisitor` on an `Operand` directly.
+    pub fn visit<T: OperandVisitor>(&self, visitor: &mut T) -> Result<T::Ok, T::Error> {
+        match self {
+            Operand::Nothing => {
+                visitor.visit_other()
+            }
+            Operand::Register(reg) => {
+                visitor.visit_reg(*reg)
+            }
+            Operand::RegDeref(reg) => {
+                visitor.visit_deref(*reg)
+            }
+            Operand::RegDisp(reg, disp) => {
+                visitor.visit_disp(*reg, *disp)
+            }
+            Operand::ImmediateI8(imm) => visitor.visit_i8(*imm),
+            Operand::ImmediateU8(imm) => visitor.visit_u8(*imm),
+            Operand::ImmediateI16(imm) => visitor.visit_i16(*imm),
+            Operand::ImmediateU16(imm) => visitor.visit_u16(*imm),
+            Operand::ImmediateI32(imm) => visitor.visit_i32(*imm),
+            Operand::ImmediateU32(imm) => visitor.visit_u32(*imm),
+            Operand::DisplacementU16(disp) => visitor.visit_abs_u16(*disp),
+            Operand::DisplacementU32(disp) => visitor.visit_abs_u32(*disp),
+            Operand::RegScale(reg, scale) => visitor.visit_reg_scale(*reg, *scale),
+            Operand::RegScaleDisp(reg, scale, disp) => visitor.visit_reg_scale_disp(*reg, *scale, *disp),
+            Operand::RegIndexBase(_, _) => { /* not actually reachable anymore */ visitor.visit_other() },
+            Operand::RegIndexBaseDisp(_, _, _) => { /* not actually reachable anymore */ visitor.visit_other() },
+            Operand::RegIndexBaseScale(base, index, scale) => visitor.visit_index_base_scale(*base, *index, *scale),
+            Operand::RegIndexBaseScaleDisp(base, index, scale, disp) => visitor.visit_index_base_scale_disp(*base, *index, *scale, *disp),
+            Operand::RegisterMaskMerge(reg, mask, merge) => visitor.visit_reg_mask_merge(*reg, *mask, *merge),
+            Operand::RegisterMaskMergeSae(reg, mask, merge, sae) => visitor.visit_reg_mask_merge_sae(*reg, *mask, *merge, *sae),
+            Operand::RegisterMaskMergeSaeNoround(reg, mask, merge) => visitor.visit_reg_mask_merge_sae_noround(*reg, *mask, *merge),
+            Operand::RegDerefMasked(reg, mask) => visitor.visit_reg_deref_masked(*reg, *mask),
+            Operand::RegDispMasked(reg, disp, mask) => visitor.visit_reg_disp_masked(*reg, *disp, *mask),
+            Operand::RegScaleMasked(reg, scale, mask) => visitor.visit_reg_scale_masked(*reg, *scale, *mask),
+            Operand::RegIndexBaseMasked(_, _, _) => { /* not actually reachable anymore */ visitor.visit_other() },
+            Operand::RegIndexBaseDispMasked(_, _, _, _) => { /* not actually reachable anymore */ visitor.visit_other() },
+            Operand::RegScaleDispMasked(base, scale, disp, mask) => visitor.visit_reg_scale_disp_masked(*base, *scale, *disp, *mask),
+            Operand::RegIndexBaseScaleMasked(base, index, scale, mask) => visitor.visit_index_base_scale_masked(*base, *index, *scale, *mask),
+            Operand::RegIndexBaseScaleDispMasked(base, index, scale, disp, mask) => visitor.visit_index_base_scale_disp_masked(*base, *index, *scale, *disp, *mask),
+            Operand::AbsoluteFarAddress { segment, address } => visitor.visit_absolute_far_address(*segment, *address),
+        }
+    }
 }
 
 #[test]