WARP: Cache type references in the view caches

This allows us to free the type references instead of keeping them around forever

3 files changed, 108 insertions, 32 deletions

diff --git a/plugins/warp/src/bin/sigem.rs b/plugins/warp/src/bin/sigem.rs
index 2da14639..d7f1f50b 100644
--- a/plugins/warp/src/bin/sigem.rs
+++ b/plugins/warp/src/bin/sigem.rs
@@ -103,10 +103,17 @@ fn data_from_view(view: &BinaryView) -> Data {
             Some(warp_ninja::cache::cached_function(&f, &llil))
         })
         .collect::<Vec<_>>();
-    data.types.extend(view.types().iter().map(|ty| {
-        let ref_ty = ty.type_object().to_owned();
-        ComputedType::new(from_bn_type(view, &ref_ty, u8::MAX))
-    }));
+
+    if let Some(ref_ty_cache) = cached_type_references(view) {
+        let referenced_types = ref_ty_cache
+            .cache
+            .iter()
+            .filter_map(|t| t.to_owned())
+            .collect::<Vec<_>>();
+
+        data.types.extend(referenced_types);
+    }
+    
     data
 }
 
diff --git a/plugins/warp/src/cache.rs b/plugins/warp/src/cache.rs
index b8a63cdd..e3f1259c 100644
--- a/plugins/warp/src/cache.rs
+++ b/plugins/warp/src/cache.rs
@@ -1,18 +1,20 @@
-use crate::convert::from_bn_symbol;
+use crate::convert::{from_bn_symbol, from_bn_type_internal};
 use crate::{build_function, function_guid};
 use binaryninja::architecture::Architecture;
-use binaryninja::binaryview::{BinaryView, BinaryViewBase, BinaryViewExt};
+use binaryninja::binaryview::{BinaryView, BinaryViewExt};
 use binaryninja::function::Function as BNFunction;
 use binaryninja::llil::{FunctionMutability, NonSSA, NonSSAVariant};
 use binaryninja::rc::Guard;
 use binaryninja::rc::Ref as BNRef;
 use binaryninja::symbol::Symbol as BNSymbol;
+use binaryninja::types::NamedTypeReference as BNNamedTypeReference;
 use binaryninja::{llil, ObjectDestructor};
 use dashmap::mapref::one::Ref;
 use dashmap::DashMap;
 use std::collections::HashSet;
 use std::hash::{DefaultHasher, Hash, Hasher};
 use std::sync::OnceLock;
+use warp::r#type::ComputedType;
 use warp::signature::function::constraints::FunctionConstraint;
 use warp::signature::function::{Function, FunctionGUID};
 
@@ -20,6 +22,7 @@ pub static MATCHED_FUNCTION_CACHE: OnceLock<DashMap<ViewID, MatchedFunctionCache
     OnceLock::new();
 pub static FUNCTION_CACHE: OnceLock<DashMap<ViewID, FunctionCache>> = OnceLock::new();
 pub static GUID_CACHE: OnceLock<DashMap<ViewID, GUIDCache>> = OnceLock::new();
