4 files changed, 354 insertions, 0 deletions
diff --git a/plugins/warp/src/cache/container.rs b/plugins/warp/src/cache/container.rs
new file mode 100644
index 00000000..59b31bad
--- /dev/null
+++ b/plugins/warp/src/cache/container.rs
@@ -0,0 +1,28 @@
+use crate::container::Container;
+use dashmap::DashMap;
+use std::ops::Deref;
+use std::sync::{Arc, OnceLock, RwLock};
+
+pub static CONTAINER_CACHE: OnceLock<DashMap<String, Arc<RwLock<Box<dyn Container>>>>> =
+    OnceLock::new();
+
+pub fn for_cached_containers(f: impl Fn(&dyn Container)) {
+    let containers_cache = CONTAINER_CACHE.get_or_init(Default::default);
+    for container in containers_cache.iter() {
+        if let Ok(guarded_container) = container.read() {
+            f(guarded_container.deref().as_ref());
+        }
+    }
+}
+
+// TODO: The static lifetime here is a little wierd... (we need it to Box)
+pub fn add_cached_container(container: impl Container + 'static) {
+    let containers_cache = CONTAINER_CACHE.get_or_init(Default::default);
+    let container_name = container.to_string();
+    containers_cache.insert(container_name, Arc::new(RwLock::new(Box::new(container))));
+}
+
+pub fn cached_containers() -> Vec<Arc<RwLock<Box<dyn Container>>>> {
+    let containers_cache = CONTAINER_CACHE.get_or_init(Default::default);
+    containers_cache.iter().map(|c| c.clone()).collect()
+}
diff --git a/plugins/warp/src/cache/function.rs b/plugins/warp/src/cache/function.rs
new file mode 100644
index 00000000..2fef05fd
--- /dev/null
+++ b/plugins/warp/src/cache/function.rs
@@ -0,0 +1,28 @@
+use binaryninja::function::{Function as BNFunction, FunctionUpdateType};
+use warp::signature::function::Function;
+
+/// Inserts a function match into the cache.
+///
+/// IMPORTANT: This will mark the function as needing updates, if you intend to fill in functions with
+/// no match (i.e. `None`), then you must change this function to prevent marking that as needing updates.
+/// However, it's perfectly valid to remove a match and need to update the function still, so be careful.
+pub fn insert_cached_function_match(function: &BNFunction, matched_function: Option<Function>) {
+    // NOTE: If we expect to run match_function multiple times on a function, we should move this elsewhere.
+    // Mark the function as needing updates so that reanalysis occurs on the function, and we apply the match.
+    function.mark_updates_required(FunctionUpdateType::FullAutoFunctionUpdate);
+    match matched_function {
+        Some(matched_function) => {
+            function.store_metadata("warp_matched_function", &matched_function.to_bytes(), false);
+        }
+        None => {
+            function.remove_metadata("warp_matched_function");
+        }
+    }
+}
+
+// TODO: This does allocations, and for every reanalysis.
+pub fn try_cached_function_match(function: &BNFunction) -> Option<Function> {
+    let metadata = function.query_metadata("warp_matched_function")?;
+    let raw_metadata = metadata.get_raw()?;
+    Function::from_bytes(&raw_metadata)
+}
diff --git a/plugins/warp/src/cache/guid.rs b/plugins/warp/src/cache/guid.rs
new file mode 100644
index 00000000..f1788a33
--- /dev/null
+++ b/plugins/warp/src/cache/guid.rs
@@ -0,0 +1,193 @@
+use crate::cache::FunctionID;
+use crate::convert::from_bn_symbol;
+use crate::function_guid;
+use binaryninja::binary_view::BinaryViewExt;
+use binaryninja::function::Function as BNFunction;
+use binaryninja::low_level_il::function::{FunctionMutability, LowLevelILFunction, NonSSA};
+use binaryninja::symbol::Symbol as BNSymbol;
+use std::collections::HashSet;
+use uuid::Uuid;
+use warp::signature::constraint::Constraint;
+use warp::signature::function::FunctionGUID;
+
+pub fn cached_function_guid<M: FunctionMutability>(
+    function: &BNFunction,
+    lifted_il: &LowLevelILFunction<M, NonSSA>,
+) -> FunctionGUID {
+    let cached_guid = try_cached_function_guid(function);
+    if let Some(cached_guid) = cached_guid {
+        return cached_guid;
+    }
+
+    let function_guid = function_guid(function, lifted_il);
+    function.store_metadata(
+        "warp_function_guid",
+        &function_guid.as_bytes().to_vec(),
+        false,
+    );
+    function_guid
+}
+
+pub fn try_cached_function_guid(function: &BNFunction) -> Option<FunctionGUID> {
+    let metadata = function.query_metadata("warp_function_guid")?;
+    let raw_metadata = metadata.get_raw()?;
+    let uuid = Uuid::from_slice(raw_metadata.as_slice()).ok()?;
+    Some(FunctionGUID::from(uuid))
+}
+
+pub fn cached_constraints<F>(function: &BNFunction, filter: F) -> HashSet<Constraint>
+where
+    F: Fn(&BNFunction) -> bool,
+{
+    // TODO: Implied constraints, symbol name, image offset
+    let cs_constraints = cached_call_site_constraints(function);
+    let adj_constraints = cached_adjacency_constraints(function, filter);
+    cs_constraints.union(&adj_constraints).cloned().collect()
+}
+
+pub fn cached_call_site_constraints(function: &BNFunction) -> HashSet<Constraint> {
+    let cache = ConstraintBuilder;
+    cache.call_site_constraints(function)
+}
+
+pub fn cached_adjacency_constraints<F>(function: &BNFunction, filter: F) -> HashSet<Constraint>
+where
+    F: Fn(&BNFunction) -> bool,
+{
+    let cache = ConstraintBuilder;
+    cache.adjacency_constraints(function, filter)
+}
+
+#[derive(Clone, Debug, Default)]
+pub struct ConstraintBuilder;
+
+impl ConstraintBuilder {
+    pub fn call_site_constraints(&self, function: &BNFunction) -> HashSet<Constraint> {
+        let view = function.view();
+        let func_id = FunctionID::from(function);
+        let func_start = function.start();
+        let func_platform = function.platform();
+        let mut constraints = HashSet::new();
+        for call_site in &function.call_sites() {
+            for cs_ref_addr in view.code_refs_from_addr(call_site.address, Some(function)) {
+                match view.function_at(&func_platform, cs_ref_addr) {
+                    Some(cs_ref_func) => {
+                        // Call site is a function, constrain on it.
+                        let cs_ref_func_id = FunctionID::from(cs_ref_func.as_ref());
+                        if cs_ref_func_id != func_id {
+                            let call_site_offset: i64 =
+                                call_site.address.wrapping_sub(func_start) as i64;
+                            // TODO: If the function is thunk we should also insert the called function.
+                            constraints.extend(
+                                self.related_function_constraint(&cs_ref_func, call_site_offset),
+                            );
+                        }
+                    }
+                    None => {
+                        // We could be dealing with an extern symbol, get the symbol as a constraint.
+                        let call_site_offset: i64 =
+                            call_site.address.wrapping_sub(func_start) as i64;
+                        if let Some(call_site_sym) = view.symbol_by_address(cs_ref_addr) {
+                            constraints.insert(
+                                self.related_symbol_constraint(&call_site_sym, call_site_offset),
+                            );
+                        }
+                    }
+                }
+            }
+        }
+        constraints
+    }
+
+    pub fn adjacency_constraints<F>(&self, function: &BNFunction, filter: F) -> HashSet<Constraint>
+    where
+        F: Fn(&BNFunction) -> bool,
+    {
+        let view = function.view();
+        let func_id = FunctionID::from(function);
+        let func_start = function.start();
+        let mut constraints = HashSet::new();
+
+        let mut func_addr_constraint = |func_start_addr| {
+            // NOTE: We could potentially have dozens of functions all at the same start address.
+            for curr_func in &view.functions_at(func_start_addr) {
+                let curr_func_id = FunctionID::from(curr_func.as_ref());
+                if curr_func_id != func_id && filter(curr_func.as_ref()) {
+                    // NOTE: For this to work the GUID has to have already been cached. If not it will just be the symbol.
+                    // Function adjacent to another function, constrain on the pattern.
+                    let curr_addr_offset = (func_start_addr as i64) - func_start as i64;
+                    constraints
+                        .extend(self.related_function_constraint(&curr_func, curr_addr_offset));
+                }
+            }
+        };
+
+        let mut before_func_start = func_start;
+        for _ in 0..2 {
+            before_func_start = view.function_start_before(before_func_start);
+            func_addr_constraint(before_func_start);
+        }
+
+        let mut after_func_start = func_start;
+        for _ in 0..2 {
+            after_func_start = view.function_start_after(after_func_start);
+            func_addr_constraint(after_func_start);
+        }
+
+        constraints
+    }
+
+    /// Construct a function constraint, must pass the offset at which it is located.
+    pub fn related_function_constraint(
+        &self,
+        function: &BNFunction,
+        offset: i64,
+    ) -> Vec<Constraint> {
+        let mut constraints = vec![];
+        if let Some(guid) = try_cached_function_guid(function) {
+            let guid_constraint = Constraint::from_function(&guid, Some(offset));
+            constraints.push(guid_constraint);
+        }
+        let symbol_constraint = self.related_symbol_constraint(&function.symbol(), offset);
+        constraints.push(symbol_constraint);
+        constraints
+    }
+
+    /// Construct a symbol constraint, must pass the offset at which it is located.
+    pub fn related_symbol_constraint(&self, symbol: &BNSymbol, offset: i64) -> Constraint {
+        let mut symbol = from_bn_symbol(symbol);
+        symbol.name = clean_symbol_name(&symbol.name);
+        Constraint::from_symbol(&symbol, Some(offset))
+    }
+}
+
+/// Cleans various internal symbol prefixes and suffixes for consistency.
+///
+/// This is very important for getting matching symbol constraints.
+///
+/// Examples:
+/// - "__imp__RemoveDirectoryW@4" -> "RemoveDirectoryW"
+/// - "__free_base" -> "free_base"
+/// - "__impl__free_base" -> "free_base"
+/// - "j___free_base" -> "free_base"
+/// - "j_free_base" -> "free_base"
+/// - "_free_base" -> "free_base"
+pub fn clean_symbol_name(symbol_name: &str) -> String {
+    // Handle MSVC-style imported symbols
+    let without_imp = symbol_name.strip_prefix("__imp__").unwrap_or(symbol_name);
+
+    // Handle jump thunk prefix
+    let without_jump = without_imp.strip_prefix("j_").unwrap_or(without_imp);
+
+    // Strip all remaining leading underscores
+    let mut result = without_jump;
+    while result.starts_with('_') {
+        result = &result[1..];
+    }
+
+    // Remove stdcall decoration (@N suffix)
+    match result.find('@') {
+        Some(pos) => result[..pos].to_string(),
+        None => result.to_string(),
+    }
+}
diff --git a/plugins/warp/src/cache/type_reference.rs b/plugins/warp/src/cache/type_reference.rs
new file mode 100644
index 00000000..0b07c139
--- /dev/null
+++ b/plugins/warp/src/cache/type_reference.rs
@@ -0,0 +1,105 @@
+use crate::cache::ViewID;
+use crate::convert::from_bn_type_internal;
+use binaryninja::binary_view::BinaryView;
+use binaryninja::confidence::MAX_CONFIDENCE;
+use binaryninja::rc::Guard;
+use binaryninja::rc::Ref as BNRef;
+use binaryninja::types::NamedTypeReference as BNNamedTypeReference;
+use dashmap::mapref::one::Ref;
+use dashmap::DashMap;
+use std::collections::HashSet;
+use std::hash::{DefaultHasher, Hasher};
+use std::sync::OnceLock;
+use warp::r#type::ComputedType;
+
+pub static TYPE_REF_CACHE: OnceLock<DashMap<ViewID, TypeRefCache>> = OnceLock::new();
+
+pub fn clear_type_ref_cache(view: &BinaryView) {
+    let view_id = ViewID::from(view);
+    if let Some(cache) = TYPE_REF_CACHE.get() {
+        cache.remove(&view_id);
+    }
+}
+
+pub fn cached_type_reference(
+    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(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 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,
+                        MAX_CONFIDENCE,
+                    ));
+                    self.cache
+                        .entry(ntr_id)
+                        .insert(Some(computed_ty))
+                        .to_owned()
+                }
+                None => self.cache.entry(ntr_id).insert(None).to_owned(),
+            },
+        }
+    }
+}
+
+/// 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())
+    }
+}