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

module Morley.Michelson.TypeCheck.Helpers
    ( hstToTs
    , eqHST
    , eqHST1
    , lengthHST

    , ensureDistinctAsc
    , handleError
    , eqType
    , onTypeCheckInstrErr
    , onScopeCheckInstrErr
    , typeCheckInstrErr
    , typeCheckInstrErr'
    , typeCheckImpl
    , typeCheckImplStripped

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

import Prelude hiding (EQ, GT, LT)

import Control.Monad.Except (MonadError, catchError, throwError)
import Data.Constraint (Dict(..), withDict)
import Data.Default (def)
import Data.Singletons (Sing, SingI(sing), demote, fromSing)
import Data.Singletons.Decide ((:~:)(Refl))
import Fmt (Buildable, (+|), (|+))

import Morley.Michelson.ErrorPos (ErrorSrcPos)
import Morley.Michelson.TypeCheck.Error (TCError(..), TCTypeError(..), TypeContext(..))
import Morley.Michelson.TypeCheck.TypeCheck
import Morley.Michelson.TypeCheck.TypeCheckedSeq (IllTypedInstr(..), TypeCheckedSeq(..))
import Morley.Michelson.TypeCheck.Types
import Morley.Michelson.Typed
  (BadTypeForScope(..), CommentType(StackTypeComment), Comparable, ExtInstr(COMMENT_ITEM),
  Instr(..), SingT(..), T(..), WellTyped, getComparableProofS, requireEq)
import Morley.Michelson.Typed.Annotation
import Morley.Michelson.Typed.Arith (Add, ArithOp(..), EDiv, Mul, Sub, UnaryArithOp(..))
import Morley.Michelson.Typed.Polymorphic
  (ConcatOp, EDivOp(..), GetOp(..), MemOp(..), SizeOp, SliceOp, UpdOp(..))
import Morley.Michelson.Untyped qualified as Un
import Morley.Michelson.Untyped.Annotation (VarAnn)
import Morley.Util.MismatchError
import Morley.Util.Sing (eqI)

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

-- | Check whether the given stack types are equal.
eqHST
  :: forall as bs. (SingI as, SingI bs)
  => HST as -> HST bs -> Either TCTypeError (as :~: bs)
eqHST (hst :: HST xs) (hst' :: HST ys) = do
  case eqI @as @bs of
    Nothing -> Left $ StackEqError MkMismatchError {meActual = hstToTs hst, meExpected = hstToTs hst'}
    Just Refl -> Right 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. (SingI st, WellTyped t)
  => HST st -> Either TCTypeError (st :~: '[t])
eqHST1 hst = do
  let hst' = sing @t -:& SNil
  case eqI @'[t] @st of
    Nothing -> Left $ StackEqError MkMismatchError {meActual = hstToTs hst, meExpected = 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, Buildable 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

-- | Flipped version of 'catchError'.
handleError :: MonadError e m => (e -> m a) -> m a -> m a
handleError = flip catchError

-- | Function @eqType@ is a simple wrapper around @Data.Singletons.decideEquality@ suited
-- for use within @Either TCTypeError a@ applicative.
eqType
  :: forall (a :: T) (b :: T). (Each '[SingI] [a, b])
  => Either TCTypeError (a :~: b)
eqType =
  requireEq @a @b $ Left ... TypeEqError

onTypeCheckInstrErr
  :: (MonadReader TypeCheckInstrEnv 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 TypeCheckInstrEnv 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 <- view tcieErrorPos
    throwError $ TCFailedOnInstr instr hst pos mContext $
      Just $ UnsupportedTypeForScope (demote @t) e

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

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

withCompareableCheck
  :: forall a m v ts. (SingI ts, MonadReader TypeCheckInstrEnv 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. SingI inp
  => TcInstrHandler
  -> Un.ExpandedOp
  -> HST inp
  -> TypeCheckInstrNoExcept (TypeCheckedSeq inp)
typeCheckOpImpl tcInstr op' hst = case op' of
  Un.WithSrcEx _ op@Un.WithSrcEx{} -> typeCheckOpImpl tcInstr op hst
  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 :: ErrorSrcPos -> Un.ExpandedInstr -> TypeCheckInstrNoExcept (TypeCheckedSeq inp)
    typeCheckPrimWithLoc loc op = local (set tcieErrorPos loc)
      (wrapWithLoc loc <$> typeCheckPrim op)

    typeCheckPrim :: Un.ExpandedInstr -> TypeCheckInstrNoExcept (TypeCheckedSeq inp)
    typeCheckPrim op = tcInstr op hst

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

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

-- | Like 'typeCheckImpl' but doesn't add a stack type comment after the
-- sequence.
typeCheckImplNoLastTypeComment
  :: forall inp . SingI 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 -> ErrorSrcPos
    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 . SingI 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 . SingI 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 <- 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 <- 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 (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 :: ErrorSrcPos -> 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)

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

-- | Generic implementation for MEMeration
memImpl
  :: forall c memKey rs inp m .
    ( MemOp c
    , SingI (MemOpKey c)
    , inp ~ (memKey : c : rs)
    , SingI rs
    , MonadReader TypeCheckInstrEnv m
    , MonadError TCError m
    )
  => HST inp
  -> VarAnn
  -> m (SomeInstr inp)
memImpl inputHST@(_ ::& _ ::& hstTail) varAnn =
  case eqType @memKey @(MemOpKey c) of
    Right Refl ->
      pure $ inputHST :/
        AnnMEM (Anns1 varAnn) ::: ((sing, Dict) ::& hstTail)
    Left m ->
      typeCheckInstrErr' uInstr (SomeHST inputHST) (Just ContainerKeyType) m
  where
    uInstr = Un.MEM varAnn

getImpl
  :: forall c getKey rs inp m .
    ( GetOp c, SingI (GetOpKey c)
    , WellTyped (GetOpVal c)
    , inp ~ (getKey : c : rs)
    , SingI rs
    , MonadReader TypeCheckInstrEnv m
    , MonadError TCError m
    )
  => HST inp
  -> SingT (GetOpVal c)
  -> VarAnn
  -> m (SomeInstr inp)
getImpl inputHST@(_ ::& _ ::& hstTail) valueSing varAnn =
  case eqType @getKey @(GetOpKey c) of
    Right Refl ->
      pure $ inputHST :/
        AnnGET (Anns1 varAnn) ::: ((STOption valueSing, Dict) ::& hstTail)
    Left m ->
      typeCheckInstrErr' uInstr (SomeHST inputHST) (Just ContainerKeyType) m
  where
    uInstr = Un.GET varAnn

updImpl
  :: forall c updKey updParams rs inp m .
    ( UpdOp c
    , SingI (UpdOpKey c), SingI (UpdOpParams c)
    , SingI rs
    , inp ~ (updKey : updParams : c : rs)
    , MonadReader TypeCheckInstrEnv m
    , MonadError TCError m
    )
  => HST inp
  -> VarAnn
  -> m (SomeInstr inp)
updImpl inputHST@(_ ::& _ ::& cTuple ::& hstTail) varAnn =
  case (eqType @updKey @(UpdOpKey c), eqType @updParams @(UpdOpParams c)) of
    (Right Refl, Right Refl) ->
      pure $ inputHST :/
        AnnUPDATE (Anns1 varAnn) ::: (cTuple ::& hstTail)
    (Left m, _) ->
      typeCheckInstrErr' uInstr (SomeHST inputHST) (Just ContainerKeyType) m
    (_, Left m) ->
      typeCheckInstrErr' uInstr (SomeHST inputHST) (Just ContainerValueType) m
  where
    uInstr = Un.UPDATE varAnn

getUpdImpl
  :: forall c updKey updParams rs inp m .
    ( UpdOp c, GetOp c
    , SingI (UpdOpKey c)
    , SingI (GetOpVal c)
    , inp ~ (updKey : updParams : c : rs)
    , SingI rs
    , GetOpKey c ~ UpdOpKey c
    , UpdOpParams c ~ 'TOption (GetOpVal c)
    , MonadReader TypeCheckInstrEnv m
    , MonadError TCError m
    )
  => HST inp
  -> VarAnn
  -> m (SomeInstr inp)
getUpdImpl inputHST@(_ ::& hst1 ::& cTuple ::& hstTail) varAnn =
  case (eqType @updKey @(UpdOpKey c), eqType @updParams @(UpdOpParams c)) of
    (Right Refl, Right Refl) ->
      pure $ inputHST :/
        AnnGET_AND_UPDATE (Anns1 varAnn) ::: (hst1 ::& cTuple ::& hstTail)
    (Left m, _) ->
      typeCheckInstrErr' uInstr (SomeHST inputHST) (Just ContainerKeyType) m
    (_, Left m) ->
      typeCheckInstrErr' uInstr (SomeHST inputHST) (Just ContainerValueType) m
  where
    uInstr = Un.GET_AND_UPDATE varAnn

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

sliceImpl
  :: (SliceOp c, inp ~ ('TNat ': 'TNat ': c ': rs), Monad m)
  => HST inp
  -> Un.VarAnn
  -> m (SomeInstr inp)
sliceImpl i@(_ ::& _ ::& (cn, Dict) ::& rs) vn = do
  let rn = STOption cn
  pure $ i :/ AnnSLICE (Anns1 vn) ::: ((rn, Dict) ::& rs)

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

concatImpl
  :: ( ConcatOp c, inp ~ (c ': c ': rs)
     , WellTyped c
     , MonadReader TypeCheckInstrEnv m
     )
  => HST inp
  -> Un.VarAnn
  -> m (SomeInstr inp)
concatImpl i@((cn1, _) ::& (_, _) ::& rs) vn = do
  pure $ i :/ AnnCONCAT (Anns1 vn) ::: ((cn1, Dict) ::& rs)

-- | Helper function to construct instructions for binary arithmetic
-- operations.
arithImpl
  :: forall aop inp m n s t.
     ( WellTyped (ArithRes aop n m)
     , inp ~ (n ': m ': s)
     , MonadReader TypeCheckInstrEnv t
     )
  => (Anns '[VarAnn] -> Instr inp (ArithRes aop n m ': s))
  -> HST inp
  -> VarAnn
  -> Un.ExpandedInstr
  -> t (SomeInstr inp)
arithImpl mkInstr i@(_ ::& _ ::& rs) vn _ =
  pure $ i :/ mkInstr (Anns1 vn) ::: ((sing, Dict) ::& rs)

addImpl
  :: forall a b inp rs m.
     ( Each '[SingI] [a, b]
     , inp ~ (a ': b ': rs)
     , SingI rs
     , MonadReader TypeCheckInstrEnv 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 AnnADD
  (STInt, STNat) -> arithImpl @Add AnnADD
  (STNat, STInt) -> arithImpl @Add AnnADD
  (STNat, STNat) -> arithImpl @Add AnnADD
  (STInt, STTimestamp) -> arithImpl @Add AnnADD
  (STTimestamp, STInt) -> arithImpl @Add AnnADD
  (STMutez, STMutez) -> arithImpl @Add AnnADD
  (STBls12381Fr, STBls12381Fr) -> arithImpl @Add AnnADD
  (STBls12381G1, STBls12381G1) -> arithImpl @Add AnnADD
  (STBls12381G2, STBls12381G2) -> arithImpl @Add AnnADD
  _ -> \i _ uInstr -> typeCheckInstrErr' uInstr (SomeHST i) (Just ArithmeticOperation) $
    NotNumericTypes (demote @a) (demote @b)

edivImpl
  :: forall a b inp rs m.
     ( SingI rs
     , Each '[SingI] [a, b]
     , inp ~ (a ': b ': rs)
     , MonadReader TypeCheckInstrEnv 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
  :: ( ArithOp EDiv n m
     , ArithRes EDiv n m ~ 'TOption ('TPair (EDivOpRes n m) (EModOpRes n m))
     , WellTyped (EModOpRes n m)
     , WellTyped (EDivOpRes n m)
     , inp ~ (n ': m ': s)
     , MonadReader TypeCheckInstrEnv t
     )
  => HST inp
  -> VarAnn
  -> Un.ExpandedInstr
  -> t (SomeInstr inp)
edivImplDo i@(_ ::& _ ::& rs) vn _ =
  pure $ i :/ AnnEDIV (Anns1 vn) ::: ((sing, Dict) ::& rs)

subImpl
  :: forall a b inp rs m.
     ( Each '[SingI] [a, b]
     , inp ~ (a ': b ': rs)
     , SingI rs
     , MonadReader TypeCheckInstrEnv 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 AnnSUB
  (STInt, STNat) -> arithImpl @Sub AnnSUB
  (STNat, STInt) -> arithImpl @Sub AnnSUB
  (STNat, STNat) -> arithImpl @Sub AnnSUB
  (STTimestamp, STTimestamp) -> arithImpl @Sub AnnSUB
  (STMutez, STMutez) -> \i _ uInstr -> typeCheckInstrErr' uInstr (SomeHST i) (Just ArithmeticOperation) $
    InvalidInstruction uInstr "Use of SUB on `mutez` operands is deprecated; use SUB_MUTEZ"
  (STTimestamp, STInt) -> arithImpl @Sub AnnSUB
  _ -> \i _ uInstr -> typeCheckInstrErr' uInstr (SomeHST i) (Just ArithmeticOperation) $
    NotNumericTypes (demote @a) (demote @b)

mulImpl
  :: forall a b inp rs m.
     ( Each '[SingI] [a, b]
     , inp ~ (a ': b ': rs)
     , SingI rs
     , MonadReader TypeCheckInstrEnv 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 AnnMUL
  (STInt, STNat) -> arithImpl @Mul AnnMUL
  (STNat, STInt) -> arithImpl @Mul AnnMUL
  (STNat, STNat) -> arithImpl @Mul AnnMUL
  (STNat, STMutez) -> arithImpl @Mul AnnMUL
  (STMutez, STNat) -> arithImpl @Mul AnnMUL
  (STInt, STBls12381Fr) -> arithImpl @Mul AnnMUL
  (STNat, STBls12381Fr) -> arithImpl @Mul AnnMUL
  (STBls12381Fr, STInt) -> arithImpl @Mul AnnMUL
  (STBls12381Fr, STNat) -> arithImpl @Mul AnnMUL
  (STBls12381Fr, STBls12381Fr) -> arithImpl @Mul AnnMUL
  (STBls12381G1, STBls12381Fr) -> arithImpl @Mul AnnMUL
  (STBls12381G2, STBls12381Fr) -> arithImpl @Mul AnnMUL
  _ -> \i _ uInstr -> typeCheckInstrErr' uInstr (SomeHST i) (Just ArithmeticOperation) $
    NotNumericTypes (demote @a) (demote @b)

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

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