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

module Michelson.TypeCheck.Helpers
    ( deriveSpecialVN
    , deriveSpecialFNs
    , deriveVN
    , deriveNsOr
    , deriveNsOption
    , convergeHSTEl
    , convergeHST
    , hstToTs
    , eqHST
    , eqHST1
    , lengthHST

    , ensureDistinctAsc
    , eqType
    , onTypeCheckInstrAnnErr
    , onTypeCheckInstrErr
    , onScopeCheckInstrErr
    , typeCheckInstrErr
    , typeCheckInstrErr'
    , typeCheckImpl
    , typeCheckImplStripped
    , matchTypes

    , memImpl
    , getImpl
    , updImpl
    , sliceImpl
    , concatImpl
    , concatImpl'
    , sizeImpl
    , arithImpl
    , addImpl
    , subImpl
    , mulImpl
    , edivImpl
    , unaryArithImpl
    , unaryArithImplAnnotated
    , withCompareableCheck
    ) where

import Prelude hiding (EQ, GT, LT)

import Control.Monad.Except (MonadError, throwError)
import Data.Constraint (Dict(..), withDict)
import Data.Default (def)
import Data.Singletons (Sing, SingI(sing), demote)
import qualified Data.Text as T
import Data.Typeable ((:~:)(..), eqT)
import Fmt ((+||), (||+))

import Michelson.ErrorPos (InstrCallStack)
import Michelson.TypeCheck.Error (TCError(..), TCTypeError(..), TypeContext(..))
import Michelson.TypeCheck.TypeCheck
import Michelson.TypeCheck.Types
import Michelson.TypeCheck.TypeCheckedSeq (IllTypedInstr(..), TypeCheckedSeq(..))
import Michelson.Typed
  (BadTypeForScope(..), CommentType(StackTypeComment), Comparable, ExtInstr(COMMENT_ITEM),
  Instr(..), KnownT, Notes(..), PackedNotes(..), SingT(..), T(..), WellTyped, converge,
  getComparableProofS, notesT, orAnn, starNotes)
import Michelson.Typed.Annotation (AnnConvergeError, isStar)
import Michelson.Typed.Arith (Add, ArithOp(..), Mul, Sub, UnaryArithOp(..))
import Michelson.Typed.Polymorphic
  (ConcatOp, EDivOp(..), GetOp(..), MemOp(..), SizeOp, SliceOp, UpdOp(..))

import qualified Michelson.Untyped as Un
import Michelson.Untyped.Annotation (Annotation(..), FieldAnn, VarAnn, ann)
import Util.Type (onFirst)