+pub static TYPE_REF_CACHE: OnceLock<DashMap<ViewID, TypeRefCache>> = OnceLock::new();
 
 pub fn register_cache_destructor() {
     pub static mut CACHE_DESTRUCTOR: CacheDestructor = CacheDestructor;
@@ -130,6 +133,30 @@ pub fn try_cached_function_guid(function: &BNFunction) -> Option<FunctionGUID> {
     guid_cache.get(&view_id)?.try_function_guid(function)
 }
 
+pub fn cached_type_reference<'a>(
+    view: &BinaryView,
+    visited_refs: &mut HashSet<TypeRefID>,
+    type_ref: &BNNamedTypeReference,
+) -> Option<ComputedType> {
+    let view_id = ViewID::from(view);
+    let type_ref_cache = TYPE_REF_CACHE.get_or_init(Default::default);
+    match type_ref_cache.get(&view_id) {
+        Some(cache) => cache.cached_type_reference(view, visited_refs, type_ref),
+        None => {
+            let cache = TypeRefCache::default();
+            let ntr = cache.cached_type_reference(view, visited_refs, type_ref);
+            type_ref_cache.insert(view_id, cache);
+            ntr
+        }
+    }
+}
+
+pub fn cached_type_references<'a>(view: &BinaryView) -> Option<Ref<ViewID, TypeRefCache>> {
+    let view_id = ViewID::from(view);
+    let type_ref_cache = TYPE_REF_CACHE.get_or_init(Default::default);
+    type_ref_cache.get(&view_id)
+}
+
 #[derive(Clone, Debug, Default)]
 pub struct MatchedFunctionCache {
     pub cache: DashMap<FunctionID, Option<Function>>,
@@ -305,6 +332,36 @@ impl GUIDCache {
     }
 }
 
+#[derive(Clone, Debug, Default)]
+pub struct TypeRefCache {
+    pub cache: DashMap<TypeRefID, Option<ComputedType>>,
+}
+
+impl TypeRefCache {
+    /// NOTE: No self-referential type must be used on this function.
+    pub fn cached_type_reference(
+        &self,
+        view: &BinaryView,
+        visited_refs: &mut HashSet<TypeRefID>,
+        type_ref: &BNNamedTypeReference,
+    ) -> Option<ComputedType> {
+        let ntr_id = TypeRefID::from(type_ref);
+        match self.cache.get(&ntr_id) {
+            Some(cache) => cache.to_owned(),
+            None => match type_ref.target(view) {
+                Some(raw_ty) => {
+                    let computed_ty = ComputedType::new(from_bn_type_internal(view, visited_refs, &raw_ty, 255));
+                    self.cache
+                        .entry(ntr_id)
+                        .insert(Some(computed_ty))
+                        .to_owned()
+                }
+                None => self.cache.entry(ntr_id).insert(None).to_owned(),
+            },
+        }
+    }
+}
+
 /// A unique view ID, used for caching.
 #[derive(Copy, Clone, Debug, Hash, PartialEq, Eq, PartialOrd, Ord)]
 pub struct ViewID(u64);
@@ -358,6 +415,30 @@ impl From<Guard<'_, BNFunction>> for FunctionID {
     }
 }
 
+/// A unique named type reference ID, used for caching.
+#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq, PartialOrd, Ord)]
+pub struct TypeRefID(u64);
+
+impl From<&BNNamedTypeReference> for TypeRefID {
+    fn from(value: &BNNamedTypeReference) -> Self {
+        let mut hasher = DefaultHasher::new();
+        hasher.write(value.id().as_bytes());
+        Self(hasher.finish())
+    }
+}
+
+impl From<BNRef<BNNamedTypeReference>> for TypeRefID {
+    fn from(value: BNRef<BNNamedTypeReference>) -> Self {
+        Self::from(value.as_ref())
+    }
+}
+
+impl From<Guard<'_, BNNamedTypeReference>> for TypeRefID {
+    fn from(value: Guard<'_, BNNamedTypeReference>) -> Self {
+        Self::from(value.as_ref())
+    }
+}
+
 pub struct CacheDestructor;
 
 impl ObjectDestructor for CacheDestructor {
@@ -373,6 +454,9 @@ impl ObjectDestructor for CacheDestructor {
         if let Some(cache) = GUID_CACHE.get() {
             cache.remove(&view_id);
         }
+        if let Some(cache) = TYPE_REF_CACHE.get() {
+            cache.remove(&view_id);
+        }
         log::debug!("Removed WARP caches for {:?}", view);
     }
 }
