Add new dictionary builders

4 files changed, 166 insertions, 44 deletions

diff --git a/bench/Main.hs b/bench/Main.hs
index 4953a06..0259b3c 100644
--- a/bench/Main.hs
+++ b/bench/Main.hs
@@ -1,7 +1,9 @@
-{-# LANGUAGE PackageImports    #-}
-{-# LANGUAGE DeriveGeneric     #-}
-{-# LANGUAGE RecordWildCards   #-}
-{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE PackageImports     #-}
+{-# LANGUAGE DeriveGeneric      #-}
+{-# LANGUAGE RecordWildCards    #-}
+{-# LANGUAGE OverloadedStrings  #-}
+{-# LANGUAGE DeriveDataTypeable #-}
+{-# LANGUAGE BangPatterns       #-}
 module Main (main) where
 
 import Control.Applicative
@@ -12,6 +14,7 @@ import qualified Data.ByteString.Lazy as BL
 import Data.List as L
 import Data.Maybe
 import Data.Monoid
+import Data.Typeable
 import System.Environment
 
 import Criterion.Main
@@ -23,6 +26,8 @@ import             Data.AttoBencode.Parser as B
 import "bencoding" Data.BEncode     as C
 import "bencoding" Data.BEncode.Internal as C
 import "bencoding" Data.BEncode.Types    as C
+import Debug.Trace
+
 
 instance NFData A.BEncode where
     rnf (A.BInt    i) = rnf i
@@ -62,47 +67,47 @@ replicate' c x
 
 data Torrent = Torrent {
     tAnnounce     :: !ByteString
-  , tInfo         :: !BDict
   , tAnnounceList :: !(Maybe ByteString)
   , tComment      :: !(Maybe ByteString)
   , tCreatedBy    :: !(Maybe ByteString)
   , tCreationDate :: !(Maybe ByteString)
   , tEncoding     :: !(Maybe ByteString)
+  , tInfo         :: !BDict
   , tPublisher    :: !(Maybe ByteString)
   , tPublisherURL :: !(Maybe ByteString)
   , tSignature    :: !(Maybe ByteString)
-  } deriving (Show, Eq)
+  } deriving (Show, Eq, Typeable)
 
 instance NFData Torrent where
   rnf Torrent {..} = ()
 
 instance C.BEncode Torrent where
-  toBEncode Torrent {..} = fromAscAssocs
-    [ "announce"      -->  tAnnounce
-    , "announce-list" -->? tAnnounceList
-    , "comment"       -->? tComment
-    , "created by"    -->? tCreatedBy
-    , "creation date" -->? tCreationDate
-    , "encoding"      -->? tEncoding
-    , "info"          -->  tInfo
-    , "publisher"     -->? tPublisher
-    , "publisher-url" -->? tPublisherURL
-    , "signature"     -->? tSignature
-    ]
-
-  fromBEncode (C.BDict d) =
-    Torrent <$> d >--  "announce"
-            <*> d >--  "info"
-            <*> d >--? "announce-list"
-            <*> d >--? "comment"
-            <*> d >--? "created by"
-            <*> d >--? "creation date"
-            <*> d >--? "encoding"
-            <*> d >--? "publisher"
-            <*> d >--? "publisher-url"
-            <*> d >--? "signature"
-
-  fromBEncode _ = decodingError "Torrent"
+  toBEncode Torrent {..} = toDict $
+         "announce"      .=! tAnnounce
+    C..: "announce-list" .=? tAnnounceList
+    C..: "comment"       .=? tComment
+    C..: "created by"    .=? tCreatedBy
+    C..: "creation date" .=? tCreationDate
+    C..: "encoding"      .=? tEncoding
+    C..: "info"          .=!  tInfo
+    C..: "publisher"     .=? tPublisher
+    C..: "publisher-url" .=? tPublisherURL
+    C..: "signature"     .=? tSignature
+    C..: endDict
+
+
+
+  fromBEncode = fromDict $ do
+    Torrent <$>! "announce"
+            <*>? "announce-list"
+            <*>? "comment"
+            <*>? "created by"
+            <*>? "creation date"
+            <*>? "encoding"
+            <*>! "info"
+            <*>? "publisher"
+            <*>? "publisher-url"
+            <*>? "signature"
 
 {-----------------------------------------------------------------------
 --  Main
@@ -175,11 +180,21 @@ main = do
                 :: List Int -> List Int)
                d
 
-       , let Right be = C.parse torrentFile
+       , let Right !be = C.parse torrentFile
              id'   x  = let t = either error id (fromBEncode x)
                         in toBEncode (t :: Torrent)
+             !test = let Right t = C.decode torrentFile
+                     in if C.decode (BL.toStrict (C.encode t))
+                           /= Right (t :: Torrent)
+                        then error "invalid instance: BEncode Torrent"
+                        else True
+
+             replFn n f = go n
+               where go 0 = id
+                     go n = f . go (pred n)
+
+         in bench "bigdict" $ nf (replFn (1000 :: Int) id') be
 
-         in bench "bigdict" $ nf
-              (appEndo $ mconcat $ L.replicate 1000 (Endo id'))
-              be
+       , let fn x = let Right t = C.decode x in t :: Torrent
+         in bench "torrent/decode" $ nf fn torrentFile
        ]
diff --git a/bencoding.cabal b/bencoding.cabal
index b6fe821..6ac642f 100644
--- a/bencoding.cabal
+++ b/bencoding.cabal
@@ -51,7 +51,7 @@ library
                      , bytestring >= 0.10.0.2
                      , text       >= 0.11
                      , pretty
-  ghc-options:         -Wall -fno-warn-unused-do-bind
+  ghc-options:         -Wall -O2 -fno-warn-unused-do-bind
 
 
 test-suite properties
diff --git a/src/Data/BEncode.hs b/src/Data/BEncode.hs
index 81010b6..427d401 100644
--- a/src/Data/BEncode.hs
+++ b/src/Data/BEncode.hs
@@ -70,10 +70,11 @@ module Data.BEncode
          -- ** Dictionaries
          -- *** Building
        , Assoc
-       , (-->)
-       , (-->?)
-       , fromAssocs
-       , fromAscAssocs
+       , (.=!)
+       , (.=?)
+       , (.:)
+       , endDict
+       , toDict
 
          -- *** Extraction
        , decodingError
@@ -81,11 +82,26 @@ module Data.BEncode
        , optKey
        , (>--)
        , (>--?)
+
+         -- *** Extraction
+       , Get
+       , fromDict
+
+       , req
+       , opt
+       , field
+
+       , (<$>!)
+       , (<$>?)
+       , (<*>!)
+       , (<*>?)
        ) where
 
 
 import Control.Applicative
 import Control.Monad
+import Control.Monad.State
+import Control.Monad.Error
 import Data.Int
 import Data.List as L
 import Data.Maybe         (mapMaybe)
@@ -546,9 +562,67 @@ instance (BEncode a, BEncode b, BEncode c, BEncode d, BEncode e)
   {-# INLINE fromBEncode #-}
 
 {--------------------------------------------------------------------
-  Building dictionaries
+--  Dictionary extraction
 --------------------------------------------------------------------}
 
+newtype Get a = Get { runGet :: StateT BDict Result a }
+  deriving (Functor, Applicative, Alternative)
+
+next :: Get BValue
+next = Get (StateT go)
+  where
+    go  Nil          = throwError "no next"
+    go (Cons _ v xs) = pure (v, xs)
+
+req :: BKey -> Get BValue
+req !key = Get (StateT search)
+  where
+    search  Nil          = Left msg
+    search (Cons k v xs) =
+      case compare k key of
+        EQ -> Right (v, xs)
+        LT -> search xs
+        GT -> Left msg
+
+    msg = "required field `" ++ BC.unpack key ++ "' not found"
+{-# INLINE req #-}
+
+opt :: BKey -> Get (Maybe BValue)
+opt = optional . req
+{-# INLINE opt #-}
+
+{-# SPECIALIZE field :: Get BValue -> Get BValue #-}
+field :: BEncode a => Get BValue -> Get a
+field m = Get $ do
+  v <- runGet m
+  either throwError pure $ fromBEncode v
+
+(<$>!) :: BEncode a => (a -> b) -> BKey -> Get b
+f <$>! k = f <$> field (req k)
+{-# INLINE (<$>!) #-}
+
+(<$>?) :: BEncode a => (Maybe a -> b) -> BKey -> Get b
+f <$>? k = f <$> optional (field (req k))
+{-# INLINE (<$>?) #-}
+
+(<*>!) :: BEncode a => Get (a -> b) -> BKey -> Get b
+f <*>! k = f <*> field (req k)
+{-# INLINE (<*>!) #-}
+
+(<*>?) :: BEncode a => Get (Maybe a -> b) -> BKey -> Get b
+f <*>? k = f <*> optional (field (req k))
+{-# INLINE (<*>?) #-}
+
+fromDict :: forall a. Typeable a => Get a -> BValue -> Result a
+fromDict m (BDict d) = evalStateT (runGet m) d
+fromDict _  _        = decodingError (show (typeOf inst))
+  where
+    inst = error "fromDict: impossible" :: a
+
+{--------------------------------------------------------------------
+  Building dictionaries
+--------------------------------------------------------------------}
+{-
 -- | /Assoc/ used to easily build dictionaries with required and
 -- optional keys. Suppose we have we following datatype we want to
 -- serialize:
@@ -598,6 +672,39 @@ fromAssocs = undefined -- BDict . M.fromList . mapMaybe unAssoc
 fromAscAssocs :: [Assoc] -> BValue
 fromAscAssocs = BDict . BD.fromAscList . mapMaybe unAssoc
 {-# INLINE fromAscAssocs #-}
+-}
+
+type BPair = (BKey, BValue)
+type Assoc = Maybe BPair
+
+-- TODO better name
+(.=!) :: BEncode a => BKey -> a -> Assoc
+k .=! v = Just (k, toBEncode v)
+{-# INLINE (.=!) #-}
+
+infix 6 .=!
+
+(.=?) :: BEncode a => BKey -> Maybe a -> Assoc
+_ .=? Nothing = Nothing
+k .=? Just v  = Just (k, toBEncode v)
+{-# INLINE (.=?) #-}
+
+infix 6 .=?
+
+(.:) :: Assoc -> BDict -> BDict
+Nothing     .: d = d
+Just (k, v) .: d = Cons k v d
+{-# INLINE (.:) #-}
+
+infixr 5 .:
+
+toDict :: BDict -> BValue
+toDict = BDict
+{-# INLINE toDict #-}
+
+endDict :: BDict
+endDict = Nil
+{-# INLINE endDict #-}
 
 {--------------------------------------------------------------------
   Dictionary extraction
diff --git a/src/Data/BEncode/BDict.hs b/src/Data/BEncode/BDict.hs
index 2884851..925027b 100644
--- a/src/Data/BEncode/BDict.hs
+++ b/src/Data/BEncode/BDict.hs
@@ -48,12 +48,12 @@ type BKey = ByteString
 
 -- | BDictMap is list of key value pairs sorted by keys.
 data BDictMap a
-  = Cons !BKey a (BDictMap a)
+  = Cons !BKey a !(BDictMap a)
   | Nil
     deriving (Show, Read, Eq, Ord)
 
 instance NFData a => NFData (BDictMap a) where
-  rnf Nil = ()
+  rnf  Nil         = ()
   rnf (Cons _ v xs)= rnf v `seq` rnf xs
 
 instance Functor BDictMap where