feat(allocator): add ReplaceWith trait#24012
Conversation
Merging this PR will not alter performance
Comparing Footnotes
|
There was a problem hiding this comment.
Pull request overview
This PR introduces an oxc_allocator::ReplaceWith trait intended to support common AST-mutation patterns (build a new node from the old node and write it back in the same slot) without the TakeIn::take_in “dummy write” that permanently consumes arena space.
Changes:
- Add
ReplaceWithtrait + tests inoxc_allocator, and re-export it from the crate root. - Extend
ast_toolsderives to generateimpl ReplaceWithfor AST schema types. - Register
ReplaceWithinoxc_ast_macros’ generated trait mapping for#[generate_derive(...)]assertions.
Reviewed changes
Copilot reviewed 5 out of 6 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| tasks/ast_tools/src/main.rs | Registers the new derive so schema codegen emits ReplaceWith impls. |
| tasks/ast_tools/src/derives/replace_with.rs | New derive generator for ReplaceWith impl blocks. |
| tasks/ast_tools/src/derives/mod.rs | Wires the new derive module into the derives registry. |
| crates/oxc_ast_macros/src/generated/derived_traits.rs | Adds ReplaceWith to the macro-side trait-path mapping used for compile-time assertions. |
| crates/oxc_allocator/src/replace_with.rs | Implements ReplaceWith::replace_with and adds unit tests. |
| crates/oxc_allocator/src/lib.rs | Adds the module and re-exports ReplaceWith. |
6c56bbc to
302522e
Compare
302522e to
fcdedec
Compare
9238bc3 to
0b14f1a
Compare
fcdedec to
043d24d
Compare
043d24d to
9fc4b0d
Compare
|
@Dunqing This is ready for review now. The bug that Copilot found is now fixed. |
|
@codex review |
|
Codex Review: Didn't find any major issues. More of your lovely PRs please. Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
Dunqing
left a comment
There was a problem hiding this comment.
I highly doubt this naive example is overthinking in practice, which has led to the implementation being more complicated and making it hard to understand the underlying logic 😅. But thank you for adding detailed handwritten comments to explain it.
Merge activity
|
## The problem
A common pattern we have when mutating an AST (e.g. in transformer) is:
- Take a node from the AST.
- Wrap it in some other node, or extract a field from that node.
- Write the result back into the same place in AST the original came from.
Rust does not make this easy, due to aliasing rules. Currently we do something like this (example from arrow function transform):
```rs
let Statement::ExpressionStatement(expr_stmt) = match stmt else { unreachable!() };
let expr = expr_stmt.expression.take_in(ctx);
*stmt = Statement::new_return_statement(expr.span(), Some(expr), ctx);
```
The problem with this is that `take_in` writes a dummy `Expression` into the AST, in order to get the owned `Expression`. This is pointless, as the whole `Statement` that contains that dummy `Expression` is overwritten in the next line - the dummy is immediately disconnected from the AST. But that dummy node lives on forever in the arena, consuming memory.
## Solution in this PR
Introduce `ReplaceWith` trait which replaces the node in place, without writing a dummy:
```rs
stmt.replace_with(|stmt| {
let Statement::ExpressionStatement(expr_stmt) = stmt else { unreachable!() };
let expr = expr_stmt.unbox().expression;
Statement::new_return_statement(expr.span(), Some(expr), ctx)
});
```
The closure receives an owned copy of the `Statement`, and can do whatever it needs to with it. The closure has to return a new `Statement` which is written to the "hole" the original came from.
This PR only adds the trait. Later PRs in this stack implement the trait on AST types, and then use it in various crates.
## Implementation
The difficulty is in making this sound even if the closure panics, and `std::panic::catch_unwind` is in use, allowing the old value to be observed.
A naive implementation would allow UB as follows:
```rs
use std::panic::{AssertUnwindSafe, catch_unwind};
use oxc_allocator::{Allocator, ArenaBox, ReplaceWith};
let allocator = Allocator::new();
let allocator = &allocator;
let mut boxed: ArenaBox<u32> = ArenaBox::new_in(1, &allocator);
let mut copy: Option<ArenaBox<u32>> = None;
let _ = catch_unwind(AssertUnwindSafe(|| {
boxed.replace_with(|old| {
// Move the copy out of the closure
copy = Some(old);
panic!("Unwind before write-back");
});
}));
let mut copy: ArenaBox<u32> = copy.unwrap();
let boxed_mut: &mut u32 = boxed.as_mut();
let copy_mut: &mut u32 = copy.as_mut();
// `boxed_mut` and `copy_mut` both point to same `u32`.
// We have 2 `&mut u32`s pointing to same place. Aliasing violation!
*boxed_mut = 2;
*copy_mut = 3;
```
The solution is to make it possible to implement `ReplaceWith` only on types that also implement `Dummy`.
If the closure panics, a guard is triggered which writes a dummy into the original slot which the type was moved out of. This happens before any other code can observe the value in the slot.
To keep the API ergonomic, we don't want to have to pass in an `&Allocator` to `replace_with`. But the dummy has to be allocated *somewhere*. So we use a global store of `Allocator`s, stored in a `static`, which live for the entire life of the process. When a panic occurs, a new `Allocator` is created, and the dummy is allocated in it.
This is a memory leak - the `Allocator` containing the dummy will never be freed until the process exits.
However, in practice this is fine, because:
#### 1. It shouldn't happen
The guard only triggers, and these leaked `Allocator`s are only created, if the closure panics - which it never should. If the closure *does* panic, that's a bug, and we'll aim to fix it swiftly.
#### 2. Panic usually exits process
It's unusual to use `catch_unwind`. Generally a panic is not caught and it causes the process to exit. So an `Allocator` can be leaked, but only for the few milliseconds before the process exits, and it's freed again.
#### 3. `panic = "abort"`
Oxc's apps (Oxlint, Oxfmt) and NAPI packages (`oxc-parser` etc) are compiled with `panic = "abort"`.
With `panic = "abort"`, panics are converted to process abort. In that case, the compiler sees that the guard can never fire. It removes the guard and all the code which creates `Allocator`s and allocates dummies into them. So all these mechanics are completely zero cost, and no leak ever occurs.
9fc4b0d to
7db7a29
Compare
Implement `ReplaceWith` on all AST types which are not `Copy`. See #24012.
Replace `take_in` with `replace_with` in as many places as possible in transformer. See #24012 for explanation of the advantage of `replace_with`. This has only a minimal perf effect according to Codspeed benchmarks (1% improvement). But it does cut the number of arena allocations on TS files by 50%, so reduces memory usage.
# Oxlint ### 🚀 Features - 7db7a29 allocator: Add `ReplaceWith` trait (#24012) (overlookmotel) - a2c97f3 linter/unicorn: Implement `explicit-timer-delay` rule (#23612) (Mikhail Baev) - 85735cb linter/unicorn: Implement `no-confusing-array-with` rule (#23638) (Shekhu☺️ ) - cb4fbb9 linter/eslint: Implement no-unreachable-loop rule (#23975) (Todor Andonov) - dc32112 linter/eslint/no-constant-binary-expression: Check relational comparisons (#24088) (camc314) - 439c344 linter/jsdoc: Added missing options to `jsdoc/require-param` rule (#23364) (kapobajza) - 62af717 linter/unicorn/filename-case: Add `lowercase` and `screamingSnakeCase` (#24045) (Boshen) - d963967 linter/unicorn/no-array-sort: Add `allowAfterSpread` option (#24043) (Boshen) - 0a75682 linter: Add per-rule timings for type-aware linting (#22488) (camchenry) - 743e222 linter/react: Add `disallowedValues` option for `forbid-dom-props` rule (#23970) (Mikhail Baev) ### 🐛 Bug Fixes - 7b80010 linter: Use direct binding symbol ids (#24216) (camc314) - 8f94b49 linter/import/no-duplicates: Don't flag a type-only import beside a side-effect import (#24030) (Boshen) - d8c3fee linter/react/rules-of-hooks: Flag `useEffectEvent` escapes (#23764) (Rayan Salhab) - 0a7312b linter/no-deprecated-functions: Map `require.requireActual` to `jest.requireActual` (#23627) (Jerry Zhao) - d9e3ab3 linter/eslint/no-useless-return: Handle switch case continuation (#23984) (camc314) - 0b25582 ast: Type binding node `typeAnnotation` as `TSTypeAnnotation | null` (#23113) (Boshen) - 122d112 linter/eslint/no-restricted-imports: Flag dynamic import() expressions (#24029) (Boshen) - 59b6b83 linter: Avoid `OnceLock` re-entry on cyclic `export *` re-exports (#23632) (Jerry Zhao) - dd09af0 linter/import/namespace: Avoid panic on destructuring of an unresolvable namespace re-export (#23626) (Jerry Zhao) - bdb51c7 linter/jest/prefer-ending-with-an-expect: Validate config patterns (#24122) (camc314) - e383843 linter/unicorn/prefer-modern-dom-apis: Skip fixer for non identifier arguments (#23630) (Jerry Zhao) - 0ac4c83 linter: Detect circular config extends (#24115) (camc314) - bae1edf linter/import/namespace: Check namespace imports after named imports (#24094) (camc314) - cd8fdfe linter/eslint/no-eval: Recognize Array.from family thisArg (#24091) (camc314) - 851ee43 linter/eslint/no-eval: Resolve this binding for functions returned from an IIFE (#23643) (Jerry Zhao) - 002ab35 linter/unicorn: Avoid prefer-array-find rest destructuring false positive (#23654) (ColemanDunn) - 01c8775 linter/unicorn/filename-case: Keep digits attached in screamingSnakeCase (#24056) (Boshen) - f256941 linter: Recognize `@effect/vitest` as a vitest import source (#24025) (Boshen) - 73eeb1d linter/import/extensions: Honor per-extension `never` for explicit extensions (#24031) (Boshen) - d4ebe1f linter: Reject non-object oxlint config files (#24026) (Boshen) - 45d607d linter/react/forbid-component-props: Make allow/disallow lists optional in schema (#24024) (Boshen) - 54076ad linter/unicorn/no-array-for-each: Suggest entries loop for index callbacks (#24004) (camc314) - d057736 linter/jsdoc: Avoid param root underflow (#23945) (camc314) - 29c76bf linter/unicorn/prefer-at: Skip object numeric-key access (#23909) (Gaurav Dubey) ### ⚡ Performance - 657a8fc linter/oxc/bad-array-method-on-arguments: Only run on member expressions instead of all identifiers (#24164) (camchenry) - 073d9e7 linter/eslint/prefer-rest-params: Run on functions instead of all identifiers (#24163) (camchenry) - e5a4162 linter/jest/no-confusing-set-timeout: Early exit fast path (#24092) (camc314) - bca7ce5 linter: Only run react-perf rules on JSX attribute nodes (#24083) (camchenry) - 6881bf6 linter: Compute `apply_overrides` rule set lazily (#23648) (Jerry Zhao) - 911c106 linter/eslint/no-obj-calls: Use resolved reference instead of scope walk (#23895) (Marius Schulz) - dc8fd9a linter/unicorn/prefer-dom-node-text-content: Change dispatch to run only on less common node types (#23897) (Connor Shea) - fdbd34d linter/eslint/no-useless-call: Fast-path static callees (#24077) (camc314) - b1be114 linter/import/extensions: Skip empty config and borrow extensions (#24075) (camc314) - 4781b2d linter/eslint/no-obj-calls: Use direct global matches (#24076) (camc314) - e6cee89 linter: Avoid node-chain allocation for non-Jest calls (#23907) (Yagiz Nizipli) - 30dc517 linter/typescript/no-restricted-types: O(1) banned-type lookups (#23827) (Yagiz Nizipli) ### 📚 Documentation - 6ca9125 linter/typescript: Clarify consistent-type-imports behavior (#23972) (camc314) # Oxfmt ### 🚀 Features - 4f4313e formatter_css: Update oxc-css-parser 0.0.5 (#24120) (leaysgur) - 0ccd8a1 formatter_graphql: Update oxc-graphql-parser 0.0.5 (#24106) (leaysgur) - 89ec3d9 formatter_core: Add literal line and root indention primitives (#24051) (leaysgur) - 213a96b formatter_core: Add no-expand-parent for multiline text (#24050) (leaysgur) - 0e5bcc9 formatter_graphql: Update oxc-graphql-parser 0.0.4 (#24039) (leaysgur) - e0b35a1 formatter_css: Update `[email protected]` (#23974) (leaysgur) ### 🐛 Bug Fixes - 1fe6546 formatter: Omit unneeded `;` for type members with `no-semi` (#24212) (leaysgur) - 0ad7316 formatter: Print space for `ForStatement`.`update` only if exists (#24211) (leaysgur) - 3abbed5 formatter: Print `;` before jsdoc type-cast parens with no-semi (#24208) (leaysgur) - 9af3833 formatter_css: Make scss formatter consistent (#24207) (leaysgur) - 46d7194 formatter_css: Use fill IR for `@forward` members (#24206) (leaysgur) - e31038f formatter_css: Keep comment inside sass config list (#24205) (leaysgur) - d3b9591 formatter: Add parens around `await/yield` with `<T>` (#24202) (leaysgur) - 2121a55 oxfmt: Reuse tinypool process during the same LSP process (#24197) (leaysgur) - 9bf4b4a formatter_css: Align CSS output to Prettier 3.9.1 (#24100) (leaysgur) - cd2452e formatter_css: Align SCSS output to Prettier 3.9.1 (#24097) (leaysgur) - 4ee8745 formatter_css: Keep selector value contain line-break without breaking line (#24055) (leaysgur) - e1ece97 formatter_graphql: Break `implements` list by print-width (#23997) (leaysgur) - 0a6b16c formatter_json: Preserve key and literal value for json-stringify (#23996) (leaysgur) - 903ab6e formatter_css: Preserve newlines in css-in-js selector list (#23992) (leaysgur) - ea5d095 oxfmt: Update `--migrate prettier` (#23963) (leaysgur) ### ⚡ Performance - 468e1e3 formatter_core: Make printer queues cursor-based (#24098) (Boshen) - c59f2fe rust: Return impl ExactSizeIterator from slice-backed accessors (#24144) (Boshen) - c292fb2 formatter: Inline fits element dispatcher (#23982) (camc314) Co-authored-by: Boshen <[email protected]>
### 🚀 Features - 260425f semantic/examples: Include unresolved references (#24214) (camc314) - 2d9b0b3 minifier: Fold boolean-literal ternary branches in value contexts (#24110) (Dunqing) - 61fbf10 ast: Implement `ReplaceWith` on all AST types (#24013) (overlookmotel) - 7db7a29 allocator: Add `ReplaceWith` trait (#24012) (overlookmotel) - 4eb074e mangler: Add `reserved` option for names that must not be mangled (#24041) (Dunqing) - 2e62012 data_structures: Add `StringExt` trait (#24006) (overlookmotel) - 60e7160 minifier: Drop side-effect-free IIFEs whose result is unused (#23967) (Dunqing) - 26dd9e2 ast: Add method to widen inherited enum ref to parent ref (#23961) (overlookmotel) ### 🐛 Bug Fixes - e8b50ee transformer: Clean up semantics for stripped TypeScript syntax (#24180) (camc314) - d966d0b react_compiler: Remove clippy allows (#24168) (Boshen) - 854ef8d react_compiler: Compile generic functions instead of over-bailing on type-param hoisting (#24158) (Boshen) - 093586c react_compiler: Align memoization cache-slot allocation with Babel (#24157) (Boshen) - 09c8f59 react_compiler: Normalize snapshot fixture paths (#24142) (camc314) - f13df97 react_compiler: Drop stray empty statement from catch bindings (#24133) (Boshen) - cb2a505 react_compiler: Codegen destructuring reassignment targets (#24131) (Boshen) - b82c394 react_compiler: Propagate codegen invariants instead of emitting empty bodies (#24128) (Boshen) - 5771982 react_compiler: Render unchanged programs as source in fixture snapshots (#24129) (Boshen) - 4b16e1a transformer/async-to-generator: Preserve direct eval scope flags (#24136) (camc314) - 4e9194f react_compiler: Lower `delete obj.prop` to Property/ComputedDelete (#24123) (Boshen) - 0b25582 ast: Type binding node `typeAnnotation` as `TSTypeAnnotation | null` (#23113) (Boshen) - 018c0e5 transformer: Hoist lowered async declarations (#22770) (camc314) - 652fbaf mangler: Keep names of destructured exported bindings (#24036) (Dunqing) - e274415 minifier: Don't drop global calls that throw despite pure arguments (#23917) (Dunqing) - 59abb30 minifier: Only merge string literals in `try_fold_add` when the inner operator is `+` (#23622) (Jerry Zhao) ### ⚡ Performance - c5ca77b transformer: Avoid cloning refresh options (#24191) (camc314) - bf1a151 react_compiler: Compile out debug printers (#24184) (Boshen) - abb44a0 transformer: Build fixed object-rest arguments (#24190) (camc314) - a4db731 isolated_declarations: Use `ReplaceWith` instead of `TakeIn` (#24016) (overlookmotel) - ff10855 transformer: Use `ReplaceWith` instead of `TakeIn` (#24015) (overlookmotel) - bd49aff ecmascript: Avoid heap-allocating Math.min/max/imul operands (#23941) (Lawrence Lin) - e4b708b react_compiler: Skip compiled files before prefilters (#24171) (Boshen) - c59f2fe rust: Return impl ExactSizeIterator from slice-backed accessors (#24144) (Boshen) - 5d6d04a codegen: SWAR-skip boring byte runs in sourcemap line/column scan (#24023) (Boshen) - a55e0be traverse: Reduce string operations in `get_var_name_from_node` (#24007) (overlookmotel) - e6d48e1 transformer/nullish_coalescing: Move cold path into separate function (#23989) (overlookmotel) - c4e35b5 transformer/object_rest_spread: Pre-allocate capacity in `Vec` (#23988) (overlookmotel) - 527b8e5 transformer/decorators: Narrow type earlier (#23987) (overlookmotel) ### 📚 Documentation - 30d17f5 allocator: Clarify docs for `TakeIn::take_in_box` (#24093) (overlookmotel) - 675e6a8 ast: Correct doc comment for `PrivateFieldExpression` (#24008) (overlookmotel) - e4c30e6 minifier: Explain what `dce` mode means (#23994) (Dunqing) - 37cbf88 ast_macros: Document fields of `StructDetails` (#23959) (overlookmotel) - 4de3e54 ast: Correct doc comment (#23948) (overlookmotel) Co-authored-by: Boshen <[email protected]>
Replace `take_in` with `replace_with` in as many places as possible in minifier. See #24012 for explanation of the advantage of `replace_with`. Likely we could use `replace_with` in many more places, if we adapted the methods through which most replacement operations flow (which also track "dirty" state). I've left this for now, as I don't really understand the minifier.
Replace one usage of `take_in` in parser with `replace_with`. See #24012 for explanation of the advantage of `replace_with`.

