Skip to content

perf(formatter): reduce AstNode size by 8 bytes using thread-local allocator #18340

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Dunqing wants to merge 1 commit into from
Closed

perf(formatter): reduce AstNode size by 8 bytes using thread-local allocator #18340

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2,721 changes: 1,033 additions & 1,688 deletions crates/oxc_formatter/src/ast_nodes/generated/ast_nodes.rs
View file
Open in desktop

Large diffs are not rendered by default.

740 changes: 229 additions & 511 deletions crates/oxc_formatter/src/ast_nodes/generated/format.rs
View file
Open in desktop

Large diffs are not rendered by default.

73 changes: 18 additions & 55 deletions crates/oxc_formatter/src/ast_nodes/iterator.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,47 +6,40 @@

use std::cmp::min;

use oxc_allocator::{Allocator, Vec};
use oxc_allocator::Vec;
use oxc_ast::ast::*;
use oxc_span::{GetSpan, Span};

use super::{AstNode, AstNodes};
use super::{AstNode, AstNodes, allocator};

pub struct AstNodeIterator<'a, T> {
inner: std::iter::Peekable<std::slice::Iter<'a, T>>,
parent: &'a AstNodes<'a>,
allocator: &'a Allocator,
}

macro_rules! impl_ast_node_vec {
($type:ty) => {
impl<'a> AstNode<'a, Vec<'a, $type>> {
pub fn iter(&self) -> AstNodeIterator<'a, $type> {
AstNodeIterator {
inner: self.inner.iter().peekable(),
parent: self.parent,
allocator: self.allocator,
}
AstNodeIterator { inner: self.inner.iter().peekable(), parent: self.parent }
}

pub fn first(&self) -> Option<&'a AstNode<'a, $type>> {
let mut inner_iter = self.inner.iter();
self.allocator
allocator()
.alloc(inner_iter.next().map(|inner| AstNode {
inner,
parent: self.parent,
allocator: self.allocator,
following_span: inner_iter.next().map(GetSpan::span),
}))
.as_ref()
}

pub fn last(&self) -> Option<&'a AstNode<'a, $type>> {
self.allocator
allocator()
.alloc(self.inner.last().map(|inner| AstNode {
inner,
parent: self.parent,
allocator: self.allocator,
following_span: None,
}))
.as_ref()
Expand All @@ -56,12 +49,10 @@ macro_rules! impl_ast_node_vec {
impl<'a> Iterator for AstNodeIterator<'a, $type> {
type Item = &'a AstNode<'a, $type>;
fn next(&mut self) -> Option<Self::Item> {
let allocator = self.allocator;
allocator
allocator()
.alloc(self.inner.next().map(|inner| AstNode {
parent: self.parent,
inner,
allocator,
following_span: self.inner.peek().copied().map(GetSpan::span),
}))
.as_ref()
Expand All @@ -75,7 +66,6 @@ macro_rules! impl_ast_node_vec {
AstNodeIterator::<$type> {
inner: self.inner.iter().peekable(),
parent: self.parent,
allocator: self.allocator,
}
}
}
Expand All @@ -86,32 +76,26 @@ macro_rules! impl_ast_node_vec_for_option {
($type:ty) => {
impl<'a> AstNode<'a, Vec<'a, $type>> {
pub fn iter(&self) -> AstNodeIterator<'a, $type> {
AstNodeIterator {
inner: self.inner.iter().peekable(),
parent: self.parent,
allocator: self.allocator,
}
AstNodeIterator { inner: self.inner.iter().peekable(), parent: self.parent }
}

pub fn first(&self) -> Option<&'a AstNode<'a, $type>> {
let mut inner_iter = self.inner.iter();
self.allocator
allocator()
.alloc(inner_iter.next().map(|inner| AstNode {
inner,
parent: self.parent,
allocator: self.allocator,
following_span:
inner_iter.next().and_then(|opt| opt.as_ref().map(GetSpan::span)),
}))
.as_ref()
}

pub fn last(&self) -> Option<&'a AstNode<'a, $type>> {
self.allocator
allocator()
.alloc(self.inner.last().map(|inner| AstNode {
inner,
parent: self.parent,
allocator: self.allocator,
following_span: None,
}))
.as_ref()
Expand All @@ -121,13 +105,11 @@ macro_rules! impl_ast_node_vec_for_option {
impl<'a> Iterator for AstNodeIterator<'a, $type> {
type Item = &'a AstNode<'a, $type>;
fn next(&mut self) -> Option<Self::Item> {
let allocator = self.allocator;
allocator
allocator()
.alloc(self.inner.next().map(|inner| {
AstNode {
parent: self.parent,
inner,
allocator,
following_span: self
.inner
.peek()
Expand All @@ -146,7 +128,6 @@ macro_rules! impl_ast_node_vec_for_option {
AstNodeIterator::<$type> {
inner: self.inner.iter().peekable(),
parent: self.parent,
allocator: self.allocator,
}
}
}
Expand Down Expand Up @@ -188,29 +169,23 @@ impl_ast_node_vec_for_option!(Option<BindingPattern<'a>>);
// <https://github.com/oxc-project/oxc/issues/10409>
impl<'a> AstNode<'a, Vec<'a, Statement<'a>>> {
pub fn iter(&self) -> AstNodeIterator<'a, Statement<'a>> {
AstNodeIterator {
inner: self.inner.iter().peekable(),
parent: self.parent,
allocator: self.allocator,
}
AstNodeIterator { inner: self.inner.iter().peekable(), parent: self.parent }
}
pub fn first(&self) -> Option<&'a AstNode<'a, Statement<'a>>> {
let mut inner_iter = self.inner.iter();
self.allocator
allocator()
.alloc(inner_iter.next().map(|inner| AstNode {
inner,
parent: self.parent,
allocator: self.allocator,
following_span: inner_iter.next().map(GetSpan::span),
}))
.as_ref()
}
pub fn last(&self) -> Option<&'a AstNode<'a, Statement<'a>>> {
self.allocator
allocator()
.alloc(self.inner.last().map(|inner| AstNode {
inner,
parent: self.parent,
allocator: self.allocator,
following_span: None,
}))
.as_ref()
Expand All @@ -219,12 +194,10 @@ impl<'a> AstNode<'a, Vec<'a, Statement<'a>>> {
impl<'a> Iterator for AstNodeIterator<'a, Statement<'a>> {
type Item = &'a AstNode<'a, Statement<'a>>;
fn next(&mut self) -> Option<Self::Item> {
let allocator = self.allocator;
allocator
allocator()
.alloc(self.inner.next().map(|inner| AstNode {
parent: self.parent,
inner,
allocator,
following_span: {
match self.inner.peek() {
// `@decorator export default class A {}`
Expand Down Expand Up @@ -271,7 +244,6 @@ impl<'a> IntoIterator for &AstNode<'a, Vec<'a, Statement<'a>>> {
AstNodeIterator::<Statement<'a>> {
inner: self.inner.iter().peekable(),
parent: self.parent,
allocator: self.allocator,
}
}
}
Expand All @@ -293,20 +265,15 @@ fn get_following_span_for_directive_parent(parent: &AstNodes<'_>) -> Option<Span
// following_span for the last directive in Program.body.
impl<'a> AstNode<'a, Vec<'a, Directive<'a>>> {
pub fn iter(&self) -> AstNodeIterator<'a, Directive<'a>> {
AstNodeIterator {
inner: self.inner.iter().peekable(),
parent: self.parent,
allocator: self.allocator,
}
AstNodeIterator { inner: self.inner.iter().peekable(), parent: self.parent }
}
pub fn first(&self) -> Option<&'a AstNode<'a, Directive<'a>>> {
let mut inner_iter = self.inner.iter();
self.allocator
allocator()
.alloc(inner_iter.next().map(|inner| {
AstNode {
inner,
parent: self.parent,
allocator: self.allocator,
following_span: inner_iter
.next()
.map(GetSpan::span)
Expand All @@ -316,11 +283,10 @@ impl<'a> AstNode<'a, Vec<'a, Directive<'a>>> {
.as_ref()
}
pub fn last(&self) -> Option<&'a AstNode<'a, Directive<'a>>> {
self.allocator
allocator()
.alloc(self.inner.last().map(|inner| AstNode {
inner,
parent: self.parent,
allocator: self.allocator,
following_span: get_following_span_for_directive_parent(self.parent),
}))
.as_ref()
Expand All @@ -329,13 +295,11 @@ impl<'a> AstNode<'a, Vec<'a, Directive<'a>>> {
impl<'a> Iterator for AstNodeIterator<'a, Directive<'a>> {
type Item = &'a AstNode<'a, Directive<'a>>;
fn next(&mut self) -> Option<Self::Item> {
let allocator = self.allocator;
allocator
allocator()
.alloc(self.inner.next().map(|inner| {
AstNode {
parent: self.parent,
inner,
allocator,
following_span: self
.inner
.peek()
Expand All @@ -354,7 +318,6 @@ impl<'a> IntoIterator for &AstNode<'a, Vec<'a, Directive<'a>>> {
AstNodeIterator::<Directive<'a>> {
inner: self.inner.iter().peekable(),
parent: self.parent,
allocator: self.allocator,
}
}
}
30 changes: 30 additions & 0 deletions crates/oxc_formatter/src/ast_nodes/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,36 @@ pub mod impls;
mod iterator;
mod node;

use std::{cell::Cell, ptr};

use oxc_allocator::Allocator;

pub use generated::ast_nodes::AstNodes;
pub use iterator::AstNodeIterator;
pub use node::AstNode;

thread_local! {
static ALLOCATOR: Cell<Option<*const Allocator>> = const { Cell::new(None) };
}

/// Sets the thread-local allocator for use during AST node operations.
///
/// # Safety
/// The caller must ensure that the allocator outlives all uses of `allocator()`.
/// Typically this is called at the start of formatting and cleared after.
#[inline]
pub fn set_allocator(allocator: &Allocator) {
ALLOCATOR.with(|cell| cell.set(Some(ptr::from_ref::<Allocator>(allocator))));
}

/// Gets a reference to the thread-local allocator.
///
/// # Panics
/// Panics if no allocator has been set via `set_allocator`.
#[inline]
pub fn allocator<'a>() -> &'a Allocator {
ALLOCATOR.with(|cell| {
// SAFETY: The caller of `set_allocator` guarantees the allocator outlives this access.
unsafe { &*cell.get().expect("Allocator not set. Call set_allocator first.") }
})
}
22 changes: 10 additions & 12 deletions crates/oxc_formatter/src/ast_nodes/node.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,16 +4,15 @@ use std::{
ops::Deref,
};

use oxc_allocator::{Allocator, Vec};
use oxc_allocator::Vec;
use oxc_ast::ast::*;
use oxc_span::{GetSpan, Span};

use super::AstNodes;
use super::{AstNodes, allocator};

pub struct AstNode<'a, T> {
pub(super) inner: &'a T,
pub parent: &'a AstNodes<'a>,
pub(super) allocator: &'a Allocator,
pub(super) following_span: Option<Span>,
}

Expand Down Expand Up @@ -43,11 +42,10 @@ impl<'a, T> AsRef<T> for AstNode<'a, T> {

impl<'a, T> AstNode<'a, Option<T>> {
pub fn as_ref(&self) -> Option<&'a AstNode<'a, T>> {
self.allocator
allocator()
.alloc(self.inner.as_ref().map(|inner| AstNode {
inner,
parent: self.parent,
allocator: self.allocator,
following_span: self.following_span,
}))
.as_ref()
Expand Down Expand Up @@ -109,8 +107,8 @@ impl<'a, T> AstNode<'a, T> {
}

impl<'a> AstNode<'a, Program<'a>> {
pub fn new(inner: &'a Program<'a>, parent: &'a AstNodes<'a>, allocator: &'a Allocator) -> Self {
AstNode { inner, parent, allocator, following_span: None }
pub fn new(inner: &'a Program<'a>, parent: &'a AstNodes<'a>) -> Self {
AstNode { inner, parent, following_span: None }
}
}

Expand Down Expand Up @@ -154,7 +152,8 @@ impl<'a> AstNode<'a, ImportExpression<'a>> {
// This allows us to reuse CallExpression's argument formatting logic when printing
// import expressions, since import(source, options) has the same structure as
// a function call with arguments.
let mut arguments = Vec::new_in(self.allocator);
let allocator = allocator();
let mut arguments = Vec::new_in(allocator);

// SAFETY: Argument inherits all Expression variants through the inherit_variants! macro,
// so Expression and Argument have identical memory layout for shared variants.
Expand All @@ -169,13 +168,12 @@ impl<'a> AstNode<'a, ImportExpression<'a>> {
}
}

let arguments_ref = self.allocator.alloc(arguments);
let arguments_ref = allocator.alloc(arguments);
let following_span = self.following_span;

self.allocator.alloc(AstNode {
allocator.alloc(AstNode {
inner: arguments_ref,
allocator: self.allocator,
parent: self.allocator.alloc(AstNodes::ImportExpression({
parent: allocator.alloc(AstNodes::ImportExpression({
// SAFETY: `self` is already allocated in Arena, so transmute from `&` to `&'a` is safe.
unsafe {
transmute::<
Expand Down
7 changes: 5 additions & 2 deletions crates/oxc_formatter/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ pub use crate::ir_transform::options::*;
pub use crate::options::*;
pub use crate::service::*;
use crate::{
ast_nodes::{AstNode, AstNodes},
ast_nodes::{AstNode, AstNodes, set_allocator},
formatter::{FormatContext, Formatted},
ir_transform::SortImportsTransform,
};
Expand Down Expand Up @@ -58,8 +58,11 @@ impl<'a> Formatter<'a> {
program: &'a Program<'a>,
external_callbacks: Option<ExternalCallbacks>,
) -> Formatted<'a> {
// Set the thread-local allocator for use during AST node operations.
set_allocator(self.allocator);

let parent = self.allocator.alloc(AstNodes::Dummy());
let program_node = AstNode::new(program, parent, self.allocator);
let program_node = AstNode::new(program, parent);

let context = FormatContext::new(
program.source_text,
Expand Down
Loading
Loading