[Rust] Support SSA form properly in low level IL

1 files changed, 14 insertions, 12 deletions

diff --git a/rust/src/low_level_il/function.rs b/rust/src/low_level_il/function.rs
index 21eebc3c..62811b1d 100644
--- a/rust/src/low_level_il/function.rs
+++ b/rust/src/low_level_il/function.rs
@@ -12,15 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use binaryninjacore_sys::BNFreeLowLevelILFunction;
-use binaryninjacore_sys::BNGetLowLevelILOwnerFunction;
-use binaryninjacore_sys::BNLowLevelILFunction;
-use binaryninjacore_sys::BNNewLowLevelILFunctionReference;
-
 use std::fmt::Debug;
 use std::hash::{Hash, Hasher};
 use std::marker::PhantomData;
 
+use binaryninjacore_sys::*;
+
 use crate::architecture::CoreArchitecture;
 use crate::basic_block::BasicBlock;
 use crate::function::Function;
@@ -160,13 +157,8 @@ where
             Function::ref_from_raw(func)
         }
     }
-}
 
-// LLIL basic blocks are not available until the function object
-// is finalized, so ensure we can't try requesting basic blocks
-// during lifting
-impl<F: FunctionForm> LowLevelILFunction<Finalized, F> {
-    pub fn basic_blocks(&self) -> Array<BasicBlock<LowLevelILBlock<Finalized, F>>> {
+    pub fn basic_blocks(&self) -> Array<BasicBlock<LowLevelILBlock<M, F>>> {
         use binaryninjacore_sys::BNGetLowLevelILBasicBlockList;
 
         unsafe {
@@ -178,6 +170,17 @@ impl<F: FunctionForm> LowLevelILFunction<Finalized, F> {
     }
 }
 
+impl<M: FunctionMutability, V: NonSSAVariant> LowLevelILFunction<M, NonSSA<V>> {
+    /// Retrieve the SSA form of the function.
+    pub fn ssa_form(&self) -> Option<Ref<LowLevelILFunction<M, SSA>>> {
+        let handle = unsafe { BNGetLowLevelILSSAForm(self.handle) };
+        if handle.is_null() {
+            return None;
+        }
+        Some(unsafe { LowLevelILFunction::ref_from_raw(handle) })
+    }
+}
+
 // Allow instantiating Lifted IL functions for querying Lifted IL from Architectures
 impl LowLevelILFunction<Mutable, NonSSA<LiftedNonSSA>> {
     // TODO: Document what happens when you pass None for `source_func`.
@@ -201,7 +204,6 @@ impl LowLevelILFunction<Mutable, NonSSA<LiftedNonSSA>> {
 
     pub fn generate_ssa_form(&self) {
         use binaryninjacore_sys::BNGenerateLowLevelILSSAForm;
-
         unsafe { BNGenerateLowLevelILSSAForm(self.handle) };
     }
 }