{-# LANGUAGE BlockArguments #-}

-- | Combinators that can be used for aggregating independent queries. See my <https://identicalsnowflake.github.io/QueryAggregation.html article> about aggregating mget queries for more information.

module HLRDB.Primitives.Aggregate
       (
         T(..)
       , type (⟿)
       , type (~~>)
       , type Query
       , aggregatePair
       , remember
       , bitraverse'
       , runT

       -- | Aggregate, atomic multi-set query (as in setting multiple things in a single query)
       , MSET
       ) where

import Data.Bitraversable
import Data.Profunctor
import Data.Profunctor.Traversing
import Control.Lens hiding (Traversing)
import Data.ByteString
import HLRDB.Internal (MSET)


-- | Abstract representation for aggregation.
newtype T x y a b = T (Traversal a b x y) deriving (a -> T x y a b -> T x y a a
(a -> b) -> T x y a a -> T x y a b
(forall a b. (a -> b) -> T x y a a -> T x y a b)
-> (forall a b. a -> T x y a b -> T x y a a) -> Functor (T x y a)
forall a b. a -> T x y a b -> T x y a a
forall a b. (a -> b) -> T x y a a -> T x y a b
forall x y a a b. a -> T x y a b -> T x y a a
forall x y a a b. (a -> b) -> T x y a a -> T x y a b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> T x y a b -> T x y a a
$c<$ :: forall x y a a b. a -> T x y a b -> T x y a a
fmap :: (a -> b) -> T x y a a -> T x y a b
$cfmap :: forall x y a a b. (a -> b) -> T x y a a -> T x y a b
Functor)

instance Profunctor (T x y) where
  {-# INLINE lmap #-}
  lmap :: (a -> b) -> T x y b c -> T x y a c
lmap a -> b
f (T Traversal b c x y
t) = Traversal a c x y -> T x y a c
forall x y a b. Traversal a b x y -> T x y a b
T \x -> f y
x -> (x -> f y) -> b -> f c
Traversal b c x y
t x -> f y
x (b -> f c) -> (a -> b) -> a -> f c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
f
  {-# INLINE rmap #-}
  rmap :: (b -> c) -> T x y a b -> T x y a c
rmap b -> c
g (T Traversal a b x y
t) = Traversal a c x y -> T x y a c
forall x y a b. Traversal a b x y -> T x y a b
T \x -> f y
x -> (b -> c) -> f b -> f c
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap b -> c
g (f b -> f c) -> (a -> f b) -> a -> f c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (x -> f y) -> a -> f b
Traversal a b x y
t x -> f y
x
  {-# INLINE dimap #-}
  dimap :: (a -> b) -> (c -> d) -> T x y b c -> T x y a d
dimap a -> b
f c -> d
g (T Traversal b c x y
t) = Traversal a d x y -> T x y a d
forall x y a b. Traversal a b x y -> T x y a b
T \x -> f y
m -> (c -> d) -> f c -> f d
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap c -> d
g (f c -> f d) -> (a -> f c) -> a -> f d
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (x -> f y) -> b -> f c
Traversal b c x y
t x -> f y
m (b -> f c) -> (a -> b) -> a -> f c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
f

instance Traversing (T x y) where
  {-# INLINE traverse' #-}
  traverse' :: T x y a b -> T x y (f a) (f b)
traverse' (T Traversal a b x y
t) = Traversal (f a) (f b) x y -> T x y (f a) (f b)
forall x y a b. Traversal a b x y -> T x y a b
T ((a -> f b) -> f a -> f (f b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ((a -> f b) -> f a -> f (f b))
-> ((x -> f y) -> a -> f b) -> (x -> f y) -> f a -> f (f b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (x -> f y) -> a -> f b
Traversal a b x y
t)

instance Applicative (T x y a) where
  {-# INLINE pure #-}
  pure :: a -> T x y a a
pure a
x = Traversal a a x y -> T x y a a
forall x y a b. Traversal a b x y -> T x y a b
T (Traversal a a x y -> T x y a a) -> Traversal a a x y -> T x y a a
forall a b. (a -> b) -> a -> b
$ \x -> f y
_ a
_ -> a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x
  {-# INLINE (<*>) #-}
  <*> :: T x y a (a -> b) -> T x y a a -> T x y a b
(<*>) (T Traversal a (a -> b) x y
f) (T Traversal a a x y
x) = Traversal a b x y -> T x y a b
forall x y a b. Traversal a b x y -> T x y a b
T \x -> f y
g a
a -> (x -> f y) -> a -> f (a -> b)
Traversal a (a -> b) x y
f x -> f y
g a
a f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (x -> f y) -> a -> f a
Traversal a a x y
x x -> f y
g a
a

-- | We can merge any two arbitrary mget queries.
{-# INLINE aggregatePair #-}
aggregatePair :: (Traversing p , Functor (p (a , a')) , Applicative (p (a , a'))) => p a b -> p a' b' -> p (a , a') (b , b')
aggregatePair :: p a b -> p a' b' -> p (a, a') (b, b')
aggregatePair p a b
x p a' b'
y =
  (,) (b -> b' -> (b, b')) -> p (a, a') b -> p (a, a') (b' -> (b, b'))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((a, a') -> a) -> p a b -> p (a, a') b
forall (p :: * -> * -> *) a b c.
Profunctor p =>
(a -> b) -> p b c -> p a c
lmap (Getting a (a, a') a -> (a, a') -> a
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting a (a, a') a
forall s t a b. Field1 s t a b => Lens s t a b
_1) p a b
x p (a, a') (b' -> (b, b')) -> p (a, a') b' -> p (a, a') (b, b')
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ((a, a') -> a') -> p a' b' -> p (a, a') b'
forall (p :: * -> * -> *) a b c.
Profunctor p =>
(a -> b) -> p b c -> p a c
lmap (Getting a' (a, a') a' -> (a, a') -> a'
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting a' (a, a') a'
forall s t a b. Field2 s t a b => Lens s t a b
_2) p a' b'
y

-- Remember could probably be a Profunctor typeclass in general (is it?)
-- | And we can remember the lookup
{-# INLINE remember #-}
remember :: T x y a b -> T x y a (a , b)
remember :: T x y a b -> T x y a (a, b)
remember (T Traversal a b x y
f) = Traversal a (a, b) x y -> T x y a (a, b)
forall x y a b. Traversal a b x y -> T x y a b
T \x -> f y
x a
a -> (,) a
a (b -> (a, b)) -> f b -> f (a, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (x -> f y) -> a -> f b
Traversal a b x y
f x -> f y
x a
a

{-# INLINABLE bitraverse' #-}
bitraverse' :: Bitraversable t => a ~~> b -> c ~~> d -> t a c ~~> t b d
bitraverse' :: (a ~~> b) -> (c ~~> d) -> t a c ~~> t b d
bitraverse' a ~~> b
x c ~~> d
y = (t a c -> () ~~> t b d) -> t a c ~~> t b d
forall a b. (a -> () ~~> b) -> a ~~> b
rev' ((a -> T ByteString (Maybe ByteString) () b)
-> (c -> T ByteString (Maybe ByteString) () d)
-> t a c
-> () ~~> t b d
forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse (((() -> a) -> (a ~~> b) -> T ByteString (Maybe ByteString) () b)
-> (a ~~> b) -> (() -> a) -> T ByteString (Maybe ByteString) () b
forall a b c. (a -> b -> c) -> b -> a -> c
flip (() -> a) -> (a ~~> b) -> T ByteString (Maybe ByteString) () b
forall (p :: * -> * -> *) a b c.
Profunctor p =>
(a -> b) -> p b c -> p a c
lmap a ~~> b
x ((() -> a) -> T ByteString (Maybe ByteString) () b)
-> (a -> () -> a) -> a -> T ByteString (Maybe ByteString) () b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> () -> a
forall a b. a -> b -> a
const) (((() -> c) -> (c ~~> d) -> T ByteString (Maybe ByteString) () d)
-> (c ~~> d) -> (() -> c) -> T ByteString (Maybe ByteString) () d
forall a b c. (a -> b -> c) -> b -> a -> c
flip (() -> c) -> (c ~~> d) -> T ByteString (Maybe ByteString) () d
forall (p :: * -> * -> *) a b c.
Profunctor p =>
(a -> b) -> p b c -> p a c
lmap c ~~> d
y ((() -> c) -> T ByteString (Maybe ByteString) () d)
-> (c -> () -> c) -> c -> T ByteString (Maybe ByteString) () d
forall b c a. (b -> c) -> (a -> b) -> a -> c
. c -> () -> c
forall a b. a -> b -> a
const))
      where
        rev' :: (a -> () ~~> b) -> a ~~> b
        rev' :: (a -> () ~~> b) -> a ~~> b
rev' a -> () ~~> b
f = Traversal a b ByteString (Maybe ByteString) -> a ~~> b
forall x y a b. Traversal a b x y -> T x y a b
T \ByteString -> f (Maybe ByteString)
g a
v -> case a -> () ~~> b
f a
v of
          T Traversal () b ByteString (Maybe ByteString)
m -> (ByteString -> f (Maybe ByteString)) -> () -> f b
Traversal () b ByteString (Maybe ByteString)
m ByteString -> f (Maybe ByteString)
g ()

instance Strong (T x y) where
  {-# INLINE first' #-}
  first' :: T x y a b -> T x y (a, c) (b, c)
first' = T x y a b -> T x y (a, c) (b, c)
forall (p :: * -> * -> *) a b c.
Traversing p =>
p a b -> p (a, c) (b, c)
firstTraversing

instance Choice (T x y) where
  {-# INLINE left' #-}
  left' :: T x y a b -> T x y (Either a c) (Either b c)
left' = T x y a b -> T x y (Either a c) (Either b c)
forall (p :: * -> * -> *) a b c.
Traversing p =>
p a b -> p (Either a c) (Either b c)
leftTraversing

-- | Reify aggregation into a target functor.
{-# INLINE runT #-}
runT :: Functor f => ([x] -> f [y]) -> T x y a b -> a -> f b
runT :: ([x] -> f [y]) -> T x y a b -> a -> f b
runT [x] -> f [y]
i (T Traversal a b x y
t) = Traversing (->) f a b x y -> LensLike f a b [x] [y]
forall (f :: * -> *) s t a b.
Functor f =>
Traversing (->) f s t a b -> LensLike f s t [a] [b]
unsafePartsOf Traversing (->) f a b x y
Traversal a b x y
t [x] -> f [y]
i


-- | A query using input of type 'a' and yielding an output of type 'b'
type (⟿) a b = T ByteString (Maybe ByteString) a b

-- | An ASCII version of ⟿
type (~~>) a b = T ByteString (Maybe ByteString) a b

-- | Non-infix alias of ⟿
type Query a b = a ⟿ b