1 files changed, 23 insertions, 13 deletions
diff --git a/src/Language/GraphQL/Execute.hs b/src/Language/GraphQL/Execute.hs
index de937ee..204d08c 100644
--- a/src/Language/GraphQL/Execute.hs
+++ b/src/Language/GraphQL/Execute.hs
@@ -7,9 +7,10 @@ module Language.GraphQL.Execute
     ) where
 
 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 Language.GraphQL.AST.Document
@@ -24,12 +25,13 @@ 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 :: Monad m
-    => NonEmpty (Schema.Resolver m) -- ^ Resolvers.
+    => HashMap Text (NonEmpty (Schema.Resolver m)) -- ^ Resolvers.
     -> Schema.Subs -- ^ Variable substitution function.
     -> 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,23 +42,24 @@ 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 :: Monad m
-    => NonEmpty (Schema.Resolver m) -- ^ Resolvers
+    => HashMap Text (NonEmpty (Schema.Resolver m)) -- ^ Resolvers
     -> Text -- ^ Operation name.
     -> Schema.Subs -- ^ Variable substitution function.
     -> 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 :: Monad m
-    => NonEmpty (Schema.Resolver 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,10 +70,17 @@ document schema (Just name) operations = case NE.dropWhile matchingName operatio
 document _ _ _ = return $ singleError "Missing operation name."
 
 operation :: Monad m
-    => NonEmpty (Schema.Resolver 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"