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

-- | By default we represent error tags using strings. This module
-- makes it possible to use numbers instead. It introduces new [error format].
--
-- There are two possible ways to use it:
-- 1. If you have just one Lorentz instruction (potentially a big one),
-- just use 'useNumericErrors' function. It will change error representation
-- there and return a map that can be used to interpret new error codes.
-- 2. If your contract consists of multiple parts, start with gathering all
-- error tags ('gatherErrorTags'). Then build 'ErrorTagMap' using
-- 'addNewErrorTags'. Pass empty map if you are building from scratch
-- (you can use 'buildErrorTagMap' shortcut) or an existing
-- map if you have one (e. g. you are upgrading a contract).

module Lorentz.Errors.Numeric.Contract
  ( ErrorTagMap
  , ErrorTagExclusions
  , gatherErrorTags
  , addNewErrorTags
  , buildErrorTagMap
  , excludeErrorTags
  , applyErrorTagMap
  , applyErrorTagMapWithExclusions
  , useNumericErrors

  , errorFromValNumeric
  , errorToValNumeric
  ) where

import Data.Bimap (Bimap)
import Data.Bimap qualified as Bimap
import Data.Default (def)
import Data.HashSet qualified as HS
import Data.Singletons (withSingI)
import Fmt (pretty)

import Lorentz.Base
import Lorentz.Errors
import Morley.Michelson.Analyzer
import Morley.Michelson.FailPattern
import Morley.Michelson.Text (MText)
import Morley.Michelson.Typed

-- | This is a bidirectional map with correspondence between numeric
-- and textual error tags.
type ErrorTagMap = Bimap Natural MText

-- | Tags excluded from map.
type ErrorTagExclusions = HashSet MText

-- | Find all textual error tags that are used in typical
-- @FAILWITH@ patterns within given instruction.
-- Map them to natural numbers.
gatherErrorTags :: inp :-> out -> HashSet MText
gatherErrorTags :: forall (inp :: [*]) (out :: [*]). (inp :-> out) -> HashSet MText
gatherErrorTags = HashMap MText () -> HashSet MText
forall a. HashMap a () -> HashSet a
HS.fromMap (HashMap MText () -> HashSet MText)
-> ((inp :-> out) -> HashMap MText ())
-> (inp :-> out)
-> HashSet MText
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HashMap MText Word -> HashMap MText ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (HashMap MText Word -> HashMap MText ())
-> ((inp :-> out) -> HashMap MText Word)
-> (inp :-> out)
-> HashMap MText ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AnalyzerRes -> HashMap MText Word
arErrorTags (AnalyzerRes -> HashMap MText Word)
-> ((inp :-> out) -> AnalyzerRes)
-> (inp :-> out)
-> HashMap MText Word
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Instr (ToTs inp) (ToTs out) -> AnalyzerRes
forall (inp :: [T]) (out :: [T]). Instr inp out -> AnalyzerRes
analyze (Instr (ToTs inp) (ToTs out) -> AnalyzerRes)
-> ((inp :-> out) -> Instr (ToTs inp) (ToTs out))
-> (inp :-> out)
-> AnalyzerRes
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (inp :-> out) -> Instr (ToTs inp) (ToTs out)
forall (inp :: [*]) (out :: [*]).
(inp :-> out) -> Instr (ToTs inp) (ToTs out)
iAnyCode

-- | Add more error tags to an existing 'ErrorTagMap'. It is useful when
-- your contract consists of multiple parts (e. g. in case of contract
-- upgrade), you have existing map for some part and want to add tags
-- from another part to it.
-- You can pass empty map as existing one if you just want to build
-- 'ErrorTagMap' from a set of textual tags. See 'buildErrorTagMap'.
addNewErrorTags :: ErrorTagMap -> HashSet MText -> ErrorTagMap
addNewErrorTags :: ErrorTagMap -> HashSet MText -> ErrorTagMap
addNewErrorTags ErrorTagMap
existingMap HashSet MText
newTags =
  (ErrorTagMap -> Element [(Natural, MText)] -> ErrorTagMap)
