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

module Test.Code.Lambda
  ( SmallMatrix (..)
  , dummyContract
  , sumLambdaCalledOnce
  , sumLambdaCalledTwice
  , lambdasSideEffects
  , lambdaInLambda1
  , lambdaInLambda2
  , lambdaOuterVarClosure
  ) where

import Indigo
import qualified Lorentz as L
import Lorentz.TypeAnns (HasTypeAnn)
import Test.Util

sumLambda :: pr :~> (Integer, Integer) => pr -> IndigoM (Var Integer)
sumLambda = defNamedPureLambda1 "add" $  \pr -> pure (fst pr + snd pr)

-- | Pure lambda called once.
-- In this case code has to be just inlined.
sumLambdaCalledOnce :: ContractCode [Integer] Integer
sumLambdaCalledOnce = compileIndigoContract $ \lst -> do
  s <- new$ 0 int
  forEach lst $
    setVar s <=< sumLambda . pair s
  storageVar =: s

-- | Pure lambda called twice.
-- In this case lambda has to be pushed on the stack and then called using @exec@.
sumLambdaCalledTwice :: ContractCode [Integer] Integer
sumLambdaCalledTwice = compileIndigoContract $ \lst -> do
  s <- (new$ 0 int) >>= sumLambda . pair (0 int)
  forEach lst $
    setVar s <=< sumLambda . pair s
  storageVar =: s

lambdaCreateContract ::
  ( ex :~> (Maybe KeyHash, Integer)
  , HasSideEffects, HasStorage Integer
  ) => ex -> IndigoM ()
lambdaCreateContract = defNamedEffLambda1 @Integer "create storage" $ \paramSt -> do
  m <- new$ 0 mutez
  _addr <- createLorentzContract dummyContract (fst paramSt) m (snd paramSt)
  storageVar @Integer += 1 int

dummyContract
  :: L.Contract Integer Integer
dummyContract = noOptimizationContract $ compileIndigoContract $ \param -> do
  a <- new$ 7 int + param
  when (param < a) do
    _c <- new$ storageVar @Integer
    return ()
  _c <- new$ param < storageVar @Integer
  return ()

-- | Lambda that creates new contract is called twice.
-- Test that original operations list is updated in lambda.
lambdasSideEffects :: ContractCode (Maybe KeyHash) Integer
lambdasSideEffects = compileIndigoContract $ \param -> do
  lambdaCreateContract (pair param $ 0 int)
  lambdaCreateContract (pair param $ 1 int)

newtype SmallMatrix = SmallMatrix [[Integer]]
  deriving stock (Generic, Show)
  deriving anyclass (IsoValue)
  deriving newtype (IterOpHs, HasTypeAnn)

sumInRowTwice :: (HasStorage Integer, lst :~> [Integer]) => lst -> IndigoM ()
sumInRowTwice = defNamedLambda1 @Integer "sum of the list" $ \lst -> do
  let storage = storageVar @Integer
  setVar storage =<< sumLambda (pair (0 int) storage) -- to call 'sumLambda' twice
  forEach lst $
    setVar storage <=< sumLambda . pair storage

-- | Define lambda inside another lambda and call it there.
lambdaInLambda1 :: ContractCode SmallMatrix Integer
lambdaInLambda1 = compileIndigoContract $ \matrix -> do
  sumInRowTwice [] -- to call 'sumInRow' twice
  forEach matrix sumInRowTwice

sumInRowOnce :: (HasStorage Integer, lst :~> [Integer]) => lst -> IndigoM ()
sumInRowOnce = defNamedLambda1 @Integer "sum of the list" $ \lst -> do
  let storage = storageVar @Integer
  forEach lst $
    setVar storage <=< sumLambda . pair storage

-- | Define lambda in another lambda but call inner one
-- only once but outer twice, consequently, inner one
-- has to be inlined but outer one is executed via @exex@.
lambdaInLambda2 :: ContractCode SmallMatrix Integer
lambdaInLambda2 = compileIndigoContract $ \matrix -> do
  sumInRowOnce [] -- to call 'sumInRowOnce' twice
  let storage = storageVar @Integer
  setVar storage =<< sumLambda (pair (0 int) storage) -- to pretend we call 'sumLambda' twice
  forEach matrix sumInRowOnce

-- | Use a variable from outer scope to check
-- that an error is raised.
-- TODO attach scopes to variables and prevent
-- variables from leaking more severely.
-- Current approach doesn't throw a proper error in the following cases:
-- * a contract param is in closure of lambda
-- * a pure lambda uses @storageVar@ or @opsVar@
lambdaOuterVarClosure :: ContractCode Integer Integer
lambdaOuterVarClosure  = compileIndigoContract $ \param -> do
  plusTwo <- new$ param + 2 int
  let lam :: pr :~> Integer => pr -> IndigoM (Var Integer)
      lam = defNamedPureLambda1 "lambda" $ \innerPar -> pure (plusTwo + innerPar)
  _v1 <- lam $ 3 int
  _v2 <- lam $ 4 int
  return ()