Sequential.hs

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

-- | 'Instruction' datatype and its compilations.
--
-- The idea behind this module is to provide an intermediate representation that
-- can:
--
-- - be generated from the frontend freer monad
-- - be compiled to the backend 'IndigoState'
-- - be easy to analyze, manipulate and modify
--
-- This is meant to be the common ground for modular optimizations on the code
-- before this is handed off to the backend and eventually translated to
-- Michelson.

module Indigo.Compilation.Sequential
  ( Block
  , Instruction (..)

  , IndigoSeqCaseClause (..)
  , CaseBranch (..)

  -- * Translations
  , SequentialHooks (..)
  , stmtHookL
  , InstrCollector (..)
  , indigoMtoSequential
  , sequentialToLorentz

  -- * Case machinery
  , updateClauses
  , mapMClauses
  ) where

import Prelude

import Data.Vinyl.Core (RMap(..))

import Indigo.Backend
import Indigo.Common.SIS
import Indigo.Common.State hiding ((>>))
import Indigo.Common.State qualified as St
import Indigo.Common.Var
import Indigo.Compilation.Sequential.Types
import Indigo.Frontend.Internal.Statement qualified as S
import Indigo.Frontend.Program (IndigoM(..), interpretProgram)
import Indigo.Lorentz hiding (comment)
import Morley.Michelson.Typed qualified as MT

----------------------------------------------------------------------------
-- Translations
----------------------------------------------------------------------------

-- | Transformation from 'IndigoM' to a 'Block' of 'Instruction's.
--
-- Requires the first non-used 'RefId' and returns the next one.
indigoMtoSequential
  :: RefId
  -> SequentialHooks
  -> IndigoM a
  -> (Block, RefId)
indigoMtoSequential refId hook code =
  let InstrCollector {..} = snd $ instrCollect refId hook code
  in (instrList, nextRef)

-- | Collects instructions starting from an 'IndigoM'.
-- Returns an 'InstrCollector' as well as the return value for that 'IndigoM'.
instrCollect :: RefId -> SequentialHooks -> IndigoM a -> (a, InstrCollector)
instrCollect ref hooks (IndigoM imCode) =
  let instrColl = InstrCollector ref [] hooks
      (res, resColl) = usingState instrColl $ interpretProgram collectStatement imCode
  in (res, InstrCollector (nextRef resColl) (reverse $ instrList resColl) hooks)

