Skip to content

Commit 68ea2b8

Browse files
authored
[red-knot] simplify "removing" in UnionBuilder::add (astral-sh#16947)
## Summary Simplify "removing" in UnionBuilder::add It's now O(m) instead of O(n + m) and easier to read. ## Test Plan cargo test (incl. mdtest)
1 parent e87fee4 commit 68ea2b8

File tree

1 file changed

+7
-18
lines changed
  • crates/red_knot_python_semantic/src/types

1 file changed

+7
-18
lines changed

crates/red_knot_python_semantic/src/types/builder.rs

Lines changed: 7 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -110,25 +110,14 @@ impl<'db> UnionBuilder<'db> {
110110
return self.collapse_to_object();
111111
}
112112
}
113-
match to_remove[..] {
114-
[] => self.elements.push(to_add),
115-
[index] => self.elements[index] = to_add,
116-
_ => {
117-
let mut current_index = 0;
118-
let mut to_remove = to_remove.into_iter();
119-
let mut next_to_remove_index = to_remove.next();
120-
self.elements.retain(|_| {
121-
let retain = if Some(current_index) == next_to_remove_index {
122-
next_to_remove_index = to_remove.next();
123-
false
124-
} else {
125-
true
126-
};
127-
current_index += 1;
128-
retain
129-
});
130-
self.elements.push(to_add);
113+
if let Some((&first, rest)) = to_remove.split_first() {
114+
self.elements[first] = to_add;
115+
// We iterate in descending order to keep remaining indices valid after `swap_remove`.
116+
for &index in rest.iter().rev() {
117+
self.elements.swap_remove(index);
131118
}
119+
} else {
120+
self.elements.push(to_add);
132121
}
133122
}
134123
}

0 commit comments

Comments
 (0)