Lambda.hs

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

{-# LANGUAGE NoMonomorphismRestriction #-}

-- | 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 ((\\))
import Fmt (fmt, pretty, (+|), (|+))

import Indigo.Backend.Expr.Compilation (compileExpr)
import Indigo.Backend.Lookup (varActionGet, varActionSet)
import Indigo.Backend.Prelude
import Indigo.Backend.Scope
import Indigo.Backend.Var
import Indigo.Common.Expr (Expr(V))
import Indigo.Common.Object (IsObject)
import Indigo.Common.SIS (SomeIndigoState, runSIS)
import Indigo.Common.State
import Indigo.Common.Var
import Indigo.Lorentz
import Lorentz.Instr qualified as L
import Lorentz.Zip (ZipInstr, ZippedStack)
import Morley.Michelson.Typed.Instr.Internal.Proofs (assocThm)
import Morley.Util.Type (KnownList, type (++))

----------------------------------------------------------------------------
-- 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 = WrappedLambda (arg : extra) (RetOutStack res ++ extra)

type CreateLambda1CGeneric extra arg res =
  ( ScopeCodeGen res, KnownValue arg, Typeable extra, KnownList 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 = operationsVar
          opsRefId = case ops of Var x -> x
          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 (WrappedLambda (arg : extra) (RetOutStack res ++ extra))
  , KnownList extra
  , ZipInstr (arg : extra)
  , KnownList (RetOutStack res ++ extra)
  , ZipInstr (RetOutStack res ++ extra)
  , Typeable (RetOutStack res ++ extra)
  , Typeable 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
      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)))
      code = getArgs #
              L.execute @_ @_ @inp #
              liftClear @res cleanup
              \\ assocThm @(RetOutStack res) @extra @inp
  in finalizeStatement @res mdStack vars code