feat: refactor to match spans and fix bug with header doc comments

Author: ericnordelo

packages/macros/src/attribute/common/mod.rs

@@ -0,0 +1 @@
+pub mod text_span;

packages/macros/src/attribute/common/text_span.rs

@@ -0,0 +1,64 @@
+use std::collections::{HashMap, VecDeque};
+
+use cairo_lang_macro::{quote, TextSpan, Token, TokenStream, TokenTree};
+use cairo_lang_parser::utils::SimpleParserDatabase;
+use cairo_lang_syntax::node::with_db::SyntaxNodeWithDb;
+
+fn make_ident(text: &str, span: TextSpan) -> TokenTree {
+    TokenTree::Ident(Token::new(text, span))
+}
+
+fn tokenize_str(s: &str, db: &SimpleParserDatabase) -> TokenStream {
+    let syntax_node = db.parse_virtual(s).unwrap();
+    let syntax_node_with_db = SyntaxNodeWithDb::new(&syntax_node, db);
+    quote! {#syntax_node_with_db}
+}
+
+/// Merge spans from `initial_tokens` into a tokenized version of `final_output`.
+/// - If a token text in `final_output` matches one from `initial_tokens` (by exact text),
+///   it inherits that token’s span (first-come, first-served).
+/// - Otherwise, it gets `TextSpan::call_site()`.
+pub fn merge_spans_from_initial(
+    initial_tokens: &[TokenTree],
+    final_output: &str,
+    db: &SimpleParserDatabase,
+) -> Vec<TokenTree> {
+    // Build a multimap: text -> queue of spans (to support duplicates, left-to-right)
+    let mut span_index: HashMap<String, VecDeque<TextSpan>> = HashMap::new();
+    for tt in initial_tokens {
+        let TokenTree::Ident(tok) = tt;
+        let text = tok.content.to_string();
+        span_index
+            .entry(text)
+            .or_default()
+            .push_back(tok.span.clone());
+    }
+
+    // Tokenize final output
+    let final_tokens = tokenize_str(final_output, db)
+        .tokens
+        .iter()
+        .map(|tt| {
+            let TokenTree::Ident(tok) = tt;
+            tok.content.to_string()
+        })
+        .collect::<Vec<_>>();
+
+    // Rebuild final tokens, reusing spans when the text matches; else call_site.
+    let mut out = Vec::with_capacity(final_tokens.len());
+    for text in final_tokens {
+        if let Some(queue) = span_index.get_mut(&text) {
+            if let Some(span) = queue.pop_front() {
+                out.push(make_ident(&text, span));
+                if queue.is_empty() {
+                    // keep the map clean
+                    span_index.remove(&text);
+                }
+                continue;
+            }
+        }
+        out.push(make_ident(&text, TextSpan::call_site()));
+    }
+
+    out
+}

packages/macros/src/attribute/mod.rs

@@ -1,4 +1,5 @@
 //! OpenZeppelin attribute macros.
 
+pub mod common;
 pub mod type_hash;
 pub mod with_components;

packages/macros/src/attribute/type_hash/definition.rs

@@ -10,6 +10,7 @@ use convert_case::{Case, Casing};
 use indoc::formatdoc;
 use regex::Regex;
 
+use crate::attribute::common::text_span::merge_spans_from_initial;
 use crate::type_hash::parser::TypeHashParser;
 
 use super::diagnostics::errors;
@@ -29,11 +30,13 @@ use super::parser::parse_string_arg;
 /// ```
 #[attribute_macro]
 pub fn type_hash(attr_stream: TokenStream, item_stream: TokenStream) -> ProcMacroResult {
+    let no_op_result = ProcMacroResult::new(item_stream.clone());
+
     // 1. Parse the attribute stream
     let config = match parse_args(&attr_stream.to_string()) {
         Ok(config) => config,
         Err(err) => {
-            return ProcMacroResult::new(TokenStream::empty()).with_diagnostics(err.into());
+            return no_op_result.with_diagnostics(err.into());
         }
     };
 
@@ -44,24 +47,30 @@ pub fn type_hash(attr_stream: TokenStream, item_stream: TokenStream) -> ProcMacr
         Ok(node) => handle_node(&db, node, &config),
         Err(err) => {
             let error = Diagnostic::error(err.format(&db));
-            return ProcMacroResult::new(TokenStream::empty()).with_diagnostics(error.into());
+            return no_op_result.with_diagnostics(error.into());
         }
     };
 
     let generated = match content {
         Ok(generated) => generated,
         Err(err) => {
-            return ProcMacroResult::new(TokenStream::empty()).with_diagnostics(err.into());
+            return no_op_result.with_diagnostics(err.into());
         }
     };
 
-    // 3. Return the result
+    // 3. Merge spans from the item stream into the content
+    // TODO!: This should be refactored when scarb APIs get improved
     let syntax_node = db.parse_virtual(generated).unwrap();
     let content_node = SyntaxNodeWithDb::new(&syntax_node, &db);
 
-    let mut result = item_stream;
+    let mut result = item_stream.clone();
     result.extend(quote! {#content_node});
-    ProcMacroResult::new(result)
+
+    let syntax_node_with_spans =
+        merge_spans_from_initial(&item_stream.tokens, &result.to_string(), &db);
+    let token_stream =
+        TokenStream::new(syntax_node_with_spans).with_metadata(item_stream.metadata().clone());
+    ProcMacroResult::new(token_stream)
 }
 
 /// This attribute macro is used to specify an override for the SNIP-12 type.

packages/macros/src/attribute/with_components/definition.rs

@@ -1,10 +1,11 @@
-use cairo_lang_formatter::format_string;
-use cairo_lang_macro::{attribute_macro, quote, Diagnostic, ProcMacroResult, TokenStream};
+use cairo_lang_macro::{attribute_macro, Diagnostic, ProcMacroResult, TokenStream};
 use cairo_lang_parser::utils::SimpleParserDatabase;
-use cairo_lang_syntax::node::with_db::SyntaxNodeWithDb;
 use regex::Regex;
 
-use crate::with_components::{components::AllowedComponents, parser::WithComponentsParser};
+use crate::{
+    attribute::common::text_span::merge_spans_from_initial,
+    with_components::{components::AllowedComponents, parser::WithComponentsParser},
+};
 
 /// Inserts multiple component dependencies into a modules codebase.
 #[attribute_macro]
@@ -13,15 +14,15 @@ pub fn with_components(attribute_stream: TokenStream, item_stream: TokenStream)
 
     // 1. Get the components info (if valid)
     let mut components_info = vec![];
-    let empty_result = ProcMacroResult::new(TokenStream::empty());
+    let no_op_result = ProcMacroResult::new(item_stream.clone());
     for arg in args {
         let maybe_component = AllowedComponents::from_str(&arg);
         match maybe_component {
             Ok(component) => {
                 components_info.push(component.get_info());
             }
             Err(err) => {
-                return empty_result.with_diagnostics(err.into());
+                return no_op_result.with_diagnostics(err.into());
             }
         }
     }
@@ -32,19 +33,16 @@ pub fn with_components(attribute_stream: TokenStream, item_stream: TokenStream)
         Ok(node) => WithComponentsParser::new(node, &components_info).parse(&db),
         Err(err) => {
             let error = Diagnostic::error(err.format(&db));
-            return empty_result.with_diagnostics(error.into());
+            return no_op_result.with_diagnostics(error.into());
         }
     };
 
-    let formatted_content = if !content.is_empty() {
-        format_string(&db, content)
-    } else {
-        content
-    };
-
-    let syntax_node = db.parse_virtual(formatted_content).unwrap();
-    let formatted_content_node = SyntaxNodeWithDb::new(&syntax_node, &db);
-    ProcMacroResult::new(quote! {#formatted_content_node}).with_diagnostics(diagnostics)
+    // 3. Merge spans from the item stream into the content
+    // TODO!: This should be refactored when scarb APIs get improved
+    let syntax_node_with_spans = merge_spans_from_initial(&item_stream.tokens, &content, &db);
+    let token_stream =
+        TokenStream::new(syntax_node_with_spans).with_metadata(item_stream.metadata().clone());
+    ProcMacroResult::new(token_stream).with_diagnostics(diagnostics)
 }
 
 /// Parses the arguments from the attribute stream.

packages/macros/src/attribute/with_components/parser.rs

@@ -9,12 +9,12 @@ use crate::{
 };
 use cairo_lang_defs::patcher::{PatchBuilder, RewriteNode};
 use cairo_lang_macro::{Diagnostic, Diagnostics};
-use cairo_lang_syntax::node::helpers::QueryAttrs;
 use cairo_lang_syntax::node::{
     ast::{self, MaybeModuleBody},
     db::SyntaxGroup,
     SyntaxNode, Terminal, TypedSyntaxNode,
 };
+use cairo_lang_syntax::node::{helpers::QueryAttrs, kind::SyntaxKind};
 use indoc::indoc;
 use regex::Regex;
 
@@ -47,9 +47,19 @@ impl<'a> WithComponentsParser<'a> {
 
         let typed = ast::SyntaxFile::from_syntax_node(db, base_node);
         let mut base_rnode = RewriteNode::from_ast(&typed);
-        let module_rnode = base_rnode
-            .modify_child(db, ast::SyntaxFile::INDEX_ITEMS)
-            .modify_child(db, 0);
+        let module_rnode = base_rnode.modify_child(db, ast::SyntaxFile::INDEX_ITEMS);
+
+        // If the module has a header doc, skip it
+        let module_rnode = if let RewriteNode::Copied(copied) = module_rnode {
+            let children = copied.get_children(db);
+            if !children.is_empty() && children[0].kind(db) == SyntaxKind::ItemHeaderDoc {
+                module_rnode.modify_child(db, 1)
+            } else {
+                module_rnode.modify_child(db, 0)
+            }
+        } else {
+            module_rnode.modify_child(db, 0)
+        };
 
         // Validate the contract module
         let (errors, mut warnings) =

packages/macros/src/tests/snapshots/openzeppelin_macros__tests__test_type_hash__invalid_type_hash_attribute.snap

@@ -5,7 +5,10 @@ snapshot_kind: text
 ---
 TokenStream:
 
-None
+pub struct MyType {
+    pub member: felt252,
+}
+
 
 Diagnostics:
 

packages/macros/src/tests/snapshots/openzeppelin_macros__tests__test_type_hash__with_inner_custom_type.snap

@@ -5,7 +5,14 @@ snapshot_kind: text
 ---
 TokenStream:
 
-None
+pub struct MyType {
+    pub name: felt252,
+    pub version: felt252,
+    pub chain_id: felt252,
+    pub revision: felt252,
+    pub member: InnerCustomType,
+}
+
 
 Diagnostics:
 

packages/macros/src/tests/snapshots/openzeppelin_macros__tests__test_with_components__with_header_doc.snap

@@ -0,0 +1,42 @@
+---
+source: src/tests/test_with_components.rs
+expression: result
+snapshot_kind: text
+---
+TokenStream:
+
+#[starknet::contract]
+pub mod MyContract {
+    use starknet::ContractAddress;
+    #[storage]
+    pub struct Storage {
+        #[substorage(v0)]
+        pub governor_votes: GovernorVotesComponent::Storage,
+    }
+    #[constructor]
+    fn constructor(ref self: ContractState, votes_token: ContractAddress) {
+        self.governor_votes.initializer(votes_token);
+    }
+    use openzeppelin_governance::governor::extensions::GovernorVotesComponent;
+
+    component!(path: GovernorVotesComponent, storage: governor_votes, event: GovernorVotesEvent);
+
+    impl GovernorVotesInternalImpl = GovernorVotesComponent::InternalImpl<ContractState>;
+    impl GovernorVotesGovernorVotes = GovernorVotesComponent::GovernorVotes<ContractState>;
+
+    #[event]
+    #[derive(Drop, starknet::Event)]
+    enum Event {
+        #[flat]
+        GovernorVotesEvent: GovernorVotesComponent::Event,
+    }
+}
+
+
+Diagnostics:
+
+None
+
+AuxData:
+
+None

packages/macros/src/tests/snapshots/openzeppelin_macros__tests__test_with_components__with_invalid_component.snap

@@ -5,7 +5,12 @@ snapshot_kind: text
 ---
 TokenStream:
 
-None
+#[starknet::contract]
+pub mod MyContract {
+    #[storage]
+    pub struct Storage {}
+}
+
 
 Diagnostics: