{-# LANGUAGE FlexibleContexts #-}

module Opaleye.With
  ( with,
    withRecursive,

    -- * Explicit versions
    withExplicit,
    withRecursiveExplicit,
  )
where

import Control.Monad.Trans.State.Strict (State)
import Data.Profunctor.Product.Default (Default, def)
import Opaleye.Binary (unionAllExplicit)
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, runSimpleQueryArr')
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 :: Select a -> (Select a -> Select b) -> Select b
with = Unpackspec a a -> Select a -> (Select a -> Select b) -> Select b
forall a b.
Unpackspec a a -> Select a -> (Select a -> Select b) -> Select b
withExplicit Unpackspec a a
forall (p :: * -> * -> *) a b. Default p a b => p a b
def

-- | @withRecursive s f@ is the smallest set of rows @r@ such that
--
-- @
-- r == s \`'unionAll'\` (r >>= f)
-- @
withRecursive :: Default Binaryspec a a => Select a -> (a -> Select a) -> Select a
withRecursive :: Select a -> (a -> Select a) -> Select a
withRecursive = Binaryspec a a -> Select a -> (a -> Select a) -> Select a
forall a. Binaryspec a a -> Select a -> (a -> Select a) -> Select a
withRecursiveExplicit Binaryspec a a
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 :: Unpackspec a a -> Select a -> (Select a -> Select b) -> Select b
withExplicit Unpackspec a a
unpackspec Select a
rhsSelect Select a -> Select b
bodySelect = State Tag (b, PrimQuery) -> Select b
forall a. State Tag (a, PrimQuery) -> Query a
productQueryArr (State Tag (b, PrimQuery) -> Select b)
-> State Tag (b, PrimQuery) -> Select b
forall a b. (a -> b) -> a -> b
$ do
  Unpackspec a a
-> Recursive
-> (Select a -> Select a)
-> (Select a -> Select b)
-> State Tag (b, PrimQuery)
forall a b.
Unpackspec a a
-> Recursive
-> (Select a -> Select a)
-> (Select a -> Select b)
-> State Tag (b, PrimQuery)
withG Unpackspec a a
unpackspec Recursive
PQ.NonRecursive (\Select a
_ -> Select a
rhsSelect) Select a -> Select b
bodySelect

withRecursiveExplicit :: Binaryspec a a -> Select a -> (a -> Select a) -> Select a
withRecursiveExplicit :: Binaryspec a a -> Select a -> (a -> Select a) -> Select a
withRecursiveExplicit Binaryspec a a
binaryspec Select a
base a -> Select a
recursive = State Tag (a, PrimQuery) -> Select a
forall a. State Tag (a, PrimQuery) -> Query a
productQueryArr (State Tag (a, PrimQuery) -> Select a)
-> State Tag (a, PrimQuery) -> Select a
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 = Binaryspec a a -> Select a -> Select a -> Select a
forall fields fields'.
Binaryspec fields fields'
-> Select fields -> Select fields -> Select fields'
unionAllExplicit Binaryspec a a
binaryspec Select a
base (Select a
selectCte Select a -> (a -> Select a) -> Select a
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= a -> Select a
recursive)

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

withG ::
  Unpackspec a a ->
  PQ.Recursive ->
  (Select a -> Select a) ->
  (Select a -> Select b) ->
  State Tag.Tag (b, PQ.PrimQuery)
withG :: Unpackspec a a
-> Recursive
-> (Select a -> Select a)
-> (Select a -> Select b)
-> State Tag (b, PrimQuery)
withG Unpackspec a a
unpackspec Recursive
recursive Select a -> Select a
rhsSelect Select a -> Select b
bodySelect = do
  (Select a
selectCte, Recursive -> PrimQuery -> (b, PrimQuery) -> (b, PrimQuery)
withCte) <- Unpackspec a a
-> State
     Tag
     (Select a,
      Recursive -> PrimQuery -> (b, PrimQuery) -> (b, PrimQuery))
forall a b.
Unpackspec a a
-> State
     Tag
     (Select a,
      Recursive -> 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) <- Select a -> () -> State Tag (a, PrimQuery)
forall a b. QueryArr a b -> a -> State Tag (b, PrimQuery)
runSimpleQueryArr' Select a
rhsSelect' ()
  (b, PrimQuery)
bodyQ <- Select b -> () -> State Tag (b, PrimQuery)
forall a b. QueryArr a b -> a -> State Tag (b, PrimQuery)
runSimpleQueryArr' Select b
bodySelect' ()

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

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

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

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

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

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

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