eben/Data/Eben.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ViewPatterns #-}

module Data.Eben where

import qualified Data.ByteString as B
import           Data.ByteString.Lazy (ByteString)
import qualified Data.ByteString.Lazy as BS
import           Data.ByteString.Builder (Builder)
import qualified Data.ByteString.Builder as BL
import           Data.Int (Int64)
import           Data.List (sortOn)
import           Data.Map.Strict (Map)
import qualified Data.Map as M
import           Data.Monoid ((<>))
import           Data.Word (Word8)

data Value
  = List [Value]
  | Dict (Map B.ByteString Value)
  | Integer Int64
  | Float Float
  | String B.ByteString
    deriving (Eq, Show, Read)

decode :: ByteString -> Maybe (Value, ByteString)
decode = go
  where go bs = case BS.uncons bs of
          Just (108, rs) -> decodeList [] rs
          Just (100, rs) -> decodeDict [] rs
          Just (105, rs) -> decodeInt rs
          Just (102, rs) -> decodeFloat rs
          Just (i  , rs)
            | isDigit (fromIntegral i) -> do
              (s, rs) <- decodeBS bs
              return (String s, rs)
            | otherwise -> Nothing
        decodeList ls bs = case BS.uncons bs of
          Just (101, rs) -> Just (List ls, rs)
          _ -> do
            (x, rs) <- go bs
            decodeList (ls ++ [x]) rs
        decodeDict ls bs = case BS.uncons bs of
          Just (101, rs) -> Just (Dict (M.fromList ls), rs)
          _ -> do
            (k, rs)  <- decodeBS bs
            (v, rs') <- go rs
            decodeDict ((k,v):ls) rs'
        decodeInt bs =
          let (is, rs) = BS.break (== 101) bs
          in return (Integer (toNum 0 is), BS.tail rs)
        decodeFloat bs =
          let (fs, rs) = BS.splitAt 4 bs
          in return (Float 0.0, BS.tail rs)
        decodeBS bs =
          let (is, rs')   = BS.break (== 58) bs
              len         = toNum 0 is
              (str, rs'') = BS.splitAt len (BS.tail rs')
          in Just (BS.toStrict str, rs'')

asDict :: Value -> Maybe (Map B.ByteString Value)
asDict (Dict ds) = Just ds
asDict _         = Nothing

asList :: Value -> Maybe [Value]
asList (List ls) = Just ls
asList _         = Nothing

asFloat :: Value -> Maybe Float
asFloat (Float f)   = Just f
asFloat (Integer i) = Just (fromIntegral i)
asFloat _           = Nothing

lookup :: B.ByteString -> Value -> Maybe Value
lookup k (Dict ds) = M.lookup k ds
lookup _ _         = Nothing

isDigit :: Word8 -> Bool
isDigit n = n >= 48 && n <= 57

toDigit :: Word8 -> Int64
toDigit n = fromIntegral n - 48

toNum :: Int64 -> ByteString -> Int64
toNum n (BS.uncons->Just(b, bs)) =
  toNum (n * 10 + toDigit b) bs
toNum n _ = n

encode :: Value -> ByteString
encode = BL.toLazyByteString . go
  where go (List vs) =
          BL.char7 'l' <> foldMap go vs <> BL.char7 'e'
        go (Dict vs) =
          BL.char7 'd'
          <> mconcat [ str k <> go v | (k, v) <- sortOn fst (M.toList vs) ]
          <> BL.char7 'e'
        go (Integer i) =
          BL.char7 'i' <> BL.string8 (show i) <> BL.char7 'e'
        go (Float f) =
          BL.char7 'f' <> BL.floatLE f <> BL.char7 'e'
        go (String bs) = str bs
        str bs =
          BL.intDec (B.length bs)
          <> BL.char7 ':'
          <> BL.byteString bs