{-# LANGUAGE EmptyDataDecls, GADTs, KindSignatures, Rank2Types, TypeOperators #-}
-- | Tables of values and keys for that tables.
--
-- Each value in the table may be accompanied with references to other tables.
module Data.HMemDb
    (
     MS,
     Multitude, Single, Multiple,
-- * Main structures
     Table, Key,
-- * Value references
     TableVarU, TableVar, TableVars, fromList, toList, readVar, readRefs,
-- * Specifications
     Spec (Spec, sRefs, sKeys),
-- ** Foreign table references
     TableRef, only, some,
     RefsC, Refs (Refs), RefsComponent, Ref, (:&:)((:&:)), splitRef,
-- ** Keys
     KeySpec, single, multiple, single_, multiple_,
     KeysC, Keys (Keys), KeysComponent, KeyRef, (:+:)((:+:)), splitKey,
-- * Table manipulation
     Created (Created),
     createTable, select, select_, selectBetween, insert, update, update_, delete,
-- * Persistence
     getTable, getTable_, getTable__,
     putTable, putTable_, putTable__
    ) where
import Control.Concurrent.STM (STM, TVar, modifyTVar', newTVar, readTVar, writeTVar)
import Control.Monad (forM, forM_, guard, liftM, liftM2, replicateM)
import Control.Monad.STM.Class (MonadSTM, liftSTM)
import Control.Monad.Trans.Class (lift)
import Control.Monad.Trans.Maybe (MaybeT (MaybeT))
import Data.Binary (Binary (get, put), Get, Put)
import Data.Functor.Identity (Identity (Identity, runIdentity))
import qualified Data.Map as M
    (Map, empty,
     elems, fromList, toList,
     alter, delete, insert, lookup, update,
     maxViewWithKey, minViewWithKey, splitLookup)
import Data.Maybe (fromMaybe)
import qualified Data.Set as S (Set, delete, fromList, insert, null, singleton, toList)
liftMaybe :: Monad m => Maybe a -> MaybeT m a
liftMaybe = MaybeT . return
-- | 'STM' that can fail.
-- Note that it doesn't revert the transaction on failure.
type MS = MaybeT STM
-- | This type specifies that we want a single value.
newtype Single = Single {sVal :: Integer} deriving (Eq, Ord)
-- | This type specifies that we want multiple values.
newtype Multiple = Multiple {mVal :: S.Set Integer} deriving Eq
-- | Closed class.
-- It's instances allow us to choose whether we want to get a single value
-- or multiple ones.
class Binary u => Multitude u where
    mToList :: u -> [Integer]
    mSingleton :: Integer -> u
    mInsert :: Integer -> u -> Maybe u -- Nothing means failure
    mDelete :: Integer -> u -> Maybe u -- Nothing means emptyness
instance Binary Single where
    get = fmap Single get
    put = put . sVal
instance Multitude Single where
    mToList = return . sVal
    mSingleton = Single
    mInsert _ _ = Nothing
    mDelete n s = guard (n /= sVal s) >> return s
instance Binary Multiple where
    get = fmap Multiple get
    put = put . mVal
instance Multitude Multiple where
    mToList = S.toList . mVal
    mSingleton = Multiple . S.singleton
    mInsert n u = return $ u {mVal = S.insert n $ mVal u}
    mDelete n u =
        let s = S.delete n $ mVal u in if S.null s then Nothing else Just (Multiple s)
-- | Base type for 'TableVar' and 'TableVars'
-- Type 't' is an abstract type, same as in the 'Table'.
-- Type 'a' is a type of value, which can be obtained with 'unVar',
-- also same as in the 'Table'.
data TableVarU t a u = TableVar {tvVal :: u} deriving (Eq, Ord)
-- | Reference to a single value in some table.
type TableVar t a = TableVarU t a Single
-- | Reference to multiple values in a single table.
type TableVars t a = TableVarU t a Multiple
-- | Function that converts a list of single-value references
-- to a single multiple-value reference.
-- Normally it should only be used in 'cInsert' statments.
fromList :: [TableVar t a] -> TableVars t a
fromList vs = TableVar $ Multiple $ S.fromList $ map (sVal . tvVal) vs
-- | Function that converts a multiple-value reference
-- to a list of single-value references.
-- Should be used with multiple-value references accompanying values in the 'Table'.
toList :: TableVars t a -> [TableVar t a]
toList v = map (TableVar . Single) $ S.toList $ mVal $ tvVal v
data KeyBack r a i u =
    KeyBack
    {
      kbMap :: TVar (M.Map i u),
      kbKey :: a -> r TableVarU -> i
    }
data PreTable t r k a =
    PreTable
    {
      tMap :: TVar (M.Map Integer (TVar (a, r TableVarU))),
      tKey :: k (KeyBack r a)
    }
-- | Class of key specifications, used in the 'sKeys' field of the 'Spec'.
class KeysC k where
    forKeys
        :: Monad m =>
           k f
        -> (forall i u. (Multitude u, Ord i) => f i u -> m (g i u))
        -> m (k g)
-- | Empty key specification.
-- It doesn't specify any key whatsoever.
data Keys (f :: * -> * -> *) = Keys
instance KeysC Keys where forKeys ~Keys _ = return Keys
-- | One key specification.
-- Note that it can't be used in the 'sKeys' field by itself,
-- but rather should be combined with 'Keys' with the ':+:' operator.
data KeyRef i u
-- | Combining operator for key specifications.
data (ks :+: k) f where (:+:) :: ks f -> f i u -> (ks :+: KeyRef i u) f
infixl 5 :+:
-- | Splitting keys.
splitKey :: (ks :+: KeyRef i u) f -> (ks f, f i u)
splitKey (ksf :+: fiu) = (ksf, fiu)
-- | Class of the part of key specification, corresponding to one key.
class KeysComponent k where
    forKeysComponent
        :: (KeysC ks, Monad m) =>
           (ks :+: k) f
        -> (forall i u. (Multitude u, Ord i) => f i u -> m (g i u))
        -> m ((ks :+: k) g)
instance (KeysC ks, KeysComponent k) => KeysC (ks :+: k) where forKeys = forKeysComponent
instance (Multitude u, Ord i) => KeysComponent (KeyRef i u) where
    forKeysComponent (ksf :+: fiu) action =
        liftM2 (:+:) (forKeys ksf action) (action fiu)
-- | Class of table reference specifications, used in the 'sRefs' field of the 'Spec'.
class RefsC r where
    putRefs
        :: Monad m =>
           r f
        -> (forall t a u. Multitude u => f t a u -> m ())
        -> m ()
    getRefs
        :: Monad m =>
           (forall t a u. Multitude u => m (f t a u))
        -> m (r f)
-- | Empty reference specification.
-- It doesn't specify any reference whatsoever.
data Refs (f :: * -> * -> * -> *) = Refs
instance RefsC Refs where
    putRefs ~Refs _ = return ()
    getRefs _ = return Refs
-- | One table reference specification.
-- Note that it can't be used in the 'sRefs' field by itself,
-- but rather should be combined with 'Refs' with the ':&:' operator.
data Ref t a u
-- | Combining operator for reference specifications.
data (rs :&: r) f where (:&:) :: rs f -> f t a u -> (rs :&: Ref t a u) f
infix 5 :&:
-- | Splitting references.
splitRef :: (rs :&: Ref t a u) f -> (rs f, f t a u)
splitRef (rsf :&: ftau) = (rsf, ftau)
-- | Class of the part of reference specification, corresponding to one reference.
class RefsComponent r where
    putRefsComponent
        :: (RefsC rs, Monad m) =>
           (rs :&: r) f
        -> (forall t a u. Multitude u => f t a u -> m ())
        -> m ()
    getRefsComponent
        :: (RefsC rs, Monad m) =>
           (forall t a u. Multitude u => m (f t a u))
        -> m ((rs :&: r) f)
instance (RefsC rs, RefsComponent r) => RefsC (rs :&: r) where
    putRefs = putRefsComponent
    getRefs = getRefsComponent
instance Multitude u => RefsComponent (Ref t a u) where
    putRefsComponent (rsf :&: ftau) action = putRefs rsf action >> action ftau
    getRefsComponent action = liftM2 (:&:) (getRefs action) action
-- | Abstract type, which represents a collection of values of type 'a',
-- possibly accompanied with some references to other 'Table's.
-- The type 't' is an abstract type, used to ensure that we don't confuse
-- different tables with values of the same type.
-- 'r' is a type of references accompanying each value.
data Table t r a where
    Table :: (KeysC k, RefsC r) => PreTable t r k a -> TVar Integer -> Table t r a
-- | Abstract type, which allows us to 'select' one or many values from the 'Table'.
-- Type 't' is an abstract type, same as in the 'Table'.
-- Type 'a' is a type of values, also same as in the 'Table'.
-- Type 'i' is a type of index values, used by this key.
-- Type 'u' is either 'Multiple' or 'Single', depending on whether this key
-- allows different values to have the same index, or not.
newtype Key t a i u = Key {kVal :: TVar (M.Map i u)}
-- | Type that is a template for the key. Used only in 'Spec's.
-- Type 't' is an abstract type, same as in the 'Table'.
-- Type 'a' is a type of values in that 'Table'.
-- Type 'i' is a type of index values, used by this key.
-- Type 'u' is either 'Multiple' or 'Single', depending on whether this key
-- allows different values to have the same index, or not.
newtype KeySpec r a i u = KeySpec {ksVal :: a -> r TableVarU -> i}
-- | This is a more generic version of 'single'.
-- The difference is that value index will be calculated based on both the value
-- and it's accompanying references.
single_ :: (a -> r TableVarU -> i) -> KeySpec r a i Single
single_ = KeySpec
-- | This is a more generic version of 'multiple'.
-- The difference is that value index will be calculated based on both the value
-- and it's accompanying references.
multiple_ :: (a -> r TableVarU -> i) -> KeySpec r a i Multiple
multiple_ = KeySpec
-- | This key will provide access to a single value within a 'Table'.
-- It's index will be calculated, based on this value alone.
single :: (a -> i) -> KeySpec r a i Single
single f = single_ $ const . f
-- | This key will provide access to multiple values in the same 'Table'.
-- Their indices will be calculated based on the value alone.
multiple :: (a -> i) -> KeySpec r a i Multiple
multiple f = multiple_ $ const . f
-- | Type that is a template for references to another table. Used only in 'Spec's.
-- Type 't' is an abstract type, same as in the 'Table'.
-- Type 'a' is a type of values in that 'Table'.
-- Type 'u' is either 'Single' or 'Multiple',
-- depending on whether the reference, accompanying the value,
-- should be single-value or multiple-value
data TableRef t a u = TableRef
-- | Each value in the table-to-be should be accompanied with a single-value reference.
only :: Table t r a -> TableRef t a Single
only = const TableRef
-- | Each value in the table-to-be should be accompanied with a multiple-value reference.
some :: Table t r a -> TableRef t a Multiple
some = const TableRef
-- | Type of table specifications.
data Spec r k a =
    Spec
    {
      sRefs :: r TableRef,
      -- ^ Other tables that should be referenced
      -- by values of this one.
      sKeys :: k (KeySpec r a) -- ^ Keys for the table-to-be
    }
-- | Output of the 'createTable' function. Contains the created table and the keys to it.
data Created r k a where Created :: Table t r a -> k (Key t a) -> Created r k a
data KeyProcess r a i u =
    KeyProcess
    {
      kpBack :: KeyBack r a i u,
      kpMap :: M.Map i u
    }
insertMap :: (Multitude u, Ord k) => Integer -> k -> M.Map k u -> Maybe (M.Map k u)
insertMap n i km =
    case M.lookup i km of
      Nothing -> return $ M.insert i (mSingleton n) km
      Just u -> flip (M.insert i) km `fmap` mInsert n u
forKeys_
    :: (KeysC k, Monad m) =>
       k f
    -> (forall i u. (Multitude u, Ord i) => f i u -> m ())
    -> m ()
forKeys_ ks action = forKeys ks (\k -> action k >> return k) >> return ()
-- | Function that creates the table (along with keys and everything) based on a 'Spec'.
createTable :: (KeysC k, RefsC r) => Spec r k a -> STM (Created r k a)
createTable s = 
    do counter <- newTVar 0
       tm <- newTVar M.empty
       tk <-
           forKeys (sKeys s) $ \ks ->
               do kbm <- newTVar M.empty
                  return KeyBack {kbMap = kbm, kbKey = ksVal ks}
       let cTable = Table PreTable {tMap = tm, tKey = tk} counter
           cKeys = runIdentity $ forKeys tk $ Identity . Key . kbMap
       return $ Created cTable cKeys
-- | Function that selects one value from a 'Key'.
-- Note that the value is not returned directly.
-- Instead, a reference to it is returned, which allows to get other references,
-- accompanying that value in the 'Table'.
select :: Ord i => Key t a i Single -> i -> MS (TableVar t a)
select k i = fmap TableVar $ lift (readTVar $ kVal k) >>= liftMaybe . M.lookup i
listUnMaybe :: Maybe [a] -> [a]
listUnMaybe Nothing = []
listUnMaybe (Just as) = as
-- | A more generic version of 'select'. Instead of one value, it returns multiple ones.
-- It can also select values with indices that are smaller or greater to the provided one,
-- depending on the third argument, which could be anything like @(>)@, @(<=)@, @(/=)@,
-- or even @return True@.
--
-- @
-- select_ k i (==) ~~ [select k i]
-- @
select_ ::
    (Multitude u, Ord i)
    => Key t a i u
    -> i
    -> (forall o. Ord o => o -> o -> Bool)
    -> STM [TableVar t a]
select_ k i c =
    do kv <- readTVar $ kVal k
       let ~(l, e, g) = M.splitLookup i kv
           lvs =
               do ~((li, _), _) <- M.minViewWithKey l
                  guard $ i `c` li
                  return $ M.elems l >>= mToList
           evs =
               do u <- e
                  guard $ i `c` i
                  return $ mToList u
           gvs =
               do ~((gi, _), _) <- M.maxViewWithKey g
                  guard $ i `c` gi
                  return $ M.elems g >>= mToList
       return $ map (TableVar . Single) $ [lvs, evs, gvs] >>= listUnMaybe
-- | A variant of 'select_', which allows to choose two bounds for the index.
-- Additional boolean arguments show whether to include bounds themselves or not.
selectBetween
    :: (Multitude u, Ord i) =>
       Key t a i u
    -> i -- ^ lower bound
    -> Bool -- ^ including lower bound?
    -> i -- ^ upper bound
    -> Bool -- ^ including upper bound?
    -> STM [TableVar t a]
selectBetween k il bl ig bg =
    do kv <- readTVar $ kVal k
       let ~(_, l, mgu) = M.splitLookup il kv
           ~(m, g, _) = M.splitLookup ig mgu
           lvs = if bl then fmap mToList l else Nothing
           mvs = return $ M.elems m >>= mToList
           gvs = if bg then fmap mToList g else Nothing
       return $ map (TableVar . Single) $ [lvs, mvs, gvs] >>= listUnMaybe
-- | Function that lets one to insert a new value to the 'Table'.
-- Of course, we have to provide accompanying references as well.
-- This function can fail if some key clashes with an already existing one.
insert :: Table t r a -> a -> r TableVarU -> MS (TableVar t a)
insert (Table pt counter) a r =
    do c <- lift $ readTVar counter
       kps <-
           forKeys (tKey pt) $ \kb ->
               do km <- lift $ readTVar $ kbMap kb
                  km' <- liftMaybe $ insertMap c (kbKey kb a r) km
                  return KeyProcess {kpBack = kb, kpMap = km'}
       lift $ do
         writeTVar counter $! c + 1
         forKeys_ kps $ \kp -> writeTVar (kbMap $ kpBack kp) $ kpMap kp
         pr <- newTVar (a, r)
         modifyTVar' (tMap pt) $ M.insert c pr
         return $ TableVar $ Single c
-- | Function that dereferences a value from table.
-- Note that we have to provide the 'Table' along with 'TableVar'.
readVar :: Table t r a -> TableVar t a -> MS a
readVar (Table pt _) v =
    do mp <- lift $ readTVar $ tMap pt
       pr <- liftMaybe $ M.lookup (sVal $ tvVal v) mp
       ~(a, _) <- lift $ readTVar pr
       return a
-- | Function that reads all references accompanying the value.
readRefs :: Table t r a -> TableVar t a -> MS (r TableVarU)
readRefs (Table pr _) v =
    fmap snd $ lift (readTVar $ tMap pr) >>=
    liftMaybe . M.lookup (sVal $ tvVal v) >>= lift . readTVar
-- | More generic version of 'update'.
-- It allows changing accompanying references as well as the value.
update_ :: Table t r a -> TableVar t a -> a -> r TableVarU -> MS ()
update_ (Table pt _) v a r =
    do let n = sVal $ tvVal v
       pr <- lift (readTVar $ tMap pt) >>= liftMaybe . M.lookup n
       ~(a', r') <- lift $ readTVar pr
       kps <-
           forKeys (tKey pt) $ \kb ->
               do km <- lift $ readTVar $ kbMap kb
                  km' <-
                      liftMaybe $
                      insertMap n (kbKey kb a r) $
                      M.update (mDelete n) (kbKey kb a' r') km
                  return KeyProcess {kpBack = kb, kpMap = km'}
       lift $ do
         forKeys_ kps $ \kp -> writeTVar (kbMap $ kpBack kp) $ kpMap kp
         writeTVar pr (a, r)
-- | Function that writes another value to the referenced place in the 'Table'.
-- It doesn't change the accompanying references.
-- In case that it fails due to some single-value key prohibiting the new value,
-- nothing is changed, and the 'Table' remains the same.
update :: Table t r a -> TableVar t a -> a -> MS ()
update t v a = readRefs t v >>= update_ t v a
-- | Function that removes the value (along with accompanying references)
-- from the 'Table'. It only fails if the value was already removed.
delete :: Table t r a -> TableVar t a -> MS ()
delete (Table pt _) v =
    do let n = sVal $ tvVal v
       tm <- lift $ readTVar $ tMap pt
       pr <- liftMaybe $ M.lookup n tm
       lift $ do
         ~(a, r) <- readTVar pr
         forKeys_ (tKey pt) $ \kb ->
             modifyTVar' (kbMap kb) $ M.update (mDelete n) (kbKey kb a r)
         writeTVar (tMap pt) $! M.delete n tm
-- | The most generic version of 'getTable'.
-- Not only it allows to change the way values are serialized,
-- it also permits side-effects during the deserialization.
-- The table is still filled in one 'STM' transaction,
-- thus avoiding any difficulties with multithreading.
getTable__ :: (Monad m, MonadSTM m) => Get (m a) -> Table t r a -> Get (m ())
getTable__ g (Table pt c) =
    do l <- get
       listM <-
           replicateM l $ do
             i <- get :: Get Integer
             ma <- g
             r <- getRefs $ liftM TableVar get
             return (i, ma, r)
       n <- get
       return $ do
         list <- forM listM $ \ ~(i, ma, r) -> liftM (\a -> (i, a, r)) ma
         let result =
                 do forKeys_ (tKey pt) $ \kb -> writeTVar (kbMap kb) M.empty
                    tm <-
                        forM list $ \ ~(i, a, r) ->
                            do pr <- newTVar (a, r)
                               forKeys_ (tKey pt) $ \kb ->
                                   modifyTVar' (kbMap kb) $
                                   flip M.alter (kbKey kb a r) $ Just . \mu ->
                                       case mu of
                                         Nothing -> mSingleton i
                                         Just u -> fromMaybe u $ mInsert i u
                               return (i, pr)
                    writeTVar (tMap pt) $ M.fromList tm
                    writeTVar c n
         liftSTM result
-- | More generic version of 'getTable'
-- that allows to change the way values are serialized.
getTable_ :: Get a -> Table t r a -> Get (STM ())
getTable_ g = getTable__ $ fmap return g
-- | Function that makes it possible to read the table from the file or other source.
-- Table should be created beforehand, as specifications are not serializable.
getTable :: Binary a => Table t r a -> Get (STM ())
getTable = getTable_ get
-- | The most generic version of 'putTable'.
-- Not only it allows to change the way values are serialized,
-- it also permits side-effects during the serialization.
-- The table is still read in one 'STM' transaction,
-- thus avoiding any difficulties with multithreading.
putTable__ :: (Monad m, MonadSTM m) => (a -> m Put) -> Table t r a -> m Put
putTable__ p (Table pt c) =
    do ~(listM, n) <-
           liftSTM $ do
             tm <- readTVar $ tMap pt
             list <-
                 forM (M.toList tm) $ \ ~(i, v) ->
                 do ~(a, r) <- readTVar v
                    return (i, a, r)
             n <- readTVar c
             return (list, n)
       list <- forM listM $ \ ~(i, a, r) -> liftM (\pa -> (i, pa, r)) $ p a
       return $ do
         put $ length list
         forM_ list $ \ ~(i, pa, r) ->
             do put i
                pa
                putRefs r $ \v -> put (tvVal v)
         put n
-- | More generic version of 'putTable'
-- that allows to change the way values are serialized.
putTable_ :: (a -> Put) -> Table t r a -> STM Put
putTable_ p = putTable__ $ return . p
-- | Function that makes it possible to write the table to the file or other storage.
putTable :: Binary a => Table t r a -> STM Put
putTable = putTable_ put