{-# LANGUAGE GeneralizedNewtypeDeriving, OverloadedStrings, CPP #-}
module Database.Selda.Query.Type where
import Control.Monad.State.Strict
    ( StateT(StateT), MonadState(put, get), State, runState )
import Data.Text (pack)
import Database.Selda.SQL ( SQL )
import Database.Selda.Exp
    ( Exp(Col), UntypedCol(..), SomeCol(Named) )
import Database.Selda.Types (ColName, mkColName, addColSuffix)

type Scope = Int
type Ident = Int

-- | A name, consisting of a scope and an identifier.
data Name = Name Scope Ident

instance Show Name where
  show :: Name -> String
show (Name Int
0 Int
n) = forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [forall a. Show a => a -> String
show Int
n]
  show (Name Int
s Int
n) = forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [forall a. Show a => a -> String
show Int
s, String
"s_", forall a. Show a => a -> String
show Int
n]

-- | An SQL query.
newtype Query s a = Query {forall s a. Query s a -> State GenState a
unQ :: State GenState a}
  deriving (forall a b. a -> Query s b -> Query s a
forall a b. (a -> b) -> Query s a -> Query s b
forall s a b. a -> Query s b -> Query s a
forall s a b. (a -> b) -> Query s a -> Query s b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> Query s b -> Query s a
$c<$ :: forall s a b. a -> Query s b -> Query s a
fmap :: forall a b. (a -> b) -> Query s a -> Query s b
$cfmap :: forall s a b. (a -> b) -> Query s a -> Query s b
Functor, forall s. Functor (Query s)
forall a. a -> Query s a
forall s a. a -> Query s a
forall a b. Query s a -> Query s b -> Query s a
forall a b. Query s a -> Query s b -> Query s b
forall a b. Query s (a -> b) -> Query s a -> Query s b
forall s a b. Query s a -> Query s b -> Query s a
forall s a b. Query s a -> Query s b -> Query s b
forall s a b. Query s (a -> b) -> Query s a -> Query s b
forall a b c. (a -> b -> c) -> Query s a -> Query s b -> Query s c
forall s a b c.
(a -> b -> c) -> Query s a -> Query s b -> Query s c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
<* :: forall a b. Query s a -> Query s b -> Query s a
$c<* :: forall s a b. Query s a -> Query s b -> Query s a
*> :: forall a b. Query s a -> Query s b -> Query s b
$c*> :: forall s a b. Query s a -> Query s b -> Query s b
liftA2 :: forall a b c. (a -> b -> c) -> Query s a -> Query s b -> Query s c
$cliftA2 :: forall s a b c.
(a -> b -> c) -> Query s a -> Query s b -> Query s c
<*> :: forall a b. Query s (a -> b) -> Query s a -> Query s b
$c<*> :: forall s a b. Query s (a -> b) -> Query s a -> Query s b
pure :: forall a. a -> Query s a
$cpure :: forall s a. a -> Query s a
Applicative, forall s. Applicative (Query s)
forall a. a -> Query s a
forall s a. a -> Query s a
forall a b. Query s a -> Query s b -> Query s b
forall a b. Query s a -> (a -> Query s b) -> Query s b
forall s a b. Query s a -> Query s b -> Query s b
forall s a b. Query s a -> (a -> Query s b) -> Query s b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
return :: forall a. a -> Query s a
$creturn :: forall s a. a -> Query s a
>> :: forall a b. Query s a -> Query s b -> Query s b
$c>> :: forall s a b. Query s a -> Query s b -> Query s b
>>= :: forall a b. Query s a -> (a -> Query s b) -> Query s b
$c>>= :: forall s a b. Query s a -> (a -> Query s b) -> Query s b
Monad)

-- | Run a query computation from an initial state.
runQueryM :: Scope -> Query s a -> (a, GenState)
runQueryM :: forall s a. Int -> Query s a -> (a, GenState)
runQueryM Int
scope = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall s a. State s a -> s -> (a, s)
runState (Int -> GenState
initState Int
scope) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s a. Query s a -> State GenState a
unQ

-- | Run a query computation in isolation, but reusing the current name supply.
isolate :: Query s a -> State GenState (GenState, a)
isolate :: forall s a. Query s a -> State GenState (GenState, a)
isolate (Query State GenState a
q) = do
  GenState
st <- forall s (m :: * -> *). MonadState s m => m s
get
  forall s (m :: * -> *). MonadState s m => s -> m ()
put forall a b. (a -> b) -> a -> b
$ (Int -> GenState
initState (GenState -> Int
nameScope GenState
st)) {nameSupply :: Int
nameSupply = GenState -> Int
nameSupply GenState
st}
  a
x <- State GenState a
q
  GenState
st' <- forall s (m :: * -> *). MonadState s m => m s
get
  forall s (m :: * -> *). MonadState s m => s -> m ()
put forall a b. (a -> b) -> a -> b
$ GenState
st {nameSupply :: Int
nameSupply = GenState -> Int
nameSupply GenState
st'}
  forall (m :: * -> *) a. Monad m => a -> m a
return (GenState
st', a
x)

-- | SQL generation internal state.
--   Contains the subqueries and static (i.e. not dependent on any subqueries)
--   restrictions of the query currently being built, as well as a name supply
--   for column renaming.
data GenState = GenState
  { GenState -> [SQL]
sources         :: ![SQL]
  , GenState -> [Exp SQL Bool]
staticRestricts :: ![Exp SQL Bool]
  , GenState -> [SomeCol SQL]
groupCols       :: ![SomeCol SQL]
  , GenState -> Int
nameSupply      :: !Int
  , GenState -> Int
nameScope       :: !Int
  }

-- | Initial state: no subqueries, no restrictions.
initState :: Int -> GenState
initState :: Int -> GenState
initState Int
scope = GenState
  { sources :: [SQL]
sources = []
  , staticRestricts :: [Exp SQL Bool]
staticRestricts = []
  , groupCols :: [SomeCol SQL]
groupCols = []
  , nameSupply :: Int
nameSupply = Int
0
  , nameScope :: Int
nameScope  = Int
scope
  }

renameAll :: [UntypedCol sql] -> State GenState [SomeCol sql]
renameAll :: forall sql. [UntypedCol sql] -> State GenState [SomeCol sql]
renameAll = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall sql. UntypedCol sql -> State GenState [SomeCol sql]
rename

-- | Generate a unique name for the given column.
rename :: UntypedCol sql -> State GenState [SomeCol sql]
rename :: forall sql. UntypedCol sql -> State GenState [SomeCol sql]
rename (Untyped Exp sql a
col) = do
    Name
n <- State GenState Name
freshId
    forall (m :: * -> *) a. Monad m => a -> m a
return [forall sql a. ColName -> Exp sql a -> SomeCol sql
Named (forall {a}. Show a => a -> ColName
newName Name
n) Exp sql a
col]
  where
    newName :: a -> ColName
newName a
ns =
      case Exp sql a
col of
        Col ColName
n -> ColName -> Text -> ColName
addColSuffix ColName
n forall a b. (a -> b) -> a -> b
$ Text
"_" forall a. Semigroup a => a -> a -> a
<> String -> Text
pack (forall a. Show a => a -> String
show a
ns)
        Exp sql a
_     -> Text -> ColName
mkColName forall a b. (a -> b) -> a -> b
$ Text
"tmp_" forall a. Semigroup a => a -> a -> a
<> String -> Text
pack (forall a. Show a => a -> String
show a
ns)

-- | Get a guaranteed unique identifier.
freshId :: State GenState Name
freshId :: State GenState Name
freshId = do
  GenState
st <- forall s (m :: * -> *). MonadState s m => m s
get
  forall s (m :: * -> *). MonadState s m => s -> m ()
put forall a b. (a -> b) -> a -> b
$ GenState
st {nameSupply :: Int
nameSupply = forall a. Enum a => a -> a
succ forall a b. (a -> b) -> a -> b
$ GenState -> Int
nameSupply GenState
st}
  forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> Int -> Name
Name (GenState -> Int
nameScope GenState
st) (GenState -> Int
nameSupply GenState
st))

-- | Get a guaranteed unique column name.
freshName :: State GenState ColName
freshName :: State GenState ColName
freshName = do
  Name
n <- State GenState Name
freshId
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Text -> ColName
mkColName forall a b. (a -> b) -> a -> b
$ Text
"tmp_" forall a. Semigroup a => a -> a -> a
<> String -> Text
pack (forall a. Show a => a -> String
show Name
n)