[ty] Avoid overcounting shared memory usage (#19773)

## Summary

Use a global tracker to avoid double counting `Arc` instances.

crates/ty_python_semantic/src/semantic_index.rs

@@ -28,7 +28,6 @@ use crate::semantic_index::scope::{
 use crate::semantic_index::symbol::ScopedSymbolId;
 use crate::semantic_index::use_def::{EnclosingSnapshotKey, ScopedEnclosingSnapshotId, UseDefMap};
 use crate::semantic_model::HasTrackedScope;
-use crate::util::get_size::untracked_arc_size;
 
 pub mod ast_ids;
 mod builder;
@@ -52,7 +51,7 @@ pub(crate) use self::use_def::{
 /// Returns the semantic index for `file`.
 ///
 /// Prefer using [`symbol_table`] when working with symbols from a single scope.
-#[salsa::tracked(returns(ref), no_eq, heap_size=get_size2::heap_size)]
+#[salsa::tracked(returns(ref), no_eq, heap_size=ruff_memory_usage::heap_size)]
 pub(crate) fn semantic_index(db: &dyn Db, file: File) -> SemanticIndex<'_> {
     let _span = tracing::trace_span!("semantic_index", ?file).entered();
 
@@ -66,7 +65,7 @@ pub(crate) fn semantic_index(db: &dyn Db, file: File) -> SemanticIndex<'_> {
 /// Using [`place_table`] over [`semantic_index`] has the advantage that
 /// Salsa can avoid invalidating dependent queries if this scope's place table
 /// is unchanged.
-#[salsa::tracked(returns(deref), heap_size=get_size2::heap_size)]
+#[salsa::tracked(returns(deref), heap_size=ruff_memory_usage::heap_size)]
 pub(crate) fn place_table<'db>(db: &'db dyn Db, scope: ScopeId<'db>) -> Arc<PlaceTable> {
     let file = scope.file(db);
     let _span = tracing::trace_span!("place_table", scope=?scope.as_id(), ?file).entered();
@@ -86,7 +85,7 @@ pub(crate) fn place_table<'db>(db: &'db dyn Db, scope: ScopeId<'db>) -> Arc<Plac
 ///
 ///   - We cannot resolve relative imports (which aren't allowed in `import` statements) without
 ///     knowing the name of the current module, and whether it's a package.
-#[salsa::tracked(returns(deref), heap_size=get_size2::heap_size)]
+#[salsa::tracked(returns(deref), heap_size=ruff_memory_usage::heap_size)]
 pub(crate) fn imported_modules<'db>(db: &'db dyn Db, file: File) -> Arc<FxHashSet<ModuleName>> {
     semantic_index(db, file).imported_modules.clone()
 }
@@ -96,8 +95,8 @@ pub(crate) fn imported_modules<'db>(db: &'db dyn Db, file: File) -> Arc<FxHashSe
 /// Using [`use_def_map`] over [`semantic_index`] has the advantage that
 /// Salsa can avoid invalidating dependent queries if this scope's use-def map
 /// is unchanged.
-#[salsa::tracked(returns(deref), heap_size=get_size2::heap_size)]
-pub(crate) fn use_def_map<'db>(db: &'db dyn Db, scope: ScopeId<'db>) -> ArcUseDefMap<'db> {
+#[salsa::tracked(returns(deref), heap_size=ruff_memory_usage::heap_size)]
+pub(crate) fn use_def_map<'db>(db: &'db dyn Db, scope: ScopeId<'db>) -> Arc<UseDefMap<'db>> {
     let file = scope.file(db);
     let _span = tracing::trace_span!("use_def_map", scope=?scope.as_id(), ?file).entered();
     let index = semantic_index(db, file);
@@ -123,7 +122,7 @@ pub(crate) fn attribute_assignments<'db, 's>(
         let member = place_table.member_id_by_instance_attribute_name(name)?;
         let use_def = &index.use_def_maps[function_scope_id];
         Some((
-            use_def.inner.all_reachable_member_bindings(member),
+            use_def.all_reachable_member_bindings(member),
             function_scope_id,
         ))
     })
@@ -147,7 +146,7 @@ pub(crate) fn attribute_declarations<'db, 's>(
         let member = place_table.member_id_by_instance_attribute_name(name)?;
         let use_def = &index.use_def_maps[function_scope_id];
         Some((
-            use_def.inner.all_reachable_member_declarations(member),
+            use_def.all_reachable_member_declarations(member),
             function_scope_id,
         ))
     })
@@ -184,7 +183,7 @@ pub(crate) fn attribute_scopes<'db, 's>(
 }
 
 /// Returns the module global scope of `file`.
-#[salsa::tracked(heap_size=get_size2::heap_size)]
+#[salsa::tracked(heap_size=ruff_memory_usage::heap_size)]
 pub(crate) fn global_scope(db: &dyn Db, file: File) -> ScopeId<'_> {
     let _span = tracing::trace_span!("global_scope", ?file).entered();
 
@@ -223,7 +222,7 @@ pub(crate) struct SemanticIndex<'db> {
     scope_ids_by_scope: IndexVec<FileScopeId, ScopeId<'db>>,
 
     /// Use-def map for each scope in this file.
-    use_def_maps: IndexVec<FileScopeId, ArcUseDefMap<'db>>,
+    use_def_maps: IndexVec<FileScopeId, Arc<UseDefMap<'db>>>,
 
     /// Lookup table to map between node ids and ast nodes.
     ///
@@ -262,7 +261,7 @@ impl<'db> SemanticIndex<'db> {
     /// Use the Salsa cached [`use_def_map()`] query if you only need the
     /// use-def map for a single scope.
     #[track_caller]
-    pub(super) fn use_def_map(&self, scope_id: FileScopeId) -> ArcUseDefMap<'_> {
+    pub(super) fn use_def_map(&self, scope_id: FileScopeId) -> Arc<UseDefMap<'_>> {
         self.use_def_maps[scope_id].clone()
     }
 
@@ -510,7 +509,6 @@ impl<'db> SemanticIndex<'db> {
                         };
                         if let Some(id) = self.enclosing_snapshots.get(&key) {
                             return self.use_def_maps[enclosing_scope]
-                                .inner
                                 .enclosing_snapshot(*id, key.nested_laziness);
                         }
                     }
