-----------------------------------------------------------------------------
-- |
-- Module      :  Control.Concurrent.TBox.TSkipList
-- Copyright   :  Peter Robinson 2010
-- License     :  LGPL
-- 
-- Maintainer  :  Peter Robinson <robinson@ecs.tuwien.ac.at>
-- Stability   :  experimental
-- Portability :  non-portable (requires STM)
--
-- Provides an implementation of a skip list in the 'AdvSTM' monad.
-- A skip list is a probabilistic data structure with map-like operations.
-- In contrast to a balanced tree, a skip list does not need any rebalancing,
-- which makes it suitable for concurrent programming.
-- See: /William Pugh. Skip Lists: A Probabilistic Alternative to Balanced Trees./
--
-- The elements of the skip list are stored in a 'TBox'.
-- When an element of the skip list is modified, the operation is relegated
-- to the corresponding 'TBox'.
--
-- For a concrete instance see module 'Control.Concurrent.TFile.TSkipList'
-----------------------------------------------------------------------------
 
module Control.Concurrent.TBox.TSkipList(-- * Data type 
                                         TSkipList,newIO,
                                         -- * Operations
                                         insert,lookup,update,delete,geq,leq,min,filter,
                                         -- * Low-level Operations
                                         insertNode,lookupNode,readAndValidate,newNode,
                                         contentTBox,key,
                                         -- * Utilities 
                                         chooseLevel,
                                         toString,
                                        ) 

where
import Control.Concurrent.TBox(TBox)
import qualified Control.Concurrent.TBox as TBox

import Control.Exception
import Control.Concurrent.AdvSTM.TVar
import Control.Concurrent.AdvSTM.TArray
import Control.Concurrent.AdvSTM
import Control.Applicative
import Control.Monad
import Control.Monad.IfElse(unlessM)

import System.Random
import Data.Maybe
--import Data.List(unlines)
import Data.Map(Map)
import qualified Data.Map as M
import Data.Array.MArray
import Prelude hiding(lookup,filter,catch,min)

type ForwardPtrs t k a = TArray Int (Node t k a)


data TSkipList t k a = TSkipList 
  { maxLevel    :: Int
  , probability :: Float
  , curLevel    :: TVar Int
  , listHead    :: ForwardPtrs t k a
  }

data Node t k a
  = Nil 
  | Node { key          :: k 
         , contentTBox  :: t k a 
         , forwardPtrs  :: ForwardPtrs t k a
         }

newNode :: TBox t k a => k -> t k a -> Int -> AdvSTM (Node t k a)
newNode k t maxLvl = Node k t `liftM` newForwardPtrs maxLvl

isNil :: Node t k a -> Bool
isNil Nil = True
isNil _   = False

-- | An empty skiplist.
newIO :: TBox t k a 
           => Float  -- ^ Probability for choosing a new level
           -> Int    -- ^ Maximum number of levels
           -> IO (TSkipList t k a)
newIO p maxLvl = 
  atomically $ new p maxLvl


-- | An empty skiplist.
new :: TBox t k a 
    => Float -- ^ Probability for choosing a new level
    -> Int   -- ^ Maximum number of levels
    -> AdvSTM (TSkipList t k a)
new p maxLvl = 
  TSkipList maxLvl p `liftM` {- newTVar stdG 
                        `ap` -} newTVar 1 
                        `ap` newForwardPtrs maxLvl


newForwardPtrs :: Int -> AdvSTM (ForwardPtrs t k a)
newForwardPtrs maxLvl = newListArray (1,maxLvl) $ replicate maxLvl Nil


-- | Returns a randomly chosen level. Is used for inserting new elements.
-- Note that this function uses 'unsafeIOToAdvSTM' to access the random
-- number generator.
chooseLevel :: TSkipList t k a -> AdvSTM Int
chooseLevel tskip = do
  stdG <- unsafeIOToAdvSTM 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) 


-- | Returns all elements that are smaller than the key. 
leq :: ({- Show k,-}Ord k, TBox t k a) => k -> TSkipList t k a -> AdvSTM (Map k a)
leq k tskip = 
  leqAcc (listHead tskip) 1 M.empty
  where
  leqAcc fwdPtrs lvl curAcc = do
    let moveDown acc _ level  = 
          leqAcc fwdPtrs (level-1) acc 
    let moveRight acc succNode level = do 
          newAcc <- addElem acc succNode
          leqAcc (forwardPtrs succNode) level newAcc  
    let onFound acc succNode _ = 
          addElem acc succNode
    traverse k fwdPtrs lvl (moveDown curAcc) (moveRight curAcc) (onFound curAcc) (moveDown curAcc) curAcc

  addElem acc succNode = 
    maybe acc (\a -> M.insert (key succNode) a acc) <$> readAndValidate tskip succNode


-- | Returns all elements that are greater than the key.
-- TODO: currently in O(n), can be made more efficient (like 'leq')
geq :: ({- Show k,-}Ord k, TBox t k a) => k -> TSkipList t k a -> AdvSTM (Map k a)
geq k = filter (\k' _ -> (k'>=k))

