Allocate token's content into an arena

commit-id:5e4b122f

Cargo.toml

@@ -35,6 +35,7 @@ anyhow = "1"
 assert_fs = "1"
 async-trait = "0.1"
 axum = { version = "0.6", features = ["http2"] }
+bumpalo = "3"
 cairo-lang-casm = { git = "https://github.com/starkware-libs/cairo", rev = "5e447a6bec62eadc97af00d84b74e3b435bff723" }
 cairo-lang-compiler = { git = "https://github.com/starkware-libs/cairo", rev = "5e447a6bec62eadc97af00d84b74e3b435bff723" }
 cairo-lang-debug = { git = "https://github.com/starkware-libs/cairo", rev = "5e447a6bec62eadc97af00d84b74e3b435bff723" }

plugins/cairo-lang-macro-stable/src/lib.rs

@@ -9,7 +9,8 @@ pub mod ffi;
 #[derive(Debug)]
 pub struct StableToken {
     pub span: StableTextSpan,
-    pub content: *mut c_char,
+    pub ptr: *const u8,
+    pub len: usize,
 }
 
 #[repr(C)]
@@ -44,6 +45,7 @@ pub type StableExpansionsList = StableSlice<StableExpansion>;
 pub struct StableTokenStream {
     pub tokens: StableSlice<StableTokenTree>,
     pub metadata: StableTokenStreamMetadata,
+    pub size_hint: usize,
 }
 
 /// Token stream metadata.

plugins/cairo-lang-macro/Cargo.toml

@@ -13,7 +13,8 @@ readme = "README.md"
 repository.workspace = true
 
 [dependencies]
-cairo-lang-macro-attributes = "0.1"
+bumpalo.workspace = true
+cairo-lang-macro-attributes = { path = "../cairo-lang-macro-attributes" }
 cairo-lang-macro-stable = { path = "../cairo-lang-macro-stable" }
 linkme.workspace = true
 serde = { workspace = true, optional = true }

plugins/cairo-lang-macro/src/lib.rs

@@ -18,17 +18,22 @@
 pub use cairo_lang_macro_attributes::*;
 #[doc(hidden)]
 pub use linkme;
+use std::cell::RefCell;
 
 use cairo_lang_macro_stable::ffi::StableSlice;
 use cairo_lang_macro_stable::{
     StableExpansionsList, StablePostProcessContext, StableProcMacroResult,
 };
 use std::ffi::{c_char, CStr, CString};
+use std::ops::Deref;
 
 mod types;
 
 pub use types::*;
 
