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

-- | Tests for 'Morley.Michelson.Typed.Util'.

module Test.Michelson.Typed.Util
  ( unit_linearizeLeft_sample
  , hprop_linearizeLeft_performance
  , unit_dfsInstr_counter
  , unit_dfsInstr_meta
  , unit_dfsInstr_ctorEffectsApp
  ) where

import Control.Monad.Writer.Strict (Writer, runWriter, tell, writer)
import Data.Default (def)
import Data.Typeable (cast)

import Hedgehog (Property)
import Test.HUnit (Assertion, assertFailure, (@?), (@?=))

import Morley.Michelson.Typed
import Test.Cleveland.Michelson (meanTimeUpperBoundPropNF, sec)

unit_linearizeLeft_sample :: Assertion
unit_linearizeLeft_sample =
  isLeftLinear (linearizeLeft (rightLinear 3)) @? "Failed to linearize"

isLeftLinear :: Instr inp out -> Bool
isLeftLinear (Seq _ (Seq {})) = False
isLeftLinear (Seq a _) = isLeftLinear a
isLeftLinear _ = True

rightLinear :: Word -> Instr '[ 'TUnit, 'TUnit ] '[ 'TUnit, 'TUnit]
rightLinear = go
  where
    go 0 = SWAP
    go n = Seq SWAP (go (n - 1))

hprop_linearizeLeft_performance :: Property
hprop_linearizeLeft_performance =
  meanTimeUpperBoundPropNF (sec 10) (linearizeLeft . rightLinear) 300000

-- Use 'dfsFoldInstr' to count the number of the @UNIT@ instructions.
unit_dfsInstr_counter :: Assertion
unit_dfsInstr_counter = do
  dfsFoldInstr (settings True) step instr @?= Sum 3
  dfsFoldInstr (settings False) step instr @?= Sum 2
  where
    settings goToValues =
      def
      { dsGoToValues = goToValues
      }

    step :: Instr a b -> Sum Word
    step = \case
      UNIT -> Sum 1
      _ -> mempty

    instr = UNIT `Seq` DUP `Seq` DIP UNIT `Seq` DROP `Seq` PUSH v

    v :: Value $ 'TPair 'TUnit $ 'TLambda 'TBool 'TBool
    v = VPair (VUnit, (VLam $ RfNormal $ UNIT `Seq` DROP))

unit_dfsInstr_meta :: Assertion
unit_dfsInstr_meta = do
  let (output, prod) = runWriter $ dfsTraverseInstr def go input
  -- We don't have an untyped version of Meta, so let's pattern match instead of
  -- checking for equality.
  case output of
    Seq CAR (Seq (Meta str' NIL) PAIR) -> matchStr str'
    _ -> assertFailure @() "Failed to match output"
  getProduct prod @?= 120
  where
    str :: String
    str = "preserve me"

    matchStr :: SomeMeta -> Assertion
    matchStr (SomeMeta a) = case cast a of
      Nothing -> assertFailure "Failed to cast string"
      Just a' -> a' @?= str

    nilOp = NIL @('TList 'TOperation)
    input =
      Meta (SomeMeta (1 :: Word)) (Seq
        (Meta (SomeMeta (2 :: Word)) CAR)
          (Meta (SomeMeta (3 :: Word)) (Seq
            (Meta (SomeMeta str) (Meta (SomeMeta (4 :: Word)) nilOp))
              (Meta (SomeMeta (5 :: Word)) PAIR))))

    go :: Instr i o -> Writer (Product Word) (Instr i o)
    go = \case
      Meta (SomeMeta m) i | Just (v :: Word) <- cast m -> i <$ tell (Product v)
      i -> pure i

unit_dfsInstr_ctorEffectsApp :: Assertion
unit_dfsInstr_ctorEffectsApp = do
  -- simple case
  dfsFoldInstr def{ dsCtorEffectsApp = ctorEffectsApp1 } step instr1
    @?= Product (2 * (2 + 3) * 5)

  -- ctor effects are applied to primitives
  dfsFoldInstr def{ dsCtorEffectsApp = ctorEffectsApp2 } step instr1
    @?= Product ((2 + 1) * (2 + 1) * (1 + 2) * 5)

  -- ctor effects are applied to intermediate nodes
  dfsFoldInstr def{ dsCtorEffectsApp = ctorEffectsApp1 } step instr2
    @?= Product (2 * (1 + 2 * (1 + 2)))
  where
    step :: Instr i o -> Product Int
    step = \case
      UNIT -> Product 2
      PUSH _ -> Product 5
      _ -> mempty

    ctorEffectsApp1 = CtorEffectsApp "custom adds for wrappers" \oldInstr ->
      let addition = case oldInstr of
            Nested{} -> 1
            DIP{} -> 3
            _ -> 0
      in writer . second (\(Product x) -> Product (x + addition)) . runWriter

    ctorEffectsApp2 = CtorEffectsApp "custom adds for leaves" \oldInstr ->
      let addition = case oldInstr of
            UNIT{} -> 1
            DROP{} -> 2
            _ -> 0
      in writer . second (\(Product x) -> Product (x + addition)) . runWriter

    instr1 = UNIT `Seq` DUP `Seq` DIP UNIT `Seq` DROP `Seq` PUSH VUnit
    instr2 = UNIT `Seq` Nested (DROP `Seq` UNIT `Seq` Nested (DROP `Seq` UNIT))