-----------------------------------------------------------------------------
-- |
-- Module      :  Data.TMap
-- Copyright   :  Peter Robinson 2009
-- License     :  LGPL
--
-- Maintainer  :  Peter Robinson <thaldyron@gmail.com>
-- Stability   :  experimental
-- Portability :  non-portable (requires STM)
--
-- Provides a thread-safe STM interface for finite map types with optional persistent
-- storage.
--
-----------------------------------------------------------------------------

module Data.TMap( -- * TMap Types
                  TMap, 
                  TFiniteMap,
                  -- * Creating a new TMap
                  newTMapIO, 
                  newTFiniteMapIO,
                  -- * Finite Map Interace
                  lookup, 
                  insert, 
                  delete, 
                  member, 
                  adjust, 
                  -- * Controlling the size of the TMap 
--                  purgeTMap, 
                  purgeTMapIO, 
                  getMaximumSize, 
                  setMaximumSize, 
                  getCurrentSize,
                  -- * Backend Communication
                  markAsDirty,
                  tryMarkAsDirty,
                  flushBackend,
                  -- * Exception Type
                  module Data.TMap.Exception,
                )
where

import Control.Concurrent.AdvSTM
import Control.Concurrent.AdvSTM.TVar
import Control.Monad( liftM, when )
import Control.Monad.Trans( MonadIO, liftIO )

import qualified Control.Exception as E
import Data.Maybe( isJust, isNothing )
import Prelude hiding (lookup,catch)

import qualified Data.TMap.Backend as B
import qualified Data.CacheStructure as C
import Data.CacheStructure.LRU(LRU) 
import Data.TMap.Exception( TMapException(..) )

import qualified Data.Edison.Assoc as M
import qualified Data.Edison.Assoc.StandardMap as FM


--------------------------------------------------------------------------------


data Entry a = Entry a         -- ^ Cache hit
             | NotInTMap       -- ^ Cache miss
             | NotInBackend    -- ^ Element exists neither in backend nor in cache
             | Exc E.SomeException -- ^ An exception occurred in the retry IO action
             deriving (Show)


instance Functor Entry where
  fmap _ NotInBackend = NotInBackend
  fmap _ NotInTMap    = NotInTMap
  fmap f (Entry a)    = Entry (f a)
  fmap _ (Exc e)      = Exc e

-- | The generic transactional map type.
data TMap map key val backendType cacheType = TMap 
  { backend  :: {- B.Backend k a b => -} backendType key val
  , sizeTVar :: TVar (Maybe Int)
  , tmapTVar :: {- (M.FiniteMapX map k, C.CacheStructure c k) 
             => -} TVar (map (Entry val),cacheType key)
  }

-- | The standard library type 'Data.Map' repackaged as a 'TMap'.
type TFiniteMap key val backendType = TMap (FM.FM key) key val backendType LRU

--------------------------------------------------------------------------------

-- | Creates a new TMap. You will need to use an apropriate backend and specify
-- the caching policy, e.g.,
--
-- @
--   import Data.TMap.Backend.Binary( BinaryBackend,mkBinaryBackend )
--   import Data.CacheStructure.LRU(LRU)
-- @
--
-- will use a binary-serialization backend for persistent storage and a \"least recently
-- used\" caching algorithm. See 'newTFiniteMapIO' for a less generic construction method.
--
-- Now, to create an unbounded map that uses the 'FM Int String' (see package EdisonCore) 
-- as the map type, you can write
--
-- @
--   backend <- mkBinaryBackend \"myworkdir\" \"mytempdir\"
--   tmap <- newTMapIO backend Nothing :: IO (TMap (FM.FM key) key val BinaryBackend LRU)
-- @
--
-- Note that 'newTFiniteMapIO' provides an easier construction method.
-- See file /Sample.hs/ for further examples.
newTMapIO :: (M.FiniteMapX map k, Ord k, B.Backend k a b,C.CacheStructure c k) 
          => b k a            -- ^ the backend
          -> Maybe Int        -- ^ maximum-size: Use 'Nothing' for unbounded size.
          -> IO (TMap map k a b c)
newTMapIO b maxsize = do 
  tvar <- newTVarIO (M.empty,C.empty)
  tvarSize <- newTVarIO maxsize
  B.initialize b
  return $ TMap b tvarSize tvar 


-- | Creates an (unbounded) 'TFiniteMap'. 
newTFiniteMapIO :: (Ord k, B.Backend k a b) 
          => b k a            -- ^ the backend
          -> IO (TFiniteMap k a b)
newTFiniteMapIO b = newTMapIO b Nothing


