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

{-# OPTIONS_GHC -Wno-redundant-constraints #-}
{-# LANGUAGE InstanceSigs #-}

-- | Instructions working on product types derived from Haskell ones.
module Morley.Michelson.Typed.Haskell.Instr.Product
  ( InstrGetFieldC
  , InstrSetFieldC
  , InstrConstructC
  , instrToField
  , instrGetField
  , instrGetFieldOpen
  , instrSetField
  , instrSetFieldOpen
  , instrConstruct
  , instrConstructStack
  , instrDeconstruct
  , InstrDeconstructC

  , GetFieldType
  , GLookupNamed
  , ConstructorFieldTypes
  , ConstructorFieldNames
  , FieldConstructor (..)
  , CastFieldConstructors (..)
  ) where

import Data.Vinyl.Core (Rec(..))
import GHC.Generics ((:*:)(..), (:+:)(..))
import GHC.Generics qualified as G
import GHC.TypeLits (ErrorMessage(..), Symbol, TypeError)

import Morley.Michelson.Text
import Morley.Michelson.Typed.Haskell.Instr.Helpers
import Morley.Michelson.Typed.Haskell.Value
import Morley.Michelson.Typed.Instr
import Morley.Michelson.Typed.Scope
import Morley.Util.Label (Label)
import Morley.Util.Named
import Morley.Util.Type

-- Fields lookup (helper)
----------------------------------------------------------------------------

-- | Result of field lookup - its type and path to it in
-- the tree.
data LookupNamedResult = LNR Type Path

type family LnrFieldType (lnr :: LookupNamedResult) where
  LnrFieldType ('LNR f _) = f
type family LnrBranch (lnr :: LookupNamedResult) :: Path where
  LnrBranch ('LNR _ p) = p

-- | Find field of some product type by its name.
--
-- Name might be either field record name, or one in 'NamedF' if
-- field is wrapped using 'Named.(:!)' or 'Named.(:?)'.
type GetNamed name a = LNRequireFound name a (GLookupNamed name (G.Rep a))

-- | Force result of 'GLookupNamed' to be 'Just'
type family LNRequireFound
  (name :: Symbol)
  (a :: Type)
  (mf :: Maybe LookupNamedResult)
    :: LookupNamedResult where
  LNRequireFound _ _ ('Just lnr) = lnr
  LNRequireFound name a 'Nothing = TypeError
    ('Text "Datatype " ':<>: 'ShowType a ':<>:
     'Text " has no field " ':<>: 'ShowType name)

type family GLookupNamed (name :: Symbol) (x :: Type -> Type)
          :: Maybe LookupNamedResult where
  GLookupNamed name (G.D1 _ x) = GLookupNamed name x
  GLookupNamed name (G.C1 _ x) = GLookupNamed name x

  GLookupNamed name (G.S1 ('G.MetaSel ('Just recName) _ _ _) (G.Rec0 a)) =
    If (name == recName)
      ('Just $ 'LNR a '[])
      'Nothing
  GLookupNamed name (G.S1 _ (G.Rec0 (NamedF f a fieldName))) =
    If (name == fieldName)
      ('Just $ 'LNR (NamedInner (NamedF f a fieldName)) '[])
      'Nothing
  GLookupNamed _ (G.S1 _ _) = 'Nothing
  GLookupNamed name (x :*: y) =
    LNMergeFound name (GLookupNamed name x) (GLookupNamed name y)
  GLookupNamed name (_ :+: _) = TypeError
    ('Text "Cannot seek for a field " ':<>: 'ShowType name ':<>:
     'Text " in sum type")
  GLookupNamed _ G.U1 = 'Nothing
  GLookupNamed _ G.V1 = TypeError
    ('Text "Cannot access fields of void-like type")

type family LNMergeFound
  (name :: Symbol)
  (f1 :: Maybe LookupNamedResult)
  (f2 :: Maybe LookupNamedResult)
    :: Maybe LookupNamedResult where
  LNMergeFound _ 'Nothing 'Nothing = 'Nothing
  LNMergeFound _ ('Just ('LNR a p)) 'Nothing = 'Just $ 'LNR a ('L ': p)
  LNMergeFound _ 'Nothing ('Just ('LNR a p)) = 'Just $ 'LNR a ('R ': p)
  LNMergeFound name ('Just _) ('Just _) = TypeError
    ('Text "Ambigous reference to datatype field: " ':<>: 'ShowType name)

-- | Get type of field by datatype it is contained in and field name.
type GetFieldType dt name = LnrFieldType (GetNamed name dt)

----------------------------------------------------------------------------
-- Value accessing instruction
----------------------------------------------------------------------------

-- | Make an instruction which accesses given field of the given datatype.
instrToField
  :: forall dt name st.
     InstrGetFieldC dt name
  => Label name -> Instr (ToT dt ': st) (ToT (GetFieldType dt name) ': st)
instrToField _ =
  gInstrToField @name @(G.Rep dt) @(LnrBranch (GetNamed name dt))
    @(GetFieldType dt name)

-- | Make an instruction which copies given field of the given datatype.
--
-- This behaves exactly as @Seq DUP (instrToField #name)@, but does not require
-- the entire datatype to be dupable (the field ofc still must be dupable).
--
-- If we follow the path from the root to the copied field in the pairs tree,
-- the more nodes contain non-dupable elements in the subtree, the less
-- efficient this function becomes.
-- Assuming that all fields are accessed equally often, the most optimal
-- representation would be to put all dupable elements to one subtree
-- of the root, and all non-dupable elements in the second subtree of the root.
instrGetField
  :: forall dt name st.
     (InstrGetFieldC dt name, DupableScope (ToT (GetFieldType dt name)))
  => Label name -> Instr (ToT dt ': st) (ToT (GetFieldType dt name) ': ToT dt ': st)
instrGetField = instrGetFieldOpen @dt DUP Nop

-- | \"Open\" version of 'instrGetField', meaning that it accepts
-- continuations that accept the copied field. This allows chaining
-- getters without requiring 'DupableScope' on the intermediate fields.
--
-- Accepted continuations:
--
-- 1. The one that leaves the field on stack (and does a duplication inside).
--    Used when the datatype had an unfortunate placement of non-dupable fields,
--    or the intermediate field contains non-dupable elements, so the duplication
--    cannot occur automatically in between.
-- 2. The one that consumes the field, used in case we managed to call @DUP@
--    earlier to make the instruction's implementation smaller, and now we
--    only need to access the field with @CAR@ and @CDR@s.
--
-- Note that only one continuation will be chosen eventually, no contract size
-- overhead is expected in this regard.
instrGetFieldOpen
  :: forall dt name res st.
     (InstrGetFieldC dt name)
  => Instr '[ToT (GetFieldType dt name)] '[res, ToT (GetFieldType dt name)]
  -> Instr '[ToT (GetFieldType dt name)] '[res]
  -> Label name
  -> Instr (ToT dt ': st) (res ': ToT dt ': st)
instrGetFieldOpen contWDup contWoDup _ =
  gInstrGetFieldOpen @name @(G.Rep dt) @(LnrBranch (GetNamed name dt))
    @(GetFieldType dt name) contWDup contWoDup

-- | Constraint for 'instrGetField'.
type InstrGetFieldC dt name =
  ( GenericIsoValue dt
  , GInstrGet name (G.Rep dt)
      (LnrBranch (GetNamed name dt))
      (LnrFieldType (GetNamed name dt))
  )

{- Note about bulkiness of `instrToField` type signature:

Read this only if you are going to modify the signature qualitatively.

It may seem that you can replace 'LnrBranch' and 'LnrFieldType' getters since
their result is already assigned to a type variable, but if you do so,
on attempt to access non-existing field GHC will prompt you a couple of huge
"cannot deduce type" errors along with desired "field is not present".
To make error messages clearer we prevent GHC from deducing 'GInstrAccess' when
no field is found via using those getters.
-}

-- | Generic traversal for 'instrToField'.
class GIsoValue x =>
  GInstrGet
    (name :: Symbol)
    (x :: Type -> Type)
    (path :: Path)
    (fieldTy :: Type) where
  gInstrToField
    :: GIsoValue x
    => Instr (GValueType x ': s) (ToT fieldTy ': s)
  gInstrGetFieldOpen
    :: (GIsoValue x)
    => Instr '[ToT fieldTy] '[res, ToT fieldTy]
    -> Instr '[ToT fieldTy] '[res]
    -> Instr (GValueType x ': s) (res ': GValueType x ': s)

-- | If the current node in the tree is dupable, this provides the more
-- efficient implementation.
gTryShortCircuitNonDup
  :: forall name x path fieldTy res s.
     (GInstrGet name x path fieldTy)
  => Instr '[ToT fieldTy] '[res]
  -> Maybe $ Instr (GValueType x ': s) (res ': GValueType x ': s)
gTryShortCircuitNonDup cont =
  case checkScope @(DupableScope $ GValueType x) of
    Left{} -> Nothing
    Right Dict -> Just $
      DUP `Seq` gInstrToField @name @x @path @fieldTy `Seq` FrameInstr Proxy cont

instance GInstrGet name x path f => GInstrGet name (G.M1 t i x) path f where
  gInstrToField = gInstrToField @name @x @path @f
  gInstrGetFieldOpen = gInstrGetFieldOpen @name @x @path @f

instance (IsoValue f, ToT f ~ ToT f') =>
         GInstrGet name (G.Rec0 f) '[] f' where
  gInstrToField = Nop
  gInstrGetFieldOpen contWDup _ = FrameInstr Proxy contWDup

instance (GInstrGet name x path f, GIsoValue y) => GInstrGet name (x :*: y) ('L ': path) f where
  gInstrToField = CAR `Seq` gInstrToField @name @x @path @f
  gInstrGetFieldOpen contWDup contWoDup =
    gTryShortCircuitNonDup @name @(x :*: y) @('L ': path) @f contWoDup ?:
    UNPAIR
      `Seq` gInstrGetFieldOpen @name @x @path @f contWDup contWoDup
      `Seq` DIP PAIR

instance (GInstrGet name y path f, GIsoValue x) => GInstrGet name (x :*: y) ('R ': path) f where
  gInstrToField = CDR `Seq` gInstrToField @name @y @path @f
  gInstrGetFieldOpen contWDup contWoDup =
    gTryShortCircuitNonDup @name @(x :*: y) @('R ': path) @f contWoDup ?:
    UNPAIR `Seq` SWAP
      `Seq` gInstrGetFieldOpen @name @y @path @f contWDup contWoDup
      `Seq` DIP (SWAP `Seq` PAIR)

-- Examples
----------------------------------------------------------------------------

type MyType1 = ("int" :! Integer, "bytes" :! ByteString, "bytes2" :? ByteString)

_getIntInstr1 :: Instr (ToT MyType1 ': s) (ToT Integer ': s)
_getIntInstr1 = instrToField @MyType1 #int

_getTextInstr1 :: Instr (ToT MyType1 ': s) (ToT (Maybe ByteString) ': s)
_getTextInstr1 = instrToField @MyType1 #bytes2

data MyType2 = MyType2
  { getInt :: Integer
  , getText :: MText
  , getUnit :: ()
  , getMyType1 :: MyType1
  } deriving stock (Generic)
    deriving anyclass (IsoValue)

_getIntInstr2 :: Instr (ToT MyType2 ': s) (ToT () ': s)
_getIntInstr2 = instrToField @MyType2 #getUnit

_getIntInstr2' :: Instr (ToT MyType2 ': s) (ToT Integer ': s)
_getIntInstr2' =
  instrToField @MyType2 #getMyType1 `Seq` instrToField @MyType1 #int

----------------------------------------------------------------------------
-- Value modification instruction
----------------------------------------------------------------------------

-- | For given complex type @dt@ and its field @fieldTy@ update the field value.
instrSetField
  :: forall dt name st.
     InstrSetFieldC dt name
  => Label name -> Instr (ToT (GetFieldType dt name) ': ToT dt ': st) (ToT dt ': st)
instrSetField =
  instrSetFieldOpen @dt (DIP DROP)

-- | \"Open\" version of 'instrSetField', meaning that it accepts a continuation
-- that describes how to update the field. This allows chaining setters with
-- zero overhead and without requiring 'DupableScope' on the intermediate fields
-- (if we supported only closed setters that work directly with one datatype, we
-- would need to duplicate intermediate fields to chain setters).
instrSetFieldOpen
  :: forall dt name new st.
     InstrSetFieldC dt name
  => Instr '[new, ToT (GetFieldType dt name)] '[ToT (GetFieldType dt name)]
  -> Label name
  -> Instr (new ': ToT dt ': st) (ToT dt ': st)
instrSetFieldOpen cont _ =
  gInstrSetFieldOpen @name @(G.Rep dt) @(LnrBranch (GetNamed name dt))
    @(GetFieldType dt name) cont

-- | Constraint for 'instrSetField'.
type InstrSetFieldC dt name =
  ( GenericIsoValue dt
  , GInstrSetField name (G.Rep dt)
      (LnrBranch (GetNamed name dt))
      (LnrFieldType (GetNamed name dt))
  )

-- | Generic traversal for 'instrSetField'.
class GIsoValue x =>
  GInstrSetField
    (name :: Symbol)
    (x :: Type -> Type)
    (path :: Path)
    (fieldTy :: Type) where
  gInstrSetFieldOpen
    :: Instr '[new, ToT fieldTy] '[ToT fieldTy]
    -> Instr (new ': GValueType x ': s) (GValueType x ': s)

instance GInstrSetField name x path f => GInstrSetField name (G.M1 t i x) path f where
  gInstrSetFieldOpen = gInstrSetFieldOpen @name @x @path @f

instance (IsoValue f, ToT f ~ ToT f') =>
         GInstrSetField name (G.Rec0 f) '[] f' where
  gInstrSetFieldOpen cont = FrameInstr Proxy cont

instance (GInstrSetField name x path f, GIsoValue y) => GInstrSetField name (x :*: y) ('L ': path) f where
  gInstrSetFieldOpen cont =
    DIP UNPAIR `Seq`
    gInstrSetFieldOpen @name @x @path @f cont `Seq`
    PAIR

instance (GInstrSetField name y path f, GIsoValue x) => GInstrSetField name (x :*: y) ('R ': path) f where
  gInstrSetFieldOpen cont =
    DIP (UNPAIR `Seq` SWAP) `Seq`
    gInstrSetFieldOpen @name @y @path @f cont `Seq`
    SWAP `Seq` PAIR

-- Examples
----------------------------------------------------------------------------

_setIntInstr1 :: Instr (ToT Integer ': ToT MyType2 ': s) (ToT MyType2 ': s)
_setIntInstr1 = instrSetField @MyType2 #getInt

----------------------------------------------------------------------------
-- Value construction instruction
----------------------------------------------------------------------------

-- | Way to construct one of the fields in a complex datatype.
newtype FieldConstructor (st :: [k]) (field :: Type) =
  FieldConstructor (Instr (ToTs' st) (ToT field ': ToTs' st))

-- | Ability to pass list of fields with the same ToTs.
-- It may be useful if you don't want to work with NamedF in ConstructorFieldTypes.
class ToTs xs ~ ToTs ys => CastFieldConstructors xs ys where
  castFieldConstructorsImpl :: Rec (FieldConstructor st) xs -> Rec (FieldConstructor st) ys

instance CastFieldConstructors '[] '[] where
  castFieldConstructorsImpl RNil = RNil

instance (CastFieldConstructors xs ys, ToTs xs ~ ToTs ys, ToT x ~ ToT y)
    => CastFieldConstructors (x ': xs) (y ': ys) where
  castFieldConstructorsImpl (FieldConstructor fctr :& xs) =
    FieldConstructor fctr :& (castFieldConstructorsImpl @xs @ys xs)

-- | For given complex type @dt@ and its field @fieldTy@ update the field value.
instrConstruct
  :: forall dt st. InstrConstructC dt
  => Rec (FieldConstructor st) (ConstructorFieldTypes dt)
  -> Instr st (ToT dt ': st)
instrConstruct = gInstrConstruct @(G.Rep dt)

instrConstructStack
  :: forall dt stack st .
  ( InstrConstructC dt
  , stack ~ ToTs (ConstructorFieldTypes dt)
  , KnownList stack
  )
  => Instr (stack ++ st) (ToT dt ': st)
instrConstructStack =
  FrameInstr (Proxy @st) (gInstrConstructStack @(G.Rep dt))

-- | Types of all fields in a datatype.
type ConstructorFieldTypes dt = GFieldTypes (G.Rep dt)

-- | Names of all fields in a datatype.
type ConstructorFieldNames dt = GFieldNames (G.Rep dt)

-- | Constraint for 'instrConstruct' and 'gInstrConstructStack'.
type InstrConstructC dt = (GenericIsoValue dt, GInstrConstruct (G.Rep dt))

-- | Retrieve field names of a constructor.
type family GFieldNames x :: [Symbol] where
  GFieldNames (G.D1 _ x) = GFieldNames x
  GFieldNames (G.C1 _ x) = GFieldNames x
  GFieldNames (x :*: y) = GFieldNames x ++ GFieldNames y
  GFieldNames (G.S1 ('G.MetaSel ('Just fieldName) _ _ _) _) = '[fieldName]
  GFieldNames (G.S1 _ (G.Rec0 (NamedF _ _ fieldName))) = '[fieldName]

  GFieldNames (G.S1 ('G.MetaSel _ _ _ _) _) = TypeError ('Text "All fields have to be named")
  GFieldNames (_ :+: _) = TypeError ('Text "Cannot get field names of sum type")
  GFieldNames G.V1 = TypeError ('Text "Must be at least one constructor")

-- | Generic traversal for 'instrConstruct'.
class GIsoValue x => GInstrConstruct (x :: Type -> Type) where
  type GFieldTypes x :: [Type]
  gInstrConstruct :: Rec (FieldConstructor st) (GFieldTypes x) -> Instr st (GValueType x ': st)
  gInstrConstructStack :: Instr (ToTs (GFieldTypes x)) '[GValueType x]

instance GInstrConstruct x => GInstrConstruct (G.M1 t i x) where
  type GFieldTypes (G.M1 t i x) = GFieldTypes x
  gInstrConstruct = gInstrConstruct @x
  gInstrConstructStack = gInstrConstructStack @x

instance ( GInstrConstruct x, GInstrConstruct y
         , RSplit (GFieldTypes x) (GFieldTypes y)
         , KnownList (ToTs (GFieldTypes x)), KnownList (ToTs (GFieldTypes y))
         , ToTs (GFieldTypes x) ++ ToTs (GFieldTypes y) ~ ToTs (GFieldTypes x ++ GFieldTypes y)
         ) => GInstrConstruct (x :*: y) where
  type GFieldTypes (x :*: y) = GFieldTypes x ++ GFieldTypes y
  gInstrConstruct fs =
    let (lfs, rfs) = rsplit fs
        linstr = gInstrConstruct @x lfs
        rinstr = gInstrConstruct @y rfs
    in linstr `Seq` DIP rinstr `Seq` PAIR

  gInstrConstructStack =
    let linstr = gInstrConstructStack @x
        rinstr = gInstrConstructStack @y
    in FrameInstr (Proxy @(ToTs (GFieldTypes y))) linstr `Seq` DIP rinstr `Seq` PAIR

instance GInstrConstruct G.U1 where
  type GFieldTypes G.U1 = '[]
  gInstrConstruct RNil = UNIT
  gInstrConstructStack = UNIT

instance ( TypeError ('Text "Cannot construct sum type")
         , GIsoValue x, GIsoValue y
         ) => GInstrConstruct (x :+: y) where
  type GFieldTypes (x :+: y) = '[]
  gInstrConstruct = error "impossible"
  gInstrConstructStack = error "impossible"

instance ( TypeError ('Text "Cannot construct void-like type")
         ) => GInstrConstruct G.V1 where
  type GFieldTypes G.V1 = '[]
  gInstrConstruct = error "impossible"
  gInstrConstructStack = error "impossible"

instance IsoValue a => GInstrConstruct (G.Rec0 a) where
  type GFieldTypes (G.Rec0 a) = '[a]
  gInstrConstruct (FieldConstructor i :& RNil) = i
  gInstrConstructStack = Nop

-- Examples
----------------------------------------------------------------------------

_constructInstr1 :: Instr (ToT MyType1 ': s) (ToT MyType2 ': ToT MyType1 ': s)
_constructInstr1 =
  instrConstruct @MyType2 $
    FieldConstructor (PUSH (toVal @Integer 5)) :&
    FieldConstructor (PUSH (toVal @MText [mt||])) :&
    FieldConstructor UNIT :&
    FieldConstructor DUP :&
    RNil

_constructInstr2 :: Instr s (ToT MyType1 ': s)
_constructInstr2 =
  instrConstruct @MyType1 $
    FieldConstructor (PUSH (toVal @Integer 5)) :&
    FieldConstructor (PUSH (toVal @ByteString "bytestring1")) :&
    FieldConstructor (PUSH (toVal @(Maybe ByteString) (Just "bytestring2"))) :&
    RNil

----------------------------------------------------------------------------
-- Value deconstruction
----------------------------------------------------------------------------

-- | Constraint for 'instrConstruct'.
type InstrDeconstructC dt = (GenericIsoValue dt, GInstrDeconstruct (G.Rep dt))

-- | For given complex type @dt@ deconstruct it to its field types.
instrDeconstruct
  :: forall dt stack st .
  (InstrDeconstructC dt
  , stack ~ ToTs (ConstructorFieldTypes dt)
  , KnownList stack
  )
  => Instr (ToT dt ': st) (stack ++ st)
instrDeconstruct = FrameInstr (Proxy @st) $ gInstrDeconstruct @(G.Rep dt)

-- | Generic traversal for 'instrDeconstruct'.
class GIsoValue x => GInstrDeconstruct (x :: Type -> Type) where
  gInstrDeconstruct :: Instr '[GValueType x] (ToTs (GFieldTypes x))

instance GInstrDeconstruct x => GInstrDeconstruct (G.M1 t i x) where
  gInstrDeconstruct = gInstrDeconstruct @x

instance ( GInstrDeconstruct x, GInstrDeconstruct y
         , t ~ (x :*: y)
         , KnownList (ToTs (GFieldTypes x)), KnownList (ToTs (GFieldTypes y))
         , ToTs (GFieldTypes x) ++ ToTs (GFieldTypes y) ~ ToTs (GFieldTypes x ++ GFieldTypes y)
         ) => GInstrDeconstruct (x :*: y) where
  gInstrDeconstruct =
    let linstr = gInstrDeconstruct @x
        rinstr = gInstrDeconstruct @y
    in UNPAIR `Seq` DIP rinstr `Seq` FrameInstr (Proxy @(ToTs (GFieldTypes y))) linstr

instance GInstrDeconstruct G.U1 where
  gInstrDeconstruct = DROP

instance ( TypeError ('Text "Cannot deconstruct sum type")
         , GIsoValue x, GIsoValue y
         ) => GInstrDeconstruct (x :+: y) where
  gInstrDeconstruct = error "impossible"

instance ( TypeError ('Text "Cannot deconstruct void-like type")
         ) => GInstrDeconstruct G.V1 where
  gInstrDeconstruct = error "impossible"

instance IsoValue a => GInstrDeconstruct (G.Rec0 a) where
  gInstrDeconstruct = Nop