1 files changed, 30 insertions, 21 deletions
diff --git a/src/Language/GraphQL/Execute.hs b/src/Language/GraphQL/Execute.hs
index 74f33f9..204d08c 100644
--- a/src/Language/GraphQL/Execute.hs
+++ b/src/Language/GraphQL/Execute.hs
@@ -6,14 +6,14 @@ module Language.GraphQL.Execute
     , executeWithName
     ) where
 
-import Control.Monad.IO.Class (MonadIO)
 import qualified Data.Aeson as Aeson
-import Data.Foldable (toList)
 import Data.List.NonEmpty (NonEmpty(..))
-import qualified Data.List.NonEmpty as NE
+import qualified Data.List.NonEmpty as NonEmpty
+import Data.HashMap.Strict (HashMap)
+import qualified Data.HashMap.Strict as HashMap
 import Data.Text (Text)
 import qualified Data.Text as Text
-import qualified Language.GraphQL.AST as AST
+import Language.GraphQL.AST.Document
 import qualified Language.GraphQL.AST.Core as AST.Core
 import qualified Language.GraphQL.Execute.Transform as Transform
 import Language.GraphQL.Error
@@ -24,13 +24,14 @@ import qualified Language.GraphQL.Schema as Schema
 --
 -- Returns the result of the query against the schema wrapped in a /data/
 -- field, or errors wrapped in an /errors/ field.
-execute :: MonadIO m
-    => NonEmpty (Schema.Resolver m) -- ^ Resolvers.
+execute :: Monad m
+    => HashMap Text (NonEmpty (Schema.Resolver m)) -- ^ Resolvers.
     -> Schema.Subs -- ^ Variable substitution function.
-    -> AST.Document -- @GraphQL@ document.
+    -> Document -- @GraphQL@ document.
     -> m Aeson.Value
 execute schema subs doc =
-    maybe transformError (document schema Nothing) $ Transform.document subs doc
+    maybe transformError (document schema Nothing)
+        $ Transform.document subs doc
   where
     transformError = return $ singleError "Schema transformation error."
 
@@ -40,24 +41,25 @@ execute schema subs doc =
 --
 -- Returns the result of the query against the schema wrapped in a /data/
 -- field, or errors wrapped in an /errors/ field.
-executeWithName :: MonadIO m
-    => NonEmpty (Schema.Resolver m) -- ^ Resolvers
+executeWithName :: Monad m
+    => HashMap Text (NonEmpty (Schema.Resolver m)) -- ^ Resolvers
     -> Text -- ^ Operation name.
     -> Schema.Subs -- ^ Variable substitution function.
-    -> AST.Document -- ^ @GraphQL@ Document.
+    -> Document -- ^ @GraphQL@ Document.
     -> m Aeson.Value
 executeWithName schema name subs doc =
-    maybe transformError (document schema $ Just name) $ Transform.document subs doc
+    maybe transformError (document schema $ Just name)
+        $ Transform.document subs doc
   where
     transformError = return $ singleError "Schema transformation error."
 
-document :: MonadIO m
-    => NonEmpty (Schema.Resolver m)
+document :: Monad m
+    => HashMap Text (NonEmpty (Schema.Resolver m))
     -> Maybe Text
     -> AST.Core.Document
     -> m Aeson.Value
 document schema Nothing (op :| []) = operation schema op
-document schema (Just name) operations = case NE.dropWhile matchingName operations of
+document schema (Just name) operations = case NonEmpty.dropWhile matchingName operations of
     [] -> return $ singleError
         $ Text.unwords ["Operation", name, "couldn't be found in the document."]
     (op:_)  -> operation schema op
@@ -67,11 +69,18 @@ document schema (Just name) operations = case NE.dropWhile matchingName operatio
     matchingName _ = False
 document _ _ _ = return $ singleError "Missing operation name."
 
-operation :: MonadIO m
-    => NonEmpty (Schema.Resolver m)
+operation :: Monad m
+    => HashMap Text (NonEmpty (Schema.Resolver m))
     -> AST.Core.Operation
     -> m Aeson.Value
-operation schema (AST.Core.Query _ flds)
-    = runCollectErrs (Schema.resolve (toList schema) flds)
-operation schema (AST.Core.Mutation _ flds)
-    = runCollectErrs (Schema.resolve (toList schema) flds)
+operation schema = schemaOperation
+  where
+    runResolver fields = runCollectErrs
+        . flip Schema.resolve fields
+        . Schema.resolversToMap
+    resolve fields queryType = maybe lookupError (runResolver fields)
+        $ HashMap.lookup queryType schema
+    lookupError = pure
+        $ singleError "Root operation type couldn't be found in the schema."
+    schemaOperation (AST.Core.Query _ fields) = resolve fields "Query"
+    schemaOperation (AST.Core.Mutation _ fields) = resolve fields "Mutation"