-- | Collects instructions starting from 'S.StatementF'.
-- IMPORTANT: the instructions are collected in the opposite order (as a stack).
collectStatement :: S.StatementF IndigoM a -> State InstrCollector a
collectStatement = \case
  S.LiftIndigoState is -> appendNewInstr $ LiftIndigoState is
  S.CalledFrom callStk iM -> do
    InstrCollector nRef _prevInstrs hooks <- get
    let (res, inner) = instrCollect nRef hooks iM
    modify $ \s -> s {nextRef = nextRef inner}
    res <$ shStmtHook hooks callStk (instrList inner)
  S.NewVar ex -> do
    var <- mkNextVar
    appendNewInstr $ AssignVar var ex
    return var
  S.SetVar vx ex -> appendNewInstr $ SetVar vx ex
  S.VarModification upd vx ey -> appendNewInstr $ VarModification upd vx ey
  S.SetField vStore fname exF -> appendNewInstr $ SetField vStore fname exF

  S.LambdaCall1 lKind lName vIm ex -> withLambdaKind lKind $ do
    (var, block, ret, retVars) <- collectInLambda vIm
    appendNewInstr $ LambdaCall1 lKind lName ex var block ret retVars
    return retVars

  S.Scope (iM :: IndigoM ret) -> do
    retVars <- allocateVars @ret mkNextVar
    (ret, block) <- collectInner iM
    appendNewInstr $ Scope block ret retVars
    return retVars
  S.If ex (iMa :: IndigoM ret) iMb -> do
    retVars <- allocateVars @ret mkNextVar
    (retA, blockA) <- collectInner iMa
    (retB, blockB) <- collectInner iMb
    appendNewInstr $ If ex blockA retA blockB retB retVars
    return retVars
  S.IfSome ex (vIMa :: Var x -> IndigoM ret) iMb -> do
    retVars <- allocateVars @ret mkNextVar
    varX <- mkNextVar
    (retA, blockA) <- collectInner $ vIMa varX
    (retB, blockB) <- collectInner iMb
    appendNewInstr $ IfSome ex varX blockA retA blockB retB retVars
    return retVars
  S.IfRight ex (vIMa :: Var x -> IndigoM ret) vIMb -> do
    retVars <- allocateVars @ret mkNextVar
    varR <- mkNextVar
    (retA, blockA) <- collectInner $ vIMa varR
    varL <- mkNextVar
    (retB, blockB) <- collectInner $ vIMb varL
    appendNewInstr $ IfRight ex varR blockA retA varL blockB retB retVars
    return retVars
  S.IfCons ex (vvIMa :: Var x -> Var (List x) -> IndigoM ret) iMb -> do
    retVars <- allocateVars @ret mkNextVar
    varX <- mkNextVar
    varLX <- mkNextVar
    (retA, blockA) <- collectInner $ vvIMa varX varLX
    (retB, blockB) <- collectInner iMb
    appendNewInstr $ IfCons ex varX varLX blockA retA blockB retB retVars
    return retVars

  S.Case grd clauses -> do
    retVars <- allocateClausesVars clauses
    blockClauses <- collectClauses clauses
    appendNewInstr $ Case grd blockClauses retVars
    return retVars
  S.EntryCase proxy grd clauses -> do
    retVars <- allocateClausesVars clauses
    blockClauses <- collectClauses clauses
    appendNewInstr $ EntryCase proxy grd blockClauses retVars
    return retVars
  S.EntryCaseSimple grd clauses -> do
    retVars <- allocateClausesVars clauses
    blockClauses <- collectClauses clauses
    appendNewInstr $ EntryCaseSimple grd blockClauses retVars
    return retVars

  S.While ex iM -> do
    ((), block) <- collectInner iM
    appendNewInstr $ While ex block
  S.WhileLeft ex vIm -> do
    varL <- mkNextVar
    varR <- mkNextVar
    ((), block) <- collectInner $ vIm varL
    appendNewInstr $ WhileLeft ex varL block varR
    return varR
  S.ForEach ex vIm -> do
    varIop <- mkNextVar
    ((), block) <- collectInner $ vIm varIop
    appendNewInstr $ ForEach ex varIop block

  S.ContractName tx iM -> do
    ((), block) <- collectInner iM
    appendNewInstr $ ContractName tx block
  S.DocGroup dg iM -> do
    ((), block) <- collectInner iM
    appendNewInstr $ DocGroup dg block
  S.ContractGeneral iM -> do
    ((), block) <- collectInner iM
    appendNewInstr $ ContractGeneral block
  S.FinalizeParamCallingDoc vIm param -> do
    varCp <- mkNextVar
    ((), block) <- collectInner $ vIm varCp
    appendNewInstr $ FinalizeParamCallingDoc varCp block param

  S.TransferTokens ex exm exc ->
    appendNewInstr $ TransferTokens ex exm exc
  S.SetDelegate ex ->
    appendNewInstr $ SetDelegate ex
  S.CreateContract ctrc exk exm exs -> do
    varAddr <- mkNextVar
    appendNewInstr $ CreateContract ctrc exk exm exs varAddr
    return varAddr
  S.SelfCalling proxy ep -> do
    varCR <- mkNextVar
    appendNewInstr $ SelfCalling proxy ep varCR
    return varCR
  S.ContractCalling proxy epRef exAddr -> do
    varMcr <- mkNextVar
    appendNewInstr $ ContractCalling proxy epRef exAddr varMcr
    return varMcr
  S.Emit tag ex ->
    appendNewInstr $ Emit tag ex

  S.Fail (_ :: Proxy ret) failure -> do
    appendNewInstr $ Fail failure
    -- Note: because this is a failing instr, this vars are effectively never used
    allocateVars @ret mkNextVar
  S.FailOver (_ :: Proxy ret) failure ex -> do
    appendNewInstr $ FailOver failure ex
    -- Note: because this is a failing instr, this vars are effectively never used
    allocateVars @ret mkNextVar