-> ErrorTagMap -> [(Natural, MText)] -> ErrorTagMap
forall t b. Container t => (b -> Element t -> b) -> b -> t -> b
foldl' (((Natural, MText) -> ErrorTagMap -> ErrorTagMap)
-> ErrorTagMap -> (Natural, MText) -> ErrorTagMap
forall a b c. (a -> b -> c) -> b -> a -> c
flip (((Natural, MText) -> ErrorTagMap -> ErrorTagMap)
 -> ErrorTagMap -> (Natural, MText) -> ErrorTagMap)
-> ((Natural, MText) -> ErrorTagMap -> ErrorTagMap)
-> ErrorTagMap
-> (Natural, MText)
-> ErrorTagMap
forall a b. (a -> b) -> a -> b
$ (Natural -> MText -> ErrorTagMap -> ErrorTagMap)
-> (Natural, MText) -> ErrorTagMap -> ErrorTagMap
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Natural -> MText -> ErrorTagMap -> ErrorTagMap
forall a b. (Ord a, Ord b) => a -> b -> Bimap a b -> Bimap a b
Bimap.tryInsert) ErrorTagMap
existingMap [(Natural, MText)]
newItems
  where
    firstUnusedNumeric :: Natural
firstUnusedNumeric
      | ErrorTagMap -> Bool
forall a b. Bimap a b -> Bool
Bimap.null ErrorTagMap
existingMap = Natural
0
      | Bool
otherwise = (Natural, MText) -> Natural
forall a b. (a, b) -> a
fst (ErrorTagMap -> (Natural, MText)
forall a b. Bimap a b -> (a, b)
Bimap.findMax ErrorTagMap
existingMap) Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
+ Natural
1

    newItems :: [(Natural, MText)]
    newItems :: [(Natural, MText)]
newItems = [Natural] -> [MText] -> [(Natural, MText)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Natural
firstUnusedNumeric .. ] (HashSet MText -> [Element (HashSet MText)]
forall t. Container t => t -> [Element t]
toList HashSet MText
newTags)

-- | Build 'ErrorTagMap' from a set of textual tags.
buildErrorTagMap :: HashSet MText -> ErrorTagMap
buildErrorTagMap :: HashSet MText -> ErrorTagMap
buildErrorTagMap = ErrorTagMap -> HashSet MText -> ErrorTagMap
addNewErrorTags ErrorTagMap
forall a b. Bimap a b
Bimap.empty

-- | Remove some error tags from map.
-- This way you say to remain these string tags intact, while others will be
-- converted to numbers when this map is applied.
--
-- Note that later you have to apply this map using
-- 'applyErrorTagMapWithExclusions', otherwise an error would be raised.
excludeErrorTags
  :: HasCallStack
  => ErrorTagExclusions -> ErrorTagMap -> ErrorTagMap
excludeErrorTags :: HasCallStack => HashSet MText -> ErrorTagMap -> ErrorTagMap
excludeErrorTags HashSet MText
toExclude ErrorTagMap
errMap =
  (ErrorTagMap -> Element (HashSet MText) -> ErrorTagMap)
-> ErrorTagMap -> HashSet MText -> ErrorTagMap
forall t b. Container t => (b -> Element t -> b) -> b -> t -> b
foldl' ((MText -> ErrorTagMap -> ErrorTagMap)
-> ErrorTagMap -> MText -> ErrorTagMap
forall a b c. (a -> b -> c) -> b -> a -> c
flip MText -> ErrorTagMap -> ErrorTagMap
forall {a} {a}.
(Ord a, Ord a, Buildable a) =>
a -> Bimap a a -> Bimap a a
deleteExistingR) ErrorTagMap
errMap HashSet MText
toExclude
  where
    deleteExistingR :: a -> Bimap a a -> Bimap a a
