{-# language StandaloneKindSignatures #-}

module Rel8.Query
  ( Query( Query )
  )
where

-- base
import Control.Applicative ( liftA2 )
import Control.Monad ( liftM2 )
import Data.Kind ( Type )
import Data.Monoid ( Any( Any ) )
import Prelude

-- opaleye
import qualified Opaleye.Internal.HaskellDB.PrimQuery as Opaleye
import qualified Opaleye.Internal.PackMap as Opaleye
import qualified Opaleye.Internal.PrimQuery as Opaleye hiding (lateral)
import qualified Opaleye.Internal.QueryArr as Opaleye
import qualified Opaleye.Internal.Tag as Opaleye

-- rel8
import Rel8.Query.Set ( unionAll )
import Rel8.Query.Opaleye ( fromOpaleye )
import Rel8.Query.Values ( values )
import Rel8.Table ( fromColumns, toColumns )
import Rel8.Table.Alternative
  ( AltTable, (<|>:)
  , AlternativeTable, emptyTable
  )
import Rel8.Table.Projection ( Projectable, apply, project )

-- semigroupoids
import Data.Functor.Apply ( Apply, (<.>) )
import Data.Functor.Bind ( Bind, (>>-) )


-- | The @Query@ monad allows you to compose a @SELECT@ query. This monad has
-- semantics similar to the list (@[]@) monad.
type Query :: Type -> Type
newtype Query a =
  Query (
    -- This is based on Opaleye's Select monad, but with two addtions. We
    -- maintain a stack of PrimExprs from parent previous subselects. In
    -- practice, these are always the results of dummy calls to random().
    --
    -- We also return a Bool that indicates to the parent subselect whether
    -- or not that stack of PrimExprs were used at any point. If they weren't,
    -- then the call to random() is never added to the query.
    --
    -- This is all needed to implement evaluate. Consider the following code:
    --
    -- do
    --   x <- values [lit 'a', lit 'b', lit 'c']
    --   y <- evaluate $ nextval "user_id_seq"
    --   pure (x, y)
    --
    -- If we just used Opaleye's Select monad directly, the SQL would come out
    -- like this:
    --
    -- SELECT
    --   a, b
    -- FROM
    --   (VALUES ('a'), ('b'), ('c')) Q1(a),
    --   LATERAL (SELECT nextval('user_id_seq')) Q2(b);
    --
    -- From the Haskell code, you would intuitively expect to get back the
    -- results of three different calls to nextval(), but from Postgres' point
    -- of view, because the Q2 subquery doesn't reference anything from the Q1
    -- query, it thinks it only needs to call nextval() once. This is actually
    -- exactly the same problem you get with the deprecated ListT IO monad from
    -- the transformers package — *> behaves differently to >>=, so
    -- using ApplicativeDo can change the results of a program. ApplicativeDo
    -- is exactly the optimisation Postgres does on a "LATERAL" query that
    -- doesn't make any references to previous subselects.
    --
    -- Rel8's solution is generate the following SQL instead:
    --
    -- SELECT
    --   a, b
    -- FROM
    --   (SELECT
    --      random() AS dummy,
    --      *
    --    FROM
    --      (VALUES ('a'), ('b'), ('c')) Q1(a)) Q1,
    --   LATERAL (SELECT
    --     CASE
    --       WHEN dummy IS NOT NULL
    --       THEN nextval('user_id_seq')
    --     END) Q2(b);
    --
    -- We use random() here as the dummy value (and not some constant) because
    -- Postgres will again optimize if it sees that a value is constant
    -- (and thus only call nextval() once), but because random() is marked as
    -- VOLATILE, this inhibits Postgres from doing that optimisation.
    --
    -- Why not just reference the a column from the previous query directly
    -- instead of adding a dummy value? Basically, even if we extract out all
    -- the bindings introduced in a PrimQuery, we can't always be sure which
    -- ones refer to constant values, so if we end up laterally referencing a
    -- constant value, then all of this would be for nothing.
    --
    -- Why not just add the call to the previous subselect directly, like so:
    --
    -- SELECT
    --   a, b
    -- FROM
    --   (SELECT
    --      nextval('user_id_seq') AS eval,
    --      *
    --    FROM
    --      (VALUES ('a'), ('b'), ('c')) Q1(a)) Q1,
    --   LATERAL (SELECT eval) Q2(b);
    --
    -- That would work in this case. But consider the following Rel8 code:
    --
    -- do
    --   x <- values [lit 'a', lit 'b', lit 'c']
    --   y <- values [lit 'd', lit 'e', lit 'f']
    --   z <- evaluate $ nextval "user_id_seq"
    --   pure (x, y, z)
    --
    -- How many calls to nextval should there be? Our Haskell intuition says
    -- nine. But that's not what you would get if you used the above
    -- technique. The problem is, which VALUES query should the nextval be
    -- added to? You can choose one or the other to get three calls to
    -- nextval, but you still need to make a superfluous LATERAL references to
    -- the other if you want nine calls. So for the above Rel8 code we generate
    -- the following SQL:
    --
    -- SELECT
    --   a, b, c
    -- FROM
    --   (SELECT
    --      random() AS dummy,
    --      *
    --    FROM
    --      (VALUES ('a'), ('b'), ('c')) Q1(a)) Q1,
    --   (SELECT
    --      random() AS dummy,
    --      *
    --    FROM
    --      (VALUES ('d'), ('e'), ('f')) Q2(b)) Q2,
    --   LATERAL (SELECT
    --     CASE
    --       WHEN Q1.dummy IS NOT NULL AND Q2.dummy IS NOT NULL
    --       THEN nextval('user_id_seq')
    --     END) Q3(c);
    --
    -- This gives nine calls to nextval() as we would expect.
    [Opaleye.PrimExpr] -> Opaleye.Select (Any, a)
  )


instance Projectable Query where
  project :: Projection a b -> Query a -> Query b
project Projection a b
f = (a -> b) -> Query a -> Query b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Columns (Transpose (Field a) b) (Context b) -> b
forall (context :: * -> *) a.
Table context a =>
Columns a context -> a
fromColumns (Columns (Transpose (Field a) b) (Context b) -> b)
-> (a -> Columns (Transpose (Field a) b) (Context b)) -> a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Projection a b -> Columns a (Context b) -> Columns b (Context b)
forall a b (context :: * -> *).
Projecting a b =>
Projection a b -> Columns a context -> Columns b context
apply Projection a b
f (Columns a (Context b)
 -> Columns (Transpose (Field a) b) (Context b))
-> (a -> Columns a (Context b))
-> a
-> Columns (Transpose (Field a) b) (Context b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Columns a (Context b)
forall (context :: * -> *) a.
Table context a =>
a -> Columns a context
toColumns)


instance Functor Query where
  fmap :: (a -> b) -> Query a -> Query b
fmap a -> b
f (Query [PrimExpr] -> Select (Any, a)
a) = ([PrimExpr] -> Select (Any, b)) -> Query b
forall a. ([PrimExpr] -> Select (Any, a)) -> Query a
Query ((Select (Any, a) -> Select (Any, b))
-> ([PrimExpr] -> Select (Any, a)) -> [PrimExpr] -> Select (Any, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (((Any, a) -> (Any, b)) -> Select (Any, a) -> Select (Any, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a -> b) -> (Any, a) -> (Any, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f)) [PrimExpr] -> Select (Any, a)
a)


instance Apply Query where
  <.> :: Query (a -> b) -> Query a -> Query b
(<.>) = Query (a -> b) -> Query a -> Query b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
(<*>)


instance Applicative Query where
  pure :: a -> Query a
pure = Select a -> Query a
forall a. Select a -> Query a
fromOpaleye (Select a -> Query a) -> (a -> Select a) -> a -> Query a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Select a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
  liftA2 :: (a -> b -> c) -> Query a -> Query b -> Query c
liftA2 = (a -> b -> c) -> Query a -> Query b -> Query c
forall (m :: * -> *) a1 a2 r.
Monad m =>
(a1 -> a2 -> r) -> m a1 -> m a2 -> m r
liftM2


instance Bind Query where
  >>- :: Query a -> (a -> Query b) -> Query b
(>>-) = Query a -> (a -> Query b) -> Query b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
(>>=)


instance Monad Query where
  Query [PrimExpr] -> Select (Any, a)
q >>= :: Query a -> (a -> Query b) -> Query b
>>= a -> Query b
f = ([PrimExpr] -> Select (Any, b)) -> Query b
forall a. ([PrimExpr] -> Select (Any, a)) -> Query a
Query (([PrimExpr] -> Select (Any, b)) -> Query b)
-> ([PrimExpr] -> Select (Any, b)) -> Query b
forall a b. (a -> b) -> a -> b
$ \[PrimExpr]
dummies -> (() -> Tag -> ((Any, b), PrimQueryArr, Tag)) -> Select (Any, b)
forall a b. (a -> Tag -> (b, PrimQueryArr, Tag)) -> QueryArr a b
Opaleye.stateQueryArr ((() -> Tag -> ((Any, b), PrimQueryArr, Tag)) -> Select (Any, b))
-> (() -> Tag -> ((Any, b), PrimQueryArr, Tag)) -> Select (Any, b)
forall a b. (a -> b) -> a -> b
$ \()
_ Tag
tag ->
    let
      qa :: Select (Any, a)
qa = [PrimExpr] -> Select (Any, a)
q [PrimExpr]
dummies
      ((Any
m, a
a), PrimQueryArr
query, Tag
tag') = Select (Any, a) -> () -> Tag -> ((Any, a), PrimQueryArr, Tag)
forall a b. QueryArr a b -> a -> Tag -> (b, PrimQueryArr, Tag)
Opaleye.runStateQueryArr Select (Any, a)
qa () Tag
tag
      Query [PrimExpr] -> Select (Any, b)
q' = a -> Query b
f a
a
      ([PrimExpr]
dummies', PrimQueryArr
query', Tag
tag'') =
        ( PrimExpr
dummy PrimExpr -> [PrimExpr] -> [PrimExpr]
forall a. a -> [a] -> [a]
: [PrimExpr]
dummies
        , PrimQueryArr
query PrimQueryArr -> PrimQueryArr -> PrimQueryArr
forall a. Semigroup a => a -> a -> a
<> Bindings PrimExpr -> PrimQueryArr
Opaleye.aRebind Bindings PrimExpr
bindings
        , Tag -> Tag
Opaleye.next Tag
tag'
        )
        where
          (PrimExpr
dummy, Bindings PrimExpr
bindings) = PM (Bindings PrimExpr) PrimExpr -> (PrimExpr, Bindings PrimExpr)
forall a r. PM [a] r -> (r, [a])
Opaleye.run (PM (Bindings PrimExpr) PrimExpr -> (PrimExpr, Bindings PrimExpr))
-> PM (Bindings PrimExpr) PrimExpr -> (PrimExpr, Bindings PrimExpr)
forall a b. (a -> b) -> a -> b
$ PrimExpr -> PM (Bindings PrimExpr) PrimExpr
forall primExpr. primExpr -> PM [(Symbol, primExpr)] PrimExpr
name PrimExpr
random
            where
              random :: PrimExpr
random = Name -> [PrimExpr] -> PrimExpr
Opaleye.FunExpr Name
"random" []
              name :: primExpr -> PM [(Symbol, primExpr)] PrimExpr
name = Name -> Tag -> primExpr -> PM [(Symbol, primExpr)] PrimExpr
forall primExpr.
Name -> Tag -> primExpr -> PM [(Symbol, primExpr)] PrimExpr
Opaleye.extractAttr Name
"dummy" Tag
tag'
      qa' :: SelectArr i (Any, b)
qa' = (i -> Select (Any, b)) -> SelectArr i (Any, b)
forall i a. (i -> Select a) -> SelectArr i a
Opaleye.lateral ((i -> Select (Any, b)) -> SelectArr i (Any, b))
-> (i -> Select (Any, b)) -> SelectArr i (Any, b)
forall a b. (a -> b) -> a -> b
$ \i
_ -> [PrimExpr] -> Select (Any, b)
q' [PrimExpr]
dummies'
      ((m' :: Any
m'@(Any Bool
needsDummies), b
b), PrimQueryArr
query'', Tag
tag''') = Select (Any, b) -> () -> Tag -> ((Any, b), PrimQueryArr, Tag)
forall a b. QueryArr a b -> a -> Tag -> (b, PrimQueryArr, Tag)
Opaleye.runStateQueryArr Select (Any, b)
forall i. SelectArr i (Any, b)
qa' () Tag
tag''
      query''' :: PrimQueryArr
query'''
        | Bool
needsDummies = PrimQueryArr
query' PrimQueryArr -> PrimQueryArr -> PrimQueryArr
forall a. Semigroup a => a -> a -> a
<> PrimQueryArr
query''
        | Bool
otherwise = PrimQueryArr
query PrimQueryArr -> PrimQueryArr -> PrimQueryArr
forall a. Semigroup a => a -> a -> a
<> PrimQueryArr
query''
      m'' :: Any
m'' = Any
m Any -> Any -> Any
forall a. Semigroup a => a -> a -> a
<> Any
m'
    in
      ((Any
m'', b
b), PrimQueryArr
query''', Tag
tag''')


-- | '<|>:' = 'unionAll'.
instance AltTable Query where
  <|>: :: Query a -> Query a -> Query a
(<|>:) = Query a -> Query a -> Query a
forall a. Table Expr a => Query a -> Query a -> Query a
unionAll


-- | 'emptyTable' = 'values' @[]@.
instance AlternativeTable Query where
  emptyTable :: Query a
emptyTable = [a] -> Query a
forall a (f :: * -> *).
(Table Expr a, Foldable f) =>
f a -> Query a
values []