Skip to content

Commit ae61600

Browse files
AlexWaygoodAlexWaygood
committed
[ty] Encapsulate constraint node interning
1 parent f40ea89 commit ae61600

4 files changed

Lines changed: 108 additions & 51 deletions

File tree

‎crates/ty_python_core/src/interned_nodes.rs‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
//! Data structures for interning interior nodes.
2+
//!
3+
//! Each graph needs stable IDs for its nodes and a way to reuse an existing ID when the same
4+
//! node is constructed again. The indexed vector owns the nodes; a hash table of IDs provides
5+
//! lookup without storing another copy of each node.
6+
7+
use std::hash::{BuildHasher, Hash};
8+
use std::ops::Index;
9+
10+
use hashbrown::{HashTable, hash_table::Entry};
11+
use ruff_index::{Idx, IndexVec};
12+
use rustc_hash::FxBuildHasher;
13+
14+
/// An indexed collection of distinct nodes with a reverse table of their IDs.
15+
///
16+
/// The vector owns each node at the index given by its stable ID. The hash table stores only
17+
/// IDs, hashing and comparing them by reading the corresponding nodes from the vector. Both
18+
/// collections are private, and [`intern`](Self::intern) updates them together: every table
19+
/// entry refers to a node in the vector, and an equal node reuses that node's ID. Once the
20+
/// graph is built, [`into_nodes`](Self::into_nodes) returns the vector and drops the table.
21+
#[derive(Debug)]
22+
pub(crate) struct InternedNodes<I: Idx, N> {
23+
nodes: IndexVec<I, N>,
24+
ids: HashTable<I>,
25+
}
26+
27+
impl<I: Idx, N> Default for InternedNodes<I, N> {
28+
fn default() -> Self {
29+
Self {
30+
nodes: IndexVec::default(),
31+
ids: HashTable::default(),
32+
}
33+
}
34+
}
35+
36+
impl<I: Idx, N> InternedNodes<I, N> {
37+
pub(crate) const fn len(&self) -> usize {
38+
self.nodes.raw.len()
39+
}
40+
41+
/// Consume `self` and return an iterator over the interned nodes.
42+
pub(crate) fn into_node_iterator(self) -> impl Iterator<Item = N> {
43+
self.nodes.into_iter()
44+
}
45+
46+
/// Consume `self` and return a boxed slice of the interned nodes.
47+
pub(crate) fn into_nodes_boxed_slice(self) -> Box<[N]> {
48+
self.nodes.raw.into_boxed_slice()
49+
}
50+
}
51+
52+
impl<I: Idx, N: Eq + Hash> InternedNodes<I, N> {
53+
pub(crate) fn find(&self, node: &N) -> Option<I> {
54+
self.ids
55+
.find(FxBuildHasher.hash_one(node), |id| self.nodes[*id].eq(node))
56+
.copied()
57+
}
58+
59+
/// Returns the node ID and whether the node was newly inserted.
60+
pub(crate) fn intern(&mut self, node: N) -> (I, bool) {
61+
let nodes = &mut self.nodes;
62+
match self.ids.entry(
63+
FxBuildHasher.hash_one(&node),
64+
|id| nodes[*id].eq(&node),
65+
|id| FxBuildHasher.hash_one(&nodes[*id]),
66+
) {
67+
Entry::Occupied(entry) => (*entry.get(), false),
68+
Entry::Vacant(entry) => {
69+
let id = nodes.push(node);
70+
entry.insert(id);
71+
(id, true)
72+
}
73+
}
74+
}
75+
}
76+
77+
impl<I: Idx, N> Index<I> for InternedNodes<I, N> {
78+
type Output = N;
79+
80+
fn index(&self, id: I) -> &N {
81+
&self.nodes[id]
82+
}
83+
}

‎crates/ty_python_core/src/lib.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ mod db;
4545
pub mod definition;
4646
pub mod expression;
4747
pub mod frozen;
48+
mod interned_nodes;
4849
pub(crate) mod member;
4950
pub mod narrowing_constraints;
5051
pub mod node_key;

‎crates/ty_python_core/src/narrowing_constraints.rs‎

Lines changed: 12 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -33,13 +33,12 @@
3333
//! `A OR (NOT A AND B)` simplifies to `A OR B`.
3434
3535
use std::cmp::Ordering;
36-
use std::hash::BuildHasher;
3736

38-
use hashbrown::hash_table::Entry;
39-
use ruff_index::{Idx, IndexVec};
40-
use rustc_hash::{FxBuildHasher, FxHashMap};
37+
use ruff_index::Idx;
38+
use rustc_hash::FxHashMap;
4139

