{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}

module Cardano.Ledger.Binary.Decoding.Sharing (
  DecShareCBOR (..),
  Interns (..),
  Intern (..),
  decShareLensCBOR,
  decSharePlusLensCBOR,
  decNoShareCBOR,
  interns,
  internsFromMap,
  internsFromVMap,
  toMemptyLens,
  decShareMonadCBOR,
)
where

import Cardano.Ledger.Binary.Decoding.DecCBOR
import Cardano.Ledger.Binary.Decoding.Decoder
import Control.Monad ((<$!>))
import Control.Monad.Trans
import Control.Monad.Trans.State.Strict
import qualified Data.Foldable as F
import Data.Kind
import qualified Data.Map.Strict as Map (size)
import Data.Map.Strict.Internal (Map (..))
import Data.Primitive.Types (Prim)
import Data.VMap (VB, VMap, VP)
import qualified Data.VMap as VMap
import Lens.Micro

-- =======================================

-- | This is an abstract interface that does the interning. In other words it
-- does the actual sharing by looking up the supplied value in some existing
-- data structure and uses that value instead. Relying on this interface gives us
-- the benefit of ignoring the type of underlying data structure and allows us
-- to compose many `Intern`s with the monoidal interface provided by `Interns`
-- wrapper. In order to create an `Intern` see the `internsFromMap` or
-- `internsFromVMap` functions.
data Intern a = Intern
  { forall a. Intern a -> a -> Maybe a
internMaybe :: a -> Maybe a
  -- ^ Function that will do the interning. If value is not available then
  -- `Nothing` is returned.
  , forall a. Intern a -> Int
internWeight :: !Int
  -- ^ Used for sorting. Normally set to the size of the underlying data
  -- structure. Keeping interns sorted with respect to how many elements
  -- is in the underlying data structure in theory gives a better chance of
  -- successful intern hit sooner rather than later.
  }

newtype Interns a = Interns [Intern a]
  deriving (Interns a
[Interns a] -> Interns a
Interns a -> Interns a -> Interns a
forall {a}. Semigroup (Interns a)
forall a. Interns a
forall a.
Semigroup a -> a -> (a -> a -> a) -> ([a] -> a) -> Monoid a
forall a. [Interns a] -> Interns a
forall a. Interns a -> Interns a -> Interns a
mconcat :: [Interns a] -> Interns a
$cmconcat :: forall a. [Interns a] -> Interns a
mappend :: Interns a -> Interns a -> Interns a
$cmappend :: forall a. Interns a -> Interns a -> Interns a
mempty :: Interns a
$cmempty :: forall a. Interns a
Monoid)

interns :: Interns k -> k -> k
interns :: forall k. Interns k -> k -> k
interns (Interns []) !k
k = k
k -- optimize for common case when there are no interns
interns (Interns [Intern k]
is) !k
k = [Intern k] -> k
go [Intern k]
is
  where
    go :: [Intern k] -> k
go [] = k
k
    go (Intern k
x : [Intern k]
xs) =
      case forall a. Intern a -> a -> Maybe a
internMaybe Intern k
x k
k of
        Just k
kx -> k
kx
        Maybe k
