-- https://www.postgresql.org/docs/current/tutorial-window.html#id-1.4.5.6.9.5
-- talks about partitions and window frames.  The window frame is the
-- way the elements of a partition are ordered for processing the
-- result row of each element of the partition.
--
-- So neither of these terms is suitable for the _whole thing_.
-- Perhaps the answer should be "Window"?  This is also attested by
-- the WINDOW declaration in a SELECT.

module Opaleye.Internal.Window where

import           Control.Applicative (Applicative, pure, (<*>), liftA2)
import           Data.Profunctor (lmap, Profunctor, dimap)
import           Data.Semigroup (Semigroup, (<>))

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
import Data.Functor.Contravariant (contramap, Contravariant)
import Control.Arrow (second)


-- | 'WindowFunction' represents expressions that contain [window
-- functions](https://www.postgresql.org/docs/current/tutorial-window.html).
-- You can choose a 'WindowFunction' from the options below, and
-- combine and manipulate them using the @Applicative@ and
-- 'Data.Profunctor.Profunctor' operations.
newtype WindowFunction a b =
  WindowFunction (PM.PackMap HPQ.WndwOp HPQ.PrimExpr a b)

instance Functor (WindowFunction a) where
  fmap :: forall a b. (a -> b) -> WindowFunction a a -> WindowFunction a b
fmap a -> b
f (WindowFunction PackMap WndwOp PrimExpr a a
w) = forall a b. PackMap WndwOp PrimExpr a b -> WindowFunction a b
WindowFunction (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f PackMap WndwOp PrimExpr a a
w)

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

instance Profunctor WindowFunction where
  dimap :: forall a b c d.
(a -> b) -> (c -> d) -> WindowFunction b c -> WindowFunction a d
dimap a -> b
f c -> d
g (WindowFunction PackMap WndwOp PrimExpr b c
w) =  forall a b. PackMap WndwOp PrimExpr a b -> WindowFunction a b
WindowFunction (forall (p :: * -> * -> *) a b c d.
Profunctor p =>
(a -> b) -> (c -> d) -> p b c -> p a d
dimap a -> b
f c -> d
g PackMap WndwOp PrimExpr b c
w)

-- | You can create @Windows@ using 'over', and combine and manipulate
-- them using the @Applicative@ and 'Data.Profunctor.Profunctor'
-- operations.
newtype Windows a b =
  Windows (PM.PackMap (HPQ.WndwOp, Window a) HPQ.PrimExpr a b)

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

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

instance Profunctor Windows where
  dimap :: forall a b c d. (a -> b) -> (c -> d) -> Windows b c -> Windows a d
dimap a -> b
f c -> d
g (Windows (PM.PackMap forall (f :: * -> *).
Applicative f =>
((WndwOp, Window b) -> f PrimExpr) -> b -> f c
pm)) =
    forall a b. PackMap (WndwOp, Window a) PrimExpr a b -> Windows a b
Windows forall a b. (a -> b) -> a -> b
$ forall a b s t.
(forall (f :: * -> *). Applicative f => (a -> f b) -> s -> f t)
-> PackMap a b s t
PM.PackMap forall a b. (a -> b) -> a -> b
$ \(WndwOp, Window a) -> f PrimExpr
h a
a ->
      forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap c -> d
g (forall (f :: * -> *).
Applicative f =>
((WndwOp, Window b) -> f PrimExpr) -> b -> f c
pm (\(WndwOp
op, Window b
w) -> (WndwOp, Window a) -> f PrimExpr
h (WndwOp
op, forall (f :: * -> *) a' a.
Contravariant f =>
(a' -> a) -> f a -> f a'
contramap a -> b
f Window b
w)) (a -> b
f a
a))

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


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


-- | A 'WindowFunction' that doesn't actually contain any window
-- function.
noWindowFunction :: (a -> b) -> WindowFunction a b
noWindowFunction :: forall a b. (a -> b) -> WindowFunction a b
noWindowFunction a -> b
f = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f (forall a b. PackMap WndwOp PrimExpr a b -> WindowFunction a b
WindowFunction (forall a b s t.
(forall (f :: * -> *). Applicative f => (a -> f b) -> s -> f t)
-> PackMap a b s t
PM.PackMap (forall a b. a -> b -> a
const forall (f :: * -> *) a. Applicative f => a -> f a
pure)))


-- | @runWindows@ runs a query composed of expressions containing
-- [window
-- functions](https://www.postgresql.org/docs/current/tutorial-window.html).
-- @runWindows@ 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 per group. In Haskell
-- terminology, 'Opaleye.aggregate' is to 'foldl' as @runWindows@ is
-- to 'scanl'.
runWindows :: Windows a b -> Q.Select a -> Q.Select b
runWindows :: forall a b. Windows a b -> Select a -> Select b
runWindows Windows a b
wndw Select a
q = forall a. State Tag (a, PrimQuery) -> Query a
Q.productQueryArr forall a b. (a -> b) -> a -> b
$ do
  (a
a, PrimQuery
primQ) <- forall a. Select a -> State Tag (a, PrimQuery)
Q.runSimpleSelect Select a
q
  Tag
tag <- State Tag Tag
T.fresh
  let
    (b
b, Bindings (WndwOp, Partition)
bindings) = forall a r. PM [a] r -> (r, [a])
PM.run (forall (f :: * -> *) a b.
Applicative f =>
Windows a b -> ((WndwOp, Window a) -> f PrimExpr) -> a -> f b
runWindows' Windows a b
wndw (forall a.
Tag
-> a
-> (WndwOp, Window a)
-> PM (Bindings (WndwOp, Partition)) PrimExpr
extractWindowFields Tag
tag a
a) a
a)
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (b
b, forall a.
Bindings (WndwOp, Partition) -> PrimQuery' a -> PrimQuery' a
PQ.Window Bindings (WndwOp, Partition)
bindings PrimQuery
primQ)


windowsApply :: Windows (Windows a b, a) b
windowsApply :: forall a b. Windows (Windows a b, a) b
windowsApply = forall a b. PackMap (WndwOp, Window a) PrimExpr a b -> Windows a b
Windows forall a b. (a -> b) -> a -> b
$ forall a b s t.
(forall (f :: * -> *). Applicative f => (a -> f b) -> s -> f t)
-> PackMap a b s t
PM.PackMap forall a b. (a -> b) -> a -> b
$ \(WndwOp, Window (Windows a b, a)) -> f PrimExpr
f (Windows a b
agg, a
a) ->
  case Windows a b
agg of
    Windows (PM.PackMap forall (f :: * -> *).
Applicative f =>
((WndwOp, Window a) -> f PrimExpr) -> a -> f b
inner) -> forall (f :: * -> *).
Applicative f =>
((WndwOp, Window a) -> f PrimExpr) -> a -> f b
inner ((WndwOp, Window (Windows a b, a)) -> f PrimExpr
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (forall (f :: * -> *) a' a.
Contravariant f =>
(a' -> a) -> f a -> f a'
contramap forall a b. (a, b) -> b
snd)) a
a


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


makeWndwField :: (HPQ.PrimExpr -> HPQ.WndwOp)
              -> WindowFunction (C.Field_ n a) (C.Field_ n' a')
makeWndwField :: forall (n :: Nullability) a (n' :: Nullability) a'.
(PrimExpr -> WndwOp) -> WindowFunction (Field_ n a) (Field_ n' a')
makeWndwField PrimExpr -> WndwOp
f = forall (p :: * -> * -> *) a b c.
Profunctor p =>
(a -> b) -> p b c -> p a c
lmap (PrimExpr -> WndwOp
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (n :: Nullability) a. Field_ n a -> PrimExpr
C.unColumn) forall (n :: Nullability) a. WindowFunction WndwOp (Field_ n a)
makeWndw


makeWndwAny :: HPQ.WndwOp -> WindowFunction a (C.Field_ n b)
makeWndwAny :: forall a (n :: Nullability) b.
WndwOp -> WindowFunction a (Field_ n b)
makeWndwAny WndwOp
op = forall (p :: * -> * -> *) a b c.
Profunctor p =>
(a -> b) -> p b c -> p a c
lmap (forall a b. a -> b -> a
const WndwOp
op) forall (n :: Nullability) a. WindowFunction WndwOp (Field_ n a)
makeWndw

-- | 'aggregatorWindowFunction' allows the use of 'A.Aggregator's in
-- 'WindowFunction's. In particular, @'aggregatorWindowFunction'
-- 'Opaleye.sum'@ gives a running total (when combined with an order
-- argument to 'over').
aggregatorWindowFunction :: A.Aggregator a b -> (a' -> a) -> WindowFunction a' b
aggregatorWindowFunction :: forall a b a'. Aggregator a b -> (a' -> a) -> WindowFunction a' b
aggregatorWindowFunction Aggregator a b
agg a' -> a
g = forall a b. PackMap WndwOp PrimExpr a b -> WindowFunction a b
WindowFunction forall a b. (a -> b) -> a -> b
$ forall a b s t.
(forall (f :: * -> *). Applicative f => (a -> f b) -> s -> f t)
-> PackMap a b s t
PM.PackMap forall a b. (a -> b) -> a -> b
$ \WndwOp -> f PrimExpr
f a'
a ->
  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 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure PrimExpr
expr
         Just (AggrOp
op, [OrderExpr]
_, AggrDistinct
_) -> WndwOp -> f PrimExpr
f (AggrOp -> PrimExpr -> WndwOp
HPQ.WndwAggregate AggrOp
op PrimExpr
expr)) a'
a
  where A.Aggregator (PM.PackMap forall (f :: * -> *).
Applicative f =>
((Maybe (AggrOp, [OrderExpr], AggrDistinct), PrimExpr)
 -> f PrimExpr)
-> a' -> f b
pm) = forall (p :: * -> * -> *) a b c.
Profunctor p =>
(a -> b) -> p b c -> p a c
lmap a' -> a
g Aggregator a b
agg


-- | 'over' applies a 'WindowFunction' on a particular 'Window'.  For
-- example,
--
-- @
-- over ('aggregatorWindowFunction' 'Opaleye.sum' salary) ('partitionBy' department) ('Opaleye.desc' salary)
-- @
--
-- If you want to use a 'Window' that consists of the entire @SELECT@
-- then supply 'mempty' for the @'Window' a@ argument.  If you don't
-- want to order the 'Window' then supply 'mempty' for the @'O.Order'
-- a@ argument.
over :: WindowFunction a b -> Window a -> O.Order a -> Windows a b
over :: forall a b.
WindowFunction a b -> Window a -> Order a -> Windows a b
over (WindowFunction PackMap WndwOp PrimExpr a b
windowFunction) Window a
partition Order a
order =
  let PM.PackMap forall (f :: * -> *).
Applicative f =>
(WndwOp -> f PrimExpr) -> a -> f b
pm = PackMap WndwOp PrimExpr a b
windowFunction
      orderPartitionBy' :: Window a
orderPartitionBy' = forall a. Order a -> Window a
orderPartitionBy Order a
order
  in forall a b. PackMap (WndwOp, Window a) PrimExpr a b -> Windows a b
Windows forall a b. (a -> b) -> a -> b
$ forall a b s t.
(forall (f :: * -> *). Applicative f => (a -> f b) -> s -> f t)
-> PackMap a b s t
PM.PackMap forall a b. (a -> b) -> a -> b
$ \(WndwOp, Window a) -> f PrimExpr
f -> forall (f :: * -> *).
Applicative f =>
(WndwOp -> f PrimExpr) -> a -> f b
pm (\WndwOp
op ->
    (WndwOp, Window a) -> f PrimExpr
f (WndwOp
op, Window a
partition forall a. Semigroup a => a -> a -> a
<> Window a
orderPartitionBy'))

-- | In PostgreSQL, window functions must specify the \"window\" over
-- which they operate. The syntax for this looks like: @SUM(salary)
-- OVER (PARTITION BY department)@. The Opaleye type 'Window'
-- represents the segment consisting of the @PARTIION BY@.
--
-- You can create a @Window@ using 'partitionBy' and combine two
-- @Windows@ in a single one which combines the partition of both by
-- using '<>'.
data Window a = Window (a -> [HPQ.PrimExpr]) (O.Order a)

instance Semigroup (Window a) where
  Window a -> [PrimExpr]
p1 Order a
o1 <> :: Window a -> Window a -> Window a
<> Window a -> [PrimExpr]
p2 Order a
o2 = forall a. (a -> [PrimExpr]) -> Order a -> Window a
Window (a -> [PrimExpr]
p1 forall a. Semigroup a => a -> a -> a
<> a -> [PrimExpr]
p2) (Order a
o1 forall a. Semigroup a => a -> a -> a
<> Order a
o2)

instance Monoid (Window a) where
  mempty :: Window a
mempty = forall a. (a -> [PrimExpr]) -> Order a -> Window a
Window forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty
  mappend :: Window a -> Window a -> Window a
mappend = forall a. Semigroup a => a -> a -> a
(<>)

instance Contravariant Window where
  contramap :: forall a' a. (a' -> a) -> Window a -> Window a'
contramap a' -> a
f (Window a -> [PrimExpr]
p Order a
o) = forall a. (a -> [PrimExpr]) -> Order a -> Window a
Window (forall (p :: * -> * -> *) a b c.
Profunctor p =>
(a -> b) -> p b c -> p a c
lmap a' -> a
f a -> [PrimExpr]
p) (forall (f :: * -> *) a' a.
Contravariant f =>
(a' -> a) -> f a -> f a'
contramap a' -> a
f Order a
o)

-- | The window where each partition shares the same value for the
-- given 'Field'.
partitionBy :: (a -> C.Field_ n b) -> Window a
partitionBy :: forall a (n :: Nullability) b. (a -> Field_ n b) -> Window a
partitionBy a -> Field_ n b
f = forall a. (a -> [PrimExpr]) -> Order a -> Window a
Window (\a
a -> [forall (n :: Nullability) a. Field_ n a -> PrimExpr
C.unColumn (a -> Field_ n b
f a
a)]) forall a. Monoid a => a
mempty

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