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 bs = go
  where go = case BS.uncons bs of
          Just (108, rs) -> decodeList
          Just (100, rs) -> decodeDict
          Just (105, rs) -> decodeInt
          Just (102, rs) -> decodeFloat
          Just (i  , rs)
            | isDigit (fromIntegral i) ->
              let (is, rs')   = BS.break (== 58) rs
                  len         = toNum (toDigit i) is
                  (str, rs'') = BS.splitAt len (BS.tail rs')
              in Just (String (BS.toStrict str), rs'')
            | otherwise -> 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 + fromIntegral 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