deleteExistingR a
k Bimap a a
m = case a -> Bimap a a -> Maybe a
forall a b (m :: * -> *).
(Ord a, Ord b, MonadThrow m) =>
b -> Bimap a b -> m a
Bimap.lookupR a
k Bimap a a
m of
      Just a
_ -> a -> Bimap a a -> Bimap a a
forall a b. (Ord a, Ord b) => b -> Bimap a b -> Bimap a b
Bimap.deleteR a
k Bimap a a
m
      Maybe a
Nothing ->
        Text -> Bimap a a
forall a. HasCallStack => Text -> a
error (Text -> Bimap a a) -> Text -> Bimap a a
forall a b. (a -> b) -> a -> b
$ Text
"Tag " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> a -> Text
forall a b. (Buildable a, FromBuilder b) => a -> b
pretty a
k Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" does not appear in the contract"

-- | For each typical 'FAILWITH' that uses a string to represent error
-- tag this function changes error tag to be a number using the
-- supplied conversion map.
-- It assumes that supplied map contains all such strings
-- (and will error out if it does not).
-- It will always be the case if you gather all error tags using
-- 'gatherErrorTags' and build 'ErrorTagMap' from them using 'addNewErrorTags'.
applyErrorTagMap :: HasCallStack => ErrorTagMap -> inp :-> out -> inp :-> out
applyErrorTagMap :: forall (inp :: [*]) (out :: [*]).
HasCallStack =>
ErrorTagMap -> (inp :-> out) -> inp :-> out
applyErrorTagMap ErrorTagMap
errorTagMap = ErrorTagMap -> HashSet MText -> (inp :-> out) -> inp :-> out
forall (inp :: [*]) (out :: [*]).
HasCallStack =>
ErrorTagMap -> HashSet MText -> (inp :-> out) -> inp :-> out
applyErrorTagMapWithExclusions ErrorTagMap
errorTagMap HashSet MText
forall a. Monoid a => a
mempty

-- | Similar to 'applyErrorTagMap', but for case when you have excluded some
-- tags from map via 'excludeErrorTags'.
-- Needed, because both 'excludeErrorTags' and this function do not tolerate
-- unknown errors in contract code (for your safety).
applyErrorTagMapWithExclusions
  :: HasCallStack
  => ErrorTagMap -> ErrorTagExclusions -> inp :-> out -> inp :-> out