-- | Function which derives special annotations
-- for PAIR instruction.
--
-- Namely, it does following transformation:
-- @
--  PAIR %@@ %@@ [ @@p.a int : @@p.b int : .. ]
--  ~
--  [ @@p (pair (int %a) (int %b) : .. ]
-- @
--
-- All relevant cases (e.g. @PAIR %myf %@@ @)
-- are handled as they should be according to spec.
deriveSpecialFNs
  :: FieldAnn -> FieldAnn
  -> VarAnn -> VarAnn
  -> (VarAnn, FieldAnn, FieldAnn)
deriveSpecialFNs "@" "@" pvn qvn = (vn, pfn, qfn)
  where
    ps = T.splitOn "." $ unAnnotation pvn
    qs = T.splitOn "." $ unAnnotation qvn
    fns = fst <$> takeWhile (uncurry (==)) (zip ps qs)
    vn = ann $ T.intercalate "." fns
    pfn = ann $ T.intercalate "." $ drop (length fns) ps
    qfn = ann $ T.intercalate "." $ drop (length fns) qs
deriveSpecialFNs "@" qfn pvn _   = (def, Un.convAnn pvn, qfn)
deriveSpecialFNs pfn "@" _ qvn   = (def, pfn, Un.convAnn qvn)
deriveSpecialFNs pfn qfn _ _     = (def, pfn, qfn)

-- | Function which derives special annotations
-- for CDR / CAR instructions.
deriveSpecialVN :: VarAnn -> FieldAnn -> VarAnn -> VarAnn
deriveSpecialVN vn elFn pairVN
  | vn == "%" = Un.convAnn elFn
  | vn == "%%" && elFn /= def = pairVN <> Un.convAnn elFn
  | otherwise = vn

-- | Append suffix to variable annotation (if it's not empty)
deriveVN :: VarAnn -> VarAnn -> VarAnn
deriveVN suffix vn = bool (suffix <> vn) def (vn == def)

-- | Function which extracts annotations for @or@ type
-- (for left and right parts).
--
-- It extracts field/type annotations and also auto-generates variable
-- annotations if variable annotation is not provided as second argument.
deriveNsOr :: Notes ('TOr a b) -> VarAnn -> (Notes a, Notes b, VarAnn, VarAnn)
deriveNsOr (NTOr _ afn bfn an bn) ovn =
  let avn = deriveVN (Un.convAnn afn `orAnn` "left") ovn
      bvn = deriveVN (Un.convAnn bfn `orAnn` "right") ovn
   in (an, bn, avn, bvn)

-- | Function which extracts annotations for @option t@ type.
--
-- It extracts field/type annotations and also auto-generates variable
-- annotation for @Some@ case if it is not provided as second argument.
deriveNsOption :: Notes ('TOption a) -> VarAnn -> (Notes a, VarAnn)
deriveNsOption (NTOption _ an) ovn =
  let avn = deriveVN "some" ovn
   in (an, avn)

convergeHSTEl
  :: (Notes t, Dict (WellTyped t), VarAnn)
  -> (Notes t, Dict (WellTyped t), VarAnn)
  -> Either AnnConvergeError (Notes t, Dict (WellTyped t), VarAnn)
convergeHSTEl (an, d@Dict, avn) (bn, _, bvn) =
  (,,) <$> converge an bn <*> pure d
      <*> pure (bool def avn $ avn == bvn)

-- | Combine annotations from two given stack types
convergeHST :: HST ts -> HST ts -> Either AnnConvergeError (HST ts)
convergeHST SNil SNil = pure SNil
convergeHST (a ::& as) (b ::& bs) =
    liftA2 (::&) (convergeHSTEl a b) (convergeHST as bs)

-- | Extract singleton for each single type of the given stack.
hstToTs :: HST st -> [T]
hstToTs = \case
  SNil -> []
  (notes, _, _) ::& hst -> notesT notes : hstToTs hst

-- | Check whether the given stack types are equal.
eqHST
  :: forall as bs.
      (Typeable as, Typeable bs)
  => HST as -> HST bs -> Either TCTypeError (as :~: bs)
eqHST (hst :: HST xs) (hst' :: HST ys) = do
  case eqT @as @bs of
    Nothing -> Left $ StackEqError (hstToTs hst) (hstToTs hst')
    Just Refl -> do
      void $ convergeHST hst hst' `onFirst` AnnError
      return Refl

-- | Check whether the given stack has size 1 and its only element matches the
-- given type. This function is a specialized version of `eqHST`.
eqHST1
  :: forall t st.
      (Typeable st, WellTyped t)
  => HST st -> Either TCTypeError (st :~: '[t])
eqHST1 hst = do
  let hst' = sing @t -:& SNil
  case eqT @'[t] @st of
    Nothing -> Left $ StackEqError (hstToTs hst') (hstToTs hst)
    Just Refl -> Right Refl

lengthHST :: HST xs -> Natural
lengthHST (_ ::& xs) = 1 + lengthHST xs
lengthHST SNil = 0

--------------------------------------------
-- Typechecker auxiliary
--------------------------------------------

-- | Check whether elements go in strictly ascending order and
-- return the original list (to keep only one pass on the original list).
ensureDistinctAsc :: (Ord b, Show a) => (a -> b) -> [a] -> Either Text [a]
ensureDistinctAsc toCmp = \case
  (e1 : e2 : l) ->
    if toCmp e1 < toCmp e2
    then (e1 :) <$> ensureDistinctAsc toCmp (e2 : l)
    else Left $ "Entries are unordered (" +|| e1 ||+ " >= " +|| e2 ||+ ")"
  l -> Right l

-- | Function @eqType@ is a simple wrapper around @Data.Typeable.eqT@ suited
-- for use within @Either TCTypeError a@ applicative.
eqType
  :: forall (a :: T) (b :: T).
      (Each '[KnownT] [a, b])
  => Either TCTypeError (a :~: b)
eqType = maybe (Left $ TypeEqError (demote @a) (demote @b)) pure eqT

onTypeCheckInstrErr
  :: (MonadReader InstrCallStack m, MonadError TCError m)
  => Un.ExpandedInstr -> SomeHST -> Maybe TypeContext
  -> Either TCTypeError a -> m a
onTypeCheckInstrErr instr hst mContext ei = do
  either (typeCheckInstrErr' instr hst mContext) return ei

onScopeCheckInstrErr
  :: forall (t :: T) m a.
      (MonadReader InstrCallStack m, MonadError TCError m, SingI t)
  => Un.ExpandedInstr -> SomeHST -> Maybe TypeContext
  -> Either BadTypeForScope a -> m a
onScopeCheckInstrErr instr hst mContext = \case
  Right a -> return a
  Left e -> do
    pos <- ask
    throwError $ TCFailedOnInstr instr hst pos mContext $
      Just $ UnsupportedTypeForScope (demote @t) e

typeCheckInstrErr
  :: (MonadReader InstrCallStack m, MonadError TCError m)
  => Un.ExpandedInstr -> SomeHST -> Maybe TypeContext
  -> m a
typeCheckInstrErr instr hst mContext = do
  pos <- ask
  throwError $ TCFailedOnInstr instr hst pos mContext Nothing

typeCheckInstrErr'
  :: (MonadReader InstrCallStack m, MonadError TCError m)
  => Un.ExpandedInstr -> SomeHST -> Maybe TypeContext
  -> TCTypeError -> m a
typeCheckInstrErr' instr hst mContext err = do
  pos <- ask
  throwError $ TCFailedOnInstr instr hst pos mContext (Just err)

onTypeCheckInstrAnnErr
  :: (MonadReader InstrCallStack m, MonadError TCError m, Typeable ts)
  => Un.ExpandedInstr -> HST ts -> Maybe TypeContext
  -> Either AnnConvergeError a -> m a
onTypeCheckInstrAnnErr instr i mContext ei =
  onTypeCheckInstrErr instr (SomeHST i) mContext (ei `onFirst` AnnError)

withCompareableCheck
  :: forall a m v ts. (Typeable ts, MonadReader InstrCallStack m, MonadError TCError m)
  => Sing a
  -> Un.ExpandedInstr
  -> HST ts
  -> (Comparable a => v)
  -> m v
withCompareableCheck sng instr i act = case getComparableProofS sng of
  Just d@Dict -> pure $ withDict d act
  Nothing -> typeCheckInstrErr instr (SomeHST i) $ Just ComparisonArguments

typeCheckOpImpl
  :: forall inp. Typeable inp
  => TcInstrHandler
  -> Un.ExpandedOp
  -> HST inp
  -> TypeCheckInstrNoExcept (TypeCheckedSeq inp)
typeCheckOpImpl tcInstr op' inputStack = case op' of
  Un.WithSrcEx _ op@Un.WithSrcEx{} -> typeCheckOpImpl tcInstr op inputStack
  Un.WithSrcEx loc (Un.PrimEx op)  -> typeCheckPrimWithLoc loc op
  Un.WithSrcEx loc (Un.SeqEx sq)   -> typeCheckSeqWithLoc loc sq
  Un.PrimEx op                     -> typeCheckPrim op
  Un.SeqEx sq                      -> typeCheckSeq sq
  where
    -- If we know source location from the untyped instruction, keep it in the typed one.
    typeCheckPrimWithLoc :: InstrCallStack -> Un.ExpandedInstr -> TypeCheckInstrNoExcept (TypeCheckedSeq inp)
    typeCheckPrimWithLoc loc op = local (const loc)
      (wrapWithLoc loc <$> typeCheckPrim op)

    typeCheckPrim :: Un.ExpandedInstr -> TypeCheckInstrNoExcept (TypeCheckedSeq inp)
    typeCheckPrim op = tcInstr op inputStack <&> mapSeq addNotes

    typeCheckSeqWithLoc :: InstrCallStack -> [Un.ExpandedOp] -> TypeCheckInstrNoExcept (TypeCheckedSeq inp)
    typeCheckSeqWithLoc loc = fmap (wrapWithLoc loc) . local (const loc) . typeCheckSeq

    typeCheckSeq :: [Un.ExpandedOp] -> TypeCheckInstrNoExcept (TypeCheckedSeq inp)
    typeCheckSeq sq = typeCheckImpl tcInstr sq inputStack
                  <&> mapSeq (addNotes . mapSomeInstr Nested)

    addNotes :: SomeInstr inp -> SomeInstr inp
    addNotes (inp :/ i ::: out) = inp :/ wrapWithNotes out i ::: out
    addNotes i = i

    wrapWithNotes :: HST d -> Instr c d -> Instr c d
    wrapWithNotes outputStack instr = case instr of
      -- Abstractions for instructions:
      Nop -> instr'
      Seq _ _ -> instr'
      Nested _ -> instr'
      DocGroup _ _ -> instr'
      Ext _ -> instr'
      FrameInstr _ _ -> instr'
      WithLoc _ _ -> instr'
      -- These two shouldn't happen, since annotations are added here.
      InstrWithNotes _ _ -> instr'
      InstrWithVarNotes _ _ -> instr'

      -- Instructions that don't produce notes:
      DROP -> instr'
      SWAP -> instr'
      DIG _ -> instr'
      DUG _ -> instr'
      IF_NONE _ _ -> instr'
      IF_LEFT _ _ -> instr'
      IF_CONS _ _ -> instr'
      ITER _ -> instr'
      IF _ _ -> instr'
      LOOP _ -> instr'
      LOOP_LEFT _ -> instr'
      DIP _ -> instr'
      DIPN _ _ -> instr'
      FAILWITH -> instr'

      -- Instructions that produce at most two notes:
      CREATE_CONTRACT _ -> case outputStack of
        ((np, _, vp) ::& (_, _, vs) ::& _) ->
          let withNotes = if isStar np then id else InstrWithNotes (PackedNotes np)
              withVarNotes = if vp == Un.noAnn && vs == Un.noAnn then id else InstrWithVarNotes (vp :| [vs])
           in withNotes . withVarNotes $ instr

      -- Instructions that produce at most one note:
      _ -> case outputStack of
        ((n, _, v) ::& _) ->
          let withNotes = if isStar n then id else InstrWithNotes (PackedNotes n)
              withVarNotes = if v == Un.noAnn then id else InstrWithVarNotes (one v)
           in withNotes . withVarNotes $ instr
        SNil -> instr
      where
        instr' = addNotesNoVarAnn outputStack instr

    addNotesNoVarAnn :: HST d -> Instr c d -> Instr c d
    addNotesNoVarAnn outputStack instr = case outputStack of
      ((n, _, _) ::& _)
        | isStar n -> instr
        | otherwise -> InstrWithNotes (PackedNotes n) instr
      SNil -> instr

-- | Like 'typeCheckImpl' but doesn't add a stack type comment after the
-- sequence.
typeCheckImplNoLastTypeComment
  :: forall inp . Typeable inp
  => TcInstrHandler
  -> [Un.ExpandedOp]
  -> HST inp
  -> TypeCheckInstrNoExcept (TypeCheckedSeq inp)
typeCheckImplNoLastTypeComment _ [] inputStack
  = pure (WellTypedSeq (inputStack :/ Nop ::: inputStack))
typeCheckImplNoLastTypeComment tcInstr [op] inputStack = do
  typeCheckOpImpl tcInstr op inputStack
      >>= mapMSeq prependStackTypeComment
typeCheckImplNoLastTypeComment tcInstr (op : ops) inputStack = do
  done <- typeCheckOpImpl tcInstr op inputStack
      >>= mapMSeq prependStackTypeComment
  continueTypeChecking tcInstr done ops

continueTypeChecking
  :: forall inp. ()
  => TcInstrHandler
  -> TypeCheckedSeq inp
  -> [Un.ExpandedOp]
  -> TypeCheckInstrNoExcept (TypeCheckedSeq inp)
continueTypeChecking tcInstr done rest = case done of
  WellTypedSeq instr -> handleFirst instr
  MixedSeq i e left -> pure (MixedSeq i e (left <> map NonTypedInstr rest))
  IllTypedSeq e left -> pure (IllTypedSeq e (left <> map NonTypedInstr rest))
  where
    handleFirst :: SomeInstr inp -> TypeCheckInstrNoExcept (TypeCheckedSeq inp)
    handleFirst packedInstr@(inputStack :/ instrAndOutputStack) = do
      case instrAndOutputStack of
        instr ::: outputStack -> do
          nextPiece <- typeCheckImplNoLastTypeComment tcInstr rest outputStack
          let combiner = combine inputStack instr
          pure case nextPiece of
            WellTypedSeq nextInstr -> WellTypedSeq (combiner nextInstr)
            MixedSeq nextInstr err left -> MixedSeq (combiner nextInstr) err left
            IllTypedSeq err left -> MixedSeq packedInstr err left
        AnyOutInstr{} -> pure case rest of
          [] -> WellTypedSeq packedInstr
          op : ops -> (MixedSeq
                        packedInstr
                        (TCUnreachableCode (extractOpPos op) (op :| ops))
                        (map NonTypedInstr ops))

    combine inp Nop (_ :/ nextPart) = inp :/ nextPart
    combine inp i1 (_ :/ nextPart) = inp :/ mapSomeInstrOut (Seq i1) nextPart

    extractOpPos :: Un.ExpandedOp -> InstrCallStack
    extractOpPos (Un.WithSrcEx loc _) = loc
    extractOpPos _ = def

-- | Like 'typeCheckImpl' but without the first and the last stack type
-- comments. Useful to reduce duplication of stack type comments.
typeCheckImplStripped
  :: forall inp . Typeable inp
  => TcInstrHandler
  -> [Un.ExpandedOp]
  -> HST inp
  -> TypeCheckInstrNoExcept (TypeCheckedSeq inp)
typeCheckImplStripped tcInstr [] inp
  = typeCheckImplNoLastTypeComment tcInstr [] inp
typeCheckImplStripped tcInstr (op : ops) inp = do
  done <- typeCheckOpImpl tcInstr op inp
  continueTypeChecking tcInstr done ops

typeCheckImpl
  :: forall inp . Typeable inp
  => TcInstrHandler
  -> [Un.ExpandedOp]
  -> HST inp
  -> TypeCheckInstrNoExcept (TypeCheckedSeq inp)
typeCheckImpl tcInstr ops inputStack = do
  tcSeq <- typeCheckImplNoLastTypeComment tcInstr ops inputStack
  mapMSeq appendTypeComment tcSeq
  where
    appendTypeComment packedI@(inp :/ iAndOut) = do
      verbose <- lift (asks tcVerbose)
      pure case (verbose, iAndOut) of
        (True, i ::: out) -> inp :/ Seq i (stackTypeComment out) ::: out
        (True, AnyOutInstr i) -> inp :/ AnyOutInstr (Seq i noStackTypeComment)
        _ -> packedI


prependStackTypeComment
  :: SomeInstr inp -> TypeCheckInstrNoExcept (SomeInstr inp)
prependStackTypeComment packedInstr@(inp :/ _) = do
  verbose <- lift (asks tcVerbose)
  pure if verbose && (not (isNop' packedInstr))
    then mapSomeInstr (Seq (stackTypeComment inp)) packedInstr
    else packedInstr

isNop' :: SomeInstr inp -> Bool
isNop' (_ :/ i ::: _) = isNop i
isNop' (_ :/ AnyOutInstr i) = isNop i

isNop :: Instr inp out -> Bool
isNop (WithLoc _ i) = isNop i
isNop (InstrWithNotes _ i) = isNop i
isNop (InstrWithVarNotes _ i) = isNop i
isNop (FrameInstr _ i) = isNop i
isNop (Seq i1 i2) = isNop i1 && isNop i2
isNop (Nested i) = isNop i
isNop Nop = True
isNop (Ext _) = True
isNop _ = False

mapMSeq
  :: Applicative f
  => (SomeInstr inp -> f (SomeInstr inp'))
  -> TypeCheckedSeq inp
  -> f (TypeCheckedSeq inp')
mapMSeq f v = case v of
  WellTypedSeq instr -> f instr <&> WellTypedSeq
  MixedSeq instr err tail' -> f instr <&> \instr' -> MixedSeq instr' err tail'
  IllTypedSeq err tail' -> pure $ IllTypedSeq err tail'

mapSeq
  :: (SomeInstr inp -> SomeInstr inp')
  -> TypeCheckedSeq inp
  -> TypeCheckedSeq inp'
mapSeq f = runIdentity . mapMSeq (Identity . f)

stackTypeComment :: HST st -> Instr st st
stackTypeComment = Ext . COMMENT_ITEM . StackTypeComment . Just . hstToTs

noStackTypeComment :: Instr st st
noStackTypeComment = Ext (COMMENT_ITEM (StackTypeComment Nothing))

wrapWithLoc :: InstrCallStack -> TypeCheckedSeq inp -> TypeCheckedSeq inp
wrapWithLoc loc = mapSeq $ \someInstr -> case someInstr of
  (_ :/ WithLoc{} ::: _) -> someInstr
  (inp :/ instr ::: out) -> inp :/ WithLoc loc instr ::: out
  (inp :/ AnyOutInstr instr) -> inp :/ (AnyOutInstr $ WithLoc loc instr)

-- | Check whether given types are structurally equal and annotations converge.
matchTypes
  :: forall t1 t2.
      (Each '[KnownT] [t1, t2])
  => Notes t1 -> Notes t2 -> Either TCTypeError (t1 :~: t2, Notes t1)
matchTypes n1 n2 = do
  Refl <- eqType @t1 @t2
  nr <- converge n1 n2 `onFirst` AnnError
  return (Refl, nr)

--------------------------------------------
-- Some generic instruction implementation
--------------------------------------------

-- | Generic implementation for MEMeration
memImpl
  :: forall c memKey rs inp m .
    ( MemOp c
    , KnownT (MemOpKey c)
    , inp ~ (memKey : c : rs)
    , MonadReader InstrCallStack m
    , MonadError TCError m
    )
  => Notes (MemOpKey c)
  -> HST inp
  -> VarAnn
  -> m (SomeInstr inp)
memImpl cKeyNotes inputHST@(hst0 ::& _ ::& hstTail) varAnn =
  case eqType @memKey  @(MemOpKey c) of
    Right Refl -> do
      _ <- onTypeCheckInstrAnnErr uInstr inputHST
        (Just ContainerKeyType) (converge memKeyNotes cKeyNotes)
      pure $ inputHST :/
        MEM ::: ((starNotes, Dict, varAnn) ::& hstTail)
    Left m ->
      typeCheckInstrErr' uInstr (SomeHST inputHST) (Just ContainerKeyType) m
  where
    (memKeyNotes, Dict, _) = hst0
    uInstr = Un.MEM varAnn

getImpl
  :: forall c getKey rs inp m .
    ( GetOp c, KnownT (GetOpKey c)
    , WellTyped (GetOpVal c)
    , inp ~ (getKey : c : rs)
    , MonadReader InstrCallStack m
    , MonadError TCError m
    )
  => Notes (GetOpKey c)
  -> HST inp
  -> Notes (GetOpVal c)
  -> VarAnn
  -> m (SomeInstr inp)
getImpl notesKeyC inputHST@(hst0 ::& _ ::& hstTail) valueNotes varAnn =
  case eqType @getKey @(GetOpKey c) of
    Right Refl -> do
      _ <- onTypeCheckInstrAnnErr uInstr inputHST
        (Just ContainerKeyType) (converge getKeyNotes notesKeyC)
      pure $ inputHST :/
        GET ::: ((NTOption def valueNotes, Dict, varAnn) ::& hstTail)
    Left m ->
      typeCheckInstrErr' uInstr (SomeHST inputHST) (Just ContainerKeyType) m
  where
    (getKeyNotes, Dict, _) = hst0
    uInstr = Un.GET varAnn

updImpl
  :: forall c updKey updParams rs inp m .
    ( UpdOp c
    , KnownT (UpdOpKey c), KnownT (UpdOpParams c)
    , inp ~ (updKey : updParams : c : rs)
    , MonadReader InstrCallStack m
    , MonadError TCError m
    )
  => Notes (UpdOpKey c)
  -> HST inp
  -> Notes (UpdOpParams c)
  -> VarAnn
  -> m (SomeInstr inp)
updImpl cKeyNotes inputHST@(hst0 ::& hst1 ::& cTuple ::& hstTail) cValueNotes varAnn =
  case (eqType @updKey @(UpdOpKey c), eqType @updParams @(UpdOpParams c)) of
    (Right Refl, Right Refl) -> do
      _ <- onTypeCheckInstrAnnErr uInstr inputHST
        (Just ContainerKeyType) (converge updKeyNotes cKeyNotes)
      _ <- onTypeCheckInstrAnnErr uInstr inputHST
        (Just ContainerValueType) (converge updValueNotes cValueNotes)
      pure $ inputHST :/
        UPDATE ::: ((cTuple & _3 .~ varAnn) ::& hstTail)
    (Left m, _) ->
      typeCheckInstrErr' uInstr (SomeHST inputHST) (Just ContainerKeyType) m
    (_, Left m) ->
      typeCheckInstrErr' uInstr (SomeHST inputHST) (Just ContainerValueType) m
  where
    (updKeyNotes, Dict, _) = hst0
    (updValueNotes, Dict, _) = hst1
    uInstr = Un.UPDATE varAnn

sizeImpl
  :: (SizeOp c, inp ~ (c ': rs), Monad m)
  => HST inp
  -> VarAnn
  -> m (SomeInstr inp)
sizeImpl i@(_ ::& rs) vn =
  pure $ i :/ SIZE ::: ((starNotes, Dict, vn) ::& rs)

sliceImpl
  :: (SliceOp c, Typeable c, inp ~ ('TNat ': 'TNat ': c ': rs), Monad m)
  => HST inp
  -> Un.VarAnn
  -> m (SomeInstr inp)
sliceImpl i@(_ ::& _ ::& (cn, Dict, cvn) ::& rs) vn = do
  let vn' = vn `orAnn` deriveVN "slice" cvn
      rn = NTOption def cn
  pure $ i :/ SLICE ::: ((rn, Dict, vn') ::& rs)

concatImpl'
  :: (ConcatOp c, WellTyped c, inp ~ ('TList c : rs), Monad m)
  => HST inp
  -> Un.VarAnn
  -> m (SomeInstr inp)
concatImpl' i@((NTList _ n, Dict, _) ::& rs) vn = do
  pure $ i :/ CONCAT' ::: ((n, Dict, vn) ::& rs)

concatImpl
  :: ( ConcatOp c, inp ~ (c ': c ': rs)
     , WellTyped c
     , MonadReader InstrCallStack m
     , MonadError TCError m
     )
  => HST inp
  -> Un.VarAnn
  -> m (SomeInstr inp)
concatImpl i@((cn1, _, _) ::& (cn2, _, _) ::& rs) vn = do
  cn <- onTypeCheckInstrAnnErr (Un.CONCAT vn) i (Just ConcatArgument) (converge cn1 cn2)
  pure $ i :/ CONCAT ::: ((cn, Dict, vn) ::& rs)

-- | Helper function to construct instructions for binary arithmetic
-- operations.
arithImpl
  :: forall aop inp m n s t.
     ( ArithOp aop n m
     , Typeable (ArithRes aop n m ': s)
     , WellTyped (ArithRes aop n m)
     , inp ~ (n ': m ': s)
     , MonadReader InstrCallStack t
     , MonadError TCError t
     )
  => Instr inp (ArithRes aop n m ': s)
  -> HST inp
  -> VarAnn
  -> Un.ExpandedInstr
  -> t (SomeInstr inp)
arithImpl mkInstr i@((an, _, _) ::& (bn, _, _) ::& rs) vn uInstr = do
  case convergeArith (Proxy @aop) an bn of
    Right cn ->
      pure $ i :/ mkInstr ::: ((cn, Dict, vn) ::& rs)
    Left err -> do
      typeCheckInstrErr' uInstr (SomeHST i) (Just ArithmeticOperation) $ AnnError err

addImpl
  :: forall a b inp rs m.
     ( Typeable rs
     , Each '[KnownT] [a, b]
     , inp ~ (a ': b ': rs)
     , MonadReader InstrCallStack m
     , MonadError TCError m
     )
  => Sing a -> Sing b
  -> HST inp
  -> VarAnn
  -> Un.ExpandedInstr
  -> m (SomeInstr inp)
addImpl t1 t2 = case (t1, t2) of
  (STInt, STInt) -> arithImpl @Add ADD
  (STInt, STNat) -> arithImpl @Add ADD
  (STNat, STInt) -> arithImpl @Add ADD
  (STNat, STNat) -> arithImpl @Add ADD
  (STInt, STTimestamp) -> arithImpl @Add ADD
  (STTimestamp, STInt) -> arithImpl @Add ADD
  (STMutez, STMutez) -> arithImpl @Add ADD
  _ -> \i _ uInstr -> typeCheckInstrErr' uInstr (SomeHST i) (Just ArithmeticOperation) $
    NotNumericTypes (demote @a) (demote @b)

edivImpl
  :: forall a b inp rs m.
     ( Typeable rs
     , Each '[KnownT] [a, b]
     , inp ~ (a ': b ': rs)
     , MonadReader InstrCallStack m
     , MonadError TCError m
     )
  => Sing a -> Sing b
  -> HST inp
  -> VarAnn
  -> Un.ExpandedInstr
  -> m (SomeInstr inp)
edivImpl t1 t2 = case (t1, t2) of
  (STInt, STInt) -> edivImplDo
  (STInt, STNat) -> edivImplDo
  (STNat, STInt) -> edivImplDo
  (STNat, STNat) -> edivImplDo
  (STMutez, STMutez) -> edivImplDo
  (STMutez, STNat) -> edivImplDo
  _ -> \i _ uInstr -> typeCheckInstrErr' uInstr (SomeHST i) (Just ArithmeticOperation) $
    NotNumericTypes (demote @a) (demote @b)

edivImplDo
  :: ( EDivOp n m
     , WellTyped (EModOpRes n m)
     , WellTyped (EDivOpRes n m)
     , inp ~ (n ': m ': s)
     , MonadReader InstrCallStack t
     , MonadError TCError t
     )
  => HST inp
  -> VarAnn
  -> Un.ExpandedInstr
  -> t (SomeInstr inp)
edivImplDo i@((an, _, _) ::& (bn, _, _) ::& rs) vn uInstr = do
  case convergeEDiv an bn of
    Right cn ->
      pure $ i :/ EDIV ::: ((cn, Dict, vn) ::& rs)
    Left err -> do
      typeCheckInstrErr' uInstr (SomeHST i) (Just ArithmeticOperation) $ AnnError err

subImpl
  :: forall a b inp rs m.
     ( Typeable rs
     , Each '[KnownT] [a, b]
     , inp ~ (a ': b ': rs)
     , MonadReader InstrCallStack m
     , MonadError TCError m
     )
  => Sing a -> Sing b
  -> HST inp
  -> VarAnn
  -> Un.ExpandedInstr
  -> m (SomeInstr inp)
subImpl t1 t2 = case (t1, t2) of
  (STInt, STInt) -> arithImpl @Sub SUB
  (STInt, STNat) -> arithImpl @Sub SUB
  (STNat, STInt) -> arithImpl @Sub SUB
  (STNat, STNat) -> arithImpl @Sub SUB
  (STTimestamp, STTimestamp) -> arithImpl @Sub SUB
  (STTimestamp, STInt) -> arithImpl @Sub SUB
  (STMutez, STMutez) -> arithImpl @Sub SUB
  _ -> \i _ uInstr -> typeCheckInstrErr' uInstr (SomeHST i) (Just ArithmeticOperation) $
    NotNumericTypes (demote @a) (demote @b)

mulImpl
  :: forall a b inp rs m.
     ( Typeable rs
     , Each '[KnownT] [a, b]
     , inp ~ (a ': b ': rs)
     , MonadReader InstrCallStack m
     , MonadError TCError m
     )
  => Sing a -> Sing b
  -> HST inp
  -> VarAnn
  -> Un.ExpandedInstr
  -> m (SomeInstr inp)
mulImpl t1 t2 = case (t1, t2) of
  (STInt, STInt) -> arithImpl @Mul MUL
  (STInt, STNat) -> arithImpl @Mul MUL
  (STNat, STInt) -> arithImpl @Mul MUL
  (STNat, STNat) -> arithImpl @Mul MUL
  (STNat, STMutez) -> arithImpl @Mul MUL
  (STMutez, STNat) -> arithImpl @Mul MUL
  _ -> \i _ uInstr -> typeCheckInstrErr' uInstr (SomeHST i) (Just ArithmeticOperation) $
    NotNumericTypes (demote @a) (demote @b)

-- | Helper function to construct instructions for unary arithmetic
-- operations.
unaryArithImpl
  :: ( Typeable (UnaryArithRes aop n ': s)
     , WellTyped (UnaryArithRes aop n)
     , inp ~ (n ': s)
     , Monad t
     )
  => Instr inp (UnaryArithRes aop n ': s)
  -> HST inp
  -> VarAnn
  -> t (SomeInstr inp)
unaryArithImpl mkInstr i@(_ ::& rs) vn = do
  pure $ i :/ mkInstr ::: ((starNotes, Dict, vn) ::& rs)

-- | Helper function to construct instructions for unary arithmetic
-- operations that should preserve annotations.
unaryArithImplAnnotated
  :: ( Typeable (UnaryArithRes aop n ': s)
     , WellTyped (UnaryArithRes aop n)
     , inp ~ (n ': s)
     , Monad t
     , n ~ UnaryArithRes aop n
     )
  => Instr inp (UnaryArithRes aop n ': s)
  -> HST inp
  -> VarAnn
  -> t (SomeInstr inp)
unaryArithImplAnnotated mkInstr i@((n, _, _) ::& rs) vn = do
  pure $ i :/ mkInstr ::: ((n, Dict, vn) ::& rs)