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

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

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

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

import Lorentz.Base
import Morley.Michelson.Typed.Haskell
import Morley.Michelson.Typed.Instr
import Morley.Util.Peano

-- | 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

justComment :: Text -> s :-> s
justComment = comment . JustComment

comment :: CommentType -> s :-> s
comment = I . Ext . COMMENT_ITEM

commentAroundFun :: Text -> (i :-> o) -> (i :-> o)
commentAroundFun funName funBody =
  comment (FunctionStarts funName) #
  funBody #
  comment (FunctionEnds funName)

commentAroundStmt :: Text -> (i :-> o) -> (i :-> o)
commentAroundStmt stmtName stmtBody =
  comment (StatementStarts stmtName) #
  stmtBody #
  comment (StatementEnds stmtName)

-- | Test an invariant, fail if it does not hold.
--
-- This won't be included into production contract and is executed only in tests.
testAssert
  :: (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 : _)