From ea7f2fc9b6be6283837365da3c8c8419e552834c Mon Sep 17 00:00:00 2001
From: Phillip Hilliard <pdh@seas.upenn.edu>
Date: Tue, 28 Sep 2021 22:02:17 -0400
Subject: [PATCH] Used existing macro code as a basis for implementing a source
 argument for the derive macro. This should allow users to derive their own
 Source using the provided trait, which was previously missing from derive.

---
 logos-derive/src/lib.rs        | 10 +++++++---
 logos-derive/src/parser/mod.rs |  8 ++++++++
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/logos-derive/src/lib.rs b/logos-derive/src/lib.rs
index 5c92971b..d256bb89 100644
--- a/logos-derive/src/lib.rs
+++ b/logos-derive/src/lib.rs
@@ -199,9 +199,13 @@ pub fn logos(input: TokenStream) -> TokenStream {
     let mut root = Fork::new();
 
     let extras = parser.extras.take();
-    let source = match parser.mode {
-        Mode::Utf8 => quote!(str),
-        Mode::Binary => quote!([u8]),
+    let source = if let Some(tokens) = parser.source.take(){
+        tokens
+    } else {
+        match parser.mode {
+            Mode::Utf8 => quote!(str),
+            Mode::Binary => quote!([u8]),
+        }
     };
 
     let error_def = match error {
diff --git a/logos-derive/src/parser/mod.rs b/logos-derive/src/parser/mod.rs
index 2f2215af..2dc1267c 100644
--- a/logos-derive/src/parser/mod.rs
+++ b/logos-derive/src/parser/mod.rs
@@ -24,6 +24,7 @@ use self::type_params::{replace_lifetime, traverse_type, TypeParams};
 pub struct Parser {
     pub errors: Errors,
     pub mode: Mode,
+    pub source: Option<TokenStream>,
     pub extras: MaybeVoid,
     pub subpatterns: Subpatterns,
     types: TypeParams,
@@ -93,6 +94,13 @@ impl Parser {
             };
 
             match (name.to_string().as_str(), value) {
+                ("source", NestedValue::Assign(value)) => {
+                    let span = value.span();
+                    if let Some(previous) = self.source.replace(value) {
+                        self.err("Source can be defined only once", span)
+                            .err("Previous definition here", previous.span());
+                    }
+                }
                 ("extras", NestedValue::Assign(value)) => {
                     let span = value.span();