{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE FlexibleInstances         #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE UndecidableInstances      #-}
module Flat.Instances.Mono
  ( sizeSequence
  , encodeSequence
  , decodeSequence
  , sizeList
  , encodeList
  , decodeList
  , sizeSet
  , encodeSet
  , decodeSet
  , sizeMap
  , encodeMap
  , decodeMap
  , AsArray(..)
  , AsList(..)
  , AsSet(..)
  , AsMap(..)
  )
where

import           Data.Containers
import qualified Data.Foldable        as F
import           Data.MonoTraversable (Element, MonoFoldable, ofoldl', otoList)
import           Data.Sequences       (IsSequence)
import qualified Data.Sequences       as S
import           Flat.Instances.Util

{- $setup
>>> import Flat.Instances.Base()
>>> import Flat.Instances.Test
>>> import Data.Word
>>> import qualified Data.Set
>>> import qualified Data.Map
-}

{-|
Sequences are defined as Arrays:

Array v = A0
        | A1 v (Array v)
        | A2 v v (Array v)
        ...
        | A255 ... (Array v)

In practice, this means that the sequence is encoded as a sequence of blocks of up to 255 elements, with every block preceded by the count of the elements in the block and a final 0-length block.

Lists are defined as:

List a ≡  Nil
        | Cons a (List a)

In practice, this means that the list elements will be prefixed with a 1 bit and followed by a final 0 bit.

The AsList/AsArray wrappers can be used to serialise sequences as Lists or Arrays.

Let's see some examples.

>>> flatBits $ AsList [True,True,True]
"1111110"

So we have Cons True (11) repeated three times, followed by a final Nil (0).

The list encoding is the default one for lists so AsList is in this case unnecessary:

>>> flatBits $ [True,True,True]
"1111110"

We can force a list to be encoded as an Array with AsArray:

>>> flatBits $ AsArray [True,True,True]
"00000011 11100000 000"

We have the initial block with a count of 3 (3 == 00000011) followed by the elements True True True (111) and then the final block of 0 elements ("00000 000").

>>> flatBits $ [AsArray [True,True,True]]
"10000001 11110000 00000"

>>> flatBits $ (True,True,True,AsArray $ replicate 7 True)
"11100000 11111111 11000000 00"

>>> flatBits $ AsArray ([]::[()])
"00000000"

>>> flatBits $ AsList ([]::[()])
"0"

>>> tst (AsList [11::Word8,22,33])
(True,28,[133,197,164,32])

>>> tst (AsSet (Data.Set.fromList [11::Word8,22,33]))
(True,28,[133,197,164,32])

>>> tst [AsArray [1..3], AsArray [4..8]]
(True,99,[129,129,2,3,0,65,66,2,131,3,132,0,0])

>>> tst $ [AsArray [(1::Word8)..3], AsArray [4..8]]
(True,99,[129,128,129,1,128,65,65,1,65,129,194,0,0])

>>> tst $ [AsArray [(1::Int)..3]]
(True,42,[129,129,2,3,0,0])
-}
newtype AsArray a =
  AsArray
    { unArray :: a
    } deriving (Show,Eq,Ord)

instance (IsSequence r, Flat (Element r)) => Flat (AsArray r) where
  size (AsArray a) = sizeSequence a
  encode (AsArray a) = encodeSequence a
  decode = AsArray <$> decodeSequence

{- |
Calculate size of an instance of IsSequence as the sum:

* of the size of all the elements

* plus the size of the array constructors (1 byte every 255 elements plus one final byte)
-}
sizeSequence
  :: (IsSequence mono, Flat (Element mono)) => mono -> NumBits -> NumBits
sizeSequence s acc =
  let (sz, len) =
          ofoldl' (\(acc, l) e -> (size e acc, l + 1)) (acc, 0 :: NumBits) s
  in  sz + arrayBits len
{-# INLINE sizeSequence #-}

-- TODO: check which one is faster
-- sizeSequence s acc = ofoldl' (flip size) acc s + arrayBits (olength s)

-- |Encode an instance of IsSequence, as an array
encodeSequence :: (Flat (Element mono), MonoFoldable mono) => mono -> Encoding
encodeSequence = encodeArray . otoList
{-# INLINE encodeSequence #-}

-- |Decode an instance of IsSequence, as an array
decodeSequence :: (Flat (Element b), IsSequence b) => Get b
decodeSequence = S.fromList <$> decodeArrayWith decode
{-# INLINE decodeSequence #-}

newtype AsList a =
  AsList
    { unList :: a
    } deriving (Show,Eq,Ord)

instance (IsSequence l, Flat (Element l)) => Flat (AsList l) where
  -- size   = sizeList . S.unpack . unList
  -- encode = encodeList . S.unpack . unList
  -- decode = AsList . S.fromList <$> decodeListotoList

  size   = sizeList . unList
  encode = encodeList . unList
  decode = AsList <$> decodeList

{-# INLINE sizeList #-}
sizeList
  :: (MonoFoldable mono, Flat (Element mono)) => mono -> NumBits -> NumBits
sizeList l sz = ofoldl' (\s e -> size e (s + 1)) (sz + 1) l

{-# INLINE encodeList #-}
encodeList :: (Flat (Element mono), MonoFoldable mono) => mono -> Encoding
encodeList = encodeListWith encode . otoList

{-# INLINE decodeList #-}
decodeList :: (IsSequence b, Flat (Element b)) => Get b
decodeList = S.fromList <$> decodeListWith decode

{-|
Sets are saved as lists of values.

>>> tstBits $ AsSet (Data.Set.fromList ([False,True,False]::[Bool]))
(True,5,"10110")
-}
newtype AsSet a =
  AsSet
    { unSet :: a
    } deriving (Show,Eq,Ord)

instance (IsSet set, Flat (Element set)) => Flat (AsSet set) where
  size   = sizeSet . unSet
  encode = encodeSet . unSet
  decode = AsSet <$> decodeSet

sizeSet :: (IsSet set, Flat (Element set)) => Size set
sizeSet l acc = ofoldl' (\acc e -> size e (acc + 1)) (acc + 1) l
{-# INLINE sizeSet #-}

encodeSet :: (IsSet set, Flat (Element set)) => set -> Encoding
encodeSet = encodeList . setToList
{-# INLINE encodeSet #-}

decodeSet :: (IsSet set, Flat (Element set)) => Get set
decodeSet = setFromList <$> decodeList
{-# INLINE decodeSet #-}

{-|
Maps are saved as lists of (key,value) tuples.

>>> tst (AsMap (Data.Map.fromList ([]::[(Word8,())])))
(True,1,[0])

>>> tst (AsMap (Data.Map.fromList [(3::Word,9::Word)]))
(True,18,[129,132,128])
-}
newtype AsMap a =
  AsMap
    { unMap :: a
    } deriving (Show,Eq,Ord)

instance (IsMap map, Flat (ContainerKey map), Flat (MapValue map)) => Flat (AsMap map) where
  size   = sizeMap . unMap
  encode = encodeMap . unMap
  decode = AsMap <$> decodeMap

{-# INLINE sizeMap #-}
sizeMap :: (Flat (ContainerKey r), Flat (MapValue r), IsMap r) => Size r
sizeMap m acc =
  F.foldl' (\acc' (k, v) -> size k (size v (acc' + 1))) (acc + 1)
    . mapToList
    $ m
-- sizeMap l sz = ofoldl' (\s (k, v) -> size k (size v (s + 1))) (sz + 1) l

{-# INLINE encodeMap #-}
-- |Encode an instance of IsMap, as a list of (Key,Value) tuples
encodeMap
  :: (Flat (ContainerKey map), Flat (MapValue map), IsMap map)
  => map
  -> Encoding
encodeMap = encodeListWith (\(k, v) -> encode k <> encode v) . mapToList

{-# INLINE decodeMap #-}
-- |Decode an instance of IsMap, as a list of (Key,Value) tuples
decodeMap
  :: (Flat (ContainerKey map), Flat (MapValue map), IsMap map) => Get map
decodeMap = mapFromList <$> decodeListWith ((,) <$> decode <*> decode)