[Rust] Refactor `AnalysisProgress` returned from `BinaryViewExt::analysis_progress`

2 files changed, 59 insertions, 14 deletions

diff --git a/rust/src/binary_view.rs b/rust/src/binary_view.rs
index 0ab2504c..d0b22103 100644
--- a/rust/src/binary_view.rs
+++ b/rust/src/binary_view.rs
@@ -62,6 +62,7 @@ use crate::variable::DataVariable;
 use crate::{Endianness, BN_FULL_CONFIDENCE};
 use std::collections::HashMap;
 use std::ffi::{c_char, c_void, CString};
+use std::fmt::{Display, Formatter};
 use std::ops::Range;
 use std::path::Path;
 use std::ptr::NonNull;
@@ -187,11 +188,57 @@ pub struct AnalysisInfo {
     pub active_info: Vec<ActiveAnalysisInfo>,
 }
 
-#[derive(Debug, Clone)]
-pub struct AnalysisProgress {
-    pub state: AnalysisState,
-    pub count: usize,
-    pub total: usize,
+#[derive(Debug, Clone, Ord, PartialOrd, Eq, PartialEq)]
+pub enum AnalysisProgress {
+    Initial,
+    Hold,
+    Idle,
+    Discovery,
+    Disassembling(usize, usize),
+    Analyzing(usize, usize),
+    ExtendedAnalysis,
+}
+
+impl Display for AnalysisProgress {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self {
+            AnalysisProgress::Initial => {
+                write!(f, "Initial")
+            }
+            AnalysisProgress::Hold => {
+                write!(f, "Hold")
+            }
+            AnalysisProgress::Idle => {
+                write!(f, "Idle")
+            }
+            AnalysisProgress::Discovery => {
+                write!(f, "Discovery")
+            }
+            AnalysisProgress::Disassembling(count, total) => {
+                write!(f, "Disassembling ({count}/{total})")
+            }
+            AnalysisProgress::Analyzing(count, total) => {
+                write!(f, "Analyzing ({count}/{total})")
+            }
+            AnalysisProgress::ExtendedAnalysis => {
+                write!(f, "Extended Analysis")
+            }
+        }
+    }
+}
+
+impl From<BNAnalysisProgress> for AnalysisProgress {
+    fn from(value: BNAnalysisProgress) -> Self {
+        match value.state {
+            BNAnalysisState::InitialState => Self::Initial,
+            BNAnalysisState::HoldState => Self::Hold,
+            BNAnalysisState::IdleState => Self::Idle,
+            BNAnalysisState::DiscoveryState => Self::Discovery,
+            BNAnalysisState::DisassembleState => Self::Disassembling(value.count, value.total),
+            BNAnalysisState::AnalyzeState => Self::Analyzing(value.count, value.total),
+            BNAnalysisState::ExtendedAnalyzeState => Self::ExtendedAnalysis,
+        }
+    }
 }
 
 pub trait BinaryViewExt: BinaryViewBase {
@@ -604,12 +651,8 @@ pub trait BinaryViewExt: BinaryViewBase {
     }
 
     fn analysis_progress(&self) -> AnalysisProgress {
-        let progress = unsafe { BNGetAnalysisProgress(self.as_ref().handle) };
-        AnalysisProgress {
-            state: progress.state,
-            count: progress.count,
-            total: progress.total,
-        }
+        let progress_raw = unsafe { BNGetAnalysisProgress(self.as_ref().handle) };
+        AnalysisProgress::from(progress_raw)
     }
 
     fn default_arch(&self) -> Option<CoreArchitecture> {
diff --git a/rust/tests/binary_view.rs b/rust/tests/binary_view.rs
index c2a98da7..7a0a3299 100644
--- a/rust/tests/binary_view.rs
+++ b/rust/tests/binary_view.rs
@@ -1,5 +1,7 @@
 use binaryninja::binary_view::search::SearchQuery;
-use binaryninja::binary_view::{AnalysisState, BinaryViewBase, BinaryViewExt, StringType};
+use binaryninja::binary_view::{
+    AnalysisProgress, AnalysisState, BinaryViewBase, BinaryViewExt, StringType,
+};
 use binaryninja::data_buffer::DataBuffer;
 use binaryninja::function::{Function, FunctionViewType};
 use binaryninja::headless::Session;
@@ -16,7 +18,7 @@ fn test_binary_loading() {
     let out_dir = env!("OUT_DIR").parse::<PathBuf>().unwrap();
     let view = binaryninja::load(out_dir.join("atox.obj")).expect("Failed to create view");
     assert!(view.has_initial_analysis(), "No initial analysis");
-    assert_eq!(view.analysis_progress().state, AnalysisState::IdleState);
+    assert_eq!(view.analysis_progress(), AnalysisProgress::Idle);
     assert_eq!(view.file().is_analysis_changed(), false);
     assert_eq!(view.file().is_database_backed(), false);
 }
@@ -188,7 +190,7 @@ fn test_deterministic_functions() {
     for file_name in TARGET_FILES {
         let path = out_dir.join(file_name);
         let view = session.load(&path).expect("Failed to load view");
-        assert_eq!(view.analysis_progress().state, AnalysisState::IdleState);
+        assert_eq!(view.analysis_progress(), AnalysisProgress::Idle);
         let functions: BTreeMap<u64, FunctionSnapshot> = view
             .functions()
             .iter()