{-# LANGUAGE CPP                   #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE LambdaCase            #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TupleSections         #-}
module Data.TableMethods where

import Data.Functor.Contravariant
import Data.Time.Clock.POSIX
import Data.Word
import qualified Data.IntMap.Strict as IntMap
         ;import Data.IntMap.Strict (IntMap)
import qualified Data.Map.Strict    as Map
         ;import Data.Map.Strict    (Map)
import qualified Data.Word64Map     as W64Map
         ;import Data.Word64Map     (Word64Map)

import Data.Wrapper.PSQ as PSQ

type Priority = POSIXTime

data OptionalPriority t tid x
    = NoPriority
    | HasPriority (Priority -> t x -> ([(tid, Priority, x)], t x))

-- | The standard lookup table methods.
data TableMethods t tid = TableMethods
    { -- | Insert a new /tid/ entry into the transaction table.
      tblInsert :: forall a. tid -> a -> Priority -> t a -> t a
      -- | Delete transaction /tid/ from the transaction table.
    , tblDelete :: forall a. tid -> t a -> t a
       -- | Lookup the value associated with transaction /tid/.
    , tblLookup :: forall a. tid -> t a -> Maybe a
    }

data QMethods t tid x = QMethods
    { qTbl           :: TableMethods t tid
    , qAtMostView    :: OptionalPriority t tid x
    }

vanillaTable :: TableMethods t tid -> QMethods t tid x
vanillaTable tbl = QMethods tbl NoPriority

priorityTable :: TableMethods t tid
                 -> (Priority -> t x -> ([(k, Priority, x)], t x))
                 -> (k -> x -> tid)
                 -> QMethods t tid x
priorityTable tbl atmost f = QMethods
    { qTbl        = tbl
    , qAtMostView = HasPriority $ \p t -> case atmost p t of
        (es,t') -> (map (\(k,p,a) -> (f k a, p, a)) es, t')
    }

-- | Methods for using 'Data.IntMap'.
intMapMethods :: TableMethods IntMap Int
intMapMethods = TableMethods
    { tblInsert     = \tid a p -> IntMap.insert tid a
    , tblDelete     = IntMap.delete
    , tblLookup     = IntMap.lookup
    }

-- | Methods for using 'Data.Word64Map'.
w64MapMethods :: TableMethods Word64Map Word64
w64MapMethods = TableMethods
    { tblInsert     = \tid a p -> W64Map.insert tid a
    , tblDelete     = W64Map.delete
    , tblLookup     = W64Map.lookup
    }

-- | Methods for using 'Data.Map'
mapMethods :: Ord tid => TableMethods (Map tid) tid
mapMethods = TableMethods
    { tblInsert     = \tid a p -> Map.insert tid a
    , tblDelete     = Map.delete
    , tblLookup     = Map.lookup
    }

-- psqMethods :: PSQKey tid => QMethods (HashPSQ tid Priority) tid x
psqMethods :: PSQKey k => (tid -> k) -> (k -> x -> tid) -> QMethods (PSQ' k Priority) tid x
psqMethods g f = priorityTable (contramap g tbl) PSQ.atMostView f
 where
    tbl :: PSQKey tid => TableMethods (PSQ' tid Priority) tid
    tbl = TableMethods
        { tblInsert = PSQ.insert'
        , tblDelete = PSQ.delete
        , tblLookup = \tid t -> case PSQ.lookup tid t of
            Just (p,a) -> Just a
            Nothing    -> Nothing
        }


-- | Change the key type for a lookup table implementation.
--
-- This can be used with 'intMapMethods' or 'mapMethods' to restrict lookups to
-- only a part of the generated /tid/ value.  This is useful for /tid/ types
-- that are especially large due their use for other purposes, such as secure
-- nonces for encryption.
instance Contravariant (TableMethods t) where
    -- contramap :: (tid -> t1) -> TableMethods t t1 -> TableMethods t tid
    contramap f (TableMethods ins del lookup) =
        TableMethods (\k p v t -> ins (f k) p v t)
                     (\k t     -> del (f k) t)
                     (\k t     -> lookup (f k) t)