-- | Returns the element with the least key, if it exists. /O(1)/.
min :: (Ord k, TBox t k a) => TSkipList t k a -> AdvSTM (Maybe a)
min tskip = do
  node <- readArray (listHead tskip) 1 
  if isNil node 
    then return Nothing
    else readAndValidate tskip node 

-- | Reads the 'TBox' of the node. If the 'TBox' is empty, the node
-- is removed from the skip list.
-- This is necessary when 'TBox's are shared between different data
-- structures.  
readAndValidate :: (Ord k, TBox t k a) 
                => TSkipList t k a -> Node t k a -> AdvSTM (Maybe a)
readAndValidate tskip succNode = do
  ma <- TBox.read (contentTBox succNode)
  case ma of
    Just a  -> return $ Just a
    Nothing -> do
      delete (key succNode) tskip
      return Nothing
  
{-    
  maxAcc (listHead tskip) 1 M.empty
  where
  maxAcc fwdPtrs level acc = do
    succNode <- readArray fwdPtrs level 
    if isNil succNode 
      then return Nothing
      else do
        if isNil succsuccNode 
          then return (key succNode)
          else maxAcc (forwardPtrs succNode) level newAcc
        newAcc <- addElem acc succNode
        filterAcc (forwardPtrs succNode) level newAcc
-}


lookupNode :: ({- Show k,-}Ord k, TBox t k a) => k -> TSkipList t k a -> AdvSTM (Maybe (Node t k a))
lookupNode k tskip = 
  lookupAcc (listHead tskip) =<< readTVar (curLevel tskip)
  where
  lookupAcc fwdPtrs lvl = do
    let moveDown _ level  = lookupAcc fwdPtrs (level-1)
    let moveRight succNode = lookupAcc (forwardPtrs succNode) 
    let onFound succNode _   = return (Just succNode)
    traverse k fwdPtrs lvl moveDown moveRight onFound moveDown Nothing


lookup :: ({- Show k,-}Ord k, TBox t k a) => k -> TSkipList t k a -> AdvSTM (Maybe a)
lookup k tskip = 
  maybe (return Nothing)
        (readAndValidate tskip) =<< lookupNode k tskip


-- | Updates an element. Throws 'AssertionFailed' if the element is not in the
-- list.
update :: ({- Show k,-}Ord k, TBox t k a) => k -> a -> TSkipList t k a -> AdvSTM ()
update k a tskip = 
  maybe (throw $ AssertionFailed "TSkipList.update: element not found!") 
        (flip TBox.write a . contentTBox) =<< lookupNode k tskip 


delete :: ({- Show k,-}Ord k, TBox t k a) => k -> TSkipList t k a -> AdvSTM ()
delete k tskip = 
  deleteAcc (listHead tskip) =<<  readTVar (curLevel tskip)
  where
  deleteAcc fwdPtrs lvl = do
    let moveDown _ level = deleteAcc fwdPtrs (level-1)
    let moveRight succNode      = deleteAcc (forwardPtrs succNode) 
    let onFound succNode level  = do
          let tbox = contentTBox succNode
          unlessM (TBox.isEmptyNotDirty tbox) $
            TBox.clear (contentTBox succNode)
          succsuccNode <- readArray (forwardPtrs succNode) level 
          writeArray fwdPtrs level succsuccNode
          moveDown succNode level
    traverse k fwdPtrs lvl moveDown moveRight onFound moveDown ()


insert :: ({- Show k,-}Ord k, TBox t k a) => k -> a -> TSkipList t k a ->  AdvSTM ()
insert k a tskip = do
  -- Make new TBox:
  tbox    <- TBox.new k a 
  newPtrs <- newForwardPtrs (maxLevel tskip)
  let node = Node k tbox newPtrs
  insertNode k node tskip


insertNode :: ({- Show k,-}Ord k, TBox t k a) => k -> Node t k a -> TSkipList t k a ->  AdvSTM ()
insertNode k node tskip = do
  newLevel <-  chooseLevel tskip
  -- Adapt current maximum level:
  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 (level-1)
    let moveRight succNode = 
          insertAcc (forwardPtrs succNode) 
    let onFound _ level = do
          writeArray fwdPtrs level node
          insertAcc fwdPtrs (level-1)
    traverse k fwdPtrs lvl moveDown moveRight onFound moveDown ()


traverse :: ({- Show k,-}Ord k, TBox t k a) 
         => k -> ForwardPtrs t k a -> Int 
         -> (Node t k a -> Int -> AdvSTM b)
         -> (Node t k a -> Int -> AdvSTM b)
         -> (Node t k a -> Int -> AdvSTM b)
         -> (Node t k a -> Int -> AdvSTM b)
         -> b
         -> AdvSTM 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


-- | Returns all elements that satisfy the predicate. O(n).
filter :: ({- Show k,-}Ord k, TBox t k a) 
      => (k -> a -> Bool) -> TSkipList t k a -> AdvSTM (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 = 
      maybe acc (\a -> if p (key succNode) a 
                         then M.insert (key succNode) a acc
                         else acc) 
            <$> readAndValidate tskip succNode


-- | Debug helper. Returns the skip list as a string.
-- All elements smaller than the given key are written to the string.
toString :: (Ord k, Show k, TBox t k a) => k -> TSkipList t k a -> AdvSTM 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 ""