-- | Continues collecting 'Instruction's from an inner 'IndigoM' (e.g. scoped).
-- This keeps advancing the ref counter as well.
collectInner :: IndigoM ret -> State InstrCollector (ret, Block)
collectInner iM = do
  iColl <- get
  let (ret, InstrCollector newRef block _) = instrCollect (nextRef iColl) (seqHooks iColl) iM
  put $ iColl {nextRef = newRef}
  return (ret, block)

-- | Just a common set of steps used by collection of single-arg lambda's values.
collectInLambda
  :: ScopeCodeGen ret
  => (Var arg -> IndigoM ret)
  -> State InstrCollector (Var arg, Block, ret, RetVars ret)
collectInLambda vIm = do
  var <- mkNextVar
  (ret :: ret, block) <- collectInner $ vIm var
  retVars <- allocateVars @ret mkNextVar
  return (var, block, ret, retVars)

-- | Append a new 'Instruction' to the head of the list in the state.
appendNewInstr :: Instruction -> State InstrCollector ()
appendNewInstr is = modify $ \iColl -> iColl {instrList = is : instrList iColl}

-- | Creates a new var. This simply advances the ref counter and updates it.
mkNextVar :: State InstrCollector (Var a)
mkNextVar = do
  iColl <- get
  let ref = nextRef iColl
  put $ iColl {nextRef = ref + 1}
  return $ Var ref

-- | Translation from a 'Block' and an initial 'MetaData' to Lorentz.
sequentialToLorentz
  :: MetaData inp
  -> (Block, RefId)
  -> inp :-> inp
sequentialToLorentz md block =
  runSIS (sequentialToSIS block) md St.cleanGenCode

-- | Translation from a 'Block' to a 'SomeIndigoState'.
sequentialToSIS :: (Block, RefId) -> SomeIndigoState inp
sequentialToSIS ([], _) = toSIS St.nopState
sequentialToSIS (x : xs, refId) = instrToSIS refId x `thenSIS` sequentialToSIS (xs, refId)

-- | Translation from a single 'Instruction' to a 'SomeIndigoState'.
instrToSIS :: RefId -> Instruction -> SomeIndigoState inp
instrToSIS nextRef = \case
  LiftIndigoState sis -> sis
  Comment txt -> toSIS $ comment $ MT.JustComment txt
  AssignVar vx ex -> toSIS $ assignVar vx ex
  SetVar vx ex -> toSIS $ setVar nextRef vx ex
  VarModification upd vx ey -> toSIS $ updateVar nextRef upd vx ey
  SetField vSt lName ex -> toSIS $ setField nextRef vSt lName ex

  LambdaCall1 lKind _lName ex var block ret retVars ->
    withLambdaKind lKind $
      toSIS $ scope (sequentialToSIS (AssignVar var ex : block, nextRef)) ret retVars
  CreateLambda1 lamMd _var body ret varLam ->
    toSIS $ createLambda1Generic varLam ret lamMd (sequentialToSIS (body, nextRef))
  ExecLambda1 lKind (Proxy :: Proxy ret) ex varLam retVars ->
    toSIS $ executeLambda1 @ret lKind nextRef retVars varLam ex

  Scope block ret retVars ->
    toSIS $ scope (sequentialToSIS (block, nextRef)) ret retVars
  If ex blockA retA blockB retB retVars ->
    toSIS $ if_ ex (sequentialToSIS (blockA, nextRef)) retA (sequentialToSIS (blockB, nextRef)) retB retVars
  IfSome ex varX blockA retA blockB retB retVars ->
    toSIS $ ifSome ex varX (sequentialToSIS (blockA, nextRef)) retA (sequentialToSIS (blockB, nextRef)) retB retVars
  IfRight ex varR blockA retA varL blockB retB retVars ->
    toSIS $ ifRight ex varR (sequentialToSIS (blockA, nextRef)) retA varL (sequentialToSIS (blockB, nextRef)) retB retVars
  IfCons ex varX varLX blockA retA blockB retB retVars ->
    toSIS $ ifCons ex varX varLX (sequentialToSIS (blockA, nextRef)) retA (sequentialToSIS (blockB, nextRef)) retB retVars

  Case grd blockClauses retVars ->
    toSIS $ caseRec grd (clausesToBackend nextRef blockClauses) retVars
  EntryCase proxy grd blockClauses retVars ->
    toSIS $ entryCaseRec proxy grd (clausesToBackend nextRef blockClauses) retVars
  EntryCaseSimple grd blockClauses retVars ->
    toSIS $ entryCaseSimpleRec grd (clausesToBackend nextRef blockClauses) retVars

  While ex block ->
    toSIS $ while ex (sequentialToSIS (block, nextRef))
  WhileLeft ex varL block varR ->
    toSIS $ whileLeft ex varL (sequentialToSIS (block, nextRef)) varR
  ForEach ex varIop block ->
    toSIS $ forEach ex varIop (sequentialToSIS (block, nextRef))

  ContractName tx block ->
    docGroup (DName tx) (sequentialToSIS (block, nextRef))
  DocGroup dg block ->
    docGroup dg (sequentialToSIS (block, nextRef))
  ContractGeneral block ->
    docGroup DGeneralInfoSection (sequentialToSIS (block, nextRef))
  FinalizeParamCallingDoc varCp block param ->
    finalizeParamCallingDoc varCp (sequentialToSIS (block, nextRef)) param

  TransferTokens ex exm exc ->
    toSIS $ transferTokens ex exm exc
  SetDelegate ex ->
    toSIS $ setDelegate ex
  CreateContract ctrc exk exm exs varAddr ->
    toSIS $ createContract ctrc exk exm exs varAddr
  SelfCalling (Proxy :: Proxy p) ep varCR ->
    toSIS $ selfCalling @p ep varCR
  ContractCalling (Proxy :: Proxy (cp, vd)) epRef exAddr varMcr ->
    toSIS $ contractCalling @cp @vd epRef exAddr varMcr
  Emit tag ex ->
    toSIS $ emit tag ex

  Fail failure -> failure
  FailOver failure ex -> failure ex