4240
use crate::ast_ids::ScopedUseId;
41+
use crate::interned_nodes::InternedNodes;
4342
use crate::predicate::ScopedPredicateId;
4443
use crate::rank::{RankBitBox, RankBitBoxVec};
4544
use crate::scope::FileScopeId;
@@ -132,10 +131,8 @@ impl NarrowingConstraints {
132131

133132
#[derive(Debug, Default)]
134133
pub struct NarrowingConstraintsBuilder {
135-
interiors: IndexVec<ScopedNarrowingConstraint, InteriorNode>,
134+
interiors: InternedNodes<ScopedNarrowingConstraint, InteriorNode>,
136135
interior_used: RankBitBoxVec,
137-
// Nodes are already stored in `interiors`; keep only their IDs in the reverse table.
138-
interior_cache: hashbrown::HashTable<ScopedNarrowingConstraint>,
139136
and_cache: FxHashMap<
140137
(ScopedNarrowingConstraint, ScopedNarrowingConstraint),
141138
ScopedNarrowingConstraint,
@@ -150,13 +147,13 @@ impl NarrowingConstraintsBuilder {
150147
pub(crate) fn build(self) -> NarrowingConstraints {
151148
if self.interior_used.first_zero().is_none() {
152149
NarrowingConstraints {
153-
used_interiors: self.interiors.raw.into_boxed_slice(),
150+
used_interiors: self.interiors.into_nodes_boxed_slice(),
154151
used_indices: None,
155152
}
156153
} else {
157154
let used_interiors = self
158155
.interiors
159-
.into_iter()
156+
.into_node_iterator()
160157
.zip(&self.interior_used)
161158
.filter_map(|(interior, used)| used.then_some(interior))
162159
.collect();
@@ -231,20 +228,11 @@ impl NarrowingConstraintsBuilder {
231228
});
232229
}
233230

234-
let interiors = &mut self.interiors;
235-
match self.interior_cache.entry(
236-
FxBuildHasher.hash_one(node),
237-
|id| interiors[*id] == node,
238-
|id| FxBuildHasher.hash_one(interiors[*id]),
239-
) {
240-
Entry::Occupied(entry) => *entry.get(),
241-
Entry::Vacant(entry) => {
242-
self.interior_used.push(false);
243-
let id = interiors.push(node);
244-
entry.insert(id);
245-
id
246-
}
231+
let (id, inserted) = self.interiors.intern(node);
232+
if inserted {
233+
self.interior_used.push(false);
247234
}
235+
id
248236
}
249237

250238
pub(crate) fn add_atom(&mut self, predicate: ScopedPredicateId) -> ScopedNarrowingConstraint {
@@ -293,13 +281,8 @@ impl NarrowingConstraintsBuilder {
293281
if_uncertain: ALWAYS_FALSE,
294282
if_false,
295283
};
296-
if let Some(cached) = self
297-
.interior_cache
298-
.find(FxBuildHasher.hash_one(node), |id| {
299-
self.interiors[*id] == node
300-
})
301-
{
302-
return *cached;
284+
if let Some(cached) = self.interiors.find(&node) {
285+
return cached;
303286
}
304287
if self.interiors.len() >= MAX_INTERIOR_NODES {
305288
return ALWAYS_TRUE;

‎crates/ty_python_core/src/reachability_constraints.rs‎

Lines changed: 12 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,11 @@
33
//! See [`crate::reachability_constraints`] for more details.
44
55
use std::cmp::Ordering;
6-
use std::hash::BuildHasher;
76

8-
use hashbrown::hash_table::Entry;
9-
use ruff_index::{Idx, IndexVec};
10-
use rustc_hash::{FxBuildHasher, FxHashMap};
7+
use ruff_index::Idx;
8+
use rustc_hash::FxHashMap;
119

10+
use crate::interned_nodes::InternedNodes;
1211
use crate::narrowing_constraints::{NarrowingConstraintsBuilder, ScopedNarrowingConstraint};
1312
use crate::predicate::ScopedPredicateId;
1413
use crate::rank::{RankBitBox, RankBitBoxVec};
@@ -176,10 +175,8 @@ impl ReachabilityConstraints {
176175

177176
#[derive(Debug, Default)]
178177
pub struct ReachabilityConstraintsBuilder {
179-
interiors: IndexVec<ScopedReachabilityConstraintId, InteriorNode>,
178+
interiors: InternedNodes<ScopedReachabilityConstraintId, InteriorNode>,
180179
interior_used: RankBitBoxVec,
181-
// Nodes are already stored in `interiors`; keep only their IDs in the reverse table.
182-
interior_cache: hashbrown::HashTable<ScopedReachabilityConstraintId>,
183180
not_cache: FxHashMap<ScopedReachabilityConstraintId, ScopedReachabilityConstraintId>,
184181
and_cache: FxHashMap<
185182
(
@@ -206,11 +203,13 @@ impl ReachabilityConstraintsBuilder {
206203
pub(crate) fn build(self) -> ReachabilityConstraints {
207204
if self.interior_used.first_zero().is_none() {
208205
ReachabilityConstraints {
209-
used_interiors: self.interiors.raw.into_boxed_slice(),
206+
used_interiors: self.interiors.into_nodes_boxed_slice(),
210207
used_indices: None,
211208
}
212209
} else {
213-
let used_interiors = (self.interiors.into_iter())
210+
let used_interiors = self
211+
.interiors
212+
.into_node_iterator()
214213
.zip(&self.interior_used)
215214
.filter_map(|(interior, used)| used.then_some(interior))
216215
.collect();
@@ -354,20 +353,11 @@ impl ReachabilityConstraintsBuilder {
354353
return node.if_true;
355354
}
356355

357-
let interiors = &mut self.interiors;
358-
match self.interior_cache.entry(
359-
FxBuildHasher.hash_one(node),
360-
|id| interiors[*id] == node,
361-
|id| FxBuildHasher.hash_one(interiors[*id]),
362-
) {
363-
Entry::Occupied(entry) => *entry.get(),
364-
Entry::Vacant(entry) => {
365-
self.interior_used.push(false);
366-
let id = interiors.push(node);
367-
entry.insert(id);
368-
id
369-
}
356+
let (id, inserted) = self.interiors.intern(node);
357+
if inserted {
358+
self.interior_used.push(false);
370359
}
360+
id
371361
}
372362

373363
/// Adds a new reachability constraint that checks a single [`super::predicate::Predicate`].

0 commit comments

Comments
 (0)