feat(span): fix memory leak by implementing inlineable string for oxc…

…_allocator (#2294)

closes #1803

This string is currently unsafe, but I want to get miri working before
introducing more changes.

I want to make a progress from memory leak to unsafe then to safety.
It's harder to do the steps in one go.

.github/workflows/miri.yml

@@ -2,6 +2,17 @@ name: Miri
 
 on:
   workflow_dispatch:
+  pull_request:
+    types: [opened, synchronize]
+    paths:
+      - 'crates/oxc_parser/**'
+      - '.github/workflows/miri.yml'
+  push:
+    branches:
+      - main
+    paths:
+      - 'crates/oxc_parser/**'
+      - '.github/workflows/miri.yml'
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
@@ -15,6 +26,12 @@ jobs:
       - name: Checkout
         uses: actions/checkout@v4
 
+      - name: Install Rust
+        uses: ./.github/actions/rustup
+        with:
+          shared-key: miri
+          save-cache: ${{ github.ref_name == 'main' }}
+
       - name: Install Miri
         run: |
           rustup toolchain install nightly --component miri

Cargo.toml

@@ -88,7 +88,6 @@ assert-unchecked          = { version = "0.1.2" }
 bpaf                      = { version = "0.9.9" }
 bitflags                  = { version = "2.4.2" }
 bumpalo                   = { version = "3.14.0" }
-compact_str               = { version = "0.7.1" }
 convert_case              = { version = "0.6.0" }
 criterion                 = { version = "0.5.1", default-features = false }
 crossbeam-channel         = { version = "0.5.11" }

crates/oxc_semantic/src/module_record/builder.rs

@@ -283,11 +283,10 @@ impl ModuleRecordBuilder {
         };
         let export_entry = ExportEntry {
             export_name: ExportExportName::Default(exported_name.span()),
-            local_name: id
-                .as_ref()
-                .map_or(ExportLocalName::Default(exported_name.span()), |ident| {
-                    ExportLocalName::Name(NameSpan::new(ident.name.clone(), ident.span))
-                }),
+            local_name: id.as_ref().map_or_else(
+                || ExportLocalName::Default(exported_name.span()),
+                |ident| ExportLocalName::Name(NameSpan::new(ident.name.clone(), ident.span)),
+            ),
             span: decl.declaration.span(),
             ..ExportEntry::default()
         };

crates/oxc_span/Cargo.toml

@@ -18,16 +18,16 @@ workspace = true
 [lib]
 doctest = false
 
-[features]
-default = []
-serde   = ["dep:serde", "compact_str/serde"]
-wasm    = ["dep:tsify", "dep:wasm-bindgen"]
-
 [dependencies]
-miette      = { workspace = true }
-compact_str = { workspace = true }
+miette           = { workspace = true }
+inlinable_string = "0.1.15"
 
 tsify        = { workspace = true, optional = true }
 wasm-bindgen = { workspace = true, optional = true }
 
 serde = { workspace = true, features = ["derive"], optional = true }
+
+[features]
+default = []
+serde   = ["dep:serde", "inlinable_string/serde"]
+wasm    = ["dep:tsify", "dep:wasm-bindgen"]

crates/oxc_span/src/atom.rs

@@ -1,17 +1,15 @@
 use std::{
     borrow::{Borrow, Cow},
-    fmt,
+    fmt, hash,
     ops::Deref,
 };
 
-use compact_str::CompactString;
 #[cfg(feature = "serde")]
-use serde::Serialize;
+use serde::{Serialize, Serializer};
 
-/// Newtype for [`CompactString`]
-#[derive(Clone, Eq, Hash)]
-#[cfg_attr(feature = "serde", derive(Serialize))]
-pub struct Atom(CompactString);
+use inlinable_string::inline_string::{InlineString, INLINE_STRING_CAPACITY};
+
+const BASE54_CHARS: &[u8; 64] = b"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ$_0123456789";
 
 #[cfg_attr(
     all(feature = "serde", feature = "wasm"),
@@ -22,16 +20,57 @@ const TS_APPEND_CONTENT: &'static str = r#"
 export type Atom = string;
 "#;
 
-const BASE54_CHARS: &[u8; 64] = b"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ$_0123456789";
+/// An inlinable string for oxc_allocator.
+///
+/// SAFETY: It is unsafe to use this string after the allocator is dropped.
+///
+#[derive(Clone, Eq, Hash)]
+pub struct Atom(AtomImpl);
+
+/// Immutable Inlinable String
+///
+/// https://github.com/fitzgen/inlinable_string/blob/master/src/lib.rs
+#[derive(Clone, Eq, PartialEq)]
+enum AtomImpl {
+    /// A arena heap-allocated string.
+    Arena(&'static str),
+    /// A heap-allocated string.
+    Heap(Box<str>),
+    /// A small string stored inline.
+    Inline(InlineString),
+}
+
+#[cfg(feature = "serde")]
+impl Serialize for Atom {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        serializer.serialize_str(self.as_str())
+    }
+}
 
 impl Atom {
-    pub const fn new_inline(s: &str) -> Self {
-        Self(CompactString::new_inline(s))
+    pub fn new_inline(s: &str) -> Self {
+        Self(AtomImpl::Inline(InlineString::from(s)))
     }
 
     #[inline]
     pub fn as_str(&self) -> &str {
-        self.0.as_str()
+        match &self.0 {
+            AtomImpl::Arena(s) => s,
+            AtomImpl::Heap(s) => s,
+            AtomImpl::Inline(s) => s.as_ref(),
+        }
+    }
+
+    #[inline]
+    pub fn into_string(self) -> String {
+        match self.0 {
+            AtomImpl::Arena(s) => String::from(s),
+            AtomImpl::Heap(s) => s.to_string(),
+            AtomImpl::Inline(s) => s.to_string(),
+        }
     }
 
     /// Get the shortest mangled name for a given n.
@@ -41,7 +80,7 @@ impl Atom {
         // Base 54 at first because these are the usable first characters in JavaScript identifiers
         // <https://tc39.es/ecma262/#prod-IdentifierStart>
         let base = 54usize;
-        let mut ret = CompactString::default();
+        let mut ret = String::new();
         ret.push(BASE54_CHARS[num % base] as char);
         num /= base;
         // Base 64 for the rest because after the first character we can also use 0-9 too
@@ -52,70 +91,94 @@ impl Atom {
             ret.push(BASE54_CHARS[num % base] as char);
             num /= base;
         }
-        Self(ret)
-    }
-}
-
-impl Deref for Atom {
-    type Target = str;
-
-    fn deref(&self) -> &Self::Target {
-        self.0.as_str()
+        Self(AtomImpl::Heap(ret.into_boxed_str()))
     }
 }
 
 impl<'a> From<&'a str> for Atom {
     fn from(s: &'a str) -> Self {
-        Self(s.into())
+        if s.len() <= INLINE_STRING_CAPACITY {
+            Self(AtomImpl::Inline(InlineString::from(s)))
+        } else {
+            // SAFETY: It is unsafe to use this string after the allocator is dropped.
+            Self(AtomImpl::Arena(unsafe { std::mem::transmute(s) }))
+        }
     }
 }
 
 impl From<String> for Atom {
     fn from(s: String) -> Self {
-        Self(s.into())
+        if s.len() <= INLINE_STRING_CAPACITY {
+            Self(AtomImpl::Inline(InlineString::from(s.as_str())))
+        } else {
+            Self(AtomImpl::Heap(s.into_boxed_str()))
+        }
     }
 }
 
 impl From<Cow<'_, str>> for Atom {
     fn from(s: Cow<'_, str>) -> Self {
-        Self(s.into())
+        if s.len() <= INLINE_STRING_CAPACITY {
+            Self(AtomImpl::Inline(InlineString::from(s.borrow())))
+        } else {
+            Self(AtomImpl::Heap(s.into()))
+        }
+    }
+}
+
+impl Deref for Atom {
+    type Target = str;
+
+    fn deref(&self) -> &Self::Target {
+        self.as_str()
     }
 }
 
 impl AsRef<str> for Atom {
     #[inline]
     fn as_ref(&self) -> &str {
-        self.0.as_str()
+        self.as_str()
     }
 }
 
 impl Borrow<str> for Atom {
     #[inline]
     fn borrow(&self) -> &str {
-        self.0.as_str()
+        self.as_str()
     }
 }
 
 impl<T: AsRef<str>> PartialEq<T> for Atom {
     fn eq(&self, other: &T) -> bool {
-        self.0.as_str() == other.as_ref()
+        self.as_str() == other.as_ref()
     }
 }
 
 impl PartialEq<Atom> for &str {
     fn eq(&self, other: &Atom) -> bool {
-        *self == other.0.as_str()
+        *self == other.as_str()
+    }
+}
+
+impl hash::Hash for AtomImpl {
+    #[inline]
+    fn hash<H: hash::Hasher>(&self, hasher: &mut H) {
+        match self {
+            Self::Arena(s) => s.hash(hasher),
+            Self::Heap(s) => s.hash(hasher),
+            Self::Inline(s) => s.hash(hasher),
+        }
     }
 }
 
 impl fmt::Debug for Atom {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        fmt::Debug::fmt(self.0.as_str(), f)
+        fmt::Debug::fmt(self.as_str(), f)
     }
 }
 
 impl fmt::Display for Atom {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        fmt::Display::fmt(self.0.as_str(), f)
+        fmt::Display::fmt(self.as_str(), f)
     }
 }