Skip to content

Commit 572e4b5

Browse files
charliermarshcharliermarsh
authored
[ty] Compact retained semantic arrays (#25454)
## Summary Reduces the memory retained by Salsa-cached semantic indexes by using a `FrozenIndexVec` for index-addressed immutable arrays, like the `FrozenMap` abstraction. Apart from a >1% memory reduction, we see a 1% improvement on `attrs`, `anyio`, `hydra-zen`, `multithreaded`; 1% regression on `pydantic` and `freqtrade`.
1 parent d475a23 commit 572e4b5

10 files changed

Lines changed: 228 additions & 86 deletions

File tree

‎crates/ruff_index/src/frozen.rs‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
use crate::{Idx, IndexSlice, IndexVec};
2+
use std::fmt::{Debug, Formatter};
3+
use std::marker::PhantomData;
4+
use std::ops::{Deref, DerefMut};
5+
6+
/// A structurally immutable sequence of `T` indexed by `I`.
7+
#[derive(Clone, PartialEq, Eq, Hash, get_size2::GetSize)]
8+
pub struct FrozenIndexVec<I, T> {
9+
raw: Box<[T]>,
10+
index: PhantomData<I>,
11+
}
12+
13+
impl<I: Idx, T> FrozenIndexVec<I, T> {
14+
#[inline]
15+
pub fn from_raw(raw: Box<[T]>) -> Self {
16+
Self {
17+
raw,
18+
index: PhantomData,
19+
}
20+
}
21+
22+
#[inline]
23+
pub fn as_slice(&self) -> &IndexSlice<I, T> {
24+
IndexSlice::from_raw(&self.raw)
25+
}
26+
27+
#[inline]
28+
pub fn as_mut_slice(&mut self) -> &mut IndexSlice<I, T> {
29+
IndexSlice::from_raw_mut(&mut self.raw)
30+
}
31+
}
32+
33+
impl<I, T> Debug for FrozenIndexVec<I, T>
34+
where
35+
T: Debug,
36+
{
37+
fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
38+
std::fmt::Debug::fmt(&self.raw, f)
39+
}
40+
}
41+
42+
impl<I: Idx, T> Deref for FrozenIndexVec<I, T> {
43+
type Target = IndexSlice<I, T>;
44+
45+
#[inline]
46+
fn deref(&self) -> &Self::Target {
47+
self.as_slice()
48+
}
49+
}
50+
51+
impl<I: Idx, T> DerefMut for FrozenIndexVec<I, T> {
52+
#[inline]
53+
fn deref_mut(&mut self) -> &mut Self::Target {
54+
self.as_mut_slice()
55+
}
56+
}
57+
58+
impl<I: Idx, T> From<IndexVec<I, T>> for FrozenIndexVec<I, T> {
59+
#[inline]
60+
fn from(vec: IndexVec<I, T>) -> Self {
61+
Self::from_raw(vec.raw.into_boxed_slice())
62+
}
63+
}
64+
65+
impl<I: Idx, T> FromIterator<T> for FrozenIndexVec<I, T> {
66+
#[inline]
67+
fn from_iter<Iter: IntoIterator<Item = T>>(iter: Iter) -> Self {
68+
Self::from_raw(iter.into_iter().collect())
69+
}
70+
}
71+
72+
// Whether `FrozenIndexVec` is `Send` depends only on the data,
73+
// not the phantom data.
74+
#[expect(unsafe_code)]
75+
unsafe impl<I: Idx, T> Send for FrozenIndexVec<I, T> where T: Send {}
76+
77+
#[expect(unsafe_code)]
78+
#[cfg(feature = "salsa")]
79+
unsafe impl<I, T> salsa::Update for FrozenIndexVec<I, T>
80+
where
81+
T: salsa::Update,
82+
{
83+
#[expect(unsafe_code)]
84+
unsafe fn maybe_update(old_pointer: *mut Self, new_value: Self) -> bool {
85+
let old_box: &mut FrozenIndexVec<I, T> = unsafe { &mut *old_pointer };
86+
unsafe { salsa::Update::maybe_update(&raw mut old_box.raw, new_value.raw) }
87+
}
88+
}

‎crates/ruff_index/src/lib.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,12 @@
33
//!
44
//! Inspired by [rustc_index](https://github.com/rust-lang/rust/blob/master/compiler/rustc_index/src/lib.rs).
55
6+
mod frozen;
67
mod idx;
78
mod slice;
89
mod vec;
910

11+
pub use frozen::FrozenIndexVec;
1012
pub use idx::Idx;
1113
pub use ruff_macros::newtype_index;
1214
pub use slice::IndexSlice;

‎crates/ty_python_core/src/ast_ids.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ fn ast_ids<'db>(db: &'db dyn Db, scope: ScopeId) -> &'db AstIds {
4343

