{-# LANGUAGE DeriveAnyClass     #-}
{-# LANGUAGE DeriveFunctor      #-}
{-# LANGUAGE DeriveGeneric      #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE FlexibleInstances  #-}
{-# LANGUAGE GeneralizedNewtypeDeriving  #-}
{-# LANGUAGE Rank2Types         #-}
{-# LANGUAGE TypeFamilies       #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- | The CBOR class 'FromCBOR' does not support access to the original bytestring that is being deserialized.
--   The 'Annotated' module recovers this ability by introducing several newtypes types that,
--   along with some new operations, recover this ability.
--
-- 1. 'ByteSpan'  A pair of indexes into a bytestring, indicating a substring.
-- 2. 'Annotated'  Used in practice to pair a value with a 'ByteSpan'.
-- 3. 'FullByteString' A newtype (around a bytestring) used to store the original bytestring being deserialized.
-- 4. 'Annotator' An explict reader monad whose environment is a 'FullByteString'
--
-- The basic idea is, for a given type @t@, where we need the original bytestring, either
--
-- 1. To complete the deserialization, or
-- 2. To combine the deserialized answer with the original bytestring.
--
-- We should proceed as follows: Define instances
-- @(FromCBOR (Annotator t))@ instead of @(FromCBOR t)@. When making this instance we may freely use
-- that both 'Decoder' and 'Annotator' are both monads, and that functions 'withSlice' and 'annotatorSlice'
-- provide access to the original bytes, or portions thereof, inside of decoders.
-- Then, to actually decode a value of type @t@, we use something similar to the following code fragment.
--
-- @
-- howToUseFullBytes bytes = do
--   Annotator f <- decodeFullDecoder \"DecodingAnnotator\" (fromCBOR :: forall s. Decoder s (Annotator t)) bytes
--   pure (f (Full bytes))
-- @
-- Decode the bytes to get an @(Annotator f)@ where f is a function that when given original bytes produces a value of type @t@, then apply @f@ to @(Full bytes)@ to get the answer.
module Cardano.Binary.Annotated
  ( Annotated(..)
  , ByteSpan(..)
  , Decoded(..)
  , annotationBytes
  , annotatedDecoder
  , slice
  , fromCBORAnnotated
  , decodeFullAnnotatedBytes
  , reAnnotate
  , Annotator (..)
  , annotatorSlice
  , decodeAnnotator
  , withSlice
  , FullByteString (..)
  , serializeEncoding
  , encodePreEncoded
  )
where

import Codec.CBOR.Read (ByteOffset)
import Data.Aeson (FromJSON(..), ToJSON(..))
import qualified Data.ByteString.Lazy as BSL

import Cardano.Binary.Deserialize (decodeFullDecoder)
import Cardano.Binary.FromCBOR
  (Decoder, DecoderError, FromCBOR(..), decodeWithByteSpan)
import Cardano.Binary.ToCBOR
  (ToCBOR(..))
import Cardano.Binary.Serialize (serialize',serializeEncoding)
import Codec.CBOR.Encoding(encodePreEncoded)
import Control.DeepSeq (NFData)
import Data.Bifunctor (Bifunctor (first, second))
import qualified Data.ByteString as BS
import Data.Function (on)
import Data.Functor ((<&>))
import Data.Kind (Type)
import Data.Text (Text)
import GHC.Generics (Generic)
import NoThunks.Class (NoThunks)

-- | Extract a substring of a given ByteString corresponding to the offsets.
slice :: BSL.ByteString -> ByteSpan -> BSL.ByteString
slice :: ByteString -> ByteSpan -> ByteString
slice ByteString
bytes (ByteSpan ByteOffset
start ByteOffset
end) =
  ByteOffset -> ByteString -> ByteString
BSL.take (ByteOffset
end forall a. Num a => a -> a -> a
- ByteOffset
start) forall a b. (a -> b) -> a -> b
$ ByteOffset -> ByteString -> ByteString
BSL.drop ByteOffset
start ByteString
bytes

-- | A pair of offsets delimiting the beginning and end of a substring of a ByteString
data ByteSpan = ByteSpan !ByteOffset !ByteOffset
  deriving (forall x. Rep ByteSpan x -> ByteSpan
forall x. ByteSpan -> Rep ByteSpan x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep ByteSpan x -> ByteSpan
$cfrom :: forall x. ByteSpan -> Rep ByteSpan x
Generic, Int -> ByteSpan -> ShowS
[ByteSpan] -> ShowS
ByteSpan -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ByteSpan] -> ShowS
$cshowList :: [ByteSpan] -> ShowS
show :: ByteSpan -> String
$cshow :: ByteSpan -> String
showsPrec :: Int -> ByteSpan -> ShowS
$cshowsPrec :: Int -> ByteSpan -> ShowS
Show)

-- Used for debugging purposes only.
instance ToJSON ByteSpan where

data Annotated b a = Annotated { forall b a. Annotated b a -> b
unAnnotated :: !b, forall b a. Annotated b a -> a
annotation :: !a }
  deriving (Annotated b a -> Annotated b a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall b a. (Eq b, Eq a) => Annotated b a -> Annotated b a -> Bool
/= :: Annotated b a -> Annotated b a -> Bool
$c/= :: forall b a. (Eq b, Eq a) => Annotated b a -> Annotated b a -> Bool
== :: Annotated b a -> Annotated b a -> Bool
$c== :: forall b a. (Eq b, Eq a) => Annotated b a -> Annotated b a -> Bool
Eq, Int -> Annotated b a -> ShowS
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall b a. (Show b, Show a) => Int -> Annotated b a -> ShowS
forall b a. (Show b, Show a) => [Annotated b a] -> ShowS
forall b a. (Show b, Show a) => Annotated b a -> String
showList :: [Annotated b a] -> ShowS
$cshowList :: forall b a. (Show b, Show a) => [Annotated b a] -> ShowS
show :: Annotated b a -> String
$cshow :: forall b a. (Show b, Show a) => Annotated b a -> String
showsPrec :: Int -> Annotated b a -> ShowS
$cshowsPrec :: forall b a. (Show b, Show a) => Int -> Annotated b a -> ShowS
Show, forall a b. a -> Annotated b b -> Annotated b a
forall a b. (a -> b) -> Annotated b a -> Annotated b b
forall b a b. a -> Annotated b b -> Annotated b a
forall b a b. (a -> b) -> Annotated b a -> Annotated b b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> Annotated b b -> Annotated b a
$c<$ :: forall b a b. a -> Annotated b b -> Annotated b a
fmap :: forall a b. (a -> b) -> Annotated b a -> Annotated b b
$cfmap :: forall b a b. (a -> b) -> Annotated b a -> Annotated b b
Functor, forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall b a x. Rep (Annotated b a) x -> Annotated b a
forall b a x. Annotated b a -> Rep (Annotated b a) x
$cto :: forall b a x. Rep (Annotated b a) x -> Annotated b a
$cfrom :: forall b a x. Annotated b a -> Rep (Annotated b a) x
Generic)
  deriving anyclass (forall a. (a -> ()) -> NFData a
forall b a. (NFData b, NFData a) => Annotated b a -> ()
rnf :: Annotated b a -> ()
$crnf :: forall b a. (NFData b, NFData a) => Annotated b a -> ()
NFData, forall a.
(Context -> a -> IO (Maybe ThunkInfo))
-> (Context -> a -> IO (Maybe ThunkInfo))
-> (Proxy a -> String)
-> NoThunks a
forall b a.
(NoThunks b, NoThunks a) =>
Context -> Annotated b a -> IO (Maybe ThunkInfo)
forall b a.
(NoThunks b, NoThunks a) =>
Proxy (Annotated b a) -> String
showTypeOf :: Proxy (Annotated b a) -> String
$cshowTypeOf :: forall b a.
(NoThunks b, NoThunks a) =>
Proxy (Annotated b a) -> String
wNoThunks :: Context -> Annotated b a -> IO (Maybe ThunkInfo)
$cwNoThunks :: forall b a.
(NoThunks b, NoThunks a) =>
Context -> Annotated b a -> IO (Maybe ThunkInfo)
noThunks :: Context -> Annotated b a -> IO (Maybe ThunkInfo)
$cnoThunks :: forall b a.
(NoThunks b, NoThunks a) =>
Context -> Annotated b a -> IO (Maybe ThunkInfo)
NoThunks)

instance Bifunctor Annotated where
  first :: forall a b c. (a -> b) -> Annotated a c -> Annotated b c
first a -> b
f (Annotated a
b c
a) = forall b a. b -> a -> Annotated b a
Annotated (a -> b
f a
b) c
a
  second :: forall b c a. (b -> c) -> Annotated a b -> Annotated a c
second = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap

instance (Eq a, Ord b) => Ord (Annotated b a) where
  compare :: Annotated b a -> Annotated b a -> Ordering
compare = forall a. Ord a => a -> a -> Ordering
compare forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` forall b a. Annotated b a -> b
unAnnotated

instance ToJSON b => ToJSON (Annotated b a) where
  toJSON :: Annotated b a -> Value
toJSON = forall a. ToJSON a => a -> Value
toJSON forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall b a. Annotated b a -> b
unAnnotated
  toEncoding :: Annotated b a -> Encoding
toEncoding = forall a. ToJSON a => a -> Encoding
toEncoding forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall b a. Annotated b a -> b
unAnnotated

instance FromJSON b => FromJSON (Annotated b ()) where
  parseJSON :: Value -> Parser (Annotated b ())
parseJSON Value
j = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall b a. b -> a -> Annotated b a
Annotated () forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. FromJSON a => Value -> Parser a
parseJSON Value
j

-- | A decoder for a value paired with an annotation specifying the start and end
-- of the consumed bytes.
annotatedDecoder :: Decoder s a -> Decoder s (Annotated a ByteSpan)
annotatedDecoder :: forall s a. Decoder s a -> Decoder s (Annotated a ByteSpan)
annotatedDecoder Decoder s a
d = forall s a. Decoder s a -> Decoder s (a, ByteOffset, ByteOffset)
decodeWithByteSpan Decoder s a
d
  forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \(a
x, ByteOffset
start, ByteOffset
end) -> forall b a. b -> a -> Annotated b a
Annotated a
x (ByteOffset -> ByteOffset -> ByteSpan
ByteSpan ByteOffset
start ByteOffset
end)

-- | A decoder for a value paired with an annotation specifying the start and end
-- of the consumed bytes.
fromCBORAnnotated :: FromCBOR a => Decoder s (Annotated a ByteSpan)
fromCBORAnnotated :: forall a s. FromCBOR a => Decoder s (Annotated a ByteSpan)
fromCBORAnnotated = forall s a. Decoder s a -> Decoder s (Annotated a ByteSpan)
annotatedDecoder forall a s. FromCBOR a => Decoder s a
fromCBOR

annotationBytes :: Functor f => BSL.ByteString -> f ByteSpan -> f BS.ByteString
annotationBytes :: forall (f :: * -> *).
Functor f =>
ByteString -> f ByteSpan -> f ByteString
annotationBytes ByteString
bytes = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (ByteString -> ByteString
BSL.toStrict forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteSpan -> ByteString
slice ByteString
bytes)

-- | Decodes a value from a ByteString, requiring that the full ByteString is consumed, and
-- replaces ByteSpan annotations with the corresponding substrings of the input string.
decodeFullAnnotatedBytes
  :: Functor f
  => Text
  -> (forall s . Decoder s (f ByteSpan))
  -> BSL.ByteString
  -> Either DecoderError (f BS.ByteString)
decodeFullAnnotatedBytes :: forall (f :: * -> *).
Functor f =>
Text
-> (forall s. Decoder s (f ByteSpan))
-> ByteString
-> Either DecoderError (f ByteString)
decodeFullAnnotatedBytes Text
lbl forall s. Decoder s (f ByteSpan)
decoder ByteString
bytes =
  forall (f :: * -> *).
Functor f =>
ByteString -> f ByteSpan -> f ByteString
annotationBytes ByteString
bytes forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a.
Text
-> (forall s. Decoder s a) -> ByteString -> Either DecoderError a
decodeFullDecoder Text
lbl forall s. Decoder s (f ByteSpan)
decoder ByteString
bytes

-- | Reconstruct an annotation by re-serialising the payload to a ByteString.
reAnnotate :: ToCBOR a => Annotated a b -> Annotated a BS.ByteString
reAnnotate :: forall a b. ToCBOR a => Annotated a b -> Annotated a ByteString
reAnnotate (Annotated a
x b
_) = forall b a. b -> a -> Annotated b a
Annotated a
x (forall a. ToCBOR a => a -> ByteString
serialize' a
x)

class Decoded t where
  type BaseType t :: Type
  recoverBytes :: t -> BS.ByteString

instance Decoded (Annotated b BS.ByteString) where
  type BaseType (Annotated b BS.ByteString) = b
  recoverBytes :: Annotated b ByteString -> ByteString
recoverBytes = forall b a. Annotated b a -> a
annotation

-------------------------------------------------------------------------
-- Annotator
-------------------------------------------------------------------------

-- | This marks the entire bytestring used during decoding, rather than the
--   piece we need to finish constructing our value.
newtype FullByteString = Full BSL.ByteString

-- | A value of type @(Annotator a)@ is one that needs access to the entire
--   bytestring used during decoding to finish construction of a vaue of type @a@. A typical use
--   is some type that stores the bytes that were used to deserialize it.
--   For example the type @Inner@ below is constructed using the helper function @makeInner@
--   which serializes and stores its bytes (using 'serializeEncoding').
--   Note how we build the
--   'Annotator' by abstracting over the full bytes, and
--   using those original bytes to fill the bytes field of the constructor @Inner@.
--   The 'ToCBOR' instance just reuses the stored bytes to produce an encoding
--   (using 'encodePreEncoded').
--
-- @
-- data Inner = Inner Int Bool LByteString
--
-- makeInner :: Int -> Bool -> Inner
-- makeInner i b = Inner i b (serializeEncoding (toCBOR i <> toCBOR b))
--
-- instance ToCBOR Inner where
--   toCBOR (Inner _ _ bytes) = encodePreEncoded bytes
--
-- instance FromCBOR (Annotator Inner) where
--   fromCBOR = do
--      int <- fromCBOR
--      trueOrFalse <- fromCBOR
--      pure (Annotator (\(Full bytes) -> Inner int trueOrFalse bytes))
-- @
--
-- if an @Outer@ type has a field of type @Inner@, with a @(ToCBOR (Annotator Inner))@ instance,
-- the @Outer@ type must also have a @(ToCBOR (Annotator Outer))@ instance.
-- The key to writing that instance is to use the operation @withSlice@ which returns a pair.
-- The first component is an @Annotator@ that can build @Inner@, the second is an @Annotator@ that given the
-- full bytes, extracts just the bytes needed to decode @Inner@.
--
-- @
-- data Outer = Outer Text Inner
--
-- instance ToCBOR Outer where
--   toCBOR (Outer t i) = toCBOR t <> toCBOR i
--
-- instance FromCBOR (Annotator Outer) where
--   fromCBOR = do
--     t <- fromCBOR
--     (Annotator mkInner, Annotator extractInnerBytes) <- withSlice fromCBOR
--     pure (Annotator (\ full -> Outer t (mkInner (Full (extractInnerBytes full)))))
-- @
--
newtype Annotator a = Annotator { forall a. Annotator a -> FullByteString -> a
runAnnotator :: FullByteString -> a }
  deriving newtype (Applicative Annotator
forall a. a -> Annotator a
forall a b. Annotator a -> Annotator b -> Annotator b
forall a b. Annotator a -> (a -> Annotator b) -> Annotator b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
return :: forall a. a -> Annotator a
$creturn :: forall a. a -> Annotator a
>> :: forall a b. Annotator a -> Annotator b -> Annotator b
$c>> :: forall a b. Annotator a -> Annotator b -> Annotator b
>>= :: forall a b. Annotator a -> (a -> Annotator b) -> Annotator b
$c>>= :: forall a b. Annotator a -> (a -> Annotator b) -> Annotator b
Monad, Functor Annotator
forall a. a -> Annotator a
forall a b. Annotator a -> Annotator b -> Annotator a
forall a b. Annotator a -> Annotator b -> Annotator b
forall a b. Annotator (a -> b) -> Annotator a -> Annotator b
forall a b c.
(a -> b -> c) -> Annotator a -> Annotator b -> Annotator c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
<* :: forall a b. Annotator a -> Annotator b -> Annotator a
$c<* :: forall a b. Annotator a -> Annotator b -> Annotator a
*> :: forall a b. Annotator a -> Annotator b -> Annotator b
$c*> :: forall a b. Annotator a -> Annotator b -> Annotator b
liftA2 :: forall a b c.
(a -> b -> c) -> Annotator a -> Annotator b -> Annotator c
$cliftA2 :: forall a b c.
(a -> b -> c) -> Annotator a -> Annotator b -> Annotator c
<*> :: forall a b. Annotator (a -> b) -> Annotator a -> Annotator b
$c<*> :: forall a b. Annotator (a -> b) -> Annotator a -> Annotator b
pure :: forall a. a -> Annotator a
$cpure :: forall a. a -> Annotator a
Applicative, forall a b. a -> Annotator b -> Annotator a
forall a b. (a -> b) -> Annotator a -> Annotator b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> Annotator b -> Annotator a
$c<$ :: forall a b. a -> Annotator b -> Annotator a
fmap :: forall a b. (a -> b) -> Annotator a -> Annotator b
$cfmap :: forall a b. (a -> b) -> Annotator a -> Annotator b
Functor)

-- | The argument is a decoder for a annotator that needs access to the bytes that
-- | were decoded. This function constructs and supplies the relevant piece.
annotatorSlice :: Decoder s (Annotator (BSL.ByteString -> a)) -> Decoder s (Annotator a)
annotatorSlice :: forall s a.
Decoder s (Annotator (ByteString -> a)) -> Decoder s (Annotator a)
annotatorSlice Decoder s (Annotator (ByteString -> a))
dec = do
  (Annotator (ByteString -> a)
k, Annotator ByteString
bytes) <- forall s a. Decoder s a -> Decoder s (a, Annotator ByteString)
withSlice Decoder s (Annotator (ByteString -> a))
dec
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Annotator (ByteString -> a)
k forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Annotator ByteString
bytes

-- | Pairs the decoder result with an annotator that can be used to construct the exact bytes used to decode the result.
withSlice :: Decoder s a -> Decoder s (a, Annotator BSL.ByteString)
withSlice :: forall s a. Decoder s a -> Decoder s (a, Annotator ByteString)
withSlice Decoder s a
dec = do
  Annotated a
r ByteSpan
byteSpan <- forall s a. Decoder s a -> Decoder s (Annotated a ByteSpan)
annotatedDecoder Decoder s a
dec
  forall (m :: * -> *) a. Monad m => a -> m a
return (a
r, forall a. (FullByteString -> a) -> Annotator a
Annotator (\(Full ByteString
bsl) -> ByteString -> ByteSpan -> ByteString
slice ByteString
bsl ByteSpan
byteSpan))

-- | Supplies the bytestring argument to both the decoder and the produced annotator.
decodeAnnotator :: Text -> (forall s. Decoder s (Annotator a)) -> BSL.ByteString -> Either DecoderError a
decodeAnnotator :: forall a.
Text
-> (forall s. Decoder s (Annotator a))
-> ByteString
-> Either DecoderError a
decodeAnnotator Text
label' forall s. Decoder s (Annotator a)
decoder ByteString
bytes =
  (\Annotator a
x -> forall a. Annotator a -> FullByteString -> a
runAnnotator Annotator a
x (ByteString -> FullByteString
Full ByteString
bytes)) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a.
Text
-> (forall s. Decoder s a) -> ByteString -> Either DecoderError a
decodeFullDecoder Text
label' forall s. Decoder s (Annotator a)
decoder ByteString
bytes