-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

{-# LANGUAGE FunctionalDependencies #-}
{-# OPTIONS_GHC -Wno-unticked-promoted-constructors #-}

module Lorentz.ADT
  ( HasField
  , HasFieldOfType
  , HasFieldsOfType
  , NamedField (..)
  , (:=)
  , HasDupableGetters
  , toField
  , toFieldNamed
  , getField
  , getFieldNamed
  , setField
  , modifyField
  , construct
  , constructT
  , constructStack
  , deconstruct
  , fieldCtor
  , wrap_
  , wrapOne
  , case_
  , caseT
  , unsafeUnwrap_
  , CaseTC
  , CaseArrow (..)
  , CaseClauseL (..)
  , InstrConstructC
  , ConstructorFieldTypes

    -- * Useful re-exports
  , Rec (..)
  , (:!)
  , (:?)

    -- * Advanced methods
  , getFieldOpen
  , setFieldOpen

    -- * Definitions used in examples
    -- $setup
  ) where

import Data.Vinyl.Core (RMap(..), Rec(..))
import GHC.TypeLits (AppendSymbol, Symbol)

import Lorentz.Base
import Lorentz.Coercions
import Lorentz.Constraints
import Morley.Michelson.Typed.Haskell.Instr
import Morley.Michelson.Typed.Haskell.Value
import Morley.Util.Generic
import Morley.Util.Label (Label)
import Morley.Util.Named
import Morley.Util.Type (type (++))
import Morley.Util.TypeTuple

{- $setup
>>> import Data.Vinyl (Rec(..))
>>> import Fmt (pretty)
>>> import Lorentz.Base ((#))
>>> import Lorentz.Instr as L
>>> import Lorentz.Zip (ZippedStackRepr(..), ZSNil(..))
>>> import Lorentz.Run.Simple ((-$), (-$?))
>>> import Morley.Michelson.Runtime.Dummy (dummySelf)
>>> import Morley.Michelson.Typed (IsoValue(..))
>>> import Morley.Michelson.Typed.Haskell.Value (Ticket(..))
>>> import Morley.Tezos.Address (Constrained(..))
>>> :{
data TestProduct = TestProduct
  { fieldA :: Bool
  , fieldB :: Integer
  , fieldC :: ()
  } deriving stock (Generic, Eq, Show)
    deriving anyclass (IsoValue)
--
data TestProductWithNonDup = TestProductWithNonDup
  { fieldTP :: TestProduct
  , fieldD  :: Ticket () -- non-dupable value
  } deriving stock (Generic, Eq, Show)
    deriving anyclass (IsoValue)
--
data TestSum
  = TestSumA Integer
  | TestSumB (Bool, ())
  deriving stock (Generic, Eq, Show)
  deriving anyclass (IsoValue)
:}

>>> let testTicket = Ticket (Constrained dummySelf) () 10
>>> let testProduct = TestProduct True 42 ()
>>> let testProductWithNonDup = TestProductWithNonDup testProduct testTicket
-}

-- | Allows field access and modification.
type HasField dt fname =
  ( InstrGetFieldC dt fname
  , InstrSetFieldC dt fname
  )

-- | Like 'HasField', but allows constrainting field type.
type HasFieldOfType dt fname fieldTy =
  ( HasField dt fname
  , GetFieldType dt fname ~ fieldTy
  )

-- | A pair of field name and type.
data NamedField = NamedField Symbol Type
type n := ty = 'NamedField n ty
infixr 0 :=

-- | Shortcut for multiple 'HasFieldOfType' constraints.
type family HasFieldsOfType (dt :: Type) (fs :: [NamedField])
             :: Constraint where
  HasFieldsOfType _ '[] = ()
  HasFieldsOfType dt ((n := ty) ': fs) =
    (HasFieldOfType dt n ty, HasFieldsOfType dt fs)

-- | This marker typeclass is a requirement for the 'getField'
-- (where it is imposed on the /datatype/), and it is supposed to be satisfied
-- in two cases:
--
-- 1. The entire datatype is 'Dupable';
-- 2. When the datatype has non-dupable fields, they are located so that
-- 'getField' remains efficient.
--
-- The problem we are trying to solve here: without special care, 'getField'
-- may become multiple times more costly, see 'instrGetField' for the
-- explanation.
-- And this typeclass imposes an invariant: if we ever use 'getField' on a
-- datatype, then we have to pay attention to the datatype's Michelson
-- representation and ensure 'getField' remains optimal.
--
-- When you are developing a contract:
-- "Lorentz.Layouts.NonDupable" module contains utilities to help you
-- provide the necessary Michelson layout. In case you want to use your
-- custom layout but still allow 'getField' for it, you can define an instance
-- for your type manually as an assurance that Michelson layout is optimal enough
-- to use 'getField' on this type.
--
-- When you are developing a library:
-- Note that 'HasDupableGetters' resolves to 'Dupable' by default, and when
-- propagating this constraint you can switch to 'Dupable' anytime but this
-- will also make your code unusable in the presence of 'Ticket's and other
-- non-dupable types.
class HasDupableGetters a
instance {-# OVERLAPPABLE #-} Dupable a => HasDupableGetters a

-- | Extract a field of a datatype replacing the value of this
-- datatype with the extracted field.
--
-- For this and the following functions you have to specify field name
-- which is either record name or name attached with @(:!)@ operator.
--
-- >>> :{
--  (toField @TestProductWithNonDup #fieldTP -$ testProductWithNonDup) == testProduct
-- :}
-- True
--
toField
  :: forall dt name st.
     InstrGetFieldC dt name
  => Label name -> dt : st :-> GetFieldType dt name : st
toField = I . instrToField @dt

-- | Like 'toField', but leaves field named.
toFieldNamed
  :: forall dt name st.
     InstrGetFieldC dt name
  => Label name -> dt : st :-> (name :! GetFieldType dt name) : st
toFieldNamed l = toField l # toNamed l

-- | Extract a field of a datatype, leaving the original datatype on stack.
--
-- >>> :{
--  (getField @TestProduct #fieldB # L.swap # toField @TestProduct #fieldB # mul) -$
--     testProduct { fieldB = 3 }
-- :}
-- 9
getField
  :: forall dt name st.
     (InstrGetFieldC dt name, Dupable (GetFieldType dt name), HasDupableGetters dt)
  => Label name -> dt : st :-> GetFieldType dt name : dt ': st
getField = I . instrGetField @dt
  where
    _needHasDupableGetters = Dict @(HasDupableGetters dt)

-- | Like 'getField', but leaves field named.
getFieldNamed
  :: forall dt name st.
     (InstrGetFieldC dt name, Dupable (GetFieldType dt name), HasDupableGetters dt)
  => Label name -> dt : st :-> (name :! GetFieldType dt name) : dt ': st
getFieldNamed l = getField l # coerceWrap

-- | Set a field of a datatype.
--
-- >>> setField @TestProduct #fieldB -$ 23 ::: testProduct
-- TestProduct {fieldA = True, fieldB = 23, fieldC = ()}
setField
  :: forall dt name st.
     InstrSetFieldC dt name
  => Label name -> (GetFieldType dt name ': dt ': st) :-> (dt ': st)
setField = I . instrSetField @dt

-- | Apply given modifier to a datatype field.
--
-- >>> modifyField @TestProduct #fieldB (dup # mul) -$ testProduct { fieldB = 8 }
-- TestProduct {fieldA = True, fieldB = 64, fieldC = ()}
modifyField
  :: forall dt name st.
     ( InstrGetFieldC dt name
     , InstrSetFieldC dt name
     , Dupable (GetFieldType dt name)
     , HasDupableGetters dt
     )
  => Label name
  -> (forall st0. (GetFieldType dt name ': st0) :-> (GetFieldType dt name ': st0))
  -> dt : st :-> dt : st
modifyField l i = getField @dt l # i # setField @dt l

-- | \"Open\" version of 'getField', an advanced method suitable for chaining
-- getters.
--
-- It accepts two continuations accepting the extracted field, one that
-- leaves the field on stack (and does a duplication of @res@ inside)
-- and another one that consumes the field. Normally these are just @getField@
-- and @toField@ for some nested field.
--
-- Unlike the straightforward chaining of 'getField'/'toField' methods,
-- @getFieldOpen@ does not require the immediate field to be dupable; rather,
-- in the best case only @res@ has to be dupable.
getFieldOpen
  :: forall dt name res st.
     (InstrGetFieldC dt name, HasDupableGetters dt)
  => '[GetFieldType dt name] :-> '[res, GetFieldType dt name]
  -> '[GetFieldType dt name] :-> '[res]
  -> Label name
  -> dt : st :-> res : dt ': st
getFieldOpen contWDup contWoDup =
  I . instrGetFieldOpen @dt (iNonFailingCode contWDup) (iNonFailingCode contWoDup)
  where
    _needHasDupableGetters = Dict @(HasDupableGetters dt)

-- | \"Open\" version of 'setField', an advanced method suitable for chaining
-- setters.
--
-- It accepts a continuation accepting the field extracted for the update and
-- the new value that is being set. Normally this continuation is just @setField@
-- for some nested field.
setFieldOpen
  :: forall dt name new st.
     InstrSetFieldC dt name
  => '[new, GetFieldType dt name] :-> '[GetFieldType dt name]
  -> Label name
  -> (new ': dt ': st) :-> (dt ': st)
setFieldOpen cont =
  I . instrSetFieldOpen @dt (iNonFailingCode cont)

-- | Make up a datatype. You provide a pack of individual fields constructors.
--
-- Each element of the accepted record should be an instruction wrapped with
-- 'fieldCtor' function. This instruction will have access to the stack at
-- the moment of calling @construct@.
-- Instructions have to output fields of the built datatype, one per instruction;
-- instructions order is expected to correspond to the order of fields in the
-- datatype.
--
-- >>> :{
--  let ctor =
--        (fieldCtor (push True)) :&
--        (fieldCtor (push (42 :: Integer))) :&
--        (fieldCtor (push ())) :&
--        RNil
-- :}
--
-- >>> construct @TestProduct ctor -$ ZSNil
-- TestProduct {fieldA = True, fieldB = 42, fieldC = ()}
construct
  :: forall dt st.
     ( InstrConstructC dt
     , RMap (ConstructorFieldTypes dt)
     )
  => Rec (FieldConstructor st) (ConstructorFieldTypes dt)
  -> st :-> dt : st
construct fctors =
  I $ instrConstruct @dt $
  rmap (\(FieldConstructor i) -> FieldConstructor i) fctors

-- | Version of 'construct' which accepts tuple of field constructors.
--
-- >>> let ctor = (fieldCtor (push True), fieldCtor (push @Integer 42), fieldCtor (push ()))
-- >>> constructT @TestProduct ctor -$ ZSNil
-- TestProduct {fieldA = True, fieldB = 42, fieldC = ()}
constructT
  :: forall dt fctors st.
     ( InstrConstructC dt
     , RMap (ConstructorFieldTypes dt)
     , fctors ~ Rec (FieldConstructor st) (ConstructorFieldTypes dt)
     , RecFromTuple fctors
     )
  => IsoRecTuple fctors
  -> st :-> dt : st
constructT = construct . recFromTuple

-- | Construct an object from fields on the stack.
--
-- >>> constructStack @TestProduct -$ True ::: 42 ::: ()
-- TestProduct {fieldA = True, fieldB = 42, fieldC = ()}
constructStack
  :: forall dt fields st .
  ( InstrConstructC dt
  , fields ~ ConstructorFieldTypes dt
  , ToTs fields ++ ToTs st ~ ToTs (fields ++ st)
  )
  => (fields ++ st) :-> dt : st
constructStack =
  I (instrConstructStack @dt @(ToTs fields))

{- | Decompose a complex object into its fields

>>> deconstruct @TestProduct # constructStack @TestProduct -$ testProduct
TestProduct {fieldA = True, fieldB = 42, fieldC = ()}

Will show human-readable error on a missing `Generic` instance:

>>> data Foo = Foo () ()
>>> deconstruct @Foo
...
... GHC.Generics.Rep Foo
... is stuck. Likely
... Generic Foo
... instance is missing or out of scope.
...

>>> data Foo = Foo () () deriving (Generic, IsoValue)
>>> pretty $ deconstruct @Foo
[UNPAIR, DIP {  }]
-}
deconstruct
  :: forall dt fields st .
  ( InstrDeconstructC dt (ToTs st)
  , fields ~ GFieldTypes (GRep dt) '[]
  , ToTs fields ++ ToTs st ~ ToTs (fields ++ st)
  )
  => dt : st :-> (fields ++ st)
deconstruct =
  I (instrDeconstruct @dt @(ToTs fields))

-- | Lift an instruction to field constructor.
fieldCtor :: HasCallStack => (st :-> f : st) -> FieldConstructor st f
fieldCtor = \case
  I i -> FieldConstructor i
  FI _ -> error "Field constructor always fails"

-- | Wrap entry in constructor. Useful for sum types.
--
-- >>> wrap_ @TestSum #cTestSumB -$ (False, ())
-- TestSumB (False,())
wrap_
  :: forall dt name st.
     InstrWrapC dt name
  => Label name -> (AppendCtorField (GetCtorField dt name) st) :-> dt : st
wrap_ =
  case appendCtorFieldAxiom @(GetCtorField dt name) @st of
    Dict -> I . instrWrap @dt

-- | Wrap entry in single-field constructor. Useful for sum types.
--
-- >>> wrapOne @TestSum #cTestSumA -$ 42
-- TestSumA 42
wrapOne
  :: forall dt name st.
     InstrWrapOneC dt name
  => Label name -> (CtorOnlyField name dt ': st) :-> dt : st
wrapOne = case appendCtorFieldAxiom @(GetCtorField dt name) @st of
  Dict -> I . instrWrapOne @dt

-- | Lorentz analogy of 'CaseClause', it works on plain 'Type' types.
data CaseClauseL (inp :: [Type]) (out :: [Type]) (param :: CaseClauseParam) where
  CaseClauseL :: AppendCtorField x inp :-> out -> CaseClauseL inp out ('CaseClauseParam ctor x)

-- | Provides "case" arrow which works on different wrappers for clauses.
class CaseArrow name body clause | clause -> name, clause -> body where
  -- | Lift an instruction to case clause.
  --
  -- You should write out constructor name corresponding to the clause
  -- explicitly. Prefix constructor name with "c" letter, otherwise
  -- your label will not be recognized by Haskell parser.
  -- Passing constructor name can be circumvented but doing so is not recomended
  -- as mentioning contructor name improves readability and allows avoiding
  -- some mistakes.
  (/->) :: Label name -> body -> clause
  infixr 0 /->

instance ( name ~ ("c" `AppendSymbol` ctor)
         , body ~ (AppendCtorField x inp :-> out)
         ) => CaseArrow name body
                (CaseClauseL inp out ('CaseClauseParam ctor x)) where
  (/->) _ = CaseClauseL

-- | Pattern match on the given sum type.
--
-- You have to provide a 'Rec' containing case branches.
-- To construct a case branch use '/->' operator.
--
-- >>> :{
--  let caseTestSum = case_ @TestSum $
--        (#cTestSumA /-> nop) :&
--        (#cTestSumB /-> L.drop # push 23) :&
--        RNil
-- :}
--
-- >>> caseTestSum -$ TestSumA 42
-- 42
-- >>> caseTestSum -$ TestSumB (False, ())
-- 23
case_
  :: forall dt out inp.
     ( InstrCaseC dt
     , RMap (CaseClauses dt)
     )
  => Rec (CaseClauseL inp out) (CaseClauses dt) -> dt : inp :-> out
case_ = LorentzInstr . instrCase @dt . rmap coerceCaseClause
  where
    coerceCaseClause
      :: forall clauses.
         CaseClauseL inp out clauses -> CaseClause (ToTs inp) (ToTs out) clauses
    coerceCaseClause (CaseClauseL (LorentzInstr cc)) =
      CaseClause $ case Proxy @clauses of
        (_ :: Proxy ('CaseClauseParam ctor cc)) ->
          case appendCtorFieldAxiom @cc @inp of Dict -> cc

-- | Like 'case_', accepts a tuple of clauses, which may be more convenient.
--
-- If user is experiencing problems with wierd errors about tuples while using
-- this function, he should take look at "Morley.Util.TypeTuple.Instances" and ensure
-- that his tuple isn't bigger than generated instances, if so, he should probably
-- extend number of generated instances.
--
-- >>> :{
--  let caseTTestSum = caseT @TestSum $
--        ( #cTestSumA /-> nop
--        , #cTestSumB /-> L.drop # push 23
--        )
-- :}
--
-- >>> caseTTestSum -$ TestSumA 42
-- 42
-- >>> caseTTestSum -$ TestSumB (False, ())
-- 23
caseT
  :: forall dt out inp clauses.
     CaseTC dt out inp clauses
  => IsoRecTuple clauses -> dt : inp :-> out
caseT = case_ @dt . recFromTuple

type CaseTC dt out inp clauses =
  ( InstrCaseC dt
  , RMap (CaseClauses dt)
  , RecFromTuple clauses
  , clauses ~ Rec (CaseClauseL inp out) (CaseClauses dt)
  )

-- | Unwrap a constructor with the given name. Useful for sum types.
--
-- >>> unsafeUnwrap_ @TestSum #cTestSumA -$? TestSumA 42
-- Right 42
-- >>> pretty $ unsafeUnwrap_ @TestSum #cTestSumA -$? TestSumB (False, ())
-- <Left: Reached FAILWITH instruction with '"BadCtor"' at line 1 char 1.>
unsafeUnwrap_
  :: forall dt name st.
     InstrUnwrapC dt name
  => Label name -> dt : st :-> (CtorOnlyField name dt ': st)
unsafeUnwrap_ =
  case appendCtorFieldAxiom @(GetCtorField dt name) @st of
    Dict -> I . unsafeInstrUnwrap @dt