Add .span() method to derive the span of an AST node

lambda-fairy · Jun 16, 2018 · b9279f7 · b9279f7
1 parent 4b25728
commit b9279f7
Show file tree

Hide file tree

Showing 2 changed files with 65 additions and 11 deletions.
diff --git a/maud_macros/src/ast.rs b/maud_macros/src/ast.rs
@@ -1,4 +1,4 @@
-use proc_macro::{Span, TokenStream};
+use proc_macro::{Span, TokenStream, TokenTree};
 
 #[derive(Debug)]
 pub enum Markup {
@@ -34,6 +34,30 @@ pub enum Markup {
     },
 }
 
+impl Markup {
+    pub fn span(&self) -> Span {
+        match *self {
+            Markup::Block(ref block) => block.span(),
+            Markup::Literal { span, .. } => span,
+            Markup::Symbol { ref symbol } => span_tokens(symbol.clone()),
+            Markup::Splice { outer_span, .. } => outer_span,
+            Markup::Element { ref name, ref body, .. } => {
+                let name_span = span_tokens(name.clone());
+                name_span.join(body.span()).unwrap_or(name_span)
+            },
+            Markup::Let { at_span, ref tokens } => {
+                at_span.join(span_tokens(tokens.clone())).unwrap_or(at_span)
+            },
+            Markup::Special { ref segments } => {
+                join_spans(segments.iter().map(|segment| segment.span()))
+            },
+            Markup::Match { at_span, arms_span, .. } => {
+                at_span.join(arms_span).unwrap_or(at_span)
+            },
+        }
+    }
+}
+
 #[derive(Debug)]
 pub struct Attrs {
     pub classes_static: Vec<ClassOrId>,
@@ -50,19 +74,41 @@ pub enum ElementBody {
     Block { block: Block },
 }
 
+impl ElementBody {
+    pub fn span(&self) -> Span {
+        match *self {
+            ElementBody::Void { semi_span } => semi_span,
+            ElementBody::Block { ref block } => block.span(),
+        }
+    }
+}
+
 #[derive(Debug)]
 pub struct Block {
     pub markups: Vec<Markup>,
     pub outer_span: Span,
 }
 
+impl Block {
+    pub fn span(&self) -> Span {
+        self.outer_span
+    }
+}
+
 #[derive(Debug)]
 pub struct Special {
     pub at_span: Span,
     pub head: TokenStream,
     pub body: Block,
 }
 
+impl Special {
+    pub fn span(&self) -> Span {
+        let body_span = self.body.span();
+        self.at_span.join(body_span).unwrap_or(self.at_span)
+    }
+}
+
 #[derive(Debug)]
 pub struct Attribute {
     pub name: TokenStream,
@@ -90,3 +136,21 @@ pub struct MatchArm {
     pub head: TokenStream,
     pub body: Block,
 }
+
+pub fn span_tokens<I: IntoIterator<Item=TokenTree>>(tokens: I) -> Span {
+    join_spans(tokens.into_iter().map(|token| token.span()))
+}
+
+pub fn join_spans<I: IntoIterator<Item=Span>>(spans: I) -> Span {
+    let mut iter = spans.into_iter();
+    let mut span = match iter.next() {
+        Some(span) => span,
+        None => return Span::call_site(),
+    };
+    for new_span in iter {
+        if let Some(joined) = span.join(new_span) {
+            span = joined;
+        }
+    }
+    span
+}
diff --git a/maud_macros/src/generate.rs b/maud_macros/src/generate.rs
@@ -224,16 +224,6 @@ fn desugar_toggler(Toggler { mut cond, cond_span }: Toggler) -> TokenStream {
     quote!(if $cond)
 }
 
-fn span_tokens<I: IntoIterator<Item=TokenTree>>(tokens: I) -> Span {
-    tokens
-        .into_iter()
-        .fold(None, |span: Option<Span>, token| Some(match span {
-            None => token.span(),
-            Some(span) => span.join(token.span()).unwrap_or(span),
-        }))
-        .unwrap_or_else(Span::def_site)
-}
-
 ////////////////////////////////////////////////////////
 
 struct Builder {