1 files changed, 25 insertions, 17 deletions

diff --git a/src/Language/GraphQL/AST/Lexer.hs b/src/Language/GraphQL/AST/Lexer.hs
index f95070c..0ba55e3 100644
--- a/src/Language/GraphQL/AST/Lexer.hs
+++ b/src/Language/GraphQL/AST/Lexer.hs
@@ -15,6 +15,7 @@ module Language.GraphQL.AST.Lexer
     , dollar
     , comment
     , equals
+    , extend
     , integer
     , float
     , lexeme
@@ -28,20 +29,16 @@ module Language.GraphQL.AST.Lexer
     , unicodeBOM
     ) where
 
-import Control.Applicative ( Alternative(..)
-                           , liftA2
-                           )
-import Data.Char ( chr
-                 , digitToInt
-                 , isAsciiLower
-                 , isAsciiUpper
-                 , ord
-                 )
+import Control.Applicative (Alternative(..), liftA2)
+import Data.Char (chr, digitToInt, isAsciiLower, isAsciiUpper, ord)
 import Data.Foldable (foldl')
 import Data.List (dropWhileEnd)
+import qualified Data.List.NonEmpty as NonEmpty
+import Data.List.NonEmpty (NonEmpty(..))
 import Data.Proxy (Proxy(..))
 import Data.Void (Void)
 import Text.Megaparsec ( Parsec
+                       , (<?>)
                        , between
                        , chunk
                        , chunkToTokens
@@ -56,11 +53,9 @@ import Text.Megaparsec ( Parsec
                        , takeWhile1P
                        , try
                        )
-import Text.Megaparsec.Char ( char
-                            , digitChar
-                            , space1
-                            )
+import Text.Megaparsec.Char (char, digitChar, space1)
 import qualified Text.Megaparsec.Char.Lexer as Lexer
+import Data.Text (Text)
 import qualified Data.Text as T
 import qualified Data.Text.Lazy as TL
 
@@ -97,8 +92,8 @@ dollar :: Parser T.Text
 dollar = symbol "$"
 
 -- | Parser for "@".
-at :: Parser Char
-at = char '@'
+at :: Parser Text
+at = symbol "@"
 
 -- | Parser for "&".
 amp :: Parser T.Text
@@ -134,7 +129,7 @@ braces = between (symbol "{") (symbol "}")
 
 -- | Parser for strings.
 string :: Parser T.Text
-string = between "\"" "\"" stringValue <* spaceConsumer 
+string = between "\"" "\"" stringValue <* spaceConsumer
   where
     stringValue = T.pack <$> many stringCharacter
     stringCharacter = satisfy isStringCharacter1
@@ -143,7 +138,7 @@ string = between "\"" "\"" stringValue <* spaceConsumer
 
 -- | Parser for block strings.
 blockString :: Parser T.Text
-blockString = between "\"\"\"" "\"\"\"" stringValue <* spaceConsumer 
+blockString = between "\"\"\"" "\"\"\"" stringValue <* spaceConsumer
   where
     stringValue = do
         byLine <- sepBy (many blockStringCharacter) lineTerminator
@@ -226,3 +221,16 @@ escapeSequence = do
 -- | Parser for the "Byte Order Mark".
 unicodeBOM :: Parser ()
 unicodeBOM = optional (char '\xfeff') >> pure ()
+
+-- | Parses "extend" followed by a 'symbol'. It is used by schema extensions.
+extend :: forall a. Text -> String -> NonEmpty (Parser a) -> Parser a
+extend token extensionLabel parsers
+    = foldr combine headParser (NonEmpty.tail parsers)
+    <?> extensionLabel
+  where
+    headParser = tryExtension $ NonEmpty.head parsers
+    combine current accumulated = accumulated <|> tryExtension current
+    tryExtension extensionParser = try
+        $ symbol "extend"
+        *> symbol token
+        *> extensionParser
\ No newline at end of file