Adapt proc macro server to handle both macro apis

software-mansion · Mar 3, 2025 · f538bd8 · f538bd8
1 parent d48ac76
commit f538bd8
Show file tree

Hide file tree

Showing 16 changed files with 697 additions and 650 deletions.
diff --git a/scarb/src/ops/proc_macro_server/methods/expand_attribute.rs b/scarb/src/ops/proc_macro_server/methods/expand_attribute.rs
@@ -1,9 +1,11 @@
 use std::sync::Arc;
 
 use anyhow::{Context, Result};
+use cairo_lang_macro_v2::{TextSpan, Token, TokenStream, TokenTree};
 use scarb_proc_macro_server_types::methods::{expand::ExpandAttribute, ProcMacroResult};
 
-use super::Handler;
+use super::{from_v1_diagnostic, from_v2_token_stream, Handler};
+use crate::compiler::plugin::proc_macro_common::VersionedPlugin;
 use crate::compiler::plugin::{collection::WorkspaceProcMacros, proc_macro_common::ExpansionKind};
 
 impl Handler for ExpandAttribute {
@@ -16,6 +18,7 @@ impl Handler for ExpandAttribute {
             attr,
             args,
             item,
+            call_site,
         } = params;
 
         let plugin = workspace_macros
@@ -34,15 +37,34 @@ impl Handler for ExpandAttribute {
             })
             .with_context(|| format!("Unsupported attribute: {attr}"))?;
 
-        let result = instance
-            .plugin()
-            .as_v1()
-            .unwrap()
-            .generate_code(attr.into(), args, item);
+        let plugin = instance.plugin();
+        let (token_stream, diagnostics) = match plugin {
+            VersionedPlugin::V2(plugin) => {
+                let result = plugin.generate_code(attr.into(), call_site, args, item);
+                (result.token_stream, result.diagnostics)
+            }
+            VersionedPlugin::V1(plugin) => {
+                let result = plugin.generate_code(
+                    attr.into(),
+                    from_v2_token_stream(args),
+                    from_v2_token_stream(item),
+                );
+                let token_stream = TokenStream::new(vec![TokenTree::Ident(Token::new(
+                    result.token_stream.to_string(),
+                    TextSpan::new(0, 0),
+                ))]);
+                let diagnostics = result
+                    .diagnostics
+                    .into_iter()
+                    .map(from_v1_diagnostic)
+                    .collect();
+                (token_stream, diagnostics)
+            }
+        };
 
         Ok(ProcMacroResult {
-            token_stream: result.token_stream,
-            diagnostics: result.diagnostics,
+            token_stream,
+            diagnostics,
         })
     }
 }
diff --git a/scarb/src/ops/proc_macro_server/methods/expand_attribute_v2.rs b/scarb/src/ops/proc_macro_server/methods/expand_attribute_v2.rs
diff --git a/scarb/src/ops/proc_macro_server/methods/expand_derive.rs b/scarb/src/ops/proc_macro_server/methods/expand_derive.rs
@@ -1,11 +1,12 @@
 use std::sync::Arc;
 
 use anyhow::{Context, Result};
-use cairo_lang_macro::TokenStream;
+use cairo_lang_macro_v2::{TextSpan, Token, TokenStream, TokenTree};
 use convert_case::{Case, Casing};
 use scarb_proc_macro_server_types::methods::{expand::ExpandDerive, ProcMacroResult};
 
