-- | Module, providing functions for conversion from
-- instruction and value representation from @Michelson.Type@ module
-- to strictly-typed GADT-based representation from @Michelson.Value@ module.
--
-- This conversion is labeled as type check because that's what we are obliged
-- to do on our way.
--
-- Type check algorithm relies on the property of Michelson language that each
-- instruction on a given input stack type produces a definite output stack
-- type.
-- Michelson contract defines concrete types for storage and parameter, from
-- which input stack type is deduced. Then this type is being combined with
-- each subsequent instruction, producing next stack type after each
-- application.
--
-- Function @typeCheck@ takes list of instructions and returns value of type
-- @Instr inp out@ along with @HST inp@ and @HST out@ all wrapped into
-- @SomeInstr@ data type. This wrapping is done to satsify Haskell type
-- system (which has no support for dependent types).
-- Functions @typeCheckInstr@, @typeCheckValue@ behave similarly.
--
-- When a recursive call is made within @typeCheck@, @typeCheckInstr@ or
-- @typeCheckValue@, result of a call is unwrapped from @SomeInstr@ and type
-- information from @HST inp@ and @HST out@ is being used to assert that
-- recursive call returned instruction of expected type
-- (error is thrown otherwise).
module Michelson.TypeCheck.Instr
    ( typeCheckContract
    , typeCheckValue
    , typeVerifyValue
    , typeCheckList
    , typeCheckStorageOrParameter
    , StorageOrParameter (..)
    ) where

import Prelude hiding (EQ, GT, LT)

import Control.Monad.Except (MonadError, liftEither, throwError)
import Data.Constraint (Dict(..))
import Data.Default (def)
import Data.Generics (everything, mkQ)
import Data.Singletons (SingI(sing), demote)
import Data.Typeable ((:~:)(..), gcast)

import Michelson.ErrorPos
import Michelson.TypeCheck.Error
import Michelson.TypeCheck.Ext
import Michelson.TypeCheck.Helpers
import Michelson.TypeCheck.TypeCheck
import Michelson.TypeCheck.Types
import Michelson.TypeCheck.Value

import Michelson.Typed
import Util.Peano

import qualified Michelson.Untyped as U
import Michelson.Untyped.Annotation (VarAnn)

typeCheckContract
  :: TcOriginatedContracts
  -> U.Contract
  -> Either TCError SomeContract
typeCheckContract cs c = runTypeCheck (U.para c) cs $ typeCheckContractImpl c

typeCheckContractImpl
  :: U.Contract
  -> TypeCheck SomeContract
typeCheckContractImpl (U.Contract mParam mStorage pCode) = do
  code <- maybe (throwError $ TCContractError "no instructions in contract code" Nothing)
                pure (nonEmpty pCode)
  withUType mParam $ \(paramS :: Sing param, paramNote :: Notes param) ->
    withUType mStorage $ \(storageS :: Sing st, storageNote :: Notes st) -> do
      (Dict, Dict) <-
        note (hasTypeError "parameter" paramS) $
          (,) <$> opAbsense paramS
              <*> nestedBigMapsAbsense paramS
      (Dict, Dict, Dict) <-
        note (hasTypeError "storage" storageS) $
          (,,) <$> opAbsense storageS
               <*> nestedBigMapsAbsense storageS
               <*> contractTypeAbsense storageS
      let inpNote = NTPair def def def paramNote storageNote
      let inp = (STPair paramS storageS, inpNote, def) ::& SNil
      inp' :/ instrOut <- typeCheckNE code inp
      let (paramNotesRaw, fcStoreNotes) = case inp' of
            (_, NTPair _ _ _ cpNotes stNotes, _) ::& SNil -> (cpNotes, stNotes)
      fcParamNotesSafe <-
        liftEither $
        mkParamNotes paramNotesRaw `onLeft`
        (TCContractError "invalid parameter declaration: " . Just . IllegalParamDecl)
      case instrOut of
        instr ::: out -> liftEither $ do
          case eqHST1 @(ContractOut1 st) out of
            Right Refl -> do
              let (_, outN, _) ::& SNil = out
              _ <- converge outN (NTPair def def def starNotes storageNote)
                      `onLeft`
                  ((TCContractError "contract output type violates convention:") . Just . AnnError)
              pure $ SomeContract FullContract
                { fcCode = instr
                , fcParamNotesSafe, fcStoreNotes
                }
            Left err -> Left $ TCContractError "contract output type violates convention:" $ Just err
        AnyOutInstr instr ->
          pure $ SomeContract FullContract
            { fcCode = instr
            , fcParamNotesSafe, fcStoreNotes
            }
  where
    hasTypeError name (_ :: proxy t) =
      TCContractError ("contract " <> name <> " type error") $
      Just $ UnsupportedTypes [demote @t]

-- | Like 'typeCheck', but for non-empty lists.
typeCheckNE
  :: (Typeable inp)
  => NonEmpty U.ExpandedOp
  -> HST inp
  -> TypeCheck (SomeInstr inp)
typeCheckNE (x :| xs) = usingReaderT def . typeCheckImpl typeCheckInstr (x : xs)

-- | Function @typeCheckList@ converts list of Michelson instructions
-- given in representation from @Michelson.Type@ module to representation
-- in strictly typed GADT.
--
-- Types are checked along the way which is neccessary to construct a
-- strictly typed value.
--
-- As a second argument, @typeCheckList@ accepts input stack type representation.
typeCheckList
  :: (Typeable inp)
  => [U.ExpandedOp]
  -> HST inp
  -> TypeCheck (SomeInstr inp)
typeCheckList = usingReaderT def ... typeCheckImpl typeCheckInstr

-- | Function @typeCheckValue@ converts a single Michelson value
-- given in representation from @Michelson.Untyped@ module hierarchy to
-- representation in strictly typed GADT.
--
-- As a second argument, @typeCheckValue@ accepts expected type of value.
--
-- Type checking algorithm pattern-matches on parse value representation,
-- expected type @t@ and constructs @Val t@ value.
--
-- If there was no match on a given pair of value and expected type,
-- that is interpreted as input of wrong type and type check finishes with
-- error.
typeCheckValue
  :: U.Value
  -> (Sing t, Notes t)
  -> TypeCheckInstr SomeNotedValue