diff --git a/plugins/warp/src/convert.rs b/plugins/warp/src/convert.rs
index c425bdfe..6f2bf7e3 100644
--- a/plugins/warp/src/convert.rs
+++ b/plugins/warp/src/convert.rs
@@ -17,6 +17,7 @@ use binaryninja::types::{
     StructureType as BNStructureType, Type as BNType, TypeClass as BNTypeClass,
 };
 
+use crate::cache::{cached_type_reference, TypeRefID};
 use warp::r#type::class::array::ArrayModifiers;
 use warp::r#type::class::function::{Location, RegisterLocation};
 use warp::r#type::class::pointer::PointerAddressing;
@@ -26,7 +27,6 @@ use warp::r#type::class::{
     EnumerationMember, FloatClass, FunctionClass, FunctionMember, IntegerClass, PointerClass,
     ReferrerClass, StructureClass, StructureMember, TypeClass,
 };
-use warp::r#type::guid::TypeGUID;
 use warp::r#type::Type;
 use warp::symbol::class::SymbolClass;
 use warp::symbol::{Symbol, SymbolModifiers};
@@ -121,9 +121,9 @@ pub fn from_bn_type(view: &BinaryView, raw_ty: &BNType, confidence: u8) -> Type
     from_bn_type_internal(view, &mut HashSet::new(), raw_ty, confidence)
 }
 
-fn from_bn_type_internal(
+pub fn from_bn_type_internal(
     view: &BinaryView,
-    visited_refs: &mut HashSet<String>,
+    visited_refs: &mut HashSet<TypeRefID>,
     raw_ty: &BNType,
     confidence: u8,
 ) -> Type {
@@ -318,32 +318,17 @@ fn from_bn_type_internal(
         }
         BNTypeClass::NamedTypeReferenceClass => {
             let raw_ntr = raw_ty.get_named_type_reference().unwrap();
-            let ref_id_str = raw_ntr.id().to_string();
-            let raw_ntr_ty = raw_ntr.target(view);
-            if raw_ntr_ty.is_none() || !visited_refs.insert(ref_id_str.clone()) {
-                let ref_class = ReferrerClass::new(None, Some(raw_ntr.name().to_string()));
-                TypeClass::Referrer(ref_class)
-            } else {
-                use dashmap::DashMap;
-                use std::sync::Arc;
-                use std::sync::OnceLock;
-                static REF_CACHE: OnceLock<Arc<DashMap<String, TypeClass>>> = OnceLock::new();
-                let ref_cache = REF_CACHE.get_or_init(|| Arc::new(DashMap::new()));
-                // Check the cache first before proceeding
-                if let Some(cached_type) = ref_cache.get(&ref_id_str) {
-                    cached_type.value().to_owned()
-                } else {
-                    let ntr_ty =
-                        from_bn_type_internal(view, visited_refs, &raw_ntr_ty.unwrap(), confidence);
-                    visited_refs.remove(&ref_id_str);
+            let ref_id = TypeRefID::from(raw_ntr.as_ref());
+            let mut ref_class = ReferrerClass::new(None, Some(raw_ntr.name().to_string()));
+            if visited_refs.insert(ref_id) {
+                // This ntr is NOT self-referential, meaning we can deduce a type GUID.
+                if let Some(computed_ty) = cached_type_reference(view, visited_refs, &raw_ntr) {
                     // NOTE: The GUID here must always equal the same for any given type for this to work effectively.
-                    let ntr_guid = TypeGUID::from(&ntr_ty);
-                    let ref_class = ReferrerClass::new(Some(ntr_guid), ntr_ty.name);
-                    let ntr_ty_class = TypeClass::Referrer(ref_class);
-                    ref_cache.insert(ref_id_str, ntr_ty_class.clone());
-                    ntr_ty_class
+                    ref_class.guid = Some(computed_ty.guid);
                 }
+                visited_refs.remove(&ref_id);
             }
+            TypeClass::Referrer(ref_class)
         }
         BNTypeClass::WideCharTypeClass => {
             let char_class = CharacterClass {