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

{-# LANGUAGE FunctionalDependencies #-}

-- | This module provides storage interfaces.
module Lorentz.StoreClass
  ( -- * Class
    StoreHasField (..)
  , StoreFieldOps (..)
  , StoreHasSubmap (..)
  , StoreSubmapOps (..)
  , StoreHasEntrypoint (..)
  , StoreEntrypointOps (..)

    -- * Useful type synonyms
  , EntrypointLambda
  , EntrypointsField

    -- * Expressing constraints on storage
  , type (~>)
  , type (::->)
  , StorageContains

    -- * Methods to work with storage
  , stToField
  , stGetField
  , stSetField
  , stMem
  , stGet
  , stUpdate
  , stDelete
  , stInsert
  , stInsertNew
  , stEntrypoint
  , stToEpLambda
  , stGetEpLambda
  , stSetEpLambda
  , stToEpStore
  , stGetEpStore
  , stSetEpStore

    -- * Implementations
  , storeFieldOpsADT
  , storeEntrypointOpsADT
  , storeEntrypointOpsFields
  , storeEntrypointOpsSubmapField
  , storeFieldOpsDeeper
  , storeSubmapOpsDeeper
  , storeEntrypointOpsDeeper
  , storeFieldOpsReferTo
  , storeSubmapOpsReferTo
  , storeEntrypointOpsReferTo
  , composeStoreFieldOps
  , composeStoreSubmapOps
  , composeStoreEntrypointOps

    -- * Storage generation
  , mkStoreEp
  ) where

import Data.Map (singleton)

import Lorentz.ADT
import Lorentz.Base
import Lorentz.Errors (failUnexpected)
import Lorentz.Constraints
import qualified Lorentz.Instr as L
import qualified Lorentz.Macro as L
import Lorentz.Value
import Michelson.Text (labelToMText)

----------------------------------------------------------------------------
-- Fields
----------------------------------------------------------------------------

-- | Datatype containing the full implementation of 'StoreHasField' typeclass.
--
-- We use this grouping because in most cases the implementation will be chosen
-- among the default ones, and initializing all methods at once is simpler
-- and more consistent.
-- (One can say that we are trying to emulate the @DerivingVia@ extension.)
data StoreFieldOps store fname ftype = StoreFieldOps
  { sopToField
      :: forall s.
         Label fname -> store : s :-> ftype : s
  , sopSetField
      :: forall s.
         Label fname -> ftype : store : s :-> store : s
  }

-- Using fundeps here for the sake of less amount of boilerplate on user side,
-- switch to type families if having any issues with that.
-- | Provides operations on fields for storage.
class StoreHasField store fname ftype | store fname -> ftype where
  storeFieldOps :: StoreFieldOps store fname ftype

-- | Pick storage field.
stToField
  :: StoreHasField store fname ftype
  => Label fname -> store : s :-> ftype : s
stToField = sopToField storeFieldOps

-- | Get storage field, preserving the storage itself on stack.
stGetField
  :: StoreHasField store fname ftype
  => Label fname -> store : s :-> ftype : store : s
stGetField l = L.dup # sopToField storeFieldOps l

-- | Update storage field.
stSetField
  :: StoreHasField store fname ftype
  => Label fname -> ftype : store : s :-> store : s
stSetField = sopSetField storeFieldOps

----------------------------------------------------------------------------
-- Virtual big maps
----------------------------------------------------------------------------

-- | Datatype containing the full implementation of 'StoreHasSubmap' typeclass.
--
-- We use this grouping because in most cases the implementation will be chosen
-- among the default ones, and initializing all methods at once is simpler
-- and more consistent.
-- (One can say that we are trying to emulate the @DerivingVia@ extension.)
data StoreSubmapOps store mname key value = StoreSubmapOps
  { sopMem
      :: forall s.
         Label mname -> key : store : s :-> Bool : s
  , sopGet
      :: forall s.
         (KnownValue value)
      => Label mname -> key : store : s :-> Maybe value : s
  , sopUpdate
      :: forall s.
         Label mname -> key : Maybe value : store : s :-> store : s

    -- Methods below are derivatives of methods above, they can be provided
    -- if for given specific storage type more efficient implementation is
    -- available.
  , sopDelete
      :: forall s.
         Maybe (Label mname -> key : store : s :-> store : s)
  , sopInsert
      :: forall s.
         Maybe (Label mname -> key : value : store : s :-> store : s)
  }

-- | Provides operations on submaps of storage.
class StoreHasSubmap store mname key value | store mname -> key value where
  storeSubmapOps :: StoreSubmapOps store mname key value

-- | Check value presence in storage.
stMem
  :: StoreHasSubmap store mname key value
  => Label mname -> key : store : s :-> Bool : s
stMem = sopMem storeSubmapOps

-- | Get value in storage.
stGet
  :: (StoreHasSubmap store mname key value, KnownValue value)
  => Label mname -> key : store : s :-> Maybe value : s
stGet = sopGet storeSubmapOps

-- | Update a value in storage.
stUpdate
  :: StoreHasSubmap store mname key value
  => Label mname -> key : Maybe value : store : s :-> store : s
stUpdate = sopUpdate storeSubmapOps

-- | Delete a value in storage.
stDelete
  :: forall store mname key value s.
     (StoreHasSubmap store mname key value, KnownValue value)
  => Label mname -> key : store : s :-> store : s
stDelete l = case sopDelete storeSubmapOps of
  Just delOp -> delOp l
  Nothing -> L.dip L.none # stUpdate l

-- | Add a value in storage.
stInsert
  :: StoreHasSubmap store mname key value
  => Label mname -> key : value : store : s :-> store : s
stInsert l = case sopInsert storeSubmapOps of
  Just insOp -> insOp l
  Nothing -> L.dip L.some # stUpdate l

-- | Add a value in storage, but fail if it will overwrite some existing entry.
stInsertNew
  :: StoreHasSubmap store mname key value
  => Label mname
  -> (forall s0 any. key : s0 :-> any)
  -> key : value : store : s
  :-> store : s
stInsertNew l doFail =
  L.duupX @3 # L.duupX @2 # stMem l #
  L.if_ doFail (stInsert l)

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

-- | 'BigMap' can be used as standalone key-value storage,
-- name of submap is not accounted in this case.
instance (key ~ key', value ~ value', NiceComparable key) =>
         StoreHasSubmap (BigMap key' value') name key value where
  storeSubmapOps = StoreSubmapOps
    { sopMem = \_label -> L.mem
    , sopGet = \_label -> L.get
    , sopUpdate = \_label -> L.update
    , sopDelete = Nothing
    , sopInsert = Nothing
    }

-- | 'Map' can be used as standalone key-value storage if very needed.
instance (key ~ key', value ~ value', NiceComparable key) =>
         StoreHasSubmap (Map key' value') name key value where
  storeSubmapOps = StoreSubmapOps
    { sopMem = \_label -> L.mem
    , sopGet = \_label -> L.get
    , sopUpdate = \_label -> L.update
    , sopDelete = Nothing
    , sopInsert = Nothing
    }

----------------------------------------------------------------------------
-- Stored Entrypoints
----------------------------------------------------------------------------

-- | Type synonym for a 'Lambda' that can be used as an entrypoint
type EntrypointLambda param store = Lambda (param, store) ([Operation], store)

-- | Type synonym of a 'BigMap' mapping 'MText' (entrypoint names) to
-- 'EntrypointLambda'.
--
-- This is useful when defining instances of 'StoreHasEntrypoint' as a storage
-- field containing one or more entrypoints (lambdas) of the same type.
type EntrypointsField param store = BigMap MText (EntrypointLambda param store)

-- | Datatype containing the full implementation of 'StoreHasEntrypoint' typeclass.
--
-- We use this grouping because in most cases the implementation will be chosen
-- among the default ones, and initializing all methods at once is simpler
-- and more consistent.
-- (One can say that we are trying to emulate the @DerivingVia@ extension.)
data StoreEntrypointOps store epName epParam epStore = StoreEntrypointOps
  { sopToEpLambda
      :: forall s.
         Label epName
      -> store : s :-> (EntrypointLambda epParam epStore) : s
  , sopSetEpLambda
      :: forall s.
         Label epName
      -> (EntrypointLambda epParam epStore) : store : s :-> store : s
  , sopToEpStore
      :: forall s.
         Label epName
      -> store : s :-> epStore : s
  , sopSetEpStore
      :: forall s.
         Label epName
      -> epStore : store : s :-> store : s
  }

-- | Provides operations on stored entrypoints.
--
-- @store@ is the storage containing both the entrypoint @epName@ (note: it has
-- to be in a 'BigMap' to take advantage of lazy evaluation) and the @epStore@
-- field this operates on.
class StoreHasEntrypoint store epName epParam epStore | store epName -> epParam epStore where
  storeEpOps :: StoreEntrypointOps store epName epParam epStore

-- | Extracts and executes the @epName@ entrypoint lambda from storage, returing
-- the updated full storage (@store@) and the produced 'Operation's.
stEntrypoint
  :: StoreHasEntrypoint store epName epParam epStore
  => Label epName -> epParam : store : s :-> ([Operation], store) : s
stEntrypoint l =
  L.dip (L.dup # stGetEpLambda l # L.swap # stToEpStore l) #
  L.pair # L.exec # L.unpair #
  L.dip (stSetEpStore l) # L.pair

-- | Pick stored entrypoint lambda.
stToEpLambda
  :: StoreHasEntrypoint store epName epParam epStore
  => Label epName -> store : s :-> (EntrypointLambda epParam epStore) : s
stToEpLambda = sopToEpLambda storeEpOps

-- | Get stored entrypoint lambda, preserving the storage itself on the stack.
stGetEpLambda
  :: StoreHasEntrypoint store epName epParam epStore
  => Label epName -> store : s :-> (EntrypointLambda epParam epStore) : store : s
stGetEpLambda l = L.dup # stToEpLambda l

-- | Stores the entrypoint lambda in the storage. Fails if already set.
stSetEpLambda
  :: StoreHasEntrypoint store epName epParam epStore
  => Label epName -> (EntrypointLambda epParam epStore) : store : s :-> store : s
stSetEpLambda = sopSetEpLambda storeEpOps

-- | Pick the sub-storage that the entrypoint operates on.
stToEpStore
  :: StoreHasEntrypoint store epName epParam epStore
  => Label epName -> store : s :-> epStore : s
stToEpStore = sopToEpStore storeEpOps

-- | Get the sub-storage that the entrypoint operates on, preserving the storage
-- itself on the stack.
stGetEpStore
  :: StoreHasEntrypoint store epName epParam epStore
  => Label epName -> store : s :-> epStore : store : s
stGetEpStore l = L.dup # stToEpStore l

-- | Update the sub-storage that the entrypoint operates on.
stSetEpStore
  :: StoreHasEntrypoint store epName epParam epStore
  => Label epName -> epStore : store : s :-> store : s
stSetEpStore = sopSetEpStore storeEpOps

----------------------------------------------------------------------------
-- Implementations
----------------------------------------------------------------------------

-- | Implementation of 'StoreHasField' for case of datatype
-- keeping a pack of fields.
storeFieldOpsADT
  :: HasFieldOfType dt fname ftype
  => StoreFieldOps dt fname ftype
storeFieldOpsADT = StoreFieldOps
  { sopToField = toField
  , sopSetField = setField
  }

-- | Implementation of 'StoreHasEntrypoint' for a datatype keeping a pack of
-- fields, among which one has contains the entrypoint and another is what such
-- entrypoint operates on.
storeEntrypointOpsADT
  :: ( HasFieldOfType store epmName (EntrypointsField epParam epStore)
     , HasFieldOfType store epsName epStore
     , KnownValue epParam, KnownValue epStore
     )
  => Label epmName -> Label epsName
  -> StoreEntrypointOps store epName epParam epStore
storeEntrypointOpsADT mapLabel fieldLabel = StoreEntrypointOps
  { sopToEpLambda = \l -> toField mapLabel # pushStEp l # L.get # someStEp l
  , sopSetEpLambda = \l -> L.dip (getField mapLabel) # setStEp mapLabel l # setField mapLabel
  , sopToEpStore = \_l -> toField fieldLabel
  , sopSetEpStore = \_l -> setField fieldLabel
  }

-- | Implementation of 'StoreHasEntrypoint' for a datatype that has a 'StoreHasField'
-- for an 'EntrypointsField' that contains the entrypoint and a 'StoreHasField'
-- for the field such entrypoint operates on.
storeEntrypointOpsFields
  :: ( StoreHasField store epmName (EntrypointsField epParam epStore)
     , StoreHasField store epsName epStore
     , KnownValue epParam, KnownValue epStore
     )
  => Label epmName -> Label epsName
  -> StoreEntrypointOps store epName epParam epStore
storeEntrypointOpsFields mapLabel fieldLabel = StoreEntrypointOps
  { sopToEpLambda = \l -> stToField mapLabel # pushStEp l # L.get # someStEp l
  , sopSetEpLambda = \l -> L.dip (stGetField mapLabel) # setStEp mapLabel l # stSetField mapLabel
  , sopToEpStore = \_l -> stToField fieldLabel
  , sopSetEpStore = \_l -> stSetField fieldLabel
  }

-- | Implementation of 'StoreHasEntrypoint' for a datatype that has a 'StoreHasSubmap'
-- that contains the entrypoint and a 'StoreHasField' for the field such
-- entrypoint operates on.
storeEntrypointOpsSubmapField
  :: ( StoreHasSubmap store epmName MText (EntrypointLambda epParam epStore)
     , StoreHasField store epsName epStore
     , KnownValue epParam, KnownValue epStore
     )
  => Label epmName -> Label epsName
  -> StoreEntrypointOps store epName epParam epStore
storeEntrypointOpsSubmapField mapLabel fieldLabel = StoreEntrypointOps
  { sopToEpLambda = \l -> pushStEp l # stGet mapLabel # someStEp l
  , sopSetEpLambda = \l -> setStEp mapLabel l
  , sopToEpStore = \_l -> stToField fieldLabel
  , sopSetEpStore = \_l -> stSetField fieldLabel
  }

-- | Implementation of 'StoreHasField' for a data type which has an
-- instance of 'StoreHasField' inside.
-- For instance, it can be used for top-level storage.
storeFieldOpsDeeper
  :: ( HasFieldOfType storage fieldsPartName fields
     , StoreHasField fields fname ftype
     )
  => Label fieldsPartName
  -> StoreFieldOps storage fname ftype
storeFieldOpsDeeper fieldsLabel =
  composeStoreFieldOps fieldsLabel storeFieldOpsADT storeFieldOps

-- | Implementation of 'StoreHasSubmap' for a data type which has an
-- instance of 'StoreHasSubmap' inside.
-- For instance, it can be used for top-level storage.
storeSubmapOpsDeeper
  :: ( HasFieldOfType storage bigMapPartName fields
     , StoreHasSubmap fields mname key value
     )
  => Label bigMapPartName
  -> StoreSubmapOps storage mname key value
storeSubmapOpsDeeper submapLabel =
  composeStoreSubmapOps submapLabel storeFieldOpsADT storeSubmapOps

-- | Implementation of 'StoreHasEntrypoint' for a data type which has an
-- instance of 'StoreHasEntrypoint' inside.
-- For instance, it can be used for top-level storage.
storeEntrypointOpsDeeper
  :: ( HasFieldOfType store nameInStore substore
     , StoreHasEntrypoint substore epName epParam epStore
     )
  => Label nameInStore
  -> StoreEntrypointOps store epName epParam epStore
storeEntrypointOpsDeeper fieldsLabel =
  composeStoreEntrypointOps fieldsLabel storeFieldOpsADT storeEpOps

{- | Pretend that given 'StoreSubmapOps' implementation is made up
for submap with name @desiredName@, not its actual name.
Logic of the implementation remains the same.

Use case: imagine that your code requires access to submap named @X@,
but in your storage that submap is called @Y@.
Then you implement the instance which makes @X@ refer to @Y@:

@
instance StoreHasSubmap Store X Key Value where
  storeSubmapOps = storeSubmapOpsReferTo #Y storeSubmapOpsForY
@
-}
storeSubmapOpsReferTo
  :: Label name
  -> StoreSubmapOps storage name key value
  -> StoreSubmapOps storage desiredName key value
storeSubmapOpsReferTo l StoreSubmapOps{..} =
  StoreSubmapOps
  { sopMem = \_l -> sopMem l
  , sopGet = \_l -> sopGet l
  , sopUpdate = \_l -> sopUpdate l
  , sopDelete = (\op _l -> op l) <$> sopDelete
  , sopInsert = (\op _l -> op l) <$> sopInsert
  }

-- | Pretend that given 'StoreFieldOps' implementation is made up
-- for field with name @desiredName@, not its actual name.
-- Logic of the implementation remains the same.
--
-- See also 'storeSubmapOpsReferTo'.
storeFieldOpsReferTo
  :: Label name
  -> StoreFieldOps storage name field
  -> StoreFieldOps storage desiredName field
storeFieldOpsReferTo l StoreFieldOps{..} =
  StoreFieldOps
  { sopToField = \_l -> sopToField l
  , sopSetField = \_l -> sopSetField l
  }

-- | Pretend that given 'StoreEntrypointOps' implementation is made up
-- for entrypoint with name @desiredName@, not its actual name.
-- Logic of the implementation remains the same.
--
-- See also 'storeSubmapOpsReferTo'.
storeEntrypointOpsReferTo
  :: Label epName
  -> StoreEntrypointOps store epName epParam epStore
  -> StoreEntrypointOps store desiredName epParam epStore
storeEntrypointOpsReferTo l StoreEntrypointOps{..} = StoreEntrypointOps
  { sopToEpLambda = \_l -> sopToEpLambda l
  , sopSetEpLambda = \_l -> sopSetEpLambda l
  , sopToEpStore = \_l -> sopToEpStore l
  , sopSetEpStore = \_l -> sopSetEpStore l
  }

-- | Chain two implementations of field operations.
--
-- Suits for a case when your store does not contain its fields directly
-- rather has a nested structure.
composeStoreFieldOps
  :: Label nameInStore
  -> StoreFieldOps store nameInStore substore
  -> StoreFieldOps substore nameInSubstore field
  -> StoreFieldOps store nameInSubstore field
composeStoreFieldOps l1 ops1 ops2 =
  StoreFieldOps
  { sopToField = \l2 ->
      sopToField ops1 l1 # sopToField ops2 l2
  , sopSetField = \l2 ->
      L.dip (L.dup # sopToField ops1 l1) #
      sopSetField ops2 l2 #
      sopSetField ops1 l1
  }

-- | Chain implementations of field and submap operations.
composeStoreSubmapOps
  :: Label nameInStore
  -> StoreFieldOps store nameInStore substore
  -> StoreSubmapOps substore mname key value
  -> StoreSubmapOps store mname key value
composeStoreSubmapOps l1 ops1 ops2 =
  StoreSubmapOps
  { sopMem = \l2 ->
      L.dip (sopToField ops1 l1) # sopMem ops2 l2
  , sopGet = \l2 ->
      L.dip (sopToField ops1 l1) # sopGet ops2 l2
  , sopUpdate = \l2 ->
      L.dip (L.dip (L.dup # sopToField ops1 l1)) #
      sopUpdate ops2 l2 #
      sopSetField ops1 l1
  , sopDelete = case sopDelete ops2 of
      Nothing -> Nothing
      Just delOp -> Just $ \l2 ->
        L.dip (L.dup # sopToField ops1 l1) #
        delOp l2 #
        sopSetField ops1 l1
  , sopInsert = case sopInsert ops2 of
      Nothing -> Nothing
      Just insOp -> Just $ \l2 ->
        L.dip (L.dip (L.dup # sopToField ops1 l1)) #
        insOp l2 #
        sopSetField ops1 l1
  }

composeStoreEntrypointOps
  :: Label nameInStore
  -> StoreFieldOps store nameInStore substore
  -> StoreEntrypointOps substore epName epParam epStore
  -> StoreEntrypointOps store epName epParam epStore
composeStoreEntrypointOps l1 ops1 ops2 = StoreEntrypointOps
  { sopToEpLambda = \l2 ->
      sopToField ops1 l1 # sopToEpLambda ops2 l2
  , sopSetEpLambda = \l2 ->
      L.dip (L.dup # sopToField ops1 l1) #
      sopSetEpLambda ops2 l2 #
      sopSetField ops1 l1
  , sopToEpStore = \l2 ->
      sopToField ops1 l1 # sopToEpStore ops2 l2
  , sopSetEpStore = \l2 ->
      L.dip (L.dup # sopToField ops1 l1) #
      sopSetEpStore ops2 l2 #
      sopSetField ops1 l1
  }

-- | Utility to 'push' the 'MText' name of and entrypoint from its 'Label'
pushStEp :: Label name -> s :-> MText : s
pushStEp = L.push . labelToMText

-- | Utility to extract an 'EntrypointLambda' from a 'Maybe', fails in case of
-- 'Nothing'.
someStEp
  :: Label epName
  -> Maybe (EntrypointLambda epParam epStore) : s :-> (EntrypointLambda epParam epStore) : s
someStEp l = L.ifSome L.nop $
  failUnexpected ([mt|unknown storage entrypoint: |] <> labelToMText l)

-- | Utility to set an 'EntrypointLambda' into a store.
-- Fails in case the entrypoint is already set.
setStEp
  :: StoreHasSubmap store epmName MText (EntrypointLambda epParam epStore)
  => Label epmName -> Label epsName
  -> (EntrypointLambda epParam epStore) : store : s :-> store : s
setStEp ml l = pushStEp l # stInsertNew ml failAlreadySetEp
  where
    failAlreadySetEp :: MText : s :-> any
    failAlreadySetEp =
      L.push [mt|Storage entrypoint already set: |] # L.concat # L.failWith

----------------------------------------------------------------------------
-- Storage generation
----------------------------------------------------------------------------

-- Note: we could make this safer with a 'StoreHasEntrypoint' constraint, but GHC
-- would flag it as redundant and we'd also need to annotate the @store@
--
-- | Utility to create 'EntrypointsField's from an entrypoint name (@epName@) and
-- an 'EntrypointLambda' implementation. Note that you need to merge multiple of
-- these (with '<>') if your field contains more than one entrypoint lambda.
mkStoreEp
  :: Label epName
  -> EntrypointLambda epParam epStore
  -> EntrypointsField epParam epStore
mkStoreEp l = BigMap . singleton (labelToMText l)


----------------------------------------------------------------------------
-- Utilities
----------------------------------------------------------------------------

-- | Indicates a submap with given key and value types.
data k ~> v
infix 9 ~>

-- | Indicates a stored entrypoint with the given @param@ and @store@ types.
data param ::-> store
infix 9 ::->

{- | Concise way to write down constraints with expected content of a storage.

Use it like follows:

@
type StorageConstraint store = StorageContains store
  [ "fieldInt" := Int
  , "fieldNat" := Nat
  , "epsToNat" := Int ::-> Nat
  , "balances" := Address ~> Int
  ]
@

-}
type family StorageContains store (content :: [NamedField]) :: Constraint where
  StorageContains _ '[] = ()
  StorageContains store ((n := k ~> v) ': ct) =
    (StoreHasSubmap store n k v, StorageContains store ct)
  StorageContains store ((n := ep ::-> es) ': ct) =
    (StoreHasEntrypoint store n ep es, StorageContains store ct)
  StorageContains store ((n := ty) ': ct) =
    (StoreHasField store n ty, StorageContains store ct)