module Opaleye.Internal.Window where

import           Control.Applicative (Applicative, pure, (<*>))

import qualified Opaleye.Internal.Aggregate as A
import qualified Opaleye.Internal.PackMap as PM
import qualified Opaleye.Internal.PrimQuery as PQ
import qualified Opaleye.Internal.QueryArr as Q
import qualified Opaleye.Internal.Tag as T
import qualified Opaleye.Internal.Column as C
import qualified Opaleye.Internal.Order as O

import qualified Opaleye.Internal.HaskellDB.PrimQuery as HPQ


-- | 'Window' is an applicative functor that represents expressions that
-- contain
-- [window functions](https://www.postgresql.org/docs/current/tutorial-window.html).
-- 'window' can be used to evaluate these expressions over a particular query.
newtype Window a =
  Window (PM.PackMap (HPQ.WndwOp, Partition) HPQ.PrimExpr () a)


instance Functor Window where
  fmap :: (a -> b) -> Window a -> Window b
fmap a -> b
f (Window PackMap (WndwOp, Partition) PrimExpr () a
g) = PackMap (WndwOp, Partition) PrimExpr () b -> Window b
forall a. PackMap (WndwOp, Partition) PrimExpr () a -> Window a
Window ((a -> b)
-> PackMap (WndwOp, Partition) PrimExpr () a
-> PackMap (WndwOp, Partition) PrimExpr () b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f PackMap (WndwOp, Partition) PrimExpr () a
g)


instance Applicative Window where
  pure :: a -> Window a
pure = PackMap (WndwOp, Partition) PrimExpr () a -> Window a
forall a. PackMap (WndwOp, Partition) PrimExpr () a -> Window a
Window (PackMap (WndwOp, Partition) PrimExpr () a -> Window a)
-> (a -> PackMap (WndwOp, Partition) PrimExpr () a)
-> a
-> Window a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> PackMap (WndwOp, Partition) PrimExpr () a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
  Window PackMap (WndwOp, Partition) PrimExpr () (a -> b)
f <*> :: Window (a -> b) -> Window a -> Window b
<*> Window PackMap (WndwOp, Partition) PrimExpr () a
x = PackMap (WndwOp, Partition) PrimExpr () b -> Window b
forall a. PackMap (WndwOp, Partition) PrimExpr () a -> Window a
Window (PackMap (WndwOp, Partition) PrimExpr () (a -> b)
f PackMap (WndwOp, Partition) PrimExpr () (a -> b)
-> PackMap (WndwOp, Partition) PrimExpr () a
-> PackMap (WndwOp, Partition) PrimExpr () b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> PackMap (WndwOp, Partition) PrimExpr () a
x)


runWindow :: Applicative f
  => Window a -> ((HPQ.WndwOp, Partition) -> f HPQ.PrimExpr) -> f a
runWindow :: Window a -> ((WndwOp, Partition) -> f PrimExpr) -> f a
runWindow (Window PackMap (WndwOp, Partition) PrimExpr () a
a) (WndwOp, Partition) -> f PrimExpr
f = PackMap (WndwOp, Partition) PrimExpr () a
-> ((WndwOp, Partition) -> f PrimExpr) -> () -> f a
forall (f :: * -> *) a b s t.
Applicative f =>
PackMap a b s t -> (a -> f b) -> s -> f t
PM.traversePM PackMap (WndwOp, Partition) PrimExpr () a
a (WndwOp, Partition) -> f PrimExpr
f ()


extractWindowFields
  :: T.Tag
  -> (HPQ.WndwOp, Partition)
  -> PM.PM (PQ.Bindings (HPQ.WndwOp, HPQ.Partition)) HPQ.PrimExpr
extractWindowFields :: Tag
-> (WndwOp, Partition)
-> PM (Bindings (WndwOp, Partition)) PrimExpr
extractWindowFields Tag
tag (WndwOp
op, Partition [PrimExpr]
ps [OrderExpr]
os) = do
  String
i <- PM (Bindings (WndwOp, Partition)) String
forall a. PM a String
PM.new
  let symbol :: Symbol
symbol = String -> Tag -> Symbol
HPQ.Symbol (String
"window" String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
i) Tag
tag
  (Symbol, (WndwOp, Partition))
-> PM (Bindings (WndwOp, Partition)) ()
forall a. a -> PM [a] ()
PM.write (Symbol
symbol, (WndwOp
op, ([PrimExpr] -> [OrderExpr] -> Partition
HPQ.Partition [PrimExpr]
ps [OrderExpr]
os)))
  PrimExpr -> PM (Bindings (WndwOp, Partition)) PrimExpr
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Symbol -> PrimExpr
HPQ.AttrExpr Symbol
symbol)