The problem
A common pattern we have when mutating an AST (e.g. in transformer) is:
Rust does not make this easy, due to aliasing rules. Currently we do something like this (example from arrow function transform):
The problem with this is that
take_inwrites a dummyExpressioninto the AST, in order to get the ownedExpression. This is pointless, as the wholeStatementthat contains that dummyExpressionis overwritten in the next line - the dummy is immediately disconnected from the AST. But that dummy node lives on forever in the arena, consuming memory.Solution in this PR
Introduce
ReplaceWithtrait which replaces the node in place, without writing a dummy:The closure receives an owned copy of the
Statement, and can do whatever it needs to with it. The closure has to return a newStatementwhich is written to the "hole" the original came from.This PR only adds the trait. Later PRs in this stack implement the trait on AST types, and then use it in various crates.
Implementation
The difficulty is in making this sound even if the closure panics, and
std::panic::catch_unwindis in use, allowing the old value to be observed.A naive implementation would allow UB as follows:
The solution is to make it possible to implement
ReplaceWithonly on types that also implementDummy.If the closure panics, a guard is triggered which writes a dummy into the original slot which the type was moved out of. This happens before any other code can observe the value in the slot.
To keep the API ergonomic, we don't want to have to pass in an
&Allocatortoreplace_with. But the dummy has to be allocated somewhere. So we use a global store ofAllocators, stored in astatic, which live for the entire life of the process. When a panic occurs, a newAllocatoris created, and the dummy is allocated in it.This is a memory leak - the
Allocatorcontaining the dummy will never be freed until the process exits.However, in practice this is fine, because:
1. It shouldn't happen
The guard only triggers, and these leaked
Allocators are only created, if the closure panics - which it never should. If the closure does panic, that's a bug, and we'll aim to fix it swiftly.2. Panic usually exits process
It's unusual to use
catch_unwind. Generally a panic is not caught and it causes the process to exit. So anAllocatorcan be leaked, but only for the few milliseconds before the process exits, and it's freed again.3.
panic = "abort"Oxc's apps (Oxlint, Oxfmt) and NAPI packages (
oxc-parseretc) are compiled withpanic = "abort".With
panic = "abort", panics are converted to process abort. In that case, the compiler sees that the guard can never fire. It removes the guard and all the code which createsAllocators and allocates dummies into them. So all these mechanics are completely zero cost, and no leak ever occurs.