diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,9 @@ +# 0.3.6 + +* Addition of `unsafeArrayLinearCopy`, `unsafeLinearCopy`, `unsafeLinearShrink`, `unsafeLinearGrow` +* Implementation of `iterateUntil` and `iterateUntilM` +* `identityMatrix` - generation of identity matrix + # 0.3.5 * Fix and export `guardNumberOfElements` diff --git a/massiv.cabal b/massiv.cabal --- a/massiv.cabal +++ b/massiv.cabal @@ -1,5 +1,5 @@ name: massiv -version: 0.3.5.0 +version: 0.3.6.0 synopsis: Massiv (Массив) is an Array Library. description: Multi-dimensional Arrays with fusion, stencils and parallel computation. homepage: https://github.com/lehins/massiv @@ -79,6 +79,10 @@ , primitive , unliftio-core , vector + + if impl(ghc < 8.4) + build-depends: ghc-prim + include-dirs: include install-includes: massiv.h diff --git a/src/Data/Massiv/Array.hs b/src/Data/Massiv/Array.hs --- a/src/Data/Massiv/Array.hs +++ b/src/Data/Massiv/Array.hs @@ -118,8 +118,11 @@ , module Data.Massiv.Array.Ops.Transform -- * Slicing , module Data.Massiv.Array.Ops.Slice - -- * Sorting + -- * Algorithms + -- ** Sorting , quicksort + -- ** Iterations + , iterateUntil -- * Conversion , module Data.Massiv.Array.Manifest.List -- * Mutable diff --git a/src/Data/Massiv/Array/Delayed/Pull.hs b/src/Data/Massiv/Array/Delayed/Pull.hs --- a/src/Data/Massiv/Array/Delayed/Pull.hs +++ b/src/Data/Massiv/Array/Delayed/Pull.hs @@ -28,7 +28,6 @@ import qualified Data.Foldable as F import Data.Massiv.Array.Ops.Fold.Internal as A import Data.Massiv.Core.Common -import Data.Massiv.Core.Index.Internal import Data.Massiv.Core.List (L, showArrayList, showsArrayPrec) import GHC.Base (build) import Prelude hiding (zipWith) diff --git a/src/Data/Massiv/Array/Manifest/Internal.hs b/src/Data/Massiv/Array/Manifest/Internal.hs --- a/src/Data/Massiv/Array/Manifest/Internal.hs +++ b/src/Data/Massiv/Array/Manifest/Internal.hs @@ -38,6 +38,8 @@ , fromRaggedArray , sizeofArray , sizeofMutableArray + , iterateUntil + , iterateUntilM ) where import Control.Exception (try) @@ -375,3 +377,131 @@ (Mutable r ix e, StrideLoad r' ix e) => r -> Stride ix -> Array r' ix e -> Array r ix e computeWithStrideAs _ = computeWithStride {-# INLINE computeWithStrideAs #-} + + + +-- | Efficiently iterate a function until a convergence condition is satisfied. If the +-- size of array doesn't change between iterations then no more than two new arrays will be +-- allocated, regardless of the number of iterations. If the size does change from one +-- iteration to another, an attempt will be made to grow/shrink the intermediate mutable +-- array instead of allocating a new one. +-- +-- ====__Example__ +-- +-- >>> import Data.Massiv.Array +-- >>> a = computeAs P $ makeLoadArrayS (Sz2 8 8) (0 :: Int) $ \ w -> w (0 :. 0) 1 >> pure () +-- >>> a +-- Array P Seq (Sz (8 :. 8)) +-- [ [ 1, 0, 0, 0, 0, 0, 0, 0 ] +-- , [ 0, 0, 0, 0, 0, 0, 0, 0 ] +-- , [ 0, 0, 0, 0, 0, 0, 0, 0 ] +-- , [ 0, 0, 0, 0, 0, 0, 0, 0 ] +-- , [ 0, 0, 0, 0, 0, 0, 0, 0 ] +-- , [ 0, 0, 0, 0, 0, 0, 0, 0 ] +-- , [ 0, 0, 0, 0, 0, 0, 0, 0 ] +-- , [ 0, 0, 0, 0, 0, 0, 0, 0 ] +-- ] +-- >>> nextPascalRow cur above = if cur == 0 then above else cur +-- >>> pascal = makeStencil (Sz2 2 2) 1 $ \ get -> nextPascalRow <$> get (0 :. 0) <*> get (-1 :. -1) + get (-1 :. 0) +-- >>> iterateUntil (\_ _ a -> (a ! (7 :. 7)) /= 0) (\ _ -> mapStencil (Fill 0) pascal) a +-- Array P Seq (Sz (8 :. 8)) +-- [ [ 1, 0, 0, 0, 0, 0, 0, 0 ] +-- , [ 1, 1, 0, 0, 0, 0, 0, 0 ] +-- , [ 1, 2, 1, 0, 0, 0, 0, 0 ] +-- , [ 1, 3, 3, 1, 0, 0, 0, 0 ] +-- , [ 1, 4, 6, 4, 1, 0, 0, 0 ] +-- , [ 1, 5, 10, 10, 5, 1, 0, 0 ] +-- , [ 1, 6, 15, 20, 15, 6, 1, 0 ] +-- , [ 1, 7, 21, 35, 35, 21, 7, 1 ] +-- ] +-- +-- @since 0.3.6 +iterateUntil :: + (Load r' ix e, Mutable r ix e) + => (Int -> Array r ix e -> Array r ix e -> Bool) + -- ^ Convergence condition. Accepts current iteration counter, array at the previous + -- state and at the current state. + -> (Int -> Array r ix e -> Array r' ix e) + -- ^ A modifying function to apply at each iteration. The size of resulting array may + -- differ if necessary + -> Array r ix e -- ^ Initial source array + -> Array r ix e +iterateUntil convergence iteration initArr0 + | convergence 0 initArr0 initArr1 = initArr1 + | otherwise = + unsafePerformIO $ do + let loadArr = iteration 1 initArr1 + marr <- unsafeNew (size loadArr) + iterateLoop + (\n a a' _ -> pure $ convergence n a a') + iteration + 1 + initArr1 + loadArr + (asArr initArr0 marr) + where + !initArr1 = compute $ iteration 0 initArr0 + asArr :: Array r ix e -> MArray s r ix e -> MArray s r ix e + asArr _ = id +{-# INLINE iterateUntil #-} + +-- | Monadic version of `iterateUntil` where at each iteration mutable version of an array +-- is available. +-- +-- @since 0.3.6 +iterateUntilM :: + (Load r' ix e, Mutable r ix e, PrimMonad m, MonadIO m, PrimState m ~ RealWorld) + => (Int -> Array r ix e -> MArray (PrimState m) r ix e -> m Bool) + -- ^ Convergence condition. Accepts current iteration counter, pure array at previous + -- state and a mutable at the current state, therefore after each iteration its contents + -- can be modifed if necessary. + -> (Int -> Array r ix e -> Array r' ix e) + -- ^ A modifying function to apply at each iteration. The size of resulting array may + -- differ if necessary. + -> Array r ix e -- ^ Initial source array + -> m (Array r ix e) +iterateUntilM convergence iteration initArr0 = do + let loadArr0 = iteration 0 initArr0 + initMArr1 <- unsafeNew (size loadArr0) + computeInto initMArr1 loadArr0 + shouldStop <- convergence 0 initArr0 initMArr1 + initArr1 <- unsafeFreeze (getComp loadArr0) initMArr1 + if shouldStop + then pure initArr1 + else do + let loadArr1 = iteration 1 initArr1 + marr <- unsafeNew (size loadArr1) + iterateLoop (\n a _ -> convergence n a) iteration 1 initArr1 loadArr1 marr +{-# INLINE iterateUntilM #-} + + +iterateLoop :: + (Load r' ix e, Mutable r ix e, PrimMonad m, MonadIO m, PrimState m ~ RealWorld) + => (Int -> Array r ix e -> Array r ix e -> MArray (PrimState m) r ix e -> m Bool) + -> (Int -> Array r ix e -> Array r' ix e) + -> Int + -> Array r ix e + -> Array r' ix e + -> MArray (PrimState m) r ix e + -> m (Array r ix e) +iterateLoop convergence iteration = go + where + go !n !arr !loadArr !marr = do + let !sz = size loadArr + !k = totalElem sz + !mk = totalElem (msize marr) + marr' <- + if k == mk + then pure marr + else if k < mk + then unsafeLinearShrink marr sz + else unsafeLinearGrow marr sz + computeInto marr' loadArr + arr' <- unsafeFreeze (getComp loadArr) marr' + shouldStop <- convergence n arr arr' marr' + if shouldStop + then pure arr' + else do + nextMArr <- unsafeThaw arr + go (n + 1) arr' (iteration (n + 1) arr') nextMArr +{-# INLINE iterateLoop #-} diff --git a/src/Data/Massiv/Array/Manifest/Primitive.hs b/src/Data/Massiv/Array/Manifest/Primitive.hs --- a/src/Data/Massiv/Array/Manifest/Primitive.hs +++ b/src/Data/Massiv/Array/Manifest/Primitive.hs @@ -26,6 +26,7 @@ , toMutableByteArray , fromMutableByteArrayM , fromMutableByteArray + , shrinkMutableByteArray , unsafeAtomicReadIntArray , unsafeAtomicWriteIntArray , unsafeCasIntArray @@ -51,6 +52,9 @@ import Data.Primitive.Types import GHC.Base (Int(..)) import GHC.Exts as GHC +#if !MIN_VERSION_base(4,11,0) +import GHC.Prim as GHC +#endif import Prelude hiding (mapM) import System.IO.Unsafe (unsafePerformIO) @@ -193,7 +197,26 @@ unsafeLinearSet (MPArray _ ma) = setByteArray ma {-# INLINE unsafeLinearSet #-} + unsafeLinearCopy (MPArray _ maFrom) iFrom (MPArray _ maTo) iTo (Sz k) = + copyMutableByteArray maTo (iTo * esz) maFrom (iFrom * esz) (k * esz) + where esz = sizeOf (undefined :: e) + {-# INLINE unsafeLinearCopy #-} + unsafeArrayLinearCopy (PArray _ _ aFrom) iFrom (MPArray _ maTo) iTo (Sz k) = + copyByteArray maTo (iTo * esz) aFrom (iFrom * esz) (k * esz) + where esz = sizeOf (undefined :: e) + {-# INLINE unsafeArrayLinearCopy #-} + + unsafeLinearShrink (MPArray _ ma) sz = do + shrinkMutableByteArray ma (totalElem sz * sizeOf (undefined :: e)) + pure $ MPArray sz ma + {-# INLINE unsafeLinearShrink #-} + + unsafeLinearGrow (MPArray _ ma) sz = + MPArray sz <$> resizeMutableByteArray ma (totalElem sz * sizeOf (undefined :: e)) + {-# INLINE unsafeLinearGrow #-} + + instance (Prim e, Index ix) => Load P ix e where size = pSize {-# INLINE size #-} @@ -458,3 +481,24 @@ mba (toLinearIndex sz ix) {-# INLINE unsafeAtomicXorIntArray #-} + + + +#if !MIN_VERSION_primitive(0,6,4) +resizeMutableByteArray :: + PrimMonad m => MutableByteArray (PrimState m) -> Int -> m (MutableByteArray (PrimState m)) +resizeMutableByteArray (MutableByteArray arr#) (I# n#) = + primitive + (\s# -> + case resizeMutableByteArray# arr# n# s# of + (# s'#, arr'# #) -> (# s'#, MutableByteArray arr'# #)) +{-# INLINE resizeMutableByteArray #-} +#endif + +shrinkMutableByteArray :: forall m. (PrimMonad m) + => MutableByteArray (PrimState m) + -> Int -- ^ new size + -> m () +{-# INLINE shrinkMutableByteArray #-} +shrinkMutableByteArray (MutableByteArray arr#) (I# n#) + = primitive_ (shrinkMutableByteArray# arr# n#) diff --git a/src/Data/Massiv/Array/Manifest/Storable.hs b/src/Data/Massiv/Array/Manifest/Storable.hs --- a/src/Data/Massiv/Array/Manifest/Storable.hs +++ b/src/Data/Massiv/Array/Manifest/Storable.hs @@ -2,6 +2,7 @@ {-# LANGUAGE CPP #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MagicHash #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} @@ -32,8 +33,11 @@ import Control.DeepSeq (NFData(..), deepseq) import Control.Monad.IO.Unlift +import Control.Monad.Primitive (unsafePrimToPrim) import Data.Massiv.Array.Delayed.Pull (eq, ord) import Data.Massiv.Array.Manifest.Internal +import Data.Massiv.Array.Manifest.Primitive (shrinkMutableByteArray) +import Data.Primitive.ByteArray (MutableByteArray(..)) import Data.Massiv.Array.Manifest.List as A import Data.Massiv.Array.Mutable import Data.Massiv.Core.Common @@ -41,8 +45,11 @@ import qualified Data.Vector.Generic.Mutable as VGM import qualified Data.Vector.Storable as VS import qualified Data.Vector.Storable.Mutable as MVS -import Foreign.ForeignPtr import Foreign.Ptr +import GHC.ForeignPtr (ForeignPtr(..), ForeignPtrContents(..)) +import Foreign.ForeignPtr (withForeignPtr) +import Foreign.Storable +import Foreign.Marshal.Array (copyArray, advancePtr) import GHC.Exts as GHC (IsList(..)) import Prelude hiding (mapM) import System.IO.Unsafe (unsafePerformIO) @@ -155,6 +162,43 @@ unsafeLinearWrite (MSArray _ mv) = INDEX_CHECK("(Mutable S ix e).unsafeLinearWrite", Sz . MVS.length, MVS.unsafeWrite) mv {-# INLINE unsafeLinearWrite #-} + + -- TODO: Try approach from `vector`, fallback on Prim for setByteArray/recursive copyArray + -- unsafeLinearSet (MSArray _ v) = setByteArray ma + -- {-# INLINE unsafeLinearSet #-} + + unsafeLinearCopy marrFrom iFrom marrTo iTo (Sz k) = do + let MSArray _ (MVS.MVector _ fpFrom) = marrFrom + MSArray _ (MVS.MVector _ fpTo) = marrTo + unsafePrimToPrim $ + withForeignPtr fpFrom $ \ ptrFrom -> + withForeignPtr fpTo $ \ ptrTo -> do + let ptrFrom' = advancePtr ptrFrom iFrom + ptrTo' = advancePtr ptrTo iTo + copyArray ptrTo' ptrFrom' k + {-# INLINE unsafeLinearCopy #-} + + unsafeArrayLinearCopy arrFrom iFrom marrTo iTo sz = do + marrFrom <- unsafeThaw arrFrom + unsafeLinearCopy marrFrom iFrom marrTo iTo sz + {-# INLINE unsafeArrayLinearCopy #-} + + unsafeLinearShrink marr@(MSArray _ mv@(MVS.MVector _ (ForeignPtr _ fpc))) sz = do + let shrinkMBA :: MutableByteArray RealWorld -> IO () + shrinkMBA mba = shrinkMutableByteArray mba (totalElem sz * sizeOf (undefined :: e)) + {-# INLINE shrinkMBA #-} + case fpc of + MallocPtr mba# _ -> do + unsafePrimToPrim $ shrinkMBA (MutableByteArray mba#) + pure $ MSArray sz mv + PlainPtr mba# -> do + unsafePrimToPrim $ shrinkMBA (MutableByteArray mba#) + pure $ MSArray sz mv + _ -> unsafeDefaultLinearShrink marr sz + {-# INLINE unsafeLinearShrink #-} + + unsafeLinearGrow (MSArray _ mv) sz = MSArray sz <$> MVS.unsafeGrow mv (totalElem sz) + {-# INLINE unsafeLinearGrow #-} instance (Index ix, VS.Storable e) => Load S ix e where diff --git a/src/Data/Massiv/Array/Manifest/Unboxed.hs b/src/Data/Massiv/Array/Manifest/Unboxed.hs --- a/src/Data/Massiv/Array/Manifest/Unboxed.hs +++ b/src/Data/Massiv/Array/Manifest/Unboxed.hs @@ -177,6 +177,8 @@ INDEX_CHECK("(Mutable U ix e).unsafeLinearWrite", Sz . MVU.length, MVU.unsafeWrite) mv {-# INLINE unsafeLinearWrite #-} + unsafeLinearGrow (MUArray _ mv) sz = MUArray sz <$> MVU.unsafeGrow mv (totalElem sz) + {-# INLINE unsafeLinearGrow #-} instance ( VU.Unbox e , IsList (Array L ix e) diff --git a/src/Data/Massiv/Array/Mutable.hs b/src/Data/Massiv/Array/Mutable.hs --- a/src/Data/Massiv/Array/Mutable.hs +++ b/src/Data/Massiv/Array/Mutable.hs @@ -82,7 +82,7 @@ -- TODO: add fromListM, et al. -import Control.Monad (unless) +import Control.Monad (when, unless) import Control.Monad.ST import Control.Scheduler import Data.Massiv.Core.Common @@ -140,8 +140,20 @@ -- -- @since 0.1.0 thaw :: forall r ix e m. (Mutable r ix e, MonadIO m) => Array r ix e -> m (MArray RealWorld r ix e) -thaw arr = liftIO $ makeMArrayLinear (getComp arr) (size arr) (pure . unsafeLinearIndexM arr) --- TODO: use faster memcpy +thaw arr = + liftIO $ do + let sz = size arr + totalLength = totalElem sz + marr <- unsafeNew sz + withScheduler_ (getComp arr) $ \scheduler -> + splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do + loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start -> + scheduleWork_ scheduler $ unsafeArrayLinearCopy arr start marr start (SafeSz chunkLength) + let slackLength = totalLength - slackStart + when (slackLength > 0) $ + scheduleWork_ scheduler $ + unsafeArrayLinearCopy arr slackStart marr slackStart (SafeSz slackLength) + pure marr {-# INLINE thaw #-} -- | Same as `thaw`, but restrict computation to sequential only. @@ -162,13 +174,13 @@ forall r ix e m. (Mutable r ix e, PrimMonad m) => Array r ix e -> m (MArray (PrimState m) r ix e) -thawS arr = makeMArrayLinearS (size arr) (pure . unsafeLinearIndexM arr) --- TODO: use faster memcpy +thawS arr = do + tmarr <- unsafeNew (size arr) + unsafeArrayLinearCopy arr 0 tmarr 0 (SafeSz (totalElem (size arr))) + pure tmarr {-# INLINE thawS #-} --- TODO: implement and benchmark parallel `thawIO` and `freezeIO` with memcpy - -- | /O(n)/ - Yield an immutable copy of the mutable array. Note that mutable representations -- have to be the same. -- @@ -189,7 +201,20 @@ => Comp -> MArray RealWorld r ix e -> m (Array r ix e) -freeze comp marr = liftIO $ generateArrayLinear comp (msize marr) (unsafeLinearRead marr) +freeze comp smarr = + liftIO $ do + let sz = msize smarr + totalLength = totalElem sz + tmarr <- unsafeNew sz + withScheduler_ comp $ \scheduler -> + splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do + loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start -> + scheduleWork_ scheduler $ unsafeLinearCopy smarr start tmarr start (SafeSz chunkLength) + let slackLength = totalLength - slackStart + when (slackLength > 0) $ + scheduleWork_ scheduler $ + unsafeLinearCopy smarr slackStart tmarr slackStart (SafeSz slackLength) + unsafeFreeze comp tmarr {-# INLINE freeze #-} @@ -201,8 +226,13 @@ forall r ix e m. (Mutable r ix e, PrimMonad m) => MArray (PrimState m) r ix e -> m (Array r ix e) -freezeS marr = generateArrayLinearS Seq (msize marr) (unsafeLinearRead marr) +freezeS smarr = do + let sz = msize smarr + tmarr <- unsafeNew sz + unsafeLinearCopy smarr 0 tmarr 0 (SafeSz (totalElem sz)) + unsafeFreeze Seq tmarr {-# INLINE freezeS #-} + newMaybeInitialized :: (Load r' ix e, Mutable r ix e, PrimMonad m) => Array r' ix e -> m (MArray (PrimState m) r ix e) diff --git a/src/Data/Massiv/Array/Mutable/Algorithms.hs b/src/Data/Massiv/Array/Mutable/Algorithms.hs --- a/src/Data/Massiv/Array/Mutable/Algorithms.hs +++ b/src/Data/Massiv/Array/Mutable/Algorithms.hs @@ -11,9 +11,11 @@ module Data.Massiv.Array.Mutable.Algorithms ( quicksortM_ , unstablePartitionM + , iterateUntilM ) where import Data.Massiv.Array.Ops.Sort +import Data.Massiv.Array.Manifest.Internal (iterateUntilM) import Data.Massiv.Core.Common diff --git a/src/Data/Massiv/Array/Numeric.hs b/src/Data/Massiv/Array/Numeric.hs --- a/src/Data/Massiv/Array/Numeric.hs +++ b/src/Data/Massiv/Array/Numeric.hs @@ -19,6 +19,7 @@ , (.^) , (|*|) , multiplyTransposed + , identityMatrix , negateA , absA , signumA @@ -64,6 +65,7 @@ ) where import Data.Massiv.Array.Delayed.Pull +import Data.Massiv.Array.Delayed.Push import Data.Massiv.Array.Manifest.Internal import Data.Massiv.Array.Ops.Fold as A import Data.Massiv.Array.Ops.Map as A @@ -179,6 +181,25 @@ SafeSz (m1 :. n1) = size arr1 SafeSz (n2 :. m2) = size arr2 {-# INLINE multiplyTransposed #-} + +-- | Create an indentity matrix. +-- +-- ====___Example__ +-- +-- >>> import Data.Massiv.Array +-- >>> identityMatrix 5 +-- Array DL Seq (Sz (5 :. 5)) +-- [ [ 1, 0, 0, 0, 0 ] +-- , [ 0, 1, 0, 0, 0 ] +-- , [ 0, 0, 1, 0, 0 ] +-- , [ 0, 0, 0, 1, 0 ] +-- , [ 0, 0, 0, 0, 1 ] +-- ] +-- +-- @since 0.3.6 +identityMatrix :: Int -> Array DL Ix2 Int +identityMatrix n = makeLoadArrayS (Sz2 n n) 0 $ \ w -> loopM_ 0 (< n) (+1) $ \ i -> w (i :. i) 1 +{-# INLINE identityMatrix #-} negateA diff --git a/src/Data/Massiv/Array/Stencil.hs b/src/Data/Massiv/Array/Stencil.hs --- a/src/Data/Massiv/Array/Stencil.hs +++ b/src/Data/Massiv/Array/Stencil.hs @@ -1,6 +1,7 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilies #-} -- | -- Module : Data.Massiv.Array.Stencil -- Copyright : (c) Alexey Kuleshevich 2018-2019 @@ -114,3 +115,4 @@ inline relStencil $ \ !ixD -> getVal (liftIndex2 (+) ix ixD) {-# INLINE stencil #-} {-# INLINE makeStencilDef #-} + diff --git a/src/Data/Massiv/Array/Stencil/Internal.hs b/src/Data/Massiv/Array/Stencil/Internal.hs --- a/src/Data/Massiv/Array/Stencil/Internal.hs +++ b/src/Data/Massiv/Array/Stencil/Internal.hs @@ -25,7 +25,6 @@ import Control.DeepSeq import Data.Massiv.Array.Delayed.Pull import Data.Massiv.Core.Common -import Data.Massiv.Core.Index.Internal -- | Stencil is abstract description of how to handle elements in the neighborhood of every array -- cell in order to compute a value for the cells in the new array. Use `Data.Array.makeStencil` and diff --git a/src/Data/Massiv/Array/Unsafe.hs b/src/Data/Massiv/Array/Unsafe.hs --- a/src/Data/Massiv/Array/Unsafe.hs +++ b/src/Data/Massiv/Array/Unsafe.hs @@ -42,6 +42,10 @@ , unsafeWrite , unsafeLinearWrite , unsafeLinearSet + , unsafeLinearCopy + , unsafeArrayLinearCopy + , unsafeLinearShrink + , unsafeLinearGrow -- * Pointer access , unsafeWithPtr , unsafeArrayToForeignPtr diff --git a/src/Data/Massiv/Core/Common.hs b/src/Data/Massiv/Core/Common.hs --- a/src/Data/Massiv/Core/Common.hs +++ b/src/Data/Massiv/Core/Common.hs @@ -37,6 +37,7 @@ , unsafeRead , unsafeWrite , unsafeLinearModify + , unsafeDefaultLinearShrink , Ragged(..) , Nested(..) , NestedStruct @@ -45,6 +46,7 @@ -- * Size , elemsCount , isEmpty + , Sz(SafeSz) -- * Indexing , (!?) , index @@ -86,6 +88,7 @@ scheduleWork, scheduleWork_) import Data.Massiv.Core.Exception import Data.Massiv.Core.Index +import Data.Massiv.Core.Index.Internal (Sz(SafeSz)) import Data.Typeable #include "massiv.h" @@ -328,12 +331,78 @@ return marr {-# INLINE initializeNew #-} - + -- | Set all cells in the mutable array within the range to a specified value. + -- + -- @since 0.3.0 unsafeLinearSet :: PrimMonad m => - MArray (PrimState m) r ix e -> Int -> Int -> e -> m () + MArray (PrimState m) r ix e -> Ix1 -> Int -> e -> m () unsafeLinearSet marr offset len e = loopM_ offset (< (offset + len)) (+1) (\i -> unsafeLinearWrite marr i e) {-# INLINE unsafeLinearSet #-} + + -- | Copy part of one mutable array into another + -- + -- @since 0.3.6 + unsafeLinearCopy :: PrimMonad m => + MArray (PrimState m) r ix e -- ^ Source mutable array + -> Ix1 -- ^ Starting index at source array + -> MArray (PrimState m) r ix e -- ^ Target mutable array + -> Ix1 -- ^ Starting index at target array + -> Sz1 -- ^ Number of elements to copy + -> m () + unsafeLinearCopy marrFrom iFrom marrTo iTo (SafeSz k) = do + let delta = iTo - iFrom + loopM_ iFrom (< k + iFrom) (+1) $ \i -> + unsafeLinearRead marrFrom i >>= unsafeLinearWrite marrTo (i + delta) + {-# INLINE unsafeLinearCopy #-} + + -- | Copy a part of a pure array into a mutable array + -- + -- @since 0.3.6 + unsafeArrayLinearCopy :: PrimMonad m => + Array r ix e -- ^ Source pure array + -> Ix1 -- ^ Starting index at source array + -> MArray (PrimState m) r ix e -- ^ Target mutable array + -> Ix1 -- ^ Starting index at target array + -> Sz1 -- ^ Number of elements to copy + -> m () + unsafeArrayLinearCopy arrFrom iFrom marrTo iTo (SafeSz k) = do + let delta = iTo - iFrom + loopM_ iFrom (< k + iFrom) (+1) $ \i -> + unsafeLinearWrite marrTo (i + delta) (unsafeLinearIndex arrFrom i) + {-# INLINE unsafeArrayLinearCopy #-} + + -- | Linearly reduce the size of an array. Total number of elements should be smaller. + -- + -- @since 0.3.6 + unsafeLinearShrink :: PrimMonad m => + MArray (PrimState m) r ix e -> Sz ix -> m (MArray (PrimState m) r ix e) + unsafeLinearShrink = unsafeDefaultLinearShrink + {-# INLINE unsafeLinearShrink #-} + + -- | Linearly increase the size of an array. Total number of elements should be larger. + -- + -- @since 0.3.6 + unsafeLinearGrow :: PrimMonad m => + MArray (PrimState m) r ix e -> Sz ix -> m (MArray (PrimState m) r ix e) + unsafeLinearGrow marr sz = do + marr' <- unsafeNew sz + unsafeLinearCopy marr 0 marr' 0 $ SafeSz (totalElem (msize marr)) + pure marr' + {-# INLINE unsafeLinearGrow #-} + + +unsafeDefaultLinearShrink :: + (Mutable r ix e, PrimMonad m) + => MArray (PrimState m) r ix e + -> Sz ix + -> m (MArray (PrimState m) r ix e) +unsafeDefaultLinearShrink marr sz = do + marr' <- unsafeNew sz + unsafeLinearCopy marr 0 marr' 0 $ SafeSz (totalElem sz) + pure marr' +{-# INLINE unsafeDefaultLinearShrink #-} + -- | Read an array element unsafeRead :: (Mutable r ix e, PrimMonad m) => diff --git a/src/Data/Massiv/Core/List.hs b/src/Data/Massiv/Core/List.hs --- a/src/Data/Massiv/Core/List.hs +++ b/src/Data/Massiv/Core/List.hs @@ -33,7 +33,6 @@ import Data.Foldable (foldr') import qualified Data.List as L import Data.Massiv.Core.Common -import Data.Massiv.Core.Index.Internal import Data.Proxy import Data.Typeable import GHC.Exts diff --git a/tests/Data/Massiv/Array/MutableSpec.hs b/tests/Data/Massiv/Array/MutableSpec.hs --- a/tests/Data/Massiv/Array/MutableSpec.hs +++ b/tests/Data/Massiv/Array/MutableSpec.hs @@ -7,6 +7,7 @@ module Data.Massiv.Array.MutableSpec (spec) where import Control.Concurrent.Async +import Control.Monad (when) import Control.Monad.ST import Data.Functor.Identity import Data.List as L @@ -43,6 +44,50 @@ arr' <- run $ generateArray (getComp arr) (size arr) (evaluateM arr) return (arr === arr') +prop_shrinkIO :: + ( Mutable r ix Int + , Resize r ix + , Source r Ix1 Int + ) + => r + -> Proxy ix + -> ArrIx r ix Int + -> Property +prop_shrinkIO _ _ (ArrIx arr ix) = + monadicIO $ + run $ do + marr <- thaw arr + sarr <- unsafeFreeze (getComp arr) =<< unsafeLinearShrink marr (Sz ix) + pure (A.foldlS (.&&.) (property True) $ A.zipWith (===) (flatten arr) (flatten sarr)) + +prop_growShrinkIO :: + ( Show (Array r ix Int) + , Eq (Array r ix Int) + , Mutable r ix Int + , Extract r ix Int + , Num ix + , Load (EltRepr r ix) ix Int + ) + => r + -> Proxy ix + -> Arr r ix Int + -> NonNegative Int + -> Property +prop_growShrinkIO _ _ (Arr arr) (NonNegative delta) = + monadicIO $ + run $ do + marr <- thaw arr + let sz = size arr + k <- getDimM (unSz sz) (dimensions sz) + -- increase the outer most dimension, just so the structure doesn't change + newSz <- Sz <$> setDimM (unSz sz) (dimensions sz) (k + delta) + gMarr <- unsafeLinearGrow marr newSz + -- Make sure we can write into the newly allocated area + when (delta > 0) $ write' gMarr (unSz newSz - 1) delta + garr <- compute . extract' 0 sz <$> unsafeFreeze (getComp arr) gMarr + sarr <- freeze (getComp arr) =<< unsafeLinearShrink gMarr sz + pure (garr === arr .&&. sarr === arr) + prop_atomicModifyIntArrayMany :: ArrIx P Ix2 Int -> Array B Ix1 Int -> Property prop_atomicModifyIntArrayMany (ArrIx arr ix) barr = monadicIO $ do @@ -58,15 +103,16 @@ prop_atomicReadIntArrayMany :: Array P Ix2 Int -> Array B Ix1 Ix2 -> Property -prop_atomicReadIntArrayMany arr bix = monadicIO $ do +prop_atomicReadIntArrayMany arr bix = + monadicIO $ run $ do - marr <- thaw arr - as :: Array N Ix1 (Maybe Int) <- forM bix (A.read marr) - as' <- forM bix (atomicReadIntArray marr) - pure (as === as') + marr <- thaw arr + as :: Array N Ix1 (Maybe Int) <- forM bix (A.read marr) + as' <- forM bix (atomicReadIntArray marr) + pure (as === as') -prop_atomicWriteIntArrayMany :: Array P Ix2 Int -> Array B Ix1 Ix2 -> (Fun Ix2 Int) -> Property +prop_atomicWriteIntArrayMany :: Array P Ix2 Int -> Array B Ix1 Ix2 -> Fun Ix2 Int -> Property prop_atomicWriteIntArrayMany arr bix f = monadicIO $ run $ do @@ -96,20 +142,30 @@ rev a = computeAs P $ backpermute' sz1 (\ix1 -> unSz sz1 - ix1 - 1) a + mutableSpec :: ( Show r - , Show (Array r Ix3 Int) , Show (Array r Ix1 Int) , Show (Array r Ix2 Int) - , Eq (Array r Ix3 Int) + , Show (Array r Ix3 Int) , Eq (Array r Ix1 Int) , Eq (Array r Ix2 Int) - , Mutable r Ix3 Int + , Eq (Array r Ix3 Int) , Mutable r Ix1 Int , Mutable r Ix2 Int - , Construct r Ix3 Int + , Mutable r Ix3 Int , Construct r Ix1 Int , Construct r Ix2 Int + , Construct r Ix3 Int + , Extract r Ix1 Int + , Extract r Ix2 Int + , Extract r Ix3 Int + , Resize r Ix1 + , Resize r Ix2 + , Resize r Ix3 + , Load (EltRepr r Ix1) Ix1 Int + , Load (EltRepr r Ix2) Ix2 Int + , Load (EltRepr r Ix3) Ix3 Int ) => r -> SpecWith () @@ -131,7 +187,14 @@ it "Ix1" $ property $ prop_generateMakeIO r (Proxy :: Proxy Ix1) it "Ix2" $ property $ prop_generateMakeIO r (Proxy :: Proxy Ix2) it "Ix3" $ property $ prop_generateMakeIO r (Proxy :: Proxy Ix3) - + describe "shrink" $ do + it "Ix1" $ property $ prop_shrinkIO r (Proxy :: Proxy Ix1) + it "Ix2" $ property $ prop_shrinkIO r (Proxy :: Proxy Ix2) + it "Ix3" $ property $ prop_shrinkIO r (Proxy :: Proxy Ix3) + describe "grow+shrink" $ do + it "Ix1" $ property $ prop_growShrinkIO r (Proxy :: Proxy Ix1) + it "Ix2" $ property $ prop_growShrinkIO r (Proxy :: Proxy Ix2) + it "Ix3" $ property $ prop_growShrinkIO r (Proxy :: Proxy Ix3) generateSpec :: Spec generateSpec = do