-- | 'window' runs a query composed of expressions containing
-- [window functions](https://www.postgresql.org/docs/current/tutorial-window.html).
-- 'window' is similar to 'Opaleye.aggregate', with the main difference being
-- that in a window query, each input row corresponds to one output row,
-- whereas aggregation queries fold the entire input query down into a single
-- row. To put this into a Haskell context, 'Opaleye.aggregate' is to 'foldl'
-- as 'window' is to 'scanl'.
window :: Q.Select (Window a) -> Q.Select a
window :: Select (Window a) -> Select a
window Select (Window a)
q = State Tag (a, PrimQuery) -> Select a
forall a. State Tag (a, PrimQuery) -> Query a
Q.productQueryArr (State Tag (a, PrimQuery) -> Select a)
-> State Tag (a, PrimQuery) -> Select a
forall a b. (a -> b) -> a -> b
$ do
  (Window a
wndw, PrimQuery
primQ) <- Select (Window a) -> () -> State Tag (Window a, PrimQuery)
forall a b. QueryArr a b -> a -> State Tag (b, PrimQuery)
Q.runSimpleQueryArr' Select (Window a)
q ()
  Tag
tag <- State Tag Tag
T.fresh
  let
    (a
a, Bindings (WndwOp, Partition)
bindings) = PM (Bindings (WndwOp, Partition)) a
-> (a, Bindings (WndwOp, Partition))
forall a r. PM [a] r -> (r, [a])
PM.run (Window a
-> ((WndwOp, Partition)
    -> PM (Bindings (WndwOp, Partition)) PrimExpr)
-> PM (Bindings (WndwOp, Partition)) a
forall (f :: * -> *) a.
Applicative f =>
Window a -> ((WndwOp, Partition) -> f PrimExpr) -> f a
runWindow Window a
wndw (Tag
-> (WndwOp, Partition)
-> PM (Bindings (WndwOp, Partition)) PrimExpr
extractWindowFields Tag
tag))
  (a, PrimQuery) -> State Tag (a, PrimQuery)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
a, Bindings (WndwOp, Partition) -> PrimQuery -> PrimQuery
forall a.
Bindings (WndwOp, Partition) -> PrimQuery' a -> PrimQuery' a
PQ.Window Bindings (WndwOp, Partition)
bindings PrimQuery
primQ)


