-- | Tests for optimizer.

module Test.Optimizer
  ( unit_Optimize_DROP_n
  , unit_Optimize_DIP_n
  , unit_Redundant_DIP
  , unit_UNPAIR_DROP
  , unit_Optimize_PUSH_PACK
  , unit_Optimizer_Tree_Independence
  , unit_Sample_optimize
  ) 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 Michelson.Untyped (CT(..))
import Util.Peano (Sing(..))

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

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_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_Optimize_PUSH_PACK :: Assertion
unit_Optimize_PUSH_PACK =
  optimize'
    (PUSH strValue `Seq` PACK `Seq` DUP) @?=
    (PUSH (T.VC $ T.CvBytes $ packValue' strValue) `Seq` DUP)
  where
    optimize' =
      optimizeWithConf @Stack0 @'[ 'T.Tc 'CBytes, 'T.Tc 'CBytes ]
      (def {ruleset = defaultRulesAndPushPack})

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

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

strValue :: T.Value ('T.Tc 'CString)
strValue = T.VC $ T.CvString str

nonOptimal :: T.Contract ('T.Tc 'CString) ('T.Tc 'CString)
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.VC $ T.CvBool False) `Seq` IF (Nop) (SWAP `Seq` SWAP)) `Seq`
  CONCAT `Seq`
  Nested (SIZE `Seq`
  -- `COMPARE` with 0 is redundant
  (PUSH (T.VC $ T.CvNat 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.Contract ('T.Tc 'CString) ('T.Tc 'CString)
expectedOptimized =
  CAR #<#
  DUP #<#
  PUSH strValue #<#
  CONCAT #<#
  SIZE #<#
  INT #<# EQ #<#
  DROP #<#
  DUP #<# CONCAT #<#
  NIL #<# PAIR