+// A thread-local allocation context for allocating tokens on proc macro side.
+thread_local!(static CONTEXT: RefCell<AllocationContext> =  RefCell::default() );
+
 #[doc(hidden)]
 #[derive(Clone)]
 pub struct ExpansionDefinition {
@@ -97,29 +102,42 @@ pub unsafe extern "C" fn expand(
     stable_attr: cairo_lang_macro_stable::StableTokenStream,
     stable_token_stream: cairo_lang_macro_stable::StableTokenStream,
 ) -> cairo_lang_macro_stable::StableResultWrapper {
-    let token_stream = TokenStream::from_stable(&stable_token_stream);
-    let attr_token_stream = TokenStream::from_stable(&stable_attr);
-    let item_name = CStr::from_ptr(item_name).to_string_lossy().to_string();
-    let fun = MACRO_DEFINITIONS_SLICE
-        .iter()
-        .find_map(|m| {
-            if m.name == item_name.as_str() {
-                Some(m.fun.clone())
-            } else {
-                None
-            }
-        })
-        .expect("procedural macro not found");
-    let result = match fun {
-        ExpansionFunc::Attr(fun) => fun(attr_token_stream, token_stream),
-        ExpansionFunc::Other(fun) => fun(token_stream),
-    };
-    let result: StableProcMacroResult = result.into_stable();
-    cairo_lang_macro_stable::StableResultWrapper {
-        input: stable_token_stream,
-        input_attr: stable_attr,
-        output: result,
-    }
+    CONTEXT.with(|ctx_cell| {
+        // Read size hint from stable token stream. This will be used to create a sufficiently
+        // large bump allocation buffer.
+        let size_hint: usize = stable_token_stream.size_hint + stable_attr.size_hint;
+        // Replace the allocation context with a new one.
+        // If there is no interned string guards, the old context will be de-allocated.
+        ctx_cell.replace(AllocationContext::with_capacity(size_hint));
+        let ctx_borrow = ctx_cell.borrow();
+        let ctx: &AllocationContext = ctx_borrow.deref();
+        // Copy the stable token stream into current context.
+        let token_stream = TokenStream::from_stable_in(&stable_token_stream, ctx);
+        let attr_token_stream = TokenStream::from_stable_in(&stable_attr, ctx);
+        let item_name = CStr::from_ptr(item_name)
+            .to_str()
+            .expect("item name must be a valid string");
+        let fun = MACRO_DEFINITIONS_SLICE
+            .iter()
+            .find_map(|m| {
+                if m.name == item_name {
+                    Some(m.fun.clone())
+                } else {
+                    None
+                }
+            })
+            .expect("procedural macro not found");
+        let result = match fun {
+            ExpansionFunc::Attr(fun) => fun(attr_token_stream, token_stream),
+            ExpansionFunc::Other(fun) => fun(token_stream),
+        };
+        let result: StableProcMacroResult = result.into_stable();
+        cairo_lang_macro_stable::StableResultWrapper {
+            input: stable_token_stream,
+            input_attr: stable_attr,
+            output: result,
+        }
+    })
 }
 
 /// Free the memory allocated for the [`StableProcMacroResult`].

plugins/cairo-lang-macro/src/types/conversion.rs

@@ -1,6 +1,7 @@
 use crate::{
-    AuxData, Diagnostic, ExpansionDefinition, FullPathMarker, PostProcessContext, ProcMacroResult,
-    Severity, TextSpan, Token, TokenStream, TokenStreamMetadata, TokenTree,
+    AllocationContext, AuxData, Diagnostic, ExpansionDefinition, FullPathMarker,
+    PostProcessContext, ProcMacroResult, Severity, TextSpan, Token, TokenStream,
+    TokenStreamMetadata, TokenTree,
 };
 use cairo_lang_macro_stable::ffi::StableSlice;
 use cairo_lang_macro_stable::{
@@ -44,6 +45,7 @@ impl ProcMacroResult {
     /// # Safety
     #[doc(hidden)]
     pub unsafe fn from_stable(result: &StableProcMacroResult) -> Self {
+        let ctx = AllocationContext::with_capacity(result.token_stream.size_hint);
         let (ptr, n) = result.diagnostics.raw_parts();
         let diagnostics = slice::from_raw_parts(ptr, n)
             .iter()
@@ -55,7 +57,7 @@ impl ProcMacroResult {
             .map(|m| from_raw_cstr(*m))
             .collect::<Vec<_>>();
         ProcMacroResult {
-            token_stream: TokenStream::from_stable(&result.token_stream),
+            token_stream: TokenStream::from_stable_in(&result.token_stream, &ctx),
             diagnostics,
             full_path_markers,
             aux_data: AuxData::from_stable(&result.aux_data),
@@ -110,10 +112,12 @@ impl Token {
     /// Convert to FFI-safe representation.
     #[doc(hidden)]
     pub fn into_stable(self) -> StableToken {
-        let cstr = CString::new(self.content.as_bytes()).unwrap();
+        let ptr = self.content.as_ptr();
+        let len = self.content.len();
         StableToken {
             span: self.span.into_stable(),
-            content: cstr.into_raw(),
+            ptr,
+            len,
         }
     }
 
@@ -123,9 +127,11 @@ impl Token {
     ///
     /// # Safety
     #[doc(hidden)]
-    pub unsafe fn from_stable(token: &StableToken) -> Self {
+    pub unsafe fn from_stable_in(token: &StableToken, ctx: &AllocationContext) -> Self {
+        let content = slice::from_raw_parts(token.ptr, token.len);
+        let content = ctx.intern(std::str::from_utf8_unchecked(content));
         Self {
-            content: from_raw_cstr(token.content),
+            content,
             span: TextSpan::from_stable(&token.span),
         }
     }
@@ -138,7 +144,6 @@ impl Token {
     /// # Safety
     #[doc(hidden)]
     pub unsafe fn free_owned_stable(token: StableToken) {
-        free_raw_cstring(token.content);
         TextSpan::free_owned_stable(token.span);
     }
 }
@@ -158,9 +163,9 @@ impl TokenTree {
     ///
     /// # Safety
     #[doc(hidden)]
-    pub unsafe fn from_stable(token_tree: &StableTokenTree) -> Self {
+    pub unsafe fn from_stable_in(token_tree: &StableTokenTree, ctx: &AllocationContext) -> Self {
         match token_tree {
-            StableTokenTree::Ident(token) => Self::Ident(Token::from_stable(token)),
+            StableTokenTree::Ident(token) => Self::Ident(Token::from_stable_in(token, ctx)),
         }
     }
 
@@ -186,14 +191,19 @@ impl TokenStream {
     /// # Safety
     #[doc(hidden)]
     pub fn into_stable(self) -> StableTokenStream {
+        let mut size_hint: usize = 0;
         let tokens = self
             .tokens
             .into_iter()
-            .map(|token| token.into_stable())
+            .map(|token| {
+                size_hint += token.size_hint();
+                token.into_stable()
+            })
             .collect::<Vec<_>>();
         StableTokenStream {
             tokens: StableSlice::new(tokens),
             metadata: self.metadata.into_stable(),
+            size_hint,
         }
     }
 
@@ -203,11 +213,14 @@ impl TokenStream {
     ///
     /// # Safety
     #[doc(hidden)]
-    pub unsafe fn from_stable(token_stream: &StableTokenStream) -> Self {
+    pub unsafe fn from_stable_in(
+        token_stream: &StableTokenStream,
+        ctx: &AllocationContext,
+    ) -> Self {
         let (ptr, n) = token_stream.tokens.raw_parts();
         let tokens = slice::from_raw_parts(ptr, n)
             .iter()
-            .map(|token_tree| TokenTree::from_stable(token_tree))
+            .map(|token_tree| TokenTree::from_stable_in(token_tree, ctx))
             .collect::<Vec<_>>();
         Self {
             tokens,

plugins/cairo-lang-macro/src/types/mod.rs

@@ -44,7 +44,7 @@ pub struct ProcMacroResult {
 ///     let token_stream = TokenStream::new(vec![
 ///         TokenTree::Ident(
 ///             Token::new(
-///                 code.clone(),
+///                 code.as_str(),
 ///                 TextSpan::new(0, code.len())
 ///             )
 ///         )