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

module Indigo.Compilation.Lambda
  ( compileLambdas
  ) where

import Prelude

import qualified Data.Map as M

import Indigo.Backend as B
import Indigo.Compilation.Sequential
import Indigo.Internal.Var

-- | Collects named lambdas that are used more than once and separates them into
-- a lambda creation and multiple lambda executions.
-- Leaves the remaining lambdas untouched, to be compiled inline.
compileLambdas :: (Block, RefId) -> (Block, RefId)
compileLambdas (block, nextRef) = (mkLambdas <> updatedBlock, newNextRef)
  where
    lambdaSet = collectNotInlinableLambdas block
    (lambdaRefDefs, newNextRef) = createLambdaRefs nextRef lambdaSet
    mkLambdas = createAllLambdas lambdaRefDefs
    updatedBlock = updateBlock block . M.fromList $ map (first ldName) lambdaRefDefs

-- | Collect all used lambdas in a computation that are called at least twice.
-- Only the outer lambdas will be gathered, for example, if we call lambda "func1"
-- from "func0", only "func0" will be considered.
collectNotInlinableLambdas :: Block -> Set Lambda1Def
collectNotInlinableLambdas = M.keysSet . M.filter (> 1) . executingState mempty . lookForLambdas

-- | Associates each given 'Lambda1Def' to a new 'RefId', starting from the given
-- one. Also returns the first unused 'RefId'
createLambdaRefs :: RefId -> Set Lambda1Def -> ([(Lambda1Def, RefId)], RefId)
createLambdaRefs nextRef =
  foldr (\lm (lst, ref) -> ((lm, ref) : lst, ref + 1)) ([], nextRef)

-- | Generates an 'Instruction' for each given tuple, to generate a lambda
-- (assigned to the respective variable) and leave it on the stack.
createAllLambdas :: [(Lambda1Def, RefId)] -> Block
createAllLambdas = map $ \(Lambda1Def {..}, lamRef) ->
  CreateLambda1 ldStack ldArgVar ldBody ldRet (Var lamRef)

