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

module Test.Typecheck.Annotation
  ( test_variableAnnotations
  ) where

import Data.Vinyl (Rec(..))
import Test.HUnit ((@?=))
import Test.Tasty (TestTree)
import Test.Tasty.HUnit (testCase)

import Data.Default (def)
import Morley.Michelson.ErrorPos (InstrCallStack(..), srcPos)
import Morley.Michelson.Parser (notes)
import Morley.Michelson.Typed
import Morley.Michelson.Untyped (noAnn)
import Test.Cleveland.Instances ()
import Test.Cleveland.Michelson.Import (importContract)
import Test.Util.Contracts (inContractsDir, (</>))

test_variableAnnotations :: [TestTree]
test_variableAnnotations =
  [ testCase "Annotations are preserved in cadr_annotation.tz" $ do
      let file = inContractsDir ("tezos_examples" </> "attic" </> "cadr_annotation.tz")
          param = "param"
          no_name = "no_name"
          name = "name"
      contract <-
        importContract @('TPair ('TPair 'TUnit 'TString) 'TBool) @'TUnit file
      stripVarAnn (cCode contract) @?=
          loc 2 7 (withNotes [notes|pair (pair %p1 unit (string %no_name)) bool|] (AnnCAR param noAnn))
          `Seq` Nested (
            withNotes [notes|pair unit (string %no_name)|] (AnnCAR noAnn noAnn)
            `Seq` AnnCDR name no_name)
          `Seq` loc 2 40 DROP
          `Seq` loc 2 46 UNIT
          `Seq` loc 2 52 NIL
          `Seq` loc 2 67 (AnnPAIR noAnn noAnn noAnn)

  , testCase "Annotations are preserved in pexec_2.tz" $ do
      let file = inContractsDir ("tezos_examples" </> "opcodes" </> "pexec_2.tz")
      contract <-
        importContract @'TInt @('TList 'TInt) file
      stripVarAnn (cCode contract) @?=
          AnnUNPAIR "p" "s" "" ""
          `Seq` loc 4 6 (LAMBDA (VLam (RfNormal (
            AnnUNPAIR "" "" "" ""
            `Seq` loc 5 24 (DIP (AnnUNPAIR "" "" "" ""))
            `Seq` loc 5 41 ADD
            `Seq` loc 5 47 MUL))))
          `Seq` loc 6 6 SWAP
          `Seq` loc 6 13 APPLY
          `Seq` loc 7 6 (PUSH (VInt 3))
          `Seq` loc 7 19 APPLY
          `Seq` loc 8 6 SWAP
          `Seq` loc 8 13 (MAP (loc 8 19 (DIP (loc 8 25 DUP))
            `Seq` loc 8 33 EXEC))
          `Seq` loc 9 6 (DIP (loc 9 12 DROP))
          `Seq` loc 10 6 NIL
          `Seq` loc 10 21 (AnnPAIR noAnn noAnn noAnn)

  -- Regression test:
  , testCase "UNPAPAIR macro generates UNPAIR instructions without annotations" $ do
      let file = inContractsDir "unpair_macro_simple.tz"
          a1 = "a1"; a2 = "a2"; a3 = "a3"
          u1 = "u1"; u2 = "u2"; u3 = "u3"
          y1 = "y1"; y2 = "y2"
      contract <- importContract @'TUnit @'TUnit file
      stripVarAnn (cCode contract) @?=
          loc 2 7 DROP
          `Seq` loc 3 7  (withNotes (NTUnit u3) (withVN a3 UNIT))
          `Seq` loc 3 21 (withNotes (NTUnit u2) (withVN a2 UNIT))
          `Seq` loc 3 35 (withNotes (NTUnit u1) (withVN a1 UNIT))
          `Seq` withNotes (NTPair noAnn y1 noAnn a1 noAnn (NTUnit u1) (NTPair noAnn y2 noAnn a2 a3 (NTUnit u2) (NTUnit u3))) (Nested (
            withNotes (NTUnit u1) (DIP (
              withNotes (NTPair noAnn y2 noAnn a2 a3 (NTUnit u2) (NTUnit u3)) (AnnPAIR noAnn y2 noAnn)))
            `Seq` withNotes (NTPair noAnn y1 noAnn a1 noAnn (NTUnit u1) (NTPair noAnn y2 noAnn a2 a3 (NTUnit u2) (NTUnit u3))) (withVN "q" (AnnPAIR noAnn y1 noAnn))))
          `Seq`
            withNotes (NTUnit u1)
              ( Nested
                ( InstrWithNotes Proxy
                    (NTUnit u1 :& NTPair noAnn y2 noAnn a2 a3 (NTUnit u2) (NTUnit u3) :& RNil)
                    UNPAIR
                  `Seq`
                  DIP UNPAIR
                )
              )
          `Seq` loc 5 7 (withNotes (NTUnit noAnn) (DIP (
            loc 5 13 (withNotes (NTUnit noAnn) DROP)
            `Seq` (loc 5 19 DROP))))
          `Seq` loc 5 27 NIL
          `Seq` loc 5 42 (withNotes (NTPair noAnn noAnn noAnn noAnn noAnn (NTList noAnn (NTOperation noAnn)) (NTUnit noAnn)) (AnnPAIR noAnn noAnn noAnn))
  ]
  where
    loc :: Word -> Word -> Instr a b -> Instr a b
    loc row col = WithLoc (InstrCallStack [] (srcPos row col))

    withNotes n = InstrWithNotes Proxy (n :& RNil)

    withVN vn = InstrWithVarNotes (one vn)

    stripVarAnn :: Instr a b -> Instr a b
    stripVarAnn = dfsModifyInstr def \case
      InstrWithVarAnns _ i -> i
      i -> i