1 files changed, 10 insertions, 23 deletions
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;