module ZkFold.Prelude where

import           Data.Aeson           (FromJSON, ToJSON, decode, encode)
import           Data.ByteString.Lazy (readFile, writeFile)
import           Data.List            (foldl', genericIndex)
import           Data.Map             (Map, lookup)
import           GHC.Num              (Natural, integerToNatural)
import           GHC.Stack            (HasCallStack)
import           Prelude              hiding (drop, lookup, readFile, replicate, take, writeFile, (!!))
import           Test.QuickCheck      (Gen, chooseInteger, shuffle)

log2ceiling :: (Integral a, Integral b) => a -> b
log2ceiling :: forall a b. (Integral a, Integral b) => a -> b
log2ceiling = forall a b. (RealFrac a, Integral b) => a -> b
ceiling @Double (Double -> b) -> (a -> Double) -> a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Double -> Double
forall a. Floating a => a -> a -> a
logBase Double
2 (Double -> Double) -> (a -> Double) -> a -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral

length :: Foldable t => t a -> Natural
length :: forall (t :: Type -> Type) a. Foldable t => t a -> Natural
length = (Natural -> a -> Natural) -> Natural -> t a -> Natural
forall b a. (b -> a -> b) -> b -> t a -> b
forall (t :: Type -> Type) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Natural
c a
_ -> Natural
c Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
+ Natural
1) Natural
0

take :: HasCallStack => Natural -> [a] -> [a]
take :: forall a. HasCallStack => Natural -> [a] -> [a]
take Natural
0 [a]
_      = []
take Natural
n (a
x:[a]
xs) = a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: Natural -> [a] -> [a]
forall a. HasCallStack => Natural -> [a] -> [a]
take (Natural
n Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
- Natural
1) [a]
xs
take Natural
_ []     = [Char] -> [a]
forall a. HasCallStack => [Char] -> a
error [Char]
"ZkFold.Prelude.take: empty list"

drop :: Natural -> [a] -> [a]
drop :: forall a. Natural -> [a] -> [a]
drop Natural
0 [a]
xs     = [a]
xs
drop Natural
n (a
_:[a]
xs) = Natural -> [a] -> [a]
forall a. Natural -> [a] -> [a]
drop (Natural
n Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
- Natural
1) [a]
xs
drop Natural
_ []     = [Char] -> [a]
forall a. HasCallStack => [Char] -> a
error [Char]
"ZkFold.Prelude.drop: empty list"

splitAt :: Natural -> [a] -> ([a], [a])
splitAt :: forall a. Natural -> [a] -> ([a], [a])
splitAt Natural
n [a]
xs = (Natural -> [a] -> [a]
forall a. HasCallStack => Natural -> [a] -> [a]
take Natural
n [a]
xs, Natural -> [a] -> [a]
forall a. Natural -> [a] -> [a]
drop Natural
n [a]
xs)

iterateM :: Monad m => Natural -> (a -> m a) -> a -> m a
iterateM :: forall (m :: Type -> Type) a.
Monad m =>
Natural -> (a -> m a) -> a -> m a
iterateM Natural
0 a -> m a
_ a
x = a -> m a
forall a. a -> m a
forall (m :: Type -> Type) a. Monad m => a -> m a
return a
x
iterateM Natural
n a -> m a
f a
x = a -> m a
f a
x m a -> (a -> m a) -> m a
forall a b. m a -> (a -> m b) -> m b
forall (m :: Type -> Type) a b. Monad m => m a -> (a -> m b) -> m b
>>= Natural -> (a -> m a) -> a -> m a
forall (m :: Type -> Type) a.
Monad m =>
Natural -> (a -> m a) -> a -> m a
iterateM (Natural
n Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
- Natural
1) a -> m a
f

replicate :: Natural -> a -> [a]
replicate :: forall a. Natural -> a -> [a]
replicate Natural
n a
x
    | Natural
n Natural -> Natural -> Bool
forall a. Eq a => a -> a -> Bool
== Natural
0    = []
    | Bool
otherwise = a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: Natural -> a -> [a]
forall a. Natural -> a -> [a]
replicate (Natural
n Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
- Natural
1) a
x

replicateA :: Applicative f => Natural -> f a -> f [a]
replicateA :: forall (f :: Type -> Type) a.
Applicative f =>
Natural -> f a -> f [a]
replicateA Natural
n f a
fx = [f a] -> f [a]
forall (t :: Type -> Type) (f :: Type -> Type) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
forall (f :: Type -> Type) a. Applicative f => [f a] -> f [a]
sequenceA (Natural -> f a -> [f a]
forall a. Natural -> a -> [a]
replicate Natural
n f a
fx)

