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

-- | Decompose a complex value into its fields
-- to be used in 'setVar'.
-- Also functionality to generate code to deconstruct storage
-- into primitive fields the storage consists of
-- and to construct it back.
module Indigo.Internal.Expr.Decompose
  ( decomposeExpr
  , deepDecomposeCompose
  , ExprDecomposition (..)
  , IsObject
  ) where

import Prelude

import Data.Constraint (Dict(..))
import Data.Vinyl.TypeLevel

import Indigo.Internal.Expr.Compilation
import Indigo.Internal.Expr.Types
import Indigo.Internal.Lookup
import Indigo.Internal.Object
import Indigo.Internal.SIS
import Indigo.Internal.State
import Indigo.Internal.Var
import Indigo.Lorentz
import qualified Lorentz.ADT as L
import qualified Lorentz.Instr as L
import Michelson.Typed.Haskell.Instr.Product (GetFieldType)
import Util.Type

-----------------------------------------
-- Object decomposition
-----------------------------------------

-- | Alike 'SomeIndigoState' datatype but without objects argument
type SIS' stk a = RefId -> StackVars stk -> (a, RefId, SomeGenCode stk)

-- | For given element on stack, generate code which
-- decomposes it to list of its deep non-decomposable fields.
-- Clean up code of 'SomeIndigoState' composes the value back.
deepDecomposeCompose
  :: forall a inp . IsObject a
  => SIS' (a : inp) (Object a)
deepDecomposeCompose
  | Just Dict <- complexObjectDict @a = \refId st ->
      let decomposedSt = fst (noRefGenCode @(FieldTypes a) $ popNoRef st) in
      withStack refId decomposedSt (decomposeComposeFields @(FieldTypes a)) $ \(result, newRefId, gc) ->
        ( Decomposed (typedToNamedRec @a typedToNamedFieldObj result)
        , newRefId
        , SomeGenCode $ GenCode
          { gcStack = gcStack gc
          , gcCode = L.deconstruct @a @(FieldTypes a) # gcCode gc
          , gcClear = gcClear gc # L.constructStack @a  @(FieldTypes a)
          }
        )
  | otherwise =
      \refId stk -> (Cell refId, refId + 1, SomeGenCode $
                      usingIndigoState (MetaData stk mempty emptyGenCodeHooks) (assignTopVar $ Var refId) )
  where
    decomposeComposeFields
      :: forall flds . (KnownList flds, AllConstrained IsObject flds)
      => SIS' (flds ++ inp) (Rec TypedFieldObj flds)
    decomposeComposeFields = case klist @flds of
      KNil -> \refId stk -> (RNil, refId, SomeGenCode $ GenCode stk L.nop L.nop)
      KCons (_ :: Proxy r) (_ :: Proxy rest) -> \refId st ->
        withStack refId (popNoRef st) (decomposeComposeFields @rest) $ \(resultRest, refId', restGc) ->
          withStack refId' (pushNoRef $ gcStack restGc) (deepDecomposeCompose @r) $ \(resultCur, newRefId, curGc) ->
            ( TypedFieldObj resultCur :& resultRest
            , newRefId
            , SomeGenCode $ GenCode
              { gcStack = gcStack curGc
              , gcCode = L.dip (gcCode restGc) # gcCode curGc
              , gcClear = gcClear curGc # L.dip (gcClear restGc)
              }
            )

withStack
  :: RefId
  -> StackVars inp
  -> SIS' inp a
  -> (forall out . (a, RefId, GenCode inp out) -> r)
  -> r
withStack refId stk sis f = case sis refId stk of
  (res, newRefId, SomeGenCode genCode) -> f (res, newRefId, genCode)

-----------------------------------------
-- Expr decomposition
-----------------------------------------

-- | Datatype representing decomposition of 'Expr'.
data ExprDecomposition inp a where
  ExprFields :: Rec Expr (FieldTypes a) -> ExprDecomposition inp a
  Deconstructed :: IndigoState inp (FieldTypes a ++ inp) -> ExprDecomposition inp a

-- | Decompose (shallowly) an expression to list of its direct fields.
decomposeExpr :: ComplexObjectC a => DecomposedObjects -> Expr a -> ExprDecomposition inp a
decomposeExpr _ (ConstructWithoutNamed _ fields) = ExprFields fields
decomposeExpr objs (V v) = withObject objs v $ decomposeObjectF namedToExpr
decomposeExpr objs (ObjMan objMan) = case runObjectManipulation objs objMan of
  StillObject obj -> decomposeObjectF unNamedFieldExpr obj
  OnStack comp -> deconstructOnStack comp
decomposeExpr _ ex = deconstructOnStack $ compileExpr ex

-- | Decompose any 'IndigoObjectF' with regards to decomposer for field.
decomposeObjectF
  :: forall a inp f . ComplexObjectC a
  => (forall name . f name -> Expr (GetFieldType a name))
  -> IndigoObjectF f a
  -> ExprDecomposition inp a
decomposeObjectF _ (Cell refId) =
  deconstructOnStack $
    IndigoState $ \md -> GenCode (pushNoRef $ mdStack md) (varActionGet @a refId $ mdStack md) L.drop
decomposeObjectF unF (Decomposed fields) =
  ExprFields $ namedToTypedRec @a unF fields

-- | Deconstruct top element of the stack and return it
-- wrapped into 'Deconstructed' constructor.
deconstructOnStack
  :: forall a inp . ComplexObjectC a
  => IndigoState inp (a : inp)
  -> ExprDecomposition inp a
deconstructOnStack fetchFld =
  Deconstructed $ IndigoState $ \md ->
    let (newSt, clean) = noRefGenCode @(FieldTypes a) (mdStack md) in
    GenCode newSt (gcCode (runIndigoState fetchFld md) # L.deconstruct @a @(FieldTypes a)) clean

-----------------------------------------
-- Helpers
-----------------------------------------

-- | Push the passed stack cells without references to them.
noRefGenCode
  :: forall rs inp . (KnownList rs, AllConstrained KnownValue rs)
  => StackVars inp -> (StackVars (rs ++ inp), (rs ++ inp) :-> inp)
noRefGenCode md = case klist @rs of
  KNil -> (md, L.nop)
  KCons Proxy (_ :: Proxy rest) -> bimap pushNoRef (L.drop #) (noRefGenCode @rest md)