ruff/crates/ruff_python_formatter/src/statement/stmt_match.rs
Ibraheem Ahmed c9dff5c7d5
[ty] AST garbage collection (#18482)
## Summary

Garbage collect ASTs once we are done checking a given file. Queries
with a cross-file dependency on the AST will reparse the file on demand.
This reduces ty's peak memory usage by ~20-30%.

The primary change of this PR is adding a `node_index` field to every
AST node, that is assigned by the parser. `ParsedModule` can use this to
create a flat index of AST nodes any time the file is parsed (or
reparsed). This allows `AstNodeRef` to simply index into the current
instance of the `ParsedModule`, instead of storing a pointer directly.

The indices are somewhat hackily (using an atomic integer) assigned by
the `parsed_module` query instead of by the parser directly. Assigning
the indices in source-order in the (recursive) parser turns out to be
difficult, and collecting the nodes during semantic indexing is
impossible as `SemanticIndex` does not hold onto a specific
`ParsedModuleRef`, which the pointers in the flat AST are tied to. This
means that we have to do an extra AST traversal to assign and collect
the nodes into a flat index, but the small performance impact (~3% on
cold runs) seems worth it for the memory savings.

Part of https://github.com/astral-sh/ty/issues/214.
2025-06-13 08:40:11 -04:00

68 lines
2.1 KiB
Rust

use ruff_formatter::{format_args, write};
use ruff_python_ast::StmtMatch;
use crate::comments::leading_alternate_branch_comments;
use crate::context::{NodeLevel, WithNodeLevel};
use crate::expression::maybe_parenthesize_expression;
use crate::expression::parentheses::Parenthesize;
use crate::prelude::*;
use crate::statement::clause::{ClauseHeader, clause_header};
#[derive(Default)]
pub struct FormatStmtMatch;
impl FormatNodeRule<StmtMatch> for FormatStmtMatch {
fn fmt_fields(&self, item: &StmtMatch, f: &mut PyFormatter) -> FormatResult<()> {
let StmtMatch {
range: _,
node_index: _,
subject,
cases,
} = item;
let comments = f.context().comments().clone();
let dangling_item_comments = comments.dangling(item);
// There can be at most one dangling comment after the colon in a match statement.
debug_assert!(dangling_item_comments.len() <= 1);
clause_header(
ClauseHeader::Match(item),
dangling_item_comments,
&format_args![
token("match"),
space(),
maybe_parenthesize_expression(subject, item, Parenthesize::IfBreaks),
],
)
.fmt(f)?;
let mut cases_iter = cases.iter();
let Some(first) = cases_iter.next() else {
return Ok(());
};
// The new level is for the `case` nodes.
let mut f = WithNodeLevel::new(NodeLevel::CompoundStatement, f);
write!(f, [block_indent(&first.format())])?;
let mut last_case = first;
for case in cases_iter {
let last_suite_in_statement = Some(case) == cases.last();
write!(
f,
[block_indent(&format_args!(
leading_alternate_branch_comments(
comments.leading(case),
last_case.body.last(),
),
case.format().with_options(last_suite_in_statement)
))]
)?;
last_case = case;
}
Ok(())
}
}