BEncode.hs

{-

 Copyright (c) 2005 Jesper Louis Andersen <jlouis@mongers.org>
                    Lemmih <lemmih@gmail.com>

 Permission to use, copy, modify, and distribute this software for any
 purpose with or without fee is hereby granted, provided that the above
 copyright notice and this permission notice appear in all copies.

 THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

-}
-----------------------------------------------------------------------------
-- |
-- Module      :  BEncode
-- Copyright   :  (c) Jesper Louis Andersen, 2005. (c) Lemmih, 2006
-- License     :  BSD-style
-- 
-- Maintainer  :  lemmih@gmail.com
-- Stability   :  believed to be stable
-- Portability :  portable
-- 
-- Provides a BEncode data type is well as functions for converting this
-- data type to and from a String.
--
-- Also supplies a number of properties which the module must satisfy.
-----------------------------------------------------------------------------
module Data.BEncode
  (
   -- * Data types
   BEncode(..),  
   -- * Functions
   bRead,
   bShow,
   bPack
  )
where

import qualified Data.Map as Map
import Data.Map (Map)
import Data.List (sort)
import Text.ParserCombinators.Parsec
import qualified Data.ByteString.Lazy.Char8 as L
import qualified Data.ByteString.Char8 as BS
import Data.Binary

import Data.BEncode.Lexer ( Token (..), lexer )


type BParser a = GenParser Token () a

{- | The B-coding defines an abstract syntax tree given as a simple
     data type here
-}
data BEncode = BInt Integer
	     | BString L.ByteString
	     | BList [BEncode]
             | BDict (Map String BEncode)
	       deriving (Eq, Ord, Show)

instance Binary BEncode where
    put e = put (BS.concat $ L.toChunks $ bPack e)
    get = do s <- get
             case bRead (L.fromChunks [s]) of
               Just e  -> return e
               Nothing -> fail "Failed to parse BEncoded data"

-- Source possition is pretty useless in BEncoded data. FIXME
updatePos :: (SourcePos -> Token -> [Token] -> SourcePos)
updatePos pos _ _ = pos

bToken :: Token -> BParser ()
bToken t = tokenPrim show updatePos fn
    where fn t' | t' == t = Just ()
          fn _ = Nothing

token' :: (Token -> Maybe a) -> BParser a
token' = tokenPrim show updatePos

tnumber :: BParser Integer
tnumber = token' fn
    where fn (TNumber i) = Just i
          fn _ = Nothing

tstring :: BParser L.ByteString
tstring = token' fn
    where fn (TString str) = Just str
          fn _ = Nothing

withToken :: Token -> BParser a -> BParser a
withToken tok
    = between (bToken tok) (bToken TEnd)

--------------------------------------------------------------
--------------------------------------------------------------

bInt :: BParser BEncode
bInt = withToken TInt $ fmap BInt tnumber

bString :: BParser BEncode
bString = fmap BString tstring

bList :: BParser BEncode
bList = withToken TList $ fmap BList (many bParse)

bDict :: BParser BEncode
bDict = withToken TDict $ fmap (BDict . Map.fromAscList) (checkList =<< many1 bAssocList)
    where checkList lst = if (lst /= sort lst)
                          then fail "dictionary not sorted"
                          else return lst
          bAssocList
              = do str <- tstring
                   value <- bParse
                   return (L.unpack str,value)

bParse :: BParser BEncode
bParse = bDict <|> bList <|> bString <|> bInt

{- | bRead is a conversion routine. It assumes a B-coded string as input
     and attempts a parse of it into a BEncode data type
-}
bRead :: L.ByteString -> Maybe BEncode
bRead str = case parse bParse "" (lexer str) of
	     Left _err -> Nothing
	     Right b   -> Just b

-- | Render a BEncode structure to a B-coded string
bShow :: BEncode -> ShowS
bShow be = bShow' be
  where
    sc = showChar
    ss = showString
    sKV (k,v) = sString k (length k) . bShow' v
    sDict dict = foldr (.) id (map sKV (Map.toAscList dict)) 
    sList list = foldr (.) id (map bShow' list)
    sString str len = shows len . sc ':' . ss str
    bShow' b =
      case b of
        BInt i    -> sc 'i' . shows i . sc 'e'
        BString s -> sString (L.unpack s) (L.length s)
        BList bl  -> sc 'l' . sList bl . sc 'e'
        BDict bd  -> sc 'd' . sDict bd . sc 'e'

bPack :: BEncode -> L.ByteString
bPack be = L.fromChunks (bPack' be [])
    where intTag = BS.pack "i"
          colonTag = BS.pack ":"
          endTag = BS.pack "e"
          listTag = BS.pack "l"
          dictTag = BS.pack "d"
          sString s r = BS.pack (show (L.length s)) : colonTag : L.toChunks s ++ r
          bPack' (BInt i) r = intTag : BS.pack (show i) : endTag : r
          bPack' (BString s) r = sString s r
          bPack' (BList bl) r = listTag : foldr bPack' (endTag : r) bl
          bPack' (BDict bd) r = dictTag : foldr (\(k,v) -> sString (L.pack k) . bPack' v) (endTag : r) (Map.toAscList bd)

--check be = bShow be "" == L.unpack (bPack be)