-- | Updates a 'Block', it looks for lambda "Calls" (defined and used in place)
-- to replace them with lambda "Exec", provided there is a known variable for an
-- already created lambda.
updateBlock :: Block -> Map String RefId -> Block
updateBlock blk lambdaMap = updateBlock' blk
  where
    updateBlock' :: Block -> Block
    updateBlock' = map $ \case
      -- Instructions not concerned, will be kept the same
      LiftIndigoState sis -> LiftIndigoState sis
      Comment txt         -> Comment txt
      AssignVar vx ex -> AssignVar vx ex
      SetVar vx ex -> SetVar vx ex
      VarModification upd vx ey -> VarModification upd vx ey
      SetField vSt lName ex -> SetField vSt lName ex

      -- Lambda instructions to check for possible replacement
      lc@(LambdaCall1 lKind lName ex _var _block _ret retVars) ->
        case M.lookup lName lambdaMap of
          Nothing -> lc
          Just ref -> ExecLambda1 lKind Proxy ex (Var ref) retVars

      -- Lambda instructions not concerned, nothing to replace here
      c@(CreateLambda1{}) -> c
      e@(ExecLambda1{}) -> e

      -- Instructions with deeper code blocks to replace as well
      Scope block ret retVars ->
        Scope (updateBlock' block) ret retVars
      If ex blockA retA blockB retB retVars ->
        If ex (updateBlock' blockA) retA (updateBlock' blockB) retB retVars
      IfSome ex varX blockA retA blockB retB retVars ->
        IfSome ex varX (updateBlock' blockA) retA (updateBlock' blockB) retB retVars
      IfRight ex varR blockA retA varL blockB retB retVars ->
        IfRight ex varR (updateBlock' blockA) retA varL (updateBlock' blockB) retB retVars
      IfCons ex varX varLX blockA retA blockB retB retVars ->
        IfCons ex varX varLX (updateBlock' blockA) retA (updateBlock' blockB) retB retVars

      Case grd blockClauses retVars ->
        Case grd (updateClauses updateBlock' blockClauses) retVars
      EntryCase proxy grd blockClauses retVars ->
        EntryCase proxy grd (updateClauses updateBlock' blockClauses) retVars
      EntryCaseSimple grd blockClauses retVars ->
        EntryCaseSimple grd (updateClauses updateBlock' blockClauses) retVars

      While ex block ->
        While ex (updateBlock' block)
      WhileLeft ex varL block varR ->
        WhileLeft ex varL (updateBlock' block) varR
      ForEach varIop ex block ->
        ForEach varIop ex (updateBlock' block)

      ContractName tx block ->
        ContractName tx (updateBlock' block)
      DocGroup dg block ->
        DocGroup dg (updateBlock' block)
      ContractGeneral block ->
        ContractGeneral (updateBlock' block)
      FinalizeParamCallingDoc varCp block param ->
        FinalizeParamCallingDoc varCp (updateBlock' block) param

      -- Instructions not concerned, will be kept the same
      TransferTokens ex exm exc -> TransferTokens ex exm exc
      SetDelegate ex -> SetDelegate ex
      CreateContract varAddr ctrc exk exm exs -> CreateContract varAddr ctrc exk exm exs
      SelfCalling proxy varCR ep -> SelfCalling proxy varCR ep
      ContractCalling varMcr pCp epRef exAddr -> ContractCalling varMcr pCp epRef exAddr

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

-- Like 'collectLambdas', but uses 'State' to collect the 'Map' of all outer
-- lambdas encountered, including those used once.
lookForLambdas :: Block -> State (Map Lambda1Def Word) ()
lookForLambdas blk = forM_ blk match
  where
    -- pva701: it's crucial to have this function 'match' instead of code like
    -- @forM_ blk match $ \case@
    --        ... cases here ...
    -- because in the case of code above compilation of this function takes about 5-6 minutes
    -- it would be nice to figure out why (inspecting generated by GHC code)
    match :: Instruction -> State (Map Lambda1Def Word) ()
    match = \case
      -- Lambda instruction to collect
      LambdaCall1 lKind ldName _ex ldArgVar ldBody ldRet _retVars -> do
        let ldStack = initLambdaStackVars lKind ldArgVar
        withLambdaKind lKind $ addLambda $ Lambda1Def {..}

      -- Instructions with deeper code block to look into
      Scope block _ _ -> lookForLambdas block
      If _ blockA _ blockB _ _ ->
        lookForLambdas blockA >> lookForLambdas blockB
      IfSome _ _ blockA _ blockB _ _ ->
        lookForLambdas blockA >> lookForLambdas blockB
      IfRight _ _ blockA _ _ blockB _ _ ->
        lookForLambdas blockA >> lookForLambdas blockB
      IfCons _ _ _ blockA _ blockB _ _ ->
        lookForLambdas blockA >> lookForLambdas blockB
      Case _ blockClauses _ ->
        mapMClauses lookForLambdas blockClauses
      EntryCase _ _ blockClauses _ ->
        mapMClauses lookForLambdas blockClauses
      EntryCaseSimple _ blockClauses _ ->
        mapMClauses lookForLambdas blockClauses
      While _ block -> lookForLambdas block
      WhileLeft _ _ block _ -> lookForLambdas block
      ForEach _ _ block -> lookForLambdas block
      ContractName _ block -> lookForLambdas block
      DocGroup _ block -> lookForLambdas block
      ContractGeneral block -> lookForLambdas block
      FinalizeParamCallingDoc _ block _ -> lookForLambdas block

      -- We skip two types of instructions:
      -- 1) Instructions without deeper code block
      -- 2) Unnamed lambdas creation/usage (like CreateLambda1, ExecLambda1, etc)

      -- Instructions without deeper code block
      LiftIndigoState {} -> return ()
      Comment {}         -> return ()
      AssignVar {}       -> return ()
      SetVar {}          -> return ()
      VarModification {} -> return ()
      SetField {}        -> return ()

      TransferTokens {}  -> return ()
      SetDelegate {}     -> return ()
      CreateContract {}  -> return ()
      SelfCalling {}     -> return ()
      ContractCalling {} -> return ()
      Fail {}            -> return ()
      FailOver {}        -> return ()

      -- Nothing to collect in the case of already unnamed lambdas creation/usage
      CreateLambda1 {}   -> return ()
      ExecLambda1 {}     -> return ()

    addLambda :: Lambda1Def -> State (Map Lambda1Def Word) ()
    addLambda lDef = modify $ M.insertWith (+) lDef 1

-- | Contains all the data necessary for the generation of a single-argument
-- lambda. Is compared only on the base of it's 'ldName'.
data Lambda1Def where
  Lambda1Def
    :: (Typeable ret, CreateLambda1CGeneric extra arg ret)
    => { ldRet     :: ret
       , ldName    :: String
       , ldBody    :: Block
       , ldArgVar  :: Var arg
       , ldStack   :: StackVars (arg : extra)
       } -> Lambda1Def

instance Eq Lambda1Def where
  (==) l1 l2 = ldName l1 == ldName l2

instance Ord Lambda1Def where
  (<=) l1 l2 = ldName l1 <= ldName l2