-use super::Handler;
+use super::{from_v1_diagnostic, from_v2_token_stream, Handler};
+use crate::compiler::plugin::proc_macro_common::VersionedPlugin;
 use crate::compiler::plugin::{
     collection::WorkspaceProcMacros,
     proc_macro_common::{Expansion, ExpansionKind},
@@ -20,9 +21,10 @@ impl Handler for ExpandDerive {
             context,
             derives,
             item,
+            call_site,
         } = params;
 
-        let mut derived_code = String::new();
+        let mut derived_code = TokenStream::empty();
         let mut all_diagnostics = vec![];
 
         for derive in derives {
@@ -38,20 +40,44 @@ impl Handler for ExpandDerive {
                 .find(|instance| instance.get_expansions().contains(&expansion))
                 .with_context(|| format!("Unsupported derive macro: {derive}"))?;
 
-            let result = instance.plugin().as_v1().unwrap().generate_code(
-                expansion.name.clone(),
-                TokenStream::empty(),
-                item.clone(),
-            );
+            let plugin = instance.plugin();
+            let (token_stream, diagnostics) = match plugin {
+                VersionedPlugin::V2(plugin) => {
+                    let result = plugin.generate_code(
+                        expansion.name.clone(),
+                        call_site.clone(),
+                        TokenStream::empty(),
+                        item.clone(),
+                    );
+                    (result.token_stream, result.diagnostics)
+                }
+                VersionedPlugin::V1(plugin) => {
+                    let result = plugin.generate_code(
+                        expansion.name.clone(),
+                        cairo_lang_macro::TokenStream::empty(),
+                        from_v2_token_stream(item.clone()),
+                    );
+                    let token_stream = TokenStream::new(vec![TokenTree::Ident(Token::new(
+                        result.token_stream.to_string(),
+                        TextSpan::new(0, 0),
+                    ))]);
+                    let diagnostics = result
+                        .diagnostics
+                        .into_iter()
+                        .map(from_v1_diagnostic)
+                        .collect();
+                    (token_stream, diagnostics)
+                }
+            };
 
             // Register diagnostics.
-            all_diagnostics.extend(result.diagnostics);
+            all_diagnostics.extend(diagnostics);
             // Add generated code.
-            derived_code.push_str(&result.token_stream.to_string());
+            derived_code.tokens.extend(token_stream.tokens);
         }
 
         Ok(ProcMacroResult {
-            token_stream: TokenStream::new(derived_code),
+            token_stream: derived_code,
             diagnostics: all_diagnostics,
         })
     }

diff --git a/scarb/src/ops/proc_macro_server/methods/expand_derive_v2.rs b/scarb/src/ops/proc_macro_server/methods/expand_derive_v2.rs
diff --git a/scarb/src/ops/proc_macro_server/methods/expand_inline.rs b/scarb/src/ops/proc_macro_server/methods/expand_inline.rs
@@ -1,10 +1,11 @@
 use std::sync::Arc;
 
 use anyhow::{Context, Result};
-use cairo_lang_macro::TokenStream;
+use cairo_lang_macro_v2::{TextSpan, Token, TokenStream, TokenTree};
 use scarb_proc_macro_server_types::methods::{expand::ExpandInline, ProcMacroResult};
 
-use super::Handler;
+use super::{from_v1_diagnostic, from_v2_token_stream, Handler};
+use crate::compiler::plugin::proc_macro_common::VersionedPlugin;
 use crate::compiler::plugin::{collection::WorkspaceProcMacros, proc_macro_common::ExpansionKind};
 
 impl Handler for ExpandInline {
@@ -16,6 +17,7 @@ impl Handler for ExpandInline {
             context,
             name,
             args,
+            call_site,
         } = params;
 
         let plugin = workspace_macros
@@ -34,15 +36,39 @@ impl Handler for ExpandInline {
             })
             .with_context(|| format!("Unsupported inline macro: {name}"))?;
 
-        let result = instance.plugin().as_v1().unwrap().generate_code(
-            name.into(),
-            TokenStream::empty(),
-            args,
-        );
+        let plugin = instance.plugin();
+        let (token_stream, diagnostics) = match plugin {
+            VersionedPlugin::V2(plugin) => {
+                let result = plugin.generate_code(
+                    name.into(),
+                    call_site.clone(),
+                    TokenStream::empty(),
+                    args,
+                );
+                (result.token_stream, result.diagnostics)
+            }
+            VersionedPlugin::V1(plugin) => {
+                let result = plugin.generate_code(
+                    name.into(),
+                    cairo_lang_macro::TokenStream::empty(),
+                    from_v2_token_stream(args),
+                );
+                let token_stream = TokenStream::new(vec![TokenTree::Ident(Token::new(
+                    result.token_stream.to_string(),
+                    TextSpan::new(0, 0),
+                ))]);
+                let diagnostics = result
+                    .diagnostics
+                    .into_iter()
+                    .map(from_v1_diagnostic)
+                    .collect();
+                (token_stream, diagnostics)
+            }
+        };
 
         Ok(ProcMacroResult {
-            token_stream: result.token_stream,
-            diagnostics: result.diagnostics,
+            token_stream,
+            diagnostics,
         })
     }
 }
diff --git a/scarb/src/ops/proc_macro_server/methods/expand_inline_v2.rs b/scarb/src/ops/proc_macro_server/methods/expand_inline_v2.rs