-- | Internal module. Not intended for public use.

module HLRDB.Internal
       ( probIO
       , primKey
       , unwrap
       , unwrapCursor
       , unwrapCreatedBool
       , unwrapCreated
       , unwrapDeleted
       , ignore
       , fixEmpty
       , fixEmpty'
       , foldM
       , decodeMInteger
       , readInt
       , Int64
       , runIdentity
       , MSET(..)
       , HLRDB.Internal.splitWith
       ) where

import Data.Functor.Identity
import Database.Redis
import Data.ByteString hiding (foldr)
import HLRDB.Primitives.Redis
import qualified Data.ByteString as S
import qualified Data.ByteString.Unsafe as B
import GHC.Int
import Control.Monad.IO.Class
import System.Random (randomRIO)


probIO :: MonadIO m => Double -> m a -> m (Maybe a)
probIO :: Double -> m a -> m (Maybe a)
probIO Double
pr m a
a =
  if Double
pr Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
1.0 then a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> m a -> m (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m a
a else do
    Double
r :: Double <- IO Double -> m Double
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Double -> m Double) -> IO Double -> m Double
forall a b. (a -> b) -> a -> b
$ (Double, Double) -> IO Double
forall a (m :: * -> *). (Random a, MonadIO m) => (a, a) -> m a
randomRIO (Double
0, Double
1.0)
    if Double
r Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
pr
       then a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> m a -> m (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m a
a
       else Maybe a -> m (Maybe a)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
forall a. Maybe a
Nothing


{-# INLINE primKey #-}
primKey :: RedisStructure v a b -> a -> ByteString
primKey :: RedisStructure v a b -> a -> ByteString
primKey (RKeyValue (E a -> ByteString
e b -> Maybe ByteString
_ Maybe ByteString -> b
_)) a
k = a -> ByteString
e a
k
primKey (RKeyValueInteger a -> ByteString
e b -> Integer
_ Integer -> b
_) a
k = a -> ByteString
e a
k
primKey (RKeyValueByteString a -> ByteString
e) a
k = a -> ByteString
e a
k
primKey (RList (E a -> ByteString
e b -> Identity ByteString
_ Identity ByteString -> b
_) Maybe TrimScheme
_) a
k = a -> ByteString
e a
k
primKey (RHSet (E a -> ByteString
e b -> Identity ByteString
_ Identity ByteString -> b
_) HSET v
_) a
k = a -> ByteString
e a
k
primKey (RSet (E a -> ByteString
e b -> Identity ByteString
_ Identity ByteString -> b
_)) a
k = a -> ByteString
e a
k
primKey (RSortedSet (E a -> ByteString
e b -> Identity ByteString
_ Identity ByteString -> b
_) Maybe TrimScheme
_) a
k = a -> ByteString
e a
k

-- Redis should never respond with errors if we are using our types consistently,
-- so transform them into exceptions
failRedis :: Reply -> Redis a
failRedis :: Reply -> Redis a
failRedis = String -> Redis a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Redis a) -> (Reply -> String) -> Reply -> Redis a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> String -> String
forall a. [a] -> [a] -> [a]
(++) String
"Unexpected Redis response: " (String -> String) -> (Reply -> String) -> Reply -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Reply -> String
forall a. Show a => a -> String
show

{-# INLINE unwrap #-}
unwrap :: MonadRedis m => Redis (Either Reply a) -> m a
unwrap :: Redis (Either Reply a) -> m a
unwrap Redis (Either Reply a)
r = do
  Either Reply a
res <- Redis (Either Reply a) -> m (Either Reply a)
forall (m :: * -> *) a. MonadRedis m => Redis a -> m a
liftRedis Redis (Either Reply a)
r
  case Either Reply a
res of
    Left Reply
e -> Redis a -> m a
forall (m :: * -> *) a. MonadRedis m => Redis a -> m a
liftRedis (Redis a -> m a) -> Redis a -> m a
forall a b. (a -> b) -> a -> b
$ Reply -> Redis a
forall a. Reply -> Redis a
failRedis Reply
e
    Right a
i -> a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
i

{-# INLINE unwrapCursor #-}
unwrapCursor :: MonadRedis m => (a -> b) -> Redis (Either Reply (Cursor , a)) -> m (Maybe Cursor , b)
unwrapCursor :: (a -> b) -> Redis (Either Reply (Cursor, a)) -> m (Maybe Cursor, b)
unwrapCursor a -> b
f =
  let g :: (Cursor, a) -> (Maybe Cursor, b)
g (Cursor
c , a
x) = (if Cursor
c Cursor -> Cursor -> Bool
forall a. Eq a => a -> a -> Bool
== Cursor
cursor0 then Maybe Cursor
forall a. Maybe a
Nothing else Cursor -> Maybe Cursor
forall a. a -> Maybe a
Just Cursor
c , a -> b
f a
x) in
  ((Cursor, a) -> (Maybe Cursor, b))
-> m (Cursor, a) -> m (Maybe Cursor, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Cursor, a) -> (Maybe Cursor, b)
g (m (Cursor, a) -> m (Maybe Cursor, b))
-> (Redis (Either Reply (Cursor, a)) -> m (Cursor, a))
-> Redis (Either Reply (Cursor, a))
-> m (Maybe Cursor, b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Redis (Either Reply (Cursor, a)) -> m (Cursor, a)
forall (m :: * -> *) a.
MonadRedis m =>
Redis (Either Reply a) -> m a
unwrap


{-# INLINE unwrapCreatedBool #-}
unwrapCreatedBool :: MonadRedis m => Redis (Either Reply Bool) -> m (ActionPerformed Creation)
unwrapCreatedBool :: Redis (Either Reply Bool) -> m (ActionPerformed Creation)
unwrapCreatedBool = (Bool -> ActionPerformed Creation)
-> m Bool -> m (ActionPerformed Creation)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Bool
b -> if Bool
b then Integer -> ActionPerformed Creation
FreshlyCreated Integer
1 else Integer -> ActionPerformed Creation
FreshlyCreated Integer
0) (m Bool -> m (ActionPerformed Creation))
-> (Redis (Either Reply Bool) -> m Bool)
-> Redis (Either Reply Bool)
-> m (ActionPerformed Creation)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Redis (Either Reply Bool) -> m Bool
forall (m :: * -> *) a.
MonadRedis m =>
Redis (Either Reply a) -> m a
unwrap

{-# INLINE unwrapCreated #-}
unwrapCreated :: MonadRedis m => Redis (Either Reply Integer) -> m (ActionPerformed Creation)
unwrapCreated :: Redis (Either Reply Integer) -> m (ActionPerformed Creation)
unwrapCreated = (Integer -> ActionPerformed Creation)
-> m Integer -> m (ActionPerformed Creation)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Integer -> ActionPerformed Creation
FreshlyCreated (m Integer -> m (ActionPerformed Creation))
-> (Redis (Either Reply Integer) -> m Integer)
-> Redis (Either Reply Integer)
-> m (ActionPerformed Creation)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Redis (Either Reply Integer) -> m Integer
forall (m :: * -> *) a.
MonadRedis m =>
Redis (Either Reply a) -> m a
unwrap

{-# INLINE unwrapDeleted #-}
unwrapDeleted :: MonadRedis m => Redis (Either Reply Integer) -> m (ActionPerformed Deletion)
unwrapDeleted :: Redis (Either Reply Integer) -> m (ActionPerformed Deletion)
unwrapDeleted = (Integer -> ActionPerformed Deletion)
-> m Integer -> m (ActionPerformed Deletion)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Integer -> ActionPerformed Deletion
Deleted (m Integer -> m (ActionPerformed Deletion))
-> (Redis (Either Reply Integer) -> m Integer)
-> Redis (Either Reply Integer)
-> m (ActionPerformed Deletion)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Redis (Either Reply Integer) -> m Integer
forall (m :: * -> *) a.
MonadRedis m =>
Redis (Either Reply a) -> m a
unwrap

{-# INLINE ignore #-}
ignore :: Functor f => f a -> f ()
ignore :: f a -> f ()
ignore = (a -> ()) -> f a -> f ()
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (() -> a -> ()
forall a b. a -> b -> a
const ())

-- Redis does not treat treat zero cases properly, so use this to fix the algebra
{-# INLINE fixEmpty #-}
fixEmpty :: (MonadRedis m , Monoid e, Traversable t) => ([ b ] -> Redis e) -> (a -> b) -> t a -> m e
fixEmpty :: ([b] -> Redis e) -> (a -> b) -> t a -> m e
fixEmpty [b] -> Redis e
f a -> b
e t a
t = case (a -> [b] -> [b]) -> [b] -> t a -> [b]
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((:) (b -> [b] -> [b]) -> (a -> b) -> a -> [b] -> [b]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
e) [] t a
t of
  [] -> e -> m e
forall (f :: * -> *) a. Applicative f => a -> f a
pure e
forall a. Monoid a => a
mempty
  [b]
xs -> Redis e -> m e
forall (m :: * -> *) a. MonadRedis m => Redis a -> m a
liftRedis (Redis e -> m e) -> Redis e -> m e
forall a b. (a -> b) -> a -> b
$ [b] -> Redis e
f [b]
xs

{-# INLINE fixEmpty' #-}
fixEmpty' :: (MonadRedis m, Traversable t, Integral i) => ([ b ] -> Redis i) -> (a -> b) -> t a -> m i
fixEmpty' :: ([b] -> Redis i) -> (a -> b) -> t a -> m i
fixEmpty' [b] -> Redis i
f a -> b
e t a
t = case (a -> [b] -> [b]) -> [b] -> t a -> [b]
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((:) (b -> [b] -> [b]) -> (a -> b) -> a -> [b] -> [b]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
e) [] t a
t of
  [] -> i -> m i
forall (f :: * -> *) a. Applicative f => a -> f a
pure i
0
  [b]
xs -> Redis i -> m i
forall (m :: * -> *) a. MonadRedis m => Redis a -> m a
liftRedis (Redis i -> m i) -> Redis i -> m i
forall a b. (a -> b) -> a -> b
$ [b] -> Redis i
f [b]
xs

{-# INLINE foldM #-}
foldM :: Foldable t => (a -> b) -> t a -> [ b ]
foldM :: (a -> b) -> t a -> [b]
foldM a -> b
f t a
t = (a -> [b] -> [b]) -> [b] -> t a -> [b]
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\a
a [b]
xs -> a -> b
f a
a b -> [b] -> [b]
forall a. a -> [a] -> [a]
: [b]
xs) [] t a
t

{-# INLINE decodeMInteger #-}
decodeMInteger :: Maybe ByteString -> Int64
decodeMInteger :: Maybe ByteString -> Int64
decodeMInteger Maybe ByteString
Nothing = Int64
0
decodeMInteger (Just ByteString
bs) = ByteString -> Int64
readInt ByteString
bs

-- From chrisdone's https://github.com/chrisdone/advent-2017-maze-rust-haskell
readInt :: ByteString -> Int64
readInt :: ByteString -> Int64
readInt ByteString
as
  | ByteString -> Bool
S.null ByteString
as = Int64
0
  | Bool
otherwise =
    case ByteString -> Word8
B.unsafeHead ByteString
as of
      Word8
45 -> Bool -> Int64 -> Int64 -> ByteString -> Int64
loop Bool
True Int64
0 Int64
0 (ByteString -> ByteString
B.unsafeTail ByteString
as)
      Word8
43 -> Bool -> Int64 -> Int64 -> ByteString -> Int64
loop Bool
False Int64
0 Int64
0 (ByteString -> ByteString
B.unsafeTail ByteString
as)
      Word8
_ -> Bool -> Int64 -> Int64 -> ByteString -> Int64
loop Bool
False Int64
0 Int64
0 ByteString
as
  where
    loop :: Bool -> Int64 -> Int64 -> S.ByteString -> Int64
    loop :: Bool -> Int64 -> Int64 -> ByteString -> Int64
loop Bool
neg !Int64
i !Int64
n !ByteString
ps
      | ByteString -> Bool
S.null ByteString
ps = Bool -> Int64 -> Int64 -> Int64
forall a p. (Eq a, Num a, Num p) => Bool -> a -> p -> p
end Bool
neg Int64
i Int64
n
      | Bool
otherwise =
        case ByteString -> Word8
B.unsafeHead ByteString
ps of
          Word8
w
            | Word8
w Word8 -> Word8 -> Bool
forall a. Ord a => a -> a -> Bool
>= Word8
0x30 Bool -> Bool -> Bool
&& Word8
w Word8 -> Word8 -> Bool
forall a. Ord a => a -> a -> Bool
<= Word8
0x39 ->
              Bool -> Int64 -> Int64 -> ByteString -> Int64
loop
                Bool
neg
                (Int64
i Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
+ Int64
1)
                (Int64
n Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
* Int64
10 Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
+ (Word8 -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
w Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
0x30))
                (ByteString -> ByteString
B.unsafeTail ByteString
ps)
            | Bool
otherwise -> Bool -> Int64 -> Int64 -> Int64
forall a p. (Eq a, Num a, Num p) => Bool -> a -> p -> p
end Bool
neg Int64
i Int64
n
    end :: Bool -> a -> p -> p
end Bool
_ a
0 p
_ = p
0
    end Bool
True a
_ p
n = p -> p
forall a. Num a => a -> a
negate p
n
    end Bool
_ a
_ p
n = p
n

type DList a = ([ a ] -> [ a ])

-- | Aggregated @mset@ query
newtype MSET = MSET { MSET -> DList (ByteString, Maybe ByteString)
runMSET :: DList (ByteString , Maybe ByteString) }

instance Semigroup MSET where
  {-# INLINE (<>) #-}
  <> :: MSET -> MSET -> MSET
(<>) (MSET DList (ByteString, Maybe ByteString)
as) (MSET DList (ByteString, Maybe ByteString)
bs) = DList (ByteString, Maybe ByteString) -> MSET
MSET (DList (ByteString, Maybe ByteString)
as DList (ByteString, Maybe ByteString)
-> DList (ByteString, Maybe ByteString)
-> DList (ByteString, Maybe ByteString)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DList (ByteString, Maybe ByteString)
bs)

instance Monoid MSET where
  {-# INLINE mempty #-}
  mempty :: MSET
mempty = DList (ByteString, Maybe ByteString) -> MSET
MSET (DList (ByteString, Maybe ByteString) -> MSET)
-> DList (ByteString, Maybe ByteString) -> MSET
forall a b. (a -> b) -> a -> b
$ [(ByteString, Maybe ByteString)]
-> DList (ByteString, Maybe ByteString)
forall a. Semigroup a => a -> a -> a
(<>) []

splitWith :: (a -> Either b c) -> [ a ] -> ([ b ] , [ c ])
splitWith :: (a -> Either b c) -> [a] -> ([b], [c])
splitWith a -> Either b c
f = (a -> ([b], [c]) -> ([b], [c])) -> ([b], [c]) -> [a] -> ([b], [c])
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> ([b], [c]) -> ([b], [c])
g ([b], [c])
forall a. Monoid a => a
mempty
  where
    g :: a -> ([b], [c]) -> ([b], [c])
g a
x ([b]
as , [c]
bs) = case a -> Either b c
f a
x of
      Left b
a -> (b
a b -> [b] -> [b]
forall a. a -> [a] -> [a]
: [b]
as , [c]
bs)
      Right c
b -> ([b]
as , c
b c -> [c] -> [c]
forall a. a -> [a] -> [a]
: [c]
bs)