[Rust] Add Assert and ForceVersion LLIL operations

2 files changed, 185 insertions, 4 deletions

diff --git a/rust/src/low_level_il/instruction.rs b/rust/src/low_level_il/instruction.rs
index 90d2f1d8..19b721af 100644
--- a/rust/src/low_level_il/instruction.rs
+++ b/rust/src/low_level_il/instruction.rs
@@ -240,6 +240,10 @@ where
     Bp(Operation<'func, M, F, operation::NoArgs>),
     Trap(Operation<'func, M, F, operation::Trap>),
     Undef(Operation<'func, M, F, operation::NoArgs>),
+    Assert(Operation<'func, M, F, operation::Assert>),
+    AssertSsa(Operation<'func, M, F, operation::AssertSsa>),
+    ForceVersion(Operation<'func, M, F, operation::ForceVersion>),
+    ForceVersionSsa(Operation<'func, M, F, operation::ForceVersionSsa>),
 
     /// The instruction is an expression.
     Value(LowLevelILExpression<'func, M, F, ValueExpr>),
@@ -332,6 +336,18 @@ where
             LLIL_UNDEF => {
                 LowLevelILInstructionKind::Undef(Operation::new(function, op, expr_index))
             }
+            LLIL_ASSERT => {
+                LowLevelILInstructionKind::Assert(Operation::new(function, op, expr_index))
+            }
+            LLIL_ASSERT_SSA => {
+                LowLevelILInstructionKind::AssertSsa(Operation::new(function, op, expr_index))
+            }
+            LLIL_FORCE_VER => {
+                LowLevelILInstructionKind::ForceVersion(Operation::new(function, op, expr_index))
+            }
+            LLIL_FORCE_VER_SSA => {
+                LowLevelILInstructionKind::ForceVersionSsa(Operation::new(function, op, expr_index))
+            }
             _ => LowLevelILInstructionKind::Value(LowLevelILExpression::new(function, expr_index)),
         }
     }
@@ -384,7 +400,8 @@ where
             }
             Value(e) => visit!(e),
             // Do not have any sub expressions.
-            Nop(_) | NoRet(_) | Goto(_) | Syscall(_) | Bp(_) | Trap(_) | Undef(_) => {}
+            Nop(_) | NoRet(_) | Goto(_) | Syscall(_) | Bp(_) | Trap(_) | Undef(_) | Assert(_)
+            | AssertSsa(_) | ForceVersion(_) | ForceVersionSsa(_) => {}
         }
 
         VisitorAction::Sibling
diff --git a/rust/src/low_level_il/operation.rs b/rust/src/low_level_il/operation.rs
index 7f31f62b..f2ce9d1a 100644
--- a/rust/src/low_level_il/operation.rs
+++ b/rust/src/low_level_il/operation.rs
@@ -13,8 +13,8 @@
 // limitations under the License.
 
 use binaryninjacore_sys::{
-    BNGetLowLevelILByIndex, BNLowLevelILFreeOperandList, BNLowLevelILGetOperandList,
-    BNLowLevelILInstruction,
+    BNGetCachedLowLevelILPossibleValueSet, BNGetLowLevelILByIndex, BNLowLevelILFreeOperandList,
+    BNLowLevelILGetOperandList, BNLowLevelILInstruction,
 };
 
 use super::*;
@@ -22,6 +22,7 @@ use crate::architecture::{
     CoreFlag, CoreFlagGroup, CoreFlagWrite, CoreIntrinsic, CoreRegister, CoreRegisterStack,
     FlagGroupId, FlagId, FlagWriteId, IntrinsicId, RegisterStackId,
 };
+use crate::variable::PossibleValueSet;
 use std::collections::BTreeMap;
 use std::fmt::{Debug, Formatter};
 use std::marker::PhantomData;
@@ -62,7 +63,7 @@ where
         self.op.address
     }
 
-    pub fn get_operand_list(&self, operand_idx: usize) -> Vec<u64> {
+    fn get_operand_list(&self, operand_idx: usize) -> Vec<u64> {
         let mut count = 0;
         let raw_list_ptr = unsafe {
             BNLowLevelILGetOperandList(
@@ -77,6 +78,16 @@ where
         unsafe { BNLowLevelILFreeOperandList(raw_list_ptr) };
         list
     }
+
+    fn get_constraint(&self, operand_idx: usize) -> PossibleValueSet {
+        let raw_pvs = unsafe {
+            BNGetCachedLowLevelILPossibleValueSet(
+                self.function.handle,
+                self.op.operands[operand_idx] as usize,
+            )
+        };
+        PossibleValueSet::from_owned_raw(raw_pvs)
+    }
 }
 
 impl<M, O> Operation<'_, M, NonSSA<LiftedNonSSA>, O>
@@ -1870,6 +1881,155 @@ where
     }
 }
 
+// LLIL_ASSERT
+pub struct Assert;
+
+impl<M, F> Operation<'_, M, F, Assert>
+where
+    M: FunctionMutability,
+    F: FunctionForm,
+{
+    pub fn size(&self) -> usize {
+        self.op.size
+    }
+
+    pub fn source_reg(&self) -> LowLevelILRegisterKind<CoreRegister> {
+        let raw_id = RegisterId(self.op.operands[0] as u32);
+        LowLevelILRegisterKind::from_raw(&self.function.arch(), raw_id).expect("Bad register ID")
+    }
+
+    pub fn constraint(&self) -> PossibleValueSet {
+        self.get_constraint(1)
+    }
+}
+
+impl<M, F> Debug for Operation<'_, M, F, Assert>
+where
+    M: FunctionMutability,
+    F: FunctionForm,
+{
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        f.debug_struct("Assert")
+            .field("size", &self.size())
+            .field("source_reg", &self.source_reg())
+            .field("constraint", &self.constraint())
+            .finish()
+    }
+}
+
+// LLIL_ASSERT_SSA
+pub struct AssertSsa;
+
+impl<M, F> Operation<'_, M, F, AssertSsa>
+where
+    M: FunctionMutability,
+    F: FunctionForm,
+{
+    pub fn size(&self) -> usize {
+        self.op.size
+    }
+
+    pub fn source_reg(&self) -> LowLevelILSSARegisterKind<CoreRegister> {
+        let raw_id = RegisterId(self.op.operands[0] as u32);
+        let reg_kind = LowLevelILRegisterKind::from_raw(&self.function.arch(), raw_id)
+            .expect("Bad register ID");
+        let version = self.op.operands[1] as u32;
+        LowLevelILSSARegisterKind::new_full(reg_kind, version)
+    }
+
+    pub fn constraint(&self) -> PossibleValueSet {
+        self.get_constraint(2)
+    }
+}
+
+impl<M, F> Debug for Operation<'_, M, F, AssertSsa>
+where
+    M: FunctionMutability,
+    F: FunctionForm,
+{
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        f.debug_struct("AssertSsa")
+            .field("size", &self.size())
+            .field("source_reg", &self.source_reg())
+            .field("constraint", &self.constraint())
+            .finish()
+    }
+}
+
+// LLIL_FORCE_VER
+pub struct ForceVersion;
+
+impl<M, F> Operation<'_, M, F, ForceVersion>
+where
+    M: FunctionMutability,
+    F: FunctionForm,
+{
+    pub fn size(&self) -> usize {
+        self.op.size
+    }
+
+    pub fn dest_reg(&self) -> LowLevelILRegisterKind<CoreRegister> {
+        let raw_id = RegisterId(self.op.operands[0] as u32);
+        LowLevelILRegisterKind::from_raw(&self.function.arch(), raw_id).expect("Bad register ID")
+    }
+}
+
+impl<M, F> Debug for Operation<'_, M, F, ForceVersion>
+where
+    M: FunctionMutability,
+    F: FunctionForm,
+{
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        f.debug_struct("ForceVersion")
+            .field("size", &self.size())
+            .field("dest_reg", &self.dest_reg())
+            .finish()
+    }
+}
+
+// LLIL_FORCE_VER_SSA
+pub struct ForceVersionSsa;
+
+impl<M, F> Operation<'_, M, F, ForceVersionSsa>
+where
+    M: FunctionMutability,
+    F: FunctionForm,
+{
+    pub fn size(&self) -> usize {
+        self.op.size
+    }
+
+    pub fn dest_reg(&self) -> LowLevelILSSARegisterKind<CoreRegister> {
+        let raw_id = RegisterId(self.op.operands[0] as u32);
+        let reg_kind = LowLevelILRegisterKind::from_raw(&self.function.arch(), raw_id)
+            .expect("Bad register ID");
+        let version = self.op.operands[1] as u32;
+        LowLevelILSSARegisterKind::new_full(reg_kind, version)
+    }
+
+    pub fn source_reg(&self) -> LowLevelILSSARegisterKind<CoreRegister> {
+        let raw_id = RegisterId(self.op.operands[2] as u32);
+        let reg_kind = LowLevelILRegisterKind::from_raw(&self.function.arch(), raw_id)
+            .expect("Bad register ID");
+        let version = self.op.operands[3] as u32;
+        LowLevelILSSARegisterKind::new_full(reg_kind, version)
+    }
+}
+
+impl<M, F> Debug for Operation<'_, M, F, ForceVersionSsa>
+where
+    M: FunctionMutability,
+    F: FunctionForm,
+{
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        f.debug_struct("ForceVersionSsa")
+            .field("size", &self.size())
+            .field("dest_reg", &self.dest_reg())
+            .field("source_reg", &self.source_reg())
+            .finish()
+    }
+}
+
 // TODO TEST_BIT
 
 pub trait OperationArguments: 'static {}
@@ -1923,3 +2083,7 @@ impl OperationArguments for DoublePrecDivOp {}
 impl OperationArguments for UnaryOp {}
 impl OperationArguments for Condition {}
 impl OperationArguments for UnimplMem {}
+impl OperationArguments for Assert {}
+impl OperationArguments for AssertSsa {}
+impl OperationArguments for ForceVersion {}
+impl OperationArguments for ForceVersionSsa {}