Block.hs

-- |
-- Module      : Crypto.Store.Block
-- License     : BSD-style
-- Maintainer  : Olivier Chéron <olivier.cheron@gmail.com>
-- Stability   : experimental
-- Portability : unknown
--
-- Minimal port of basement @Block@ data type.  Provides a typed interface on
-- top of @ByteArray#@.
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UnboxedTuples #-}
module Crypto.Store.Block
    ( Block, Offset(..), CountOf(..), PrimType(..)
    , createWithPtr, Crypto.Store.Block.map, unsafeCast, unsafeIndex
    ) where

import Data.Memory.Endian

import Data.Proxy
import Data.Word

import Foreign.Ptr (castPtr)

import GHC.Exts
import GHC.IO (IO(..))
import GHC.Word

import System.IO.Unsafe

data Block a = Block ByteArray#
data MutableBlock a = MutableBlock (MutableByteArray# RealWorld)

newtype CountOf a = CountOf Int deriving (Show, Eq, Ord)
newtype Offset a = Offset Int deriving (Show, Eq, Ord, Num)

class PrimType a where
    primSizeInBytes :: Proxy a -> CountOf Word8
    primBaIndex :: ByteArray# -> Offset a -> a
    primMbaWrite :: MutableByteArray# RealWorld -> Offset a -> a -> IO ()

instance PrimType Word8 where
    primSizeInBytes _ = CountOf 1
    {-# INLINE primSizeInBytes #-}
    primBaIndex ba (Offset (I# n)) = W8# (indexWord8Array# ba n)
    {-# INLINE primBaIndex #-}
    primMbaWrite mba (Offset (I# n)) (W8# w) = IO $ \s -> (# writeWord8Array# mba n w s, () #)
    {-# INLINE primMbaWrite #-}

instance PrimType Word16 where
    primSizeInBytes _ = CountOf 2
    {-# INLINE primSizeInBytes #-}
    primBaIndex ba (Offset (I# n)) = W16# (indexWord16Array# ba n)
    {-# INLINE primBaIndex #-}
    primMbaWrite mba (Offset (I# n)) (W16# w) = IO $ \s -> (# writeWord16Array# mba n w s, () #)
    {-# INLINE primMbaWrite #-}

instance PrimType a => PrimType (LE a) where
    primSizeInBytes _ = primSizeInBytes (Proxy :: Proxy a)
    {-# INLINE primSizeInBytes #-}
    primBaIndex ba (Offset i) = LE $ primBaIndex ba (Offset i)
    {-# INLINE primBaIndex #-}
    primMbaWrite mba (Offset i) (LE x) = primMbaWrite mba (Offset i) x
    {-# INLINE primMbaWrite #-}

create :: PrimType a => CountOf a -> (Offset a -> a) -> Block a
create n@(CountOf !sz) f = unsafeDupablePerformIO $ do
    mb <- new n
    loop mb 0
    unsafeFreeze mb
  where
    loop !mb i
        | i == sz = pure ()
        | otherwise =
            let off = Offset i
             in unsafeWrite mb off (f off) >> loop mb (i + 1)
{-# INLINE create #-}

createWithPtr :: CountOf Word8 -> (Ptr p -> IO a) -> Block Word8
createWithPtr n f = unsafeDupablePerformIO $ do
    b <- unsafeNewPinned n >>= unsafeFreeze
    f (castPtr $ unsafeBlockPtr b) *> touch b
    return b
{-# INLINE createWithPtr #-}

length :: forall a. PrimType a => Block a -> CountOf a
length (Block ba) = CountOf (I# (sizeofByteArray# ba) `quot` sz)
  where CountOf sz = primSizeInBytes (Proxy :: Proxy a)
{-# INLINE length #-}

map :: (PrimType a, PrimType b) => (a -> b) -> Block a -> Block b
map f b = create (CountOf len) $ \(Offset i) -> f (unsafeIndex b (Offset i))
  where CountOf len = Crypto.Store.Block.length b
{-# INLINE map #-}

new :: forall a. PrimType a => CountOf a -> IO (MutableBlock a)
new (CountOf n) = IO $ \s1 ->
    case newByteArray# bytes s1 of
        (# s2, mba #) -> (# s2, MutableBlock mba #)
  where
    !(I# bytes) = n * sz
    CountOf sz = primSizeInBytes (Proxy :: Proxy a)
{-# INLINE new #-}

touch :: Block a -> IO ()
touch (Block ba) = IO $ \s1 -> case touch# ba s1 of { s2 -> (# s2, () #) }

unsafeBlockPtr :: Block a -> Ptr a
unsafeBlockPtr (Block ba) = Ptr (byteArrayContents# ba)
{-# INLINE unsafeBlockPtr #-}

unsafeCast :: Block a -> Block b
unsafeCast (Block ba) = Block ba
{-# INLINE unsafeCast #-}

unsafeFreeze :: MutableBlock a -> IO (Block a)
unsafeFreeze (MutableBlock mba) = IO $ \s1 ->
    case unsafeFreezeByteArray# mba s1 of
        (# s2, ba #) -> (# s2, Block ba #)
{-# INLINE unsafeFreeze #-}

unsafeIndex :: PrimType a => Block a -> Offset a -> a
unsafeIndex (Block ba) = primBaIndex ba
{-# INLINE unsafeIndex #-}

unsafeNewPinned :: CountOf Word8 -> IO (MutableBlock a)
unsafeNewPinned (CountOf (I# bytes)) = IO $ \s1 ->
    case newAlignedPinnedByteArray# bytes 8# s1 of
        (# s2, mba #) -> (# s2, MutableBlock mba #)
{-# INLINE unsafeNewPinned #-}

unsafeWrite :: PrimType a => MutableBlock a -> Offset a -> a -> IO ()
unsafeWrite (MutableBlock mba) = primMbaWrite mba
{-# INLINE unsafeWrite #-}