[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.
This commit is contained in:
Ibraheem Ahmed
2025-06-13 08:40:11 -04:00
committed by GitHub
parent 76d9009a6e
commit c9dff5c7d5
824 changed files with 25243 additions and 804 deletions

View File

@@ -6,9 +6,9 @@ use rustc_hash::{FxBuildHasher, FxHashSet};
use ruff_python_ast::name::Name;
use ruff_python_ast::{
self as ast, AnyStringFlags, BoolOp, CmpOp, ConversionFlag, Expr, ExprContext, FString,
InterpolatedStringElement, InterpolatedStringElements, IpyEscapeKind, Number, Operator,
OperatorPrecedence, StringFlags, TString, UnaryOp,
self as ast, AnyStringFlags, AtomicNodeIndex, BoolOp, CmpOp, ConversionFlag, Expr, ExprContext,
FString, InterpolatedStringElement, InterpolatedStringElements, IpyEscapeKind, Number,
Operator, OperatorPrecedence, StringFlags, TString, UnaryOp,
};
use ruff_text_size::{Ranged, TextLen, TextRange, TextSize};
@@ -305,6 +305,7 @@ impl<'src> Parser<'src> {
op: bin_op,
right: Box::new(right.expr),
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
})
}
};
@@ -472,6 +473,7 @@ impl<'src> Parser<'src> {
range: identifier.range,
id: identifier.id,
ctx,
node_index: AtomicNodeIndex::dummy(),
}
}
@@ -487,13 +489,21 @@ impl<'src> Parser<'src> {
let TokenValue::Name(name) = self.bump_value(TokenKind::Name) else {
unreachable!();
};
return ast::Identifier { id: name, range };
return ast::Identifier {
id: name,
range,
node_index: AtomicNodeIndex::dummy(),
};
}
if self.current_token_kind().is_soft_keyword() {
let id = Name::new(self.src_text(range));
self.bump_soft_keyword_as_name();
return ast::Identifier { id, range };
return ast::Identifier {
id,
range,
node_index: AtomicNodeIndex::dummy(),
};
}
if self.current_token_kind().is_keyword() {
@@ -508,7 +518,11 @@ impl<'src> Parser<'src> {
let id = Name::new(self.src_text(range));
self.bump_any();
ast::Identifier { id, range }
ast::Identifier {
id,
range,
node_index: AtomicNodeIndex::dummy(),
}
} else {
self.add_error(
ParseErrorType::OtherError("Expected an identifier".into()),
@@ -518,6 +532,7 @@ impl<'src> Parser<'src> {
ast::Identifier {
id: Name::empty(),
range: self.missing_node_range(),
node_index: AtomicNodeIndex::dummy(),
}
}
}
@@ -537,6 +552,7 @@ impl<'src> Parser<'src> {
Expr::NumberLiteral(ast::ExprNumberLiteral {
value: Number::Float(value),
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
})
}
TokenKind::Complex => {
@@ -546,6 +562,7 @@ impl<'src> Parser<'src> {
Expr::NumberLiteral(ast::ExprNumberLiteral {
value: Number::Complex { real, imag },
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
})
}
TokenKind::Int => {
@@ -555,6 +572,7 @@ impl<'src> Parser<'src> {
Expr::NumberLiteral(ast::ExprNumberLiteral {
value: Number::Int(value),
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
})
}
TokenKind::True => {
@@ -562,6 +580,7 @@ impl<'src> Parser<'src> {
Expr::BooleanLiteral(ast::ExprBooleanLiteral {
value: true,
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
})
}
TokenKind::False => {
@@ -569,18 +588,21 @@ impl<'src> Parser<'src> {
Expr::BooleanLiteral(ast::ExprBooleanLiteral {
value: false,
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
})
}
TokenKind::None => {
self.bump(TokenKind::None);
Expr::NoneLiteral(ast::ExprNoneLiteral {
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
})
}
TokenKind::Ellipsis => {
self.bump(TokenKind::Ellipsis);
Expr::EllipsisLiteral(ast::ExprEllipsisLiteral {
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
})
}
TokenKind::Name => Expr::Name(self.parse_name()),
@@ -608,6 +630,7 @@ impl<'src> Parser<'src> {
range: self.missing_node_range(),
id: Name::empty(),
ctx: ExprContext::Invalid,
node_index: AtomicNodeIndex::dummy(),
})
}
}
@@ -650,6 +673,7 @@ impl<'src> Parser<'src> {
func: Box::new(func),
arguments,
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
}
}
@@ -679,6 +703,7 @@ impl<'src> Parser<'src> {
arg: None,
value: value.expr,
range: parser.node_range(argument_start),
node_index: AtomicNodeIndex::dummy(),
});
seen_keyword_unpacking = true;
@@ -743,6 +768,7 @@ impl<'src> Parser<'src> {
ast::Identifier {
id: ident_expr.id,
range: ident_expr.range,
node_index: AtomicNodeIndex::dummy(),
}
} else {
// TODO(dhruvmanila): Parser shouldn't drop the `parsed_expr` if it's
@@ -755,6 +781,7 @@ impl<'src> Parser<'src> {
ast::Identifier {
id: Name::empty(),
range: parsed_expr.range(),
node_index: AtomicNodeIndex::dummy(),
}
};
@@ -764,6 +791,7 @@ impl<'src> Parser<'src> {
arg: Some(arg),
value: value.expr,
range: parser.node_range(argument_start),
node_index: AtomicNodeIndex::dummy(),
});
} else {
if !parsed_expr.is_unparenthesized_starred_expr() {
@@ -788,6 +816,7 @@ impl<'src> Parser<'src> {
let arguments = ast::Arguments {
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
args: args.into_boxed_slice(),
keywords: keywords.into_boxed_slice(),
};
@@ -829,9 +858,11 @@ impl<'src> Parser<'src> {
range: slice_range,
id: Name::empty(),
ctx: ExprContext::Invalid,
node_index: AtomicNodeIndex::dummy(),
})),
ctx: ExprContext::Load,
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
};
}
@@ -851,6 +882,7 @@ impl<'src> Parser<'src> {
ctx: ExprContext::Load,
range: self.node_range(slice_start),
parenthesized: false,
node_index: AtomicNodeIndex::dummy(),
});
} else if slice.is_starred_expr() {
// If the only slice element is a starred expression, that is represented
@@ -861,6 +893,7 @@ impl<'src> Parser<'src> {
ctx: ExprContext::Load,
range: self.node_range(slice_start),
parenthesized: false,
node_index: AtomicNodeIndex::dummy(),
});
}
@@ -909,6 +942,7 @@ impl<'src> Parser<'src> {
slice: Box::new(slice),
ctx: ExprContext::Load,
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
}
}
@@ -1002,6 +1036,7 @@ impl<'src> Parser<'src> {
Expr::Slice(ast::ExprSlice {
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
lower,
upper,
step,
@@ -1032,6 +1067,7 @@ impl<'src> Parser<'src> {
op,
operand: Box::new(operand.expr),
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
}
}
@@ -1056,6 +1092,7 @@ impl<'src> Parser<'src> {
attr,
ctx: ExprContext::Load,
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
}
}
@@ -1099,6 +1136,7 @@ impl<'src> Parser<'src> {
values,
op,
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
}
}
@@ -1178,6 +1216,7 @@ impl<'src> Parser<'src> {
ops: operators.into_boxed_slice(),
comparators: comparators.into_boxed_slice(),
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
}
}
@@ -1229,18 +1268,22 @@ impl<'src> Parser<'src> {
StringType::Str(string) => Expr::StringLiteral(ast::ExprStringLiteral {
value: ast::StringLiteralValue::single(string),
range,
node_index: AtomicNodeIndex::dummy(),
}),
StringType::Bytes(bytes) => Expr::BytesLiteral(ast::ExprBytesLiteral {
value: ast::BytesLiteralValue::single(bytes),
range,
node_index: AtomicNodeIndex::dummy(),
}),
StringType::FString(fstring) => Expr::FString(ast::ExprFString {
value: ast::FStringValue::single(fstring),
range,
node_index: AtomicNodeIndex::dummy(),
}),
StringType::TString(tstring) => Expr::TString(ast::ExprTString {
value: ast::TStringValue::single(tstring),
range,
node_index: AtomicNodeIndex::dummy(),
}),
},
_ => self.handle_implicitly_concatenated_strings(strings, range),
@@ -1307,6 +1350,7 @@ impl<'src> Parser<'src> {
return Expr::from(ast::ExprBytesLiteral {
value: ast::BytesLiteralValue::concatenated(values),
range,
node_index: AtomicNodeIndex::dummy(),
});
}
Ordering::Greater => unreachable!(),
@@ -1346,6 +1390,7 @@ impl<'src> Parser<'src> {
return Expr::from(ast::ExprStringLiteral {
value: ast::StringLiteralValue::concatenated(values),
range,
node_index: AtomicNodeIndex::dummy(),
});
}
@@ -1367,6 +1412,7 @@ impl<'src> Parser<'src> {
return Expr::from(ast::ExprTString {
value: ast::TStringValue::concatenated(parts),
range,
node_index: AtomicNodeIndex::dummy(),
});
}
@@ -1387,6 +1433,7 @@ impl<'src> Parser<'src> {
Expr::from(ast::ExprFString {
value: ast::FStringValue::concatenated(parts),
range,
node_index: AtomicNodeIndex::dummy(),
})
}
@@ -1422,6 +1469,7 @@ impl<'src> Parser<'src> {
value: Box::new([]),
range,
flags: ast::BytesLiteralFlags::from(flags).with_invalid(),
node_index: AtomicNodeIndex::dummy(),
})
} else {
// test_err invalid_string_literal
@@ -1431,6 +1479,7 @@ impl<'src> Parser<'src> {
value: "".into(),
range,
flags: ast::StringLiteralFlags::from(flags).with_invalid(),
node_index: AtomicNodeIndex::dummy(),
})
}
}
@@ -1604,6 +1653,7 @@ impl<'src> Parser<'src> {
ast::InterpolatedStringLiteralElement {
value: "".into(),
range,
node_index: AtomicNodeIndex::dummy(),
}
}),
)
@@ -1759,6 +1809,7 @@ impl<'src> Parser<'src> {
Some(Box::new(ast::InterpolatedStringFormatSpec {
range: self.node_range(spec_start),
elements,
node_index: AtomicNodeIndex::dummy(),
}))
} else {
None
@@ -1810,6 +1861,7 @@ impl<'src> Parser<'src> {
conversion,
format_spec,
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
}
}
@@ -1839,6 +1891,7 @@ impl<'src> Parser<'src> {
elts: vec![],
ctx: ExprContext::Load,
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
});
}
@@ -1890,6 +1943,7 @@ impl<'src> Parser<'src> {
return Expr::Dict(ast::ExprDict {
items: vec![],
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
});
}
@@ -2000,6 +2054,7 @@ impl<'src> Parser<'src> {
elts: vec![],
ctx: ExprContext::Load,
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
parenthesized: true,
})
.into();
@@ -2088,6 +2143,7 @@ impl<'src> Parser<'src> {
elts,
ctx: ExprContext::Load,
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
parenthesized: parenthesized.is_yes(),
}
}
@@ -2116,6 +2172,7 @@ impl<'src> Parser<'src> {
elts,
ctx: ExprContext::Load,
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
}
}
@@ -2164,6 +2221,7 @@ impl<'src> Parser<'src> {
ast::ExprSet {
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
elts,
}
}
@@ -2206,6 +2264,7 @@ impl<'src> Parser<'src> {
ast::ExprDict {
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
items,
}
}
@@ -2273,6 +2332,7 @@ impl<'src> Parser<'src> {
ast::Comprehension {
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
target: target.expr,
iter: iter.expr,
ifs,
@@ -2302,6 +2362,7 @@ impl<'src> Parser<'src> {
elt: Box::new(element),
generators,
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
parenthesized: parenthesized.is_yes(),
}
}
@@ -2322,6 +2383,7 @@ impl<'src> Parser<'src> {
elt: Box::new(element),
generators,
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
}
}
@@ -2343,6 +2405,7 @@ impl<'src> Parser<'src> {
value: Box::new(value),
generators,
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
}
}
@@ -2362,6 +2425,7 @@ impl<'src> Parser<'src> {
elt: Box::new(element),
generators,
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
}
}
@@ -2398,6 +2462,7 @@ impl<'src> Parser<'src> {
value: Box::new(parsed_expr.expr),
ctx: ExprContext::Load,
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
}
}
@@ -2420,6 +2485,7 @@ impl<'src> Parser<'src> {
ast::ExprAwait {
value: Box::new(parsed_expr.expr),
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
}
}
@@ -2468,6 +2534,7 @@ impl<'src> Parser<'src> {
Expr::Yield(ast::ExprYield {
value,
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
})
}
@@ -2507,6 +2574,7 @@ impl<'src> Parser<'src> {
Expr::YieldFrom(ast::ExprYieldFrom {
value: Box::new(expr),
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
})
}
@@ -2547,6 +2615,7 @@ impl<'src> Parser<'src> {
target: Box::new(target),
value: Box::new(value.expr),
range,
node_index: AtomicNodeIndex::dummy(),
}
}
@@ -2594,6 +2663,7 @@ impl<'src> Parser<'src> {
body: Box::new(body.expr),
parameters,
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
}
}
@@ -2618,6 +2688,7 @@ impl<'src> Parser<'src> {
test: Box::new(test.expr),
orelse: Box::new(orelse.expr),
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
}
}
@@ -2643,6 +2714,7 @@ impl<'src> Parser<'src> {
let command = ast::ExprIpyEscapeCommand {
range: self.node_range(start),
node_index: AtomicNodeIndex::dummy(),
kind,
value,
};
@@ -2901,6 +2973,7 @@ impl From<InterpolatedStringData> for FString {
elements: value.elements,
range: value.range,
flags: value.flags.into(),
node_index: AtomicNodeIndex::dummy(),
}
}
}
@@ -2911,6 +2984,7 @@ impl From<InterpolatedStringData> for TString {
elements: value.elements,
range: value.range,
flags: value.flags.into(),
node_index: AtomicNodeIndex::dummy(),
}
}
}