feat(napi/oxlint): pass AST in buffer to JS

Author: overlookmotel

.github/generated/ast_changes_watch_list.yml

@@ -63,6 +63,8 @@ src:
   - 'crates/oxc_syntax/src/serialize.rs'
   - 'crates/oxc_syntax/src/symbol.rs'
   - 'crates/oxc_traverse/src/generated/scopes_collector.rs'
+  - 'napi/oxlint2/src/generated/constants.cjs'
+  - 'napi/oxlint2/src/generated/raw_transfer_constants.rs'
   - 'napi/parser/generated/constants.js'
   - 'napi/parser/generated/deserialize/js.js'
   - 'napi/parser/generated/deserialize/ts.js'

Cargo.lock

@@ -9,28 +9,30 @@ use crate::*;
 
 #[cfg(target_pointer_width = "64")]
 const _: () = {
-    // Padding: 6 bytes
+    // Padding: 2 bytes
     assert!(size_of::<RawTransferMetadata2>() == 32);
     assert!(align_of::<RawTransferMetadata2>() == 8);
     assert!(offset_of!(RawTransferMetadata2, data_offset) == 16);
-    assert!(offset_of!(RawTransferMetadata2, is_ts) == 24);
-    assert!(offset_of!(RawTransferMetadata2, id) == 20);
-    assert!(offset_of!(RawTransferMetadata2, can_be_freed) == 25);
-    assert!(offset_of!(RawTransferMetadata2, alloc_ptr) == 8);
-    assert!(offset_of!(RawTransferMetadata2, _padding) == 0);
+    assert!(offset_of!(RawTransferMetadata2, is_ts) == 28);
+    assert!(offset_of!(RawTransferMetadata2, source_len) == 20);
+    assert!(offset_of!(RawTransferMetadata2, id) == 24);
+    assert!(offset_of!(RawTransferMetadata2, can_be_freed) == 29);
+    assert!(offset_of!(RawTransferMetadata2, alloc_ptr) == 0);
+    assert!(offset_of!(RawTransferMetadata2, chunk_ptr) == 8);
 };
 
 #[cfg(target_pointer_width = "32")]
 const _: () = {
     // Padding: 2 bytes
     assert!(size_of::<RawTransferMetadata2>() == 24);
-    assert!(align_of::<RawTransferMetadata2>() == 8);
-    assert!(offset_of!(RawTransferMetadata2, data_offset) == 12);
+    assert!(align_of::<RawTransferMetadata2>() == 4);
+    assert!(offset_of!(RawTransferMetadata2, data_offset) == 8);
     assert!(offset_of!(RawTransferMetadata2, is_ts) == 20);
+    assert!(offset_of!(RawTransferMetadata2, source_len) == 12);
     assert!(offset_of!(RawTransferMetadata2, id) == 16);
     assert!(offset_of!(RawTransferMetadata2, can_be_freed) == 21);
-    assert!(offset_of!(RawTransferMetadata2, alloc_ptr) == 8);
-    assert!(offset_of!(RawTransferMetadata2, _padding) == 0);
+    assert!(offset_of!(RawTransferMetadata2, alloc_ptr) == 0);
+    assert!(offset_of!(RawTransferMetadata2, chunk_ptr) == 4);
 };
 
 #[cfg(not(any(target_pointer_width = "64", target_pointer_width = "32")))]

crates/oxc_allocator/src/generated/assert_layouts.rs

@@ -2,6 +2,7 @@
 // To edit this generated file you have to edit `tasks/ast_tools/src/generators/raw_transfer.rs`.
 
 #![expect(clippy::unreadable_literal)]
+#![allow(dead_code)]
 
 pub const BUFFER_SIZE: usize = 2147483632;
 pub const BUFFER_ALIGN: usize = 4294967296;

crates/oxc_allocator/src/generated/fixed_size_constants.rs

@@ -80,14 +80,25 @@ mod pool_fixed_size;
     target_endian = "little"
 ))]
 use pool_fixed_size as pool;
-// Import so that `generated/assert_layouts.rs` can access it
+// Import under original name so that `generated/assert_layouts.rs` can access it
 #[cfg(all(
     feature = "fixed_size",
     not(feature = "disable_fixed_size"),
     target_pointer_width = "64",
     target_endian = "little"
 ))]
 use pool_fixed_size::RawTransferMetadata2;
