Add a hook to allow BinaryView subclasses to run code after snapshot data is applied

Snapshot data is applied when loading from a database, rebasing the
view, etc.

1 files changed, 14 insertions, 0 deletions

diff --git a/rust/src/custom_binary_view.rs b/rust/src/custom_binary_view.rs
index 4d21ce89..8c880fc1 100644
--- a/rust/src/custom_binary_view.rs
+++ b/rust/src/custom_binary_view.rs
@@ -460,6 +460,7 @@ pub unsafe trait CustomBinaryView: 'static + BinaryViewBase + Sync + Sized {
 
     fn new(handle: &BinaryView, args: &Self::Args) -> Result<Self>;
     fn init(&mut self, args: Self::Args) -> Result<()>;
+    fn on_after_snapshot_data_applied(&mut self) {}
 }
 
 /// Represents a partially initialized custom `BinaryView` that should be returned to the core
@@ -600,6 +601,18 @@ impl<'a, T: CustomBinaryViewType> CustomViewBuilder<'a, T> {
             })
         }
 
+        extern "C" fn cb_on_after_snapshot_data_applied<V>(ctxt: *mut c_void)
+        where
+            V: CustomBinaryView,
+        {
+            ffi_wrap!("BinaryViewBase::onAfterSnapshotDataApplied", unsafe {
+                let context = &mut *(ctxt as *mut CustomViewContext<V>);
+                if let CustomViewContextState::Initialized { view } = &mut context.state {
+                    view.on_after_snapshot_data_applied();
+                }
+            })
+        }
+
         extern "C" fn cb_free_object<V>(ctxt: *mut c_void)
         where
             V: CustomBinaryView,
@@ -890,6 +903,7 @@ impl<'a, T: CustomBinaryViewType> CustomViewBuilder<'a, T> {
             isRelocatable: Some(cb_relocatable::<V>),
             getAddressSize: Some(cb_address_size::<V>),
             save: Some(cb_save::<V>),
+            onAfterSnapshotDataApplied: Some(cb_on_after_snapshot_data_applied::<V>),
         };
 
         let view_name = view_name.to_cstr();