Implement aux_data callback (#1186)

commit-id:a9ba1335

---

**Stack**:
- #1200
- #1195
- #1194
- #1192
- #1186 ⬅


⚠️ *Part of a stack created by [spr](https://github.com/ejoffe/spr). Do
not merge manually using the UI - doing so may have unexpected results.*

Cargo.toml

@@ -82,6 +82,7 @@ io_tee = "0.1"
 itertools = "0.12"
 libc = "0.2"
 libloading = "0.8.3"
+linkme = "0.3"
 log = "0.4"
 ntest = "0.9"
 num-bigint = { version = "0.4", features = ["rand"] }

plugins/cairo-lang-macro-attributes/Cargo.toml

@@ -18,3 +18,4 @@ proc-macro = true
 [dependencies]
 quote.workspace = true
 syn = { workspace = true, features = ["full", "extra-traits"] }
+scarb-stable-hash = { path = "../../utils/scarb-stable-hash" }

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

@@ -1,5 +1,6 @@
 use proc_macro::TokenStream;
 use quote::quote;
+use scarb_stable_hash::short_hash;
 use syn::{parse_macro_input, ItemFn};
 
 /// Inline macro helper.
@@ -28,3 +29,40 @@ pub fn attribute_macro(_args: TokenStream, input: TokenStream) -> TokenStream {
     };
     TokenStream::from(expanded)
 }
+
+/// This macro can be used to construct the auxiliary data collection callback.
+///
+/// The procedural macro can emit additional auxiliary data alongside the generated [`TokenStream`]
+/// during the code expansion. This data can be used to collect additional information from the
+/// source code of a project that is being compiled during the macro execution.
+/// For instance, you can create a procedural macro that collects some information stored by
+/// the Cairo programmer as attributes in the project source code.
+///
+/// This should be used to implement a collection callback for the auxiliary data.
+/// This callback will be called after the source code compilation (and thus after all the procedural
+/// macro executions). All auxiliary data emitted by the procedural macro during source code compilation
+/// will be passed to the callback as an argument.
+///
+/// The callback can be used to process or persist the data collected during the compilation.
+///
+/// This macro hides the conversion to stable ABI structs from the user.
+///
+/// # Safety
+#[proc_macro_attribute]
+pub fn aux_data_collection_callback(_args: TokenStream, input: TokenStream) -> TokenStream {
+    let mut item: ItemFn = parse_macro_input!(input as ItemFn);
+    // Rename item to hide it from the macro source code.
+    let id = short_hash(item.sig.ident.to_string());
+    let item_name = format!("{}_{}", item.sig.ident, id);
+    item.sig.ident = syn::Ident::new(item_name.as_str(), item.sig.ident.span());
+
+    let item_name = &item.sig.ident;
+    let expanded = quote! {
+        #item
+
+        #[::cairo_lang_macro::linkme::distributed_slice(::cairo_lang_macro::AUX_DATA_CALLBACKS)]
+        #[linkme(crate = ::cairo_lang_macro::linkme)]
+        static AUX_DATA_CALLBACK_DESERIALIZE: fn(Vec<AuxData>) = #item_name;
+    };
+    TokenStream::from(expanded)
+}

plugins/cairo-lang-macro/Cargo.toml

@@ -15,3 +15,4 @@ repository.workspace = true
 [dependencies]
 cairo-lang-macro-attributes = { path = "../cairo-lang-macro-attributes" }
 cairo-lang-macro-stable = { path = "../cairo-lang-macro-stable" }
+linkme.workspace = true

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

@@ -1,4 +1,7 @@
 pub use cairo_lang_macro_attributes::*;
+#[doc(hidden)]
+pub use linkme;
+
 use cairo_lang_macro_stable::ffi::StableSlice;
 use cairo_lang_macro_stable::{
     StableAuxData, StableDiagnostic, StableProcMacroResult, StableSeverity, StableTokenStream,
@@ -24,6 +27,30 @@ pub unsafe extern "C" fn free_result(result: StableProcMacroResult) {
     ProcMacroResult::from_owned_stable(result);
 }
 
+#[doc(hidden)]
+#[linkme::distributed_slice]
+pub static AUX_DATA_CALLBACKS: [fn(Vec<AuxData>)];
+
+#[no_mangle]
+#[doc(hidden)]
+pub unsafe extern "C" fn aux_data_callback(
+    stable_aux_data: StableSlice<StableAuxData>,
+) -> StableSlice<StableAuxData> {
+    if !AUX_DATA_CALLBACKS.is_empty() {
+        // Callback has been defined, applying the aux data collection.
+        let (ptr, n) = stable_aux_data.raw_parts();
+        let aux_data: &[StableAuxData] = slice::from_raw_parts(ptr, n);
+        let aux_data = aux_data
+            .iter()
+            .filter_map(|a| AuxData::from_stable(a))
+            .collect::<Vec<_>>();
+        for fun in AUX_DATA_CALLBACKS {
+            fun(aux_data.clone());
+        }
+    }
+    stable_aux_data
+}
+
 #[derive(Debug)]
 pub enum ProcMacroResult {
     /// Plugin has not taken any action.
@@ -55,7 +82,7 @@ impl Display for TokenStream {
 }
 
 /// Auxiliary data returned by procedural macro.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct AuxData(Vec<u8>);
 
 impl AuxData {

scarb/src/compiler/plugin/proc_macro/ffi.rs

@@ -1,15 +1,19 @@
 use crate::core::{Config, Package, PackageId};
 use anyhow::{Context, Result};
 use cairo_lang_defs::patcher::PatchBuilder;
-use cairo_lang_macro::{ProcMacroResult, TokenStream};
-use cairo_lang_macro_stable::{StableProcMacroResult, StableResultWrapper, StableTokenStream};
+use cairo_lang_macro::{AuxData, ProcMacroResult, TokenStream};
+use cairo_lang_macro_stable::{
+    StableAuxData, StableProcMacroResult, StableResultWrapper, StableTokenStream,
+};
 use cairo_lang_syntax::node::db::SyntaxGroup;
 use cairo_lang_syntax::node::{ast, TypedSyntaxNode};
 use camino::Utf8PathBuf;
 use libloading::{Library, Symbol};
 use std::fmt::Debug;
 
 use crate::compiler::plugin::proc_macro::compilation::SharedLibraryProvider;
+use crate::compiler::plugin::proc_macro::ProcMacroAuxData;
+use cairo_lang_macro_stable::ffi::StableSlice;
 #[cfg(not(windows))]
 use libloading::os::unix::Symbol as RawSymbol;
 #[cfg(windows)]
@@ -91,14 +95,32 @@ impl ProcMacroInstance {
         // Return obtained result.
         result
     }
+
+    pub(crate) fn aux_data_callback(&self, aux_data: Vec<ProcMacroAuxData>) {
+        // Convert to stable aux data.
+        let aux_data: Vec<AuxData> = aux_data.into_iter().map(Into::into).collect();
+        let aux_data = aux_data
+            .into_iter()
+            .map(|a| a.into_stable())
+            .collect::<Vec<_>>();
+        // Create stable slice representation from vector.
+        // Note this needs to be freed manually.
+        let aux_data = StableSlice::new(aux_data);
+        // Actual call to FFI interface for aux data callback.
+        let aux_data = (self.plugin.vtable.aux_data_callback)(aux_data);
+        // Free the memory allocated by vec.
+        let _ = aux_data.into_owned();
+    }
 }
 
 type ExpandCode = extern "C" fn(StableTokenStream) -> StableResultWrapper;
 type FreeResult = extern "C" fn(StableProcMacroResult);
+type AuxDataCallback = extern "C" fn(StableSlice<StableAuxData>) -> StableSlice<StableAuxData>;
 
 struct VTableV0 {
     expand: RawSymbol<ExpandCode>,
     free_result: RawSymbol<FreeResult>,
+    aux_data_callback: RawSymbol<AuxDataCallback>,
 }
 
 impl VTableV0 {
@@ -111,9 +133,14 @@ impl VTableV0 {
             .get(b"free_result\0")
             .context("failed to load free_result function for procedural macro")?;
         let free_result = free_result.into_raw();
+        let aux_data_callback: Symbol<'_, AuxDataCallback> = library
+            .get(b"aux_data_callback\0")
+            .context("failed to load aux_data_callback function for procedural macro")?;
+        let aux_data_callback = aux_data_callback.into_raw();
         Ok(VTableV0 {
             expand,
             free_result,
+            aux_data_callback,
         })
     }
 }

scarb/src/compiler/plugin/proc_macro/host.rs

@@ -17,6 +17,7 @@ use itertools::Itertools;
 use smol_str::SmolStr;
 use std::any::Any;
 use std::sync::Arc;
+use tracing::{debug, trace_span};
 
 /// A Cairo compiler plugin controlling the procedural macro execution.
 ///
@@ -171,7 +172,19 @@ impl ProcMacroHostPlugin {
                 }
             }
         }
-        let _aux_data = data.into_iter().into_group_map_by(|d| d.macro_package_id);
+        let aux_data = data.into_iter().into_group_map_by(|d| d.macro_package_id);
+        for instance in self.macros.iter() {
+            let _ = trace_span!(
+                "aux_data_collection_callback",
+                instance = %instance.package_id()
+            )
+            .entered();
+            let data = aux_data.get(&instance.package_id()).cloned();
+            if let Some(data) = data {
+                debug!("calling aux data callback with: {data:?}");
+                instance.aux_data_callback(data.clone());
+            }
+        }
         Ok(())
     }
 }