@@ -530,9 +528,7 @@ impl<'db> SemanticIndex<'db> {
         let Some(id) = self.enclosing_snapshots.get(&key) else {
             return EnclosingSnapshotResult::NotFound;
         };
-        self.use_def_maps[enclosing_scope]
-            .inner
-            .enclosing_snapshot(*id, key.nested_laziness)
+        self.use_def_maps[enclosing_scope].enclosing_snapshot(*id, key.nested_laziness)
     }
 
     pub(crate) fn semantic_syntax_errors(&self) -> &[SemanticSyntaxError] {
@@ -540,28 +536,6 @@ impl<'db> SemanticIndex<'db> {
     }
 }
 
-#[derive(Debug, PartialEq, Eq, Clone, salsa::Update, get_size2::GetSize)]
-pub(crate) struct ArcUseDefMap<'db> {
-    #[get_size(size_fn = untracked_arc_size)]
-    inner: Arc<UseDefMap<'db>>,
-}
-
-impl<'db> std::ops::Deref for ArcUseDefMap<'db> {
-    type Target = UseDefMap<'db>;
-
-    fn deref(&self) -> &Self::Target {
-        &self.inner
-    }
-}
-
-impl<'db> ArcUseDefMap<'db> {
-    pub(crate) fn new(inner: UseDefMap<'db>) -> Self {
-        Self {
-            inner: Arc::new(inner),
-        }
-    }
-}
-
 pub(crate) struct AncestorsIter<'a> {
     scopes: &'a IndexSlice<FileScopeId, Scope>,
     next_id: Option<FileScopeId>,