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

{-# OPTIONS_GHC -Wno-unticked-promoted-constructors #-}

-- | Lorentz wrappers over instructions from Morley extension.
module Lorentz.Ext
  ( stackRef
  , printComment
  , testAssert
  , stackType
  ) where

import Data.Singletons (SingI)
import GHC.TypeNats (Nat)

import Lorentz.Base
import Michelson.Typed.Haskell
import Michelson.Typed.Instr
import Util.Peano hiding (Nat)

-- | Include a value at given position on stack into comment produced
-- by 'printComment'.
--
-- >>> stackRef @0
-- <includes the top of the stack>
stackRef
  :: forall (gn :: Nat) st n.
      (n ~ ToPeano gn, SingI n, RequireLongerThan st n)
  => PrintComment st
stackRef = PrintComment . one . Right $ mkStackRef @gn

-- | Print a comment. It will be visible in tests.
--
-- >>> printComment "Hello world!"
-- >>> printComment $ "On top of the stack I see " <> stackRef @0
printComment :: PrintComment (ToTs s) -> s :-> s
printComment = I . Ext . PRINT

-- | Test an invariant, fail if it does not hold.
--
-- This won't be included into production contract and is executed only in tests.
testAssert
  :: (SingI (ToTs out), HasCallStack)
  => Text -> PrintComment (ToTs inp) -> inp :-> Bool : out -> inp :-> inp
testAssert msg comment = \case
  I instr -> I . Ext . TEST_ASSERT $ TestAssert msg comment instr
  FI _ -> error "test assert branch always fails"

-- | Fix the current type of the stack to be given one.
--
-- >>> stackType @'[Natural]
-- >>> stackType @(Integer : Natural : s)
-- >>> stackType @'["balance" :! Integer, "toSpend" :! Integer, BigMap Address Integer]
--
-- Note that you can omit arbitrary parts of the type.
--
-- >>> stackType @'["balance" :! Integer, "toSpend" :! _, BigMap _ _]
stackType :: forall s. s :-> s
stackType = I Nop

_sample1 :: s ~ (a : s') => s :-> s
_sample1 = printComment $ "Head is " <> stackRef @0

_sample2 :: Integer : Natural : s :-> Integer : Natural : s
_sample2 = stackType @(Integer : _)