feat(ast): GetSpanMut trait (#4609)

Closes #4606.

Introduce `GetSpanMut` trait and implement it on all AST node types.

This has to be a separate trait, rather than adding `span_mut` method to `GetSpan` because `AstKind` implements `GetSpan`, and it only has an immutable `&` ref to AST node it contains.

crates/oxc_span/src/lib.rs

@@ -9,5 +9,5 @@ mod span;
 pub use crate::{
     atom::{Atom, CompactStr, MAX_INLINE_LEN as ATOM_MAX_INLINE_LEN},
     source_type::{Language, LanguageVariant, ModuleKind, SourceType, VALID_EXTENSIONS},
-    span::{GetSpan, Span, SPAN},
+    span::{GetSpan, GetSpanMut, Span, SPAN},
 };

crates/oxc_span/src/span.rs

@@ -352,12 +352,26 @@ impl From<Span> for LabeledSpan {
 pub trait GetSpan {
     fn span(&self) -> Span;
 }
+
+/// Get mutable ref to span for an AST node
+pub trait GetSpanMut {
+    fn span_mut(&mut self) -> &mut Span;
+}
+
 impl GetSpan for Span {
+    #[inline]
     fn span(&self) -> Span {
         *self
     }
 }
 
+impl GetSpanMut for Span {
+    #[inline]
+    fn span_mut(&mut self) -> &mut Span {
+        self
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::Span;

tasks/ast_codegen/src/generators/impl_get_span.rs

@@ -1,4 +1,3 @@
-use itertools::Itertools;
 use proc_macro2::TokenStream;
 use quote::quote;
 use syn::Variant;
@@ -44,7 +43,7 @@ impl Generator for ImplGetSpanGenerator {
                 endl!();
 
                 use crate::ast::*;
-                use oxc_span::{GetSpan, Span};
+                use oxc_span::{GetSpan, GetSpanMut, Span};
 
                 #(#impls)*
             },
@@ -55,11 +54,13 @@ impl Generator for ImplGetSpanGenerator {
 fn impl_enum(it @ REnum { item, .. }: &REnum) -> TokenStream {
     let typ = it.as_type();
     let generics = &item.generics;
-    let matches: Vec<TokenStream> = item
+    let (matches, matches_mut): (Vec<TokenStream>, Vec<TokenStream>) = item
         .variants
         .iter()
-        .map(|Variant { ident, .. }| quote!(Self :: #ident(it) => it.span()))
-        .collect_vec();
+        .map(|Variant { ident, .. }| {
+            (quote!(Self :: #ident(it) => it.span()), quote!(Self :: #ident(it) => it.span_mut()))
+        })
+        .unzip();
 
     quote! {
         endl!();
@@ -70,6 +71,15 @@ fn impl_enum(it @ REnum { item, .. }: &REnum) -> TokenStream {
                 }
             }
         }
+        endl!();
+
+        impl #generics GetSpanMut for #typ {
+            fn span_mut(&mut self) -> &mut Span {
+                match self {
+                    #(#matches_mut),*
+                }
+            }
+        }
     }
 }
 
@@ -78,18 +88,26 @@ fn impl_struct(it @ RStruct { item, .. }: &RStruct) -> TokenStream {
     let generics = &item.generics;
     let inner_span_hint =
         item.fields.iter().find(|it| it.attrs.iter().any(|a| a.path().is_ident("span")));
-    let span = if let Some(span_field) = inner_span_hint {
+    let (span, span_mut) = if let Some(span_field) = inner_span_hint {
         let ident = span_field.ident.as_ref().unwrap();
-        quote!(#ident.span())
+        (quote!(self.#ident.span()), quote!(self.#ident.span_mut()))
     } else {
-        quote!(span)
+        (quote!(self.span), quote!(&mut self.span))
     };
     quote! {
         endl!();
         impl #generics GetSpan for #typ {
             #[inline]
             fn span(&self) -> Span {
-                self.#span
+                #span
+            }
+        }
+        endl!();
+
+        impl #generics GetSpanMut for #typ {
+            #[inline]
+            fn span_mut(&mut self) -> &mut Span {
+                #span_mut
             }
         }
     }