----------------------------------------------------------------------------
-- Case machinery
----------------------------------------------------------------------------

-- | Convert clauses from their "sequential" representation to the "backend" one.
clausesToBackend
  :: forall ret dt . RMap dt
  => RefId
  -> Rec (IndigoSeqCaseClause ret) dt
  -> Rec (IndigoCaseClauseL ret) dt
clausesToBackend nextRef = rmap $
  \(OneFieldIndigoSeqCaseClause cName (CaseBranch vx block ret)) ->
    cName /-> (IndigoClause vx (sequentialToSIS (block, nextRef)) ret)

-- | Allocate vars for the return value(s) of a clause-like 'Instruction'.
allocateClausesVars
  :: forall ret dt. ReturnableValue ret
  => Rec (S.IndigoMCaseClauseL IndigoM ret) dt
  -> State InstrCollector (RetVars ret)
allocateClausesVars _ = allocateVars @ret mkNextVar

-- | Collects clauses of a case-like statement.
collectClauses
  :: Rec (S.IndigoMCaseClauseL IndigoM ret) dt
  -> State InstrCollector (Rec (IndigoSeqCaseClause ret) dt)
collectClauses RNil = return RNil
collectClauses ((S.OneFieldIndigoMCaseClauseL cName clause) :& xs) = do
  varX <- mkNextVar
  (ret, block) <- collectInner $ clause varX
  let clauseX = OneFieldIndigoSeqCaseClause cName (CaseBranch varX block ret)
  clauseXs <- collectClauses xs
  return $ clauseX :& clauseXs

-- | Applies the given 'Block' to 'Block' transformation to the inner code block
-- of every case clause.
updateClauses
  :: (Block -> Block)
  -> Rec (IndigoSeqCaseClause ret) dt
  -> Rec (IndigoSeqCaseClause ret) dt
updateClauses _ RNil = RNil
updateClauses f (x :& xs) = case x of
  OneFieldIndigoSeqCaseClause cName (CaseBranch vx block ret) ->
    OneFieldIndigoSeqCaseClause cName (CaseBranch vx (f block) ret)
      :& updateClauses f xs

-- | Applies the given monadic function giving it the inner code block of each
-- case clause, in order.
mapMClauses :: Monad m => (Block -> m ()) -> Rec (IndigoSeqCaseClause ret) dt -> m ()
mapMClauses _ RNil = return ()
mapMClauses f (x :& xs) = case x of
  OneFieldIndigoSeqCaseClause _cName (CaseBranch _ block _) ->
    f block >> mapMClauses f xs