Support nested fragments in any order

Fix #19.

4 files changed, 77 insertions, 44 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8f54098..741a206 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -13,12 +13,8 @@ All notable changes to this project will be documented in this file.
 - Make `Language.GraphQL.AST.Core.Object` is now just a HashMap.
 - `Language.GraphQL.AST.Transform` is now isn't exposed publically anymore.
 
-### Fixed
-- (Unsupported) nested fragments doesn't throw a runtime error but return a
-  transformation error.
-
 ### Added
-  - Nested fragments support without forward lookup.
+  - Nested fragment support.
 
 ## [0.5.1.0] - 2019-10-22
 ### Deprecated
diff --git a/src/Language/GraphQL/AST/Transform.hs b/src/Language/GraphQL/AST/Transform.hs
index 118165b..93fb557 100644
--- a/src/Language/GraphQL/AST/Transform.hs
+++ b/src/Language/GraphQL/AST/Transform.hs
@@ -1,4 +1,5 @@
 {-# LANGUAGE TupleSections #-}
+{-# LANGUAGE ExplicitForAll #-}
 
 -- | After the document is parsed, before getting executed the AST is
 --   transformed into a similar, simpler AST. This module is responsible for
@@ -7,9 +8,11 @@ module Language.GraphQL.AST.Transform
     ( document
     ) where
 
-import Control.Monad (foldM)
+import Control.Arrow (first)
+import Control.Monad (foldM, unless)
 import Control.Monad.Trans.Class (lift)
-import Control.Monad.Trans.Reader (ReaderT, asks, runReaderT)
+import Control.Monad.Trans.Reader (ReaderT, ask, runReaderT)
+import Control.Monad.Trans.State (StateT, evalStateT, gets, modify)
 import Data.HashMap.Strict (HashMap)
 import qualified Data.HashMap.Strict as HashMap
 import Data.List.NonEmpty (NonEmpty)
@@ -19,26 +22,27 @@ import qualified Language.GraphQL.AST.Core as Core
 import qualified Language.GraphQL.Schema as Schema
 
 -- | Associates a fragment name with a list of 'Core.Field's.
-type Fragments = HashMap Core.Name (NonEmpty Core.Selection)
-
 data Replacement = Replacement
-    { substitute :: Schema.Subs
-    , fragments :: Fragments
+    { fragments :: HashMap Core.Name (NonEmpty Core.Selection)
+    , fragmentDefinitions :: HashMap Full.Name Full.FragmentDefinition
     }
 
-type TransformT a = ReaderT Replacement Maybe a
+type TransformT a = StateT Replacement (ReaderT Schema.Subs Maybe) a
 
 -- | Rewrites the original syntax tree into an intermediate representation used
 -- for query execution.
 document :: Schema.Subs -> Full.Document -> Maybe Core.Document
-document subs doc = do
-    fragmentMap <- foldr go (Just HashMap.empty) fragments'
-    runReaderT (operations operations') $ Replacement subs fragmentMap
+document subs document' =
+    flip runReaderT subs
+        $ evalStateT (collectFragments >> operations operationDefinitions)
+        $ Replacement HashMap.empty fragmentTable
   where
-    (fragments', operations') = foldr defragment ([], []) doc
-    go fragDef (Just fragmentsMap) =
-        runReaderT (fragmentDefinition fragDef) (Replacement subs fragmentsMap)
-    go _ Nothing = Nothing
+    (fragmentTable, operationDefinitions) = foldr defragment mempty document'
+    defragment (Full.DefinitionOperation definition) acc =
+        (definition :) <$> acc
+    defragment (Full.DefinitionFragment definition) acc =
+        let (Full.FragmentDefinition name _ _ _) = definition
+         in first (HashMap.insert name definition) acc
 
 -- * Operation
 
@@ -46,7 +50,7 @@ document subs doc = do
 operations :: [Full.OperationDefinition] -> TransformT Core.Document
 operations operations' = do
     coreOperations <- traverse operation operations'
-    lift $ NonEmpty.nonEmpty coreOperations
+    lift . lift $ NonEmpty.nonEmpty coreOperations
 
 operation :: Full.OperationDefinition -> TransformT Core.Operation
 operation (Full.OperationSelectionSet sels) =
@@ -62,8 +66,14 @@ selection ::
     TransformT (Either (NonEmpty Core.Selection) Core.Selection)
 selection (Full.SelectionField fld) = Right . Core.SelectionField <$> field fld
 selection (Full.SelectionFragmentSpread (Full.FragmentSpread name _)) = do
-    fragments' <- asks fragments
-    lift $ Left <$> HashMap.lookup name fragments'
+    fragments' <- gets fragments
+    Left <$> maybe lookupDefinition liftJust (HashMap.lookup name fragments')
+  where
+    lookupDefinition :: TransformT (NonEmpty Core.Selection)
+    lookupDefinition = do
+        fragmentDefinitions' <- gets fragmentDefinitions
+        found <- lift . lift $ HashMap.lookup name fragmentDefinitions'
+        fragmentDefinition found
 selection (Full.SelectionInlineFragment fragment)
     | (Full.InlineFragment (Just typeCondition) _ selectionSet) <- fragment
         = Right
@@ -75,25 +85,26 @@ selection (Full.SelectionInlineFragment fragment)
 
 -- * Fragment replacement
 
--- | Extract fragments into a single 'HashMap' and operation definitions.
-defragment ::
-    Full.Definition ->
-    ([Full.FragmentDefinition], [Full.OperationDefinition]) ->
-    ([Full.FragmentDefinition], [Full.OperationDefinition])
-defragment (Full.DefinitionOperation op) (fragments', operations') =
-    (fragments', op : operations')
-defragment (Full.DefinitionFragment fragDef) (fragments', operations') =
-    (fragDef : fragments', operations')
-
-fragmentDefinition :: Full.FragmentDefinition -> TransformT Fragments
+-- | Extract fragment definitions into a single 'HashMap'.
+collectFragments :: TransformT ()
+collectFragments = do
+    fragDefs <- gets fragmentDefinitions
+    let nextValue = head $ HashMap.elems fragDefs
+    unless (HashMap.null fragDefs) $ do
+        _ <- fragmentDefinition nextValue
+        collectFragments
+
+fragmentDefinition :: Full.FragmentDefinition -> TransformT (NonEmpty Core.Selection)
 fragmentDefinition (Full.FragmentDefinition name _tc _dirs sels) = do
-    newValue <- emitValue
-    fragments' <- asks fragments
-    lift . Just $ HashMap.insert name newValue fragments'
+    selections <- traverse selection sels
+    let newValue = either id pure =<< selections
+    modify $ moveFragment newValue
+    liftJust newValue
   where
-    emitValue = do
-        selections <- traverse selection sels
-        pure $ selections >>= either id pure
+    moveFragment newValue (Replacement fullFragments emptyFragDefs) =
+        let newFragments = HashMap.insert name newValue fullFragments
+            newDefinitions = HashMap.delete name emptyFragDefs
+         in Replacement newFragments newDefinitions
 
 field :: Full.Field -> TransformT Core.Field
 field (Full.Field a n args _dirs sels) = do
@@ -106,8 +117,8 @@ argument (Full.Argument n v) = Core.Argument n <$> value v
 
 value :: Full.Value -> TransformT Core.Value
 value (Full.Variable n) = do
-    substitute' <- asks substitute
-    lift $ substitute' n
+    substitute' <- lift ask
+    lift . lift $ substitute' n
 value (Full.Int i) = pure $ Core.Int i
 value (Full.Float f) = pure $ Core.Float f
 value (Full.String x) = pure $ Core.String x
@@ -137,4 +148,7 @@ appendSelection ::
     TransformT (NonEmpty Core.Selection)
 appendSelection fullSelection = do
     coreSelection <-appendSelectionOpt fullSelection
-    lift $ NonEmpty.nonEmpty coreSelection
+    lift . lift $ NonEmpty.nonEmpty coreSelection
+
+liftJust :: forall a. a -> TransformT a
+liftJust = lift . lift . Just
diff --git a/stack.yaml b/stack.yaml
index 959ed58..5aff6f9 100644
--- a/stack.yaml
+++ b/stack.yaml
@@ -1,4 +1,4 @@
-resolver: lts-14.13
+resolver: lts-14.14
 
 packages:
 - .
diff --git a/tests/Test/FragmentSpec.hs b/tests/Test/FragmentSpec.hs
index 6a49eb6..a102104 100644
--- a/tests/Test/FragmentSpec.hs
+++ b/tests/Test/FragmentSpec.hs
@@ -98,14 +98,37 @@ spec = describe "Inline fragment executor" $ do
     it "evaluates nested fragments" $ do
         let query = [r|
           {
-            ...hatFragment
+            ...circumferenceFragment
+          }
+
+          fragment circumferenceFragment on Hat {
+            circumference
           }
 
           fragment hatFragment on Hat {
             ...circumferenceFragment
           }
+        |]
+
+        actual <- graphql (circumference :| []) query
+        let expected = object
+                [ "data" .= object
+                    [ "circumference" .= (60 :: Int)
+                    ]
+                ]
+         in actual `shouldBe` expected
+
+    it "evaluates fragments defined in any order" $ do
+        let query = [r|
+          {
+            ...circumferenceFragment
+          }
 
           fragment circumferenceFragment on Hat {
+            ...hatFragment
+          }
+
+          fragment hatFragment on Hat {
             circumference
           }
         |]