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

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

module Lorentz.ADT
  ( HasField
  , HasFieldOfType
  , HasFieldsOfType
  , NamedField (..)
  , (:=)
  , 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 (..)
  , (:!)
  , (:?)
  ) 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.Michelson.Typed.Instr
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)

-- | 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.
--
-- TODO: [#585] Make this and all depending functions require only
-- @Dupable (GetFieldType dt name)@
getField
  :: forall dt name st.
     (InstrGetFieldC dt name, Dupable dt)
  => Label name -> dt : st :-> GetFieldType dt name : dt ': st
getField l = dup # toField l
  where
    dup :: forall a s. Dupable a => a : s :-> a : a : s
    dup = I DUP \\ dupableEvi @a

-- | Like 'getField', but leaves field named.
getFieldNamed
  :: forall dt name st.
     (InstrGetFieldC dt name, Dupable 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 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

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