starkware-libs · orizi · May 1, 2024 · May 1, 2024
diff --git a/crates/cairo-lang-defs/src/patcher.rs b/crates/cairo-lang-defs/src/patcher.rs
@@ -225,12 +225,32 @@ pub struct ModifiedNode {
 
 pub struct PatchBuilder<'a> {
     pub db: &'a dyn SyntaxGroup,
-    pub code: String,
-    pub code_mappings: Vec<CodeMapping>,
+    code: String,
+    code_mappings: Vec<CodeMapping>,
+    origin: CodeOrigin,
 }
 impl<'a> PatchBuilder<'a> {
-    pub fn new(db: &'a dyn SyntaxGroup) -> Self {
-        Self { db, code: String::default(), code_mappings: vec![] }
+    /// Creates a new patch builder, originating from `origin` node.
+    pub fn new(db: &'a dyn SyntaxGroup, origin: &impl TypedSyntaxNode) -> Self {
+        Self {
+            db,
+            code: String::default(),
+            code_mappings: vec![],
+            origin: CodeOrigin::Span(origin.as_syntax_node().span_without_trivia(db)),
+        }
+    }
+
+    /// Builds the resulting code and code mappings.
+    pub fn build(mut self) -> (String, Vec<CodeMapping>) {
+        // Adds the mapping to the original node from all code not previously mapped.
+        self.code_mappings.push(CodeMapping {
+            span: TextSpan {
+                start: TextOffset::default(),
+                end: TextOffset::default().add_width(TextWidth::from_str(&self.code)),
+            },
+            origin: self.origin,
+        });
+        (self.code, self.code_mappings)
     }
 
     pub fn add_char(&mut self, c: char) {

diff --git a/crates/cairo-lang-plugins/src/plugins/config.rs b/crates/cairo-lang-plugins/src/plugins/config.rs
@@ -36,11 +36,12 @@ impl MacroPlugin for ConfigPlugin {
         } else if let Some(builder) =
             handle_undropped_item(db, metadata.cfg_set, item_ast, &mut diagnostics)
         {
+            let (content, code_mappings) = builder.build();
             PluginResult {
                 code: Some(PluginGeneratedFile {
                     name: "config".into(),
-                    content: builder.code,
-                    code_mappings: builder.code_mappings,
+                    content,
+                    code_mappings,
                     aux_data: None,
                 }),
                 diagnostics,
@@ -104,7 +105,7 @@ fn handle_undropped_item<'a>(
         ast::ModuleItem::Trait(trait_item) => {
             let body = try_extract_matches!(trait_item.body(db), ast::MaybeTraitBody::Some)?;
             let items = get_kept_items_nodes(db, cfg_set, &body.items_vec(db), diagnostics)?;
-            let mut builder = PatchBuilder::new(db);
+            let mut builder = PatchBuilder::new(db, &trait_item);
             builder.add_node(trait_item.attributes(db).as_syntax_node());
             builder.add_node(trait_item.trait_kw(db).as_syntax_node());
             builder.add_node(trait_item.name(db).as_syntax_node());
@@ -119,7 +120,7 @@ fn handle_undropped_item<'a>(
         ast::ModuleItem::Impl(impl_item) => {
             let body = try_extract_matches!(impl_item.body(db), ast::MaybeImplBody::Some)?;
             let items = get_kept_items_nodes(db, cfg_set, &body.items_vec(db), diagnostics)?;
-            let mut builder = PatchBuilder::new(db);
+            let mut builder = PatchBuilder::new(db, &impl_item);
             builder.add_node(impl_item.attributes(db).as_syntax_node());
             builder.add_node(impl_item.impl_kw(db).as_syntax_node());
             builder.add_node(impl_item.name(db).as_syntax_node());

diff --git a/crates/cairo-lang-plugins/src/plugins/generate_trait.rs b/crates/cairo-lang-plugins/src/plugins/generate_trait.rs
@@ -51,7 +51,7 @@ fn generate_trait_for_impl(db: &dyn SyntaxGroup, impl_ast: ast::ItemImpl) -> Plu
     };
 
     let mut diagnostics = vec![];
-    let mut builder = PatchBuilder::new(db);
+    let mut builder = PatchBuilder::new(db, &impl_ast);
     let leading_trivia = impl_ast
         .attributes(db)
         .elements(db)
@@ -207,11 +207,12 @@ fn generate_trait_for_impl(db: &dyn SyntaxGroup, impl_ast: ast::ItemImpl) -> Plu
             builder.add_node(body.rbrace(db).as_syntax_node());
         }
     }
+    let (content, code_mappings) = builder.build();
     PluginResult {
         code: Some(PluginGeneratedFile {
             name: "generate_trait".into(),
-            content: builder.code,
-            code_mappings: builder.code_mappings,
+            content,
+            code_mappings,
             aux_data: None,
         }),
         diagnostics,

diff --git a/crates/cairo-lang-plugins/src/plugins/panicable.rs b/crates/cairo-lang-plugins/src/plugins/panicable.rs
@@ -67,7 +67,6 @@ fn generate_panicable_code(
         ));
         return PluginResult { code: None, diagnostics, remove_original_item: false };
     }
-    let attr = attrs.swap_remove(0);
 
     let signature = declaration.signature(db);
     let Some((inner_ty, success_variant, failure_variant)) =
@@ -80,6 +79,8 @@ fn generate_panicable_code(
         return PluginResult { code: None, diagnostics, remove_original_item: false };
     };
 
+    let attr = attrs.swap_remove(0);
+    let mut builder = PatchBuilder::new(db, &attr);
     let attr = attr.structurize(db);
 
     let Some((err_value, panicable_name)) = parse_arguments(db, &attr) else {
@@ -89,7 +90,6 @@ fn generate_panicable_code(
         ));
         return PluginResult { code: None, diagnostics, remove_original_item: false };
     };
-    let mut builder = PatchBuilder::new(db);
     builder.add_node(visibility.as_syntax_node());
     builder.add_node(declaration.function_kw(db).as_syntax_node());
     builder.add_modified(RewriteNode::new_trimmed(panicable_name.as_syntax_node()));
@@ -137,11 +137,12 @@ fn generate_panicable_code(
         .into(),
     ));
 
+    let (content, code_mappings) = builder.build();
     PluginResult {
         code: Some(PluginGeneratedFile {
             name: "panicable".into(),
-            content: builder.code,
-            code_mappings: builder.code_mappings,
+            content,
+            code_mappings,
             aux_data: None,
         }),
         diagnostics,

diff --git a/crates/cairo-lang-semantic/src/diagnostic_test.rs b/crates/cairo-lang-semantic/src/diagnostic_test.rs
@@ -79,7 +79,7 @@ impl MacroPlugin for AddInlineModuleDummyPlugin {
     ) -> PluginResult {
         match item_ast {
             ast::ModuleItem::FreeFunction(func) if func.has_attr(db, "test_change_return_type") => {
-                let mut builder = PatchBuilder::new(db);
+                let mut builder = PatchBuilder::new(db, &func);
                 let mut new_func = RewriteNode::from_ast(&func);
                 if matches!(
                     func.declaration(db).signature(db).ret_ty(db),
@@ -113,11 +113,12 @@ impl MacroPlugin for AddInlineModuleDummyPlugin {
                     &[("func".to_string(), new_func)].into(),
                 ));
 
+                let (content, code_mappings) = builder.build();
                 PluginResult {
                     code: Some(PluginGeneratedFile {
                         name: "virt2".into(),
-                        content: builder.code,
-                        code_mappings: builder.code_mappings,
+                        content,
+                        code_mappings,
                         aux_data: None,
                     }),
                     diagnostics: vec![],

diff --git a/crates/cairo-lang-semantic/src/expr/expansion_test_data/inline_macros b/crates/cairo-lang-semantic/src/expr/expansion_test_data/inline_macros
@@ -80,6 +80,6 @@ array![]
 
 //! > diagnostics
 error: Type annotations needed. Failed to infer ?0.
- --> lib.cairo[array_inline_macro]:2:67
-            let mut __array_builder_macro_result__ = core::array::ArrayTrait::new();
-                                                                  ^********^
+ --> lib.cairo:2:1
+array![]
+^******^
diff --git a/crates/cairo-lang-semantic/src/expr/test_data/inline_macros b/crates/cairo-lang-semantic/src/expr/test_data/inline_macros
@@ -602,9 +602,9 @@ error: Plugin diagnostic: Macro `write` does not support this bracket type.
            ^
 
 error: Wrong number of arguments. Expected 1, found: 2
- --> lib.cairo[format_macro]:3:5
-    core::result::ResultTrait::<(), core::fmt::Error>::unwrap(
-    ^********************************************************^
+ --> lib.cairo:32:5
+    format!;
+    ^*****^
 
 warning[E0001]: Unused variable. Consider ignoring by prefixing with `_`.
  --> lib.cairo:2:9
@@ -720,9 +720,9 @@ error: Plugin diagnostic: Macro `write` does not support this bracket type.
           ^
 
 error: Wrong number of arguments. Expected 1, found: 2
- --> lib.cairo[print_macro]:3:5
-    core::result::ResultTrait::<(), core::fmt::Error>::unwrap(
-    ^********************************************************^
+ --> lib.cairo:32:5
+    print!;
+    ^****^
 
 warning[E0001]: Unused variable. Consider ignoring by prefixing with `_`.
  --> lib.cairo:2:9
@@ -838,9 +838,9 @@ error: Plugin diagnostic: Macro `writeln` does not support this bracket type.
             ^
 
 error: Wrong number of arguments. Expected 1, found: 2
- --> lib.cairo[println_macro]:3:5
-    core::result::ResultTrait::<(), core::fmt::Error>::unwrap(
-    ^********************************************************^
+ --> lib.cairo:32:5
+    println!;
+    ^******^
 
 warning[E0001]: Unused variable. Consider ignoring by prefixing with `_`.
  --> lib.cairo:2:9

diff --git a/crates/cairo-lang-semantic/src/inline_macros/array.rs b/crates/cairo-lang-semantic/src/inline_macros/array.rs
@@ -20,7 +20,7 @@ impl InlineMacroExprPlugin for ArrayMacro {
         let ast::WrappedArgList::BracketedArgList(args) = syntax.arguments(db) else {
             return unsupported_bracket_diagnostic(db, syntax);
         };
-        let mut builder = PatchBuilder::new(db);
+        let mut builder = PatchBuilder::new(db, syntax);
         builder.add_str(
             "{
             let mut __array_builder_macro_result__ = core::array::ArrayTrait::new();",
@@ -36,11 +36,12 @@ impl InlineMacroExprPlugin for ArrayMacro {
             "\n            __array_builder_macro_result__
         }",
         );
+        let (content, code_mappings) = builder.build();
         InlinePluginResult {
             code: Some(PluginGeneratedFile {
                 name: "array_inline_macro".into(),
-                content: builder.code,
-                code_mappings: builder.code_mappings,
+                content,
+                code_mappings,
                 aux_data: None,
             }),
             diagnostics: vec![],

diff --git a/crates/cairo-lang-semantic/src/inline_macros/assert.rs b/crates/cairo-lang-semantic/src/inline_macros/assert.rs
@@ -46,7 +46,7 @@ impl InlineMacroExprPlugin for AssertMacro {
         };
         let f = "__formatter_for_assert_macro__";
         let value_escaped = escape_node(db, value.as_syntax_node());
-        let mut builder = PatchBuilder::new(db);
+        let mut builder = PatchBuilder::new(db, syntax);
         builder.add_modified(RewriteNode::interpolate_patched(
             &formatdoc! {
                 r#"
@@ -98,11 +98,12 @@ impl InlineMacroExprPlugin for AssertMacro {
                 }}
             ",
         });
+        let (content, code_mappings) = builder.build();
         InlinePluginResult {
             code: Some(PluginGeneratedFile {
                 name: format!("{}_macro", Self::NAME).into(),
-                content: builder.code,
-                code_mappings: builder.code_mappings,
+                content,
+                code_mappings,
                 aux_data: None,
             }),
             diagnostics: vec![],

diff --git a/crates/cairo-lang-semantic/src/inline_macros/format.rs b/crates/cairo-lang-semantic/src/inline_macros/format.rs
@@ -20,7 +20,7 @@ impl InlineMacroExprPlugin for FormatMacro {
         syntax: &ast::ExprInlineMacro,
     ) -> InlinePluginResult {
         let arguments = syntax.arguments(db);
-        let mut builder = PatchBuilder::new(db);
+        let mut builder = PatchBuilder::new(db, syntax);
         builder.add_modified(RewriteNode::interpolate_patched(
             &formatdoc! {
                 "
@@ -52,11 +52,12 @@ impl InlineMacroExprPlugin for FormatMacro {
             ]
             .into(),
         ));
+        let (content, code_mappings) = builder.build();
         InlinePluginResult {
             code: Some(PluginGeneratedFile {
                 name: format!("{}_macro", Self::NAME).into(),
-                content: builder.code,
-                code_mappings: builder.code_mappings,
+                content,
+                code_mappings,
                 aux_data: None,
             }),
             diagnostics: vec![],

diff --git a/crates/cairo-lang-semantic/src/inline_macros/panic.rs b/crates/cairo-lang-semantic/src/inline_macros/panic.rs
@@ -23,7 +23,7 @@ impl InlineMacroExprPlugin for PanicMacro {
         let WrappedArgList::ParenthesizedArgList(arguments_syntax) = syntax.arguments(db) else {
             return unsupported_bracket_diagnostic(db, syntax);
         };
-        let mut builder = PatchBuilder::new(db);
+        let mut builder = PatchBuilder::new(db, syntax);
         let arguments = arguments_syntax.arguments(db).elements(db);
         if arguments.is_empty() {
             builder.add_str(r#"core::panics::panic_with_byte_array(@"")"#);
@@ -63,11 +63,12 @@ impl InlineMacroExprPlugin for PanicMacro {
                 .into(),
             ));
         }
+        let (content, code_mappings) = builder.build();
         InlinePluginResult {
             code: Some(PluginGeneratedFile {
                 name: format!("{}_macro", Self::NAME).into(),
-                content: builder.code,
-                code_mappings: builder.code_mappings,
+                content,
+                code_mappings,
                 aux_data: None,
             }),
             diagnostics: vec![],

diff --git a/crates/cairo-lang-semantic/src/inline_macros/print.rs b/crates/cairo-lang-semantic/src/inline_macros/print.rs
@@ -47,7 +47,7 @@ fn generate_code_inner(
     with_newline: bool,
 ) -> InlinePluginResult {
     let arguments = syntax.arguments(db);
-    let mut builder = PatchBuilder::new(db);
+    let mut builder = PatchBuilder::new(db, syntax);
     builder.add_modified(RewriteNode::interpolate_patched(
         &formatdoc! {
             "
@@ -80,11 +80,12 @@ fn generate_code_inner(
         ]
         .into(),
     ));
+    let (content, code_mappings) = builder.build();
     InlinePluginResult {
         code: Some(PluginGeneratedFile {
             name: format!("{}_macro", get_macro_name(with_newline)).into(),
-            content: builder.code,
-            code_mappings: builder.code_mappings,
+            content,
+            code_mappings,
             aux_data: None,
         }),
         diagnostics: vec![],

diff --git a/crates/cairo-lang-semantic/src/inline_macros/write.rs b/crates/cairo-lang-semantic/src/inline_macros/write.rs
@@ -52,17 +52,18 @@ fn generate_code_inner(
         Ok(info) => info,
         Err(diagnostics) => return InlinePluginResult { code: None, diagnostics },
     };
-    let mut builder = PatchBuilder::new(db);
+    let mut builder = PatchBuilder::new(db, syntax);
     let mut diagnostics = vec![];
     info.add_to_formatter(&mut builder, &mut diagnostics, with_newline);
     if !diagnostics.is_empty() {
         return InlinePluginResult { code: None, diagnostics };
     }
+    let (content, code_mappings) = builder.build();
     InlinePluginResult {
         code: Some(PluginGeneratedFile {
             name: format!("{}_macro", get_macro_name(with_newline)).into(),
-            content: builder.code,
-            code_mappings: builder.code_mappings,
+            content,
+            code_mappings,
             aux_data: None,
         }),
         diagnostics: vec![],