3 files changed, 81 insertions, 61 deletions

diff --git a/src/Remote/KRPC.hs b/src/Remote/KRPC.hs
index 22dbf3aa..719b9a25 100644
--- a/src/Remote/KRPC.hs
+++ b/src/Remote/KRPC.hs
@@ -26,10 +26,12 @@ import Control.Monad
 import Control.Monad.Trans.Control
 import Control.Monad.IO.Class
 import Data.BEncode
+import Data.ByteString.Char8 as BC
 import Data.List as L
 import Data.Map  as M
 import Data.Set  as S
 import Data.Text as T
+import Data.Text.Encoding as T
 import Data.Typeable
 import Network
 
@@ -47,14 +49,25 @@ type RemoteAddr = KRemoteAddr
 
 
 queryCall :: BEncodable param
+          => Extractable param
           => KRemote -> KRemoteAddr
           -> Method param result -> param -> IO ()
 queryCall sock addr m arg = sendMessage q addr sock
   where
-    q = kquery (L.head (methodName m)) [(L.head (methodParams m), toBEncode arg)]
+    q = kquery (methodName m) (mkVals (methodParams m) (injector arg))
+    mkVals = L.zip
+
+
+extractArgs :: [ParamName] -> Map ParamName BEncode -> Result [BEncode]
+extractArgs as d = mapM f as
+        where
+          f x | Just y <- M.lookup x d = return y
+              | otherwise            = Left ("not found key " ++ BC.unpack x)
+{-# INLINE extractArgs #-}
 
 -- TODO check scheme
 getResult :: BEncodable result
+          => Extractable result
           => KRemote -> KRemoteAddr
           -> Method param result -> IO result
 getResult sock addr m = do
@@ -62,15 +75,9 @@ getResult sock addr m = do
   case resp of
     Left e -> throw (RPCException e)
     Right (respVals -> dict) -> do
-      let valName = L.head (methodVals m)
-      case M.lookup valName dict of
-        Just val | Right res <- fromBEncode val -> return res
-        Nothing     -> throw (RPCException (ProtocolError msg))
-          where
-            msg = T.concat
-              [ "Unable to find return value: ", T.pack (show valName), "\n"
-              , "in response: ", T.pack (show dict)
-              ]
+      case extractArgs (methodVals m) dict >>= extractor of
+        Right vals -> return vals
+        Left  e    -> throw (RPCException (ProtocolError (T.pack e)))
 
 -- TODO async call
 -- | Makes remote procedure call. Throws RPCException if server
@@ -78,6 +85,7 @@ getResult sock addr m = do
 --
 call :: (MonadBaseControl IO host, MonadIO host)
      => (BEncodable param, BEncodable result)
+     => (Extractable param, Extractable result)
      => RemoteAddr
      -> Method param result
      -> param
@@ -92,6 +100,7 @@ newtype Async result = Async { waitResult :: IO result }
 -- TODO document errorneous usage
 async :: MonadIO host
       => (BEncodable param, BEncodable result)
+      => (Extractable param, Extractable result)
       => RemoteAddr
       -> Method param result
       -> param
@@ -104,52 +113,50 @@ async addr m arg = do
 
 await :: MonadIO host => Async result -> host result
 await = liftIO . waitResult
+{-# INLINE await #-}
+
 
-type HandlerBody remote = (BEncode -> remote (Result BEncode), KResponseScheme)
+type HandlerBody remote = KQuery -> remote (Either KError KResponse)
 
-type MethodHandler remote = (KQueryScheme, HandlerBody remote)
+type MethodHandler remote = (MethodName, HandlerBody remote)
 
 
 -- we can safely erase types in (==>)
 (==>) :: forall (remote :: * -> *) (param :: *) (result :: *).
-           (BEncodable param, BEncodable result)
+           (BEncodable param,  BEncodable result)
+        => (Extractable param, Extractable result)
         => Monad remote
         => Method param result
         -> (param -> remote result)
         -> MethodHandler remote
-m ==> body = (methodQueryScheme m, (newbody, methodRespScheme m))
+{-# INLINE (==>) #-}
+m ==> body = (methodName m, newbody)
   where
-    newbody x = case fromBEncode x of
-                    Right a -> liftM (Right . toBEncode) (body a)
-                    Left e  -> return (Left e)
+    {-# INLINE newbody #-}
+    newbody q =
+      case extractArgs (methodParams m) (queryArgs q) >>= extractor of
+        Left  e -> return (Left (ProtocolError (T.pack e)))
+        Right a -> do
+          r <- body a
+          return (Right (kresponse (mkVals (methodVals m) (injector r))))
 
+    mkVals :: [ValName] -> [BEncode] -> [(ValName, BEncode)]
+    mkVals = L.zip
 
 -- TODO: allow forkIO
--- TODO: allow overloading
 server :: (MonadBaseControl IO remote, MonadIO remote)
        => PortNumber
        -> [MethodHandler remote]
        -> remote ()
 server servport handlers = do
     remoteServer servport $ \_ q -> do
-      case dispatch (scheme q) of
-        Nothing -> return (Left (MethodUnknown "method"))
-        Just (m, rsc) -> do
-          let arg = snd (L.head (M.toList (queryArgs q)))
-
-          res <- invoke m arg
-          let valName = L.head (S.toList (rscVals rsc))
-          return $ bimap (ProtocolError . T.pack)
-                         (kresponse . return .  (,) valName) res
+      case dispatch (queryMethod q) of
+        Nothing -> return $ Left $ MethodUnknown (decodeUtf8 (queryMethod q))
+        Just  m -> invoke m q
   where
     handlerMap = M.fromList handlers
-
---    dispatch :: KQueryScheme -> MethodHandler remote
-    dispatch s | Just m <-  M.lookup s handlerMap = return m
-               | otherwise                        = Nothing
-
---    invoke :: MethodHandler remote -> BEncode -> remote BEncode
-    invoke m args = m args
+    dispatch s = M.lookup s handlerMap
+    invoke m q = m q
 
     bimap f _ (Left  x) = Left (f x)
     bimap _ g (Right x) = Right (g x)
\ No newline at end of file
diff --git a/src/Remote/KRPC/Method.hs b/src/Remote/KRPC/Method.hs
index 420ceacf..8aa6ddc9 100644
--- a/src/Remote/KRPC/Method.hs
+++ b/src/Remote/KRPC/Method.hs
@@ -6,6 +6,7 @@
 --   Portability :  portable
 --
 {-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE FlexibleInstances, UndecidableInstances #-}
 module Remote.KRPC.Method
        ( Method(methodName, methodParams, methodVals)
        , methodQueryScheme, methodRespScheme
@@ -14,13 +15,17 @@ module Remote.KRPC.Method
        , method
 
          -- * Predefined methods
-       , idM, composeM
+       , idM
+
+         -- * Internal
+       , Extractable(..)
        ) where
 
 import Prelude hiding ((.), id)
 import Control.Applicative
 import Control.Category
 import Control.Monad
+import Data.BEncode
 import Data.ByteString as B
 import Data.List as L
 import Data.Set as S
@@ -38,7 +43,7 @@ import Remote.KRPC.Protocol
 --
 data Method param result = Method {
     -- | Name used in query and
-    methodName   :: [MethodName]
+    methodName   :: MethodName
 
     -- | Description of each parameter in /right to left/ order.
   , methodParams :: [ParamName]
@@ -46,17 +51,8 @@ data Method param result = Method {
     -- | Description of each return value in /right to left/ order.
   , methodVals   :: [ValName]
   }
-
-instance Category Method where
-  {-# SPECIALIZE instance Category Method #-}
-  id  = idM
-  {-# INLINE id #-}
-
-  (.) = composeM
-  {-# INLINE (.) #-}
-
 methodQueryScheme :: Method a b -> KQueryScheme
-methodQueryScheme = KQueryScheme <$> B.intercalate "." . methodName
+methodQueryScheme = KQueryScheme <$> methodName
                                  <*> S.fromList . methodParams
 {-# INLINE methodQueryScheme #-}
 
@@ -75,19 +71,32 @@ idM :: Method a a
 idM = method "id" ["x"] ["y"]
 {-# INLINE idM #-}
 
--- | Pipelining of two or more methods.
---
---   NOTE: composed methods will work only with this implementation of
---   KRPC, so both server and client should use this implementation,
---   otherwise you more likely get the 'ProtocolError'.
---
-composeM :: Method b c -> Method a b -> Method a c
-composeM g h = Method (methodName g ++ methodName h)
-                      (methodParams h)
-                      (methodVals g)
-{-# INLINE composeM #-}
+method :: MethodName -> [ParamName] -> [ValName] -> Method param result
+method = Method
+{-# INLINE method #-}
 
 
-method :: MethodName -> [ParamName] -> [ValName] -> Method param result
-method name = Method [name]
-{-# INLINE method #-}
\ No newline at end of file
+
+class Extractable a where
+  injector :: a -> [BEncode]
+  extractor :: [BEncode] -> Result a
+
+instance (BEncodable a, BEncodable b) => Extractable (a, b) where
+  {- SPECIALIZE instance (BEncodable a, BEncodable b) => Extractable (a, b) -}
+  injector (a, b) = [toBEncode a, toBEncode b]
+  {-# INLINE injector #-}
+
+  extractor [a, b] = (,) <$> fromBEncode a <*> fromBEncode b
+  extractor _      = decodingError "unable to match pair"
+  {-# INLINE extractor #-}
+{-
+instance BEncodable a => Extractable a where
+  {-# SPECIALIZE instance BEncodable a => Extractable a #-}
+
+  injector x = [toBEncode x]
+  {-# INLINE injector #-}
+
+  extractor [x] = fromBEncode x
+  extractor _   = decodingError "unable to match single value"
+  {-# INLINE extractor #-}
+-}
\ No newline at end of file
diff --git a/src/Remote/KRPC/Protocol.hs b/src/Remote/KRPC/Protocol.hs
index 625aba25..98674c51 100644
--- a/src/Remote/KRPC/Protocol.hs
+++ b/src/Remote/KRPC/Protocol.hs
@@ -74,7 +74,7 @@ class KMessage message scheme | message -> scheme where
   validate = (==) . scheme
   {-# INLINE validate #-}
 
-
+-- TODO Text -> ByteString
 -- TODO document that it is and how transferred
 data KError
     -- | Some error doesn't fit in any other category.
@@ -213,9 +213,13 @@ type KRemote = Socket
 withRemote :: (MonadBaseControl IO m, MonadIO m) => (KRemote -> m a) -> m a
 withRemote = bracket (liftIO (socket AF_INET Datagram defaultProtocol))
                      (liftIO .  sClose)
+{-# SPECIALIZE withRemote :: (KRemote -> IO a) -> IO a #-}
+
 
 maxMsgSize :: Int
 maxMsgSize = 16 * 1024
+{-# INLINE maxMsgSize #-}
+
 
 -- TODO eliminate toStrict
 sendMessage :: BEncodable msg => msg -> KRemoteAddr -> KRemote -> IO ()