perf(ast, ast_codegen): optimize AST structs' memory layouts

[rzvxa]

No description provided.

[graphite-app]

Your org has enabled the Graphite merge queue for merging into main

Add the label “merge” to the PR and Graphite will automatically add it to the merge queue when it’s ready to merge. Or use the label “hotfix” to add to the merge queue as a hot fix.

You must have a Graphite account and log in to Graphite in order to use the merge queue. Sign up using this link.

[rzvxa]

Warning

This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
Learn more

• refactor(ast_codegen): more versatile generation pipeline. #4540
• perf(ast, ast_codegen): optimize AST structs' memory layouts #4404 👈
• chore: fix mistake in configuring deny.toml V2 #4619
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

Join @rzvxa and the rest of your teammates on Graphite

[codspeed-hq]

CodSpeed Performance Report

Merging #4404 will not alter performance

_{Comparing 07-22-feat_ast_codegen_add_alignment_and_size_data_to_the_schema (80e3195) with main (0c52c0d)}

Summary

✅ 32 untouched benchmarks

[overlookmotel]

For Span and Atom, can we mark them #[ast] and have codegen parse the files in oxc_span, same as it does for oxc_ast? We need both types to be made #[repr(C)] anyway.

[rzvxa]

I strongly believe we shouldn't extend the ast marker to types outside of the AST crate. Those one-off types can always be made repr(C) without all this infra. Span is already in the ideal order, And Vec type doesn't expose its fields so it can be reordered manually. Atom is a transparent type so no need for this stuff whatsoever.

But that's just my preference, I still haven't seen any real need to process those types.

[rzvxa]

With that said, I have nothing against hard coding them as opposed to using a layout of method. I just did it for the sake of consistency if we regenerate on multiple architectures.

[overlookmotel]

For Span and Atom we can mark them #[repr(C)] and hard-code them if you prefer. But:

1. AST transfer needs to know the field order for Span (so we'd have to hard-code that too).
2. We can't avoid extending outside of oxc_ast crate, because there are various types in oxc_syntax crate which are part of the AST e.g. AssignmentOperator. These types need #[repr(C)] and explicit discriminants. So it seems to me that simplest way to do that is #[ast] + codegen.

What's the reason for your opposition to codegen reaching into other crates?

[rzvxa]

I've mulled over this, I guess you are right it would be easier if we do so.
I wanted to keep the code generation's input limited to AST types, Vec, Box, Span, and Atom obviously aren't AST types. But your argument for types defined in oxc_syntax is valid.

I'm not familiar with all AST-related stuff in the oxc_syntax crate, But All I can remember is unit enums. So they should be easy to make repr(C).

Initially, I was slightly biased against doing this but now I'm natural. It sure helps not to have to worry about those types. Even if they all are simple enums.
I'll do it in a follow-up PR.

[rzvxa]

@overlookmotel I have a problem understanding something with repr C and option types. Can you help me with it?

Take a look at this struct definition:

oxc/crates/oxc_ast/src/ast/ts.rs

Lines 829 to 836 in b1b66e2

	pub struct TSCallSignatureDeclaration<'a> {
	#[serde(flatten)]
	pub span: Span,
	pub this_param: Option<TSThisParameter<'a>>,
	pub params: Box<'a, FormalParameters<'a>>,
	pub return_type: Option<Box<'a, TSTypeAnnotation<'a>>>,
	pub type_parameters: Option<Box<'a, TSTypeParameterDeclaration<'a>>>,
	}

If you get the size of the this_param field you'll see the option flag is folded into the TSThisParameter type which is 40 bytes. After investigating I found out that the option flag is getting folded onto the TSThisParameter::this field, That field contains an Atom which is a string slice and that type has 1 bit of free space in its pointer to accept the option.

Is it correct in the first place? If my assumptions are right; How does it work on FFI? I guess option doesn't mean nullable for that pointer types are a better fit, But packing this information in structures wouldn't break the C compatibility? Does it only happen on pointers because there is free space in the address? Right now I only fold the option if there is room at the end of type(either the last field of struct or if all enum variants have 1 bit headroom). Should I also consider all the padding in between fields as viable space for folding options?

---

Edit:

I went with my intuition but it might very well be wrong so please let me know if I'm doing it incorrectly.

[overlookmotel]

Option type itself is not repr(C) so it's unaffected. However its layout is specified: https://rust-lang.github.io/unsafe-code-guidelines/layout/enums.html#discriminant-elision-on-option-like-enums

To clarify one thing:

Option<T> will be same size as T if T has a "niche". This isn't quite the same thing as "1 bit of free space". A NonNull pointer doesn't have a spare bit - the memory address uses all 64 bits (not 63). But it does still have a "niche".

"Niche" means a pattern of bits which is illegal for the type to use. In the case of NonNull, that pattern is all 64 bits = zero, because a NonNull can't be null (aka 0).

I think the idea for FFI is that if you're getting a pointer from C, that C pointer can be either null or non-null. That maps cleanly onto Rust's Option<NonNull> - a null C pointer has same bit pattern as None, and a non-null C pointer has same bit pattern as Some(valid_ptr).

The rules around niches as far as I understand them are:

Box and Vec use a NonNull for their pointer, so they have a niche. &T must always point to a valid T, which can't be stored at null address. So &T is basically same as a NonNull, and has a niche (that includes &str).

Same thing for all the NonZero* integer types - NonZeroU32, NonZeroUsize etc. They all have a niche value of zero.

Some types have multiple niches. e.g. bool has 254 niches - 0 and 1 are the only legal values, and 2-255 are available as niches. So Option<bool> is 1 byte.

enum Foo { A: u32, B: u32, C: u32, D: u32 } has 252 niches because the discriminant only has 4 legal values, which leaves 252 spare values - even though the u32 "payloads" for all the variants have no niches.

(By the way, these complications around niches in enums is what motivated the "squashed enums" PR #3115. It was intractable to figure out what niches and discriminants Rust would produce when you have multiple enums, each with niches, nested enum-in-enum-in-enum).

If a struct has a field which has a niche, then the struct inherits that niche too (as you've found). The hard part is when a struct has more than 1 field with a niche - which niche does the struct use? I don't know the answer to that one - but we don't need an answer anyway right now - if T has any niches, Option<T> will be same size as T.

Padding cannot be used as a niche - I don't entirely understand that one, but pretty sure it's true.

NB: One oddity is that Cell<T> does NOT have a niche, even if T does (ditto UnsafeCell and RefCell).

One last thing: #[rustc_layout(debug)] can be helpful for figuring out what Rust is doing with type layouts. See https://www.ralfj.de/blog/2020/04/04/layout-debugging.html

[rzvxa]

Thank you, I'll rework it to consider niches correctly.

[overlookmotel]

Just to say, please don't delete this branch. We will want to return to this after we have added node_id field to all AST types.