-- 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
  ( LambdaPure1
  , createLambdaPure1
  , CreateLambdaPure1C
  , executeLambdaPure1
  , ExecuteLambdaPure1C
  , initMetaDataPure

  , Lambda1
  , createLambda1
  , CreateLambda1C
  , executeLambda1
  , ExecuteLambda1C
  , initMetaData

  , LambdaEff1
  , createLambdaEff1
  , CreateLambdaEff1C
  , executeLambdaEff1
  , ExecuteLambdaEff1C
  , initMetaDataEff

  , Lambda1Generic
  , LambdaExecutor
  , LambdaCreator
  ) where

import Data.Constraint (Dict(..))

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

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

type LambdaPure1 arg res = Lambda1Generic '[] arg res

type CreateLambdaPure1C arg res = CreateLambda1CGeneric '[] arg res

-- | Create a lambda, that takes only one argument, from the given computation.
-- The lambda is not allowed to modify storage and emit operations.
createLambdaPure1
  :: forall res arg inp out . CreateLambdaPure1C arg res
  => LambdaCreator '[] arg res inp out
createLambdaPure1 = createLambda1Generic initMetaDataPure

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
  => LambdaExecutor '[] arg res inp
executeLambdaPure1 = executeLambda1Generic @res (return ())

initMetaDataPure :: KnownValue arg => (Var arg, MetaData '[arg])
initMetaDataPure = let v = Cell 0 in (v, MetaData (Ref 0 :& RNil) 1)

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

type Lambda1 st arg res = Lambda1Generic '[st] arg res

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

-- | Create a lambda, that takes only one argument, from the given computation.
-- The lambda is not allowed to emit operations.
createLambda1
  :: forall st res arg inp out . CreateLambda1C st arg res
  => LambdaCreator '[st] arg res inp out
createLambda1 = createLambda1Generic initMetaData

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.
executeLambda1
  :: forall st res arg inp . ExecuteLambda1C st arg res
  => LambdaExecutor '[st] arg res inp
executeLambda1 =
  executeLambda1Generic @res
    -- 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 ->
      let GenCode _ newMd alloc _ = usingIndigoState md $ compileExpr (V (storageVar @st)) in
      let GenCode _ _ cleanup _   = usingIndigoState newMd (makeTopVar >>= (setVar (storageVar @st) . V)) in
      GenCode () newMd alloc (cleanup # L.drop)
    )

initMetaData :: (KnownValue arg, KnownValue st) => (Var arg, MetaData '[arg, st])
initMetaData =
  -- This numeration is intentional.
  -- We have to provide HasStorage for a lambda.
  -- To avoid excessive 'given' calls with new indexes,
  -- we just refer to storage variable with the same index.
  let argm = Cell 2 in
  (argm, MetaData (Ref 2 :& Ref 1 :& RNil) 3)

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

type LambdaEff1 st arg res = Lambda1Generic '[st, Ops] arg res

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

-- | Create a lambda, that takes only one argument, from the given computation,
-- and return a variable referring to this lambda.
-- The lambda is allowed to modify storage and emit operations.
createLambdaEff1
  :: forall st res arg inp out . CreateLambdaEff1C st arg res
  => LambdaCreator '[st, Ops] arg res inp out
createLambdaEff1 = createLambda1Generic initMetaDataEff

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 st res arg inp . ExecuteLambdaEff1C st arg res
  => LambdaExecutor '[st, Ops] arg res inp
executeLambdaEff1 =
  executeLambda1Generic @res
    -- 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 ->
      let GenCode _ newMd alloc _ =
              usingIndigoState md (do
                  compileExpr (V operationsVar)
                  compileExpr (V (storageVar @st))) in
      let (newStoreVar, newMdStore) = pushRefMd (pushNoRefMd md) in
      let (newOpsVar, newMdOps) = pushRefMd md in
      let cleanup =
            gcCode (usingIndigoState newMdStore $ setVar (storageVar @st) (V newStoreVar)) #
            L.drop #
            gcCode (usingIndigoState newMdOps $ setVar operationsVar (V newOpsVar)) #
            L.drop
      in GenCode () newMd alloc cleanup
    )

initMetaDataEff :: (KnownValue arg, KnownValue st) => (Var arg, MetaData '[arg, st, Ops])
initMetaDataEff =
  let argm = Cell 2 in
  (argm, MetaData (Ref 2 :& Ref 1 :& Ref 0 :& RNil) 3)

----------------------------------------------------------------------------
-- Common lambda functionality
----------------------------------------------------------------------------

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

type LambdaCreator extra arg res inp out
  = (Var arg -> IndigoState (arg & extra) out res)
  -> IndigoState inp (Lambda1Generic extra arg res & inp) (Var (Lambda1Generic extra arg res))

-- | 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 out . CreateLambda1CGeneric extra arg res
  => (Var arg, MetaData (arg & extra))
  -> (Var arg -> IndigoState (arg & extra) out res)
  -> IndigoState inp (Lambda1Generic extra arg res & inp) (Var (Lambda1Generic extra arg res))
createLambda1Generic (varArg, initMd) act = IndigoState $ \md ->
  let gc = runIndigoState (act varArg) initMd in
  let (var, md1) = pushRefMd md in
  GenCode var md1 (L.lambda (compileScope gc # liftClear @res @extra @(arg & extra) L.drop)) L.drop


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) (RetVars 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.
executeLambda1Generic
  :: forall res arg extra inp . ExecuteLambda1CGeneric extra arg res
  => IndigoState inp (extra ++ inp) ()
  -> Var (Lambda1Generic extra arg res)
  -> Expr arg
  -> IndigoState inp (RetOutStack res ++ inp) (RetVars res)
executeLambda1Generic allocateCleanup varF argm = IndigoState $ \md ->
  let GenCode _ allocMd allocate cleanup = runIndigoState allocateCleanup md in
  let getArgs =
        allocate #
        (gcCode $
          usingIndigoState allocMd $ do
              compileExpr argm
              compileToExpr (V varF)) in
  case listOfTypesConcatAssociativityAxiom @(RetOutStack res) @extra @inp of
    Dict ->
      let code = getArgs #
                 L.execute @_ @_ @inp #
                 liftClear @res cleanup
      in finalizeStatement @res md code