{-# 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,
  internsFromSet,
  internsFromMap,
  internsFromVMap,
  internMap,
  internSet,
  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 qualified Data.Set as Set (size)
import qualified Data.Set.Internal as Set (Set (..))
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)

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

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 #-}

internMap :: Ord k => k -> Map k a -> Maybe k
internMap :: forall k a. Ord k => k -> Map k a -> Maybe k
internMap k
k = Map k a -> Maybe k
go
  where
    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

internSet :: Ord a => a -> Set.Set a -> Maybe a
internSet :: forall a. Ord a => a -> Set a -> Maybe a
internSet a
k = Set a -> Maybe a
go
  where
    go :: Set a -> Maybe a
go Set a
Set.Tip = forall a. Maybe a
Nothing
    go (Set.Bin Int
_ a
kx Set a
l Set a
r) =
      case forall a. Ord a => a -> a -> Ordering
compare a
k a
kx of
        Ordering
LT -> Set a -> Maybe a
go Set a
l
        Ordering
GT -> Set a -> Maybe a
go Set a
r
        Ordering
EQ -> forall a. a -> Maybe a
Just a
kx

internsFromSet :: Ord k => Set.Set k -> Interns k
internsFromSet :: forall k. Ord k => Set k -> Interns k
internsFromSet Set k
s
  | forall a. Set a -> Int
Set.size Set k
s forall a. Eq a => a -> a -> Bool
== Int
0 = forall a. Monoid a => a
mempty
  | Bool
otherwise =
      forall a. [Intern a] -> Interns a
Interns
        [ Intern
            { internMaybe :: k -> Maybe k
internMaybe = (forall a. Ord a => a -> Set a -> Maybe a
`internSet` Set k
s)
            , internWeight :: Int
internWeight = forall a. Set a -> Int
Set.size Set k
s
            }
        ]

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 k a. Map k a -> Int
Map.size Map k a
m forall a. Eq a => a -> a -> Bool
== Int
0 = forall a. Monoid a => a
mempty
  | Bool
otherwise =
      forall a. [Intern a] -> Interns a
Interns
        [ Intern
            { internMaybe :: k -> Maybe k
internMaybe = (forall k a. Ord k => k -> Map k a -> Maybe k
`internMap` 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 (kv :: * -> *) k (vv :: * -> *) v.
Vector kv k =>
VMap kv vv k v -> Int
VMap.size VMap VB kv k a
m forall a. Eq a => a -> a -> Bool
== Int
0 = forall a. Monoid a => a
mempty
  | Bool
otherwise =
      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
            }
        ]

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) => DecShareCBOR (Set.Set k) where
  type Share (Set.Set k) = Interns k
  decShareCBOR :: forall s. Share (Set k) -> Decoder s (Set k)
decShareCBOR Share (Set k)
kis = forall a s. Ord a => Decoder s a -> Decoder s (Set a)
decodeSet (forall k. Interns k -> k -> k
interns Share (Set k)
kis forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a s. DecCBOR a => Decoder s a
decCBOR)
  getShare :: Set k -> Share (Set k)
getShare = forall k. Ord k => Set k -> Interns k
internsFromSet

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)