Introduce BDictMap

3 files changed, 118 insertions, 21 deletions

diff --git a/bencoding.cabal b/bencoding.cabal
index b4cd89f..bbbf001 100644
--- a/bencoding.cabal
+++ b/bencoding.cabal
@@ -39,6 +39,8 @@ library
   default-extensions:  PatternGuards
   hs-source-dirs:      src
   exposed-modules:     Data.BEncode
+                     , Data.BEncode.BDict
+
   build-depends:       base       == 4.*
                      , ghc-prim
                      , deepseq    == 1.3.*
diff --git a/src/Data/BEncode.hs b/src/Data/BEncode.hs
index b6897ec..8858730 100644
--- a/src/Data/BEncode.hs
+++ b/src/Data/BEncode.hs
@@ -106,6 +106,7 @@ import Control.Applicative
 import Control.DeepSeq
 import Control.Monad
 import Data.Int
+import Data.List as L
 import Data.Maybe         (mapMaybe)
 import Data.Monoid
 import Data.Foldable      (foldMap)
@@ -135,12 +136,13 @@ import qualified Text.ParserCombinators.ReadP as ReadP
 import GHC.Generics
 #endif
 
+import Data.BEncode.BDict as BD
+
 
 type BInteger = Integer
 type BString  = ByteString
 type BList    = [BValue]
-type BDict    = Map BKey BValue
-type BKey     = ByteString
+type BDict    = BDictMap BValue
 
 -- | 'BEncode' is straightforward ADT for b-encoded values. Please
 -- note that since dictionaries are sorted, in most cases we can
@@ -295,7 +297,7 @@ gfromM1S :: forall c. Selector c
          => GBEncodable f BValue
          => BDict -> Result (M1 i c f p)
 gfromM1S dict
-  | Just va <- M.lookup (BC.pack (selRename name)) dict = M1 <$> gfrom va
+  | Just va <- BD.lookup (BC.pack (selRename name)) dict = M1 <$> gfrom va
   | otherwise = decodingError $ "generic: Selector not found " ++ show name
   where
     name = selName (error "gfromM1S: impossible" :: M1 i c f p)
