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

-- | Tests for optimizer.

module Test.Optimizer
  ( unit_Optimize_DROP_n
  , unit_Optimize_DIP_n
  , unit_Redundant_DIP
  , unit_Adjacent_DIPs
  , unit_Nested_adjacent_DIPs
  , unit_UNPAIR_DROP
  , unit_Specific_PUSH
  , unit_Optimize_PUSH_PACK
  , unit_Optimizer_Tree_Independence
  , unit_Sample_optimize
  , unit_Pair_Unpair
  ) where

import Prelude hiding (EQ)

import Data.Default (def)
import Test.HUnit (Assertion, (@?=))

import Michelson.Interpret.Pack (packValue')
import Michelson.Optimizer
import Michelson.Text
import qualified Michelson.Typed as T
import Michelson.Typed.Instr
import Util.Peano (SingNat(..))

-- Sample stacks of length 0, 1…
type Stack0 = '[ ]
type Stack1 = '[ 'T.TUnit ]
type Stack1Int = '[ 'T.TInt ]
type Stack1Pair = '[ 'T.TPair 'T.TUnit 'T.TUnit ]
type Stack2 = '[ 'T.TUnit, 'T.TUnit ]
type Stack2UnitInt = '[ 'T.TUnit, 'T.TInt ]

unit_Optimizer_Tree_Independence :: Assertion
unit_Optimizer_Tree_Independence = do
  optimize @Stack1Pair @Stack1 (DUP `Seq` (CAR `Seq` (DIP CDR `Seq` DROP))) @?= (DUP `Seq` CAR `Seq` DROP `Seq` CDR)
  optimize @Stack1Pair @Stack1 ((DUP `Seq` CAR) `Seq` (DIP CDR `Seq` DROP)) @?= (DUP `Seq` CAR `Seq` DROP `Seq` CDR)
  optimize @Stack1Pair @Stack1 (((DUP `Seq` CAR) `Seq` DIP CDR) `Seq` DROP) @?= (DUP `Seq` CAR `Seq` DROP `Seq` CDR)

unit_Optimize_DROP_n :: Assertion
unit_Optimize_DROP_n = do
  optimize @Stack0 @Stack0 (DROPN SZ) @?= Nop
  -- Sadly it is not optimized (yet).
  optimize @Stack1 @Stack0 (DROPN (SS SZ)) @?= DROPN (SS SZ)

unit_Optimize_DIP_n :: Assertion
unit_Optimize_DIP_n = do
  optimize @Stack1 @Stack0 (DIPN SZ DROP) @?= DROP
  -- Sadly it is not optimized (yet).
  optimize @Stack1 @Stack2 (DIPN (SS SZ) UNIT) @?= (DIPN (SS SZ) UNIT)

unit_Redundant_DIP :: Assertion
unit_Redundant_DIP = do
  optimize @Stack1Int @Stack1 (DIP UNIT `Seq` DROP) @?= (DROP `Seq` UNIT)
  optimize @Stack1Int @Stack2UnitInt (UNIT `Seq` DIP (DUP `Seq` MUL)) @?= (DUP `Seq` MUL `Seq` UNIT)

unit_Adjacent_DIPs :: Assertion
unit_Adjacent_DIPs = do
  optimize (DIP (PUSH strValue) `Seq` DIP (PUSH strValue)) @?= (DIP (PUSH strValue `Seq` PUSH strValue))
  optimize
    (DIP (PUSH strValue) `Seq`
     DIP UNIT `Seq`
     DIP PAIR)
    @?= (DIP (PUSH strValue `Seq` UNIT `Seq` PAIR))

unit_Nested_adjacent_DIPs :: Assertion
unit_Nested_adjacent_DIPs = do
  optimize
    (IF_NONE (PUSH strValue) (DIP (PUSH strValue) `Seq` DIP (PUSH strValue) `Seq` DIP CONCAT `Seq` CONCAT))
    @?= (IF_NONE (PUSH strValue) (DIP (PUSH strValue `Seq` PUSH strValue `Seq` CONCAT) `Seq` CONCAT))

