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

4 files changed, 168 insertions, 51 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,
         }
     }
 }
diff --git a/rust/src/low_level_il/function.rs b/rust/src/low_level_il/function.rs
index 837b45ef..bd5332ea 100644
--- a/rust/src/low_level_il/function.rs
+++ b/rust/src/low_level_il/function.rs
@@ -299,19 +299,16 @@ impl<M: FunctionMutability> LowLevelILFunction<M, SSA> {
     #[must_use]
     pub fn get_ssa_register_uses<R: ArchReg>(
         &self,
-        reg: LowLevelILSSARegisterKind<R>,
+        reg: impl AsRef<LowLevelILSSARegister<R>>,
     ) -> Vec<LowLevelILInstruction<'_, M, SSA>> {
         use binaryninjacore_sys::BNGetLowLevelILSSARegisterUses;
-        let register_id = match reg {
-            LowLevelILSSARegisterKind::Full { kind, .. } => kind.id(),
-            LowLevelILSSARegisterKind::Partial { partial_reg, .. } => partial_reg.id(),
-        };
+        let reg = reg.as_ref();
         let mut count = 0;
         let instrs = unsafe {
             BNGetLowLevelILSSARegisterUses(
                 self.handle,
-                register_id.into(),
-                reg.version() as usize,
+                reg.id().into(),
+                reg.version as usize,
                 &mut count,
             )
         };
@@ -327,19 +324,12 @@ impl<M: FunctionMutability> LowLevelILFunction<M, SSA> {
     #[must_use]
     pub fn get_ssa_register_definition<R: ArchReg>(
         &self,
-        reg: &LowLevelILSSARegisterKind<R>,
+        reg: impl AsRef<LowLevelILSSARegister<R>>,
     ) -> Option<LowLevelILInstruction<'_, M, SSA>> {
         use binaryninjacore_sys::BNGetLowLevelILSSARegisterDefinition;
-        let register_id = match reg {
-            LowLevelILSSARegisterKind::Full { kind, .. } => kind.id(),
-            LowLevelILSSARegisterKind::Partial { partial_reg, .. } => partial_reg.id(),
-        };
+        let reg = reg.as_ref();
         let instr_idx = unsafe {
-            BNGetLowLevelILSSARegisterDefinition(
-                self.handle,
-                register_id.into(),
-                reg.version() as usize,
-            )
+            BNGetLowLevelILSSARegisterDefinition(self.handle, reg.id().into(), reg.version as usize)
         };
         self.instruction_from_index(LowLevelInstructionIndex(instr_idx))
     }
@@ -348,14 +338,11 @@ impl<M: FunctionMutability> LowLevelILFunction<M, SSA> {
     #[must_use]
     pub fn get_ssa_register_value<R: ArchReg>(
         &self,
-        reg: &LowLevelILSSARegisterKind<R>,
+        reg: impl AsRef<LowLevelILSSARegister<R>>,
     ) -> Option<RegisterValue> {
-        let register_id = match reg {
-            LowLevelILSSARegisterKind::Full { kind, .. } => kind.id(),
-            LowLevelILSSARegisterKind::Partial { partial_reg, .. } => partial_reg.id(),
-        };
+        let reg = reg.as_ref();
         let value = unsafe {
-            BNGetLowLevelILSSARegisterValue(self.handle, register_id.into(), reg.version() as usize)
+            BNGetLowLevelILSSARegisterValue(self.handle, reg.id().into(), reg.version as usize)
         };
         if value.state == BNRegisterValueType::UndeterminedValue {
             return None;
diff --git a/rust/src/low_level_il/operation.rs b/rust/src/low_level_il/operation.rs
index fa320060..2f081b5e 100644
--- a/rust/src/low_level_il/operation.rs
+++ b/rust/src/low_level_il/operation.rs
@@ -413,11 +413,11 @@ where
         let full_raw_id = RegisterId(self.op.operands[0] as u32);
         let version = self.op.operands[1] as u32;
         let partial_raw_id = RegisterId(self.op.operands[2] as u32);
-        let full_reg =
-            CoreRegister::new(self.function.arch(), full_raw_id).expect("Bad register ID");
+        let full_reg_kind = LowLevelILRegisterKind::from_raw(&self.function.arch(), full_raw_id)
+            .expect("Bad register ID");
         let partial_reg =
             CoreRegister::new(self.function.arch(), partial_raw_id).expect("Bad register ID");
-        LowLevelILSSARegisterKind::new_partial(full_reg, partial_reg, version)
+        LowLevelILSSARegisterKind::new_partial(full_reg_kind, version, partial_reg)
     }
 
     pub fn source_expr(&self) -> LowLevelILExpression<'func, M, F, ValueExpr> {
@@ -868,11 +868,11 @@ where
         let full_raw_id = RegisterId(self.op.operands[0] as u32);
         let version = self.op.operands[1] as u32;
         let partial_raw_id = RegisterId(self.op.operands[2] as u32);
-        let full_reg =
-            CoreRegister::new(self.function.arch(), full_raw_id).expect("Bad register ID");
+        let full_reg_kind = LowLevelILRegisterKind::from_raw(&self.function.arch(), full_raw_id)
+            .expect("Bad register ID");
         let partial_reg =
             CoreRegister::new(self.function.arch(), partial_raw_id).expect("Bad register ID");
-        LowLevelILSSARegisterKind::new_partial(full_reg, partial_reg, version)
+        LowLevelILSSARegisterKind::new_partial(full_reg_kind, version, partial_reg)
     }
 }
 
diff --git a/rust/tests/low_level_il.rs b/rust/tests/low_level_il.rs
index 79a79964..e6f8d8ff 100644
--- a/rust/tests/low_level_il.rs
+++ b/rust/tests/low_level_il.rs
@@ -259,13 +259,19 @@ fn test_llil_ssa() {
         LowLevelILInstructionKind::SetRegSsa(op) => {
             assert_eq!(op.size(), 4);
             match op.dest_reg() {
-                LowLevelILSSARegisterKind::Full { kind, version } => {
-                    assert_eq!(kind.name(), "edi");
-                    assert_eq!(version, 1);
+                LowLevelILSSARegisterKind::Full(reg) => {
+                    assert_eq!(reg.name(), "edi");
+                    assert_eq!(reg.version, 1);
                 }
                 _ => panic!("Expected LowLevelILSSARegisterKind::Full"),
             }
             assert_eq!(op.source_expr().index, LowLevelExpressionIndex(0));
+
+            // Verify dest_reg does not have a use, so let's verify the ssa register definition.
+            let dest_reg_def = llil_ssa_function
+                .get_ssa_register_definition(op.dest_reg())
+                .expect("Valid ssa reg def");
+            assert_eq!(dest_reg_def.address(), ssa_instr_0.address());
         }
         _ => panic!("Expected SetRegSsa"),
     }
@@ -287,6 +293,23 @@ fn test_llil_ssa() {
             let dest_memory_version = op.dest_memory_version();
             assert_eq!(dest_memory_version, 1);
             assert_eq!(dest_expr.index, LowLevelExpressionIndex(4));
+
+            // Grab the SP register so we can verify its use.
+            let dest_expr_kind = dest_expr.kind();
+            let sub_expr = dest_expr_kind.as_binary_op().unwrap();
+            match sub_expr.left().kind() {
+                LowLevelILExpressionKind::RegSsa(reg) => {
+                    // Verify esp#0 has a single use in the next instruction (same address however).
+                    let sp_0_uses = llil_ssa_function.get_ssa_register_uses(reg.source_reg());
+                    println!("{:?}", sp_0_uses);
+                    assert_eq!(sp_0_uses.len(), 2);
+                    let _next_instr_use = sp_0_uses
+                        .iter()
+                        .find(|inst| inst.index != ssa_instr_1.index)
+                        .expect("Failed to get next instructions use of sp");
+                }
+                _ => panic!("Expected RegSsa"),
+            }
         }
         _ => panic!("Expected StoreSsa"),
     }