zipWithDefault :: (a -> b -> c) -> a -> b -> [a] -> [b] -> [c]
zipWithDefault :: forall a b c. (a -> b -> c) -> a -> b -> [a] -> [b] -> [c]
zipWithDefault a -> b -> c
_ a
_ b
_ [] []         = []
zipWithDefault a -> b -> c
f a
x b
_ [] [b]
bs         = (b -> c) -> [b] -> [c]
forall a b. (a -> b) -> [a] -> [b]
map (a -> b -> c
f a
x) [b]
bs
zipWithDefault a -> b -> c
f a
_ b
y [a]
as []         = (a -> c) -> [a] -> [c]
forall a b. (a -> b) -> [a] -> [b]
map (a -> b -> c
`f` b
y) [a]
as
zipWithDefault a -> b -> c
f a
x b
y (a
a:[a]
as) (b
b:[b]
bs) = a -> b -> c
f a
a b
b c -> [c] -> [c]
forall a. a -> [a] -> [a]
: (a -> b -> c) -> a -> b -> [a] -> [b] -> [c]
forall a b c. (a -> b -> c) -> a -> b -> [a] -> [b] -> [c]
zipWithDefault a -> b -> c
f a
x b
y [a]
as [b]
bs

elemIndex :: Eq a => a -> [a] -> Maybe Natural
elemIndex :: forall a. Eq a => a -> [a] -> Maybe Natural
elemIndex a
x = Natural -> [a] -> Maybe Natural
go Natural
0
    where
        go :: Natural -> [a] -> Maybe Natural
go Natural
_ [] = Maybe Natural
forall a. Maybe a
Nothing
        go Natural
i (a
y:[a]
ys)
            | a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
y    = Natural -> Maybe Natural
forall a. a -> Maybe a
Just Natural
i
            | Bool
otherwise = Natural -> [a] -> Maybe Natural
go (Natural
i Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
+ Natural
1) [a]
ys

(!!) :: [a] -> Natural -> a
!! :: forall a. [a] -> Natural -> a
(!!) = [a] -> Natural -> a
forall i a. Integral i => [a] -> i -> a
genericIndex

(!) :: Ord k => Map k a -> k -> a
! :: forall k a. Ord k => Map k a -> k -> a
(!) Map k a
m k
k = case k -> Map k a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
lookup k
k Map k a
m of
    Just a
x  -> a
x
    Maybe a
Nothing -> [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"ZkFold.Prelude.!: key not found"

writeFileJSON :: ToJSON a => FilePath -> a -> IO ()
writeFileJSON :: forall a. ToJSON a => [Char] -> a -> IO ()
writeFileJSON [Char]
file = [Char] -> ByteString -> IO ()
writeFile [Char]
file (ByteString -> IO ()) -> (a -> ByteString) -> a -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> ByteString
forall a. ToJSON a => a -> ByteString
encode

readFileJSON :: FromJSON a => FilePath -> IO a
readFileJSON :: forall a. FromJSON a => [Char] -> IO a
readFileJSON [Char]
file = do
    ByteString
content <- [Char] -> IO ByteString
readFile [Char]
file
    case ByteString -> Maybe a
forall a. FromJSON a => ByteString -> Maybe a
decode ByteString
content of
        Maybe a
Nothing -> [Char] -> IO a
forall a. HasCallStack => [Char] -> a
error [Char]
"ZkFold.Prelude.readFileJSON: invalid JSON"
        Just a
x  -> a -> IO a
forall a. a -> IO a
forall (m :: Type -> Type) a. Monad m => a -> m a
return a
x

assert :: Show a => Bool -> a -> x -> x
assert :: forall a x. Show a => Bool -> a -> x -> x
assert Bool
statement a
obj x
x = if Bool
statement then x
x else [Char] -> x
forall a. HasCallStack => [Char] -> a
error ([Char] -> x) -> [Char] -> x
forall a b. (a -> b) -> a -> b
$ a -> [Char]
forall a. Show a => a -> [Char]
show a
obj

chooseNatural :: (Natural, Natural) -> Gen Natural
chooseNatural :: (Natural, Natural) -> Gen Natural
chooseNatural (Natural
lo, Natural
hi) = Integer -> Natural
integerToNatural (Integer -> Natural) -> Gen Integer -> Gen Natural
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> (Integer, Integer) -> Gen Integer
chooseInteger (Natural -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Natural
lo, Natural -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Natural
hi)

genSubset :: Natural -> [a] -> Gen [a]
genSubset :: forall a. Natural -> [a] -> Gen [a]
genSubset Natural
l [a]
as = Natural -> [a] -> [a]
forall a. HasCallStack => Natural -> [a] -> [a]
take Natural
l ([a] -> [a]) -> Gen [a] -> Gen [a]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> [a] -> Gen [a]
forall a. [a] -> Gen [a]
shuffle [a]
as