{-
 - Deactivated --- Can cause a non-terminating retry-loop when used with lookup k:
 - When key 'k' is not found lookup retries.  (Cond 1)
 - But this causes the creation of the tmap to be rolled back too, and so 
 - (Cond 1) holds forever.
newTMap :: (M.FiniteMapX map k, Ord k, B.Backend k a b, MonadAdvSTM stm) 
            => b -> stm (TMap map k a b)
newTMap b = do 
  tvar <- newTVar M.empty 
  return $ TMap b tvar 
-}


-- | Looks for a given key in the map and (if necessary) in the persistent storage 
-- and updates the map if necessary.
lookup :: (M.FiniteMapX map k, MonadAdvSTM stm, Ord k, B.Backend k a b,C.CacheStructure c k)
       => k -> TMap map k a b c -> stm (Maybe a) 
lookup k tmap = do
  (themap,accSeq) <- readTVar (tmapTVar tmap)
  case M.lookupWithDefault NotInTMap k themap of
    Entry v      -> do 
      writeTVar (tmapTVar tmap) (themap,C.hit k accSeq)
--      onCommit $ print ("OldAccess List: ",C.toList accSeq)
--      onCommit $ print ("NewAccess List: ",C.toList $ C.hit k accSeq)
      return $ Just v
    NotInBackend -> return Nothing
    Exc e        -> E.throw e
    NotInTMap    -> retryWith $ E.handle onExc $ do 
--      print $ "Lookup: Didn't find key retrying lookup..."
      result <- B.lookup (backend tmap) k 
      case result of 
        Nothing -> do 
--          print "Entry not in backend"
          atomically $ do 
            (themap',accSeq') <- readTVar (tmapTVar tmap)
            writeTVar (tmapTVar tmap) (M.insert k NotInBackend themap', accSeq')
        Just v  -> do 
