[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

@@ -7,14 +7,17 @@ input_file: crates/ruff_python_parser/resources/valid/statement/assignment.py
```
Module(
ModModule {
node_index: AtomicNodeIndex(..),
range: 0..729,
body: [
Assign(
StmtAssign {
node_index: AtomicNodeIndex(..),
range: 0..13,
targets: [
Name(
ExprName {
node_index: AtomicNodeIndex(..),
range: 0..1,
id: Name("x"),
ctx: Store,
@@ -23,10 +26,12 @@ Module(
],
value: Tuple(
ExprTuple {
node_index: AtomicNodeIndex(..),
range: 4..13,
elts: [
NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 5..6,
value: Int(
1,
@@ -35,6 +40,7 @@ Module(
),
NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 8..9,
value: Int(
2,
@@ -43,6 +49,7 @@ Module(
),
NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 11..12,
value: Int(
3,
@@ -58,14 +65,17 @@ Module(
),
Assign(
StmtAssign {
node_index: AtomicNodeIndex(..),
range: 15..33,
targets: [
Tuple(
ExprTuple {
node_index: AtomicNodeIndex(..),
range: 15..21,
elts: [
Name(
ExprName {
node_index: AtomicNodeIndex(..),
range: 16..17,
id: Name("x"),
ctx: Store,
@@ -73,6 +83,7 @@ Module(
),
Name(
ExprName {
node_index: AtomicNodeIndex(..),
range: 19..20,
id: Name("y"),
ctx: Store,
@@ -86,10 +97,12 @@ Module(
],
value: Tuple(
ExprTuple {
node_index: AtomicNodeIndex(..),
range: 24..33,
elts: [
NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 25..26,
value: Int(
1,
@@ -98,6 +111,7 @@ Module(
),
NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 28..29,
value: Int(
2,
@@ -106,6 +120,7 @@ Module(
),
NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 31..32,
value: Int(
3,
@@ -121,14 +136,17 @@ Module(
),
Assign(
StmtAssign {
node_index: AtomicNodeIndex(..),
range: 35..53,
targets: [
List(
ExprList {
node_index: AtomicNodeIndex(..),
range: 35..41,
elts: [
Name(
ExprName {
node_index: AtomicNodeIndex(..),
range: 36..37,
id: Name("x"),
ctx: Store,
@@ -136,6 +154,7 @@ Module(
),
Name(
ExprName {
node_index: AtomicNodeIndex(..),
range: 39..40,
id: Name("y"),
ctx: Store,
@@ -148,10 +167,12 @@ Module(
],
value: Tuple(
ExprTuple {
node_index: AtomicNodeIndex(..),
range: 44..53,
elts: [
NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 45..46,
value: Int(
1,
@@ -160,6 +181,7 @@ Module(
),
NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 48..49,
value: Int(
2,
@@ -168,6 +190,7 @@ Module(
),
NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 51..52,
value: Int(
3,
@@ -183,13 +206,16 @@ Module(
),
Assign(
StmtAssign {
node_index: AtomicNodeIndex(..),
range: 55..70,
targets: [
Attribute(
ExprAttribute {
node_index: AtomicNodeIndex(..),
range: 55..58,
value: Name(
ExprName {
node_index: AtomicNodeIndex(..),
range: 55..56,
id: Name("x"),
ctx: Load,
@@ -198,6 +224,7 @@ Module(
attr: Identifier {
id: Name("y"),
range: 57..58,
node_index: AtomicNodeIndex(..),
},
ctx: Store,
},
@@ -205,10 +232,12 @@ Module(
],
value: Tuple(
ExprTuple {
node_index: AtomicNodeIndex(..),
range: 61..70,
elts: [
NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 62..63,
value: Int(
1,
@@ -217,6 +246,7 @@ Module(
),
NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 65..66,
value: Int(
2,
@@ -225,6 +255,7 @@ Module(
),
NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 68..69,
value: Int(
3,
@@ -240,13 +271,16 @@ Module(
),
Assign(
StmtAssign {
node_index: AtomicNodeIndex(..),
range: 72..88,
targets: [
Subscript(
ExprSubscript {
node_index: AtomicNodeIndex(..),
range: 72..76,
value: Name(
ExprName {
node_index: AtomicNodeIndex(..),
range: 72..73,
id: Name("x"),
ctx: Load,
@@ -254,6 +288,7 @@ Module(
),
slice: Name(
ExprName {
node_index: AtomicNodeIndex(..),
range: 74..75,
id: Name("y"),
ctx: Load,
@@ -265,10 +300,12 @@ Module(
],
value: Tuple(
ExprTuple {
node_index: AtomicNodeIndex(..),
range: 79..88,
elts: [
NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 80..81,
value: Int(
1,
@@ -277,6 +314,7 @@ Module(
),
NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 83..84,
value: Int(
2,
@@ -285,6 +323,7 @@ Module(
),
NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 86..87,
value: Int(
3,
@@ -300,14 +339,17 @@ Module(
),
Assign(
StmtAssign {
node_index: AtomicNodeIndex(..),
range: 90..109,
targets: [
Tuple(
ExprTuple {
node_index: AtomicNodeIndex(..),
range: 90..97,
elts: [
Name(
ExprName {
node_index: AtomicNodeIndex(..),
range: 91..92,
id: Name("x"),
ctx: Store,
@@ -315,9 +357,11 @@ Module(
),
Starred(
ExprStarred {
node_index: AtomicNodeIndex(..),
range: 94..96,
value: Name(
ExprName {
node_index: AtomicNodeIndex(..),
range: 95..96,
id: Name("y"),
ctx: Store,
@@ -334,10 +378,12 @@ Module(
],
value: Tuple(
ExprTuple {
node_index: AtomicNodeIndex(..),
range: 100..109,
elts: [
NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 101..102,
value: Int(
1,
@@ -346,6 +392,7 @@ Module(
),
NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 104..105,
value: Int(
2,
@@ -354,6 +401,7 @@ Module(
),
NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 107..108,
value: Int(
3,
@@ -369,14 +417,17 @@ Module(
),
Assign(
StmtAssign {
node_index: AtomicNodeIndex(..),
range: 259..280,
targets: [
List(
ExprList {
node_index: AtomicNodeIndex(..),
range: 259..268,
elts: [
Name(
ExprName {
node_index: AtomicNodeIndex(..),
range: 260..261,
id: Name("x"),
ctx: Store,
@@ -384,6 +435,7 @@ Module(
),
Name(
ExprName {
node_index: AtomicNodeIndex(..),
range: 263..264,
id: Name("y"),
ctx: Store,
@@ -391,6 +443,7 @@ Module(
),
Name(
ExprName {
node_index: AtomicNodeIndex(..),
range: 266..267,
id: Name("z"),
ctx: Store,
@@ -403,10 +456,12 @@ Module(
],
value: List(
ExprList {
node_index: AtomicNodeIndex(..),
range: 271..280,
elts: [
NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 272..273,
value: Int(
1,
@@ -415,6 +470,7 @@ Module(
),
NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 275..276,
value: Int(
2,
@@ -423,6 +479,7 @@ Module(
),
NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 278..279,
value: Int(
3,
@@ -437,14 +494,17 @@ Module(
),
Assign(
StmtAssign {
node_index: AtomicNodeIndex(..),
range: 282..303,
targets: [
Tuple(
ExprTuple {
node_index: AtomicNodeIndex(..),
range: 282..291,
elts: [
Name(
ExprName {
node_index: AtomicNodeIndex(..),
range: 283..284,
id: Name("x"),
ctx: Store,
@@ -452,6 +512,7 @@ Module(
),
Name(
ExprName {
node_index: AtomicNodeIndex(..),
range: 286..287,
id: Name("y"),
ctx: Store,
@@ -459,6 +520,7 @@ Module(
),
Name(
ExprName {
node_index: AtomicNodeIndex(..),
range: 289..290,
id: Name("z"),
ctx: Store,
@@ -472,10 +534,12 @@ Module(
],
value: Tuple(
ExprTuple {
node_index: AtomicNodeIndex(..),
range: 294..303,
elts: [
NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 295..296,
value: Int(
1,
@@ -484,6 +548,7 @@ Module(
),
NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 298..299,
value: Int(
2,
@@ -492,6 +557,7 @@ Module(
),
NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 301..302,
value: Int(
3,
@@ -507,13 +573,16 @@ Module(
),
Assign(
StmtAssign {
node_index: AtomicNodeIndex(..),
range: 304..313,
targets: [
Subscript(
ExprSubscript {
node_index: AtomicNodeIndex(..),
range: 304..308,
value: Name(
ExprName {
node_index: AtomicNodeIndex(..),
range: 304..305,
id: Name("x"),
ctx: Load,
@@ -521,6 +590,7 @@ Module(
),
slice: NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 306..307,
value: Int(
0,
@@ -533,6 +603,7 @@ Module(
],
value: NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 311..313,
value: Int(
42,
@@ -543,13 +614,16 @@ Module(
),
Assign(
StmtAssign {
node_index: AtomicNodeIndex(..),
range: 410..419,
targets: [
Subscript(
ExprSubscript {
node_index: AtomicNodeIndex(..),
range: 410..414,
value: NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 410..411,
value: Int(
5,
@@ -558,6 +632,7 @@ Module(
),
slice: NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 412..413,
value: Int(
0,
@@ -570,6 +645,7 @@ Module(
],
value: NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 417..419,
value: Int(
42,
@@ -580,13 +656,16 @@ Module(
),
Assign(
StmtAssign {
node_index: AtomicNodeIndex(..),
range: 420..433,
targets: [
Subscript(
ExprSubscript {
node_index: AtomicNodeIndex(..),
range: 420..426,
value: Name(
ExprName {
node_index: AtomicNodeIndex(..),
range: 420..421,
id: Name("x"),
ctx: Load,
@@ -594,10 +673,12 @@ Module(
),
slice: Slice(
ExprSlice {
node_index: AtomicNodeIndex(..),
range: 422..425,
lower: Some(
NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 422..423,
value: Int(
1,
@@ -608,6 +689,7 @@ Module(
upper: Some(
NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 424..425,
value: Int(
2,
@@ -624,10 +706,12 @@ Module(
],
value: List(
ExprList {
node_index: AtomicNodeIndex(..),
range: 429..433,
elts: [
NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 430..432,
value: Int(
42,
@@ -642,13 +726,16 @@ Module(
),
Assign(
StmtAssign {
node_index: AtomicNodeIndex(..),
range: 529..542,
targets: [
Subscript(
ExprSubscript {
node_index: AtomicNodeIndex(..),
range: 529..535,
value: NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 529..530,
value: Int(
5,
@@ -657,10 +744,12 @@ Module(
),
slice: Slice(
ExprSlice {
node_index: AtomicNodeIndex(..),
range: 531..534,
lower: Some(
NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 531..532,
value: Int(
1,
@@ -671,6 +760,7 @@ Module(
upper: Some(
NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 533..534,
value: Int(
2,
@@ -687,10 +777,12 @@ Module(
],
value: List(
ExprList {
node_index: AtomicNodeIndex(..),
range: 538..542,
elts: [
NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 539..541,
value: Int(
42,
@@ -705,13 +797,16 @@ Module(
),
Assign(
StmtAssign {
node_index: AtomicNodeIndex(..),
range: 544..556,
targets: [
Attribute(
ExprAttribute {
node_index: AtomicNodeIndex(..),
range: 544..551,
value: Name(
ExprName {
node_index: AtomicNodeIndex(..),
range: 544..547,
id: Name("foo"),
ctx: Load,
@@ -720,6 +815,7 @@ Module(
attr: Identifier {
id: Name("bar"),
range: 548..551,
node_index: AtomicNodeIndex(..),
},
ctx: Store,
},
@@ -727,6 +823,7 @@ Module(
],
value: NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 554..556,
value: Int(
42,
@@ -737,18 +834,22 @@ Module(
),
Assign(
StmtAssign {
node_index: AtomicNodeIndex(..),
range: 658..670,
targets: [
Attribute(
ExprAttribute {
node_index: AtomicNodeIndex(..),
range: 658..665,
value: StringLiteral(
ExprStringLiteral {
node_index: AtomicNodeIndex(..),
range: 658..663,
value: StringLiteralValue {
inner: Single(
StringLiteral {
range: 658..663,
node_index: AtomicNodeIndex(..),
value: "foo",
flags: StringLiteralFlags {
quote_style: Double,
@@ -763,6 +864,7 @@ Module(
attr: Identifier {
id: Name("y"),
range: 664..665,
node_index: AtomicNodeIndex(..),
},
ctx: Store,
},
@@ -770,6 +872,7 @@ Module(
],
value: NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 668..670,
value: Int(
42,
@@ -780,10 +883,12 @@ Module(
),
Assign(
StmtAssign {
node_index: AtomicNodeIndex(..),
range: 672..680,
targets: [
Name(
ExprName {
node_index: AtomicNodeIndex(..),
range: 672..675,
id: Name("foo"),
ctx: Store,
@@ -792,6 +897,7 @@ Module(
],
value: NumberLiteral(
ExprNumberLiteral {
node_index: AtomicNodeIndex(..),
range: 678..680,
value: Int(
42,
@@ -802,10 +908,12 @@ Module(
),
Assign(
StmtAssign {
node_index: AtomicNodeIndex(..),
range: 682..695,
targets: [
List(
ExprList {
node_index: AtomicNodeIndex(..),
range: 682..684,
elts: [],
ctx: Store,
@@ -814,13 +922,16 @@ Module(
],
value: Tuple(
ExprTuple {
node_index: AtomicNodeIndex(..),
range: 687..695,
elts: [
Starred(
ExprStarred {
node_index: AtomicNodeIndex(..),
range: 688..693,
value: Name(
ExprName {
node_index: AtomicNodeIndex(..),
range: 689..693,
id: Name("data"),
ctx: Load,
@@ -838,10 +949,12 @@ Module(
),
Assign(
StmtAssign {
node_index: AtomicNodeIndex(..),
range: 696..709,
targets: [
Tuple(
ExprTuple {
node_index: AtomicNodeIndex(..),
range: 696..698,
elts: [],
ctx: Store,
@@ -851,13 +964,16 @@ Module(
],
value: Tuple(
ExprTuple {
node_index: AtomicNodeIndex(..),
range: 701..709,
elts: [
Starred(
ExprStarred {
node_index: AtomicNodeIndex(..),
range: 702..707,
value: Name(
ExprName {
node_index: AtomicNodeIndex(..),
range: 703..707,
id: Name("data"),
ctx: Load,
@@ -875,14 +991,17 @@ Module(
),
Assign(
StmtAssign {
node_index: AtomicNodeIndex(..),
range: 710..719,
targets: [
Tuple(
ExprTuple {
node_index: AtomicNodeIndex(..),
range: 710..714,
elts: [
Name(
ExprName {
node_index: AtomicNodeIndex(..),
range: 710..711,
id: Name("a"),
ctx: Store,
@@ -890,6 +1009,7 @@ Module(
),
Name(
ExprName {
node_index: AtomicNodeIndex(..),
range: 713..714,
id: Name("b"),
ctx: Store,
@@ -903,6 +1023,7 @@ Module(
],
value: Name(
ExprName {
node_index: AtomicNodeIndex(..),
range: 717..719,
id: Name("ab"),
ctx: Load,
@@ -912,10 +1033,12 @@ Module(
),
Assign(
StmtAssign {
node_index: AtomicNodeIndex(..),
range: 720..729,
targets: [
Name(
ExprName {
node_index: AtomicNodeIndex(..),
range: 720..721,
id: Name("a"),
ctx: Store,
@@ -923,6 +1046,7 @@ Module(
),
Name(
ExprName {
node_index: AtomicNodeIndex(..),
range: 724..725,
id: Name("b"),
ctx: Store,
@@ -931,6 +1055,7 @@ Module(
],
value: Name(
ExprName {
node_index: AtomicNodeIndex(..),
range: 728..729,
id: Name("c"),
ctx: Load,