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

-- | This module implements the ability to put
-- Indigo computations on the stack as a lambda and execute them.
module Indigo.Backend.Lambda
  ( LambdaKind (..)
  , withLambdaKind
  , executeLambda1
  , initLambdaStackVars

  -- * Functionality for Frontend
  , CreateLambdaPure1C
  , ExecuteLambdaPure1C
  , CreateLambda1C
  , ExecuteLambda1C
  , CreateLambdaEff1C
  , ExecuteLambdaEff1C

  -- * Functionality for Sequential
  , CreateLambda1CGeneric
  , createLambda1Generic
  , Lambda1Generic
  ) where

import Data.Constraint (Dict(..))
import Fmt (fmt, pretty, (+|), (|+))

import Indigo.Backend.Prelude
import Indigo.Backend.Scope
import Indigo.Backend.Var
import Indigo.Internal hiding ((+), (<>))
import Indigo.Lorentz
import qualified Lorentz.Instr as L
import Lorentz.Zip (ZipInstr, ZippedStack)
import Util.Type (type (++), KnownList, listOfTypesConcatAssociativityAxiom)

----------------------------------------------------------------------------
-- External interface
----------------------------------------------------------------------------

-- | Describes kind of lambda: pure, modifying storage, effectfull
data LambdaKind st arg res extra where
  PureLambda ::
    (ExecuteLambdaPure1C arg res, CreateLambda1CGeneric '[] arg res, Typeable res)
    => LambdaKind st arg res '[]
  StorageLambda ::
    (ExecuteLambda1C st arg res, CreateLambda1CGeneric '[st] arg res, Typeable res)
    => Proxy st
    -> LambdaKind st arg res '[st]
  EffLambda
    :: (ExecuteLambdaEff1C st arg res, CreateLambda1CGeneric '[st, Ops] arg res, Typeable res)
    => Proxy st
    -> LambdaKind st arg res '[st, Ops]

-- | Provide common constraints that are presented in all constructors of 'LambdaKind'
withLambdaKind
  :: LambdaKind st arg res extra
  -> ((ScopeCodeGen res, KnownValue arg, Typeable res, CreateLambda1CGeneric extra arg res) => r)
  -> r
withLambdaKind PureLambda r = r
withLambdaKind (StorageLambda _) r = r
withLambdaKind (EffLambda _) r = r

-- | Execute lambda depending on its 'LambdaKind'
executeLambda1
  :: forall res st arg extra inp .
     LambdaKind st arg res extra
  -- ^ Kind of lambda
  -> RefId
  -- ^ Next free variable reference
  -> RetVars res
  -- ^ Variable that will be assigned to the resulting value
  -> LambdaExecutor extra arg res inp
executeLambda1 lambdaKind refId retVars lambdaVar argEx = withLambdaKind lambdaKind $
  let execStmt = fmt $ "executeLambda (lambdaVar = " +| lambdaVar |+ ", arg = " +| argEx |+ ")" in
  stmtHookState (prettyAssign @res retVars execStmt) $
    case lambdaKind of
      PureLambda      -> executeLambdaPure1 @res       retVars lambdaVar argEx
      StorageLambda _ -> executeLambdaSt1   @res refId retVars lambdaVar argEx
      EffLambda _     -> executeLambdaEff1  @res refId retVars lambdaVar argEx

-- | Create initial stack vars depending on 'LambdaKind'
initLambdaStackVars :: LambdaKind st arg res extra -> Var arg -> StackVars (arg : extra)
initLambdaStackVars PureLambda = initStackVarsPure
initLambdaStackVars (StorageLambda _) = initStackVars
initLambdaStackVars (EffLambda _) = initStackVarsEff

type Lambda1Generic extra arg res = (arg : extra) :-> (RetOutStack res ++ extra)

type CreateLambda1CGeneric extra arg res =
  ( ScopeCodeGen res, KnownValue arg, Typeable extra
  , ZipInstr (arg : extra)
  , KnownValue (ZippedStack (arg ': extra))
  , KnownValue (ZippedStack (RetOutStack res ++ extra))
  , ZipInstr (RetOutStack res ++ extra)
  , Typeable (RetOutStack res ++ extra)
  )

-- | Create a lambda, that takes only one argument, from the given computation,
-- and return a variable referring to this lambda.
createLambda1Generic
  :: forall arg res extra inp . CreateLambda1CGeneric extra arg res
  => Var (Lambda1Generic extra arg res)
  -> res
  -> StackVars (arg : extra)
  -> SomeIndigoState (arg : extra)
  -> IndigoState inp (Lambda1Generic extra arg res : inp)
createLambda1Generic var ret initMd act = IndigoState $ \md@MetaData{..} ->
  -- Decomposed objects are passed as mempty here because in the lambda
  -- we don't decompose storage value (but we might be doing it as an optimisation)
  -- so we just have it as an stack cell
  runSIS act (MetaData initMd mempty mdHooks) $ \lambdaBody ->
    let gcStack = pushRef var mdStack
        gcCode =
          stmtHook md (prettyAssign @(Var (Lambda1Generic extra arg res)) var "createLambda") $
            L.lambda (compileScope (replStkMd md) lambdaBody ret #
                      liftClear @res @extra @(arg : extra) L.drop)
        gcClear = L.drop
    in GenCode {..}

----------------------------------------------------------------------------
-- Pure lambdas
----------------------------------------------------------------------------

type CreateLambdaPure1C arg res = CreateLambda1CGeneric '[] arg res

type ExecuteLambdaPure1C arg res = ExecuteLambda1CGeneric '[] arg res

-- | Execute a lambda, which accepts only one argument, on passed expression.
executeLambdaPure1
  :: forall res arg inp. ExecuteLambdaPure1C arg res
  => RetVars res
  -- ^ Variable(s) that will be assigned to the resulting value(s)
  -> LambdaExecutor '[] arg res inp
executeLambdaPure1 retVars = executeLambda1Generic @res retVars nopState

initStackVarsPure :: KnownValue arg => Var arg -> StackVars '[arg]
initStackVarsPure var = pushRef var emptyStack

----------------------------------------------------------------------------
-- Impure lambda (modifying storage only)
----------------------------------------------------------------------------

type CreateLambda1C st arg res = (KnownValue st, CreateLambda1CGeneric '[st] arg res)

type ExecuteLambda1C st arg res =
  ( IsObject st
  , HasStorage st
  , ExecuteLambda1CGeneric '[st] arg res
  )

-- | Execute a lambda that accepts only one argument on the given expression.
executeLambdaSt1
  :: forall res st arg inp. ExecuteLambda1C st arg res
  => RefId
  -> RetVars res
  -- ^ Variable(s) that will be assigned to the resulting value(s)
  -> LambdaExecutor '[st] arg res inp
executeLambdaSt1 nextRef retVars = executeLambda1Generic @res retVars $
    IndigoState $ \md ->
      let storage = storageVar @st
          -- TODO this @compileExpr (V (storageVar @st))@ call materialises the whole decomposed storage.
          -- This is pretty expensive operation and it has to be fixed:
          -- we have to materialise only fields used in the lambda
          GenCode resStack fetchStorage _ =
            usingIndigoState md $ exprHookState "Computing storage" $ compileExpr (V storage)
          tmpVar = Var nextRef
          setStorage =
            auxiliaryHook md "Update storage with returned from lambda" $
              gcCode (usingIndigoState (pushRefMd tmpVar md) $
                      setVar (nextRef + 1) storage (V tmpVar))
              # L.drop
      in GenCode resStack fetchStorage setStorage


initStackVars :: (HasStorage st, KnownValue arg) => Var arg -> StackVars '[arg, st]
initStackVars var = emptyStack
  & pushRef storageVar
  & pushRef var
  -- This 'storageVar' usage is intentional.
  -- We have to provide 'HasStorage' for a lambda.
  -- To avoid excessive 'given' calls with new indexes,

----------------------------------------------------------------------------
-- Lambda with side effects (might emit operations)
----------------------------------------------------------------------------

type CreateLambdaEff1C st arg res = (KnownValue st, CreateLambda1CGeneric '[st, Ops] arg res)

type ExecuteLambdaEff1C st arg res =
  ( HasStorage st
  , HasSideEffects
  , IsObject st
  , ExecuteLambda1CGeneric '[st, Ops] arg res
  )

-- | Execute a lambda that accepts only one argument on the given expression.
-- Also updates the storage and operations with the values returned from the lambda.
executeLambdaEff1
  :: forall res st arg inp. ExecuteLambdaEff1C st arg res
  => RefId
  -> RetVars res
  -- ^ Variable(s) that will be assigned to the resulting value(s)
  -> LambdaExecutor '[st, Ops] arg res inp
executeLambdaEff1 nextRef retVars =
  executeLambda1Generic @res retVars $
    -- TODO this @compileExpr (V (storageVar @st))@ call materialises the whole decomposed storage.
    -- This is pretty expensive operation and it has to be fixed:
    -- we have to materialise only fields used in the lambda
    IndigoState $ \md@MetaData{..} ->
      let storage = storageVar @st
          ops@(Var opsRefId) = operationsVar
          gcStack = pushRef storage $ pushRef ops mdStack
          fetchCode =
            auxiliaryHook md "Fetching operations" (varActionGet opsRefId mdStack) #
            (exprHook md "Computing storage" $
              gcCode $ usingIndigoState (replStkMd md sPlus) $ compileExpr (V storage))
          sPlus = pushNoRef mdStack
          tmpVar = Var nextRef
          setStorage =
            auxiliaryHook md "Update storage with returned from lambda" $
              gcCode (usingIndigoState (replStkMd md (pushRef tmpVar sPlus)) $ setVar (nextRef + 1) storage (V tmpVar))
              # L.drop
          gcClear = setStorage #
                    auxiliaryHook md "Update operations with returned from lambda" (varActionSet opsRefId mdStack)
      in GenCode {gcCode=fetchCode,..}

initStackVarsEff
  :: (HasSideEffects, HasStorage st, KnownValue arg)
  => Var arg -> StackVars '[arg, st, Ops]
initStackVarsEff var = emptyStack
  & pushRef operationsVar
  & pushRef storageVar
  & pushRef var

----------------------------------------------------------------------------
-- Generic functionality of lambda execution
----------------------------------------------------------------------------

type ExecuteLambda1CGeneric extra arg res =
  ( ScopeCodeGen res, KnownValue arg
  , KnownValue ((arg : extra) :-> (RetOutStack res ++ extra))
  , KnownList extra
  , ZipInstr (arg : extra)
  , KnownList (RetOutStack res ++ extra)
  , ZipInstr (RetOutStack res ++ extra)
  , Typeable (RetOutStack res ++ extra)
  , KnownValue (ZippedStack (RetOutStack res ++ extra))
  )

type LambdaExecutor extra arg res inp
   = Var (Lambda1Generic extra arg res)
  -> Expr arg
  -> IndigoState inp (RetOutStack res ++ inp)

-- | Execute a lambda that accepts only one argument on the given expression.
-- Also updates the storage and operations with the values returned from the lambda.
executeLambda1Generic
  :: forall res arg extra inp . ExecuteLambda1CGeneric extra arg res
  => RetVars res
  -> IndigoState inp (extra ++ inp)
  -> Var (Lambda1Generic extra arg res)
  -> Expr arg
  -> IndigoState inp (RetOutStack res ++ inp)
executeLambda1Generic vars allocateCleanup varF argm = IndigoState $ \md@MetaData{..} ->
  let GenCode allocStk allocate cleanup = usingIndigoState md allocateCleanup in
  let getArgs =
        auxiliaryHook md "Computing implicit lambda arguments" allocate #
        (gcCode $
          usingIndigoState (replStkMd md allocStk) $ do
              exprHookState ("Computing lambda parameter: " <> pretty argm) (compileExpr argm)
              exprHookState "Fetching lambda" (compileExpr (V varF))) in
  case listOfTypesConcatAssociativityAxiom @(RetOutStack res) @extra @inp of
    Dict ->
      let code = getArgs #
                 L.execute @_ @_ @inp #
                 liftClear @res cleanup
      in finalizeStatement @res mdStack vars code