Nothing -> [Intern k] -> k
go [Intern k]
xs
{-# INLINE interns #-}

internsFromMap :: Ord k => Map k a -> Interns k
internsFromMap :: forall k a. Ord k => Map k a -> Interns k
internsFromMap Map k a
m =
  forall a. [Intern a] -> Interns a
Interns
    [ Intern
        { internMaybe :: k -> Maybe k
internMaybe = \k
k ->
            let go :: Map k a -> Maybe k
go Map k a
Tip = forall a. Maybe a
Nothing
                go (Bin Int
_ k
kx a
_ Map k a
l Map k a
r) =
                  case forall a. Ord a => a -> a -> Ordering
compare k
k k
kx of
                    Ordering
LT -> Map k a -> Maybe k
go Map k a
l
                    Ordering
GT -> Map k a -> Maybe k
go Map k a
r
                    Ordering
EQ -> forall a. a -> Maybe a
Just k
kx
             in Map k a -> Maybe k
go Map k a
m
        , internWeight :: Int
internWeight = forall k a. Map k a -> Int
Map.size Map k a
m
        }
    ]

internsFromVMap :: Ord k => VMap VB kv k a -> Interns k
internsFromVMap :: forall k (kv :: * -> *) a. Ord k => VMap VB kv k a -> Interns k
internsFromVMap VMap VB kv k a
m =
  forall a. [Intern a] -> Interns a
Interns
    [ Intern
        { internMaybe :: k -> Maybe k
internMaybe = \k
k -> forall (kv :: * -> *) k (vv :: * -> *) v.
(Vector kv k, Ord k) =>
k -> VMap kv vv k v -> Maybe k
VMap.internMaybe k
k VMap VB kv k a
m
        , internWeight :: Int
internWeight = forall (kv :: * -> *) k (vv :: * -> *) v.
Vector kv k =>
VMap kv vv k v -> Int
VMap.size VMap VB kv k a
m
        }
    ]

instance Semigroup (Interns a) where
  <> :: Interns a -> Interns a -> Interns a
(<>) Interns a
is1 (Interns []) = Interns a
is1
  (<>) (Interns []) Interns a
is2 = Interns a
is2
  (<>) (Interns [Intern a]
is1) (Interns [Intern a]
is2) =
    forall a. [Intern a] -> Interns a
Interns (forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
F.foldr forall {a}. Intern a -> [Intern a] -> [Intern a]
insertIntoSortedInterns [Intern a]
is2 [Intern a]
is1)
    where
      insertIntoSortedInterns :: Intern a -> [Intern a] -> [Intern a]
insertIntoSortedInterns Intern a
i [] = [Intern a
i]
      insertIntoSortedInterns Intern a
i (Intern a
a : [Intern a]
as)
        | forall a. Intern a -> Int
internWeight Intern a
a forall a. Ord a => a -> a -> Bool
> forall a. Intern a -> Int
internWeight Intern a
i = Intern a
a forall a. a -> [a] -> [a]
: Intern a -> [Intern a] -> [Intern a]
insertIntoSortedInterns Intern a
i [Intern a]
as
        | Bool
otherwise = Intern a
i forall a. a -> [a] -> [a]
: Intern a
a forall a. a -> [a] -> [a]
: [Intern a]
as

class Monoid (Share a) => DecShareCBOR a where
  {-# MINIMAL (decShareCBOR | decSharePlusCBOR) #-}
  type Share a :: Type
  type Share a = ()

  -- | Whenever `fromShareCBOR` is being used for defining the instance this
  -- function should return the state that can be added whenever user invokes
  -- `decSharePlusCBOR`. `mempty` is returned by default.
  getShare :: a -> Share a
  getShare a
_ = forall a. Monoid a => a
mempty

  -- | Utilize sharing when decoding, but do not add anything to the state for
  -- future sharing.
  decShareCBOR :: Share a -> Decoder s a
  decShareCBOR Share a
s = forall a b. (a, b) -> a
fst forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT forall a s. DecShareCBOR a => StateT (Share a) (Decoder s) a
decSharePlusCBOR Share a
s

  -- | Deserialize with sharing and add to the state that is used for sharing. Default
  -- implementation will add value returned by `getShare` for adding to the
  -- state.
  decSharePlusCBOR :: StateT (Share a) (Decoder s) a
  decSharePlusCBOR = do
    Share a
s <- forall (m :: * -> *) s. Monad m => StateT s m s
get
    a
x <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall a s. DecShareCBOR a => Share a -> Decoder s a
decShareCBOR Share a
s
    a
x forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ forall (m :: * -> *) s. Monad m => s -> StateT s m ()
put (forall a. DecShareCBOR a => a -> Share a
getShare a
x forall a. Semigroup a => a -> a -> a
<> Share a
s)

decShareLensCBOR ::
  DecShareCBOR b =>
  SimpleGetter bs (Share b) ->
  StateT bs (Decoder s) b
decShareLensCBOR :: forall b bs s.
DecShareCBOR b =>
SimpleGetter bs (Share b) -> StateT bs (Decoder s) b
decShareLensCBOR SimpleGetter bs (Share b)
l = do
  bs
s <- forall (m :: * -> *) s. Monad m => StateT s m s
get
  forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall a s. DecShareCBOR a => Share a -> Decoder s a
decShareCBOR (bs
s forall s a. s -> Getting a s a -> a
^. SimpleGetter bs (Share b)
l)

-- | Using this function it is possible to compose two lenses. One will extract
-- a value and another will used it for placing it into a empty monoid. Here is
-- an example of how a second element of a tuple can be projected on the third
-- element of a 3-tuple.
--
-- > toMemptyLens _3 _2 == lens (\(_, b) -> (mempty, mempty, b)) (\(a, _) (_, _, b) -> (a, b))
--
-- Here is an example where we extract a second element of a tuple and insert it at
-- third position of a three tuple while all other elements are set to `mempty`:
--
-- >>> import Lens.Micro
-- >>> ("foo","bar") ^. toMemptyLens _3 _2 :: (Maybe String, (), String)
-- (Nothing,(),"bar")
--
-- In the opposite direction of extracting the third element of a 3-tuple and
-- replacing the second element of the tuple the setter is being applied to
--
-- >>> ("foo","bar") & toMemptyLens _3 _2 .~ (Just "baz", (), "booyah") :: (String, String)
-- ("foo","booyah")
toMemptyLens :: Monoid a => Lens' a b -> Lens' c b -> Lens' c a
toMemptyLens :: forall a b c. Monoid a => Lens' a b -> Lens' c b -> Lens' c a
toMemptyLens Lens' a b
lto Lens' c b
lfrom =
  forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens (\c
s -> forall a. Monoid a => a
mempty forall a b. a -> (a -> b) -> b
& Lens' a b
lto forall s t a b. ASetter s t a b -> b -> s -> t
.~ (c
s forall s a. s -> Getting a s a -> a
^. Lens' c b
lfrom)) (\c
s a
a -> c
s forall a b. a -> (a -> b) -> b
& Lens' c b
lfrom forall s t a b. ASetter s t a b -> b -> s -> t
.~ (a
a forall s a. s -> Getting a s a -> a
^. Lens' a b
lto))

-- | Just like `decSharePlusCBOR`, except allows to transform the shared state
-- with a lens.
decSharePlusLensCBOR ::
  DecShareCBOR b =>
  Lens' bs (Share b) ->
  StateT bs (Decoder s) b
decSharePlusLensCBOR :: forall b bs s.
DecShareCBOR b =>
Lens' bs (Share b) -> StateT bs (Decoder s) b
decSharePlusLensCBOR Lens' bs (Share b)
l = do
  bs
s <- forall (m :: * -> *) s. Monad m => StateT s m s
get
  (b
x, Share b
k) <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT forall a s. DecShareCBOR a => StateT (Share a) (Decoder s) a
decSharePlusCBOR (bs
s forall s a. s -> Getting a s a -> a
^. Lens' bs (Share b)
l)
  b
x forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ forall (m :: * -> *) s. Monad m => s -> StateT s m ()
put (bs
s forall a b. a -> (a -> b) -> b
& Lens' bs (Share b)
l forall s t a b. ASetter s t a b -> b -> s -> t
.~ Share b
k)

-- | Use `DecShareCBOR` class while ignoring sharing
decNoShareCBOR :: DecShareCBOR a => Decoder s a
decNoShareCBOR :: forall a s. DecShareCBOR a => Decoder s a
decNoShareCBOR = forall a s. DecShareCBOR a => Share a -> Decoder s a
decShareCBOR forall a. Monoid a => a
mempty

instance (Ord k, DecCBOR k, DecCBOR v) => DecShareCBOR (Map k v) where
  type Share (Map k v) = (Interns k, Interns v)
  decShareCBOR :: forall s. Share (Map k v) -> Decoder s (Map k v)
decShareCBOR (Interns k
kis, Interns v
vis) = do
    forall k s v.
Ord k =>
Decoder s k -> Decoder s v -> Decoder s (Map k v)
decodeMap (forall k. Interns k -> k -> k
interns Interns k
kis forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a s. DecCBOR a => Decoder s a
decCBOR) (forall k. Interns k -> k -> k
interns Interns v
vis forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a s. DecCBOR a => Decoder s a
decCBOR)
  getShare :: Map k v -> Share (Map k v)
getShare !Map k v
m = (forall k a. Ord k => Map k a -> Interns k
internsFromMap Map k v
m, forall a. Monoid a => a
mempty)

instance (Ord k, DecCBOR k, DecCBOR v) => DecShareCBOR (VMap VB VB k v) where
  type Share (VMap VB VB k v) = (Interns k, Interns v)
  decShareCBOR :: forall s. Share (VMap VB VB k v) -> Decoder s (VMap VB VB k v)
decShareCBOR (Interns k
kis, Interns v
vis) = do
    forall (kv :: * -> *) k (vv :: * -> *) v s.
(Vector kv k, Vector vv v, Ord k) =>
Decoder s k -> Decoder s v -> Decoder s (VMap kv vv k v)
decodeVMap (forall k. Interns k -> k -> k
interns Interns k
kis forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a s. DecCBOR a => Decoder s a
decCBOR) (forall k. Interns k -> k -> k
interns Interns v
vis forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a s. DecCBOR a => Decoder s a
decCBOR)
  getShare :: VMap VB VB k v -> Share (VMap VB VB k v)
getShare !VMap VB VB k v
m = (forall k (kv :: * -> *) a. Ord k => VMap VB kv k a -> Interns k
internsFromVMap VMap VB VB k v
m, forall a. Monoid a => a
mempty)

instance (Ord k, DecCBOR k, DecCBOR v, Prim v) => DecShareCBOR (VMap VB VP k v) where
  type Share (VMap VB VP k v) = Interns k
  decShareCBOR :: forall s. Share (VMap VB VP k v) -> Decoder s (VMap VB VP k v)
decShareCBOR Share (VMap VB VP k v)
kis = do
    forall (kv :: * -> *) k (vv :: * -> *) v s.
(Vector kv k, Vector vv v, Ord k) =>
Decoder s k -> Decoder s v -> Decoder s (VMap kv vv k v)
decodeVMap (forall k. Interns k -> k -> k
interns Share (VMap VB VP k v)
kis forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a s. DecCBOR a => Decoder s a
decCBOR) forall a s. DecCBOR a => Decoder s a
decCBOR
  getShare :: VMap VB VP k v -> Share (VMap VB VP k v)
getShare !VMap VB VP k v
m = forall k (kv :: * -> *) a. Ord k => VMap VB kv k a -> Interns k
internsFromVMap VMap VB VP k v
m

-- | Share every item in a functor, have deserializing it
decShareMonadCBOR :: (DecCBOR (f b), Monad f) => Interns b -> Decoder s (f b)
decShareMonadCBOR :: forall (f :: * -> *) b s.
(DecCBOR (f b), Monad f) =>
Interns b -> Decoder s (f b)
decShareMonadCBOR Interns b
kis = do
  f b
sm <- forall a s. DecCBOR a => Decoder s a
decCBOR
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall k. Interns k -> k -> k
interns Interns b
kis forall (m :: * -> *) a b. Monad m => (a -> b) -> m a -> m b
<$!> f b
sm)