[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:
@@ -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(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user