From dc10ab81bd85b8af49ee2fe8ed4a6f767d37d6e0 Mon Sep 17 00:00:00 2001
From: Micha Reiser <micha@reiser.io>
Date: Tue, 22 Jul 2025 20:39:39 +0200
Subject: [PATCH] [ty] Use `ThinVec` for sub segments in `PlaceExpr` (#19470)

---
 Cargo.lock                                    |  1 +
 Cargo.toml                                    |  1 +
 crates/ty_python_semantic/Cargo.toml          |  1 +
 .../src/semantic_index/place.rs               | 29 ++++++++++++++-----
 4 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index b0d761cfee..cc31e68ba0 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4306,6 +4306,7 @@ dependencies = [
  "strum_macros",
  "tempfile",
  "test-case",
+ "thin-vec",
  "thiserror 2.0.12",
  "tracing",
  "ty_python_semantic",
diff --git a/Cargo.toml b/Cargo.toml
index da9383c64f..089b201fdc 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -166,6 +166,7 @@ strum_macros = { version = "0.27.0" }
 syn = { version = "2.0.55" }
 tempfile = { version = "3.9.0" }
 test-case = { version = "3.3.1" }
+thin-vec = { version = "0.2.14" }
 thiserror = { version = "2.0.0" }
 tikv-jemallocator = { version = "0.6.0" }
 toml = { version = "0.9.0" }
diff --git a/crates/ty_python_semantic/Cargo.toml b/crates/ty_python_semantic/Cargo.toml
index 370b0a4154..d4961a8a8a 100644
--- a/crates/ty_python_semantic/Cargo.toml
+++ b/crates/ty_python_semantic/Cargo.toml
@@ -36,6 +36,7 @@ indexmap = { workspace = true }
 itertools = { workspace = true }
 ordermap = { workspace = true }
 salsa = { workspace = true, features = ["compact_str"] }
+thin-vec = { workspace = true }
 thiserror = { workspace = true }
 tracing = { workspace = true }
 rustc-hash = { workspace = true }
diff --git a/crates/ty_python_semantic/src/semantic_index/place.rs b/crates/ty_python_semantic/src/semantic_index/place.rs
index 7b16d377d4..4cf81896f7 100644
--- a/crates/ty_python_semantic/src/semantic_index/place.rs
+++ b/crates/ty_python_semantic/src/semantic_index/place.rs
@@ -41,7 +41,16 @@ impl PlaceExprSubSegment {
 #[derive(Eq, PartialEq, Debug, get_size2::GetSize)]
 pub struct PlaceExpr {
     root_name: Name,
-    sub_segments: SmallVec<[PlaceExprSubSegment; 1]>,
+    #[get_size(size_fn=sub_segments_size)]
+    sub_segments: thin_vec::ThinVec<PlaceExprSubSegment>,
+}
+
+fn sub_segments_size(segments: &thin_vec::ThinVec<PlaceExprSubSegment>) -> usize {
+    segments.capacity() * std::mem::size_of::<PlaceExprSubSegment>()
+        + segments
+            .iter()
+            .map(get_size2::GetSize::get_heap_size)
+            .sum::<usize>()
 }
 
 impl std::fmt::Display for PlaceExpr {
@@ -162,10 +171,10 @@ impl TryFrom<ast::ExprRef<'_>> for PlaceExpr {
 }
 
 impl PlaceExpr {
-    pub(crate) const fn name(name: Name) -> Self {
+    pub(crate) fn name(name: Name) -> Self {
         Self {
             root_name: name,
-            sub_segments: SmallVec::new_const(),
+            sub_segments: thin_vec::ThinVec::new(),
         }
     }
 
@@ -652,6 +661,8 @@ pub struct PlaceTable {
 impl PlaceTable {
     fn shrink_to_fit(&mut self) {
         self.places.shrink_to_fit();
+        self.place_set
+            .shrink_to_fit(|id| PlaceTable::hash_place_expr(&self.places[*id].expr));
     }
 
     pub(crate) fn place_expr(&self, place_id: impl Into<ScopedPlaceId>) -> &PlaceExprWithFlags {
@@ -775,7 +786,7 @@ impl std::fmt::Debug for PlaceTable {
 pub(super) struct PlaceTableBuilder {
     table: PlaceTable,
 
-    associated_place_ids: IndexVec<ScopedPlaceId, Vec<ScopedPlaceId>>,
+    associated_place_ids: IndexVec<ScopedPlaceId, SmallVec<[ScopedPlaceId; 4]>>,
 }
 
 impl PlaceTableBuilder {
@@ -794,14 +805,17 @@ impl PlaceTableBuilder {
 
                 let id = self.table.places.push(symbol);
                 entry.insert(id);
-                let new_id = self.associated_place_ids.push(vec![]);
+                let new_id = self.associated_place_ids.push(SmallVec::new_const());
                 debug_assert_eq!(new_id, id);
                 (id, true)
             }
         }
     }
 
-    pub(super) fn add_place(&mut self, place_expr: PlaceExprWithFlags) -> (ScopedPlaceId, bool) {
+    pub(super) fn add_place(
+        &mut self,
+        mut place_expr: PlaceExprWithFlags,
+    ) -> (ScopedPlaceId, bool) {
         let hash = PlaceTable::hash_place_expr(&place_expr.expr);
         let entry = self.table.place_set.entry(
             hash,
@@ -812,9 +826,10 @@ impl PlaceTableBuilder {
         match entry {
             Entry::Occupied(entry) => (*entry.get(), false),
             Entry::Vacant(entry) => {
+                place_expr.expr.sub_segments.shrink_to_fit();
                 let id = self.table.places.push(place_expr);
                 entry.insert(id);
-                let new_id = self.associated_place_ids.push(vec![]);
+                let new_id = self.associated_place_ids.push(SmallVec::new_const());
                 debug_assert_eq!(new_id, id);
                 for root in self.table.places[id].expr.root_exprs() {
                     if let Some(root_id) = self.table.place_id_by_expr(root) {