Add TokenOrShebang

efmt_core/src/parse.rs

@@ -4,7 +4,7 @@ use std::path::PathBuf;
 use std::sync::Arc;
 
 pub use self::token_stream::TokenStream;
-pub use self::tokenizer::Tokenizer;
+pub use self::tokenizer::{TokenOrShebang, Tokenizer};
 
 /// A procedural macro to derive [Parse].
 pub use efmt_derive::Parse;

efmt_core/src/parse/tokenizer.rs

@@ -1,15 +1,29 @@
-use erl_tokenize::{Position, Token};
+use erl_tokenize::{Position, PositionRange, Token};
 use std::path::Path;
 
 #[derive(Debug)]
 pub struct Tokenizer {
+    shebang_end_position: Option<Position>,
     inner: erl_tokenize::Tokenizer<String>,
 }
 
 impl Tokenizer {
     pub fn new(text: String) -> Self {
+        let mut inner = erl_tokenize::Tokenizer::new(text);
+        let shebang_end_position = if inner.text().starts_with("#!") && inner.text().contains('\n')
+        {
+            let mut line_end_position = Position::new();
+            while inner.consume_char() != Some('\n') {
+                line_end_position = inner.next_position();
+            }
+            inner.set_position(line_end_position.clone());
+            Some(line_end_position)
+        } else {
+            None
+        };
         Tokenizer {
-            inner: erl_tokenize::Tokenizer::new(text),
+            shebang_end_position,
+            inner,
         }
     }
 
@@ -27,9 +41,36 @@ impl Tokenizer {
 }
 
 impl Iterator for Tokenizer {
-    type Item = erl_tokenize::Result<Token>;
+    type Item = erl_tokenize::Result<TokenOrShebang>;
 
     fn next(&mut self) -> Option<Self::Item> {
-        self.inner.next()
+        if let Some(position) = self.shebang_end_position.take() {
+            return Some(Ok(TokenOrShebang::Shebang(position)));
+        }
+        self.inner
+            .next()
+            .map(|result| result.map(TokenOrShebang::Token))
+    }
+}
+
+#[derive(Debug, Clone)]
+pub enum TokenOrShebang {
+    Token(Token),
+    Shebang(Position),
+}
+
+impl PositionRange for TokenOrShebang {
+    fn start_position(&self) -> Position {
+        match self {
+            Self::Token(token) => token.start_position(),
+            Self::Shebang(_) => Position::new(),
+        }
+    }
+
+    fn end_position(&self) -> Position {
+        match self {
+            Self::Token(token) => token.end_position(),
+            Self::Shebang(position) => position.clone(),
+        }
     }
 }