{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
module Data.Binary.Serialise.CBOR.Extra
    ( encodeListWith
    , encodeMaybeWith
    , encodeRecordFields
    , encodeUnion
    , decodeUnion
    , decodeListWith
    , decodeMaybeWith
    , (<$!>)
    ) where

import           Codec.Serialise.Decoding
import           Codec.Serialise.Encoding
import           Data.List (foldl1')
import qualified Data.Text                      as T

#if MIN_VERSION_base(4,8,0)
import           Control.Monad ((<$!>))
#else
-- | Strict version of '<$>', which is available in base >= 4.8.0
(<$!>) :: Monad m => (a -> b) -> m a -> m b
{-# INLINE (<$!>) #-}
f <$!> m = do
  x <- m
  let z = f x
  z `seq` return z
infixl 4 <$!>
#endif

encodeListWith :: (a -> Encoding) -> [a] -> Encoding
encodeListWith :: (a -> Encoding) -> [a] -> Encoding
encodeListWith a -> Encoding
_ [] = Word -> Encoding
encodeListLen Word
0
encodeListWith a -> Encoding
f [a]
xs = Encoding
encodeListLenIndef
                        Encoding -> Encoding -> Encoding
forall a. Semigroup a => a -> a -> a
<> (a -> Encoding -> Encoding) -> Encoding -> [a] -> Encoding
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\a
x Encoding
r -> a -> Encoding
f a
x Encoding -> Encoding -> Encoding
forall a. Semigroup a => a -> a -> a
<> Encoding
r) Encoding
encodeBreak [a]
xs

encodeMaybeWith :: (a -> Encoding) -> Maybe a -> Encoding
encodeMaybeWith :: (a -> Encoding) -> Maybe a -> Encoding
encodeMaybeWith a -> Encoding
_ Maybe a
Nothing  = Word -> Encoding
encodeListLen Word
0
encodeMaybeWith a -> Encoding
f (Just a
x) = Word -> Encoding
encodeListLen Word
1 Encoding -> Encoding -> Encoding
forall a. Semigroup a => a -> a -> a
<> a -> Encoding
f a
x

-- We can assume the record has at least 1 field.
encodeRecordFields :: [Encoding] -> Encoding
encodeRecordFields :: [Encoding] -> Encoding
encodeRecordFields [Encoding]
l = (Encoding -> Encoding -> Encoding) -> [Encoding] -> Encoding
forall a. (a -> a -> a) -> [a] -> a
foldl1' Encoding -> Encoding -> Encoding
forall a. Semigroup a => a -> a -> a
(<>) [Encoding]
l

-- | Encode an element of a union as single-element map from a field
-- name to a value.
encodeUnion :: T.Text -> Encoding -> Encoding
encodeUnion :: Text -> Encoding -> Encoding
encodeUnion Text
t Encoding
e = Word -> Encoding
encodeMapLen Word
1 Encoding -> Encoding -> Encoding
forall a. Semigroup a => a -> a -> a
<> Text -> Encoding
encodeString Text
t Encoding -> Encoding -> Encoding
forall a. Semigroup a => a -> a -> a
<> Encoding
e

decodeUnion :: [(T.Text, Decoder s a)] -> Decoder s a
decodeUnion :: [(Text, Decoder s a)] -> Decoder s a
decodeUnion [(Text, Decoder s a)]
ds = do
    Int
_   <- Decoder s Int
forall s. Decoder s Int
decodeMapLen -- should always be 1
    Text
dfn <- Decoder s Text
forall s. Decoder s Text
decodeString
    case Text -> [(Text, Decoder s a)] -> Maybe (Decoder s a)
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Text
dfn [(Text, Decoder s a)]
ds of
      Maybe (Decoder s a)
Nothing -> String -> Decoder s a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Unexpected field in union in CBOR"
      Just Decoder s a
d -> Decoder s a
d

decodeListWith :: Decoder s a -> Decoder s [a]
decodeListWith :: Decoder s a -> Decoder s [a]
decodeListWith Decoder s a
dec = do
    Maybe Int
mn <- Decoder s (Maybe Int)
forall s. Decoder s (Maybe Int)
decodeListLenOrIndef
    case Maybe Int
mn of
      Maybe Int
Nothing -> ([a] -> a -> [a])
-> [a] -> ([a] -> [a]) -> Decoder s a -> Decoder s [a]
forall r a r' s.
(r -> a -> r) -> r -> (r -> r') -> Decoder s a -> Decoder s r'
decodeSequenceLenIndef ((a -> [a] -> [a]) -> [a] -> a -> [a]
forall a b c. (a -> b -> c) -> b -> a -> c
flip (:)) [] [a] -> [a]
forall a. [a] -> [a]
reverse   Decoder s a
dec
      Just Int
n  -> ([a] -> a -> [a])
-> [a] -> ([a] -> [a]) -> Int -> Decoder s a -> Decoder s [a]
forall r a r' s.
(r -> a -> r)
-> r -> (r -> r') -> Int -> Decoder s a -> Decoder s r'
decodeSequenceLenN     ((a -> [a] -> [a]) -> [a] -> a -> [a]
forall a b c. (a -> b -> c) -> b -> a -> c
flip (:)) [] [a] -> [a]
forall a. [a] -> [a]
reverse Int
n Decoder s a
dec

decodeMaybeWith :: Decoder s a -> Decoder s (Maybe a)
decodeMaybeWith :: Decoder s a -> Decoder s (Maybe a)
decodeMaybeWith Decoder s a
dec = do
    Int
n <- Decoder s Int
forall s. Decoder s Int
decodeListLen
    case Int
n of
      Int
0 -> Maybe a -> Decoder s (Maybe a)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
forall a. Maybe a
Nothing
      Int
1 -> do !a
x <- Decoder s a
dec
              Maybe a -> Decoder s (Maybe a)
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> Maybe a
forall a. a -> Maybe a
Just a
x)
      Int
_ -> String -> Decoder s (Maybe a)
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"unknown tag"