applyErrorTagMapWithExclusions :: forall (inp :: [*]) (out :: [*]).
HasCallStack =>
ErrorTagMap -> HashSet MText -> (inp :-> out) -> inp :-> out
applyErrorTagMapWithExclusions ErrorTagMap
errorTagMap HashSet MText
exclusions =
  (forall (o' :: [T]). Instr (ToTs inp) o' -> Instr (ToTs inp) o')
-> (inp :-> out) -> inp :-> out
forall (i1 :: [*]) (i2 :: [*]) (o :: [*]).
(forall (o' :: [T]). Instr (ToTs i1) o' -> Instr (ToTs i2) o')
-> (i1 :-> o) -> i2 :-> o
iMapAnyCode (ErrorTagMap
-> HashSet MText -> Instr (ToTs inp) o' -> Instr (ToTs inp) o'
forall (inp :: [T]) (out :: [T]).
HasCallStack =>
ErrorTagMap -> HashSet MText -> Instr inp out -> Instr inp out
applyErrorTagMapWithExcT ErrorTagMap
errorTagMap HashSet MText
exclusions)

-- | This function implements the simplest scenario of using this
-- module's functionality:
-- 1. Gather all error tags from a single instruction.
-- 2. Turn them into error conversion map.
-- 3. Apply this conversion.
useNumericErrors ::
  HasCallStack => inp :-> out -> (inp :-> out, ErrorTagMap)
useNumericErrors :: forall (inp :: [*]) (out :: [*]).
HasCallStack =>
(inp :-> out) -> (inp :-> out, ErrorTagMap)
useNumericErrors inp :-> out
instr = (ErrorTagMap -> (inp :-> out) -> inp :-> out
forall (inp :: [*]) (out :: [*]).
HasCallStack =>
ErrorTagMap -> (inp :-> out) -> inp :-> out
applyErrorTagMap ErrorTagMap
errorTagMap inp :-> out
instr, ErrorTagMap
errorTagMap)
  where
    errorTagMap :: ErrorTagMap
errorTagMap = HashSet MText -> ErrorTagMap
buildErrorTagMap (HashSet MText -> ErrorTagMap) -> HashSet MText -> ErrorTagMap
forall a b. (a -> b) -> a -> b
$ (inp :-> out) -> HashSet MText
forall (inp :: [*]) (out :: [*]). (inp :-> out) -> HashSet MText
gatherErrorTags inp :-> out
instr

-- This function works with 'Morley.Michelson.Typed' representation, not with Lorentz.
applyErrorTagMapWithExcT ::
     HasCallStack
  => ErrorTagMap
  -> ErrorTagExclusions
  -> Instr inp out
  -> Instr inp out
applyErrorTagMapWithExcT :: forall (inp :: [T]) (out :: [T]).
HasCallStack =>
ErrorTagMap -> HashSet MText -> Instr inp out -> Instr inp out
applyErrorTagMapWithExcT ErrorTagMap
errorTagMap HashSet MText
exclusions Instr inp out
instr =
  DfsSettings Identity
-> (forall (i :: [T]) (o :: [T]). Instr i o -> Instr i o)
-> Instr inp out
-> Instr inp out
forall (inp :: [T]) (out :: [T]).
DfsSettings Identity
-> (forall (i :: [T]) (o :: [T]). Instr i o -> Instr i o)
-> Instr inp out
-> Instr inp out
dfsModifyInstr DfsSettings Identity
dfsSettings forall (inp :: [T]) (out :: [T]).
HasCallStack =>
Instr inp out -> Instr inp out
forall (i :: [T]) (o :: [T]). Instr i o -> Instr i o
step Instr inp out
instr
  where
    dfsSettings :: DfsSettings Identity
    dfsSettings :: DfsSettings Identity
dfsSettings = DfsSettings Identity
forall a. Default a => a
def
      { dsGoToValues :: Bool
dsGoToValues = Bool
True
      }

    tagToNatValue :: HasCallStack => MText -> SomeConstant
    tagToNatValue :: HasCallStack => MText -> SomeConstant
tagToNatValue MText
tag =
      case (MText -> HashSet MText -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
HS.member MText
tag HashSet MText
exclusions, MText -> ErrorTagMap -> Maybe Natural
forall a b (m :: * -> *).
(Ord a, Ord b, MonadThrow m) =>
b -> Bimap a b -> m a
Bimap.lookupR MText
tag ErrorTagMap
errorTagMap) of
        (Bool
True, Maybe Natural
_) -> Value 'TString -> SomeConstant
forall (t :: T). ConstantScope t => Value t -> SomeConstant
SomeConstant (MText -> Value 'TString
forall (instr :: [T] -> [T] -> *). MText -> Value' instr 'TString
VString MText
tag)
        -- It will be applied to textual tags detected by 'modifyTypicalFailWith'.
        -- Here we assume that all of them are discovered by the analyzer.
        -- If this error ever happens, it means that someone used
        -- 'applyErrorTagMap' with incomplete 'ErrorTagMap' or there is an
        -- internal bug somewhere.
        (Bool
False, Maybe Natural
Nothing) -> Text -> SomeConstant
forall a. HasCallStack => Text -> a
error (Text -> SomeConstant) -> Text -> SomeConstant
forall a b. (a -> b) -> a -> b
$ Text
"Can't find a tag: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> MText -> Text
forall a b. (Buildable a, FromBuilder b) => a -> b
pretty MText
tag
        (Bool
False, Just Natural
n) -> Value 'TNat -> SomeConstant
forall (t :: T). ConstantScope t => Value t -> SomeConstant
SomeConstant (Natural -> Value 'TNat
forall (instr :: [T] -> [T] -> *). Natural -> Value' instr 'TNat
VNat Natural
n)

    step :: HasCallStack => Instr inp out -> Instr inp out
    step :: forall (inp :: [T]) (out :: [T]).
HasCallStack =>
Instr inp out -> Instr inp out
step = (MText -> SomeConstant) -> Instr inp out -> Instr inp out
forall (inp :: [T]) (out :: [T]).
HasCallStack =>
(MText -> SomeConstant) -> Instr inp out -> Instr inp out
modifyTypicalFailWith HasCallStack => MText -> SomeConstant
MText -> SomeConstant
tagToNatValue

-- | If you apply numeric error representation in your contract, 'errorFromVal'
-- will stop working because it doesn't know about this
-- transformation.
-- This function takes this transformation into account.
-- If a number is used as a tag, but it is not found in the passed
-- map, we conservatively preserve that number (because this whole
-- approach is rather a heuristic).
errorFromValNumeric ::
  (SingI t, IsError e) => ErrorTagMap -> Value t -> Either Text e
errorFromValNumeric :: forall (t :: T) e.
(SingI t, IsError e) =>
ErrorTagMap -> Value t -> Either Text e
errorFromValNumeric ErrorTagMap
errorTagMap Value t
v =
  case Value t
v of
    VNat Natural
tag
      | Just MText
textualTag <- Natural -> ErrorTagMap -> Maybe MText
forall a b (m :: * -> *).
(Ord a, Ord b, MonadThrow m) =>
a -> Bimap a b -> m b
Bimap.lookup Natural
tag ErrorTagMap
errorTagMap ->
        Value 'TString -> Either Text e
forall e (t :: T). (IsError e, SingI t) => Value t -> Either Text e
errorFromVal (Value 'TString -> Either Text e)
-> Value 'TString -> Either Text e
forall a b. (a -> b) -> a -> b
$ MText -> Value 'TString
forall (instr :: [T] -> [T] -> *). MText -> Value' instr 'TString
VString MText
textualTag
    (VPair (VNat Natural
tag, Value' Instr r
something) :: Value pair)
      | Just MText
textualTag <- Natural -> ErrorTagMap -> Maybe MText
forall a b (m :: * -> *).
(Ord a, Ord b, MonadThrow m) =>
a -> Bimap a b -> m b
Bimap.lookup Natural
tag ErrorTagMap
errorTagMap ->
        case forall {k} (a :: k). SingI a => Sing a
forall (a :: T). SingI a => Sing a
sing @pair of
          STPair Sing n1
l Sing n2
r -> Sing 'TNat -> (SingI 'TNat => Either Text e) -> Either Text e
forall {k} (n :: k) r. Sing n -> (SingI n => r) -> r
withSingI Sing n1
Sing 'TNat
l ((SingI 'TNat => Either Text e) -> Either Text e)
-> (SingI 'TNat => Either Text e) -> Either Text e
forall a b. (a -> b) -> a -> b
$ Sing r -> (SingI r => Either Text e) -> Either Text e
forall {k} (n :: k) r. Sing n -> (SingI n => r) -> r
withSingI Sing r
Sing n2
r ((SingI r => Either Text e) -> Either Text e)
-> (SingI r => Either Text e) -> Either Text e
forall a b. (a -> b) -> a -> b
$
            Value ('TPair 'TString r) -> Either Text e
forall e (t :: T). (IsError e, SingI t) => Value t -> Either Text e
errorFromVal (Value ('TPair 'TString r) -> Either Text e)
-> Value ('TPair 'TString r) -> Either Text e
forall a b. (a -> b) -> a -> b
$ (Value 'TString, Value' Instr r) -> Value ('TPair 'TString r)
forall (l :: T) (r :: T) (instr :: [T] -> [T] -> *).
(Value' instr l, Value' instr r) -> Value' instr ('TPair l r)
VPair (MText -> Value 'TString
forall (instr :: [T] -> [T] -> *). MText -> Value' instr 'TString
VString MText
textualTag, Value' Instr r
something)
    Value t
_ -> Value t -> Either Text e
forall e (t :: T). (IsError e, SingI t) => Value t -> Either Text e
errorFromVal Value t
v

-- | If you apply numeric error representation in your contract, 'errorToVal'
-- will stop working because it doesn't know about this
-- transformation.
-- This function takes this transformation into account.
-- If a string is used as a tag, but it is not found in the passed
-- map, we conservatively preserve that string (because this whole
-- approach is rather a heuristic).
errorToValNumeric ::
  IsError e => ErrorTagMap -> e -> (forall t. ConstantScope t => Value t -> r) -> r
errorToValNumeric :: forall e r.
IsError e =>
ErrorTagMap
-> e -> (forall (t :: T). ConstantScope t => Value t -> r) -> r
errorToValNumeric ErrorTagMap
errorTagMap e
e forall (t :: T). ConstantScope t => Value t -> r
cont =
  e -> (forall (t :: T). ConstantScope t => Value t -> r) -> r
forall e r.
IsError e =>
e -> (forall (t :: T). ErrorScope t => Value t -> r) -> r
errorToVal e
e ((forall (t :: T). ConstantScope t => Value t -> r) -> r)
-> (forall (t :: T). ConstantScope t => Value t -> r) -> r
forall a b. (a -> b) -> a -> b
$ \case
    VString MText
textualTag
      | Just Natural
tag  <- MText -> ErrorTagMap -> Maybe Natural
forall a b (m :: * -> *).
(Ord a, Ord b, MonadThrow m) =>
b -> Bimap a b -> m a
Bimap.lookupR MText
textualTag ErrorTagMap
errorTagMap -> Value 'TNat -> r
forall (t :: T). ConstantScope t => Value t -> r
cont (Natural -> Value 'TNat
forall (instr :: [T] -> [T] -> *). Natural -> Value' instr 'TNat
VNat Natural
tag)

    (VPair (VString MText
textualTag, Value' Instr r
something) :: Value pair)
      | Just Natural
tag <- MText -> ErrorTagMap -> Maybe Natural
forall a b (m :: * -> *).
(Ord a, Ord b, MonadThrow m) =>
b -> Bimap a b -> m a
Bimap.lookupR MText
textualTag ErrorTagMap
errorTagMap ->
          case forall {k} (a :: k). SingI a => Sing a
forall (a :: T). SingI a => Sing a
sing @pair of
            STPair Sing n1
l Sing n2
r ->
              Sing 'TString -> (SingI 'TString => r) -> r
forall {k} (n :: k) r. Sing n -> (SingI n => r) -> r
withSingI Sing n1
Sing 'TString
l ((SingI 'TString => r) -> r) -> (SingI 'TString => r) -> r
forall a b. (a -> b) -> a -> b
$ Sing r -> (SingI r => r) -> r
forall {k} (n :: k) r. Sing n -> (SingI n => r) -> r
withSingI Sing r
Sing n2
r ((SingI r => r) -> r) -> (SingI r => r) -> r
forall a b. (a -> b) -> a -> b
$ Value ('TPair 'TNat r) -> r
forall (t :: T). ConstantScope t => Value t -> r
cont ((Value 'TNat, Value' Instr r) -> Value ('TPair 'TNat r)
forall (l :: T) (r :: T) (instr :: [T] -> [T] -> *).
(Value' instr l, Value' instr r) -> Value' instr ('TPair l r)
VPair (Natural -> Value 'TNat
forall (instr :: [T] -> [T] -> *). Natural -> Value' instr 'TNat
VNat Natural
tag, Value' Instr r
something))
    Value t
v -> Value t -> r
forall (t :: T). ConstantScope t => Value t -> r
cont Value t
v