[Rust] Add `LowLevelILSSARegister` type (#7745)

1 files changed, 126 insertions, 19 deletions

diff --git a/rust/src/low_level_il.rs b/rust/src/low_level_il.rs
index bb4614ae..78d4ce9c 100644
--- a/rust/src/low_level_il.rs
+++ b/rust/src/low_level_il.rs
@@ -14,7 +14,7 @@
 
 use std::borrow::Cow;
 use std::fmt;
-
+use std::fmt::{Debug, Display};
 // TODO : provide some way to forbid emitting register reads for certain registers
 // also writing for certain registers (e.g. zero register must prohibit il.set_reg and il.reg
 // (replace with nop or const(0) respectively)
@@ -83,12 +83,18 @@ impl LowLevelILTempRegister {
     }
 }
 
-impl fmt::Debug for LowLevelILTempRegister {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+impl Debug for LowLevelILTempRegister {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
         write!(f, "temp{}", self.temp_id)
     }
 }
 
+impl Display for LowLevelILTempRegister {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        Debug::fmt(self, f)
+    }
+}
+
 impl TryFrom<RegisterId> for LowLevelILTempRegister {
     type Error = ();
 
@@ -142,11 +148,20 @@ impl<R: ArchReg> LowLevelILRegisterKind<R> {
     }
 }
 
-impl<R: ArchReg> fmt::Debug for LowLevelILRegisterKind<R> {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+impl<R: ArchReg> Debug for LowLevelILRegisterKind<R> {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
         match *self {
             LowLevelILRegisterKind::Arch(ref r) => r.fmt(f),
-            LowLevelILRegisterKind::Temp(id) => id.fmt(f),
+            LowLevelILRegisterKind::Temp(ref id) => Debug::fmt(id, f),
+        }
+    }
+}
+
+impl<R: ArchReg> Display for LowLevelILRegisterKind<R> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match *self {
+            LowLevelILRegisterKind::Arch(ref r) => write!(f, "{}", r.name()),
+            LowLevelILRegisterKind::Temp(ref id) => Display::fmt(id, f),
         }
     }
 }
@@ -158,35 +173,127 @@ impl From<LowLevelILTempRegister> for LowLevelILRegisterKind<CoreRegister> {
 }
 
 #[derive(Copy, Clone, Debug)]
+pub struct LowLevelILSSARegister<R: ArchReg> {
+    pub reg: LowLevelILRegisterKind<R>,
+    /// The SSA version of the register.
+    pub version: u32,
+}
+
+impl<R: ArchReg> LowLevelILSSARegister<R> {
+    pub fn new(reg: LowLevelILRegisterKind<R>, version: u32) -> Self {
+        Self { reg, version }
+    }
+
+    pub fn name(&self) -> Cow<'_, str> {
+        self.reg.name()
+    }
+
+    pub fn id(&self) -> RegisterId {
+        self.reg.id()
+    }
+}
+
+impl<R: ArchReg> Display for LowLevelILSSARegister<R> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}#{}", self.reg, self.version)
+    }
+}
+
+/// The kind of SSA register.
+///
+/// An SSA register can exist in two states:
+///
+/// - Full, e.g. `eax` on x86
+/// - Partial, e.g. `al` on x86
+///
+/// If you intend to query for the ssa uses or definitions you must retrieve the physical register
+/// using the function [`LowLevelILSSARegisterKind::physical_reg`] which will give you the actual
+/// [`LowLevelILSSARegister`].
+#[derive(Copy, Clone, Debug)]
 pub enum LowLevelILSSARegisterKind<R: ArchReg> {
-    Full {
-        kind: LowLevelILRegisterKind<R>,
-        version: u32,
-    },
+    /// A full register is one that is not aliasing another, such as `eax` on x86 or `rax` on x86_64.
+    Full(LowLevelILSSARegister<R>),
     Partial {
-        full_reg: CoreRegister,
+        /// This is the non-aliased register.
+        ///
+        /// This register is what is used for dataflow, otherwise the backing storage of aliased registers
+        /// like `al` on x86 would contain separate value information from the physical register `eax`.
+        ///
+        /// NOTE: While this is a [`LowLevelILSSARegister`] temporary registers are not allowed in partial
+        /// assignments, so this will always be an actual architecture register.
+        full_reg: LowLevelILSSARegister<R>,
+        /// This is the aliased register.
+        ///
+        /// On x86 if the register `al` is used that would be considered a partial register, with the
+        /// full register `eax` being used as the backing storage.
         partial_reg: CoreRegister,
-        version: u32,
     },
 }
 
 impl<R: ArchReg> LowLevelILSSARegisterKind<R> {
     pub fn new_full(kind: LowLevelILRegisterKind<R>, version: u32) -> Self {
-        Self::Full { kind, version }
+        Self::Full(LowLevelILSSARegister::new(kind, version))
     }
 
-    pub fn new_partial(full_reg: CoreRegister, partial_reg: CoreRegister, version: u32) -> Self {
+    pub fn new_partial(
+        full_reg: LowLevelILRegisterKind<R>,
+        version: u32,
+        partial_reg: CoreRegister,
+    ) -> Self {
         Self::Partial {
-            full_reg,
+            full_reg: LowLevelILSSARegister::new(full_reg, version),
             partial_reg,
-            version,
         }
     }
 
-    pub fn version(&self) -> u32 {
+    /// This is the non-aliased register used. This should be called when you intend to actually
+    /// query for SSA dataflow information, as a partial register is prohibited from being used.
+    ///
+    /// # Example
+    ///
+    /// On x86 `al` in the LLIL SSA will have a physical register of `eax`.
+    pub fn physical_reg(&self) -> LowLevelILSSARegister<R> {
         match *self {
-            LowLevelILSSARegisterKind::Full { version, .. }
-            | LowLevelILSSARegisterKind::Partial { version, .. } => version,
+            LowLevelILSSARegisterKind::Full(reg) => reg,
+            LowLevelILSSARegisterKind::Partial { full_reg, .. } => full_reg,
+        }
+    }
+
+    /// Gets the displayable register, for partial this will be the partial register name.
+    ///
+    /// # Example
+    ///
+    /// On x86 this will display "al" not "eax".
+    pub fn name(&self) -> Cow<'_, str> {
+        match *self {
+            LowLevelILSSARegisterKind::Full(ref reg) => reg.reg.name(),
+            LowLevelILSSARegisterKind::Partial {
+                ref partial_reg, ..
+            } => partial_reg.name(),
+        }
+    }
+}
+
+impl<R: ArchReg> AsRef<LowLevelILSSARegister<R>> for LowLevelILSSARegisterKind<R> {
+    fn as_ref(&self) -> &LowLevelILSSARegister<R> {
+        match self {
+            LowLevelILSSARegisterKind::Full(reg) => reg,
+            LowLevelILSSARegisterKind::Partial { full_reg, .. } => full_reg,
+        }
+    }
+}
+
+impl<R: ArchReg> From<LowLevelILSSARegister<R>> for LowLevelILSSARegisterKind<R> {
+    fn from(value: LowLevelILSSARegister<R>) -> Self {
+        LowLevelILSSARegisterKind::Full(value)
+    }
+}
+
+impl<R: ArchReg> From<LowLevelILSSARegisterKind<R>> for LowLevelILSSARegister<R> {
+    fn from(value: LowLevelILSSARegisterKind<R>) -> Self {
+        match value {
+            LowLevelILSSARegisterKind::Full(reg) => reg,
+            LowLevelILSSARegisterKind::Partial { full_reg, .. } => full_reg,
         }
     }
 }