+// Export so can be used in `napi/oxlint2`
+#[cfg(all(
+    feature = "fixed_size",
+    not(feature = "disable_fixed_size"),
+    target_pointer_width = "64",
+    target_endian = "little"
+))]
+#[doc(hidden)]
+pub use pool_fixed_size::{
+    ALLOC_LAYOUT as FIXED_SIZE_ALLOC_LAYOUT, RawTransferMetadata2 as RawTransferMetadata,
+};
 
 pub use pool::{AllocatorGuard, AllocatorPool};
 

crates/oxc_allocator/src/lib.rs

@@ -113,8 +113,10 @@ pub struct RawTransferMetadata2 {
     pub data_offset: u32,
     /// `true` if AST is TypeScript.
     pub is_ts: bool,
+    /// Length of source text in UTF-8 bytes.
+    pub source_len: u32,
     /// Unique ID of this allocator.
-    pub(crate) id: u32,
+    pub id: u32,
     /// `true` if memory allocation backing the `FixedSizeAllocator` can be freed.
     /// Is `true` if allocator is currently owned by only Rust or JS, but not both.
     ///
@@ -124,25 +126,26 @@ pub struct RawTransferMetadata2 {
     ///   Memory will be freed when the `FixedSizeAllocator` is dropped on Rust side.
     /// * Also be set to `true` if `FixedSizeAllocator` is dropped on Rust side.
     ///   Memory will be freed in finalizer when JS garbage collector collects the buffer.
-    pub(crate) can_be_freed: AtomicBool,
+    pub can_be_freed: AtomicBool,
     /// Pointer to start of original allocation backing the `Allocator`.
-    pub(crate) alloc_ptr: NonNull<u8>,
-    /// Padding to pad struct to size 32.
-    pub(crate) _padding: u64,
+    pub alloc_ptr: NonNull<u8>,
+    /// Pointer to start of `Allocator` chunk.
+    pub chunk_ptr: NonNull<u8>,
 }
 use RawTransferMetadata2 as RawTransferMetadata;
 
 const METADATA_SIZE: usize = size_of::<RawTransferMetadata>();
 
 impl RawTransferMetadata {
-    fn new(id: u32, alloc_ptr: NonNull<u8>) -> Self {
+    fn new(id: u32, alloc_ptr: NonNull<u8>, chunk_ptr: NonNull<u8>) -> Self {
         Self {
             data_offset: 0,
             is_ts: false,
+            source_len: 0,
             id,
             can_be_freed: AtomicBool::new(true),
             alloc_ptr,
-            _padding: 0,
+            chunk_ptr,
         }
     }
 }
@@ -165,7 +168,8 @@ impl RawTransferMetadata {
 const ALLOC_SIZE: usize = BUFFER_SIZE + TWO_GIB;
 const ALLOC_ALIGN: usize = TWO_GIB;
 
-const ALLOC_LAYOUT: Layout = match Layout::from_size_align(ALLOC_SIZE, ALLOC_ALIGN) {
+/// Layout of backing allocations for fixed-size allocators.
+pub const ALLOC_LAYOUT: Layout = match Layout::from_size_align(ALLOC_SIZE, ALLOC_ALIGN) {
     Ok(layout) => layout,
     Err(_) => unreachable!(),
 };
@@ -222,7 +226,7 @@ impl FixedSizeAllocator {
         let allocator = unsafe { Allocator::from_raw_parts(chunk_ptr, CHUNK_SIZE) };
 
         // Store metadata after allocator chunk
-        let metadata = RawTransferMetadata::new(id, alloc_ptr);
+        let metadata = RawTransferMetadata::new(id, alloc_ptr, chunk_ptr);
         // SAFETY: `chunk_ptr` is at least `BUFFER_SIZE` bytes from the end of the allocation.
         // `CHUNK_SIZE` had the size of `RawTransferMetadata` subtracted from it.
         // So there is space within the allocation for `RawTransferMetadata` after the `Allocator`'s chunk.
@@ -251,13 +255,28 @@ impl FixedSizeAllocator {
 }
 
 impl Drop for FixedSizeAllocator {
+    // TODO: Delete logging
     fn drop(&mut self) {
         // Get pointer to start of allocation backing this `FixedSizeAllocator`
         let alloc_ptr = {
-            let metadata_ptr = self.allocator.end_ptr().cast::<RawTransferMetadata>();
             // SAFETY: `FixedSizeAllocator` is being dropped, so no other references to data in it may exist.
             // `FixedSizeAllocator::new` wrote `RawTransferMetadata` to the location pointed to by `end_ptr`.
-            let metadata = unsafe { metadata_ptr.as_ref() };
+            let metadata = unsafe { self.allocator.metadata() };
+
+            // If `can_be_freed` is already `true`, either this `Allocator` was never sent to JS side,
+            // or JS garbage collector already collected it. We can deallocate the memory.
+            // If not, set it to `true` and exit. Memory will be freed when JS garbage collector
+            // collects the buffer.
+            //
+            // Maybe a more relaxed `Ordering` would be OK, but I (@overlookmotel) am not sure,
+            // so going with `Ordering::SeqCst` to be on safe side.
+            // Deallocation only happens at the end of the whole process, so it shouldn't matter much.
+            // TODO: Figure out if can use `Ordering::Relaxed`.
+            let can_be_freed = metadata.can_be_freed.fetch_or(true, Ordering::SeqCst);
+            if !can_be_freed {
+                return;
+            }
+
             metadata.alloc_ptr
         };
 
@@ -269,3 +288,42 @@ impl Drop for FixedSizeAllocator {
 // SAFETY: `Allocator` is `Send`.
 // Moving `alloc_ptr: NonNull<u8>` across threads along with the `Allocator` is safe.
 unsafe impl Send for FixedSizeAllocator {}
+
+impl Allocator {
+    /// Get reference to [`RawTransferMetadata`] for this `Allocator`.
+    ///
+    /// # SAFETY
+    ///
+    /// * This [`Allocator`] must have been created by a fixed-size [`AllocatorPool`].
+    /// * No mutable references to [`RawTransferMetadata`] for this `Allocator` can exist
+    ///   while the reference returned by this method lives.
+    pub unsafe fn metadata(&self) -> &RawTransferMetadata {
+        // SAFETY:
+        // * Caller guarantees that this `Allocator` was created by a fixed-size `AllocatorPool`.
+        //   `FixedSizeAllocator::new` wrote `RawTransferMetadata` to the location pointed to by `end_ptr`.
+        // * Caller guarantees no mutable references to `RawTransferMetadata` exist
+        //   while this reference lives.
+        unsafe { self.end_ptr().cast::<RawTransferMetadata>().as_ref() }
+    }
+
+    /// Get mutable reference to [`RawTransferMetadata`] for this `Allocator`.
+    ///
+    /// # SAFETY
+    ///
+    /// * This [`Allocator`] must have been created by a fixed-size [`AllocatorPool`].
+    /// * No other references to [`RawTransferMetadata`] for this `Allocator` can exist
+    ///   while the reference returned by this method lives.
+    ///
+    /// Note: This means we must be sure that buffer cannot be garbage-collected on JS side,
+    /// as that creates a `&` ref to the `RawTransferMetadata`, and the finalizer that runs
+    /// when that happens can run on another thread.
+    #[expect(clippy::mut_from_ref)]
+    pub unsafe fn metadata_mut(&self) -> &mut RawTransferMetadata {
+        // SAFETY:
+        // * Caller guarantees that this `Allocator` was created by a fixed-size `AllocatorPool`.
+        //   `FixedSizeAllocator::new` wrote `RawTransferMetadata` to the location pointed to by `end_ptr`.
+        // * Caller guarantees no other references to `RawTransferMetadata` exist
+        //   while this reference lives.
+        unsafe { self.end_ptr().cast::<RawTransferMetadata>().as_mut() }
+    }
+}

crates/oxc_allocator/src/pool_fixed_size.rs

@@ -135,7 +135,7 @@ pub static STRUCTS: phf::Map<&'static str, StructDetails> = ::phf::Map {
         ("Decorator", StructDetails { field_order: None }),
         ("CharacterClass", StructDetails { field_order: Some(&[0, 2, 3, 4, 1]) }),
         ("TemplateElementValue", StructDetails { field_order: None }),
-        ("RawTransferMetadata", StructDetails { field_order: Some(&[2, 4, 3, 5, 1, 0]) }),
+        ("RawTransferMetadata", StructDetails { field_order: Some(&[2, 5, 3, 4, 6, 0, 1]) }),
         ("TSTypeParameter", StructDetails { field_order: None }),
         ("SourceType", StructDetails { field_order: None }),
         ("ErrorLabel", StructDetails { field_order: Some(&[1, 0]) }),
@@ -266,7 +266,7 @@ pub static STRUCTS: phf::Map<&'static str, StructDetails> = ::phf::Map {
         ("ImportAttribute", StructDetails { field_order: None }),
         ("TSConditionalType", StructDetails { field_order: None }),
         ("TSNamespaceExportDeclaration", StructDetails { field_order: None }),
-        ("RawTransferMetadata2", StructDetails { field_order: Some(&[2, 4, 3, 5, 1, 0]) }),
+        ("RawTransferMetadata2", StructDetails { field_order: Some(&[2, 5, 3, 4, 6, 0, 1]) }),
         ("AssignmentTargetWithDefault", StructDetails { field_order: None }),
         ("RegExpLiteral", StructDetails { field_order: None }),
         ("CapturingGroup", StructDetails { field_order: None }),

crates/oxc_ast_macros/src/generated/structs.rs

@@ -2,6 +2,8 @@ use std::{fmt::Debug, pin::Pin, sync::Arc};
 
 use serde::{Deserialize, Serialize};
 
+use oxc_allocator::Allocator;
+
 pub type ExternalLinterLoadPluginCb = Arc<
     dyn Fn(
             String,
@@ -17,7 +19,11 @@ pub type ExternalLinterLoadPluginCb = Arc<
 >;
 
 pub type ExternalLinterCb = Arc<
-    dyn Fn(String, Vec<u32>) -> Result<Vec<LintResult>, Box<dyn std::error::Error + Send + Sync>>
+    dyn Fn(
+            String,
+            Vec<u32>,
+            &Allocator,
+        ) -> Result<Vec<LintResult>, Box<dyn std::error::Error + Send + Sync>>
         + Sync
         + Send,
 >;

crates/oxc_linter/src/external_linter.rs

@@ -1,7 +1,7 @@
 #![expect(clippy::self_named_module_files)] // for rules.rs
 #![allow(clippy::literal_string_with_formatting_args)]
 
-use std::{path::Path, rc::Rc, sync::Arc};
+use std::{path::Path, ptr, rc::Rc, sync::Arc};
 
 use oxc_allocator::Allocator;
 use oxc_diagnostics::OxcDiagnostic;
@@ -118,7 +118,7 @@ impl Linter {
         path: &Path,
         semantic: Rc<Semantic<'a>>,
         module_record: Arc<ModuleRecord>,
-        _allocator: &Allocator,
+        allocator: &Allocator,
     ) -> Vec<Message<'a>> {
         let ResolvedLinterState { rules, config, external_rules } = self.config.resolve(path);
 
@@ -206,10 +206,29 @@ impl Linter {
 
         if !external_rules.is_empty() {
             if let Some(external_linter) = self.external_linter.as_ref() {
+                // Write offset of `Program` and source text length in metadata at end of buffer
+                #[expect(clippy::missing_panics_doc, reason = "program() always returns `Some`")]
+                let program = semantic.nodes().program().unwrap();
+                let program_offset = ptr::from_ref(program) as u32;
+                #[expect(clippy::cast_possible_truncation)]
+                let source_len = program.source_text.len() as u32;
+                {
+                    // SAFETY: TODO
+                    // Is this safe? If buffer is GC-ed on JS side, it would create a `&RawTransferMetadata`
+                    // which could alias this `&mut` ref. That shouldn't be possible as all buffers
+                    // are stored in an array on JS side, so they don't get GC-ed until process exits.
+                    // But this feels a bit danngerous.
+                    let metadata = unsafe { allocator.metadata_mut() };
+                    metadata.data_offset = program_offset;
+                    metadata.source_len = source_len;
+                }
+
+                // Pass AST and rule IDs to JS
                 let result = (external_linter.run)(
                     #[expect(clippy::missing_panics_doc)]
                     path.to_str().unwrap().to_string(),
                     external_rules.iter().map(|(rule_id, _)| rule_id.raw()).collect(),
+                    allocator,
                 );
                 match result {
                     Ok(diagnostics) => {

crates/oxc_linter/src/lib.rs

@@ -22,6 +22,7 @@ test = false
 doctest = false
 
 [dependencies]
+oxc_allocator = { workspace = true, features = ["fixed_size"] }
 oxlint = { workspace = true, features = ["oxlint2", "allocator"] }
 
 napi = { workspace = true, features = ["async"] }