@@ -303,7 +305,7 @@ gfromM1S dict
 instance (Selector s, GBEncodable f BValue)
        => GBEncodable (M1 S s f) BDict where
   {-# INLINE gto #-}
-  gto s @ (M1 x) = BC.pack (selRename (selName s)) `M.singleton` gto x
+  gto s @ (M1 x) = BC.pack (selRename (selName s)) `BD.singleton` gto x
 
   {-# INLINE gfrom #-}
   gfrom = gfromM1S
@@ -377,14 +379,15 @@ instance BEncodable BList where
   fromBEncode _          = decodingError "BList"
   {-# INLINE fromBEncode #-}
 
-instance BEncodable BDict where
-  toBEncode = BDict
+-}
+
+instance BEncode BDict where
+  toBEncode   = BDict
   {-# INLINE toBEncode #-}
 
   fromBEncode (BDict d) = pure d
   fromBEncode _         = decodingError "BDict"
   {-# INLINE fromBEncode #-}
--}
 
 {--------------------------------------------------------------------
 --  Integral instances
@@ -401,7 +404,7 @@ toBEncodeIntegral = BInteger . fromIntegral
 fromBEncodeIntegral :: forall a. Typeable a => Integral a => BValue -> Result a
 fromBEncodeIntegral (BInteger i) = pure (fromIntegral i)
 fromBEncodeIntegral _
-  = decodingError $ show $ typeOf (undefined :: a)
+  = decodingError $ show $ typeOf (error "fromBEncodeIntegral: imposible" :: a)
 {-# INLINE fromBEncodeIntegral #-}
 
 
@@ -500,16 +503,17 @@ instance BEncode Text where
 
 instance BEncode a => BEncode [a] where
   {-# SPECIALIZE instance BEncode BList #-}
-  toBEncode = BList . map toBEncode
+  toBEncode = BList . L.map toBEncode
   {-# INLINE toBEncode #-}
 
   fromBEncode (BList xs) = mapM fromBEncode xs
   fromBEncode _          = decodingError "list"
   {-# INLINE fromBEncode #-}
 
+{-
 instance BEncode a => BEncode (Map BKey a) where
-  {-# SPECIALIZE instance BEncode BDict #-}
-  toBEncode = BDict . M.map toBEncode
+  {-# SPECIALIZE instance BEncode (Map BKey BValue) #-}
+  toBEncode = BDict .  -- BD.map toBEncode
   {-# INLINE toBEncode #-}
 
   fromBEncode (BDict d) = traverse fromBEncode d
@@ -524,7 +528,7 @@ instance (Eq a, BEncode a) => BEncode (Set a) where
   fromBEncode (BList xs) = S.fromAscList <$> traverse fromBEncode xs
   fromBEncode _          = decodingError "Data.Set"
   {-# INLINE fromBEncode #-}
-
+-}
 instance BEncode Version where
   toBEncode = toBEncode . BC.pack . showVersion
   {-# INLINE toBEncode #-}
@@ -639,13 +643,13 @@ key -->? mval = Assoc $ ((,) key . toBEncode) <$> mval
 
 -- | Build BEncode dictionary using key -> value description.
 fromAssocs :: [Assoc] -> BValue
-fromAssocs = BDict . M.fromList . mapMaybe unAssoc
+fromAssocs = undefined -- BDict . M.fromList . mapMaybe unAssoc
 {-# INLINE fromAssocs #-}
 
 -- | A faster version of 'fromAssocs'. Should be used only when keys
 -- in builder list are sorted by ascending.
 fromAscAssocs :: [Assoc] -> BValue
-fromAscAssocs = BDict . M.fromAscList . mapMaybe unAssoc
+fromAscAssocs = BDict . BD.fromAscList . mapMaybe unAssoc
 {-# INLINE fromAscAssocs #-}
 
 {--------------------------------------------------------------------
@@ -670,7 +674,7 @@ fromAscAssocs = BDict . M.fromAscList . mapMaybe unAssoc
 --  then whole destructuring fail.
 reqKey :: BEncode a => BDict -> BKey -> Result a
 reqKey d key
-  | Just b <- M.lookup key d = fromBEncode b
+  | Just b <- BD.lookup key d = fromBEncode b
   |        otherwise         = Left msg
   where
     msg = "required field `" ++ BC.unpack key ++ "' not found"
@@ -679,7 +683,7 @@ reqKey d key
 -- 'Nothing'.
 optKey :: BEncode a => BDict -> BKey -> Result (Maybe a)
 optKey d key
-  | Just b <- M.lookup key d
+  | Just b <- BD.lookup key d
   , Right r <- fromBEncode b = return (Just r)
   | otherwise                = return Nothing
 
@@ -759,7 +763,7 @@ builder = go
                         foldMap go l <>
                         B.word8 (c2w 'e')
       go (BDict    d) = B.word8 (c2w 'd') <>
-                        foldMap mkKV (M.toAscList d) <>
+                        foldMap mkKV (BD.toAscList d) <>
                         B.word8 (c2w 'e')
           where
             mkKV (k, v) = buildString k <> go v
@@ -786,7 +790,7 @@ parser = valueP
               'l' -> P.anyChar *> ((BList    <$> listBody) <* P.anyChar)
               'd' -> do
                      P.anyChar
-                     (BDict . M.fromDistinctAscList <$>
+                     (BDict . BD.fromAscList <$>
                           many ((,) <$> stringP <*> valueP))
                        <* P.anyChar
               t   -> fail ("bencode unknown tag: " ++ [t])
@@ -819,15 +823,16 @@ parser = valueP
 --------------------------------------------------------------------}
 
 ppBS :: ByteString -> Doc
-ppBS = text . map w2c . B.unpack
+ppBS = text . L.map w2c . B.unpack
 
 -- | Convert to easily readable JSON-like document. Typically used for
 -- debugging purposes.
 ppBEncode :: BValue -> Doc
 ppBEncode (BInteger i) = int $ fromIntegral i
 ppBEncode (BString  s) = ppBS s
-ppBEncode (BList    l) = brackets $ hsep $ punctuate comma $ map ppBEncode l
+ppBEncode (BList    l)
+    = brackets $ hsep $ punctuate comma $ L.map ppBEncode l
 ppBEncode (BDict    d)
-    = braces $ vcat $ punctuate comma $ map ppKV $ M.toAscList d
+    = braces $ vcat $ punctuate comma $ L.map ppKV $ BD.toAscList d
   where
     ppKV (k, v) = ppBS k <+> colon <+> ppBEncode v
diff --git a/src/Data/BEncode/BDict.hs b/src/Data/BEncode/BDict.hs
new file mode 100644
index 0000000..c674b21
--- /dev/null
+++ b/src/Data/BEncode/BDict.hs
@@ -0,0 +1,90 @@
+module Data.BEncode.BDict
+       ( BKey
+       , BDictMap
+
+         -- * Construction
+       , Data.BEncode.BDict.empty
+       , Data.BEncode.BDict.singleton
+
+         -- * Query
+       , Data.BEncode.BDict.lookup
+
+         -- * Combine
+       , Data.BEncode.BDict.union
+
+         -- * Transformations
+       , Data.BEncode.BDict.map
+
+         -- * Conversion
+       , Data.BEncode.BDict.fromAscList
+       , Data.BEncode.BDict.toAscList
+       ) where
+
+import Control.DeepSeq
+import Data.ByteString as BS
+import Data.Monoid
+
+
+type BKey = ByteString
+
+-- STRICTNESS NOTE: the BKey is always evaluated since we either use a
+-- literal or compare before insert to the dict
+--
+-- LAYOUT NOTE: we don't use [StrictPair BKey a] since it introduce
+-- one more constructor per cell
+--
+
+-- | BDictMap is list of key value pairs sorted by keys.
+data BDictMap a
+  = Cons !BKey a (BDictMap a)
+  | Nil
+    deriving (Show, Read, Eq, Ord)
+
+instance NFData a => NFData (BDictMap a) where
+  rnf Nil = ()
+  rnf (Cons _ v xs)= rnf v `seq` rnf xs
+
+instance Functor BDictMap where
+  fmap = Data.BEncode.BDict.map
+  {-# INLINE fmap #-}
+
+--instance Foldable BDictMap where
+
+instance Monoid (BDictMap a) where
+  mempty  = Data.BEncode.BDict.empty
+  mappend = Data.BEncode.BDict.union
+
+empty :: BDictMap a
+empty = Nil
+{-# INLINE empty #-}
+
+singleton :: BKey -> a -> BDictMap a
+singleton k v = Cons k v Nil
+{-# INLINE singleton #-}
+
+lookup :: BKey -> BDictMap a -> Maybe a
+lookup x = go
+  where
+    go Nil = Nothing
+    go (Cons k v xs)
+      |   k == x  = Just v
+      | otherwise = go xs
+{-# INLINE lookup #-}
+
+union :: BDictMap a -> BDictMap a -> BDictMap a
+union = undefined
+
+map :: (a -> b) -> BDictMap a -> BDictMap b
+map f = go
+  where
+    go Nil = Nil
+    go (Cons k v xs) = Cons k (f v) (go xs)
+{-# INLINE map #-}
+
+fromAscList :: [(BKey, a)] -> BDictMap a
+fromAscList [] = Nil
+fromAscList ((k, v) : xs) = Cons k v (fromAscList xs)
+
+toAscList :: BDictMap a -> [(BKey, a)]
+toAscList Nil = []
+toAscList (Cons k v xs) = (k, v) : toAscList xs
\ No newline at end of file