From eb3cb8d4b2a9c6e761d4914703e4907dfdb01213 Mon Sep 17 00:00:00 2001 From: David Peter Date: Mon, 13 Jan 2025 11:10:42 +0100 Subject: [PATCH] [red-knot] Use `BitSet::union` for merging of declarations (#15451) ## Summary In `SymbolState` merging, use `BitSet::union` instead of inserting declarations one by one. This used to be the case but was changed in https://github.com/astral-sh/ruff/pull/15019 because we had to iterate over declarations anyway. This is an alternative to https://github.com/astral-sh/ruff/pull/15419 by @MichaReiser. It's similar in performance, but a bit more declarative and less imperative. --- .../src/semantic_index/use_def/bitset.rs | 66 +++++++++++++++++++ .../semantic_index/use_def/symbol_state.rs | 21 +++--- 2 files changed, 77 insertions(+), 10 deletions(-) diff --git a/crates/red_knot_python_semantic/src/semantic_index/use_def/bitset.rs b/crates/red_knot_python_semantic/src/semantic_index/use_def/bitset.rs index bf7bb01365..c44b553a1c 100644 --- a/crates/red_knot_python_semantic/src/semantic_index/use_def/bitset.rs +++ b/crates/red_knot_python_semantic/src/semantic_index/use_def/bitset.rs @@ -93,6 +93,19 @@ impl BitSet { } } + /// Union in-place with another [`BitSet`]. + pub(super) fn union(&mut self, other: &BitSet) { + let mut max_len = self.blocks().len(); + let other_len = other.blocks().len(); + if other_len > max_len { + max_len = other_len; + self.resize_blocks(max_len); + } + for (my_block, other_block) in self.blocks_mut().iter_mut().zip(other.blocks()) { + *my_block |= other_block; + } + } + /// Return an iterator over the values (in ascending order) in this [`BitSet`]. pub(super) fn iter(&self) -> BitSetIterator<'_, B> { let blocks = self.blocks(); @@ -222,6 +235,59 @@ mod tests { assert_bitset(&b1, &[89]); } + #[test] + fn union() { + let mut b1 = BitSet::<1>::with(2); + let b2 = BitSet::<1>::with(4); + + b1.union(&b2); + assert_bitset(&b1, &[2, 4]); + } + + #[test] + fn union_mixed_1() { + let mut b1 = BitSet::<1>::with(4); + let mut b2 = BitSet::<1>::with(4); + b1.insert(89); + b2.insert(5); + + b1.union(&b2); + assert_bitset(&b1, &[4, 5, 89]); + } + + #[test] + fn union_mixed_2() { + let mut b1 = BitSet::<1>::with(4); + let mut b2 = BitSet::<1>::with(4); + b1.insert(23); + b2.insert(89); + + b1.union(&b2); + assert_bitset(&b1, &[4, 23, 89]); + } + + #[test] + fn union_heap() { + let mut b1 = BitSet::<1>::with(4); + let mut b2 = BitSet::<1>::with(4); + b1.insert(89); + b2.insert(90); + + b1.union(&b2); + assert_bitset(&b1, &[4, 89, 90]); + } + + #[test] + fn union_heap_2() { + let mut b1 = BitSet::<1>::with(89); + let mut b2 = BitSet::<1>::with(89); + b1.insert(91); + b2.insert(90); + + b1.union(&b2); + assert_bitset(&b1, &[89, 90, 91]); + } + #[test] fn multiple_blocks() { let mut b = BitSet::<2>::with(120); diff --git a/crates/red_knot_python_semantic/src/semantic_index/use_def/symbol_state.rs b/crates/red_knot_python_semantic/src/semantic_index/use_def/symbol_state.rs index 0a440d6fa2..5a5a06c5b0 100644 --- a/crates/red_knot_python_semantic/src/semantic_index/use_def/symbol_state.rs +++ b/crates/red_knot_python_semantic/src/semantic_index/use_def/symbol_state.rs @@ -316,6 +316,9 @@ impl SymbolState { }; std::mem::swap(&mut a, self); + self.declarations + .live_declarations + .union(&b.declarations.live_declarations); let mut a_defs_iter = a.bindings.live_bindings.iter(); let mut b_defs_iter = b.bindings.live_bindings.iter(); @@ -449,10 +452,8 @@ impl SymbolState { let mut opt_a_decl: Option = a_decls_iter.next(); let mut opt_b_decl: Option = b_decls_iter.next(); - let push = |decl, - vis_constraints_iter: &mut VisibilityConstraintsIntoIterator, + let push = |vis_constraints_iter: &mut VisibilityConstraintsIntoIterator, merged: &mut Self| { - merged.declarations.live_declarations.insert(decl); let vis_constraints = vis_constraints_iter .next() .expect("declarations and visibility_constraints length mismatch"); @@ -466,15 +467,15 @@ impl SymbolState { match (opt_a_decl, opt_b_decl) { (Some(a_decl), Some(b_decl)) => match a_decl.cmp(&b_decl) { std::cmp::Ordering::Less => { - push(a_decl, &mut a_vis_constraints_iter, self); + push(&mut a_vis_constraints_iter, self); opt_a_decl = a_decls_iter.next(); } std::cmp::Ordering::Greater => { - push(b_decl, &mut b_vis_constraints_iter, self); + push(&mut b_vis_constraints_iter, self); opt_b_decl = b_decls_iter.next(); } std::cmp::Ordering::Equal => { - push(a_decl, &mut b_vis_constraints_iter, self); + push(&mut b_vis_constraints_iter, self); let a_vis_constraint = a_vis_constraints_iter .next() @@ -487,12 +488,12 @@ impl SymbolState { opt_b_decl = b_decls_iter.next(); } }, - (Some(a_decl), None) => { - push(a_decl, &mut a_vis_constraints_iter, self); + (Some(_), None) => { + push(&mut a_vis_constraints_iter, self); opt_a_decl = a_decls_iter.next(); } - (None, Some(b_decl)) => { - push(b_decl, &mut b_vis_constraints_iter, self); + (None, Some(_)) => { + push(&mut b_vis_constraints_iter, self); opt_b_decl = b_decls_iter.next(); } (None, None) => break,