283 lines
9.4 KiB
Haskell
283 lines
9.4 KiB
Haskell
{-# LANGUAGE BangPatterns #-}
|
|
{-# LANGUAGE DeriveDataTypeable #-}
|
|
{-# LANGUAGE TypeFamilies #-}
|
|
{-# LANGUAGE RebindableSyntax #-}
|
|
{-# LANGUAGE NoImplicitPrelude #-}
|
|
-- |
|
|
-- Module : Basement.Sized.Block
|
|
-- License : BSD-style
|
|
-- Maintainer : Haskell Foundation
|
|
--
|
|
-- A Nat-sized version of Block
|
|
{-# LANGUAGE AllowAmbiguousTypes #-}
|
|
{-# LANGUAGE CPP #-}
|
|
{-# LANGUAGE ConstraintKinds #-}
|
|
{-# LANGUAGE DataKinds #-}
|
|
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
|
|
{-# LANGUAGE ScopedTypeVariables #-}
|
|
{-# LANGUAGE TypeApplications #-}
|
|
{-# LANGUAGE TypeOperators #-}
|
|
#if __GLASGOW_HASKELL__ >= 806
|
|
{-# LANGUAGE NoStarIsType #-}
|
|
#endif
|
|
|
|
module Basement.Sized.Block
|
|
( BlockN
|
|
, MutableBlockN
|
|
, length
|
|
, lengthBytes
|
|
, toBlockN
|
|
, toBlock
|
|
, new
|
|
, newPinned
|
|
, singleton
|
|
, replicate
|
|
, thaw
|
|
, freeze
|
|
, index
|
|
, indexStatic
|
|
, map
|
|
, foldl'
|
|
, foldr
|
|
, cons
|
|
, snoc
|
|
, elem
|
|
, sub
|
|
, uncons
|
|
, unsnoc
|
|
, splitAt
|
|
, all
|
|
, any
|
|
, find
|
|
, reverse
|
|
, sortBy
|
|
, intersperse
|
|
, withPtr
|
|
, withMutablePtr
|
|
, withMutablePtrHint
|
|
, cast
|
|
, mutableCast
|
|
) where
|
|
|
|
import Data.Proxy (Proxy(..))
|
|
import Basement.Compat.Base
|
|
import Basement.Numerical.Additive (scale)
|
|
import Basement.Block (Block, MutableBlock(..), unsafeIndex)
|
|
import qualified Basement.Block as B
|
|
import qualified Basement.Block.Base as B
|
|
import Basement.Monad (PrimMonad, PrimState)
|
|
import Basement.Nat
|
|
import Basement.Types.OffsetSize
|
|
import Basement.NormalForm
|
|
import Basement.PrimType (PrimType, PrimSize, primSizeInBytes)
|
|
|
|
-- | Sized version of 'Block'
|
|
--
|
|
newtype BlockN (n :: Nat) a = BlockN { unBlock :: Block a }
|
|
deriving (NormalForm, Eq, Show, Data, Ord)
|
|
|
|
newtype MutableBlockN (n :: Nat) ty st = MutableBlockN { unMBlock :: MutableBlock ty st }
|
|
|
|
toBlockN :: forall n ty . (PrimType ty, KnownNat n, Countable ty n) => Block ty -> Maybe (BlockN n ty)
|
|
toBlockN b
|
|
| expected == B.length b = Just (BlockN b)
|
|
| otherwise = Nothing
|
|
where
|
|
expected = toCount @n
|
|
|
|
length :: forall n ty
|
|
. (KnownNat n, Countable ty n)
|
|
=> BlockN n ty
|
|
-> CountOf ty
|
|
length _ = toCount @n
|
|
|
|
lengthBytes :: forall n ty
|
|
. PrimType ty
|
|
=> BlockN n ty
|
|
-> CountOf Word8
|
|
lengthBytes = B.lengthBytes . unBlock
|
|
|
|
toBlock :: BlockN n ty -> Block ty
|
|
toBlock = unBlock
|
|
|
|
cast :: forall n m a b
|
|
. ( PrimType a, PrimType b
|
|
, KnownNat n, KnownNat m
|
|
, ((PrimSize b) * m) ~ ((PrimSize a) * n)
|
|
)
|
|
=> BlockN n a
|
|
-> BlockN m b
|
|
cast (BlockN b) = BlockN (B.unsafeCast b)
|
|
|
|
mutableCast :: forall n m a b st
|
|
. ( PrimType a, PrimType b
|
|
, KnownNat n, KnownNat m
|
|
, ((PrimSize b) * m) ~ ((PrimSize a) * n)
|
|
)
|
|
=> MutableBlockN n a st
|
|
-> MutableBlockN m b st
|
|
mutableCast (MutableBlockN b) = MutableBlockN (B.unsafeRecast b)
|
|
|
|
-- | Create a new unpinned mutable block of a specific N size of 'ty' elements
|
|
--
|
|
-- If the size exceeds a GHC-defined threshold, then the memory will be
|
|
-- pinned. To be certain about pinning status with small size, use 'newPinned'
|
|
new :: forall n ty prim
|
|
. (PrimType ty, KnownNat n, Countable ty n, PrimMonad prim)
|
|
=> prim (MutableBlockN n ty (PrimState prim))
|
|
new = MutableBlockN <$> B.new (toCount @n)
|
|
|
|
-- | Create a new pinned mutable block of a specific N size of 'ty' elements
|
|
newPinned :: forall n ty prim
|
|
. (PrimType ty, KnownNat n, Countable ty n, PrimMonad prim)
|
|
=> prim (MutableBlockN n ty (PrimState prim))
|
|
newPinned = MutableBlockN <$> B.newPinned (toCount @n)
|
|
|
|
singleton :: PrimType ty => ty -> BlockN 1 ty
|
|
singleton a = BlockN (B.singleton a)
|
|
|
|
replicate :: forall n ty . (KnownNat n, Countable ty n, PrimType ty) => ty -> BlockN n ty
|
|
replicate a = BlockN (B.replicate (toCount @n) a)
|
|
|
|
thaw :: (KnownNat n, PrimMonad prim, PrimType ty) => BlockN n ty -> prim (MutableBlockN n ty (PrimState prim))
|
|
thaw b = MutableBlockN <$> B.thaw (unBlock b)
|
|
|
|
freeze :: (PrimMonad prim, PrimType ty, Countable ty n) => MutableBlockN n ty (PrimState prim) -> prim (BlockN n ty)
|
|
freeze b = BlockN <$> B.freeze (unMBlock b)
|
|
|
|
indexStatic :: forall i n ty . (KnownNat i, CmpNat i n ~ 'LT, PrimType ty, Offsetable ty i) => BlockN n ty -> ty
|
|
indexStatic b = unsafeIndex (unBlock b) (toOffset @i)
|
|
|
|
index :: forall i n ty . PrimType ty => BlockN n ty -> Offset ty -> ty
|
|
index b ofs = B.index (unBlock b) ofs
|
|
|
|
map :: (PrimType a, PrimType b) => (a -> b) -> BlockN n a -> BlockN n b
|
|
map f b = BlockN (B.map f (unBlock b))
|
|
|
|
foldl' :: PrimType ty => (a -> ty -> a) -> a -> BlockN n ty -> a
|
|
foldl' f acc b = B.foldl' f acc (unBlock b)
|
|
|
|
foldr :: PrimType ty => (ty -> a -> a) -> a -> BlockN n ty -> a
|
|
foldr f acc b = B.foldr f acc (unBlock b)
|
|
|
|
cons :: PrimType ty => ty -> BlockN n ty -> BlockN (n+1) ty
|
|
cons e = BlockN . B.cons e . unBlock
|
|
|
|
snoc :: PrimType ty => BlockN n ty -> ty -> BlockN (n+1) ty
|
|
snoc b = BlockN . B.snoc (unBlock b)
|
|
|
|
sub :: forall i j n ty
|
|
. ( (i <=? n) ~ 'True
|
|
, (j <=? n) ~ 'True
|
|
, (i <=? j) ~ 'True
|
|
, PrimType ty
|
|
, KnownNat i
|
|
, KnownNat j
|
|
, Offsetable ty i
|
|
, Offsetable ty j )
|
|
=> BlockN n ty
|
|
-> BlockN (j-i) ty
|
|
sub block = BlockN (B.sub (unBlock block) (toOffset @i) (toOffset @j))
|
|
|
|
uncons :: forall n ty . (CmpNat 0 n ~ 'LT, PrimType ty, KnownNat n, Offsetable ty n)
|
|
=> BlockN n ty
|
|
-> (ty, BlockN (n-1) ty)
|
|
uncons b = (indexStatic @0 b, BlockN (B.sub (unBlock b) 1 (toOffset @n)))
|
|
|
|
unsnoc :: forall n ty . (CmpNat 0 n ~ 'LT, KnownNat n, PrimType ty, Offsetable ty n)
|
|
=> BlockN n ty
|
|
-> (BlockN (n-1) ty, ty)
|
|
unsnoc b =
|
|
( BlockN (B.sub (unBlock b) 0 (toOffset @n `offsetSub` 1))
|
|
, unsafeIndex (unBlock b) (toOffset @n `offsetSub` 1))
|
|
|
|
splitAt :: forall i n ty . (CmpNat i n ~ 'LT, PrimType ty, KnownNat i, Countable ty i) => BlockN n ty -> (BlockN i ty, BlockN (n-i) ty)
|
|
splitAt b =
|
|
let (left, right) = B.splitAt (toCount @i) (unBlock b)
|
|
in (BlockN left, BlockN right)
|
|
|
|
elem :: PrimType ty => ty -> BlockN n ty -> Bool
|
|
elem e b = B.elem e (unBlock b)
|
|
|
|
all :: PrimType ty => (ty -> Bool) -> BlockN n ty -> Bool
|
|
all p b = B.all p (unBlock b)
|
|
|
|
any :: PrimType ty => (ty -> Bool) -> BlockN n ty -> Bool
|
|
any p b = B.any p (unBlock b)
|
|
|
|
find :: PrimType ty => (ty -> Bool) -> BlockN n ty -> Maybe ty
|
|
find p b = B.find p (unBlock b)
|
|
|
|
reverse :: PrimType ty => BlockN n ty -> BlockN n ty
|
|
reverse = BlockN . B.reverse . unBlock
|
|
|
|
sortBy :: PrimType ty => (ty -> ty -> Ordering) -> BlockN n ty -> BlockN n ty
|
|
sortBy f b = BlockN (B.sortBy f (unBlock b))
|
|
|
|
intersperse :: (CmpNat n 1 ~ 'GT, PrimType ty) => ty -> BlockN n ty -> BlockN (n+n-1) ty
|
|
intersperse sep b = BlockN (B.intersperse sep (unBlock b))
|
|
|
|
toCount :: forall n ty . (KnownNat n, Countable ty n) => CountOf ty
|
|
toCount = natValCountOf (Proxy @n)
|
|
|
|
toOffset :: forall n ty . (KnownNat n, Offsetable ty n) => Offset ty
|
|
toOffset = natValOffset (Proxy @n)
|
|
|
|
-- | Get a Ptr pointing to the data in the Block.
|
|
--
|
|
-- Since a Block is immutable, this Ptr shouldn't be
|
|
-- to use to modify the contents
|
|
--
|
|
-- If the Block is pinned, then its address is returned as is,
|
|
-- however if it's unpinned, a pinned copy of the Block is made
|
|
-- before getting the address.
|
|
withPtr :: (PrimMonad prim, KnownNat n)
|
|
=> BlockN n ty
|
|
-> (Ptr ty -> prim a)
|
|
-> prim a
|
|
withPtr b = B.withPtr (unBlock b)
|
|
|
|
-- | Create a pointer on the beginning of the MutableBlock
|
|
-- and call a function 'f'.
|
|
--
|
|
-- The mutable block can be mutated by the 'f' function
|
|
-- and the change will be reflected in the mutable block
|
|
--
|
|
-- If the mutable block is unpinned, a trampoline buffer
|
|
-- is created and the data is only copied when 'f' return.
|
|
--
|
|
-- it is all-in-all highly inefficient as this cause 2 copies
|
|
withMutablePtr :: (PrimMonad prim, KnownNat n)
|
|
=> MutableBlockN n ty (PrimState prim)
|
|
-> (Ptr ty -> prim a)
|
|
-> prim a
|
|
withMutablePtr mb = B.withMutablePtr (unMBlock mb)
|
|
|
|
-- | Same as 'withMutablePtr' but allow to specify 2 optimisations
|
|
-- which is only useful when the MutableBlock is unpinned and need
|
|
-- a pinned trampoline to be called safely.
|
|
--
|
|
-- If skipCopy is True, then the first copy which happen before
|
|
-- the call to 'f', is skipped. The Ptr is now effectively
|
|
-- pointing to uninitialized data in a new mutable Block.
|
|
--
|
|
-- If skipCopyBack is True, then the second copy which happen after
|
|
-- the call to 'f', is skipped. Then effectively in the case of a
|
|
-- trampoline being used the memory changed by 'f' will not
|
|
-- be reflected in the original Mutable Block.
|
|
--
|
|
-- If using the wrong parameters, it will lead to difficult to
|
|
-- debug issue of corrupted buffer which only present themselves
|
|
-- with certain Mutable Block that happened to have been allocated
|
|
-- unpinned.
|
|
--
|
|
-- If unsure use 'withMutablePtr', which default to *not* skip
|
|
-- any copy.
|
|
withMutablePtrHint :: forall n ty prim a . (PrimMonad prim, KnownNat n)
|
|
=> Bool -- ^ hint that the buffer doesn't need to have the same value as the mutable block when calling f
|
|
-> Bool -- ^ hint that the buffer is not supposed to be modified by call of f
|
|
-> MutableBlockN n ty (PrimState prim)
|
|
-> (Ptr ty -> prim a)
|
|
-> prim a
|
|
withMutablePtrHint skipCopy skipCopyBack (MutableBlockN mb) f =
|
|
B.withMutablePtrHint skipCopy skipCopyBack mb f
|