4444
/// Uniquely identifies a use of a name in a [`crate::FileScopeId`].
4545
#[newtype_index]
46-
#[derive(get_size2::GetSize)]
46+
#[derive(Ord, PartialOrd, get_size2::GetSize)]
4747
pub struct ScopedUseId;
4848

4949
pub trait HasScopedUseId {

‎crates/ty_python_core/src/builder.rs‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2239,16 +2239,16 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
22392239
semantic_syntax_errors.shrink_to_fit();
22402240

22412241
SemanticIndex {
2242-
place_tables,
2243-
scopes: self.scopes,
2242+
place_tables: place_tables.into(),
2243+
scopes: self.scopes.into(),
22442244
definitions_by_node: self.definitions_by_node,
22452245
expressions_by_node: self.expressions_by_node,
22462246
statements_by_node: self.statements_by_node,
2247-
scope_ids_by_scope: self.scope_ids_by_scope,
2247+
scope_ids_by_scope: self.scope_ids_by_scope.into(),
22482248
ast_ids,
22492249
scopes_by_expression: self.scopes_by_expression.build(),
22502250
scopes_by_node: self.scopes_by_node,
2251-
use_def_maps,
2251+
use_def_maps: use_def_maps.into(),
22522252
enclosing_lambda_statements: self.enclosing_lambda_statements,
22532253
collections_by_use: self.collections_by_use,
22542254
uses_by_collection: self.uses_by_collection,

‎crates/ty_python_core/src/lib.rs‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ use std::sync::Arc;
55
use itertools::Itertools;
66
use ruff_db::files::File;
77
use ruff_db::parsed::parsed_module;
8-
use ruff_index::{IndexSlice, IndexVec};
8+
use ruff_index::{FrozenIndexVec, IndexSlice, IndexVec};
99
use ruff_python_ast::NodeIndex;
1010
use ruff_python_parser::semantic_errors::SemanticSyntaxError;
1111
use ruff_text_size::TextRange;
@@ -260,10 +260,10 @@ pub enum EnclosingSnapshotResult<'map, 'db> {
260260
#[derive(Debug, Update, get_size2::GetSize)]
261261
pub struct SemanticIndex<'db> {
262262
/// List of all place tables in this file, indexed by scope.
263-
place_tables: IndexVec<FileScopeId, Arc<PlaceTable>>,
263+
place_tables: FrozenIndexVec<FileScopeId, Arc<PlaceTable>>,
264264

265265
/// List of all scopes in this file.
266-
scopes: IndexVec<FileScopeId, Scope>,
266+
scopes: FrozenIndexVec<FileScopeId, Scope>,
267267

268268
/// Map expressions to their corresponding scope.
269269
scopes_by_expression: ExpressionsScopeMap,
@@ -290,10 +290,10 @@ pub struct SemanticIndex<'db> {
290290
uses_by_collection: FxHashMap<Definition<'db>, Vec<(Statement<'db>, ExpressionNodeKey)>>,
291291

292292
/// Map from the file-local [`FileScopeId`] to the salsa-ingredient [`ScopeId`].
293-
scope_ids_by_scope: IndexVec<FileScopeId, ScopeId<'db>>,
293+
scope_ids_by_scope: FrozenIndexVec<FileScopeId, ScopeId<'db>>,
294294

295295
/// Use-def map for each scope in this file.
296-
use_def_maps: IndexVec<FileScopeId, Arc<UseDefMap<'db>>>,
296+
use_def_maps: FrozenIndexVec<FileScopeId, Arc<UseDefMap<'db>>>,
297297

298298
/// Lookup table to map between node ids and ast nodes.
299299
///

‎crates/ty_python_core/src/predicate.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
//! static reachability of a binding, and the reachability of a statement or expression.
99
1010
use ruff_db::files::File;
11-
use ruff_index::{Idx, IndexVec};
11+
use ruff_index::{FrozenIndexVec, Idx, IndexVec};
1212
use ruff_python_ast::{Singleton, name::Name};
1313

1414
use crate::db::Db;
@@ -51,7 +51,7 @@ impl Idx for ScopedPredicateId {
5151
}
5252

5353
// A collection of predicates for a given scope.
54-
pub type Predicates<'db> = IndexVec<ScopedPredicateId, Predicate<'db>>;
54+
pub type Predicates<'db> = FrozenIndexVec<ScopedPredicateId, Predicate<'db>>;
5555

5656
#[derive(Debug, Default)]
5757
pub(crate) struct PredicatesBuilder<'db> {
@@ -68,7 +68,7 @@ impl<'db> PredicatesBuilder<'db> {
6868

6969
pub(crate) fn build(mut self) -> Predicates<'db> {
7070
self.predicates.shrink_to_fit();
71-
self.predicates
71+
self.predicates.into()
7272
}
7373
}
7474

0 commit comments

Comments
 (0)