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

module Indigo.Compilation.Lambda
       ( CompiledLambda (..)
       , Lambda1Def (..)
       , collectLambdas
       ) where

import Prelude
import qualified Data.Map as M

import Indigo.Backend as B
import Indigo.Frontend.Program (IndigoM(..), interpretProgram)
import Indigo.Frontend.Statement
import Indigo.Internal.Object
import Indigo.Internal.SIS
import Indigo.Internal.State hiding ((>>))
import Indigo.Lorentz

data CompiledLambda where
  CompiledLambda
    :: (Typeable arg, Typeable res, Typeable extra)
    => { _clProxyRes :: Proxy res
       , _clName :: String
       , _clVarLam :: Var (B.Lambda1Generic extra arg res)
       } -> CompiledLambda

data Lambda1Def where
  LambdaPure1Def
    :: (Typeable res, CreateLambdaPure1C arg res)
    => { _ldProxy :: Proxy (_stUnit, arg, res)
       , _ldName :: String
       , _ldBody :: Var arg -> IndigoM res
       } -> Lambda1Def

  Lambda1Def
    :: (Typeable res, CreateLambda1C st arg res)
    => { _ldProxy :: Proxy (st, arg, res)
       , _ldName :: String
       , _ldBody :: Var arg -> IndigoM res
       } -> Lambda1Def

  LambdaEff1Def
    :: (Typeable res, CreateLambdaEff1C st arg res)
    => { _ldProxy :: Proxy (st, arg, res)
       , _ldName :: String
       , _ldBody :: Var arg -> IndigoM res
       } -> Lambda1Def

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

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

-- | This is a hack, which prevents using
-- a variable from an outer scope in a body of the lambda.
-- This is not needed when a lambda is defined as top level function,
-- but made just in case, if one wanted to define something like this:
--
-- @
-- f :: Var Storage -> IndigoM ()
-- f storage = do
--     field <- getStorageField
--     let lambda = defNamedLambda1 $ \arg -> ... using field here ...
-- @
-- The idea is that when we pass this variable in
-- a bind it will be propagated in all expressions,
-- including the ones that are in the lambdas.
-- An error will be raised during a variable lookup.
-- This hack will be rewritten later.
leakedVar :: KnownValue a => Var a
leakedVar = Cell $
  error "In a scope of function you are using a variable from an outer scope. Closures are not supported yet."

leakedScopeVariableAllocator :: KnownValue a => MetaData _inp -> (Var a, MetaData (a & _inp))
leakedScopeVariableAllocator (MetaData stk cnt) =
  let v = leakedVar
  in (v, MetaData (Ref cnt :& stk) (cnt + 1))

allocateVarsLeaked :: forall a . ReturnableValue a => RetVars a
allocateVarsLeaked = fst (allocateVars @a leakedScopeVariableAllocator emptyMetadata)

allocateVarsLeakedM :: forall a m . (Monad m, ReturnableValue a) => m a -> m (RetVars a)
allocateVarsLeakedM ma = allocateVarsLeaked @a <$ ma

-- | Collect all used lambdas in a computation
-- (which might be either a contract body or another function body),
-- which are called at least twice.
-- Only outer functions will be gathered, for instance,
-- if we call lambda func1 from func0, only func0 will be taken.
collectLambdas :: forall a . IndigoM a -> Set Lambda1Def
collectLambdas indigoM =
  M.keysSet $ M.filter (> 1) $ executingState mempty (lookForLambdas indigoM)
  where
    lookForLambdas :: IndigoM x -> State (Map Lambda1Def Word) x
    lookForLambdas (IndigoM program) = interpretProgram inspectLambda program

    inspectLambda :: StatementF IndigoM x -> State (Map Lambda1Def Word) x
    inspectLambda (LambdaPure1Call name (body :: (Var arg -> IndigoM res)) _) =
      allocateVarsLeaked @res <$ modify (addLambda (LambdaPure1Def (Proxy @((), arg, res)) name body))

    inspectLambda (Lambda1Call (_ :: Proxy st) name (body :: (Var arg -> IndigoM res)) _) =
      allocateVarsLeaked @res <$ modify (addLambda (Lambda1Def (Proxy @(st, arg, res)) name body))

    inspectLambda (LambdaEff1Call (_ :: Proxy st) name (body :: (Var arg -> IndigoM res)) _) =
      allocateVarsLeaked @res <$ modify (addLambda (LambdaEff1Def (Proxy @(st, arg, res)) name body))

    inspectLambda (Scope cd) = allocateVarsLeakedM $ lookForLambdas cd
    inspectLambda (If _ tb fb) = allocateVarsLeakedM $ lookForLambdas tb >> lookForLambdas fb
    inspectLambda (IfSome _ tb fb) = allocateVarsLeakedM $ lookForLambdas (tb leakedVar) >> lookForLambdas fb
    inspectLambda (IfRight _ rb lb) = allocateVarsLeakedM $ lookForLambdas (rb leakedVar) >> lookForLambdas (lb leakedVar)
    inspectLambda (IfCons _ tb fb) = allocateVarsLeakedM $ lookForLambdas (tb leakedVar leakedVar) >> lookForLambdas fb
    inspectLambda (Case _ clauses) = rmapClauses clauses
    inspectLambda (EntryCase _ _ clauses) = rmapClauses clauses
    inspectLambda (EntryCaseSimple _ clauses) = rmapClauses clauses
    inspectLambda (While _ body) = lookForLambdas body
    inspectLambda (WhileLeft _ body) = lookForLambdas (body leakedVar) >> pure leakedVar
    inspectLambda (ForEach _ body) = lookForLambdas $ body leakedVar
    inspectLambda (ContractName _ contr) = lookForLambdas contr
    inspectLambda (DocGroup _ ii) = lookForLambdas ii
    inspectLambda (ContractGeneral contr) = lookForLambdas contr
    inspectLambda (FinalizeParamCallingDoc entrypoint _) = lookForLambdas (entrypoint leakedVar)

    -- Not recursive simple statements. They are terminal ones
    inspectLambda (LiftIndigoState cd) = pure $ runSIS cd emptyMetadata gcOut
    inspectLambda (NewVar _) = pure leakedVar
    inspectLambda (SetVar _ _) = pure ()
    inspectLambda (SetField {}) = pure ()
    inspectLambda (VarModification {}) = pure ()
    inspectLambda (TransferTokens {}) = pure ()
    inspectLambda (SetDelegate _) = pure ()
    inspectLambda (CreateContract{}) = pure leakedVar
    inspectLambda (ContractCalling{}) = pure leakedVar

    inspectLambda (FailWith ex) = pure $ gcOut $ runIndigoState (B.failWith ex) emptyMetadata
    inspectLambda (Assert _ _) = pure ()
    inspectLambda (FailCustom tag ex) = pure $ gcOut $ runIndigoState (B.failCustom tag ex) emptyMetadata

    rmapClauses:: forall ret cs . ReturnableValue ret
       => Rec (IndigoMCaseClauseL IndigoM ret) cs
       -> State (Map Lambda1Def Word) (RetVars ret)
    rmapClauses RNil = pure (allocateVarsLeaked @ret)
    rmapClauses ((OneFieldIndigoMCaseClauseL _ clause) :& rs) =
      lookForLambdas (clause leakedVar) >> rmapClauses rs

    addLambda :: Lambda1Def -> Map Lambda1Def Word -> Map Lambda1Def Word
    addLambda  =
      M.alter (\case
        Nothing -> Just 1
        Just x -> Just (x + 1)
      )