-- 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
  ) where

import Data.Constraint ((\\))
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.Label (Label)
import Morley.Util.Named
import Morley.Util.Type (KnownList, type (++))
import Morley.Util.TypeTuple

-- | 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
  :: 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
  :: 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
  \\ dupableEvi @(GetFieldType dt name)
  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
  :: 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
  :: 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.
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.
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
  :: forall dt fields st .
  ( InstrConstructC dt
  , fields ~ ConstructorFieldTypes dt
  , KnownList fields
  )
  => (fields ++ st) :-> dt : st
constructStack =
  I (instrConstructStack @dt @(ToTs fields))
      \\ totsKnownLemma @fields
      \\ totsAppendLemma @fields @st

-- | Decompose a complex object into its fields
deconstruct
  :: forall dt fields st .
  ( InstrDeconstructC dt
  , KnownList fields
  , fields ~ ConstructorFieldTypes dt
  )
  => dt : st :-> (fields ++ st)
deconstruct =
  I (instrDeconstruct @dt @(ToTs fields))
    \\ totsKnownLemma @fields
    \\ totsAppendLemma @fields @st

-- | 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_
  :: 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
  :: 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.
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.
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_
  :: 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