Rework mlil 2

rust/Cargo.toml

@@ -13,6 +13,7 @@ log = "0.4"
 libc = "0.2"
 rayon = { version = "1.8", optional = true }
 binaryninjacore-sys = { path = "binaryninjacore-sys" }
+strum = { version = "0.26.2", features = ["derive"] }
 
 [patch.crates-io]
 # Patched pdb crate to implement some extra structures

rust/examples/mlil_visitor/src/main.rs

@@ -1,26 +1,26 @@
 use std::env;
 
 use binaryninja::binaryview::BinaryViewExt;
-use binaryninja::mlil::MediumLevelILLiftedOperand;
-use binaryninja::mlil::{MediumLevelILFunction, MediumLevelILLiftedInstruction};
+use binaryninja::mlil;
+use binaryninja::mlil::{MediumLevelILLiftedOperand, MediumLevelILFunction, MediumLevelILLiftedInstruction};
 use binaryninja::types::Variable;
 
 fn print_indent(indent: usize) {
     print!("{:<indent$}", "")
 }
 
-fn print_operation(operation: &MediumLevelILLiftedInstruction) {
+fn print_operation<I: mlil::Form>(operation: &MediumLevelILLiftedInstruction<I>) {
     print!("{}", operation.name());
 }
 
-fn print_variable(func: &MediumLevelILFunction, var: &Variable) {
+fn print_variable<I: mlil::Form>(func: &MediumLevelILFunction<I>, var: &Variable) {
     print!("{}", func.get_function().get_variable_name(var));
 }
 
-fn print_il_expr(instr: &MediumLevelILLiftedInstruction, mut indent: usize) {
+fn print_il_expr<I: mlil::Form>(instr: &MediumLevelILLiftedInstruction<I>, mut indent: usize) {
     print_indent(indent);
     print_operation(instr);
-    println!("");
+    println!();
 
     indent += 1;
 

rust/src/function.rs

@@ -251,7 +251,7 @@ impl Function {
         }
     }
 
-    pub fn medium_level_il(&self) -> Result<Ref<mlil::MediumLevelILFunction>, ()> {
+    pub fn medium_level_il(&self) -> Result<Ref<mlil::MediumLevelILFunction<mlil::NonSSA>>, ()> {
         unsafe {
             let mlil = BNGetFunctionMediumLevelIL(self.handle);
 

rust/src/functionrecognizer.rs

@@ -15,11 +15,11 @@ pub trait FunctionRecognizer {
         false
     }
 
-    fn recognize_medium_level_il(
+    fn recognize_medium_level_il<I: mlil::Form>(
         &self,
         _bv: &BinaryView,
         _func: &Function,
-        _mlil: &mlil::MediumLevelILFunction,
+        _mlil: &mlil::MediumLevelILFunction<I>,
     ) -> bool {
         false
     }
@@ -70,7 +70,7 @@ where
         let custom_handler = unsafe { &*(ctxt as *mut R) };
         let bv = unsafe { BinaryView::from_raw(BNNewViewReference(bv)) };
         let func = unsafe { Function::from_raw(BNNewFunctionReference(func)) };
-        let mlil = unsafe { mlil::MediumLevelILFunction::ref_from_raw(mlil) };
+        let mlil = unsafe { mlil::MediumLevelILFunction::<mlil::NonSSA>::ref_from_raw(mlil) };
         custom_handler.recognize_medium_level_il(bv.as_ref(), func.as_ref(), &mlil)
     }
 

rust/src/mlil/block.rs

@@ -5,15 +5,15 @@ use binaryninjacore_sys::BNGetMediumLevelILIndexForInstruction;
 use crate::basicblock::{BasicBlock, BlockContext};
 use crate::rc::Ref;
 
-use super::{MediumLevelILFunction, MediumLevelILInstruction};
+use super::{Form, MediumLevelILFunction, MediumLevelILInstruction};
 
-pub struct MediumLevelILBlockIter {
-    function: Ref<MediumLevelILFunction>,
+pub struct MediumLevelILBlockIter<I> {
+    function: Ref<MediumLevelILFunction<I>>,
     range: Range<u64>,
 }
 
-impl Iterator for MediumLevelILBlockIter {
-    type Item = MediumLevelILInstruction;
+impl<I: Form> Iterator for MediumLevelILBlockIter<I> {
+    type Item = MediumLevelILInstruction<I>;
 
     fn next(&mut self) -> Option<Self::Item> {
         self.range
@@ -25,36 +25,36 @@ impl Iterator for MediumLevelILBlockIter {
     }
 }
 
-pub struct MediumLevelILBlock {
-    pub(crate) function: Ref<MediumLevelILFunction>,
+pub struct MediumLevelILBlock<I> {
+    pub(crate) function: Ref<MediumLevelILFunction<I>>,
 }
 
-impl core::fmt::Debug for MediumLevelILBlock {
+impl<I> core::fmt::Debug for MediumLevelILBlock<I> {
     fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
         write!(f, "mlil_bb {:?}", self.function)
     }
 }
 
-impl BlockContext for MediumLevelILBlock {
-    type Iter = MediumLevelILBlockIter;
-    type Instruction = MediumLevelILInstruction;
+impl<I: Form> BlockContext for MediumLevelILBlock<I> {
+    type Iter = MediumLevelILBlockIter<I>;
+    type Instruction = MediumLevelILInstruction<I>;
 
-    fn start(&self, block: &BasicBlock<Self>) -> MediumLevelILInstruction {
+    fn start(&self, block: &BasicBlock<Self>) -> MediumLevelILInstruction<I> {
         let expr_idx = unsafe {
             BNGetMediumLevelILIndexForInstruction(self.function.handle, block.raw_start() as usize)
         };
         MediumLevelILInstruction::new(self.function.to_owned(), expr_idx)
     }
 
-    fn iter(&self, block: &BasicBlock<Self>) -> MediumLevelILBlockIter {
+    fn iter(&self, block: &BasicBlock<Self>) -> MediumLevelILBlockIter<I> {
         MediumLevelILBlockIter {
             function: self.function.to_owned(),
             range: block.raw_start()..block.raw_end(),
         }
     }
 }
 
-impl Clone for MediumLevelILBlock {
+impl<I> Clone for MediumLevelILBlock<I> {
     fn clone(&self) -> Self {
         MediumLevelILBlock {
             function: self.function.to_owned(),

rust/src/mlil/function.rs

@@ -15,36 +15,56 @@ use crate::function::Function;
 use crate::function::Location;
 use crate::rc::{Array, Ref, RefCountable};
 
-use super::{MediumLevelILBlock, MediumLevelILInstruction, MediumLevelILLiftedInstruction};
+use super::{
+    Form, MediumLevelILBlock, MediumLevelILInstruction, MediumLevelILLiftedInstruction, NonSSA, SSA,
+};
 
-pub struct MediumLevelILFunction {
+pub struct MediumLevelILFunction<I> {
     pub(crate) handle: *mut BNMediumLevelILFunction,
+    _form: std::marker::PhantomData<I>,
 }
 
-unsafe impl Send for MediumLevelILFunction {}
-unsafe impl Sync for MediumLevelILFunction {}
+unsafe impl<I> Send for MediumLevelILFunction<I> {}
+unsafe impl<I> Sync for MediumLevelILFunction<I> {}
 
-impl Eq for MediumLevelILFunction {}
-impl PartialEq for MediumLevelILFunction {
+impl<I> Eq for MediumLevelILFunction<I> {}
+impl<I> PartialEq for MediumLevelILFunction<I> {
     fn eq(&self, rhs: &Self) -> bool {
         self.get_function().eq(&rhs.get_function())
     }
 }
 
-impl Hash for MediumLevelILFunction {
+impl<I> Hash for MediumLevelILFunction<I> {
     fn hash<H: Hasher>(&self, state: &mut H) {
         self.get_function().hash(state)
     }
 }
 
-impl MediumLevelILFunction {
+impl<I> MediumLevelILFunction<I> {
     pub(crate) unsafe fn ref_from_raw(handle: *mut BNMediumLevelILFunction) -> Ref<Self> {
         debug_assert!(!handle.is_null());
 
-        Self { handle }.to_owned()
+        Self {
+            handle,
+            _form: std::marker::PhantomData,
+        }
+        .to_owned()
+    }
+
+    pub fn instruction_count(&self) -> usize {
+        unsafe { BNGetMediumLevelILInstructionCount(self.handle) }
     }
 
-    pub fn instruction_at<L: Into<Location>>(&self, loc: L) -> Option<MediumLevelILInstruction> {
+    pub fn get_function(&self) -> Ref<Function> {
+        unsafe {
+            let func = BNGetMediumLevelILOwnerFunction(self.handle);
+            Function::from_raw(func)
+        }
+    }
+}
+
+impl<I: Form> MediumLevelILFunction<I> {
+    pub fn instruction_at<L: Into<Location>>(&self, loc: L) -> Option<MediumLevelILInstruction<I>> {
         let loc: Location = loc.into();
         let arch_handle = loc.arch.unwrap();
 
@@ -58,15 +78,21 @@ impl MediumLevelILFunction {
         }
     }
 
-    pub fn instruction_from_idx(&self, expr_idx: usize) -> MediumLevelILInstruction {
+    pub fn instruction_from_idx(&self, expr_idx: usize) -> MediumLevelILInstruction<I> {
         MediumLevelILInstruction::new(self.to_owned(), expr_idx)
     }
 
-    pub fn lifted_instruction_from_idx(&self, expr_idx: usize) -> MediumLevelILLiftedInstruction {
+    pub fn lifted_instruction_from_idx(
+        &self,
+        expr_idx: usize,
+    ) -> MediumLevelILLiftedInstruction<I> {
         self.instruction_from_idx(expr_idx).lift()
     }
 
-    pub fn instruction_from_instruction_idx(&self, instr_idx: usize) -> MediumLevelILInstruction {
+    pub fn instruction_from_instruction_idx(
+        &self,
+        instr_idx: usize,
+    ) -> MediumLevelILInstruction<I> {
         MediumLevelILInstruction::new(self.to_owned(), unsafe {
             BNGetMediumLevelILIndexForInstruction(self.handle, instr_idx)
         })
@@ -75,28 +101,11 @@ impl MediumLevelILFunction {
     pub fn lifted_instruction_from_instruction_idx(
         &self,
         instr_idx: usize,
-    ) -> MediumLevelILLiftedInstruction {
+    ) -> MediumLevelILLiftedInstruction<I> {
         self.instruction_from_instruction_idx(instr_idx).lift()
     }
 
-    pub fn instruction_count(&self) -> usize {
-        unsafe { BNGetMediumLevelILInstructionCount(self.handle) }
-    }
-
-    pub fn ssa_form(&self) -> MediumLevelILFunction {
-        let ssa = unsafe { BNGetMediumLevelILSSAForm(self.handle) };
-        assert!(!ssa.is_null());
-        MediumLevelILFunction { handle: ssa }
-    }
-
-    pub fn get_function(&self) -> Ref<Function> {
-        unsafe {
-            let func = BNGetMediumLevelILOwnerFunction(self.handle);
-            Function::from_raw(func)
-        }
-    }
-
-    pub fn basic_blocks(&self) -> Array<BasicBlock<MediumLevelILBlock>> {
+    pub fn basic_blocks(&self) -> Array<BasicBlock<MediumLevelILBlock<I>>> {
         let mut count = 0;
         let blocks = unsafe { BNGetMediumLevelILBasicBlockList(self.handle, &mut count) };
         let context = MediumLevelILBlock {
@@ -107,18 +116,30 @@ impl MediumLevelILFunction {
     }
 }
 
-impl ToOwned for MediumLevelILFunction {
+impl MediumLevelILFunction<NonSSA> {
+    pub fn ssa_form(&self) -> MediumLevelILFunction<SSA> {
+        let ssa = unsafe { BNGetMediumLevelILSSAForm(self.handle) };
+        assert!(!ssa.is_null());
+        MediumLevelILFunction {
+            handle: ssa,
+            _form: std::marker::PhantomData,
+        }
+    }
+}
+
+impl<I> ToOwned for MediumLevelILFunction<I> {
     type Owned = Ref<Self>;
 
     fn to_owned(&self) -> Self::Owned {
         unsafe { RefCountable::inc_ref(self) }
     }
 }
 
-unsafe impl RefCountable for MediumLevelILFunction {
+unsafe impl<I> RefCountable for MediumLevelILFunction<I> {
     unsafe fn inc_ref(handle: &Self) -> Ref<Self> {
         Ref::new(Self {
             handle: BNNewMediumLevelILFunctionReference(handle.handle),
+            _form: std::marker::PhantomData,
         })
     }
 
@@ -127,7 +148,7 @@ unsafe impl RefCountable for MediumLevelILFunction {
     }
 }
 
-impl core::fmt::Debug for MediumLevelILFunction {
+impl<I> core::fmt::Debug for MediumLevelILFunction<I> {
     fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
         write!(f, "<mlil func handle {:p}>", self.handle)
     }