-- SPDX-FileCopyrightText: 2020 Tocqueville Group
--
-- SPDX-License-Identifier: LicenseRef-MIT-TQ

{-# OPTIONS_GHC -Wno-redundant-constraints #-}

-- | Impementation of @Store@ - object incapsulating multiple 'BigMap's.
--
-- This module also provides template for the contract storage -
-- 'StorageSkeleton'.
--
-- We represent 'Store' as @big_map bytes (a | b | ...)@.
--
-- Key of this map is formed as @(index, orig_key)@, where @index@ is
-- zero-based index of emulated map, @orig_key@ is key of this emulated map.
--
-- Value of this map is just a union of emulated map's values.

{- Note on store inner representation (@martoon)

I see an alternative approach - representing store as
@big_map bytes (option a, option b, ...)@.

This would allow for saner implementation and more convenient interface.
An obvious shortcoming here is gas consumption. But this overhead seems
insignificant against the background of some other instructions.

-}
module Lorentz.Store
  {-# DEPRECATED "Contract storage can contain multiple big_maps starting from Michelson 005" #-}
  ( -- * Store and related type definitions
    Store (..)
  , type (|->)

    -- ** Type-lookup-by-name
  , GetStoreKey
  , GetStoreValue

    -- ** Instructions
  , storeMem
  , storeGet
  , storeUpdate
  , storeInsert
  , storeInsertNew
  , storeDelete

    -- ** Instruction constraints
  , StoreMemC
  , StoreGetC
  , StoreUpdateC
  , StoreInsertC
  , StoreDeleteC
  , HasStore
  , HasStoreForAllIn

    -- * Storage skeleton
  , StorageSkeleton (..)
  , storageUnpack
  , storagePack
  , storageMem
  , storageGet
  , storageInsert
  , storageInsertNew
  , storageDelete

    -- * Store management from Haskell
  , storePiece
  , storeKeyValueList
  , storeLookup

    -- ** Function constraints
  , StorePieceC
  ) where

import Data.Default (Default)
import qualified Data.Kind as Kind
import qualified Data.Map as Map
import Data.Type.Bool (If, type (||))
import Data.Type.Equality (type (==))
import GHC.Generics ((:+:))
import qualified GHC.Generics as G
import GHC.TypeLits (AppendSymbol, ErrorMessage(..), Symbol, TypeError)
import GHC.TypeNats (type (+), Nat)
import Type.Reflection ((:~:)(Refl))

import Lorentz.ADT
import Lorentz.Base
import Lorentz.Coercions
import Lorentz.Constraints
import Lorentz.Instr as L
import Lorentz.Macro
import Lorentz.Pack
import Lorentz.StoreClass
import Michelson.Typed.Haskell.Instr.Sum
import Michelson.Typed.Haskell.Value
import Michelson.Typed.Instr
import Util.Label (Label)

{-# ANN module ("HLint: ignore Use 'natVal' from Universum" :: Text) #-}

----------------------------------------------------------------------------
-- Store
----------------------------------------------------------------------------

-- | Gathers multple 'BigMap's under one object.
--
-- Type argument of this datatype stands for a "map template" -
-- a datatype with multiple constructors, each containing an object of
-- type '|->' and corresponding to single virtual 'BigMap'.
-- It's also possible to parameterize it with a larger type which is
-- a sum of types satisfying the above property.
--
-- Inside it keeps only one 'BigMap' thus not violating Michelson limitations.
--
-- See examples below.
newtype Store a = Store { unStore :: BigMap ByteString a }
  deriving stock (Eq, Show)
  deriving newtype (Default, Semigroup, Monoid, IsoValue)

-- | Describes one virtual big map.
data k |-> v = BigMapImage v
  deriving stock Generic
  deriving anyclass IsoValue

{- Again we use generic magic to implement methods for 'Store'
(and thus 'Store' type constructor accepts a datatype, not a type-level list).

There are two reasons for this:

1. This gives us expected balanced tree of 'Or's for free.

2. This allows us selecting a map by constructor name, not by
e.g. type of map value. This is subjective, but looks like a good thing
for me (@martoon). On the other hand, it prevents us from sharing the
same interface between maps and 'Store'.

-}

-- | Position of a constructor in the corresponding datatype declaration.
type CtorIdx = Nat

-- | Number of datatype constructors.
type CtorsNum = Nat

-- | Type arguments of '|->'.
data MapSignature = MapSignature Kind.Type Kind.Type CtorIdx

-- Again, we will use these getters instead of binding types within
-- 'MapSignature' using type equality because getters does not produce extra
-- compile errors on "field not found" cases.
type family MSKey ms where
  MSKey ('MapSignature k _ _) = k
type family MSValue ms where
  MSValue ('MapSignature _ v _) = v
type family MSCtorIdx ms where
  MSCtorIdx ('MapSignature _ _ ci) = ci

-- | Get map signature from the constructor with a given name.
type GetStore name a = MSRequireFound name a (GLookupStore name (G.Rep a))

data MapLookupRes
  = MapFound MapSignature
  | MapAbsent CtorsNum

type family MSRequireFound
  (name :: Symbol)
  (a :: Kind.Type)
  (mlr :: MapLookupRes)
    :: MapSignature where
  MSRequireFound _ _ ('MapFound ms) = ms
  MSRequireFound name a ('MapAbsent _) = TypeError
    ('Text "Failed to find store template: datatype " ':<>: 'ShowType a ':<>:
     'Text " has no constructor " ':<>: 'ShowType name)

-- | Prepend a constructor name with a lower-case character so that you
-- could make a label with @OverloadedLabels@ extension matching
-- resulting thing.
type CtorNameToLabel name = "c" `AppendSymbol` name

type family GLookupStore (name :: Symbol) (x :: Kind.Type -> Kind.Type)
              :: MapLookupRes where
  GLookupStore name (G.D1 _ x) = GLookupStore name x
  GLookupStore name (x :+: y) = LSMergeFound name (GLookupStore name x)
                                                  (GLookupStore name y)
  -- When we encounter a constructor there are two cases we are interested in:
  -- 1. This constructor has one field with type `|->`. Then we check its name
  -- and return 'MapFound' if it matches and 'MapAbsent' otherwise (storing
  -- information that we've found one constructor).
  -- 2. This constructor has one field with a different type. Then we expect
  -- this field to store '|->' somewhere deeper and try to find it there.
  GLookupStore name (G.C1 ('G.MetaCons ctorName _ _) x) =
    If (IsLeafCtor x)
      (If (name == ctorName || name == CtorNameToLabel ctorName)
         ('MapFound $ GExtractMapSignature ctorName x)
         ('MapAbsent 1)
      )
      (GLookupStoreDeeper name x)
  GLookupStore _ G.V1 = 'MapAbsent 0

-- Helper type family to check whether ADT constructor has one field
-- with type `|->`.
type family IsLeafCtor (x :: Kind.Type -> Kind.Type) :: Bool where
  IsLeafCtor (G.S1 _ (G.Rec0 (_ |-> _))) = 'True
  IsLeafCtor _ = 'False

-- Helper type family to go deeper during type-level store lookup.
type family GLookupStoreDeeper (name :: Symbol) (x :: Kind.Type -> Kind.Type)
              :: MapLookupRes where
  GLookupStoreDeeper name (G.S1 _ (G.Rec0 y))  = GLookupStore name (G.Rep y)
  GLookupStoreDeeper name _ = TypeError
    ('Text "Attempt to go deeper failed while looking for" ':<>: 'ShowType name
    ':$$:
    'Text "Make sure that all constructors have exactly one field inside.")

type family LSMergeFound (name :: Symbol)
  (f1 :: MapLookupRes) (f2 :: MapLookupRes)
  :: MapLookupRes where
  LSMergeFound _ ('MapAbsent n1) ('MapAbsent n2) = 'MapAbsent (n1 + n2)
  LSMergeFound _ ('MapFound ms) ('MapAbsent _) = 'MapFound ms
  LSMergeFound _ ('MapAbsent n) ('MapFound ('MapSignature k v i)) =
    'MapFound ('MapSignature k v (n + i))
  -- It's possible that there are two constructors with the same name,
  -- because main template pattern may be a sum of smaller template
  -- patterns with same constructor names.
  LSMergeFound ctor ('MapFound _) ('MapFound _) = TypeError
    ('Text "Found more than one constructor matching " ':<>: 'ShowType ctor)

type family GExtractMapSignature (ctor :: Symbol) (x :: Kind.Type -> Kind.Type)
             :: MapSignature where
  GExtractMapSignature _ (G.S1 _ (G.Rec0 (k |-> v))) = 'MapSignature k v 0
  GExtractMapSignature ctor _ = TypeError
    ('Text "Expected exactly one field of type `k |-> v`" ':$$:
     'Text "In constructor " ':<>: 'ShowType ctor)

type GetStoreKey store name = MSKey (GetStore name store)
type GetStoreValue store name = MSValue (GetStore name store)

packKey
  :: forall (idx :: CtorIdx) a s.
     (KnownNat idx, NicePackedValue a)
  => (a : s) :-> (ByteString : s)
packKey =
  withDict (nicePackedValueEvi @a) $
    push (natVal $ Proxy @idx) #
    pair @Natural @a #
    pack

wrapBigMapImage :: (v : s) :-> ((k |-> v) : s)
wrapBigMapImage = forcedCoerce_

unwrapBigMapImage :: ((k |-> v) : s) :-> (v : s)
unwrapBigMapImage = forcedCoerce_

type StoreOpC store name =
  ( NicePackedValue (MSKey (GetStore name store))
  , KnownNat (MSCtorIdx (GetStore name store))
  )

{- Note on store initialization:

It's not possible to create an empty store, because Michelson provides no way
to create a new empty @big_map@.
-}

storeMem
  :: forall store name s.
     (StoreMemC store name)
  => Label name
  -> GetStoreKey store name : Store store : s :-> Bool : s
storeMem _ =
  packKey @(MSCtorIdx (GetStore name store)) #
  I MEM

type StoreMemC store name = StoreOpC store name

storeGet
  :: forall store name s.
     StoreGetC store name
  => Label name
  -> GetStoreKey store name : Store store : s
       :-> Maybe (GetStoreValue store name) : s
storeGet label =
  packKey @(MSCtorIdx (GetStore name store)) #
  I GET #
  ifNone none (unwrapUnsafe_ @store label # unwrapBigMapImage # L.some)

type StoreGetC store name =
  ( StoreOpC store name
  , InstrUnwrapC store name
  , KnownValue (GetStoreValue store name)
  , CtorHasOnlyField name store
      (GetStoreKey store name |-> GetStoreValue store name)
  )

storeUpdate
  :: forall store name s.
     StoreUpdateC store name
  => Label name
  -> GetStoreKey store name
      : Maybe (GetStoreValue store name)
      : Store store
      : s
  :-> Store store : s
storeUpdate label =
  packKey @(MSCtorIdx (GetStore name store)) #
  dip (ifNone none (wrapBigMapImage # wrap_ @store label # L.some)) #
  I UPDATE

type StoreUpdateC store name =
  ( KnownValue store
  , StoreOpC store name
  , InstrWrapC store name
  , CtorHasOnlyField name store
      (GetStoreKey store name |-> GetStoreValue store name)
  )

storeInsert
  :: forall store name s.
     StoreInsertC store name
  => Label name
  -> GetStoreKey store name
      : GetStoreValue store name
      : Store store
      : s
  :-> Store store : s
storeInsert label =
  packKey @(MSCtorIdx (GetStore name store)) #
  dip (wrapBigMapImage # wrap_ @store label # L.some) #
  I UPDATE

type StoreInsertC store name =
  ( StoreOpC store name
  , InstrWrapC store name
  , CtorHasOnlyField name store
      (GetStoreKey store name |-> GetStoreValue store name)
  )

-- | Insert a key-value pair, but fail if it will overwrite some existing entry.
storeInsertNew
  :: forall store name s.
     StoreInsertC store name
  => Label name
  -> (forall s0 any. GetStoreKey store name : s0 :-> any)
  -> GetStoreKey store name
      : GetStoreValue store name
      : Store store
      : s
  :-> Store store : s
storeInsertNew label doFail =
  duupX @3 # duupX @2 # storeMem label #
  if_ doFail
      (storeInsert label)

storeDelete
  :: forall store name s.
     ( StoreDeleteC store name
     )
  => Label name
  -> GetStoreKey store name : Store store : s
     :-> Store store : s
storeDelete _ =
  packKey @(MSCtorIdx (GetStore name store)) #
  dip (none @store) #
  I UPDATE

type StoreDeleteC store name =
  ( StoreOpC store name
  , KnownValue store
  )

-- | This constraint can be used if a function needs to work with
-- /big/ store, but needs to know only about some part(s) of it.
--
-- It can use all Store operations for a particular name, key and
-- value without knowing whole template.
type HasStore name key value store =
   ( StoreGetC store name
   , StoreInsertC store name
   , StoreDeleteC store name
   , GetStoreKey store name ~ key
   , GetStoreValue store name ~ value
   , StorePieceC store name key value
   )

-- | Write down all sensisble constraints which given @store@ satisfies
-- and apply them to @constrained@.
--
-- This store should have '|->' datatype in its immediate fields,
-- no deep inspection is performed.
type HasStoreForAllIn store constrained =
  GForAllHasStore constrained (G.Rep store)

type family GForAllHasStore (store :: Kind.Type) (x :: Kind.Type -> Kind.Type)
            :: Constraint where
  GForAllHasStore store (G.D1 _ x) = GForAllHasStore store x
  GForAllHasStore store (x :+: y) = ( GForAllHasStore store x
                                    , GForAllHasStore store y )
  GForAllHasStore store (G.C1 ('G.MetaCons ctorName _ _)
                          (G.S1 _ (G.Rec0 (key |-> value)))) =
    HasStore (CtorNameToLabel ctorName) key value store
  GForAllHasStore _ (G.C1 _ _) = ()
  GForAllHasStore _ G.V1 = ()

-- Instances
----------------------------------------------------------------------------

instance ( StoreMemC store name, StoreGetC store name
         , StoreUpdateC store name
         , key ~ GetStoreKey store name, value ~ GetStoreValue store name
         ) =>
         StoreHasSubmap (Store store) name key value where
  storeSubmapOps = StoreSubmapOps
    { sopMem = storeMem
    , sopGet = storeGet
    , sopUpdate = storeUpdate
    , sopDelete = Just storeDelete
    , sopInsert = Just storeInsert
    }

-- Examples
----------------------------------------------------------------------------

data MyStoreTemplate
  = IntsStore (Integer |-> ())
  | BytesStore (ByteString |-> ByteString)
  deriving stock Generic
  deriving anyclass IsoValue

type MyStore = Store MyStoreTemplate

_sample1 :: Integer : MyStore : s :-> MyStore : s
_sample1 = storeDelete @MyStoreTemplate #cIntsStore

_sample2 :: ByteString : ByteString : MyStore : s :-> MyStore : s
_sample2 = storeInsert @MyStoreTemplate #cBytesStore

data MyStoreTemplate2
  = BoolsStore (Bool |-> Bool)
  | IntsStore2 (Integer |-> Integer)
  | IntsStore3 (Integer |-> Bool)
  deriving stock Generic
  deriving anyclass IsoValue

-- You must derive 'Generic' instance for all custom types, even
-- newtypes.
newtype MyNatural = MyNatural Natural
  deriving stock Generic
  deriving newtype (IsoValue)

data MyStoreTemplate3 = MyStoreTemplate3 (Natural |-> MyNatural)
  deriving stock Generic
  deriving anyclass IsoValue

data MyStoreTemplateBig
  = BigTemplatePart1 MyStoreTemplate
  | BigTemplatePart2 MyStoreTemplate2
  | BigTemplatePart3 MyStoreTemplate3
  deriving stock Generic
  deriving anyclass IsoValue

_MyStoreTemplateBigTextsStore ::
  GetStore "cBytesStore" MyStoreTemplateBig :~: 'MapSignature ByteString ByteString 1
_MyStoreTemplateBigTextsStore = Refl

_MyStoreTemplateBigBoolsStore ::
  GetStore "cBoolsStore" MyStoreTemplateBig :~: 'MapSignature Bool Bool 2
_MyStoreTemplateBigBoolsStore = Refl

_MyStoreTemplateBigMyStoreTemplate3 ::
  GetStore "cMyStoreTemplate3" MyStoreTemplateBig :~: 'MapSignature Natural MyNatural 5
_MyStoreTemplateBigMyStoreTemplate3 = Refl

_MyStoreBigHasAllStores
  :: HasStoreForAllIn MyStoreTemplate store
  => Dict ( HasStore "cIntsStore" Integer () store
          , HasStore "cBytesStore" ByteString ByteString store
          )
_MyStoreBigHasAllStores = Dict

type MyStoreBig = Store MyStoreTemplateBig

_sample3 :: Integer : MyStoreBig : s :-> MyStoreBig : s
_sample3 = storeDelete @MyStoreTemplateBig #cIntsStore2

_sample4 :: ByteString : MyStoreBig : s :-> Bool : s
_sample4 = storeMem @MyStoreTemplateBig #cBytesStore

_sample5 :: Natural : MyNatural : MyStoreBig : s :-> MyStoreBig : s
_sample5 = storeInsert @MyStoreTemplateBig #cMyStoreTemplate3

-- Example of 'HasStoreForAllIn' use.
-- This function will work with any @store@ which has 'MyStoreTemplate3' inside.
_sample6
  :: forall store s.
      HasStoreForAllIn MyStoreTemplate3 store
  => Natural : MyNatural : Store store : s :-> Store store : s
_sample6 = storeInsert @store #cMyStoreTemplate3

-- For instance, 'sample6' works for 'MyStoreBig'.
_sample6' :: Natural : MyNatural : MyStoreBig : s :-> MyStoreBig : s
_sample6' = _sample6

----------------------------------------------------------------------------
-- Storage skeleton
----------------------------------------------------------------------------

-- | Contract storage with @big_map@.
--
-- Due to Michelson constraints it is the only possible layout containing
-- @big_map@.
data StorageSkeleton storeTemplate other = StorageSkeleton
  { sMap :: Store storeTemplate
  , sFields :: other
  } deriving stock (Eq, Show, Generic)
    deriving anyclass Default

deriving anyclass instance (WellTypedIsoValue st, WellTypedIsoValue o) => IsoValue (StorageSkeleton st o)

-- | Unpack 'StorageSkeleton' into a pair.
storageUnpack :: StorageSkeleton store fields : s :-> (Store store, fields) : s
storageUnpack = forcedCoerce_

-- | Pack a pair into 'StorageSkeleton'.
storagePack :: (Store store, fields) : s :-> StorageSkeleton store fields : s
storagePack = forcedCoerce_

storageMem
  :: forall store name fields s.
     (StoreMemC store name)
  => Label name
  -> GetStoreKey store name : StorageSkeleton store fields : s :-> Bool : s
storageMem label = dip (storageUnpack # car) # storeMem label

storageGet
  :: forall store name fields s.
     StoreGetC store name
  => Label name
  -> GetStoreKey store name : StorageSkeleton store fields : s
       :-> Maybe (GetStoreValue store name) : s
storageGet label = dip (storageUnpack # car) # storeGet label

storageInsert
  :: forall store name fields s.
     StoreInsertC store name
  => Label name
  -> GetStoreKey store name
      : GetStoreValue store name
      : StorageSkeleton store fields
      : s
  :-> StorageSkeleton store fields : s
storageInsert label =
  dip (dip (storageUnpack # dup # car # dip cdr)) #
  storeInsert label #
  pair # storagePack

-- | Insert a key-value pair, but fail if it will overwrite some existing entry.
storageInsertNew
  :: forall store name fields s.
     StoreInsertC store name
  => Label name
  -> (forall s0 any. GetStoreKey store name : s0 :-> any)
  -> GetStoreKey store name
      : GetStoreValue store name
      : StorageSkeleton store fields
      : s
  :-> StorageSkeleton store fields : s
storageInsertNew label doFail =
  dip (dip (storageUnpack # dup # car # dip cdr)) #
  storeInsertNew label doFail #
  pair # storagePack

storageDelete
  :: forall store name fields s.
     ( StoreDeleteC store name
     )
  => Label name
  -> GetStoreKey store name : StorageSkeleton store fields : s
     :-> StorageSkeleton store fields : s
storageDelete label =
  dip (storageUnpack # dup # car # dip cdr) #
  storeDelete label #
  pair # storagePack

-- Instances
----------------------------------------------------------------------------

instance (StoreHasField other fname ftype, IsoValue store, IsoValue other) =>
         StoreHasField (StorageSkeleton store other) fname ftype where
  storeFieldOps = storeFieldOpsDeeper #sFields

instance ( StoreMemC store name, StoreGetC store name
         , StoreUpdateC store name
         , key ~ GetStoreKey store name, value ~ GetStoreValue store name
         , IsoValue other
         ) =>
         StoreHasSubmap (StorageSkeleton store other) name key value where
  storeSubmapOps = storeSubmapOpsDeeper #sMap

-- Examples
----------------------------------------------------------------------------

type MyStorage = StorageSkeleton MyStoreTemplate (Integer, ByteString)

-- You can access both Store...
_storageSample1 :: Integer : MyStorage : s :-> MyStorage : s
_storageSample1 = storageDelete @MyStoreTemplate #cIntsStore

-- and other fields of the storage created with 'StorageSkeleton'.
_storageSample2 :: MyStorage : s :-> Integer : s
_storageSample2 = toField #sFields # car

----------------------------------------------------------------------------
-- Store construction from Haskell
----------------------------------------------------------------------------

packHsKey
  :: forall ctorIdx key.
     (NicePackedValue key, KnownNat ctorIdx)
  => key -> ByteString
packHsKey key =
  lPackValue (natVal (Proxy @ctorIdx), key)

-- | Lift a key-value pair to 'Store'.
--
-- Further you can use 'Monoid' instance of @Store@ to make up large stores.
storePiece
  :: forall name store key value.
     StorePieceC store name key value
  => Label name
  -> key
  -> value
  -> Store store
storePiece label key val =
  Store . BigMap $ one
    ( packHsKey @(MSCtorIdx (GetStore name store)) key
    , hsWrap @store label (BigMapImage val)
    )

storeKeyValueList
  :: forall name store key value.
     StorePieceC store name key value
  => Label name
  -> [(key, value)]
  -> Store store
storeKeyValueList label keyValues =
  Store . BigMap . Map.fromList $
  Prelude.map (\(key, val) ->
                 ( packHsKey @(MSCtorIdx (GetStore name store)) key
                 , hsWrap @store label (BigMapImage val)
                 )) keyValues

type StorePieceC store name key value =
  ( key ~ GetStoreKey store name
  , value ~ GetStoreValue store name
  , NicePackedValue key
  , KnownNat (MSCtorIdx (GetStore name store))
  , InstrWrapC store name, Generic store
  , ExtractCtorField (GetCtorField store name) ~ (key |-> value)
  )

-- | Get a value from store by key.
--
-- It expects map to be consistent, otherwise call to this function fails
-- with error.
storeLookup
  :: forall name store key value ctorIdx.
     ( key ~ GetStoreKey store name
     , value ~ GetStoreValue store name
     , ctorIdx ~ MSCtorIdx (GetStore name store)
     , NicePackedValue key
     , KnownNat ctorIdx
     , InstrUnwrapC store name, Generic store
     , CtorOnlyField name store ~ (key |-> value)
     )
  => Label name
  -> key
  -> Store store
  -> Maybe value
storeLookup label key (Store (BigMap m)) =
  Map.lookup (packHsKey @ctorIdx key) m <&> \val ->
    case hsUnwrap label val of
      Nothing -> error "Invalid store, keys and values types \
                       \correspondence is violated"
      Just (BigMapImage x) -> x

-- Examples
----------------------------------------------------------------------------

_storeSample :: Store MyStoreTemplate
_storeSample = mconcat
  [ storePiece #cIntsStore 1 ()
  , storePiece #cBytesStore "a" "b"
  ]

_lookupSample :: Maybe ByteString
_lookupSample = storeLookup #cBytesStore "a" _storeSample

_storeSampleBig :: Store MyStoreTemplateBig
_storeSampleBig = mconcat
  [ storePiece #cIntsStore 1 ()
  , storePiece #cBoolsStore True True
  , storePiece #cIntsStore3 2 False
  ]

_lookupSampleBig :: Maybe Bool
_lookupSampleBig = storeLookup #cIntsStore3 2 _storeSampleBig