unit_UNPAIR_DROP :: Assertion
unit_UNPAIR_DROP = do
  -- The left is `unpair # DROP` which is essentially just `cdr`,
  -- but we do not optimize it fully yet (we only remove redundant DIP).
  optimize @Stack1Pair @Stack1 (DUP `Seq` CAR `Seq` DIP CDR `Seq` DROP) @?= (DUP `Seq` CAR `Seq` DROP `Seq` CDR)

unit_Specific_PUSH :: Assertion
unit_Specific_PUSH = do
  optimize (PUSH (T.VMap @'T.TInt @'T.TUnit mempty)) @?= EMPTY_MAP
  optimize (PUSH (T.VSet @'T.TInt mempty) `Seq` NOW) @?= (EMPTY_SET `Seq` NOW)
  optimize (PUSH T.VUnit) @?= UNIT

unit_Optimize_PUSH_PACK :: Assertion
unit_Optimize_PUSH_PACK =
  optimize'
    (PUSH strValue `Seq` PACK `Seq` DUP) @?=
    (PUSH (T.VBytes $ packValue' strValue) `Seq` DUP)
  where
    optimize' =
      optimizeWithConf @Stack0 @'[ 'T.TBytes, 'T.TBytes ]
      (def {ruleset = defaultRulesAndPushPack})

unit_Sample_optimize :: Assertion
unit_Sample_optimize = optimize nonOptimal @?= expectedOptimized

unit_Pair_Unpair :: Assertion
unit_Pair_Unpair =
  optimize
    (PAIR `Seq` UNPAIR `Seq` UNPAIR `Seq` PAIR) @?=
    Nop

str :: MText
str = [mt|aa|]

strValue :: T.Value 'T.TString
strValue = T.VString str

nonOptimal :: T.ContractCode 'T.TString 'T.TString
nonOptimal =
  CAR `Seq`
  -- `PUSH; DROP` is erased
  -- We also arbitrarily group two instructions here to make
  -- structure definitely non-linear.
  (PUSH strValue `Seq` SWAP `Seq` SWAP `Seq` DROP) `Seq`
  -- If we PUSH and then DIP, DIP is not necessary
  PUSH strValue `Seq`
  -- `DUP; DROP` is also erased
  DIP (DUP `Seq` DUP `Seq` DROP) `Seq`
  -- `SWAP; SWAP` is erased, along with surrounding redundant instructions and outer `DIP`
  DIP (PUSH (T.VBool False) `Seq` IF (Nop) (SWAP `Seq` SWAP)) `Seq`
  CONCAT `Seq`
  Nested (SIZE `Seq`
  -- `COMPARE` with 0 is redundant
  (PUSH (T.VNat 0) `Seq` COMPARE) `Seq` EQ `Seq`
  -- Here both bodys of `IF` can be erased and then `IF` can be replaced with `DROP`
  IF (DUP `Seq` DROP) (UNIT `Seq` DROP) `Seq`
  -- `LEFT` followed by `IF_LEFT` can be optimized
  LEFT @('T.TKey) `Seq` IF_LEFT Nop (UNIT `Seq` FAILWITH) `Seq`
  -- SWAP is redundant after DUP
  DUP `Seq` SWAP `Seq` CONCAT `Seq`
  -- `DIP Nop` is thrown away
  DIP (UNIT `Seq` DROP) `Seq`
  -- Finish, nothing to optimize here
  NIL `Seq` PAIR)

-- Auxiliary operator to produce right linear sequence.  We do not use
-- it above because input instruction can have arbitrary structure,
-- but we know that the output is right balanced. In practice we can't
-- check it though, because that's how our 'Eq' is defined.
(#<#) :: T.Instr a b -> T.Instr b c -> T.Instr a c
(#<#) = Seq
infixr 1 #<#

-- Expected output of the optimizer.
expectedOptimized :: T.ContractCode 'T.TString 'T.TString
expectedOptimized =
  CAR #<#
  DUP #<#
  PUSH strValue #<#
  CONCAT #<#
  SIZE #<#
  INT #<# EQ #<#
  DROP #<#
  DUP #<# CONCAT #<#
  NIL #<# PAIR