--          print "Found entry in backend"
          atomically $ do 
            (themap',accSeq') <- readTVar (tmapTVar tmap)
            writeTVar (tmapTVar tmap) ( M.insert k (Entry v) themap'
                                      , C.hit k accSeq')
--            onCommit $ print ("Access List: ",C.toList $ C.hit k accSeq')
    where
      -- Sends backend exceptions back to the STM monad:
      onExc (e::E.SomeException) = do
        -- print "BACKEND EXCEPTION!"
        atomically $ do
          (themap',accSeq') <- readTVar (tmapTVar tmap)
          writeTVar (tmapTVar tmap) (M.insert k (Exc e) themap', accSeq')


-- | Checks whether the given key is in the map.
member :: (M.FiniteMapX map k, MonadAdvSTM stm, Ord k, B.Backend k a b,C.CacheStructure c k) 
       => k -> TMap map k a b c -> stm Bool
member k = liftM isJust . lookup k 


-- | Adds a key-value mapping to the map. Can throw a 'DuplicateEntry'
-- exception.
insert :: (M.FiniteMapX map k, MonadAdvSTM stm, Ord k, B.Backend k a b, C.CacheStructure c k) 
       => k -> a -> TMap map k a b c -> stm () 
insert k a tmap = do
  res <- lookup k tmap 
  when (isJust res) $ E.throw DuplicateEntry -- (show (k,v))
  (themap,accSeq) <- readTVar (tmapTVar tmap)
  writeTVar (tmapTVar tmap) (M.insert k (Entry a) themap, C.hit k accSeq)
  onCommit $ B.insert (backend tmap) k a


-- | Applies a function to the element identified by the key. Can throw an 'EntryNotFound' exception.
adjust :: (M.FiniteMapX map k, MonadAdvSTM stm, Ord k, B.Backend k a b, C.CacheStructure c k) 
       => (a -> a) -> k -> TMap map k a b c -> stm () 
adjust f k tmap = do
  res <- lookup k tmap 
  when (isNothing res) $ E.throw EntryNotFound 
  (themap,accSeq) <- readTVar (tmapTVar tmap)
  writeTVar (tmapTVar tmap) (M.adjust (fmap f) k themap, C.hit k accSeq)
  onCommit $ B.adjust (backend tmap) f k 


-- | Removes a key from the map. Can throw an 'EntryNotFound' exception.
delete :: (M.FiniteMapX map k, MonadAdvSTM stm, Ord k, B.Backend k a b,C.CacheStructure c k) 
       => k -> TMap map k a b c -> stm () 
delete k tmap = do
  res <- lookup k tmap 
  when (isNothing res) $ E.throw EntryNotFound 
  (themap,accSeq) <- readTVar (tmapTVar tmap)
  writeTVar (tmapTVar tmap) (M.insert k NotInBackend themap, accSeq)
  onCommit $ B.delete (backend tmap) k 

--------------------------------------------------------------------------------

-- | Reduces the map to the appropriate size if the maximum size was exceeded.
-- Calls /Data.TMap.Backend.flush/ if the map is purged.
-- Runs in /O(1)/ if the map size is within bounds, otherwise /O(n)/.
purgeTMapIO :: (M.FiniteMapX map k, MonadIO io, Ord k, B.Backend k a b, C.CacheStructure c k) 
            => TMap map k a b c -> io ()
purgeTMapIO = liftIO . atomically . purgeTMap 


-- | Reduces the map to the appropriate size if the maximum size was exceeded.
-- Calls /Data.TMap.Backend.flush/ if the map is purged.
-- Runs in /O(1)/ if the map size is within bounds, otherwise /O(n)/. 
-- /Warning:/ This function should only be called at the end of a transaction to
-- prevent nonterminating retry-loops!
purgeTMap :: (M.FiniteMapX map k, MonadAdvSTM stm, Ord k, B.Backend k a b, C.CacheStructure c k) 
          => TMap map k a b c -> stm () 
purgeTMap tmap =   do
  mSize <- readTVar (sizeTVar tmap)
  case mSize of
    Just maxSize -> do
      (themap,accSeq) <- readTVar (tmapTVar tmap)
--      onCommit $ print ("Old List: ",C.toList accSeq)
      when (C.size accSeq > maxSize) $ do
        let (restSeq,delSeq) = C.popMany (C.size accSeq - maxSize) accSeq
        writeTVar (tmapTVar tmap) (foldr M.delete themap delSeq, restSeq)
        onCommit $ B.flush (backend tmap)
--        onCommit $ print ("Purged List: ",C.toList restSeq)
    Nothing -> return ()

--------------------------------------------------------------------------------

-- | Sets the maximum size of the map. /O(1)/. Note that the size of the TMap needs
-- to be reduced manually to the maximum size by calling /purgeTMap/.
setMaximumSize :: (M.FiniteMapX map k, MonadAdvSTM stm, Ord k, B.Backend k a b, C.CacheStructure c k) 
               => TMap map k a b  c -> Int -> stm () 
setMaximumSize tmap maxSize
  | maxSize <= 0 = E.throw $ TMapDefaultExc "setMaximumSize: Invalid size specified."
  | otherwise    = writeTVar (sizeTVar tmap) $ Just maxSize

-- | Gets the maximum size of the map. /O(1)/.
getMaximumSize :: (M.FiniteMapX map k, MonadAdvSTM stm, Ord k, B.Backend k a b, C.CacheStructure c k) 
               => TMap map k a b c 
               -> stm (Maybe Int)
getMaximumSize tmap 
  | otherwise = readTVar (sizeTVar tmap) 


-- | Gets the current size of the map. /O(1)/.
getCurrentSize :: (M.FiniteMapX map k, MonadAdvSTM stm, Ord k, B.Backend k a b, C.CacheStructure c k) 
               => TMap map k a b c 
               -> stm Int
getCurrentSize tmap = do
  (_,accSeq) <- readTVar (tmapTVar tmap) 
  return $ C.size accSeq


-- | Causes the element to be reread from the backend on the next 'lookup'.
-- Throws an 'EntryNotFound' exception if the entry does not exist.
markAsDirty ::  (M.FiniteMapX map k, Ord k, B.Backend k a b, C.CacheStructure c k) 
          => k -> TMap map k a b c -> IO ()
markAsDirty k tmap = atomically $ do
  res <- lookup k tmap 
  when (isNothing res) $ E.throw EntryNotFound 
  (themap,accSeq) <- readTVar (tmapTVar tmap)
  writeTVar (tmapTVar tmap) (M.insert k NotInTMap themap, accSeq)

-- | Causes the element to be reread from the backend on the next 'lookup'. Does
-- not throw an error when the element does not exist.
tryMarkAsDirty ::  (M.FiniteMapX map k, Ord k, B.Backend k a b, C.CacheStructure c k) 
               => k -> TMap map k a b c -> IO ()
tryMarkAsDirty k tmap = 
  markAsDirty k tmap `E.catch` (\(e::TMapException) -> 
    if e == EntryNotFound then return ()
                          else E.throw e)

--------------------------------------------------------------------------------

-- | Sends a /B.flush/ request to the backend. Useful for asynchronous backend
-- implementations.
flushBackend :: (M.FiniteMapX map k, Ord k, B.Backend k a b, C.CacheStructure c k) 
             => TMap map k a b c -> IO ()
flushBackend = B.flush . backend