primal-memory-0.3.0.0: src/Data/Prim/Memory/Fold.hs
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}
-- |
-- Module : Data.Prim.Memory.Fold
-- Copyright : (c) Alexey Kuleshevich 2020
-- License : BSD3
-- Maintainer : Alexey Kuleshevich <alexey@kuleshevi.ch>
-- Stability : experimental
-- Portability : non-portable
--
module Data.Prim.Memory.Fold where
import Data.Prim
import Data.Prim.Memory.Internal
foldlMem ::
forall e a mr. (Prim e, MemRead mr)
=> (a -> e -> a)
-- ^ Folding function
-> a
-- ^ Initial accumulator
-> mr
-- ^ Memory region to iterate over
-> a
foldlMem f = ifoldlMem (\a _ -> f a)
{-# INLINE foldlMem #-}
ifoldlMem ::
forall e a mr. (Prim e, MemRead mr)
=> (a -> Off e -> e -> a)
-- ^ Folding function
-> a
-- ^ Initial accumulator
-> mr
-- ^ Memory region to iterate over
-> a
ifoldlMem f initAcc mem = ifoldlOffMem 0 (countMem mem :: Count e) f initAcc mem
{-# INLINE ifoldlMem #-}
ifoldlOffMem ::
forall e a mr. (Prim e, MemRead mr)
=> Off e
-- ^ Initial offset to start at
-> Count e
-- ^ Total number of elements to iterate through
-> (a -> Off e -> e -> a)
-- ^ Folding function
-> a
-- ^ Initial accumulator
-> mr
-- ^ Memory region to iterate over
-> a
ifoldlOffMem off count f initAcc mem = loop initAcc off
where
k = countToOff count + off
loop !acc i
| i >= k = acc
| otherwise = loop (f acc i (indexOffMem mem i)) (i + 1)
{-# INLINE ifoldlOffMem #-}
foldlLazyMem ::
forall e a mr. (Prim e, MemRead mr)
=> (a -> e -> a)
-- ^ Folding function
-> a
-- ^ Initial accumulator
-> mr
-- ^ Memory region to iterate over
-> a
foldlLazyMem f = ifoldlLazyMem (\a _ -> f a)
{-# INLINE foldlLazyMem #-}
ifoldlLazyMem ::
forall e a mr. (Prim e, MemRead mr)
=> (a -> Off e -> e -> a)
-- ^ Folding function
-> a
-- ^ Initial accumulator
-> mr
-- ^ Memory region to iterate over
-> a
ifoldlLazyMem f initAcc mem = ifoldlLazyOffMem 0 (countMem mem :: Count e) f initAcc mem
{-# INLINE ifoldlLazyMem #-}
ifoldlLazyOffMem ::
forall e a mr. (Prim e, MemRead mr)
=> Off e
-- ^ Initial offset to start at
-> Count e
-- ^ Total number of elements to iterate through
-> (a -> Off e -> e -> a)
-- ^ Folding function
-> a
-- ^ Initial accumulator
-> mr
-- ^ Memory region to iterate over
-> a
ifoldlLazyOffMem off count f initAcc mem = loop initAcc off
where
k = countToOff count + off
loop acc i
| i >= k = acc
| otherwise = loop (f acc i (indexOffMem mem i)) (i + 1)
{-# INLINE ifoldlLazyOffMem #-}
foldrMem ::
forall e a mr. (Prim e, MemRead mr)
=> (e -> a -> a)
-- ^ Folding function
-> a
-- ^ Initial accumulator
-> mr
-- ^ Memory region to iterate over
-> a
foldrMem f = ifoldrMem (const f)
{-# INLINE foldrMem #-}
ifoldrMem ::
forall e a mr. (Prim e, MemRead mr)
=> (Off e -> e -> a -> a)
-- ^ Folding function
-> a
-- ^ Initial accumulator
-> mr
-- ^ Memory region to iterate over
-> a
ifoldrMem f initAcc mem = ifoldrOffMem 0 (countMem mem :: Count e) f initAcc mem
{-# INLINE ifoldrMem #-}
ifoldrOffMem ::
forall e a mr. (Prim e, MemRead mr)
=> Off e
-- ^ Initial offset to start at
-> Count e
-- ^ Total number of elements to iterate through
-> (Off e -> e -> a -> a)
-- ^ Folding function
-> a
-- ^ Initial accumulator
-> mr
-- ^ Memory region to iterate over
-> a
ifoldrOffMem off count f initAcc mem = loop initAcc off
where
k = countToOff count + off
loop !acc i
| i >= k = acc
| otherwise = f i (indexOffMem mem i) (loop acc (i + 1))
{-# INLINE ifoldrOffMem #-}
-- | Right fold with a lazy accumulator
--
-- @since 0.3.0
foldrLazyMem ::
forall e a mr. (Prim e, MemRead mr)
=> (e -> a -> a)
-- ^ Folding function
-> a
-- ^ Initial accumulator
-> mr
-- ^ Memory region to iterate over
-> a
foldrLazyMem f = ifoldrLazyMem (const f)
{-# INLINE foldrLazyMem #-}
-- | Right fold with a lazy accumulator using an offset aware function
--
-- @since 0.3.0
ifoldrLazyMem ::
forall e a mr. (Prim e, MemRead mr)
=> (Off e -> e -> a -> a)
-- ^ Folding function
-> a
-- ^ Initial accumulator
-> mr
-- ^ Memory region to iterate over
-> a
ifoldrLazyMem f initAcc mem =
ifoldrLazyOffMem 0 (countMem mem :: Count e) f initAcc mem
{-# INLINE ifoldrLazyMem #-}
ifoldrLazyOffMem ::
forall e a mr. (Prim e, MemRead mr)
=> Off e
-- ^ Initial offset to start at
-> Count e
-- ^ Total number of elements to iterate through
-> (Off e -> e -> a -> a)
-- ^ Folding function
-> a
-- ^ Initial accumulator
-> mr
-- ^ Memory region to iterate over
-> a
ifoldrLazyOffMem off count f initAcc mem = loop initAcc off
where
k = countToOff count + off
loop acc i
| i >= k = acc
| otherwise = f i (indexOffMem mem i) (loop acc (i + 1))
{-# INLINE ifoldrLazyOffMem #-}
foldMapOffMem ::
forall e m mr. (Prim e, MemRead mr, Monoid m)
=> Off e
-> Count e
-> (e -> m)
-> mr
-> m
foldMapOffMem off count f = ifoldrLazyOffMem off count (\_ e acc -> f e `mappend` acc) mempty
{-# INLINE foldMapOffMem #-}
ifoldMapOffMem ::
forall e m mr. (Prim e, MemRead mr, Monoid m)
=> Off e
-> Count e
-> (Off e -> e -> m)
-> mr
-> m
ifoldMapOffMem off count f =
ifoldrLazyOffMem off count (\i e acc -> f i e `mappend` acc) mempty
{-# INLINE ifoldMapOffMem #-}
anyOffMem ::
forall e mr. (Prim e, MemRead mr)
=> Off e
-> Count e
-> (e -> Bool)
-> mr
-> Bool
anyOffMem off count p = getAny #. foldMapOffMem off count (Any #. p)
{-# INLINE anyOffMem #-}
ianyOffMem ::
forall e mr. (Prim e, MemRead mr)
=> Off e
-> Count e
-> (Off e -> e -> Bool)
-> mr
-> Bool
ianyOffMem off count p = getAny #. ifoldMapOffMem off count (\i -> Any #. p i)
{-# INLINE ianyOffMem #-}
anyMem :: forall e mr . (Prim e, MemRead mr) => (e -> Bool) -> mr -> Bool
anyMem p xs = anyOffMem 0 (countMem xs :: Count e) p xs
{-# INLINE anyMem #-}
ianyMem :: forall e mr . (Prim e, MemRead mr) => (Off e -> e -> Bool) -> mr -> Bool
ianyMem p xs = ianyOffMem 0 (countMem xs :: Count e) p xs
{-# INLINE ianyMem #-}
allOffMem ::
forall e mr. (Prim e, MemRead mr)
=> Off e
-> Count e
-> (e -> Bool)
-> mr
-> Bool
allOffMem off count p = getAll #. foldMapOffMem off count (All #. p)
{-# INLINE allOffMem #-}
iallOffMem ::
forall e mr. (Prim e, MemRead mr)
=> Off e
-> Count e
-> (Off e -> e -> Bool)
-> mr
-> Bool
iallOffMem off count p = getAll #. ifoldMapOffMem off count (\i -> All #. p i)
{-# INLINE iallOffMem #-}
allMem :: forall e mr . (Prim e, MemRead mr) => (e -> Bool) -> mr -> Bool
allMem p xs = allOffMem 0 (countMem xs :: Count e) p xs
{-# INLINE allMem #-}
iallMem :: forall e mr . (Prim e, MemRead mr) => (Off e -> e -> Bool) -> mr -> Bool
iallMem p xs = iallOffMem 0 (countMem xs :: Count e) p xs
{-# INLINE iallMem #-}
---------
-- Dangerous: ignores the slack
eqMem :: forall e mr . (Prim e, Eq e, MemRead mr) => mr -> mr -> Bool
eqMem m1 m2
| isSameMem m1 m2 = True
| otherwise = n == countMem m2 && eqOffMem m1 0 m2 0 n
where
n = countMem m1 :: Count e
{-# INLINE eqMem #-}
-- | Check two regions of memory for equality using the `Eq` instance. It will return
-- `True` whenever both regions hold exactly the same elements and `False` as soon as the
-- first pair of mismatched elements is discovered in the two regions. It is safe for both
-- regions to refer to the same part of memory.
--
-- [Unsafe] When any precondition for either of the offsets @memOff1@, @memOff2@ or the
-- element count @memCount@ is violated the result is either unpredictable output or
-- failure with a segfault.
--
-- @since 0.3.0
eqOffMem ::
(Prim e, Eq e, MemRead mr1, MemRead mr2)
=> mr1 -- ^ /memRead1/ - First region of memory
-> Off e
-- ^ /memOff1/ - Offset for @memRead1@ in number of elements
--
-- /__Precondition:__/
--
-- > 0 <= memOff1
-> mr2 -- ^ /memRead2/ - Second region of memory
-> Off e
-- ^ /memOff2/ - Offset for @memRead1@ in number of elements
--
-- /__Precondition:__/
--
-- > 0 <= memOff2
-> Count e
-- ^ /memCount/ - Number of elements of type __@e@__ to compare
--
-- /__Preconditions:__/
--
-- > 0 <= memCount
--
-- > offToCount memOff1 + memCount < countMem memRead1
--
-- > offToCount memOff2 + memCount < countMem memRead2
-> Bool
eqOffMem m1 off1 m2 off2 count = loop off1
where
doff = off2 - off1
k = countToOff count + off1
loop !i
| i < k = indexOffMem m1 i == indexOffMem m2 (i + doff) && loop (i + 1)
| otherwise = True
{-# INLINE[1] eqOffMem #-}
eqOffMemBinary ::
forall e mr1 mr2. (Prim e, MemRead mr1, MemRead mr2)
=> mr1
-> Off e
-> mr2
-> Off e
-> Count e
-> Bool
eqOffMemBinary m1 off1 m2 off2 count =
eqByteOffMem m1 (toByteOff off1) m2 (toByteOff off2) (toByteCount count)
{-# INLINE eqOffMemBinary #-}
{-# RULES
"eqOffMem/Char" forall mr1 (off1 :: Off Char) . eqOffMem mr1 off1 = eqOffMemBinary mr1 off1
"eqOffMem/Word" forall mr1 (off1 :: Off Word) . eqOffMem mr1 off1 = eqOffMemBinary mr1 off1
"eqOffMem/Word8" eqOffMem = eqByteOffMem
"eqOffMem/Word16" forall mr1 (off1 :: Off Word16) . eqOffMem mr1 off1 = eqOffMemBinary mr1 off1
"eqOffMem/Word32" forall mr1 (off1 :: Off Word32) . eqOffMem mr1 off1 = eqOffMemBinary mr1 off1
"eqOffMem/Word64" forall mr1 (off1 :: Off Word64) . eqOffMem mr1 off1 = eqOffMemBinary mr1 off1
"eqOffMem/Int" forall mr1 (off1 :: Off Int) . eqOffMem mr1 off1 = eqOffMemBinary mr1 off1
"eqOffMem/Int8" forall mr1 (off1 :: Off Int8) . eqOffMem mr1 off1 = eqOffMemBinary mr1 off1
"eqOffMem/Int16" forall mr1 (off1 :: Off Int16) . eqOffMem mr1 off1 = eqOffMemBinary mr1 off1
"eqOffMem/Int32" forall mr1 (off1 :: Off Int32) . eqOffMem mr1 off1 = eqOffMemBinary mr1 off1
"eqOffMem/Int64" forall mr1 (off1 :: Off Int64) . eqOffMem mr1 off1 = eqOffMemBinary mr1 off1
#-}
eqOffMutMem ::
forall e ma1 ma2 m s. (Prim e, Eq e, MonadPrim s m, MemWrite ma1, MemWrite ma2)
=> ma1 s
-> Off e
-> ma2 s
-> Off e
-> Count e
-> m Bool
eqOffMutMem mm1 off1 mm2 off2 count = loop off1
where
doff = off2 - off1
k = countToOff count + off1
loop !i
| i < k = do
e1 <- readOffMutMem mm1 i
e2 <- readOffMutMem mm2 (i + doff)
if e1 == e2
then loop (i + 1)
else pure False
| otherwise = pure True
{-# INLINE eqOffMutMem #-}
-- | Compare two mutable memory regions for element equality. Regions themselves are not
-- modified, as such it is semantically similar to `eqMem` which works on immutable
-- regions.
eqMutMem ::
forall e ma m s. (Prim e, Eq e, MonadPrim s m, MemAlloc ma)
=> ma s
-> ma s
-> m Bool
eqMutMem mm1 mm2
| isSameMutMem mm1 mm2 = pure True
| otherwise = do
n1 <- getCountMutMem mm1
n2 <- getCountMutMem mm2
if n1 /= n2
then pure False
else eqOffMutMem mm1 0 mm2 0 (n1 :: Count e)
{-# INLINE eqMutMem #-}
-- | Compare two regions using the `Ord` instance. It will return `EQ` whenever both
-- regions hold exactly the same elements and `LT` or `GT` as soon as the first discovered
-- element that is less than or greater than respectfully in the first region when
-- compared to the second one. It is safe for both regions to refer to the same part of
-- memory.
--
-- @since 0.3.0
compareMem ::
forall e mr. (Prim e, Ord e, MemRead mr)
=> mr -- ^ /memRead1/ - First region of memory
-> mr -- ^ /memRead2/ - Second region of memory
-> Ordering
compareMem m1 m2
| isSameMem m1 m2 = EQ
| otherwise = compareOffMem m1 0 m2 0 (min n1 n2) <> compare n1 n2
where
n1 = countMem m1 :: Count e
n2 = countMem m2 :: Count e
{-# INLINE compareMem #-}
-- | Compare two regions using the `Ord` instance. It will return `EQ` whenever both
-- regions hold exactly the same elements and `LT` or `GT` as soon as the first discovered
-- element that is less than or greater than respectfully in the first region when
-- compared to the second one. It is safe for both regions to refer to the same part of
-- memory.
--
-- [Unsafe] When any precondition for either of the offsets @memOff1@, @memOff2@ or the
-- element count @memCount@ is violated the result is either unpredictable output or
-- failure with a segfault.
--
-- @since 0.3.0
compareOffMem ::
(Prim e, Ord e, MemRead mr1, MemRead mr2)
=> mr1 -- ^ /memRead1/ - First region of memory
-> Off e
-- ^ /memOff1/ - Offset for @memRead1@ in number of elements
--
-- /__Precondition:__/
--
-- > 0 <= memOff1
-> mr2 -- ^ /memRead2/ - Second region of memory
-> Off e
-- ^ /memOff2/ - Offset for @memRead1@ in number of elements
--
-- /__Precondition:__/
--
-- > 0 <= memOff2
-> Count e
-- ^ /memCount/ - Number of elements of type __@e@__ to compare
--
-- /__Preconditions:__/
--
-- > 0 <= memCount
--
-- > offToCount memOff1 + memCount < countMem memRead1
--
-- > offToCount memOff2 + memCount < countMem memRead2
-> Ordering
compareOffMem m1 off1 m2 off2 count = loop off1
where
doff = off2 - off1
k = countToOff count + off1
kRem = countToOff count `rem` 4
k4 = k - kRem
-- Some loop unrolling to get an extra 25% kick for smaller types
loop !i
| i < k4 =
let !i' = i + 1
!i'' = i + 2
!i''' = i + 3
in compare (indexOffMem m1 i) (indexOffMem m2 (i + doff)) <>
compare (indexOffMem m1 i') (indexOffMem m2 (i' + doff)) <>
compare (indexOffMem m1 i'') (indexOffMem m2 (i'' + doff)) <>
compare (indexOffMem m1 i''') (indexOffMem m2 (i''' + doff)) <>
loop (i + 4)
| i < k = compare (indexOffMem m1 i) (indexOffMem m2 (i + doff)) <> loop (i + 1)
| otherwise = EQ
{-# INLINE [1] compareOffMem #-}
{-# RULES
"compareOffMem/Word8" compareOffMem = compareByteOffMem
#-}