typeCheckValue = typeCheckValImpl typeCheckInstr

-- | Like 'typeCheckValue', but returns value of a desired type.
typeVerifyValue
  :: forall t.
      (Typeable t, SingI t)
  => U.Value -> TypeCheckInstr (Value t)
typeVerifyValue uval =
  typeCheckValue uval (sing @t, starNotes) <&> \case
    val :::: _ -> gcast val ?: error "Typechecker produced value of wrong type"

typeCheckStorageOrParameter
  :: StorageOrParameter
  -> U.Value -> TcOriginatedContracts -> U.Contract -> Either TCError SomeValue
typeCheckStorageOrParameter storOrParam valueU originatedContracts U.Contract{..} =
  runTypeCheck para originatedContracts $ usingReaderT (def :: InstrCallStack) $
  typeCheckStorageImpl
  where
    typeCheckStorageImpl =
      let paraOrStor = bool para stor (storOrParam == Storage) in
      withSomeSingT (fromUType paraOrStor) $ \(_ :: Sing t) -> do
          SomeValue <$> typeVerifyValue @t valueU

-- Helper data type we use to typecheck DROPN.
data TCDropHelper inp where
  TCDropHelper ::
    forall (n :: Peano) inp out.
    (Typeable out, SingI n, KnownPeano n, LongerOrSameLength inp n, Drop n inp ~ out) =>
    Sing n -> HST out -> TCDropHelper inp

-- Helper data type we use to typecheck DIG.
data TCDigHelper inp where
  TCDigHelper ::
    forall (n :: Peano) inp out a.
    (Typeable out, ConstraintDIG n inp out a) =>
    Sing n -> HST out -> TCDigHelper inp

-- Helper data type we use to typecheck DUG.
data TCDugHelper inp where
  TCDugHelper ::
    forall (n :: Peano) inp out a.
    (Typeable out, ConstraintDUG n inp out a) =>
    Sing n -> HST out -> TCDugHelper inp

-- | Function @typeCheckInstr@ converts a single Michelson instruction
-- given in representation from @Michelson.Type@ module to representation
-- in strictly typed GADT.
--
-- As a second argument, @typeCheckInstr@ accepts input stack type representation.
--
-- Type checking algorithm pattern-matches on given instruction, input stack
-- type and constructs strictly typed GADT value, checking necessary type
-- equalities when neccessary.
--
-- If there was no match on a given pair of instruction and input stack,
-- that is interpreted as input of wrong type and type check finishes with
-- error.
typeCheckInstr :: TcInstrHandler
typeCheckInstr (U.EXT ext) si =
  typeCheckExt typeCheckList ext si

typeCheckInstr U.DROP i@(_ ::& rs) = pure (i :/ DROP ::: rs)

typeCheckInstr uInstr@(U.DROPN nTotal) inputHST =
  go nTotal inputHST <&> \case
    TCDropHelper s out -> inputHST :/ DROPN s ::: out
  where
    go :: forall inp. Typeable inp
      => Word
      -> HST inp
      -> TypeCheckInstr (TCDropHelper inp)
    go n i = case (n, i) of
      (0, _) -> pure (TCDropHelper SZ i)

      (_, SNil) -> onTypeCheckInstrErr uInstr
        (SomeHST SNil) "" (Left (NotEnoughItemsOnStack nTotal NotEnoughDrop))

      (_, (_ ::& iTail)) -> do
        go (n - 1) iTail <&> \case TCDropHelper s out -> TCDropHelper (SS s) out

typeCheckInstr (U.DUP _vn) i@(a ::& rs) =
  pure (i :/ DUP ::: (a ::& a::& rs))

typeCheckInstr U.SWAP (i@(a ::& b ::& rs)) =
  pure (i :/ SWAP ::: (b ::& a ::& rs))

typeCheckInstr uInstr@(U.DIG nTotal) inputHST =
  go nTotal inputHST <&> \case
    TCDigHelper s out -> inputHST :/ DIG s ::: out
  where
    go :: forall inp. Typeable inp
      => Word
      -> HST inp
      -> TypeCheckInstr (TCDigHelper inp)
    go n i = case (n, i) of
      -- Even 'DIG 0' is invalid on empty stack (so it is not strictly `Nop`).
      (_, SNil) -> onTypeCheckInstrErr uInstr
        (SomeHST SNil) "" (Left (NotEnoughItemsOnStack nTotal NotEnoughDig))

      (0, (_ ::& _)) -> pure (TCDigHelper SZ i)

      (_, (b ::& iTail)) ->
        go (n - 1) iTail <&> \case
        TCDigHelper s (a ::& resTail) -> TCDigHelper (SS s) (a ::& b ::& resTail)

typeCheckInstr uInstr@(U.DUG nTotal) inputHST =
  go nTotal inputHST <&> \case
    TCDugHelper s out -> inputHST :/ DUG s ::: out
  where
    go :: forall inp. Typeable inp
      => Word
      -> HST inp
      -> TypeCheckInstr (TCDugHelper inp)
    go n i = case (n, i) of
      (0, (_ ::& _)) -> pure (TCDugHelper SZ i)

      (_, (a ::& b ::& iTail)) ->
        go (n - 1) (a ::& iTail) <&> \case
        TCDugHelper s resTail -> TCDugHelper (SS s) (b ::& resTail)

      -- Two cases:
      -- 1. Input stack is empty.
      -- 2. n > 0 and input stack has exactly 1 item.
      _ -> onTypeCheckInstrErr uInstr
        (SomeHST i) "" (Left (NotEnoughItemsOnStack nTotal NotEnoughDug))

typeCheckInstr instr@(U.PUSH vn mt mval) i =
  withUType mt $ \(t' :: Sing t', nt' :: Notes t') -> do
    val :::: (t :: Sing t, nt) <- typeCheckValue mval (t', nt')
    proofOp <-
      maybe (onTypeCheckInstrErr instr (SomeHST i)
             "Operations in constant are not allowed"
             $ Left $ UnsupportedTypes [demote @t])
      pure (opAbsense t)
    proofBigMap <-
      maybe (onTypeCheckInstrErr instr (SomeHST i)
             "BigMaps in constant are not allowed"
             $ Left $ UnsupportedTypes [demote @t])
      pure (bigMapAbsense t)
    proofContract <-
      maybe (onTypeCheckInstrErr instr (SomeHST i)
             "Contract type in constant is not allowed"
             $ Left $ UnsupportedTypes [demote @t])
      pure (contractTypeAbsense t)
    case (proofOp, proofBigMap, proofContract) of
      (Dict, Dict, Dict) -> pure $ i :/ PUSH val ::: ((t, nt, vn) ::& i)

