{-# Language GADTs #-}
module Network.AMQP.ChannelAllocator where

import qualified Data.Vector.Mutable as V
import Control.Exception (throwIO)
import Data.Word
import Data.Bits

import Network.AMQP.Types

data ChannelAllocator = ChannelAllocator Int -- highest permitted channel id

                                         (V.IOVector Word64)


newChannelAllocator :: Int -> IO ChannelAllocator
newChannelAllocator :: Int -> IO ChannelAllocator
newChannelAllocator Int
maxChannel =
    Int -> IOVector Word64 -> ChannelAllocator
ChannelAllocator Int
maxChannel (IOVector Word64 -> ChannelAllocator)
-> IO (IOVector Word64) -> IO ChannelAllocator
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Word64 -> IO (MVector (PrimState IO) Word64)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MVector (PrimState m) a)
V.replicate Int
1024 Word64
0

allocateChannel :: ChannelAllocator -> IO Int
allocateChannel :: ChannelAllocator -> IO Int
allocateChannel (ChannelAllocator Int
maxChannel IOVector Word64
c) = do
    Maybe Int
maybeIx <- IOVector Word64 -> IO (Maybe Int)
findFreeIndex IOVector Word64
c
    case Maybe Int
maybeIx of
        Just Int
chunk -> do
            Word64
word <- MVector (PrimState IO) Word64 -> Int -> IO Word64
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
V.read IOVector Word64
MVector (PrimState IO) Word64
c Int
chunk
            let offset :: Int
offset = Word64 -> Int
findUnsetBit Word64
word
            let channelID :: Int
channelID = Int
chunkInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
64 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
offset
            if Int
channelID Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
maxChannel
                then AMQPException -> IO Int
forall e a. Exception e => e -> IO a
throwIO (AMQPException -> IO Int) -> AMQPException -> IO Int
forall a b. (a -> b) -> a -> b
$ Int -> AMQPException
AllChannelsAllocatedException Int
maxChannel
                else do
                    MVector (PrimState IO) Word64 -> Int -> Word64 -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
V.write IOVector Word64
MVector (PrimState IO) Word64
c Int
chunk (Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
setBit Word64
word Int
offset)
                    Int -> IO Int
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
channelID
        Maybe Int
Nothing -> AMQPException -> IO Int
forall e a. Exception e => e -> IO a
throwIO (AMQPException -> IO Int) -> AMQPException -> IO Int
forall a b. (a -> b) -> a -> b
$ Int -> AMQPException
AllChannelsAllocatedException Int
maxChannel

freeChannel :: ChannelAllocator -> Int -> IO Bool
freeChannel :: ChannelAllocator -> Int -> IO Bool
freeChannel (ChannelAllocator Int
_maxChannel IOVector Word64
c) Int
ix = do
    let (Int
chunk, Int
offset) = Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
divMod Int
ix Int
64
    Word64
word <- MVector (PrimState IO) Word64 -> Int -> IO Word64
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
V.read IOVector Word64
MVector (PrimState IO) Word64
c Int
chunk
    if Word64 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Word64
word Int
offset
        then do
            MVector (PrimState IO) Word64 -> Int -> Word64 -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
V.write IOVector Word64
MVector (PrimState IO) Word64
c Int
chunk (Word64 -> IO ()) -> Word64 -> IO ()
forall a b. (a -> b) -> a -> b
$ Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
clearBit Word64
word Int
offset
            Bool -> IO Bool
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
        else Bool -> IO Bool
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False

findUnsetBit :: Word64 -> Int
findUnsetBit :: Word64 -> Int
findUnsetBit Word64
w = Int -> Int
go Int
0
  where
    go :: Int -> Int
go Int
65 = [Char] -> Int
forall a. HasCallStack => [Char] -> a
error [Char]
"findUnsetBit"
    go Int
ix | Bool -> Bool
not (Word64 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Word64
w Int
ix) = Int
ix
    go Int
ix = Int -> Int
go (Int
ixInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)

findFreeIndex :: V.IOVector Word64 -> IO (Maybe Int)
findFreeIndex :: IOVector Word64 -> IO (Maybe Int)
findFreeIndex IOVector Word64
vec = Int -> IO (Maybe Int)
go Int
0
  where
    -- TODO: make this faster

    go :: Int -> IO (Maybe Int)
go Int
1024 = Maybe Int -> IO (Maybe Int)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Int
forall a. Maybe a
Nothing
    go Int
ix = do
        Word64
v <- MVector (PrimState IO) Word64 -> Int -> IO Word64
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
V.read IOVector Word64
MVector (PrimState IO) Word64
vec Int
ix
        if Word64
v Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word64
0xffffffffffffffff
            then Maybe Int -> IO (Maybe Int)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Int -> IO (Maybe Int)) -> Maybe Int -> IO (Maybe Int)
forall a b. (a -> b) -> a -> b
$ Int -> Maybe Int
forall a. a -> Maybe a
Just Int
ix
            else Int -> IO (Maybe Int)
go (Int -> IO (Maybe Int)) -> Int -> IO (Maybe Int)
forall a b. (a -> b) -> a -> b
$! Int
ixInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1