{-# LANGUAGE FlexibleContexts #-}

module Opaleye.With
  ( with,
    withMaterialized,
    withRecursive,
    withRecursiveDistinct,

    -- * Explicit versions
    withExplicit,
    withMaterializedExplicit,
    withRecursiveExplicit,
    withRecursiveDistinctExplicit,
  )
where

import Control.Category ((>>>))
import Control.Monad.Trans.State.Strict (State)
import Data.Profunctor.Product.Default (Default, def)
import Opaleye.Binary (unionAllExplicit, unionExplicit)
import Opaleye.Internal.Binary (Binaryspec (..))
import qualified Opaleye.Internal.HaskellDB.PrimQuery as HPQ
import Opaleye.Internal.PackMap (PackMap (..))
import qualified Opaleye.Internal.PackMap as PM
import qualified Opaleye.Internal.PrimQuery as PQ
import Opaleye.Internal.QueryArr (Select, productQueryArr, runSimpleSelect)
import Opaleye.Internal.Rebind (rebindExplicitPrefixNoStar)
import qualified Opaleye.Internal.Sql as Sql
import qualified Opaleye.Internal.Tag as Tag
import Opaleye.Internal.Unpackspec (Unpackspec (..), runUnpackspec)

with :: Default Unpackspec a a => Select a -> (Select a -> Select b) -> Select b
with :: forall a b.
Default Unpackspec a a =>
Select a -> (Select a -> Select b) -> Select b
with = forall a b.
Unpackspec a a -> Select a -> (Select a -> Select b) -> Select b
withExplicit forall (p :: * -> * -> *) a b. Default p a b => p a b
def

withMaterialized :: Default Unpackspec a a => Select a -> (Select a -> Select b) -> Select b
withMaterialized :: forall a b.
Default Unpackspec a a =>
Select a -> (Select a -> Select b) -> Select b
withMaterialized = forall a b.
Unpackspec a a -> Select a -> (Select a -> Select b) -> Select b
withMaterializedExplicit forall (p :: * -> * -> *) a b. Default p a b => p a b
def

-- | Denotionally, @withRecursive s f@ is the smallest set of rows @r@ such
-- that
--
-- @
-- r == s \`'unionAll'\` (r >>= f)
-- @
--
-- Operationally, @withRecursive s f@ takes each row in an initial set @s@ and
-- supplies it to @f@, resulting in a new generation of rows which are added
-- to the result set. Each row from this new generation is then fed back to
-- @f@, and this process is repeated until a generation comes along for which
-- @f@ returns an empty set for each row therein.
withRecursive :: Default Binaryspec a a => Select a -> (a -> Select a) -> Select a
withRecursive :: forall a.
Default Binaryspec a a =>
Select a -> (a -> Select a) -> Select a
withRecursive = forall a. Binaryspec a a -> Select a -> (a -> Select a) -> Select a
withRecursiveExplicit forall (p :: * -> * -> *) a b. Default p a b => p a b
def

-- | Denotationally, @withRecursiveDistinct s f@ is the smallest set of rows
-- @r@ such that
--
-- @
-- r == s \`'union'\` (r >>= f)
-- @
--
-- Operationally, @withRecursiveDistinct s f@ takes each /distinct/ row in an
-- initial set @s@ and supplies it to @f@, resulting in a new generation of
-- rows. Any rows returned by @f@ that already exist in the result set are not
-- considered part of this new generation by `withRecursiveDistinct` (in
-- contrast to `withRecursive`). This new generation is then added to the
-- result set, and each row therein is then fed back to @f@, and this process
-- is repeated until a generation comes along for which @f@ returns no rows
-- that don't already exist in the result set.
withRecursiveDistinct :: Default Binaryspec a a => Select a -> (a -> Select a) -> Select a
withRecursiveDistinct :: forall a.
Default Binaryspec a a =>
Select a -> (a -> Select a) -> Select a
withRecursiveDistinct = forall a. Binaryspec a a -> Select a -> (a -> Select a) -> Select a
withRecursiveDistinctExplicit forall (p :: * -> * -> *) a b. Default p a b => p a b
def

withExplicit :: Unpackspec a a -> Select a -> (Select a -> Select b) -> Select b
withExplicit :: forall a b.
Unpackspec a a -> Select a -> (Select a -> Select b) -> Select b
withExplicit Unpackspec a a
unpackspec Select a
rhsSelect Select a -> Select b
bodySelect = forall a. State Tag (a, PrimQuery) -> Query a
productQueryArr forall a b. (a -> b) -> a -> b
$ do
  forall a b.
Unpackspec a a
-> Recursive
-> Maybe Materialized
-> (Select a -> Select a)
-> (Select a -> Select b)
-> State Tag (b, PrimQuery)
withG Unpackspec a a
unpackspec Recursive
PQ.NonRecursive forall a. Maybe a
Nothing (\Select a
_ -> forall {a}. SelectArr a a -> SelectArr a a
rebind Select a
rhsSelect) Select a -> Select b
bodySelect
  where
    rebind :: SelectArr a a -> SelectArr a a
rebind = (forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> forall a b. String -> Unpackspec a b -> SelectArr a b
rebindExplicitPrefixNoStar String
"rebind" Unpackspec a a
unpackspec)

withMaterializedExplicit :: Unpackspec a a -> Select a -> (Select a -> Select b) -> Select b
withMaterializedExplicit :: forall a b.
Unpackspec a a -> Select a -> (Select a -> Select b) -> Select b
withMaterializedExplicit Unpackspec a a
unpackspec Select a
rhsSelect Select a -> Select b
bodySelect = forall a. State Tag (a, PrimQuery) -> Query a
productQueryArr forall a b. (a -> b) -> a -> b
$ do
  forall a b.
Unpackspec a a
-> Recursive
-> Maybe Materialized
-> (Select a -> Select a)
-> (Select a -> Select b)
-> State Tag (b, PrimQuery)
withG Unpackspec a a
unpackspec Recursive
PQ.NonRecursive (forall a. a -> Maybe a
Just Materialized
PQ.Materialized) (\Select a
_ -> forall {a}. SelectArr a a -> SelectArr a a
rebind Select a
rhsSelect) Select a -> Select b
bodySelect
  where
    rebind :: SelectArr a a -> SelectArr a a
rebind = (forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> forall a b. String -> Unpackspec a b -> SelectArr a b
rebindExplicitPrefixNoStar String
"rebind" Unpackspec a a
unpackspec)

withRecursiveExplicit :: Binaryspec a a -> Select a -> (a -> Select a) -> Select a
withRecursiveExplicit :: forall a. Binaryspec a a -> Select a -> (a -> Select a) -> Select a
withRecursiveExplicit Binaryspec a a
binaryspec Select a
base a -> Select a
recursive = forall a. State Tag (a, PrimQuery) -> Query a
productQueryArr forall a b. (a -> b) -> a -> b
$ do
  let bodySelect :: p -> p
bodySelect p
selectCte = p
selectCte
  let rhsSelect :: Select a -> Select a
rhsSelect Select a
selectCte = forall fields fields'.
Binaryspec fields fields'
-> Select fields -> Select fields -> Select fields'
unionAllExplicit Binaryspec a a
binaryspec Select a
base (Select a
selectCte forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= a -> Select a
recursive)

  forall a b.
Unpackspec a a
-> Recursive
-> Maybe Materialized
-> (Select a -> Select a)
-> (Select a -> Select b)
-> State Tag (b, PrimQuery)
withG Unpackspec a a
unpackspec Recursive
PQ.Recursive forall a. Maybe a
Nothing Select a -> Select a
rhsSelect forall {p}. p -> p
bodySelect
  where
    unpackspec :: Unpackspec a a
unpackspec = forall a. Binaryspec a a -> Unpackspec a a
binaryspecToUnpackspec Binaryspec a a
binaryspec

withRecursiveDistinctExplicit :: Binaryspec a a -> Select a -> (a -> Select a) -> Select a
withRecursiveDistinctExplicit :: forall a. Binaryspec a a -> Select a -> (a -> Select a) -> Select a
withRecursiveDistinctExplicit Binaryspec a a
binaryspec Select a
base a -> Select a
recursive = forall a. State Tag (a, PrimQuery) -> Query a
productQueryArr forall a b. (a -> b) -> a -> b
$ do
  let bodySelect :: p -> p
bodySelect p
selectCte = p
selectCte
  let rhsSelect :: Select a -> Select a
rhsSelect Select a
selectCte = forall fields fields'.
Binaryspec fields fields'
-> Select fields -> Select fields -> Select fields'
unionExplicit Binaryspec a a
binaryspec Select a
base (Select a
selectCte forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= a -> Select a
recursive)

  forall a b.
Unpackspec a a
-> Recursive
-> Maybe Materialized
-> (Select a -> Select a)
-> (Select a -> Select b)
-> State Tag (b, PrimQuery)
withG Unpackspec a a
unpackspec Recursive
PQ.Recursive forall a. Maybe a
Nothing Select a -> Select a
rhsSelect forall {p}. p -> p
bodySelect
  where
    unpackspec :: Unpackspec a a
unpackspec = forall a. Binaryspec a a -> Unpackspec a a
binaryspecToUnpackspec Binaryspec a a
binaryspec

withG ::
  Unpackspec a a ->
  PQ.Recursive ->
  Maybe PQ.Materialized ->
  (Select a -> Select a) ->
  (Select a -> Select b) ->
  State Tag.Tag (b, PQ.PrimQuery)
withG :: forall a b.
Unpackspec a a
-> Recursive
-> Maybe Materialized
-> (Select a -> Select a)
-> (Select a -> Select b)
-> State Tag (b, PrimQuery)
withG Unpackspec a a
unpackspec Recursive
recursive Maybe Materialized
materialized Select a -> Select a
rhsSelect Select a -> Select b
bodySelect = do
  (Select a
selectCte, Recursive
-> Maybe Materialized
-> PrimQuery
-> (b, PrimQuery)
-> (b, PrimQuery)
withCte) <- forall a b.
Unpackspec a a
-> State
     Tag
     (Select a,
      Recursive
      -> Maybe Materialized
      -> PrimQuery
      -> (b, PrimQuery)
      -> (b, PrimQuery))
freshCte Unpackspec a a
unpackspec

  let rhsSelect' :: Select a
rhsSelect' = Select a -> Select a
rhsSelect Select a
selectCte
  let bodySelect' :: Select b
bodySelect' = Select a -> Select b
bodySelect Select a
selectCte

  (a
_, PrimQuery
rhsQ) <- forall a. Select a -> State Tag (a, PrimQuery)
runSimpleSelect Select a
rhsSelect'
  (b, PrimQuery)
bodyQ <- forall a. Select a -> State Tag (a, PrimQuery)
runSimpleSelect Select b
bodySelect'

  forall (f :: * -> *) a. Applicative f => a -> f a
pure (Recursive
-> Maybe Materialized
-> PrimQuery
-> (b, PrimQuery)
-> (b, PrimQuery)
withCte Recursive
recursive Maybe Materialized
materialized PrimQuery
rhsQ (b, PrimQuery)
bodyQ)

freshCte ::
  Unpackspec a a ->
  State
    Tag.Tag
    ( Select a,
      PQ.Recursive -> Maybe PQ.Materialized -> PQ.PrimQuery -> (b, PQ.PrimQuery) -> (b, PQ.PrimQuery)
    )
freshCte :: forall a b.
Unpackspec a a
-> State
     Tag
     (Select a,
      Recursive
      -> Maybe Materialized
      -> PrimQuery
      -> (b, PrimQuery)
      -> (b, PrimQuery))
freshCte Unpackspec a a
unpackspec = do
  Symbol
cteName <- String -> Tag -> Symbol
HPQ.Symbol String
"cte" forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> State Tag Tag
Tag.fresh

  -- TODO: Make a function that explicitly ignores its argument
  (a
cteColumns, [(Symbol, PrimExpr)]
cteBindings) <- do
    Tag
startTag <- State Tag Tag
Tag.fresh
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
      forall a r. PM [a] r -> (r, [a])
PM.run forall a b. (a -> b) -> a -> b
$
        forall (f :: * -> *) columns b.
Applicative f =>
Unpackspec columns b -> (PrimExpr -> f PrimExpr) -> columns -> f b
runUnpackspec Unpackspec a a
unpackspec (forall primExpr.
String -> Tag -> primExpr -> PM [(Symbol, primExpr)] PrimExpr
PM.extractAttr String
"cte" Tag
startTag) (forall a. HasCallStack => String -> a
error String
"freshCte")

  let selectCte :: Select a
selectCte = forall a. State Tag (a, PrimQuery) -> Query a
productQueryArr forall a b. (a -> b) -> a -> b
$ do
        Tag
tag <- State Tag Tag
Tag.fresh
        let (a
renamedCte, [(Symbol, PrimExpr)]
renameCte) =
              forall a r. PM [a] r -> (r, [a])
PM.run forall a b. (a -> b) -> a -> b
$
                forall (f :: * -> *) columns b.
Applicative f =>
Unpackspec columns b -> (PrimExpr -> f PrimExpr) -> columns -> f b
runUnpackspec Unpackspec a a
unpackspec (forall primExpr.
String -> Tag -> primExpr -> PM [(Symbol, primExpr)] PrimExpr
PM.extractAttr String
"cte_renamed" Tag
tag) a
cteColumns

        forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
renamedCte, forall a. TableIdentifier -> [(Symbol, PrimExpr)] -> PrimQuery' a
PQ.BaseTable (Maybe String -> String -> TableIdentifier
PQ.TableIdentifier forall a. Maybe a
Nothing (Symbol -> String
Sql.sqlSymbol Symbol
cteName)) [(Symbol, PrimExpr)]
renameCte)

  forall (f :: * -> *) a. Applicative f => a -> f a
pure
    ( Select a
selectCte,
      \Recursive
recursive Maybe Materialized
materialized PrimQuery
withQ (b
withedCols, PrimQuery
withedQ) ->
        (b
withedCols, forall a.
Recursive
-> Maybe Materialized
-> Symbol
-> [Symbol]
-> PrimQuery' a
-> PrimQuery' a
-> PrimQuery' a
PQ.With Recursive
recursive Maybe Materialized
materialized Symbol
cteName (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(Symbol, PrimExpr)]
cteBindings) PrimQuery
withQ PrimQuery
withedQ)
    )

binaryspecToUnpackspec :: Binaryspec a a -> Unpackspec a a
binaryspecToUnpackspec :: forall a. Binaryspec a a -> Unpackspec a a
binaryspecToUnpackspec (Binaryspec (PackMap forall (f :: * -> *).
Applicative f =>
((PrimExpr, PrimExpr) -> f PrimExpr) -> (a, a) -> f a
spec)) =
  forall fields fields'.
PackMap PrimExpr PrimExpr fields fields'
-> Unpackspec fields fields'
Unpackspec 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
PackMap forall a b. (a -> b) -> a -> b
$ \PrimExpr -> f PrimExpr
f a
a -> forall (f :: * -> *).
Applicative f =>
((PrimExpr, PrimExpr) -> f PrimExpr) -> (a, a) -> f a
spec (\(PrimExpr
pe, PrimExpr
_) -> PrimExpr -> f PrimExpr
f PrimExpr
pe) (a
a, a
a)