{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- | Implementation for the array builtins of CIP-0156 that needs more than a one-liner.
module PlutusCore.Arrays
  ( multiIndexArray
  , maximumIndexCount
  ) where

import PlutusCore.Builtin.Result (BuiltinResult, builtinResultFailure, emit)

import Data.Text (pack)
import Data.Vector.Strict (Vector)
import Data.Vector.Strict qualified as Vector

{- Note [Index count limitation for multiIndexArray]
Looking up a list of indices means walking that list, and the time to walk a
list is not determined by its length alone, so execution time stops being
predictable from the index count beyond some size.  We therefore limit the
number of indices, and the cost model is fitted over exactly that range.
Keeping the range short also lets a single quadratic follow the measurements
closely across all of it, so a call is charged near what it costs.  The limit
is still far above realistic use, which reads tens of elements rather than
thousands.  It may be raised once costing can bound the cost without it, but
note that doing so would need a second variant of the builtin so that existing
scripts keep their current behaviour.

Compare Note [Input length limitation for IntegerToByteString], which limits
those builtins for the same underlying reason.

The count is checked during the traversal rather than up front, because taking
the length first would walk the list twice. -}
maximumIndexCount :: Int
maximumIndexCount :: Int
maximumIndexCount = Int
1024
{-# INLINE maximumIndexCount #-}

{-| Look up every index of the given list in the given array.

Fails if any index is out of bounds, or if there are more than
'maximumIndexCount' of them; see Note [Index count limitation for
multiIndexArray].

The elements are read eagerly, as in @indexArray@, so that the read is paid for
inside this builtin rather than wherever the element is later demanded.  Forcing
is safe: the vector is strict, so the element is already in normal form as far
as this function is concerned. -}
multiIndexArray :: forall a. Vector a -> [Integer] -> BuiltinResult [a]
multiIndexArray :: forall a. Vector a -> [Integer] -> BuiltinResult [a]
multiIndexArray Vector a
vec = Int -> [Integer] -> BuiltinResult [a]
go Int
0
  where
    !len :: Integer
len = Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Vector a -> Int
forall a. Vector a -> Int
Vector.length Vector a
vec)

    go :: Int -> [Integer] -> BuiltinResult [a]
    go :: Int -> [Integer] -> BuiltinResult [a]
go !Int
_ [] = [a] -> BuiltinResult [a]
forall a. a -> BuiltinResult a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
    go !Int
n (Integer
i : [Integer]
is)
      | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
maximumIndexCount = do
          Text -> BuiltinResult ()
emit (Text -> BuiltinResult ())
-> ([Char] -> Text) -> [Char] -> BuiltinResult ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Text
pack ([Char] -> BuiltinResult ()) -> [Char] -> BuiltinResult ()
forall a b. (a -> b) -> a -> b
$
            [Char]
"multiIndexArray: too many indices (maximum is "
              [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
maximumIndexCount
              [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
")"
          BuiltinResult [a]
forall a. BuiltinResult a
builtinResultFailure
      | Integer
0 Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
i Bool -> Bool -> Bool
&& Integer
i Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
len = do
          let !x :: a
x = Vector a -> Int -> a
forall a. Vector a -> Int -> a
Vector.unsafeIndex Vector a
vec (Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
i)
          (a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
:) ([a] -> [a]) -> BuiltinResult [a] -> BuiltinResult [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> [Integer] -> BuiltinResult [a]
go (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [Integer]
is
      | Bool
otherwise = do
          Text -> BuiltinResult ()
emit Text
"multiIndexArray: array index out of bounds"
          Text -> BuiltinResult ()
emit (Text -> BuiltinResult ()) -> Text -> BuiltinResult ()
forall a b. (a -> b) -> a -> b
$ Text
"Index: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> ([Char] -> Text
pack ([Char] -> Text) -> (Integer -> [Char]) -> Integer -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> [Char]
forall a. Show a => a -> [Char]
show (Integer -> Text) -> Integer -> Text
forall a b. (a -> b) -> a -> b
$ Integer
i)
          BuiltinResult [a]
forall a. BuiltinResult a
builtinResultFailure
{-# INLINE multiIndexArray #-}