{-# LANGUAGE BangPatterns , MagicHash , UnboxedTuples #-}
module Utilities (
    -- * Utility Chans
    -- ** Indexed MVars
      IndexedMVar()
    , newIndexedMVar, putMVarIx, readMVarIx, tryReadMVarIx
    -- * Other stuff
    , nextHighestPowerOfTwo
    , touchIORef
    ) where

import Control.Concurrent.MVar
import Control.Exception
import Control.Applicative
import Data.Bits
import Data.Word
import Data.Atomics
import Data.IORef
import GHC.Prim(touch#)
import GHC.IORef(IORef(..))
import GHC.STRef(STRef(..))
import GHC.Base(IO(..))
import Prelude

-- For now: a reverse-ordered assoc list; an IntMap might be better
newtype IndexedMVar a = IndexedMVar (IORef [(Int, MVar a)])

newIndexedMVar :: IO (IndexedMVar a)
newIndexedMVar :: IO (IndexedMVar a)
newIndexedMVar = IORef [(Int, MVar a)] -> IndexedMVar a
forall a. IORef [(Int, MVar a)] -> IndexedMVar a
IndexedMVar (IORef [(Int, MVar a)] -> IndexedMVar a)
-> IO (IORef [(Int, MVar a)]) -> IO (IndexedMVar a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Int, MVar a)] -> IO (IORef [(Int, MVar a)])
forall a. a -> IO (IORef a)
newIORef []



-- these really suck; sorry.

readMVarIx :: IndexedMVar a -> Int -> IO a
{-# INLINE readMVarIx #-}
readMVarIx :: IndexedMVar a -> Int -> IO a
readMVarIx IndexedMVar a
mvIx Int
i = do
    MVar a -> IO a
forall a. MVar a -> IO a
readMVar (MVar a -> IO a) -> IO (MVar a) -> IO a
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IndexedMVar a -> Int -> IO (MVar a)
forall a. IndexedMVar a -> Int -> IO (MVar a)
getMVarIx IndexedMVar a
mvIx Int
i

tryReadMVarIx :: IndexedMVar a -> Int -> IO (Maybe a)
{-# INLINE tryReadMVarIx #-}
tryReadMVarIx :: IndexedMVar a -> Int -> IO (Maybe a)
tryReadMVarIx IndexedMVar a
mvIx Int
i = do
    MVar a -> IO (Maybe a)
forall a. MVar a -> IO (Maybe a)
tryReadMVar (MVar a -> IO (Maybe a)) -> IO (MVar a) -> IO (Maybe a)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IndexedMVar a -> Int -> IO (MVar a)
forall a. IndexedMVar a -> Int -> IO (MVar a)
getMVarIx IndexedMVar a
mvIx Int
i

putMVarIx :: IndexedMVar a -> Int -> a -> IO ()
{-# INLINE putMVarIx #-}
putMVarIx :: IndexedMVar a -> Int -> a -> IO ()
putMVarIx IndexedMVar a
mvIx Int
i a
a = do
    (MVar a -> a -> IO ()) -> a -> MVar a -> IO ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip MVar a -> a -> IO ()
forall a. MVar a -> a -> IO ()
putMVar a
a (MVar a -> IO ()) -> IO (MVar a) -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IndexedMVar a -> Int -> IO (MVar a)
forall a. IndexedMVar a -> Int -> IO (MVar a)
getMVarIx IndexedMVar a
mvIx Int
i

-- NOTE: this uses atomic actions to stay async exception safe:
getMVarIx :: IndexedMVar a -> Int -> IO (MVar a)
{-# INLINE getMVarIx #-}
getMVarIx :: IndexedMVar a -> Int -> IO (MVar a)
getMVarIx (IndexedMVar IORef [(Int, MVar a)]
v) Int
i = do
    -- We're right to optimistically create this for readMVarIx, but throw this
    -- away for most putMVarIx (from writers), probably.
    MVar a
mv <- IO (MVar a)
forall a. IO (MVar a)
newEmptyMVar
    Ticket [(Int, MVar a)]
tk0 <- IORef [(Int, MVar a)] -> IO (Ticket [(Int, MVar a)])
forall a. IORef a -> IO (Ticket a)
readForCAS IORef [(Int, MVar a)]
v
    let go :: Ticket [(Int, MVar a)] -> IO (MVar a)
go Ticket [(Int, MVar a)]
tk = do
            let !xs :: [(Int, MVar a)]
xs = Ticket [(Int, MVar a)] -> [(Int, MVar a)]
forall a. Ticket a -> a
peekTicket Ticket [(Int, MVar a)]
tk
            case Int -> MVar a -> [(Int, MVar a)] -> Either (MVar a) [(Int, MVar a)]
forall mvar.
Int -> mvar -> [(Int, mvar)] -> Either mvar [(Int, mvar)]
findInsert Int
i MVar a
mv [(Int, MVar a)]
xs of
                 Left MVar a
alreadyPresentMVar -> MVar a -> IO (MVar a)
forall (m :: * -> *) a. Monad m => a -> m a
return MVar a
alreadyPresentMVar
                 Right [(Int, MVar a)]
xs' -> do 
                    (Bool
success,Ticket [(Int, MVar a)]
newTk) <- IORef [(Int, MVar a)]
-> Ticket [(Int, MVar a)]
-> [(Int, MVar a)]
-> IO (Bool, Ticket [(Int, MVar a)])
forall a. IORef a -> Ticket a -> a -> IO (Bool, Ticket a)
casIORef IORef [(Int, MVar a)]
v Ticket [(Int, MVar a)]
tk [(Int, MVar a)]
xs'
                    if Bool
success 
                        then MVar a -> IO (MVar a)
forall (m :: * -> *) a. Monad m => a -> m a
return MVar a
mv
                        else Ticket [(Int, MVar a)] -> IO (MVar a)
go Ticket [(Int, MVar a)]
newTk
    Ticket [(Int, MVar a)] -> IO (MVar a)
go Ticket [(Int, MVar a)]
tk0

-- Reverse-sorted:
findInsert :: Int -> mvar -> [(Int,mvar)] -> Either mvar [(Int,mvar)]
{-# INLINE findInsert #-}
findInsert :: Int -> mvar -> [(Int, mvar)] -> Either mvar [(Int, mvar)]
findInsert Int
i mvar
mv = [(Int, mvar)] -> Either mvar [(Int, mvar)]
ins where
    ins :: [(Int, mvar)] -> Either mvar [(Int, mvar)]
ins [] = [(Int, mvar)] -> Either mvar [(Int, mvar)]
forall a b. b -> Either a b
Right [(Int
i,mvar
mv)] 
    ins xss :: [(Int, mvar)]
xss@((Int
i',mvar
x):[(Int, mvar)]
xs) = 
               case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
i Int
i' of
                    Ordering
GT -> [(Int, mvar)] -> Either mvar [(Int, mvar)]
forall a b. b -> Either a b
Right ([(Int, mvar)] -> Either mvar [(Int, mvar)])
-> [(Int, mvar)] -> Either mvar [(Int, mvar)]
forall a b. (a -> b) -> a -> b
$ (Int
i,mvar
mv)(Int, mvar) -> [(Int, mvar)] -> [(Int, mvar)]
forall a. a -> [a] -> [a]
:[(Int, mvar)]
xss
                    Ordering
EQ -> mvar -> Either mvar [(Int, mvar)]
forall a b. a -> Either a b
Left mvar
x
                    Ordering
LT -> ([(Int, mvar)] -> [(Int, mvar)])
-> Either mvar [(Int, mvar)] -> Either mvar [(Int, mvar)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Int
i',mvar
x)(Int, mvar) -> [(Int, mvar)] -> [(Int, mvar)]
forall a. a -> [a] -> [a]
:) (Either mvar [(Int, mvar)] -> Either mvar [(Int, mvar)])
-> Either mvar [(Int, mvar)] -> Either mvar [(Int, mvar)]
forall a b. (a -> b) -> a -> b
$ [(Int, mvar)] -> Either mvar [(Int, mvar)]
ins [(Int, mvar)]
xs


-- Not particularly fast; if needs moar fast see
--   http://graphics.stanford.edu/~seander/bithacks.html#RoundUpPowerOf2
-- 
nextHighestPowerOfTwo :: Int -> Int
nextHighestPowerOfTwo :: Int -> Int
nextHighestPowerOfTwo Int
0 = Int
1
nextHighestPowerOfTwo Int
n 
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
maxPowerOfTwo = [Char] -> Int
forall a. HasCallStack => [Char] -> a
error ([Char] -> Int) -> [Char] -> Int
forall a b. (a -> b) -> a -> b
$ [Char]
"The next power of two greater than "[Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++(Int -> [Char]
forall a. Show a => a -> [Char]
show Int
n)[Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++[Char]
" exceeds the highest value representable by Int."
    | Bool
otherwise = 
        let !nhp2 :: Int
nhp2 = Int
2 Int -> Int -> Int
forall a b. (Num a, Integral b) => a -> b -> a
^ (Float -> Int
forall a b. (RealFrac a, Integral b) => a -> b
ceiling (Float -> Float -> Float
forall a. Floating a => a -> a -> a
logBase Float
2 (Float -> Float) -> Float -> Float
forall a b. (a -> b) -> a -> b
$ Int -> Float
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Float) -> Int -> Float
forall a b. (a -> b) -> a -> b
$ Int -> Int
forall a. Num a => a -> a
abs Int
n :: Float) :: Int)
         -- ensure return value is actually a positive power of 2:
         in Bool -> Int -> Int
forall a. HasCallStack => Bool -> a -> a
assert (Int
nhp2 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Word -> Int
forall a. Bits a => a -> Int
popCount (Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
nhp2 :: Word) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1)
              Int
nhp2

  where maxPowerOfTwo :: Int
maxPowerOfTwo = (Float -> Int
forall a b. (RealFrac a, Integral b) => a -> b
floor (Float -> Int) -> Float -> Int
forall a b. (a -> b) -> a -> b
$ Float -> Float
forall a. Floating a => a -> a
sqrt (Float -> Float) -> Float -> Float
forall a b. (a -> b) -> a -> b
$ (Int -> Float
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
forall a. Bounded a => a
maxBound :: Int)::Float)) Int -> Int -> Int
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2::Int)

-- I'm not sure what happens if we try to use touch from
-- Control.Monad.Primitive on our boxed IORef (if it gets unboxed), so we do
-- this:
touchIORef :: IORef a -> IO ()
touchIORef :: IORef a -> IO ()
touchIORef (IORef (STRef MutVar# RealWorld a
v)) = (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, () #)) -> IO ())
-> (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s -> 
    case MutVar# RealWorld a -> State# RealWorld -> State# RealWorld
forall a. a -> State# RealWorld -> State# RealWorld
touch# MutVar# RealWorld a
v State# RealWorld
s of 
         State# RealWorld
s' -> (# State# RealWorld
s', () #)