-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

-- TODO [#712]: Remove this next major release
{-# OPTIONS_GHC -Wno-deprecations #-}

-- | Tests for optimizer.

module Test.Optimizer
  ( unit_Optimize_Nop
  , unit_Optimize_DROP_n
  , unit_Optimize_DIP_n
  , unit_Redundant_DIP
  , unit_Adjacent_DIPs
  , unit_Adjacent_DROPs
  , unit_Nested_adjacent_DIPs
  , unit_UNPAIR_DROP
  , unit_Specific_PUSH
  , unit_Optimize_PUSH_PACK
  , unit_Optimizer_Tree_Independence
  , unit_Sample_optimize
  , unit_Pair_Unpair
  , unit_Optimize_PAIR_CADR
  , unit_Optimize_following_DROPs
  , unit_Optimize_following_DROPs_double
  , unit_Optimize_multiple_following_DROPs
  , unit_Optimize_branch_shortcuts_with_PUSH

  , unit_golden
  ) where

import Prelude hiding (EQ)

import Data.Default (def)
import Data.Text.IO.Utf8 qualified as Utf8 (readFile)
import Data.Text.Lazy (strip)
import Test.HUnit (Assertion, (@?=))

import Morley.Michelson.Interpret.Pack (packValue')
import Morley.Michelson.Optimizer
import Morley.Michelson.Printer (printSomeContract)
import Morley.Michelson.Runtime.Import (importUsing, readSomeContractExt)
import Morley.Michelson.Text
import Morley.Michelson.TypeCheck (mapSomeContract)
import Morley.Michelson.Typed qualified as T
import Morley.Michelson.Typed.Instr
import Morley.Michelson.Untyped.Annotation (noAnn)
import Morley.Util.PeanoNatural (PeanoNatural(..))
import Test.Cleveland.Instances ()

import Test.Util.Contracts
import Test.Util.HUnit


-- 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))) @?= CDR
  optimize @Stack1Pair @Stack1 ((DUP `Seq` CAR) `Seq` (DIP CDR `Seq` DROP)) @?= CDR
  optimize @Stack1Pair @Stack1 (((DUP `Seq` CAR) `Seq` DIP CDR) `Seq` DROP) @?= CDR
  optimize (
    PUSH (T.VBool True) `Seq`
    IF (PUSH (T.VOption $ Just strValue)) (PUSH (T.VOption Nothing)) `Seq`
    IF_NONE (PUSH (T.VBool False)) (DROP `Seq` PUSH (T.VBool True))
    `Seq` DROP) @?= Nop
  optimize (
    PUSH (T.VBool True) :#
    IF (PUSH (T.VOption $ Just strValue)) (PUSH (T.VOption Nothing)) :#
    IF_NONE (PUSH (T.VBool False)) (DROP :# PUSH (T.VBool True)) :#
    DROP) @?= Nop

unit_Optimize_Nop :: Assertion
unit_Optimize_Nop = do
  -- NOTE: at the moment of writing this test does not actually test anything
  -- because our 'Eq' on typed 'Instr' is not strict, it compares via casting
  -- the instruction to untyped representation. To be resolved in #108.
  optimize @Stack1Int @Stack1Int (Nop `Seq` DUP `Seq` Nop `Seq` ADD `Seq` Nop) @?= DUP `Seq` ADD

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

unit_Optimize_DIP_n :: Assertion
unit_Optimize_DIP_n = do
  optimize @Stack1 @Stack0 (DIPN Zero DROP) @?= DROP
  -- Sadly it is not optimized (yet).
  optimize @Stack1 @Stack2 (DIPN (Succ Zero) UNIT) @?= (DIPN (Succ Zero) 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_Adjacent_DROPs :: Assertion
unit_Adjacent_DROPs = do
  optimize (DROP `Seq` DROP) @?=
    DROPN Two
  optimize (UNIT `Seq` DROP `Seq` DROP `Seq` DROP `Seq` NOW) @?=
    (DROPN Two `Seq` NOW)
  -- TODO #299: We don't optimize more than 2 DROPs in a row yet :(
  -- optimize (DROP `Seq` DROP `Seq` DROP `Seq` UNIT) @?=
  --   (DROPN (Succ Two) `Seq` UNIT)

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
  optimize @Stack1Pair @Stack2 (DUP :# CAR :# DIP CDR) @?= UNPAIR
  optimize @Stack1Pair @Stack1 (((DUP :# CAR) :# DIP CDR) :# DROP) @?= CDR
  optimize @Stack1Pair @Stack1 (DUP :# CAR :# DIP CDR :# DROP) @?= 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 {ocRuleset = 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

unit_Optimize_PAIR_CADR :: Assertion
unit_Optimize_PAIR_CADR = do
  optimize (PAIR `Seq` CDR) @?= DROP
  optimize (PAIR `Seq` CAR) @?= (SWAP `Seq` DROP)

unit_Optimize_following_DROPs :: Assertion
unit_Optimize_following_DROPs = do
  optimize (PUSH strValue `Seq` SIZE @'T.TString `Seq` INT `Seq` DROP) @?= Nop
  optimize (PUSH intValue `Seq` NEG `Seq` ABS `Seq` DROP) @?= Nop
  optimize (PUSH strValue `Seq` PACK `Seq` UNPACK @'T.TString `Seq` DROP) @?= Nop
  optimize (PUSH (T.VPair (strValue, strValue)) `Seq` GETN (Succ Zero) `Seq` DROP) @?= Nop
  optimize (SELF_ADDRESS `Seq` (CONTRACT (T.NTUnit noAnn) T.DefEpName) `Seq` DROP) @?= Nop
  optimize (SELF (T.sepcPrimitive @'T.TUnit) `Seq` ADDRESS `Seq` DROP) @?= Nop
  optimize (PUSH strValue `Seq` PACK `Seq` SHA256 `Seq` DROP) @?= Nop
  optimize (PUSH intValue `Seq` CAST @'T.TInt `Seq` (AnnLEFT @'T.TString noAnn noAnn noAnn) `Seq` DROP) @?= Nop

unit_Optimize_following_DROPs_double :: Assertion
unit_Optimize_following_DROPs_double = do
  optimize (PUSH intValue `Seq` PUSH intValue `Seq` ADD `Seq` DROP) @?= Nop
  optimize (EMPTY_MAP @'T.TInt @'T.TInt `Seq` PUSH intValue `Seq` GET `Seq` DROP) @?= Nop
  optimize (UNIT `Seq` PUSH (T.VNat 1) `Seq` UNIT `Seq` TICKET `Seq` DROP `Seq` DROP) @?= Nop
  optimize (
    UNIT `Seq` DUP `Seq` (AnnPAIR noAnn noAnn noAnn) `Seq`
    BALANCE `Seq` AMOUNT `Seq` COMPARE `Seq`
    DROP `Seq` DROP) @?= Nop

unit_Optimize_multiple_following_DROPs :: Assertion
unit_Optimize_multiple_following_DROPs = do
  optimize (PUSH strValue `Seq` PUSH strValue `Seq` DROP `Seq` DROP) @?= Nop
  optimize (NOW `Seq` PUSH strValue `Seq` DUP `Seq` DROP `Seq` DROP `Seq` DROP) @?= Nop
  optimize (NIL @'T.TInt `Seq` UNIT `Seq` NONE @'T.TUnit `Seq` PUSH strValue `Seq` DROP `Seq` DROP `Seq` DROP `Seq` DROP) @?= Nop

unit_Optimize_branch_shortcuts_with_PUSH :: Assertion
unit_Optimize_branch_shortcuts_with_PUSH = do
  -- TODO #300: Optimization stops after one step.
  -- As a result there is (PUSH _ `Seq` PUSH _ `Seq` DROP `Seq` DROP)
  -- on stack instead of Nop.
  -- optimize (PUSH strValue `Seq` PUSH (T.VList $ replicate 5 intValue) `Seq` IF_CONS (DROP `Seq` DROP) Nop `Seq` DROP) @?= Nop
  optimize (PUSH (T.VOr @'T.TString @'T.TUnit $ Left strValue) `Seq` IF_LEFT DROP DROP) @?= Nop
  optimize (PUSH strValue `Seq` PUSH (T.VOption @'T.TUnit Nothing) `Seq` IF_NONE Nop DROP `Seq` DROP) @?= Nop

str :: MText
str = "aa"

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

intValue :: T.Value 'T.TInt
intValue = T.VInt 1

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
  (AnnLEFT @('T.TKey) noAnn noAnn noAnn) `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 #<# CONCAT #<#
  NIL #<# PAIR

-- [TODO] After introduction of WithLoc the test below started to fail.
-- Probably the cause is the same as in #300. As a temporary measure we drop
-- all 'WithLoc' before making optimizations.
dropWithLoc :: T.Instr i o -> T.Instr i o
dropWithLoc = T.dfsModifyInstr def $ \case
  T.WithLoc _ i -> i
  i -> i

unit_golden :: Assertion
unit_golden = do
  contracts <- getContractsWithReferences ".mtz" (inContractsDir "optimizer") "opt"
  mapM_ optimizerTest contracts
  where
    optimizerTest :: (FilePath, FilePath) -> Assertion
    optimizerTest (srcPath, optPath) = do
      someContract <- importUsing readSomeContractExt srcPath
      let
        optimized :: T.SomeContract
        optimized = mapSomeContract (optimize . dropWithLoc) someContract

      expectedOptimizedTxt <- strip . fromStrict <$> Utf8.readFile optPath
      assertEqualBuild
        ("Optimizing " <> srcPath <> " does not match the expected format")
        expectedOptimizedTxt
        (printSomeContract False optimized)