feat: cache code objects in global registry

Author: tzanko-matev

.agents/tasks/2025/08/21-0939-codetype-interface

@@ -38,3 +38,5 @@ Also please add usage examples to the design documentation
 Implement the CodeObjectWrapper as designed. Update the Tracer trait as well as the callback_xxx functions accordingly. Write a comprehensive unit tests for CodeObjectWrapper.
 --- FOLLOW UP TASK ---
 There is an issue in the current implementation. We don't use caching effectively, since we create a new CodeObjectWrapper at each callback_xxx call. We need a global cache, probably keyed by the code object id. Propose design changes and update the design documents. Don't implement the changes themselves before I approve them.
+--- FOLLOW UP TASK ---
+Implement the global code object registry.

codetracer-python-recorder/Cargo.lock

@@ -19,6 +19,7 @@ pyo3 = { version = "0.25.1" }
 runtime_tracing = "0.14.0"
 bitflags = "2.4"
 once_cell = "1.19"
+dashmap = "5.5"
 
 [dev-dependencies]
 pyo3 = { version = "0.25.1", features = ["auto-initialize"] }

codetracer-python-recorder/Cargo.toml

@@ -1,6 +1,8 @@
 use once_cell::sync::OnceCell;
 use pyo3::prelude::*;
 use pyo3::types::PyCode;
+use dashmap::DashMap;
+use std::sync::Arc;
 
 /// A wrapper around Python `code` objects providing cached access to
 /// common attributes and line information.
@@ -128,3 +130,35 @@ impl CodeObjectWrapper {
         }
     }
 }
+
+/// Global registry caching `CodeObjectWrapper` instances by code object id.
+#[derive(Default)]
+pub struct CodeObjectRegistry {
+    map: DashMap<usize, Arc<CodeObjectWrapper>>,
+}
+
+impl CodeObjectRegistry {
+    /// Retrieve the wrapper for `code`, inserting a new one if needed.
+    pub fn get_or_insert(
+        &self,
+        py: Python<'_>,
+        code: &Bound<'_, PyCode>,
+    ) -> Arc<CodeObjectWrapper> {
+        let id = code.as_ptr() as usize;
+        self.map
+            .entry(id)
+            .or_insert_with(|| Arc::new(CodeObjectWrapper::new(py, code)))
+            .clone() //AI? Why do we need to clone here?
+    }
+
+    /// Remove the wrapper for a given code id, if present.
+    pub fn remove(&self, id: usize) {
+        self.map.remove(&id);
+    }
+
+    /// Clear all cached wrappers.
+    pub fn clear(&self) {
+        self.map.clear();
+    }
+}
+

codetracer-python-recorder/src/code_object.rs

@@ -5,7 +5,7 @@ use pyo3::prelude::*;
 
 pub mod code_object;
 pub mod tracer;
-pub use crate::code_object::CodeObjectWrapper;
+pub use crate::code_object::{CodeObjectRegistry, CodeObjectWrapper};
 pub use crate::tracer::{install_tracer, uninstall_tracer, EventSet, Tracer};
 
 /// Global flag tracking whether tracing is active.