makeWndw :: HPQ.WndwOp -> Window (C.Field_ n a)
makeWndw :: WndwOp -> Window (Field_ n a)
makeWndw WndwOp
op = PackMap (WndwOp, Partition) PrimExpr () (Field_ n a)
-> Window (Field_ n a)
forall a. PackMap (WndwOp, Partition) PrimExpr () a -> Window a
Window ((forall (f :: * -> *).
 Applicative f =>
 ((WndwOp, Partition) -> f PrimExpr) -> () -> f (Field_ n a))
-> PackMap (WndwOp, Partition) PrimExpr () (Field_ n a)
forall a b s t.
(forall (f :: * -> *). Applicative f => (a -> f b) -> s -> f t)
-> PackMap a b s t
PM.PackMap (\(WndwOp, Partition) -> f PrimExpr
f ()
_ -> PrimExpr -> Field_ n a
forall (n :: Nullability) sqlType. PrimExpr -> Field_ n sqlType
C.Column (PrimExpr -> Field_ n a) -> f PrimExpr -> f (Field_ n a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (WndwOp, Partition) -> f PrimExpr
f (WndwOp
op, Partition
forall a. Monoid a => a
mempty)))


-- | 'cumulative' allows the use of aggregation functions in 'Window'
-- expressions. In particular, @'cumulative' 'Opaleye.sum'@
-- (when combined with 'orderPartitionBy') gives a running total,
-- also known as a \"cumulative sum\", hence the name @cumulative@.
cumulative :: A.Aggregator a b -> a -> Window b
cumulative :: Aggregator a b -> a -> Window b
cumulative (A.Aggregator (PM.PackMap forall (f :: * -> *).
Applicative f =>
((Maybe (AggrOp, [OrderExpr], AggrDistinct), PrimExpr)
 -> f PrimExpr)
-> a -> f b
pm)) a
a = PackMap (WndwOp, Partition) PrimExpr () b -> Window b
forall a. PackMap (WndwOp, Partition) PrimExpr () a -> Window a
Window (PackMap (WndwOp, Partition) PrimExpr () b -> Window b)
-> PackMap (WndwOp, Partition) PrimExpr () b -> Window b
forall a b. (a -> b) -> a -> b
$ (forall (f :: * -> *).
 Applicative f =>
 ((WndwOp, Partition) -> f PrimExpr) -> () -> f b)
-> PackMap (WndwOp, Partition) PrimExpr () b
forall a b s t.
(forall (f :: * -> *). Applicative f => (a -> f b) -> s -> f t)
-> PackMap a b s t
PM.PackMap ((forall (f :: * -> *).
  Applicative f =>
  ((WndwOp, Partition) -> f PrimExpr) -> () -> f b)
 -> PackMap (WndwOp, Partition) PrimExpr () b)
-> (forall (f :: * -> *).
    Applicative f =>
    ((WndwOp, Partition) -> f PrimExpr) -> () -> f b)
-> PackMap (WndwOp, Partition) PrimExpr () b
forall a b. (a -> b) -> a -> b
$ \(WndwOp, Partition) -> f PrimExpr
f ()
_ ->
  ((Maybe (AggrOp, [OrderExpr], AggrDistinct), PrimExpr)
 -> f PrimExpr)
-> a -> f b
forall (f :: * -> *).
Applicative f =>
((Maybe (AggrOp, [OrderExpr], AggrDistinct), PrimExpr)
 -> f PrimExpr)
-> a -> f b
pm (\(Maybe (AggrOp, [OrderExpr], AggrDistinct)
mop, PrimExpr
expr) -> case Maybe (AggrOp, [OrderExpr], AggrDistinct)
mop of
         Maybe (AggrOp, [OrderExpr], AggrDistinct)
Nothing -> PrimExpr -> f PrimExpr
forall (f :: * -> *) a. Applicative f => a -> f a
pure PrimExpr
expr
         Just (AggrOp
op, [OrderExpr]
_, AggrDistinct
_) -> (WndwOp, Partition) -> f PrimExpr
f (AggrOp -> PrimExpr -> WndwOp
HPQ.WndwAggregate AggrOp
op PrimExpr
expr, Partition
forall a. Monoid a => a
mempty)) a
a


-- | 'over' adds a 'Partition' to a 'Window' expression.
--
-- @
-- 'cumulative' 'Opaleye.sum' salary \`'over'\` 'partitionBy' department <> 'orderPartitionBy' salary ('Opaleye.desc' id)
-- @
over :: Window a -> Partition -> Window a
over :: Window a -> Partition -> Window a
over (Window (PM.PackMap forall (f :: * -> *).
Applicative f =>
((WndwOp, Partition) -> f PrimExpr) -> () -> f a
pm)) Partition
partition =
  PackMap (WndwOp, Partition) PrimExpr () a -> Window a
forall a. PackMap (WndwOp, Partition) PrimExpr () a -> Window a
Window (PackMap (WndwOp, Partition) PrimExpr () a -> Window a)
-> PackMap (WndwOp, Partition) PrimExpr () a -> Window a
forall a b. (a -> b) -> a -> b
$ (forall (f :: * -> *).
 Applicative f =>
 ((WndwOp, Partition) -> f PrimExpr) -> () -> f a)
-> PackMap (WndwOp, Partition) PrimExpr () a
forall a b s t.
(forall (f :: * -> *). Applicative f => (a -> f b) -> s -> f t)
-> PackMap a b s t
PM.PackMap ((forall (f :: * -> *).
  Applicative f =>
  ((WndwOp, Partition) -> f PrimExpr) -> () -> f a)
 -> PackMap (WndwOp, Partition) PrimExpr () a)
-> (forall (f :: * -> *).
    Applicative f =>
    ((WndwOp, Partition) -> f PrimExpr) -> () -> f a)
-> PackMap (WndwOp, Partition) PrimExpr () a
forall a b. (a -> b) -> a -> b
$ \(WndwOp, Partition) -> f PrimExpr
f -> ((WndwOp, Partition) -> f PrimExpr) -> () -> f a
forall (f :: * -> *).
Applicative f =>
((WndwOp, Partition) -> f PrimExpr) -> () -> f a
pm (((WndwOp, Partition) -> f PrimExpr) -> () -> f a)
-> ((WndwOp, Partition) -> f PrimExpr) -> () -> f a
forall a b. (a -> b) -> a -> b
$ \(WndwOp
op, Partition
partition') ->
    (WndwOp, Partition) -> f PrimExpr
f (WndwOp
op, Partition
partition' Partition -> Partition -> Partition
forall a. Semigroup a => a -> a -> a
<> Partition
partition)
infixl 1 `over`


-- | In PostgreSQL, window functions must specify the \"window\" or
-- \"partition\" over which they operate. The syntax for this looks like:
-- @SUM(salary) OVER (PARTITION BY department)@. The Opaleye type 'Partition'
-- represents everything that comes after @OVER@.
--
-- 'Partition' is a 'Monoid', so 'Partition's created with 'partitionBy' and
-- 'orderPartitionBy' can be combined using '<>'.
data Partition = Partition ![HPQ.PrimExpr] ![HPQ.OrderExpr]


instance Semigroup Partition where
  Partition [PrimExpr]
ps [OrderExpr]
os <> :: Partition -> Partition -> Partition
<> Partition [PrimExpr]
ps' [OrderExpr]
os' = [PrimExpr] -> [OrderExpr] -> Partition
Partition ([PrimExpr]
ps [PrimExpr] -> [PrimExpr] -> [PrimExpr]
forall a. Semigroup a => a -> a -> a
<> [PrimExpr]
ps') ([OrderExpr]
os [OrderExpr] -> [OrderExpr] -> [OrderExpr]
forall a. Semigroup a => a -> a -> a
<> [OrderExpr]
os')


instance Monoid Partition where
  mempty :: Partition
mempty = [PrimExpr] -> [OrderExpr] -> Partition
Partition [] []


-- | Restricts a window function to operate only the group of rows that share
-- the same value(s) for the given expression(s).
partitionBy :: C.Field_ n a -> Partition
partitionBy :: Field_ n a -> Partition
partitionBy (C.Column PrimExpr
expr) = [PrimExpr] -> [OrderExpr] -> Partition
Partition [PrimExpr
expr] []


-- | Controls the order in which rows are processed by window functions. This
-- does not need to match the ordering of the overall query.
orderPartitionBy :: a -> O.Order a -> Partition
orderPartitionBy :: a -> Order a -> Partition
orderPartitionBy a
a Order a
ordering = [PrimExpr] -> [OrderExpr] -> Partition
Partition [] (a -> Order a -> [OrderExpr]
forall a. a -> Order a -> [OrderExpr]
O.orderExprs a
a Order a
ordering)