1 files changed, 44 insertions, 28 deletions
diff --git a/src/Language/GraphQL/AST/Transform.hs b/src/Language/GraphQL/AST/Transform.hs
index 99e0f3e..3aa31b0 100644
--- a/src/Language/GraphQL/AST/Transform.hs
+++ b/src/Language/GraphQL/AST/Transform.hs
@@ -1,21 +1,25 @@
 {-# LANGUAGE OverloadedStrings #-}
+
+-- | After the document is parsed, before getting executed the AST is
+--   transformed into a similar, simpler AST. This module is responsible for
+--   this transformation.
 module Language.GraphQL.AST.Transform
     ( document
     ) where
 
 import Control.Applicative (empty)
-import Control.Monad ((<=<))
 import Data.Bifunctor (first)
 import Data.Either (partitionEithers)
 import Data.Foldable (fold, foldMap)
+import Data.List.NonEmpty (NonEmpty)
 import qualified Data.List.NonEmpty as NonEmpty
 import Data.Monoid (Alt(Alt,getAlt), (<>))
 import qualified Language.GraphQL.AST as Full
 import qualified Language.GraphQL.AST.Core as Core
 import qualified Language.GraphQL.Schema as Schema
 
--- | Replaces a fragment name by a list of 'Field'. If the name doesn't match an
---   empty list is returned.
+-- | Replaces a fragment name by a list of 'Core.Field'. If the name doesn't
+--   match an empty list is returned.
 type Fragmenter = Core.Name -> [Core.Field]
 
 -- | Rewrites the original syntax tree into an intermediate representation used
@@ -40,34 +44,38 @@ operations
   -> Fragmenter
   -> [Full.OperationDefinition]
   -> Maybe Core.Document
-operations subs fr = NonEmpty.nonEmpty <=< traverse (operation subs fr)
+operations subs fr = NonEmpty.nonEmpty . fmap (operation subs fr)
 
 operation
   :: Schema.Subs
   -> Fragmenter
   -> Full.OperationDefinition
-  -> Maybe Core.Operation
+  -> Core.Operation
 operation subs fr (Full.OperationSelectionSet sels) =
   operation subs fr $ Full.OperationDefinition Full.Query empty empty empty sels
 -- TODO: Validate Variable definitions with substituter
-operation subs fr (Full.OperationDefinition operationType name _vars _dirs sels)
-    = case operationType of
-        Full.Query -> Core.Query name <$> node
-        Full.Mutation -> Core.Mutation name <$> node
-  where
-    node = traverse (hush . selection subs fr) sels
+operation subs fr (Full.OperationDefinition Full.Query name _vars _dirs sels) =
+    Core.Query name $ appendSelection subs fr sels
+operation subs fr (Full.OperationDefinition Full.Mutation name _vars _dirs sels) =
+    Core.Mutation name $ appendSelection subs fr sels
 
 selection
   :: Schema.Subs
   -> Fragmenter
   -> Full.Selection
-  -> Either [Core.Field] Core.Field
-selection subs fr (Full.SelectionField fld) =
-  Right $ field subs fr fld
-selection _    fr (Full.SelectionFragmentSpread (Full.FragmentSpread n _dirs)) =
-  Left $ fr n
-selection _ _  (Full.SelectionInlineFragment _)  =
-  error "Inline fragments not supported yet"
+  -> Either [Core.Selection] Core.Selection
+selection subs fr (Full.SelectionField fld)
+    = Right $ Core.SelectionField $ field subs fr fld
+selection _ fr (Full.SelectionFragmentSpread (Full.FragmentSpread name _))
+    = Left $ Core.SelectionField <$> fr name
+selection subs fr  (Full.SelectionInlineFragment fragment)
+    | (Full.InlineFragment (Just typeCondition) _ selectionSet) <- fragment
+        = Right
+        $ Core.SelectionFragment
+        $ Core.Fragment typeCondition
+        $ appendSelection subs fr selectionSet
+    | (Full.InlineFragment Nothing _ selectionSet) <- fragment
+        = Left $ NonEmpty.toList $ appendSelection subs fr selectionSet
 
 -- * Fragment replacement
 
@@ -83,19 +91,22 @@ defrag subs (Full.DefinitionFragment fragDef) =
   Left $ fragmentDefinition subs fragDef
 
 fragmentDefinition :: Schema.Subs -> Full.FragmentDefinition -> Fragmenter
-fragmentDefinition subs (Full.FragmentDefinition name _tc _dirs sels) name' =
-  -- TODO: Support fragments within fragments. Fold instead of map.
-  if name == name'
-  then either id pure =<< NonEmpty.toList (selection subs mempty <$> sels)
-  else empty
+fragmentDefinition subs (Full.FragmentDefinition name _tc _dirs sels) name'
+    | name == name' = selection' <$> do
+        selections <- NonEmpty.toList $ selection subs mempty <$> sels
+        either id pure selections
+    | otherwise = empty
+  where
+    selection' (Core.SelectionField field') = field'
+    selection' _ = error "Fragments within fragments are not supported yet"
 
 field :: Schema.Subs -> Fragmenter -> Full.Field -> Core.Field
 field subs fr (Full.Field a n args _dirs sels) =
     Core.Field a n (fold $ argument subs `traverse` args) (foldr go empty sels)
   where
-    go :: Full.Selection -> [Core.Field] -> [Core.Field]
-    go (Full.SelectionFragmentSpread (Full.FragmentSpread name _dirs)) = (fr name <>)
-    go sel =  (either id pure (selection subs fr sel) <>)
+    go :: Full.Selection -> [Core.Selection] -> [Core.Selection]
+    go (Full.SelectionFragmentSpread (Full.FragmentSpread name _dirs)) = ((Core.SelectionField <$> fr name) <>)
+    go sel = (either id pure (selection subs fr sel) <>)
 
 argument :: Schema.Subs -> Full.Argument -> Maybe Core.Argument
 argument subs (Full.Argument n v) = Core.Argument n <$> value subs v
@@ -116,5 +127,10 @@ value subs (Full.ValueObject   o) =
 objectField :: Schema.Subs -> Full.ObjectField -> Maybe Core.ObjectField
 objectField subs (Full.ObjectField n v) = Core.ObjectField n <$> value subs v
 
-hush :: Either a b -> Maybe b
-hush = either (const Nothing) Just
+appendSelection ::
+    Schema.Subs ->
+    Fragmenter ->
+    NonEmpty Full.Selection ->
+    NonEmpty Core.Selection
+appendSelection subs fr = NonEmpty.fromList
+    . foldr (either (++) (:) . selection subs fr) []