typeCheckInstr (U.SOME tn vn) i@((at, an, _) ::& rs) = do
  pure (i :/ SOME ::: ((STOption at, NTOption tn an, vn) ::& rs))

typeCheckInstr (U.NONE tn vn elMt) i =
  withUType elMt $ \(elSing :: Sing t, elNotes :: Notes t) ->
    pure $ i :/ NONE ::: ((STOption elSing, NTOption tn elNotes, vn) ::& i)

typeCheckInstr (U.UNIT tn vn) i = do
  pure $ i :/ UNIT ::: ((STUnit, NTUnit tn, vn) ::& i)

typeCheckInstr (U.IF_NONE mp mq) i@((STOption a, ons, ovn) ::& rs) = do
  let (an, avn) = deriveNsOption ons ovn
  genericIf IF_NONE U.IF_NONE mp mq rs ((a, an, avn) ::& rs) i

typeCheckInstr (U.PAIR tn vn pfn qfn) i@((a, an, avn) ::&
                                             (b, bn, bvn) ::& rs) = do
  let (vn', pfn', qfn') = deriveSpecialFNs pfn qfn avn bvn
      ns = NTPair tn pfn' qfn' an bn
  pure (i :/ PAIR ::: ((STPair a b, ns, vn `orAnn` vn') ::& rs))

typeCheckInstr instr@(U.CAR vn fn)
            (i@(( STPair a b
                       , NTPair pairTN pfn qfn pns qns
                       , pairVN ) ::& rs)) = do
  pfn' <- onTypeCheckInstrAnnErr instr i "wrong car type:" (convergeAnns fn pfn)
  let vn' = deriveSpecialVN vn pfn' pairVN
      i' = ( STPair a b
           , NTPair pairTN pfn' qfn pns qns
           , pairVN ) ::& rs
  pure $ i' :/ AnnCAR fn ::: ((a, pns, vn') ::& rs)

typeCheckInstr instr@(U.CDR vn fn)
          (i@(( STPair a b
                      , NTPair pairTN pfn qfn pns qns
                      , pairVN ) ::& rs)) = do
  qfn' <- onTypeCheckInstrAnnErr instr i "wrong cdr type:" (convergeAnns fn qfn)

  let vn' = deriveSpecialVN vn qfn' pairVN
      i' = ( STPair a b
           , NTPair pairTN pfn qfn' pns qns
           , pairVN ) ::& rs
  pure $ i' :/ AnnCDR fn ::: ((b, qns, vn') ::& rs)

typeCheckInstr (U.LEFT tn vn pfn qfn bMt) i@((a, an, _) ::& rs) =
  withUType bMt $ \(b :: Sing b, bn :: Notes b) -> do
    let ns = NTOr tn pfn qfn an bn
    pure (i :/ LEFT ::: ((STOr a b, ns, vn) ::& rs))

typeCheckInstr (U.RIGHT tn vn pfn qfn aMt) i@((b, bn, _) ::& rs) =
  withUType aMt $ \(a :: Sing a, an :: Notes a) -> do
    let ns = NTOr tn pfn qfn an bn
    pure (i :/ RIGHT ::: ((STOr a b, ns, vn) ::& rs))

typeCheckInstr (U.IF_LEFT mp mq) i@((STOr a b, ons, ovn) ::& rs) = do
  let (an, bn, avn, bvn) = deriveNsOr ons ovn
      ait = (a, an, avn) ::& rs
      bit = (b, bn, bvn) ::& rs
  genericIf IF_LEFT U.IF_LEFT mp mq ait bit i

typeCheckInstr (U.NIL tn vn elMt) i =
  withUType elMt $ \(elT :: Sing elT, elNotes :: Notes elT) ->
    pure $ i :/ NIL ::: ((STList elT, NTList tn elNotes, vn) ::& i)

