module Control.Concurrent.STM.TSkipList(
TSkipList,newIO,new,
insert,lookup,update,delete,geq,leq,filter,
chooseLevel,
toString,
)
where
import Control.Concurrent.STM
import GHC.Conc
import Control.Applicative
import Control.Monad
import Control.Exception
import System.Random
import Data.Array.MArray
import Data.Map(Map)
import qualified Data.Map as M
import Prelude hiding(filter,lookup)
type ForwardPtrs k a = TArray Int (Node k a)
data TSkipList k a = TSkipList
{ maxLevel :: Int
, probability :: Float
, curLevel :: TVar Int
, listHead :: ForwardPtrs k a
}
data Node k a
= Nil
| Node { key :: k
, contentTVar :: TVar a
, forwardPtrs :: ForwardPtrs k a
}
isNil :: Node k a -> Bool
isNil Nil = True
isNil _ = False
newIO :: Float
-> Int
-> IO (TSkipList k a)
newIO p maxLvl = atomically $ new p maxLvl
new :: Float
-> Int
-> STM (TSkipList k a)
new p maxLvl =
TSkipList maxLvl p `liftM` newTVar 1
`ap` newForwardPtrs maxLvl
newForwardPtrs :: Int -> STM (ForwardPtrs k a)
newForwardPtrs maxLvl = newListArray (1,maxLvl) $ replicate maxLvl Nil
chooseLevel :: TSkipList k a -> STM Int
chooseLevel tskip = do
stdG <- unsafeIOToSTM newStdGen
let rs :: StdGen -> [Float]
rs g = x : rs g' where (x,g') = randomR (0,1) g
let samples = take (maxLevel tskip 1) (rs stdG)
return $ 1 + length (takeWhile (probability tskip <) samples)
leq :: (Ord k) => k -> TSkipList k a -> STM (Map k a)
leq k tskip =
leqAcc (listHead tskip) 1 M.empty
where
leqAcc fwdPtrs lvl curAcc = do
let moveDown acc _ level =
leqAcc fwdPtrs (level1) acc
let moveRight acc succNode level =
addElem acc succNode >>=
leqAcc (forwardPtrs succNode) level
let onFound acc succNode _ =
addElem acc succNode
traverse k fwdPtrs lvl (moveDown curAcc) (moveRight curAcc) (onFound curAcc) (moveDown curAcc) curAcc
addElem acc succNode = do
a <- readTVar (contentTVar succNode)
return $ M.insert (key succNode) a acc
geq :: (Ord k) => k -> TSkipList k a -> STM (Map k a)
geq k = filter (\k' _ -> (k'>=k))
lookupNode :: (Ord k) => k -> TSkipList k a -> STM (Maybe (Node k a))
lookupNode k tskip =
lookupAcc (listHead tskip) =<< readTVar (curLevel tskip)
where
lookupAcc fwdPtrs lvl = do
let moveDown _ level = lookupAcc fwdPtrs (level1)
let moveRight succNode = lookupAcc (forwardPtrs succNode)
let onFound succNode _ = return (Just succNode)
traverse k fwdPtrs lvl moveDown moveRight onFound moveDown Nothing
lookup :: (Ord k) => k -> TSkipList k a -> STM (Maybe a)
lookup k tskip =
maybe (return Nothing)
(\n -> Just <$> readTVar (contentTVar n)) =<< lookupNode k tskip
update :: (Ord k) => k -> a -> TSkipList k a -> STM ()
update k a tskip =
maybe (throw $ AssertionFailed "TSkipList.update: element not found!")
(flip writeTVar a . contentTVar) =<< lookupNode k tskip
delete :: (Ord k) => k -> TSkipList k a -> STM ()
delete k tskip =
deleteAcc (listHead tskip) =<< readTVar (curLevel tskip)
where
deleteAcc fwdPtrs lvl = do
let moveDown _ level = deleteAcc fwdPtrs (level1)
let moveRight succNode = deleteAcc (forwardPtrs succNode)
let onFound succNode level = do
succsuccNode <- readArray (forwardPtrs succNode) level
writeArray fwdPtrs level succsuccNode
moveDown succNode level
traverse k fwdPtrs lvl moveDown moveRight onFound moveDown ()
insert :: (Ord k) => k -> a -> TSkipList k a -> STM ()
insert k a tskip = do
mNode <- lookupNode k tskip
case mNode of
Just node -> writeTVar (contentTVar node) a
Nothing -> do
tvar <- newTVar a
newPtrs <- newForwardPtrs (maxLevel tskip)
let node = Node k tvar newPtrs
insertNode k node tskip
insertNode :: (Ord k) => k -> Node k a -> TSkipList k a -> STM ()
insertNode k node tskip = do
newLevel <- chooseLevel tskip
curLvl <- readTVar (curLevel tskip)
when (curLvl < newLevel) $
writeTVar (curLevel tskip) newLevel
insertAcc (listHead tskip) newLevel
where
insertAcc fwdPtrs lvl = do
let moveDown succNode level = do
writeArray (forwardPtrs node) level succNode
writeArray fwdPtrs level node
insertAcc fwdPtrs (level1)
let moveRight succNode =
insertAcc (forwardPtrs succNode)
let onFound _ level = do
writeArray fwdPtrs level node
insertAcc fwdPtrs (level1)
traverse k fwdPtrs lvl moveDown moveRight onFound moveDown ()
traverse :: (Ord k)
=> k -> ForwardPtrs k a -> Int
-> (Node k a -> Int -> STM b)
-> (Node k a -> Int -> STM b)
-> (Node k a -> Int -> STM b)
-> (Node k a -> Int -> STM b)
-> b
-> STM b
traverse k fwdPtrs level onLT onGT onFound onNil def
| level < 1 = return def
| otherwise = do
succNode <- readArray fwdPtrs level
if isNil succNode
then onNil succNode level
else case k `compare` key succNode of
GT -> onGT succNode level
LT -> onLT succNode level
EQ -> onFound succNode level
filter :: (Ord k)
=> (k -> a -> Bool) -> TSkipList k a -> STM (Map k a)
filter p tskip =
filterAcc (listHead tskip) 1 M.empty
where
filterAcc fwdPtrs level acc = do
succNode <- readArray fwdPtrs level
if isNil succNode
then return acc
else do
newAcc <- addElem acc succNode
filterAcc (forwardPtrs succNode) level newAcc
addElem acc succNode = do
a <- readTVar (contentTVar succNode)
return $ if p (key succNode) a
then M.insert (key succNode) a acc
else acc
toString :: (Show k,Ord k) => k -> TSkipList k a -> STM String
toString k tskip = do
curLvl <- readTVar (curLevel tskip)
ls <- forM (reverse [1..curLvl]) $ printAcc (listHead tskip) []
return $ unlines ls
where
printAcc fwdPtrs acc curLvl = do
let moveDown succNode level =
if (isNil succNode)
then return acc
else printAcc (forwardPtrs succNode) acc level
let moveRight succNode level = do
let n = (' ':show (key succNode))
printAcc (forwardPtrs succNode) (acc++n) level
let onFound succNode level = do
let n = (' ':show (key succNode))
printAcc (forwardPtrs succNode) (acc++n) level
traverse k fwdPtrs curLvl moveDown moveRight onFound moveDown ""