scarb/tests/build_cairo_plugin.rs

@@ -84,6 +84,8 @@ fn simple_project_with_code(t: &impl PathChild, code: impl ToString) {
         [dependencies]
         cairo-lang-macro = {{ path = {macro_lib_path}}}
         cairo-lang-macro-stable = {{ path = {macro_stable_lib_path}}}
+        serde = {{ version = "*", features = ["derive"] }}
+        serde_json = "*"
         "#},
         )
         .build(t);
@@ -452,3 +454,82 @@ fn can_replace_original_node() {
             Run completed successfully, returning [34]
         "#});
 }
+
+#[test]
+fn can_return_aux_data_from_plugin() {
+    let temp = TempDir::new().unwrap();
+    let t = temp.child("some");
+    simple_project_with_code(
+        &t,
+        indoc! {r##"
+        use cairo_lang_macro::{ProcMacroResult, TokenStream, attribute_macro, AuxData, aux_data_collection_callback};
+        use serde::{Serialize, Deserialize};
+
+        #[derive(Debug, Serialize, Deserialize)]
+        struct SomeMacroDataFormat {
+            msg: String
+        }
+
+        #[attribute_macro]
+        pub fn some_macro(token_stream: TokenStream) -> ProcMacroResult {
+            let token_stream = TokenStream::new(
+                token_stream
+                    .to_string()
+                    // Remove macro call to avoid infinite loop.
+                    .replace("#[some]", "")
+                    .replace("12", "34")
+            );
+            let value = SomeMacroDataFormat { msg: "Hello from some macro!".to_string() };
+            let value = serde_json::to_string(&value).unwrap();
+            let value: Vec<u8> = value.into_bytes();
+            let aux_data = AuxData::new(value);
+
+            ProcMacroResult::Replace {
+                token_stream,
+                aux_data: Some(aux_data),
+                diagnostics: Vec::new()
+            }
+        }
+
+        #[aux_data_collection_callback]
+        pub fn callback(aux_data: Vec<AuxData>) {
+            let aux_data = aux_data.into_iter()
+                .map(|aux_data| {
+                    let value: Vec<u8> = aux_data.into();
+                    let aux_data: SomeMacroDataFormat = serde_json::from_slice(&value).unwrap();
+                    aux_data
+                })
+                .collect::<Vec<_>>();
+            println!("{:?}", aux_data);
+        }
+        "##},
+    );
+
+    let project = temp.child("hello");
+    ProjectBuilder::start()
+        .name("hello")
+        .version("1.0.0")
+        .dep_starknet()
+        .dep("some", &t)
+        .lib_cairo(indoc! {r#"
+            #[some]
+            fn main() -> felt252 { 12 }
+        "#})
+        .build(&project);
+
+    Scarb::quick_snapbox()
+        .arg("cairo-run")
+        // Disable output from Cargo.
+        .env("CARGO_TERM_QUIET", "true")
+        .current_dir(&project)
+        .assert()
+        .success()
+        .stdout_matches(indoc! {r#"
+            [..]Compiling some v1.0.0 ([..]Scarb.toml)
+            [..]Compiling hello v1.0.0 ([..]Scarb.toml)
+            [SomeMacroDataFormat { msg: "Hello from some macro!" }]
+            [..]Finished release target(s) in [..]
+            [..]Running hello
+            [..]Run completed successfully, returning [..]
+        "#});
+}