Skip to content

Commit

Permalink
Avoid quadratic membership check in import fixes (#15576)
Browse files Browse the repository at this point in the history
## Summary

This leads to an explosion in runtime for (admittedly absurd) cases with
tens of thousands of imports.
  • Loading branch information
charliermarsh authored Jan 18, 2025
1 parent 6004c8c commit b8e5b95
Showing 1 changed file with 9 additions and 11 deletions.
20 changes: 9 additions & 11 deletions crates/ruff_linter/src/fix/codemods.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ use libcst_native::{
Codegen, CodegenState, Expression, ImportNames, NameOrAttribute, ParenthesizableWhitespace,
SmallStatement, Statement,
};
use rustc_hash::FxHashSet;
use smallvec::{smallvec, SmallVec};
use unicode_normalization::UnicodeNormalization;

Expand Down Expand Up @@ -80,14 +81,11 @@ pub(crate) fn remove_imports<'a>(
// Preserve the trailing comma (or not) from the last entry.
let trailing_comma = aliases.last().and_then(|alias| alias.comma.clone());

for member in member_names {
let alias_index = aliases
.iter()
.position(|alias| member == qualified_name_from_name_or_attribute(&alias.name));
if let Some(index) = alias_index {
aliases.remove(index);
}
}
// Remove any imports that are specified in the `imports` iterator.
let member_names = member_names.collect::<FxHashSet<_>>();
aliases.retain(|alias| {
!member_names.contains(qualified_name_from_name_or_attribute(&alias.name).as_str())
});

// But avoid destroying any trailing comments.
if let Some(alias) = aliases.last_mut() {
Expand Down Expand Up @@ -144,10 +142,10 @@ pub(crate) fn retain_imports(
// Preserve the trailing comma (or not) from the last entry.
let trailing_comma = aliases.last().and_then(|alias| alias.comma.clone());

// Retain any imports that are specified in the `imports` iterator.
let member_names = member_names.iter().copied().collect::<FxHashSet<_>>();
aliases.retain(|alias| {
member_names
.iter()
.any(|member| *member == qualified_name_from_name_or_attribute(&alias.name))
member_names.contains(qualified_name_from_name_or_attribute(&alias.name).as_str())
});

// But avoid destroying any trailing comments.
Expand Down

0 comments on commit b8e5b95

Please sign in to comment.