{-# LANGUAGE FunctionalDependencies #-}

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

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

    -- * Methods to work with storage
  , stToField
  , stGetField
  , stSetField
  , stMem
  , stGet
  , stUpdate
  , stDelete
  , stInsert
  , stInsertNew

    -- * Implementations
  , storeFieldOpsADT
  , storeFieldOpsDeeper
  , storeSubmapOpsDeeper
  , storeFieldOpsReferTo
  , storeSubmapOpsReferTo
  , composeStoreFieldOps
  , composeStoreSubmapOps
  ) where

import Data.Vinyl.Derived (Label)

import Lorentz.ADT
import Lorentz.Base
import Lorentz.Constraints
import qualified Lorentz.Instr as L
import qualified Lorentz.Macro as L
import Lorentz.Value
import Michelson.Typed.Haskell

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

-- | Datatype containing the full implementation of 'StoreHasField' typeclass.
--
-- We use this grouping because in most cases 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 benefits of @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 'StoreHasField' typeclass.
--
-- We use this grouping because in most cases 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 @DerivingVia@ extension.)
data StoreSubmapOps store mname key value = StoreSubmapOps
  { sopMem
      :: forall s.
         Label mname -> key : store : s :-> Bool : s
  , sopGet
      :: forall s.
         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 fields for 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
  => 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', IsComparable 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', IsComparable 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
    }


-- 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 '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

{- | 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 'StoreSubmapOps' implementation is made up
-- for submap 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
  }

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

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

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

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

Use it like follows:

@
type StorageConstraint = StorageContains
  [ "fieldInt" := Int
  , "fieldNat" := 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 := ty) ': ct) =
    (StoreHasField store n ty, StorageContains store ct)