typeCheckInstr instr@(U.CONS vn) i@((((at :: Sing a), an, _)
                              ::& (STList (_ :: Sing a'), ln, _) ::& rs)) =
  case eqType @a @a' of
    Right Refl -> do
      n <- onTypeCheckInstrAnnErr instr i "wrong cons type:" (converge ln (NTList def an))
      pure $ i :/ CONS ::: ((STList at, n, vn) ::& rs)
    Left m -> onTypeCheckInstrErr instr (SomeHST i)
                ("list element type is different from one "
                <> "that is being CONSed:") (Left m)

typeCheckInstr (U.IF_CONS mp mq) i@((STList a, ns, vn) ::& rs)  = do
  let NTList _ an = ns
      ait =
        (a, an, vn <> "hd") ::& (STList a, ns, vn <> "tl") ::& rs
  genericIf IF_CONS U.IF_CONS mp mq ait rs i

typeCheckInstr (U.SIZE vn) i@((STList _, _, _) ::& _)  = sizeImpl i vn
typeCheckInstr (U.SIZE vn) i@((STSet _, _, _) ::& _)  = sizeImpl i vn
typeCheckInstr (U.SIZE vn) i@((STMap _ _, _, _) ::& _)  = sizeImpl i vn
typeCheckInstr (U.SIZE vn) i@((STc SCString, _, _) ::& _)  = sizeImpl i vn
typeCheckInstr (U.SIZE vn) i@((STc SCBytes, _, _) ::& _)  = sizeImpl i vn

typeCheckInstr (U.EMPTY_SET tn vn (U.Comparable mk ktn)) i =
  withSomeSingCT mk $ \k ->
    pure $ i :/ EMPTY_SET ::: ((STSet k, NTSet tn ktn, vn) ::& i)

typeCheckInstr (U.EMPTY_MAP tn vn (U.Comparable mk ktn) mv) i =
  withUType mv $ \(v :: Sing v, vns :: Notes v) ->
  withSomeSingCT mk $ \k -> do
    let ns = NTMap tn ktn vns
    pure $ i :/ EMPTY_MAP ::: ((STMap k v, ns, vn) ::& i)

typeCheckInstr (U.EMPTY_BIG_MAP tn vn (U.Comparable mk ktn) mv) i =
  withUType mv $ \(v :: Sing v, vns :: Notes v) ->
  withSomeSingCT mk $ \k -> do
    let ns = NTBigMap tn ktn vns
    pure $ i :/ EMPTY_BIG_MAP ::: ((STBigMap k v, ns, vn) ::& i)

typeCheckInstr instr@(U.MAP vn mp) i@((STList _, NTList _ vns, _vn) ::& _) = do
  mapImpl vns instr mp i
    (\rt rn -> (::&) (STList rt, NTList def rn, vn))

typeCheckInstr
  instr@(U.MAP vn mp)
  i@((STMap k _, NTMap _ kns vns, _vn) ::& _)
    = do
  let pns = NTPair def def def (NTc kns) vns
  mapImpl pns instr mp i
    (\rt rn -> (::&) (STMap k rt, NTMap def kns rn, vn))

-- case `U.HSTER []` is wrongly typed by definition
-- (as it is required to at least drop an element), so we don't consider it

typeCheckInstr instr@(U.ITER (i1 : ir)) i@((STSet _, NTSet _ en, _) ::& _) = do
  iterImpl (NTc en) instr (i1 :| ir) i
typeCheckInstr
  instr@(U.ITER (i1 : ir))
  i@((STList _, NTList _ en, _) ::& _)
    = do
  iterImpl en instr (i1 :| ir) i
typeCheckInstr
  instr@(U.ITER (i1 : ir))
  i@((STMap _ _, NTMap _ kns vns, _) ::& _)
    = do
  let en = NTPair def def def (NTc kns) vns
  iterImpl en instr (i1 :| ir) i

typeCheckInstr instr@(U.MEM vn)
           i@((STc _, _, _) ::& (STSet _, _, _) ::& _) = memImpl instr i vn
typeCheckInstr instr@(U.MEM vn)
           i@((STc _, _, _) ::& (STMap _ _, _, _) ::& _) = memImpl instr i vn
typeCheckInstr instr@(U.MEM vn)
           i@((STc _, _, _) ::& (STBigMap _ _, _, _) ::& _) = memImpl instr i vn

typeCheckInstr instr@(U.GET vn) i@(_ ::& (STMap _ vt, NTMap _ _ v, _) ::& _) =
  getImpl instr i vt v vn
typeCheckInstr
  instr@(U.GET vn)
  i@(_ ::& (STBigMap _ vt, NTBigMap _ _ v, _) ::& _)
    =
  getImpl instr i vt v vn

typeCheckInstr instr@(U.UPDATE vn) i@(_ ::& _ ::& (STMap _ _, _, _) ::& _) =
  updImpl instr i vn
typeCheckInstr instr@(U.UPDATE vn) i@(_ ::& _ ::& (STBigMap _ _, _, _) ::& _) =
  updImpl instr i vn
typeCheckInstr instr@(U.UPDATE vn) i@(_ ::& _ ::& (STSet _, _, _) ::& _) =
  updImpl instr i vn

typeCheckInstr (U.IF mp mq) i@((STc SCBool, _, _) ::& rs)  =
  genericIf IF U.IF mp mq rs rs i

typeCheckInstr instr@(U.LOOP is)
           i@((STc SCBool, _, _) ::& (rs :: HST rs)) = do
  _ :/ tp <- lift $ typeCheckList is rs
  case tp of
    subI ::: (o :: HST o) -> do
      case eqHST o (sing @('Tc 'CBool) -:& rs) of
        Right Refl -> do
          let _ ::& rs' = o
          pure $ i :/ LOOP subI ::: rs'
        Left m -> onTypeCheckInstrErr instr (SomeHST i)
                    "iteration expression has wrong output stack type:" (Left m)
    AnyOutInstr subI ->
      pure $ i :/ LOOP subI ::: rs

typeCheckInstr instr@(U.LOOP_LEFT is)
           i@((STOr (at :: Sing a) (bt :: Sing b), ons, ovn)
                      ::& (rs :: HST rs)) = do
  let (an, bn, avn, bvn) = deriveNsOr ons ovn
      ait = (at, an, avn) ::& rs
  _ :/ tp <- lift $ typeCheckList is ait
  case tp of
    subI ::: o -> do
      case (eqHST o (sing @('TOr a b) -:& rs), o) of
        (Right Refl, ((STOr _ bt', ons', ovn') ::& rs')) -> do
            let (_, bn', _, bvn') = deriveNsOr ons' ovn'
            br <- onTypeCheckInstrAnnErr instr i "wrong LOOP_LEFT input type:" $
                    convergeHSTEl (bt, bn, bvn) (bt', bn', bvn')
            pure $ i :/ LOOP_LEFT subI ::: (br ::& rs')
        (Left m, _) -> onTypeCheckInstrErr instr (SomeHST i)
                         "iteration expression has wrong output stack type:" (Left m)
    AnyOutInstr subI -> do
      let br = (bt, bn, bvn)
      pure $ i :/ LOOP_LEFT subI ::: (br ::& rs)

typeCheckInstr instr@(U.LAMBDA vn imt omt is) i = do
  withUType imt $ \(it :: Sing it, ins :: Notes ins) ->
    withUType omt $ \(ot :: Sing ot, ons :: Notes ons) ->
      -- further processing is extracted into another function because
      -- otherwise I encountered some weird GHC error with that code
      -- located right here
      lamImpl instr is vn it ins ot ons i

typeCheckInstr instr@(U.EXEC vn) i@(((_ :: Sing t1), _, _)
                              ::& ( STLambda (_ :: Sing t1') t2
                                  , NTLambda _ _ t2n
                                  , _
                                  )
                              ::& rs) = do
  Refl <- onTypeCheckInstrErr instr (SomeHST i)
                "lambda is given argument with wrong type:" (eqType @t1 @t1')
  pure $ i :/ EXEC ::: ((t2, t2n, vn) ::& rs)

typeCheckInstr instr@(U.APPLY vn)
               i@(((_ :: Sing a'), _, _)
                  ::& (STLambda (STPair (sa :: Sing a) (sb :: Sing b)) sc, ln, _)
                  ::& rs) = do
  let
    l2n = case ln of
      NTLambda vann (NTPair _ _ _ _ bn) cn ->
        NTLambda vann bn cn

  proofArgEq <- onTypeCheckInstrErr instr (SomeHST i)
                "lambda is given argument with wrong type:" (eqType @a' @a)
  proofOp <-
    maybe (onTypeCheckInstrErr instr (SomeHST i)
            "Operations in argument of partially applied lambda are not allowed"
            $ Left $ UnsupportedTypes [demote @a])
    pure (opAbsense sa)
  proofBigMap <-
    maybe (onTypeCheckInstrErr instr (SomeHST i)
            "BigMaps in argument of partially applied lambda are not allowed"
            $ Left $ UnsupportedTypes [demote @a])
    pure (bigMapAbsense sa)
  proofContract <-
    maybe (onTypeCheckInstrErr instr (SomeHST i)
            "Contract type in constant is not allowed"
            $ Left $ UnsupportedTypes [demote @a])
    pure (contractTypeAbsense sa)
  case (proofArgEq, proofOp, proofBigMap, proofContract) of
    (Refl, Dict, Dict, Dict) ->
      pure $ i :/ (APPLY @a) ::: ((STLambda sb sc, l2n, vn) ::& rs)

typeCheckInstr instr@(U.DIP is) i@(a ::& (s :: HST s)) = do
  typeCheckDipBody instr is s $ \subI t -> pure $ i :/ DIP subI ::: (a ::& t)

typeCheckInstr uInstr@(U.DIPN nTotal instructions) inputHST =
  go nTotal inputHST <&> \case
  TCDipHelper s subI out -> inputHST :/ DIPN s subI ::: out
  where
    go :: forall inp. Typeable inp
      => Word
      -> HST inp
      -> TypeCheckInstr (TCDipHelper inp)
    go n curHST = case (n, curHST) of
      (0, _) -> typeCheckDipBody uInstr instructions curHST $ \subI t ->
        pure (TCDipHelper SZ subI t)
      (_, SNil) -> onTypeCheckInstrErr uInstr
        (SomeHST SNil) "" (Left (NotEnoughItemsOnStack nTotal NotEnoughDip))
      (_, hstHead ::& hstTail) ->
        go (n - 1) hstTail <&> \case
        TCDipHelper s subI out -> TCDipHelper (SS s) subI (hstHead ::& out)

typeCheckInstr U.FAILWITH i@(_ ::& _) =
  pure $ i :/ AnyOutInstr FAILWITH

typeCheckInstr instr@(U.CAST vn castToType)
           i@(((e :: Sing e), (en :: Notes e), evn) ::& rs) =
  withUType castToType $
    \(_ :: Sing castToType, castToNotes :: Notes castToType) -> do
      Refl <- errM (eqType @e @castToType)
      void $ errM (converge en castToNotes `onLeft` AnnError)
      pure $ i :/ CAST ::: ((e, castToNotes, vn `orAnn` evn) ::& rs)
  where
    errM :: (MonadReader InstrCallStack m, MonadError TCError m) => Either TCTypeError a -> m a
    errM = onTypeCheckInstrErr instr (SomeHST i) "cast to incompatible type:"

typeCheckInstr (U.RENAME vn) i@((at, an, _) ::& rs) =
  pure $ i :/ RENAME ::: ((at, an, vn) ::& rs)

typeCheckInstr instr@(U.UNPACK tn vn mt) i@((STc SCBytes, _, _) ::& rs) =
  withUType mt $ \(t :: Sing t, tns :: Notes t) -> do
    let ns = NTOption tn tns
    proofOp <-
      maybe (typeCheckInstrErr instr (SomeHST i)
             "Operation cannot appear in serialized data")
      pure (opAbsense t)
    proofBigMap <-
      maybe (typeCheckInstrErr instr (SomeHST i)
             "BigMap cannot appear in serialized data")
      pure (bigMapAbsense t)
    proofContract <-
      maybe (typeCheckInstrErr instr (SomeHST i)
             "UNPACK should not contain 'contract' type")
      pure (contractTypeAbsense t)
    case (proofOp, proofBigMap, proofContract) of
      (Dict, Dict, Dict) -> pure $ i :/ UNPACK ::: ((STOption t, ns, vn) ::& rs)

typeCheckInstr instr@(U.PACK vn) i@(((a :: Sing a), _, _) ::& rs) = do
  proofOp <-
    maybe (onTypeCheckInstrErr instr (SomeHST i)
           "Operations cannot appear in serialized data"
           $ Left $ UnsupportedTypes [demote @a])
    pure (opAbsense a)
  proofBigMap <-
    maybe (onTypeCheckInstrErr instr (SomeHST i)
           "BigMap cannot appear in serialized data"
           $ Left $ UnsupportedTypes [demote @a])
    pure (bigMapAbsense a)
  case (proofOp, proofBigMap) of
    (Dict, Dict) -> pure $ i :/ PACK ::: ((STc SCBytes, starNotes, vn) ::& rs)

typeCheckInstr (U.CONCAT vn) i@((STc SCBytes, _, _) ::&
                                    (STc SCBytes, _, _) ::& _) = concatImpl i vn
typeCheckInstr (U.CONCAT vn) i@((STc SCString, _, _) ::&
                                    (STc SCString, _, _) ::& _) = concatImpl i vn
typeCheckInstr (U.CONCAT vn) i@((STList (STc SCBytes), _, _) ::& _) = concatImpl' i vn
typeCheckInstr (U.CONCAT vn) i@((STList (STc SCString), _, _) ::& _) = concatImpl' i vn
typeCheckInstr (U.SLICE vn) i@((STc SCNat, _, _) ::&
                                   (STc SCNat, _, _) ::&
                                   (STc SCString, _, _) ::& _) = sliceImpl i vn
typeCheckInstr (U.SLICE vn) i@((STc SCNat, _, _) ::&
                                   (STc SCNat, _, _) ::&
                                   (STc SCBytes, _, _) ::& _) = sliceImpl i vn

typeCheckInstr (U.ISNAT vn') i@((STc SCInt, _, oldVn) ::& rs) = do
  let vn = vn' `orAnn` oldVn
  pure $ i :/ ISNAT ::: ((STOption (STc SCNat), starNotes, vn) ::& rs)

typeCheckInstr (U.ADD vn) i@((STc a, _, _) ::&
                                 (STc b, _, _) ::& _) = addImpl a b i vn

typeCheckInstr (U.SUB vn) i@((STc a, _, _) ::&
                                 (STc b, _, _) ::& _) = subImpl a b i vn

typeCheckInstr (U.MUL vn) i@((STc a, _, _) ::&
                                 (STc b, _, _) ::& _) = mulImpl a b i vn

typeCheckInstr (U.EDIV vn) i@((STc a, _, _) ::&
                                  (STc b, _, _) ::& _) = edivImpl a b i vn

typeCheckInstr (U.ABS vn) i@((STc SCInt, _, _) ::& _) = unaryArithImpl @Abs ABS i vn

typeCheckInstr (U.NEG vn) (i@((STc SCInt, _, _) ::& _)) = unaryArithImpl @Neg NEG i vn

typeCheckInstr (U.NEG vn) (i@((STc SCNat, _, _) ::& _)) = unaryArithImpl @Neg NEG i vn

typeCheckInstr (U.LSL vn) i@((STc SCNat, _, _) ::&
                         (STc SCNat, _, _) ::& _) = arithImpl @Lsl LSL i vn

typeCheckInstr (U.LSR vn) i@((STc SCNat, _, _) ::&
                         (STc SCNat, _, _) ::& _) = arithImpl @Lsr LSR i vn

typeCheckInstr (U.OR vn) i@((STc SCBool, _, _) ::&
                         (STc SCBool, _, _) ::& _) = arithImpl @Or OR i vn
typeCheckInstr (U.OR vn) i@((STc SCNat, _, _) ::&
                         (STc SCNat, _, _) ::& _) = arithImpl @Or OR i vn

typeCheckInstr (U.AND vn) i@((STc SCInt, _, _) ::&
                         (STc SCNat, _, _) ::& _) = arithImpl @And AND i vn
typeCheckInstr (U.AND vn) i@((STc SCNat, _, _) ::&
                         (STc SCNat, _, _) ::& _) = arithImpl @And AND i vn
typeCheckInstr (U.AND vn) i@((STc SCBool, _, _) ::&
                         (STc SCBool, _, _) ::& _) = arithImpl @And AND i vn

typeCheckInstr (U.XOR vn) i@((STc SCBool, _, _) ::&
                         (STc SCBool, _, _) ::& _) = arithImpl @Xor XOR i vn
typeCheckInstr (U.XOR vn) i@((STc SCNat, _, _) ::&
                         (STc SCNat, _, _) ::& _) = arithImpl @Xor XOR i vn

typeCheckInstr (U.NOT vn) i@((STc SCNat, _, _) ::& _) = unaryArithImpl @Not NOT i vn
typeCheckInstr (U.NOT vn) i@((STc SCBool, _, _) ::& _) = unaryArithImpl @Not NOT i vn
typeCheckInstr (U.NOT vn) i@((STc SCInt, _, _) ::& _) = unaryArithImpl @Not NOT i vn

typeCheckInstr instr@(U.COMPARE vn)
  i@((STc (_ :: Sing aCT), (an :: Notes ('Tc aCT)), _) ::&
     (STc (_ :: Sing bCT), (bn :: Notes ('Tc bCT)), _) ::& _
    ) = do
  Refl <- errM $ eqType @('Tc aCT) @('Tc bCT)
  void . errConv $ converge an bn
  compareImpl (sing @aCT) i vn
  where
    errConv :: (MonadReader InstrCallStack m, MonadError TCError m) => Either AnnConvergeError a -> m a
    errConv = onTypeCheckInstrAnnErr instr i "comparing types with different annotations:"
    errM :: (MonadReader InstrCallStack m, MonadError TCError m) => Either TCTypeError a -> m a
    errM = onTypeCheckInstrErr instr (SomeHST i) "comparing different types:"

typeCheckInstr (U.EQ vn) i@((STc SCInt, _, _) ::& _) = unaryArithImpl @Eq' EQ i vn

typeCheckInstr (U.NEQ vn) i@((STc SCInt, _, _) ::& _) = unaryArithImpl @Neq NEQ i vn

typeCheckInstr (U.LT vn) i@((STc SCInt, _, _) ::& _) = unaryArithImpl @Lt LT i vn

typeCheckInstr (U.GT vn) i@((STc SCInt, _, _) ::& _) = unaryArithImpl @Gt GT i vn

typeCheckInstr (U.LE vn) i@((STc SCInt, _, _) ::& _) = unaryArithImpl @Le LE i vn

typeCheckInstr (U.GE vn) i@((STc SCInt, _, _) ::& _) = unaryArithImpl @Ge GE i vn

typeCheckInstr (U.INT vn) i@((STc SCNat, _, _) ::& rs) =
  pure $ i :/ INT ::: ((STc SCInt, starNotes, vn) ::& rs)

typeCheckInstr instr@(U.SELF vn) i = do
  cpType <- gets tcContractParam
  let t = fromUType cpType
  withUType cpType $ \(singcp :: Sing cp, nt :: Notes cp) -> do
    proofOp <-
      maybe (onTypeCheckInstrErr instr (SomeHST i)
              "contract param type cannot contain operation"
              $ Left $ UnsupportedTypes [t])
      pure (opAbsense singcp)
    proofBigMap <-
      maybe (onTypeCheckInstrErr instr (SomeHST i)
              "contract param type cannot contain nested big_map"
              $ Left $ UnsupportedTypes [t])
      pure (nestedBigMapsAbsense singcp)
    case (proofOp, proofBigMap) of
      (Dict, Dict) -> do
        pure $ i :/ SELF @cp ::: ((sing @('TContract cp), NTContract U.noAnn nt, vn) ::& i)

typeCheckInstr instr@(U.CONTRACT vn fn mt)
           i@((STc SCAddress, _, _) ::& rs) =
  withUType mt $ \(t :: Sing t, tns :: Notes t) -> do
    proofOp <-
      maybe (onTypeCheckInstrErr instr (SomeHST i)
              "contract param type cannot contain operation"
              $ Left $ UnsupportedTypes [fromUType mt])
      pure (opAbsense t)
    proofBigMap <-
      maybe (onTypeCheckInstrErr instr (SomeHST i)
              "contract param type cannot contain nested big_map"
              $ Left $ UnsupportedTypes [fromUType mt])
      pure (nestedBigMapsAbsense t)
    let ns = NTOption def $ NTContract def tns
    epName <- onTypeCheckInstrErr instr (SomeHST i) "Bad annotation:" $
                epNameFromRefAnn fn `onLeft` IllegalEntryPoint
    case (proofOp, proofBigMap) of
      (Dict, Dict) ->
        pure $ i :/ CONTRACT tns epName ::: ((STOption $ STContract t, ns, vn) ::& rs)

typeCheckInstr instr@(U.TRANSFER_TOKENS vn) i@(((_ :: Sing p'), _, _)
  ::& (STc SCMutez, _, _) ::& (STContract (p :: Sing p), _, _) ::& rs) = do
  proofOp <-
    maybe (onTypeCheckInstrErr instr (SomeHST i)
            "contract param type cannot contain operation"
            $ Left $ UnsupportedTypes [demote @p])
    pure (opAbsense p)
  proofBigMap <-
    maybe (onTypeCheckInstrErr instr (SomeHST i)
            "contract param type cannot contain big_map"
            $ Left $ UnsupportedTypes [demote @p])
    pure (nestedBigMapsAbsense p)
  case (eqType @p @p', proofOp, proofBigMap) of
    (Right Refl, Dict, Dict) ->
      pure $ i :/ TRANSFER_TOKENS ::: ((STOperation, starNotes, vn) ::& rs)
    (Left m, _, _) ->
      onTypeCheckInstrErr instr (SomeHST i)
        "mismatch of contract param type:" (Left m)

typeCheckInstr (U.SET_DELEGATE vn)
           i@((STOption (STc SCKeyHash), _, _) ::& rs) = do
  pure $ i :/ SET_DELEGATE ::: ((STOperation, starNotes, vn) ::& rs)

typeCheckInstr instr@(U.CREATE_CONTRACT ovn avn contract)
           i@((STOption (STc SCKeyHash), _, _)
             ::& (STc SCMutez, _, _)
             ::& ((_ :: Sing g), gn, _) ::& rs) = do
  (SomeContract (FullContract (contr :: Contract p' g') paramNotes storeNotes))
    <- lift $ typeCheckContractImpl contract
  Refl <- checkEqT @g @g' instr i "contract storage type mismatch"
  void $ onTypeCheckInstrAnnErr instr i "contract storage type mismatch" (converge gn storeNotes)
  pure $ i :/ CREATE_CONTRACT (FullContract contr paramNotes storeNotes) :::
         ((STOperation, starNotes, ovn) ::& (STc SCAddress, starNotes, avn) ::& rs)

typeCheckInstr (U.IMPLICIT_ACCOUNT vn)
           i@((STc SCKeyHash, _, _) ::& rs) =
  pure $ i :/ IMPLICIT_ACCOUNT ::: ((STContract STUnit, starNotes, vn) ::& rs)

typeCheckInstr (U.NOW vn) i =
  pure $ i :/ NOW ::: ((STc SCTimestamp, starNotes, vn) ::& i)

typeCheckInstr (U.AMOUNT vn) i =
  pure $ i :/ AMOUNT ::: ((STc SCMutez, starNotes, vn) ::& i)

typeCheckInstr (U.BALANCE vn) i =
  pure $ i :/ BALANCE ::: ((STc SCMutez, starNotes, vn) ::& i)

typeCheckInstr (U.CHECK_SIGNATURE vn)
           i@((STKey, _, _)
             ::& (STSignature, _, _) ::& (STc SCBytes, _, _) ::& rs) =
  pure $ i :/ CHECK_SIGNATURE ::: ((STc SCBool, starNotes, vn) ::& rs)

typeCheckInstr (U.SHA256 vn)
           i@((STc SCBytes, _, _) ::& rs) =
  pure $ i :/ SHA256 ::: ((STc SCBytes, starNotes, vn) ::& rs)

typeCheckInstr (U.SHA512 vn)
           i@((STc SCBytes, _, _) ::& rs) =
  pure $ i :/ SHA512 ::: ((STc SCBytes, starNotes, vn) ::& rs)

typeCheckInstr (U.BLAKE2B vn)
           i@((STc SCBytes, _, _) ::& rs) =
  pure $ i :/ BLAKE2B ::: ((STc SCBytes, starNotes, vn) ::& rs)

typeCheckInstr (U.HASH_KEY vn)
           i@((STKey, _, _) ::& rs) =
  pure $ i :/ HASH_KEY ::: ((STc SCKeyHash, starNotes, vn) ::& rs)

typeCheckInstr (U.STEPS_TO_QUOTA vn) i =
  pure $ i :/ STEPS_TO_QUOTA ::: ((STc SCNat, starNotes, vn) ::& i)

typeCheckInstr (U.SOURCE vn) i =
  pure $ i :/ SOURCE ::: ((STc SCAddress, starNotes, vn) ::& i)

typeCheckInstr (U.SENDER vn) i =
  pure $ i :/ SENDER ::: ((STc SCAddress, starNotes, vn) ::& i)

typeCheckInstr (U.ADDRESS vn) i@((STContract _, _, _) ::& rs) =
  pure $ i :/ ADDRESS ::: ((STc SCAddress, starNotes, vn) ::& rs)

typeCheckInstr (U.CHAIN_ID vn) i =
  pure $ i :/ CHAIN_ID ::: ((STChainId, starNotes, vn) ::& i)

typeCheckInstr instr sit = typeCheckInstrErr instr (SomeHST sit) "unknown expression"

-- | Helper function for two-branch if where each branch is given a single
-- value.
genericIf
  :: forall bti bfi cond rs .
    (Typeable bti, Typeable bfi)
  => (forall s'.
        Instr bti s' ->
        Instr bfi s' ->
        Instr (cond ': rs) s'
     )
  -> ([U.ExpandedOp] -> [U.ExpandedOp] -> U.ExpandedInstr)
  -> [U.ExpandedOp]
  -> [U.ExpandedOp]
  -> HST bti
  -> HST bfi
  -> HST (cond ': rs)
  -> TypeCheckInstr (SomeInstr (cond ': rs))
genericIf cons mCons mbt mbf bti bfi i@(_ ::& _) = do
  _ :/ pinstr <- lift $ typeCheckList mbt bti
  _ :/ qinstr <- lift $ typeCheckList mbf bfi
  fmap (i :/) $ case (pinstr, qinstr) of
    (p ::: po, q ::: qo) -> do
      let instr = mCons mbt mbf
      Refl <- checkEqHST qo po instr i
                    "branches have different output stack types:"
      o <- onTypeCheckInstrAnnErr instr i "branches have different output stack types:" (convergeHST po qo)
      pure $ cons p q ::: o
    (AnyOutInstr p, q ::: qo) -> do
      pure $ cons p q ::: qo
    (p ::: po, AnyOutInstr q) -> do
      pure $ cons p q ::: po
    (AnyOutInstr p, AnyOutInstr q) ->
      pure $ AnyOutInstr (cons p q)

mapImpl
  :: forall c rs .
    ( MapOp c
    , SingI (MapOpInp c)
    , Typeable (MapOpInp c)
    , Typeable (MapOpRes c)
    )
  => Notes (MapOpInp c)
  -> U.ExpandedInstr
  -> [U.ExpandedOp]
  -> HST (c ': rs)
  -> (forall v' . (Typeable v', SingI v') =>
        Sing v' -> Notes v' -> HST rs -> HST (MapOpRes c v' ': rs))
  -> TypeCheckInstr (SomeInstr (c ': rs))
mapImpl vn instr mp i@(_ ::& rs) mkRes = do
  _ :/ subp <- lift $ typeCheckList mp ((sing, vn, def) ::& rs)
  case subp of
    sub ::: subo ->
      case subo of
        (b, bn, _bvn) ::& rs' -> do
          Refl <- checkEqHST rs rs' instr i $
                      "map expression has changed not only top of the stack"
          pure $ i :/ MAP sub ::: mkRes b bn rs'
        _ -> typeCheckInstrErr instr (SomeHST i) "map expression has wrong output stack type (empty stack)"
    AnyOutInstr _ ->
      typeCheckInstrErr instr (SomeHST i) "MAP code block always fails, which is not allowed"

iterImpl
  :: forall c rs .
    ( IterOp c
    , SingI (IterOpEl c)
    , Typeable (IterOpEl c)
    )
  => Notes (IterOpEl c)
  -> U.ExpandedInstr
  -> NonEmpty U.ExpandedOp
  -> HST (c ': rs)
  -> TypeCheckInstr (SomeInstr (c ': rs))
iterImpl en instr mp i@((_, _, lvn) ::& rs) = do
  let evn = deriveVN "elt" lvn
  _ :/ subp <- lift $ typeCheckNE mp ((sing, en, evn) ::& rs)
  case subp of
    subI ::: o -> do
      Refl <- checkEqHST o rs instr i
                "iteration expression has wrong output stack type"
      pure $ i :/ ITER subI ::: o
    AnyOutInstr _ ->
      typeCheckInstrErr instr (SomeHST i) "ITER code block always fails, which is not allowed"

lamImpl
  :: forall it ot ts .
    ( Typeable it, Typeable ts, Typeable ot
    , SingI it, SingI ot
    )
  => U.ExpandedInstr
  -> [U.ExpandedOp]
  -> VarAnn
  -> Sing it -> Notes it
  -> Sing ot -> Notes ot
  -> HST ts
  -> TypeCheckInstr (SomeInstr ts)
lamImpl instr is vn it ins ot ons i = do
  when (any hasSelf is) $
    typeCheckInstrErr instr (SomeHST i) "The SELF instruction cannot appear in a lambda"
  _ :/ lamI <- lift $ typeCheckList is ((it, ins, def) ::& SNil)
  let lamNotes onsr = NTLambda def ins onsr
  let lamSt onsr = (STLambda it ot, lamNotes onsr, vn) ::& i
  fmap (i :/) $ case lamI of
    lam ::: lo -> do
      case eqHST1 @ot lo of
        Right Refl -> do
            let (_, ons', _) ::& SNil = lo
            onsr <- onTypeCheckInstrAnnErr instr i "wrong output type of lambda's expression:" (converge ons ons')
            pure (LAMBDA (VLam $ RfNormal lam) ::: lamSt onsr)
        Left m -> onTypeCheckInstrErr instr (SomeHST i)
                    "wrong output type of lambda's expression:" (Left m)
    AnyOutInstr lam ->
      pure (LAMBDA (VLam $ RfAlwaysFails lam) ::: lamSt ons)
  where
    hasSelf :: U.ExpandedOp -> Bool
    hasSelf = everything (||)
      (mkQ False
       (\case
           (U.SELF _ :: U.InstrAbstract U.ExpandedOp) -> True
           _ -> False
       )
      )

----------------------------------------------------------------------------
-- Helpers for DIP (n) typechecking
----------------------------------------------------------------------------

-- Helper data type we use to typecheck DIPN.
data TCDipHelper inp where
  TCDipHelper ::
    forall (n :: Peano) inp out s s'.
    (Typeable out, ConstraintDIPN n inp out s s') =>
    Sing n -> Instr s s' -> HST out -> TCDipHelper inp

typeCheckDipBody ::
     forall inp r. Typeable inp
  => U.ExpandedInstr
  -> [U.ExpandedOp]
  -> HST inp
  -> (forall out. Typeable out =>
                    Instr inp out -> HST out -> TypeCheckInstr r)
  -> TypeCheckInstr r
typeCheckDipBody mainInstr instructions inputHST callback = do
  _ :/ tp <- lift (typeCheckList instructions inputHST)
  case tp of
    AnyOutInstr _ ->
      -- This may seem like we throw error because of despair, but in fact,
      -- the reference implementation seems to behave exactly in this way -
      -- if output stack of code block within @DIP@ occurs to be any, an
      -- error "FAILWITH must be at tail position" is raised.
      -- It is not allowed even in `DIP 0`.
      typeCheckInstrErr mainInstr (SomeHST inputHST)
          "Code within DIP instruction always fails, which is not allowed"
    subI ::: t -> callback subI t