diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,285 @@ +# 1.0.5 + +* Add `Functor` instance for `Border` +* Improve performance and reduce allocations during computation of higher dimension `DW` arrays [#142](https://github.com/lehins/massiv/issues/142) + +# 1.0.4 + +* Improve performance of sorting algorithm and its parallelization. Fix huge slow down on + CPUs with at least 16 cores. + +# 1.0.3 + +* Deprecated `indexWith` in favor of `indexAssert` +* Addition of scans: `sscanl`, `sscanl1`, `sprescanl`, `spostscanl` and `spostscanlAcc` +* Expose `unsafePrefIndex` + +# 1.0.2 + +* Addition of `Iterator` type class and related fucntions: + * Addition of `RowMajor`, `RowMajorLinear` and `RowMajorUnbalanced` iterators. + * Switch parallel left fold to new iterator +* Improvements to functions that do the looping: + * Addition of `loopNextA_` and `loopNextM` + * Deprecate `loopM_` in favor of `loopA_` + * Addition of `loopA` and `loopF` for applicative iterators + * Addition of `iloopA_` and `iloopM` + * Addition of `nextMaybeF` + * Addition of `iterLinearST_`, `iterLinearAccST_` and `iterLinearAccST` + * Addition of an optimized `scheduleMassivWork` for internal use + * Addition of a new workhorse: `splitWorkWithFactorST` + * Addition of a new workhorse: `splitWorkWithFactorST` +* Changes to `Index` class: + * Deprecate `iterM_` in favor of `iterA_` + * Adititon of sequential iterators: `iterTargetRowMajorA_`, + `iterTargetRowMajorAccM` needed for `RowMajor` iterator + * Addition of parallelizable iterators: `iterRowMajorST`, + `iterTargetRowMajorAccST`, `iterTargetRowMajorAccST_` needed + for `RowMajor` iterator + * Addition of `iterF` for using with applicative iterators. + * Addition of `stepNextMF` for streaming iteration of multi-dimensional + arrays. +* Addition of `repr`. +* Addition of `quicksortAs`, `quicksortAsBy`, `quicksortAsByM` +* Fix backwards compatibility with ghc-8.0 +* Get rid of dependency on `genvalidity`: too many compatibility issues for + little gain +* Introduce `PrefIndex` and `unsafePrefIndex`: a preference when indexing into + multidimensional `Source` arrays. Adopt it where possible for left and right + folds, traversals, but not zipping +* Improve multi-dimensional indices for streams. Improve `steps` and `isteps` +* Get rid of build warnings for ghc-9.4 + * Make `Monoid` and `Monad` instances forward compatible +* Compatibility with `vector-0.13`: + * Fix `Unbox` instance for `IxN` + * Fix safety of boxed vector conversion: `toBoxedVector`/`fromBoxedVector` +* Re-export Manifest realetd functions from `Data.Massiv.Array.Manifest` + as a migration strategy for the move in the next minor version bump. + +# 1.0.1 + +* Relax constraint on `computeInto` by removing requirement for `Size` +* Fix `BL`, which due to a forgotten `seq` was not lazy. + +# 1.0.0 + +* Addition of `sumArrays'`, `sumArraysM` and `productArrays'`, `productArraysM`. +* Remove `Num`/`Fractional`/`Floating` instances for `D` and `DI` arrays. This was done to + prevent surprises as in: [#97](https://github.com/lehins/massiv/issues/97) +* Remove helper class `Nested` and type family `NestedStuct` +* Make `negate` in `Num` instance throw error for `Sz` in order to avoid surprising + behavior reported in: [#114](https://github.com/lehins/massiv/issues/114) +* Add of `munsafeResize` +* Add `uniformArray` and `uniformRangeArray` +* Replace `isNonEmpty` with `isNotZeroSz` and added `isZeroSz` +* Consolidate `Construct` class into `Load` +* Introduce `Shape`, the parent of `Size` +* Move `size` from `Load` into new class `Size` +* Consolidate `Resize` into `Size` +* Removed `maxSize` and replaced it with `maxLinearSize` +* Remove specialized `DW` instances that used tuples as indices. +* Get rid of `M` representation +* Remove `R` type family and `Slice`, `InnerSlice` and `Extract` classes in favor of `D`. +* Consolidate `OuterSlice` into `Source` +* Add `Strategy` and move `setComp` (from `Construct`) and `getComp` (from `Load`) in there. +* Remove `ix` from `Mutable`, `Manifest`, `Source` +* Remove `liftArray2`. Instead add `liftArray2'` and `liftArray2M` that don't behave + like a `map` for singleton argument. +* Expose `liftNumArray2M` +* Prevent `showsArrayPrec` from changing index type +* Change function argument to monadic action for `unstablePartitionM` and `unsafeUnstablePartitionM` +* Replace `snull` with a more generic `isNull` +* Switch `DL` loading function to run in `ST` monad, rather than in any `Monad m`. +* Rename `msize` -> `sizeOfMArray` +* Add `unsafeResizeMArray` and `unsafeLinearSliceMArray` +* Rename: + * `loadArrayM` -> `iterArrayLinearM_` + * `loadArrayWithSetM` -> `iterArrayLinearWithSetM_`. + * `loadArrayWithStrideM` -> `iterArrayLinearWithStrideM_`. +* Add `iterArrayLinearST_` and `iterArrayLinearWithSetST_` to `Load` class instead + of `loadArrayM` and `loadArrayWithSetM`. +* Add `iterArrayLinearWithStrideST_` to `LoadStride` class instead of `loadArrayWithStrideM`. +* Add new mutable functions: + * `resizeMArrayM` and `flattenMArray`, + * `outerSliceMArrayM` and `outerSlicesMArray`, + * `for2PrimM_` and `ifor2PrimM_`, + * `zipSwapM_` +* Switch effectful mapping functions to use the representation specific + iteration. This means that they are now restricted to `Load` instead of + `Source`. Functions affected: + * `mapIO_`, `imapIO_`, `forIO_` and `iforIO_` + * `mapIO`, `imapIO`, `forIO` and `iforIO` +* Add `Uniform`, `UniformRange` and `Random` instances for `Ix2`, `IxN`, `Dim`, `Sz` and `Stride` +* Consolidate `Mutable` into `Manifest` type class and move the `MArray` data + family outside of the class. +* Make sure empty arrays are always equal, regardless of their size. +* Remove `LN` representation in favor of a standalone `List` newtype wrapper + around lists. + +# 0.6.1 + +* Addition of `withLoadMArray_`, `withLoadMArrayS`, `withLoadMArrayS_`, + `withLoadMArrayST`, `withLoadMArrayST_` +* Addition of `replaceSlice` and `replaceOuterSlice` +* Addition of `quicksortBy`, `quicksortByM` and `quicksortByM_` +* Fix performance regression for `quicksort` and `quicksortM_` introduced in previous release. + + +# 0.6.0 + +* Fix semantics of `Applicative`, `Num` and `Fractional` instance for `D` arrays: + mismatched sizes will throw an error. +* 20% speed improvement of matrix multiplication: `multiplyMatrices`, `.><.` and + `!><.` and `>`, `|*|`, `multiplyTransposed`, `fromIntegerA`, + `fromRationalA`, `piA` +* Addition of `BL` representation and related functionality, fix for [#111](https://github.com/lehins/massiv/issues/111). + * Addition of functions: `wrapLazyArray`, `unwrapLazyArray`, `toLazyArray`, + `evalLazyArray`, `forceLazyArray`, `unwrapMutableLazyArray`, `fromBoxedVector`, + `fromBoxedMVector`. + * Rename: + * `unsafeNormalBoxedArray` -> `coerceNormalBoxedArray` + * `unsafeBoxedArray` -> `coerceBoxedArray` + * Remove `unsafeFromBoxedVector` + * Conversion from vector with `castFromVector` will return `BL` representation for boxed vector + * Change type `B` -> `BL` for functions: `toBoxedVector` and `toBoxedMVector` +* Rename `N` -> `BN` and add backwards compatibility shim. +* Make `replicate` a function in `Construct` class +* Add `newMArray`, `newMArray'` and deprecate `new` +* Add custom implementation for `<$` in `Functor` instances for `BL` and `B`. + +# 0.5.9 + +* Add `mallocCompute`, `mallocCopy` and `unsafeMallocMArray` +* Fix `.><.`, `><.` and `.><` on empty matrices. Result is now guaranteed to be empty too. +* Add `unwrapByteArrayOffset` and `unwrapMutableByteArrayOffset` +* Add `fromByteArrayOffsetM` and `fromMutableByteArrayOffsetM` + +# 0.5.8 + +* Improve loading of push arrays by adding `loadArrayWithSetM` and deprecating `defaultElement`. + +# 0.5.9 + +* Add `mallocCompute`, `mallocCopy` and `unsafeMallocMArray` + +# 0.5.8 + +* Improve loading of push arrays by adding `loadArrayWithSetM` and deprecating `defaultElement`. + +# 0.5.7 + +* Improve performance of `><.` and `><.` and `!><` and `!><` + * New vector-matrix multiplication: `><.` and `>` + +# 0.5.5 + +* Add `takeWhile`, `dropWhile` and `findIndex` +* Improve performance of `any`, `and`, `or`, `all` +* Add `elem` + +# 0.5.4 + +* Addition of `unsafeTransformStencil` +* Add `zip4`, `unzip4`, `zipWith4` and `izipWith4` +* Make `Resize` a superclass of `Source` +* Addition of `outerSlices`, `innerSlices`, `withinSlices` and `withinSlicesM` +* Addition of `stackSlicesM`, `stackOuterSlicesM` and `stackInnerSlicesM` +* Addition of `computeP` +* Fix perfomrmance issue of folding functions applied to arrays with `Seq` computation + strategy. + +# 0.5.3 + +* Fix `tanA` and `tanhA`. [#96](https://github.com/lehins/massiv/pull/96) +* Relax argument of `snoc` and `cons` constraint to `Load` vectors +* Improve `unsnocM` and `unconsM` by switching to `unsafeLinearSlice`, instead of delaying + the array. +* Fix parallelization for windowed array when computed with stride +* Fix massiv doctests not being able to find massiv.h under NixOS + +# 0.5.2 + +* Addition of `lowerTriangular` and `upperTriangular` +* Relax `identityMatrix` type to return an array of any `Num` type, not just `Int`. +* Addition of `unsafeMakeLoadArrayAdjusted` +* Add matrix-vector product (`(#>)`) +* Addition of `siterate` + +# 0.5.1 + +* Fix `sfromListN` accepting a plain `Int` instead of `Sz1`, as well as switch to upper bound. +* Fix order of argumetns in `iforM` +* Restrict `szip*`, `szipWith*` and `sizipWith*` functions to flat vectors. +* Addition of `unsafeSUnfoldrN`, `unsafeSUnfoldrNM` and `unsafeSFromListN` +* Fix `sunfoldrN`, `sunfoldrNM` and `sfromListN` to not trust the supplied size. +* Move `isEmpty` into `Load` class +* Add `isNotEmpty` + +# 0.5.0 + +* Remove `Show` instance from `Value`. +* Addition of `unsafeCreateArray`, `unsafeCreateArray_` and `unsafeCreateArrayS` +* Remove `Comp` argument from functions that ignore it and set it to `Seq`: + * `createArrayS_`, `createArrayS`, `createArrayST_`, `createArrayST` + * `unfoldrPrimM_`, `iunfoldrPrimM_`, `unfoldrPrimM`, `iunfoldrPrimM` + * `unfoldlPrimM_`, `iunfoldlPrimM_`, `unfoldlPrimM`, `iunfoldlPrimM` +* Addition of `fromStorableVector` and `fromStorableMVector` +* Modify `toMutableByteArray` to produce a copy if dealing with slice. +* Addition of `toByteArrayM`, `toMutableByteArrayM` +* Change `replicate` to produce delayed load array `DL` +* Export unsafe stencil functions from `Data.Array.Massiv.Unsafe`, rather than from + `Data.Massiv.Array.Stencil.Unsafe`. +* Implement `unsafeMapStencil` and deprecate `mapStencilUnsafe` and `forStencilUnsafe` +* Addition of `castToBuilder` +* Addition of conversion functions: + * `unwrapNormalForm` and `evalNormalForm` + * `toBoxedVector`, `toBoxedMVector`, `evalBoxedVector` and `evalBoxedMVector` + * `unwrapByteArray` and `unwrapMutableByteArray` + * `toPrimitiveVector`, `toPrimitiveMVector`, `fromPrimitiveVector` and + `fromPrimitiveMVector` + * `toStorableVector`, `toStorableMVector`, `fromStorableVector` and `fromStorableMVector` + * `fromUnboxedVector` and `fromUnboxedMVector` + * `unsafeBoxedArray`, `unsafeNormalBoxedArray`, `unsafeFromBoxedVector` +* Removed deprecated `traverseAR`, `itraverseAR`, `traversePrimR` and `itraversePrimR` +* Removed: `imapMR`, `imapMR`, `iforMR`, and `iforMR` +* Renamed: + * `withMArray` to `withMArray_`, + * `withMArrayS` to `withMArrayS_` and + * `withMArrayST` to `withMArrayST_` +* Added versions that keep the artifact of mutable action: `withMArray`, `withMArrayS`, + `withMArrayST`. + # 0.4.5 * Addition of `computeIO` and `computePrimM` @@ -27,7 +309,6 @@ * Addition of `foldOuterSlice`, `ifoldOuterSlice`, `foldInnerSlice` and `ifoldInnerSlice`. Fix for [#56](https://github.com/lehins/massiv/issues/56) - # 0.4.2 * Fix loading empty `DS` stream arrays of unknown size. Fix for [#83](https://github.com/lehins/massiv/issues/83). @@ -89,6 +370,8 @@ * `generateArrayS` * Redefined most of the numeric operators with `Numeric` and `NumericFloat`. Will be required for SIMD operations. +* `Num`, `Fractional` and `Applicative` for `D` and `DI` changed behavior: instead of treating + singleton as a special array of any size it is treated as singleton. # 0.3.6 diff --git a/LICENSE b/LICENSE --- a/LICENSE +++ b/LICENSE @@ -1,4 +1,4 @@ -Copyright Alexey Kuleshevich (c) 2017-2019 +Copyright Alexey Kuleshevich (c) 2017-2022 All rights reserved. diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -1,10 +1,678 @@ # massiv -Efficient Haskell Arrays featuring Parallel computation +`massiv` is a Haskell library for array manipulation. Performance is one of its main goals, thus it +is capable of seamless parallelization of most of the operations provided by the library -There is a decent introduction to the library with some examples in the main -[README](https://github.com/lehins/massiv/blob/master/README.md) on github. +The name for this library comes from the Russian word Massiv (Масси́в), which means an Array. -See [massiv-io](https://hackage.haskell.org/package/massiv-io) for ability to read/write images. +## Status +| Language | Github Actions | Coveralls | Gitter.im | +|:--------:|:--------------:|:---------:|:---------:| +| ![GitHub top language][GHL] | [![GA-CI][GA-B]][GA-L] | [![Coveralls][Co-B]][Co-L] | [![Gitter][Gi-B]][Gi-L] +[GHL]: https://img.shields.io/github/languages/top/lehins/massiv.svg +[GA-B]: https://github.com/lehins/massiv/actions/workflows/haskell.yml/badge.svg?branch=master +[GA-L]: https://github.com/lehins/massiv/actions/workflows/haskell.yml +[Co-B]: https://coveralls.io/repos/github/lehins/massiv/badge.svg?branch=master +[Co-L]: https://coveralls.io/github/lehins/massiv?branch=master +[Gi-B]: https://badges.gitter.im/haskell-massiv/Lobby.svg +[Gi-L]: https://app.gitter.im/#/room/#haskell-massiv_Lobby:gitter.im + +| Package | Hackage | Nightly | LTS | +|:-------------------|:-------:|:-------:|:---:| +| [`massiv`](https://github.com/lehins/massiv/tree/master/massiv)| [![Hackage](https://img.shields.io/hackage/v/massiv.svg)](https://hackage.haskell.org/package/massiv)| [![Nightly](https://www.stackage.org/package/massiv/badge/nightly)](https://www.stackage.org/nightly/package/massiv)| [![LTS](https://www.stackage.org/package/massiv/badge/lts)](https://www.stackage.org/lts/package/massiv)| +| [`massiv-test`](https://github.com/lehins/massiv/tree/master/massiv-test)| [![Hackage](https://img.shields.io/hackage/v/massiv-test.svg)](https://hackage.haskell.org/package/massiv-test)| [![Nightly](https://www.stackage.org/package/massiv-test/badge/nightly)](https://www.stackage.org/nightly/package/massiv-test)| [![LTS](https://www.stackage.org/package/massiv-test/badge/lts)](https://www.stackage.org/lts/package/massiv-test)| +| [`haskell-scheduler`](https://github.com/lehins/haskell-scheduler)| [![Hackage](https://img.shields.io/hackage/v/scheduler.svg)](https://hackage.haskell.org/package/scheduler)| [![Nightly](https://www.stackage.org/package/scheduler/badge/nightly)](https://www.stackage.org/nightly/package/scheduler)| [![LTS](https://www.stackage.org/package/scheduler/badge/lts)](https://www.stackage.org/lts/package/scheduler)| + +## Introduction + +Everything in the library revolves around an `Array r ix e` - a data family for anything that can be +thought of as an array. The type variables, from the end, are: + +* `e` - element of an array. +* `ix` - an index that will map to an actual element. The index must be an instance of the `Index` + class with the default one being an `Ix n` type family and an optional being tuples of `Int`s. +* `r` - underlying representation. There are two main categories of representations described below. + +### Manifest + +These are your classical arrays that are located in memory and allow constant time lookup of +elements. Another main property they share is that they have a mutable interface. An `Array` with +manifest representation can be thawed into a mutable `MArray` and then frozen back into its +immutable counterpart after some destructive operation is applied to the mutable copy. The +differences among representations below is in the way that elements are being accessed in memory: + + * `P` - Array with elements that are an instance of `Prim` type class, i.e. common Haskell + primitive types: `Int`, `Word`, `Char`, etc. It is backed by unpinned memory and based on + [`ByteArray`](https://hackage.haskell.org/package/primitive/docs/Data-Primitive-ByteArray.html#t:ByteArray). + * `U` - Unboxed arrays. The elements are instances of the + [`Unbox`](https://hackage.haskell.org/package/vector/docs/Data-Vector-Unboxed.html#t:Vector) + type class. Usually just as fast as `P`, but has a slightly wider range of data types that it + can work with. Notable data types that can be stored as elements are `Bool`, tuples and `Ix n`. + * `S` - Storable arrays. Backed by pinned memory and based on `ForeignPtr`, while elements are + instances of the `Storable` type class. + * `B` - Boxed arrays that don't have restrictions on their elements, since they are represented + as pointers to elements, thus making them the slowest type of array, but also the most + general. Arrays of this representation are element strict, in other words its elements are + kept in Weak-Head Normal Form (WHNF). + * `BN` - Also boxed arrays, but unlike the other representation `B`, its elements are in Normal + Form, i.e. in a fully evaluated state and no thunks or memory leaks are possible. It does + require an `NFData` instance for the elements though. + * `BL` - Boxed lazy array. Just like `B` and `BN`, except values are evaluated on demand. + +### Delayed + +Main trait of delayed arrays is that they do not exist in memory and instead describe the contents +of an array as a function or a composition of functions. In fact all of the fusion capabilities in +`massiv` can be attributed to delayed arrays. + + * `D` - Delayed "pull" array is just a function from an index to an element: `(ix -> + e)`. Therefore indexing into this type of array is not possible, instead elements are evaluated + with the `evaluateM` function each time when applied to an index. It gives us a nice ability to + compose functions together when applied to an array and possibly even fold over without ever + allocating intermediate manifest arrays. + * `DW` - Delayed windowed array is very similar to the version above, except it has two functions + that describe it, one for the near border elements and one for the interior, aka. the + window. This is used for [`Stencil`](stencil) computation and things that derive from it, such as + convolution, for instance. + * `DL` - Delayed "push" array contains a monadic action that describes how an array can be loaded + into memory. This is most useful for composing arrays together. + * `DS` - Delayed stream array is a sequence of elements, possibly even an infinite one. This is + most useful for situations when we don't know the size of our resulting array ahead of time, + which is common in operations such as `filter`, `mapMaybe`, `unfold` etc. Naturally, in the end + we can only load such an array into a flat vector. + * `DI` - Is just like `D`, except loading is interleaved and is useful for parallel loading + arrays with unbalanced computation, such as Mandelbrot set or ray tracing, for example. + +## Construct + +Creating a delayed type of array allows us to fuse any future operations we decide to perform on +it. Let's look at this example: + +```haskell +λ> import Data.Massiv.Array as A +λ> makeVectorR D Seq 10 id +Array D Seq (Sz1 10) + [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9 ] +``` + +Here we created a delayed vector of size 10, which is in reality just an `id` function from its +index to an element (see the [Computation](#computation) section for the meaning of `Seq`). So let's +go ahead and square its elements + +```haskell +λ> vec = makeVectorR D Seq 10 id +λ> evaluateM vec 4 +4 +λ> vec2 = A.map (^ (2 :: Int)) vec +λ> evaluateM vec2 4 +16 +``` + +It's not that exciting, since every time we call `evaluateM` it will recompute the element, __every +time__, therefore this function should be avoided at all costs! Instead we can use all of the +functions that take `Source` like arrays and then fuse that computation together by calling +`compute`, or a handy `computeAs` function and only afterwards apply an `indexM` function or its +partial synonym: `(!)`. Any delayed array can also be reduced using one of the folding functions, +thus completely avoiding any memory allocation, or converted to a list, if that's what you need: + +```haskell +λ> vec2U = computeAs U vec2 +λ> vec2U +Array U Seq (Sz1 10) + [ 0, 1, 4, 9, 16, 25, 36, 49, 64, 81 ] +λ> vec2U ! 4 +16 +λ> toList vec2U +[0,1,4,9,16,25,36,49,64,81] +λ> A.sum vec2U +285 +``` + +There is a whole multitude of ways to construct arrays: + * by using one of many helper functions: `makeArray`, `range`, `rangeStepFrom`, `enumFromN`, etc. + * through conversion: from lists, from `Vector`s in `vector` library, from `ByteString`s in + `bytestring`; + * with a mutable interface in `PrimMonad` (`IO`, `ST`, etc.), eg: `makeMArray`, + `generateArray`, `unfoldrPrim`, etc. + +It's worth noting that, in the next example, nested lists will be loaded into an unboxed manifest +array and the sum of its elements will be computed in parallel on all available cores. + +```haskell +λ> A.sum (fromLists' Par [[0,0,0,0,0],[0,1,2,3,4],[0,2,4,6,8]] :: Array U Ix2 Double) +30.0 +``` + +The above wouldn't run in parallel in ghci of course, as the program would have to be compiled with +`ghc` using `-threaded -with-rtsopts=-N` flags in order to use all available cores. Alternatively we +could compile with the `-threaded` flag and then pass the number of capabilities directly to the +runtime with `+RTS -N`, where `` is the number of cores you'd like to utilize. + +## Index + +The main `Ix n` closed type family can be somewhat confusing, but there is no need to fully +understand how it works in order to start using it. GHC might ask you for the `DataKinds` language +extension if `IxN n` is used in a type signature, but there are type and pattern synonyms for the +first five dimensions: `Ix1`, `Ix2`, `Ix3`, `Ix4` and `Ix5`. + +There are three distinguishable constructors for the index: + +* The first one is simply an int: `Ix1 = Ix 1 = Int`, therefore vectors can be indexed in a usual way + without some extra wrapping data type, just as it was demonstrated in a previous section. +* The second one is `Ix2` for operating on 2-dimensional arrays and has a constructor `:.` + +```haskell +λ> makeArrayR D Seq (Sz (3 :. 5)) (\ (i :. j) -> i * j) +Array D Seq (Sz (3 :. 5)) + [ [ 0, 0, 0, 0, 0 ] + , [ 0, 1, 2, 3, 4 ] + , [ 0, 2, 4, 6, 8 ] + ] +``` + +* The third one is `IxN n` and is designed for working with N-dimensional arrays, and has a similar + looking constructor `:>`, except that it can be chained indefinitely on top of `:.` + +```haskell +λ> arr3 = makeArrayR P Seq (Sz (3 :> 2 :. 5)) (\ (i :> j :. k) -> i * j + k) +λ> :t arr3 +arr3 :: Array P (IxN 3) Int +λ> arr3 +Array P Seq (Sz (3 :> 2 :. 5)) + [ [ [ 0, 1, 2, 3, 4 ] + , [ 0, 1, 2, 3, 4 ] + ] + , [ [ 0, 1, 2, 3, 4 ] + , [ 1, 2, 3, 4, 5 ] + ] + , [ [ 0, 1, 2, 3, 4 ] + , [ 2, 3, 4, 5, 6 ] + ] + ] +λ> arr3 ! (2 :> 1 :. 4) +6 +λ> ix10 = 10 :> 9 :> 8 :> 7 :> 6 :> 5 :> 4 :> 3 :> 2 :. 1 +λ> :t ix10 +ix10 :: IxN 10 +λ> ix10 -- 10-dimensional index +10 :> 9 :> 8 :> 7 :> 6 :> 5 :> 4 :> 3 :> 2 :. 1 +``` + +Here is how we can construct a 4-dimensional array and sum its elements in constant memory: + +```haskell +λ> arr = makeArrayR D Seq (Sz (10 :> 20 :> 30 :. 40)) $ \ (i :> j :> k :. l) -> (i * j + k) * k + l +λ> :t arr -- a 4-dimensional array +arr :: Array D (IxN 4) Int +λ> A.sum arr +221890000 +``` + +There are quite a few helper functions that can operate on indices, but these are only needed when +writing functions that work for arrays of arbitrary dimension, as such they are scarcely used: + +```haskell +λ> pullOutDim' ix10 5 +(5,10 :> 9 :> 8 :> 7 :> 6 :> 4 :> 3 :> 2 :. 1) +λ> unconsDim ix10 +(10,9 :> 8 :> 7 :> 6 :> 5 :> 4 :> 3 :> 2 :. 1) +λ> unsnocDim ix10 +(10 :> 9 :> 8 :> 7 :> 6 :> 5 :> 4 :> 3 :. 2,1) +``` + +All of the `Ix n` indices are instances of `Num` so basic numeric operations are made easier: + +```haskell +λ> (1 :> 2 :. 3) + (4 :> 5 :. 6) +5 :> 7 :. 9 +λ> 5 :: Ix4 +5 :> 5 :> 5 :. 5 +``` + +It is important to note that the size type is distinct from the index by the newtype wrapper `Sz +ix`. There is a constructor `Sz`, which will make sure that none of the dimensions are negative: + +```haskell +λ> Sz (2 :> 3 :. 4) +Sz (2 :> 3 :. 4) +λ> Sz (10 :> 2 :> -30 :. 4) +Sz (10 :> 2 :> 0 :. 4) +``` + +Same as with indices, there are helper pattern synonyms: `Sz1`, `Sz2`, `Sz3`, `Sz4` and `Sz5`. + +```haskell +λ> Sz3 2 3 4 +Sz (2 :> 3 :. 4) +λ> Sz4 10 2 (-30) 4 +Sz (10 :> 2 :> 0 :. 4) +``` + +As well as the `Num` instance: + +```haskell +λ> 4 :: Sz5 +Sz (4 :> 4 :> 4 :> 4 :. 4) +λ> (Sz2 1 2) + 3 +Sz (4 :. 5) +λ> (Sz2 1 2) - 3 +Sz (0 :. 0) +``` + +Alternatively tuples of `Int`s can be used for working with arrays, up to and including 5-tuples +(type synonyms: `Ix2T` .. `Ix5T`), but since tuples are polymorphic it is necessary to restrict the +resulting array type. Not all operations in the library support tuples, so it is advised to avoid +them for indexing. + +```haskell +λ> makeArray Seq (4, 20) (uncurry (*)) :: Array P Ix2T Int +(Array P Seq ((4,20)) + [ [ 0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0 ] + , [ 0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19 ] + , [ 0,2,4,6,8,10,12,14,16,18,20,22,24,26,28,30,32,34,36,38 ] + , [ 0,3,6,9,12,15,18,21,24,27,30,33,36,39,42,45,48,51,54,57 ] + ]) +λ> :i Ix2T +type Ix2T = (Int, Int) +``` + +There are helper functions that can go back and forth between tuples and `Ix n` indices. + +```haskell +λ> fromIx4 (3 :> 4 :> 5 :. 6) +(3,4,5,6) +λ> toIx5 (3, 4, 5, 6, 7) +3 :> 4 :> 5 :> 6 :. 7 +``` + +## Slicing + +In order to get a subsection of an array there is no need to recompute it, unless we want to free up +the no longer memory, of course. So, there are a few slicing, resizing and extraction operators that +can do it all in constant time, modulo the index manipulation: + +```haskell +λ> arr = makeArrayR U Seq (Sz (4 :> 2 :. 6)) fromIx3 +λ> arr !> 3 !> 1 +Array M Seq (Sz1 6) + [ (3,1,0), (3,1,1), (3,1,2), (3,1,3), (3,1,4), (3,1,5) ] +``` + +As you might suspect all of the slicing, indexing, extracting, resizing operations are partial, and +those are frowned upon in Haskell. So there are matching functions that can do the same operations +safely by using `MonadThrow` and thus returning `Nothing`, `Left SomeException` or throwing an +exception in case of `IO` on failure, for example: + +```haskell +λ> arr !?> 3 ??> 1 +Array M Seq (Sz1 6) + [ (3,1,0), (3,1,1), (3,1,2), (3,1,3), (3,1,4), (3,1,5) ] +λ> arr !?> 3 ??> 1 ?? 0 :: Maybe (Int, Int, Int) +Just (3,1,0) +``` + +In above examples we first take a slice at the 4th page (index 3, since we start at 0), then another +one at the 2nd row (index 1). While in the last example we also take 1st element at +position 0. Pretty neat, huh? Naturally, by doing a slice we always reduce dimension by one. We can +do slicing from the outside as well as from the inside: + +```haskell +λ> Ix1 1 ... 9 +Array D Seq (Sz1 10) + [ 1, 2, 3, 4, 5, 6, 7, 8, 9 ] +λ> a <- resizeM (Sz (3 :> 2 :. 4)) $ Ix1 11 ... 34 +λ> a +Array D Seq (Sz (3 :> 2 :. 4)) + [ [ [ 11, 12, 13, 14 ] + , [ 15, 16, 17, 18 ] + ] + , [ [ 19, 20, 21, 22 ] + , [ 23, 24, 25, 26 ] + ] + , [ [ 27, 28, 29, 30 ] + , [ 31, 32, 33, 34 ] + ] + ] +λ> a !> 0 +Array D Seq (Sz (2 :. 4)) + [ [ 11, 12, 13, 14 ] + , [ 15, 16, 17, 18 ] + ] +λ> a a (Dim 2, 0) +Array D Seq (Sz (3 :. 4)) + [ [ 11, 12, 13, 14 ] + , [ 19, 20, 21, 22 ] + , [ 27, 28, 29, 30 ] + ] +``` + +In order to extract sub-array while preserving dimensionality we can use `extractM` or `extractFromToM`. + +```haskell +λ> extractM (0 :> 1 :. 1) (Sz (3 :> 1 :. 2)) a +Array D Seq (Sz (3 :> 1 :. 2)) + [ [ [ 16, 17 ] + ] + , [ [ 24, 25 ] + ] + , [ [ 32, 33 ] + ] + ] +λ> extractFromToM (1 :> 0 :. 1) (3 :> 2 :. 4) a +Array D Seq (Sz (2 :> 2 :. 3)) + [ [ [ 20, 21, 22 ] + , [ 24, 25, 26 ] + ] + , [ [ 28, 29, 30 ] + , [ 32, 33, 34 ] + ] + ] +``` + +## Computation and parallelism + +There is a data type `Comp` that controls how elements will be computed when calling the `compute` +function. It has a few constructors, although most of the time either `Seq` or `Par` will be +sufficient: + +* `Seq` - computation will be done sequentially on one core (capability in ghc). +* `ParOn [Int]` - perform computation in parallel while pinning the workers to particular + cores. Providing an empty list will result in the computation being distributed over all + available cores, or better known in Haskell as capabilities. +* `ParN Word16` - similar to `ParOn`, except it simply specifies the number of cores to + use, with `0` meaning all cores. +* `Par` - isn't really a constructor but a `pattern` for constructing `ParOn []`, which + will result in Scheduler using all cores, thus should be used instead of `ParOn`. +* `Par'` - similar to `Par`, except it uses `ParN 0` underneath. + +Just to make sure a simple novice mistake is prevented, which I have seen in the past, make sure +your source code is compiled with `ghc -O2 -threaded -with-rtsopts=-N`, otherwise no parallelization +and poor performance are waiting for you. Also a bit later you might notice the `{-# INLINE funcName +#-}` pragma being used, oftentimes it is a good idea to do that, but not always required. It is +worthwhile to benchmark and experiment. + +## Stencil + +Instead of manually iterating over a multi-dimensional array and applying a function to each element, +while reading its neighboring elements (as you would do in an imperative language) in a functional +language it is much more efficient to apply a stencil function and let the library take care of all +of bounds checking and iterating in a cache friendly manner. + +What's a [stencil](https://en.wikipedia.org/wiki/Stencil_code)? It is a declarative way of +specifying a pattern for how elements of an array in a neighborhood will be used in order to update +each element of the newly created array. In massiv a `Stencil` is a function that can read the +neighboring elements of the stencil's _center_ (the zero index), and only those, and then outputs a +new value for the center element. + +![stencil](massiv-examples/files/stencil.png) + +Let's create a simple, but somewhat meaningful array and create an averaging stencil. There is +nothing special about the array itself, but the averaging filter is a stencil that sums the elements +in a [Moore neighborhood](https://en.wikipedia.org/wiki/Moore_neighborhood) and divides the result +by 9, i.e. finds the average of a 3 by 3 square. + +```haskell +arrLightIx2 :: Comp -> Sz Ix2 -> Array D Ix2 Double +arrLightIx2 comp arrSz = makeArray comp arrSz $ \ (i :. j) -> sin (fromIntegral (i * i + j * j)) +{-# INLINE arrLightIx2 #-} + +average3x3Filter :: Fractional a => Stencil Ix2 a a +average3x3Filter = makeStencil (Sz (3 :. 3)) (1 :. 1) $ \ get -> + ( get (-1 :. -1) + get (-1 :. 0) + get (-1 :. 1) + + get ( 0 :. -1) + get ( 0 :. 0) + get ( 0 :. 1) + + get ( 1 :. -1) + get ( 1 :. 0) + get ( 1 :. 1) ) / 9 +{-# INLINE average3x3Filter #-} +``` + +Here is what it would look like in GHCi. We create a delayed array with some funky periodic +function, and make sure it is computed prior to mapping an average stencil over it: + +```haskell +λ> arr = computeAs U $ arrLightIx2 Par (Sz (600 :. 800)) +λ> :t arr +arr :: Array U Ix2 Double +λ> :t mapStencil Edge average3x3Filter arr +mapStencil Edge average3x3Filter arr :: Array DW Ix2 Double +``` + +As you can see, that operation produced an array of the earlier mentioned representation Delayed +Windowed `DW`. In its essence `DW` is an array type that does no bounds checking in order to gain +performance, except when it's near the border, where it uses a border resolution technique supplied +by the user (`Edge` in the example above). Currently it is used only in stencils and not much else +can be done to an array of this type besides further computing it into a manifest representation. + +This example will be continued in the next section, but before that I would like to mention that +some might notice that it looks very much like convolution, and in fact convolution can be +implemented with a stencil. There is a helper function `makeConvolutionStencil` that lets +you do just that. For the sake of example we'll do a sum of all neighbors by hand instead: + +```haskell +sum3x3Filter :: Fractional a => Stencil Ix2 a a +sum3x3Filter = makeConvolutionStencil (Sz (3 :. 3)) (1 :. 1) $ \ get -> + get (-1 :. -1) 1 . get (-1 :. 0) 1 . get (-1 :. 1) 1 . + get ( 0 :. -1) 1 . get ( 0 :. 0) 1 . get ( 0 :. 1) 1 . + get ( 1 :. -1) 1 . get ( 1 :. 0) 1 . get ( 1 :. 1) 1 +{-# INLINE sum3x3Filter #-} +``` + +There is not a single plus or multiplication sign, that is because convolutions is actually +summation of elements multiplied by a kernel element, so instead we have composition of functions +applied to an offset index and a multiplier. After we map that stencil, we can further divide each +element of the array by 9 in order to get the average. Yeah, I lied a bit, `Array DW ix` is an +instance of `Functor` class, so we can map functions over it, which will be fused as with a regular +`D`elayed array: + +```haskell +computeAs U $ fmap (/9) $ mapStencil Edge sum3x3Filter arr +``` + +If you are still confused of what a stencil is, but you are familiar with [Conway's Game of +Life](https://en.wikipedia.org/wiki/Conway%27s_Game_of_Life) this should hopefully clarify it a +bit more. The function `life` below is a single iteration of Game of Life: + +```haskell +lifeRules :: Word8 -> Word8 -> Word8 +lifeRules 0 3 = 1 +lifeRules 1 2 = 1 +lifeRules 1 3 = 1 +lifeRules _ _ = 0 + +lifeStencil :: Stencil Ix2 Word8 Word8 +lifeStencil = makeStencil (Sz (3 :. 3)) (1 :. 1) $ \ get -> + lifeRules (get (0 :. 0)) $ get (-1 :. -1) + get (-1 :. 0) + get (-1 :. 1) + + get ( 0 :. -1) + get ( 0 :. 1) + + get ( 1 :. -1) + get ( 1 :. 0) + get ( 1 :. 1) + +life :: Array S Ix2 Word8 -> Array S Ix2 Word8 +life = compute . mapStencil Wrap lifeStencil +``` + + + +The full working example that uses GLUT and OpenGL is located in +[GameOfLife](massiv-examples/GameOfLife/app/GameOfLife.hs). You can run it if you have the GLUT +dependencies installed: + +```bash +$ cd massiv-examples && stack run GameOfLife +``` + +# massiv-io + +In order to do anything useful with arrays we often need to be able to read some data from a +file. Considering that most common array-like files are images, +[massiv-io](https://github.com/lehins/massiv-io) provides an interface to read, write and display +images in common formats using Haskell native JuicyPixels and Netpbm packages. + +[Color](https://github.com/lehins/Color) package provides a variety of color spaces and conversions +between them, which are used by `massiv-io` package as pixels during reading and writing images. + +An earlier example wasn't particularly interesting, since we couldn't visualize what is actually +going on, so let's expand on it: + +```haskell +import Data.Massiv.Array +import Data.Massiv.Array.IO + +main :: IO () +main = do + let arr = computeAs S $ arrLightIx2 Par (600 :. 800) + toImage :: + (Functor (Array r Ix2), Load r Ix2 (Pixel (Y' SRGB) Word8)) + => Array r Ix2 Double + -> Image S (Y' SRGB) Word8 + toImage = computeAs S . fmap (PixelY' . toWord8) + lightPath = "files/light.png" + lightImage = toImage $ delay arr + lightAvgPath = "files/light_avg.png" + lightAvgImage = toImage $ mapStencil Edge (avgStencil 3) arr + lightSumPath = "files/light_sum.png" + lightSumImage = toImage $ mapStencil Edge (sumStencil 3) arr + writeImage lightPath lightImage + putStrLn $ "written: " ++ lightPath + writeImage lightAvgPath lightAvgImage + putStrLn $ "written: " ++ lightAvgPath + writeImage lightSumPath lightSumImage + putStrLn $ "written: " ++ lightSumPath + displayImageUsing defaultViewer True . computeAs S + =<< concatM 1 [lightAvgImage, lightImage, lightSumImage] +``` + +`massiv-examples/vision/files/light.png`: + +![Light](massiv-examples/vision/files/light.png) + +`massiv-examples/vision/files/light_avg.png`: + +![Light Average](massiv-examples/vision/files/light_avg.png) + + +The full example is in the example [vision](massiv-examples/vision/app/AvgSum.hs) package and if you +have `stack` installed you can run it as: + +```bash +$ cd massiv-examples && stack run avg-sum +``` + +# Other libraries + +A natural question might come to mind: Why even bother with a new array library when we already have +a few really good ones in the Haskell world? The main reasons for me are performance and +usability. I personally felt like there was much room for improvement before I even started working on +this package, and it seems like it turned out to be true. For example, the most common goto library +for dealing with multidimensional arrays and parallel computation used to be +[Repa](https://hackage.haskell.org/package/repa), which I personally was a big fan of for quite some +time, to the point that I even wrote a [Haskell Image +Processing](https://hackage.haskell.org/package/hip) library based on top of it. + +Here is a quick summary of how `massiv` is better than `Repa`: + +* It is actively maintained. +* Much more sophisticated scheduler. It is resumable and is capable of handling nested parallel + computation. +* Improved indexing data types. +* Safe stencils for arbitrary dimensions, not only 2D convolution. Stencils are composable +* Improved performance on almost all operations. +* Structural parallel folds (i.e. left/right - direction is preserved) +* Super easy slicing. +* Extensive mutable interface +* More fusion capabilities with delayed stream and push array representations. +* Delayed arrays aren't indexable, only Manifest are (saving user from common pitfall in Repa of + trying to read elements of delayed array) + +As far as usability of the library goes, it is very subjective, thus I'll let you be a judge of +that. When talking about performance it is the facts that do matter. Thus, let's not continue this +discussion in pure abstract words, below is a glimpse into benchmarks against Repa library running +with GHC 8.8.4 on Intel® Core™ i7-3740QM CPU @ 2.70GHz × 8 + +[Matrix multiplication](https://en.wikipedia.org/wiki/Matrix_multiplication_algorithm): + +``` +benchmarking Repa/MxM U Double - (500x800 X 800x500)/Par +time 120.5 ms (115.0 ms .. 127.2 ms) + 0.998 R² (0.996 R² .. 1.000 R²) +mean 124.1 ms (121.2 ms .. 127.3 ms) +std dev 5.212 ms (2.422 ms .. 6.620 ms) +variance introduced by outliers: 11% (moderately inflated) + +benchmarking Massiv/MxM U Double - (500x800 X 800x500)/Par +time 41.46 ms (40.67 ms .. 42.45 ms) + 0.998 R² (0.994 R² .. 0.999 R²) +mean 38.45 ms (37.22 ms .. 39.68 ms) +std dev 2.342 ms (1.769 ms .. 3.010 ms) +variance introduced by outliers: 19% (moderately inflated) +``` + +[Sobel operator](https://en.wikipedia.org/wiki/Sobel_operator): +``` +benchmarking Sobel/Par/Operator - Repa +time 17.82 ms (17.30 ms .. 18.32 ms) + 0.997 R² (0.994 R² .. 0.998 R²) +mean 17.42 ms (17.21 ms .. 17.69 ms) +std dev 593.0 μs (478.1 μs .. 767.5 μs) +variance introduced by outliers: 12% (moderately inflated) + +benchmarking Sobel/Par/Operator - Massiv +time 7.421 ms (7.230 ms .. 7.619 ms) + 0.994 R² (0.991 R² .. 0.997 R²) +mean 7.537 ms (7.422 ms .. 7.635 ms) +std dev 334.3 μs (281.3 μs .. 389.9 μs) +variance introduced by outliers: 20% (moderately inflated) +``` + +Sum all elements of a 2D array: + +``` +benchmarking Sum/Seq/Repa +time 539.7 ms (523.2 ms .. 547.9 ms) + 1.000 R² (1.000 R² .. 1.000 R²) +mean 540.1 ms (535.7 ms .. 543.2 ms) +std dev 4.727 ms (2.208 ms .. 6.609 ms) +variance introduced by outliers: 19% (moderately inflated) + +benchmarking Sum/Seq/Vector +time 16.95 ms (16.78 ms .. 17.07 ms) + 0.999 R² (0.998 R² .. 1.000 R²) +mean 17.23 ms (17.13 ms .. 17.43 ms) +std dev 331.4 μs (174.1 μs .. 490.0 μs) + +benchmarking Sum/Seq/Massiv +time 16.78 ms (16.71 ms .. 16.85 ms) + 1.000 R² (1.000 R² .. 1.000 R²) +mean 16.80 ms (16.76 ms .. 16.88 ms) +std dev 127.8 μs (89.95 μs .. 186.2 μs) + +benchmarking Sum/Par/Repa +time 81.76 ms (78.52 ms .. 84.37 ms) + 0.997 R² (0.990 R² .. 1.000 R²) +mean 79.20 ms (78.03 ms .. 80.91 ms) +std dev 2.613 ms (1.565 ms .. 3.736 ms) + +benchmarking Sum/Par/Massiv +time 8.102 ms (7.971 ms .. 8.216 ms) + 0.999 R² (0.998 R² .. 1.000 R²) +mean 7.967 ms (7.852 ms .. 8.028 ms) +std dev 236.4 μs (168.4 μs .. 343.2 μs) +variance introduced by outliers: 11% (moderately inflated) +``` + +Here is also a blog post that compares [Performance of Haskell Array libraries through Canny edge detection](https://alexey.kuleshevi.ch/blog/2020/07/10/canny-benchmarks/) + +# Further resources on learning `massiv`: + +* [2021 - Haskell eXchange - Multi-dimensional Arrays that Do Not Exist](#2021---haskell-exchange---multi-dimensional-arrays-that-do-not-exist) +* [2019 - Monadic Party - Haskell arrays with Massiv](https://github.com/lehins/talks#2019---monadic-party---haskell-arrays-with-massiv) +* [2018 - Monadic Warsaw #14 - Haskell arrays that are easy and fast](https://github.com/lehins/talks#2018---monadic-warsaw-14---haskell-arrays-that-are-easy-and-fast) diff --git a/Setup.hs b/Setup.hs --- a/Setup.hs +++ b/Setup.hs @@ -1,33 +1,4 @@ -{-# LANGUAGE CPP #-} -{-# OPTIONS_GHC -Wall #-} -module Main (main) where - -#ifndef MIN_VERSION_cabal_doctest -#define MIN_VERSION_cabal_doctest(x,y,z) 0 -#endif - -#if MIN_VERSION_cabal_doctest(1,0,0) - -import Distribution.Extra.Doctest ( defaultMainWithDoctests ) -main :: IO () -main = defaultMainWithDoctests "doctests" - -#else - -#ifdef MIN_VERSION_Cabal --- If the macro is defined, we have new cabal-install, --- but for some reason we don't have cabal-doctest in package-db --- --- Probably we are running cabal sdist, when otherwise using new-build --- workflow -#warning You are configuring this package without cabal-doctest installed. \ - The doctests test-suite will not work as a result. \ - To fix this, install cabal-doctest before configuring. -#endif - import Distribution.Simple main :: IO () main = defaultMain - -#endif diff --git a/include/massiv.h b/include/massiv.h --- a/include/massiv.h +++ b/include/massiv.h @@ -3,9 +3,9 @@ #define MASSIV_INCLUDE #if MASSIV_UNSAFE_CHECKS -#define INDEX_CHECK(name, s, f) (indexWith __FILE__ __LINE__ (name) (s) (f)) +#define HAS_CALL_STACK (HasCallStack) #else -#define INDEX_CHECK(name, s, f) ((f)) +#define HAS_CALL_STACK () #endif #endif diff --git a/massiv.cabal b/massiv.cabal --- a/massiv.cabal +++ b/massiv.cabal @@ -1,5 +1,5 @@ name: massiv -version: 0.4.5.0 +version: 1.0.5.0 synopsis: Massiv (Массив) is an Array Library. description: Multi-dimensional Arrays with fusion, stencils and parallel computation. homepage: https://github.com/lehins/massiv @@ -7,50 +7,51 @@ license-file: LICENSE author: Alexey Kuleshevich maintainer: alexey@kuleshevi.ch -copyright: 2018-2019 Alexey Kuleshevich -category: Data, Data Structures, Parallelism -build-type: Custom +copyright: 2018-2022 Alexey Kuleshevich +category: Array, Data, Data Structures, Parallelism +build-type: Simple extra-source-files: README.md , CHANGELOG.md cabal-version: >=1.10 -tested-with: GHC == 8.4.3 - , GHC == 8.4.4 - , GHC == 8.6.3 - , GHC == 8.6.4 - , GHC == 8.6.5 - , GHC == 8.8.1 +tested-with: GHC == 8.0.2 + , GHC == 8.2.2 + , GHC == 8.4.4 + , GHC == 8.6.5 + , GHC == 8.8.4 + , GHC == 8.10.7 + , GHC == 9.0.2 + , GHC == 9.2.8 + , GHC == 9.4.8 + , GHC == 9.6.6 + , GHC == 9.8.4 + , GHC == 9.10.1 + , GHC == 9.12.1 flag unsafe-checks description: Enable all the bounds checks for unsafe functions at the cost of - significant performance penalty + performance penalty default: False manual: True -custom-setup - setup-depends: - base - , Cabal - , cabal-doctest >=1.0.6 - library hs-source-dirs: src exposed-modules: Data.Massiv.Array , Data.Massiv.Array.Delayed , Data.Massiv.Array.Manifest , Data.Massiv.Array.Manifest.Vector - , Data.Massiv.Array.Manifest.Vector.Stream , Data.Massiv.Array.Mutable , Data.Massiv.Array.Mutable.Algorithms , Data.Massiv.Array.Mutable.Atomic , Data.Massiv.Array.Numeric , Data.Massiv.Array.Numeric.Integral , Data.Massiv.Array.Stencil - , Data.Massiv.Array.Stencil.Unsafe , Data.Massiv.Array.Unsafe , Data.Massiv.Core , Data.Massiv.Core.Index , Data.Massiv.Core.List , Data.Massiv.Core.Operations + , Data.Massiv.Vector + , Data.Massiv.Vector.Stream other-modules: Data.Massiv.Array.Delayed.Interleaved , Data.Massiv.Array.Delayed.Pull @@ -63,6 +64,7 @@ , Data.Massiv.Array.Manifest.Primitive , Data.Massiv.Array.Manifest.Storable , Data.Massiv.Array.Manifest.Unboxed + , Data.Massiv.Array.Mutable.Internal , Data.Massiv.Array.Ops.Construct , Data.Massiv.Array.Ops.Fold , Data.Massiv.Array.Ops.Fold.Internal @@ -72,22 +74,26 @@ , Data.Massiv.Array.Ops.Transform , Data.Massiv.Array.Stencil.Convolution , Data.Massiv.Array.Stencil.Internal + , Data.Massiv.Array.Stencil.Unsafe , Data.Massiv.Core.Common , Data.Massiv.Core.Exception , Data.Massiv.Core.Index.Internal , Data.Massiv.Core.Index.Ix , Data.Massiv.Core.Index.Stride , Data.Massiv.Core.Index.Tuple - , Data.Massiv.Core.Iterator + , Data.Massiv.Core.Index.Iterator + , Data.Massiv.Core.Loop + , Data.Massiv.Vector.Unsafe build-depends: base >= 4.9 && < 5 , bytestring - , data-default-class , deepseq , exceptions - , scheduler >= 1.4.0 - , primitive + , scheduler >= 2.0.0 && < 3 + , primitive >= 0.7.1.0 + , random >= 1.2.0 , unliftio-core - , vector + , vector >= 0.12 + , vector-stream include-dirs: include install-includes: massiv.h @@ -98,23 +104,33 @@ default-language: Haskell2010 ghc-options: -Wall -Wincomplete-record-updates - -Wincomplete-uni-patterns -Wredundant-constraints + if impl(ghc >= 8.2) + ghc-options: + -Wno-simplifiable-class-constraints + -Wincomplete-uni-patterns + else + ghc-options: + -Wno-incomplete-patterns + -Wno-unused-imports + -Wno-unrecognised-pragmas test-suite doctests type: exitcode-stdio-1.0 hs-source-dirs: tests main-is: doctests.hs - build-depends: base + build-depends: base >= 4.9 && < 5 , doctest >=0.15 - , QuickCheck - , massiv - , mersenne-random-pure64 - , random - , splitmix >= 0.0.1 - , template-haskell + if impl(ghc >= 8.2) && impl(ghc < 8.10) + build-depends: QuickCheck + , massiv + , mersenne-random-pure64 + , random >= 1.2.0 + , mwc-random >= 0.15.0.1 + , splitmix >= 0.0.1 default-language: Haskell2010 source-repository head type: git location: https://github.com/lehins/massiv + subdir: massiv 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 @@ -1,6 +1,8 @@ +{-# OPTIONS_GHC -fno-warn-duplicate-exports #-} + -- | -- Module : Data.Massiv.Array --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental @@ -19,10 +21,13 @@ -- element is a pointer to the actual value, therefore it is also the slowest -- representation. Elements are kept in a Weak Head Normal Form (WHNF). -- --- * `N` - Similar to `B`, is also a boxed type, except it's elements are always kept in a Normal +-- * `BN` - Similar to `B`, it is also a boxed type, except its elements are always kept in a Normal -- Form (NF). This property is very useful for parallel processing, i.e. when calling -- `compute` you do want all of your elements to be fully evaluated. -- +-- * `BL` - Similar to `B`, it is also a boxed type, but lazy. Its elements are not evaluated when +-- array is computed. +-- -- * `S` - Is a type of array that is backed by pinned memory, therefore pointers to those arrays -- can be passed to FFI calls, because Garbage Collector (GC) is guaranteed not to move -- it. Elements must be an instance of `Storable` class. It is just as efficient as `P` and @@ -33,9 +38,6 @@ -- * `P` - Array that can hold Haskell primitives, such as `Int`, `Word`, `Double`, etc. Any element -- must be an instance of `Prim` class. -- --- * `M` - General manifest array type, that any of the above representations can be converted to in --- constant time using `toManifest`. --- -- There are also array representations that only describe how values for its elements can be -- computed or loaded into memory, as such, they are represented by functions and do not impose the -- memory overhead, that is normally associated with arrays. They are needed for proper fusion and @@ -61,102 +63,113 @@ -- -- Other Array types: -- --- * `L` and `LN` - those types aren't particularly useful on their own, but because of their unique --- ability to be converted to and from nested lists in constant time, provide a perfect --- intermediary for lists \<-> array conversion. +-- * `L` - this type isn't particularly useful on its own, but because it has unique ability to be +-- converted to and from nested lists in constant time, it provides a perfect intermediary for +-- conversion of nested lists into manifest arrays. -- -- Most of the `Manifest` arrays are capable of in-place mutation. Check out -- "Data.Massiv.Array.Mutable" module for available functionality. -- --- Many of the function names exported by this package will clash with the ones --- from "Prelude", hence it can be more convenient to import like this: +-- Many of the function names exported by this package will clash with the ones from "Prelude", +-- hence it can be more convenient to import like this: -- -- @ -- import Prelude as P -- import Data.Massiv.Array as A -- @ --- -module Data.Massiv.Array - ( -- * Construct - module Data.Massiv.Array.Ops.Construct +module Data.Massiv.Array ( + -- * Construct + module Data.Massiv.Array.Ops.Construct, + -- * Compute - , getComp - , setComp - , compute - , computeS - , computeIO - , computePrimM - , computeAs - , computeProxy - , computeSource - , computeWithStride - , computeWithStrideAs - , clone - , convert - , convertAs - , convertProxy - , fromRaggedArrayM - , fromRaggedArray' + getComp, + setComp, + appComp, + compute, + computeS, + computeP, + computeIO, + computePrimM, + computeAs, + computeProxy, + computeSource, + computeWithStride, + computeWithStrideAs, + clone, + convert, + convertAs, + convertProxy, + fromRaggedArrayM, + fromRaggedArray', + + -- * Vector + module Data.Massiv.Vector, + -- * Size - , size - , elemsCount - , isEmpty + size, + elemsCount, + isEmpty, + isNotEmpty, + isNull, + isNotNull, + -- * Indexing - , (!?) - , (!) - , (??) - , indexM - , index - , index' - , defaultIndex - , borderIndex - , evaluateM - , evaluate' + (!?), + (!), + (??), + indexM, + index, + index', + defaultIndex, + borderIndex, + evaluateM, + evaluate', + -- * Mapping - , module Data.Massiv.Array.Ops.Map - -- * Filtering - -- ** Maybe - , mapMaybeS - , imapMaybeS - , mapMaybeM - , imapMaybeM - -- ** Predicate - , filterS - , ifilterS - , filterM - , ifilterM - -- * Folding + module Data.Massiv.Array.Ops.Map, + -- * Folding -- $folding + module Data.Massiv.Array.Ops.Fold, - , module Data.Massiv.Array.Ops.Fold -- * Transforming - , module Data.Massiv.Array.Ops.Transform + module Data.Massiv.Array.Ops.Transform, + -- * Slicing - , module Data.Massiv.Array.Ops.Slice + module Data.Massiv.Array.Ops.Slice, + -- * Algorithms + -- ** Sorting - , quicksort - , tally + quicksort, + quicksortBy, + quicksortByM, + tally, + -- ** Iterations - , iterateUntil + iterateUntil, + -- * Conversion - , module Data.Massiv.Array.Manifest.List + module Data.Massiv.Array.Manifest.List, + -- * Mutable - , module Data.Massiv.Array.Mutable + module Data.Massiv.Array.Mutable, + -- * Core - , module Data.Massiv.Core + module Data.Massiv.Core, + -- * Representations - , module Data.Massiv.Array.Delayed - , module Data.Massiv.Array.Manifest + module Data.Massiv.Array.Delayed, + module Data.Massiv.Array.Manifest, + -- * Stencil - , module Data.Massiv.Array.Stencil + module Data.Massiv.Array.Stencil, + -- * Numeric Operations - , module Data.Massiv.Array.Numeric - ) where + module Data.Massiv.Array.Numeric, +) where import Data.Massiv.Array.Delayed -import Data.Massiv.Array.Delayed.Stream import Data.Massiv.Array.Manifest import Data.Massiv.Array.Manifest.Internal import Data.Massiv.Array.Manifest.List @@ -166,62 +179,31 @@ import Data.Massiv.Array.Ops.Fold import Data.Massiv.Array.Ops.Map import Data.Massiv.Array.Ops.Slice -import Data.Massiv.Array.Ops.Sort (quicksort, tally) +import Data.Massiv.Array.Ops.Sort import Data.Massiv.Array.Ops.Transform import Data.Massiv.Array.Stencil import Data.Massiv.Core import Data.Massiv.Core.Common -import Prelude as P hiding (all, and, any, enumFromTo, foldl, foldr, mapM, - mapM_, maximum, minimum, or, product, replicate, splitAt, - sum, zip) - - --- | Similar to `mapMaybeM`, but map with an index aware function. --- --- @since 0.4.1 -imapMaybeS :: Source r ix a => (ix -> a -> Maybe b) -> Array r ix a -> Array DS Ix1 b -imapMaybeS f arr = - mapMaybeS (uncurry f) $ makeArrayR D (getComp arr) (size arr) $ \ ix -> (ix, unsafeIndex arr ix) -{-# INLINE imapMaybeS #-} - --- | Similar to `mapMaybeM`, but map with an index aware function. --- --- @since 0.4.1 -imapMaybeM :: - (Source r ix a, Applicative f) => (ix -> a -> f (Maybe b)) -> Array r ix a -> f (Array DS Ix1 b) -imapMaybeM f arr = - mapMaybeM (uncurry f) $ makeArrayR D (getComp arr) (size arr) $ \ ix -> (ix, unsafeIndex arr ix) -{-# INLINE imapMaybeM #-} - --- | Similar to `filterS`, but map with an index aware function. --- --- @since 0.4.1 -ifilterS :: Source r ix a => (ix -> a -> Bool) -> Array r ix a -> Array DS Ix1 a -ifilterS f = - imapMaybeS $ \ix e -> - if f ix e - then Just e - else Nothing -{-# INLINE ifilterS #-} - +import Data.Massiv.Vector +import Prelude as P hiding ( + all, + and, + any, + enumFromTo, + foldl, + foldr, + mapM, + mapM_, + maximum, + minimum, + or, + product, + replicate, + splitAt, + sum, + zip, + ) --- | Similar to `filterM`, but map with an index aware function. +-- $folding -- --- @since 0.4.1 -ifilterM :: - (Source r ix a, Applicative f) => (ix -> a -> f Bool) -> Array r ix a -> f (Array DS Ix1 a) -ifilterM f = - imapMaybeM $ \ix e -> - (\p -> - if p - then Just e - else Nothing) <$> - f ix e -{-# INLINE ifilterM #-} - - -{- $folding - -All folding is done in a row-major order. - --} +-- All folding is done in a row-major order. diff --git a/src/Data/Massiv/Array/Delayed.hs b/src/Data/Massiv/Array/Delayed.hs --- a/src/Data/Massiv/Array/Delayed.hs +++ b/src/Data/Massiv/Array/Delayed.hs @@ -1,39 +1,45 @@ -- | -- Module : Data.Massiv.Array.Delayed --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Delayed - ( -- * Delayed +module Data.Massiv.Array.Delayed ( + -- * Delayed + -- ** Delayed Pull Array - D(..) - , delay + D (..), + delay, + liftArray2', + liftArray2M, + -- ** Delayed Push Array - , DL(..) - , toLoadArray - , makeLoadArrayS - , makeLoadArray - , fromStrideLoad + DL (..), + toLoadArray, + makeLoadArrayS, + makeLoadArray, + fromStrideLoad, + -- ** Delayed Stream Array - , DS(..) - , toStreamArray - , toSteps - , fromSteps + DS (..), + toStreamArray, + toSteps, + fromSteps, + -- ** Delayed Interleaved Array - , DI(..) - , toInterleaved - , fromInterleaved + DI (..), + toInterleaved, + fromInterleaved, + -- ** Delayed Windowed Array - , DW(..) - , Window(..) - , insertWindow - , getWindow - , dropWindow - , makeWindowedArray - ) where + DW (..), + Window (..), + insertWindow, + getWindow, + dropWindow, + makeWindowedArray, +) where import Data.Massiv.Array.Delayed.Interleaved import Data.Massiv.Array.Delayed.Pull diff --git a/src/Data/Massiv/Array/Delayed/Interleaved.hs b/src/Data/Massiv/Array/Delayed/Interleaved.hs --- a/src/Data/Massiv/Array/Delayed/Interleaved.hs +++ b/src/Data/Massiv/Array/Delayed/Interleaved.hs @@ -1,81 +1,82 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} + -- | -- Module : Data.Massiv.Array.Delayed.Interleaved --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Delayed.Interleaved - ( DI(..) - , toInterleaved - , fromInterleaved - ) where +module Data.Massiv.Array.Delayed.Interleaved ( + DI (..), + Array (..), + toInterleaved, + fromInterleaved, +) where import Data.Massiv.Array.Delayed.Pull import Data.Massiv.Core.Common import Data.Massiv.Core.List (L, showArrayList, showsArrayPrec) - -- | Delayed array that will be loaded in an interleaved fashion during parallel -- computation. +-- +-- /Warning/ - Will be deprecated in the next major version update. data DI = DI newtype instance Array DI ix e = DIArray { diArray :: Array D ix e - } deriving (Eq, Ord, Functor, Applicative, Foldable, Num, Floating, Fractional) + } + deriving (Eq, Ord, Functor, Applicative, Foldable) instance (Ragged L ix e, Show e) => Show (Array DI ix e) where showsPrec = showsArrayPrec diArray showList = showArrayList -instance Index ix => Construct DI ix e where - setComp c arr = arr { diArray = (diArray arr) { dComp = c } } +instance Strategy DI where + setComp c arr = arr{diArray = (diArray arr){dComp = c}} {-# INLINE setComp #-} + getComp = dComp . diArray + {-# INLINE getComp #-} + repr = DI - makeArray c sz = DIArray . makeArray c sz - {-# INLINE makeArray #-} +instance Index ix => Shape DI ix where + maxLinearSize = Just . SafeSz . elemsCount + {-# INLINE maxLinearSize #-} -instance Index ix => Resize DI ix where +instance Size DI where + size (DIArray arr) = size arr + {-# INLINE size #-} unsafeResize sz = DIArray . unsafeResize sz . diArray {-# INLINE unsafeResize #-} -instance Index ix => Extract DI ix e where - unsafeExtract sIx newSz = DIArray . unsafeExtract sIx newSz . diArray - {-# INLINE unsafeExtract #-} - - instance Index ix => Load DI ix e where - size (DIArray arr) = size arr - {-# INLINE size #-} - getComp = dComp . diArray - {-# INLINE getComp #-} - loadArrayM scheduler (DIArray (DArray _ sz f)) uWrite = - loopM_ 0 (< numWorkers scheduler) (+ 1) $ \ !start -> - scheduleWork scheduler $ - iterLinearM_ sz start (totalElem sz) (numWorkers scheduler) (<) $ \ !k -> uWrite k . f - {-# INLINE loadArrayM #-} + makeArray c sz = DIArray . makeArray c sz + {-# INLINE makeArray #-} + iterArrayLinearST_ scheduler (DIArray darr@(DArray _ sz _)) uWrite = + loopA_ 0 (< numWorkers scheduler) (+ 1) $ \ !start -> + scheduleWork_ scheduler $ + iterLinearM_ sz start (totalElem sz) (numWorkers scheduler) (<) $ \ !k -> + uWrite k . unsafeIndex darr + {-# INLINE iterArrayLinearST_ #-} instance Index ix => StrideLoad DI ix e where - loadArrayWithStrideM scheduler stride resultSize arr uWrite = - let strideIx = unStride stride - DIArray (DArray _ _ f) = arr - in loopM_ 0 (< numWorkers scheduler) (+ 1) $ \ !start -> - scheduleWork scheduler $ - iterLinearM_ resultSize start (totalElem resultSize) (numWorkers scheduler) (<) $ - \ !i ix -> uWrite i (f (liftIndex2 (*) strideIx ix)) - {-# INLINE loadArrayWithStrideM #-} + iterArrayLinearWithStrideST_ scheduler stride resultSize (DIArray arr) uWrite = + loopA_ 0 (< numWorkers scheduler) (+ 1) $ \ !start -> + scheduleWork_ scheduler $ + iterLinearM_ resultSize start (totalElem resultSize) (numWorkers scheduler) (<) $ + \ !i ix -> uWrite i (unsafeIndex arr (liftIndex2 (*) (unStride stride) ix)) + {-# INLINE iterArrayLinearWithStrideST_ #-} -- | Convert a source array into an array that, when computed, will have its elemets evaluated out -- of order (interleaved amongst cores), hence making unbalanced computation better parallelizable. -toInterleaved :: Source r ix e => Array r ix e -> Array DI ix e +toInterleaved :: (Index ix, Source r e) => Array r ix e -> Array DI ix e toInterleaved = DIArray . delay {-# INLINE toInterleaved #-} 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 @@ -3,121 +3,150 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE RecordWildCards #-} {-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UndecidableInstances #-} + -- | -- Module : Data.Massiv.Array.Delayed.Pull --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Delayed.Pull - ( D(..) - , Array(..) - , delay - , eq - , ord - ) where +module Data.Massiv.Array.Delayed.Pull ( + D (..), + Array (..), + delay, + eqArrays, + compareArrays, + imap, + liftArray2', + liftArray2M, + unsafeExtract, + unsafeSlice, + unsafeInnerSlice, + zipWithInternal, +) where +import Control.Applicative import qualified Data.Foldable as F import Data.Massiv.Array.Ops.Fold.Internal as A -import Data.Massiv.Array.Manifest.Vector.Stream as S (steps) -import Data.Massiv.Core.Common -import Data.Massiv.Core.Operations +import Data.Massiv.Core.Common as A import Data.Massiv.Core.List (L, showArrayList, showsArrayPrec) +import Data.Massiv.Core.Operations +import qualified Data.Massiv.Vector.Stream as S import GHC.Base (build) import Prelude hiding (zipWith) #include "massiv.h" -- | Delayed representation. -data D = D deriving Show +data D + = D + deriving (Show) -data instance Array D ix e = DArray { dComp :: !Comp - , dSize :: !(Sz ix) - , dIndex :: ix -> e } +data instance Array D ix e = DArray + { dComp :: !Comp + , dSize :: !(Sz ix) + , dPrefIndex :: !(PrefIndex ix e) + } instance (Ragged L ix e, Show e) => Show (Array D ix e) where showsPrec = showsArrayPrec id showList = showArrayList -instance Index ix => Resize D ix where +instance Index ix => Shape D ix where + maxLinearSize = Just . SafeSz . elemsCount + {-# INLINE maxLinearSize #-} + +instance Size D where + size = dSize + {-# INLINE size #-} unsafeResize !sz !arr = - DArray (dComp arr) sz $ \ !ix -> - unsafeIndex arr (fromLinearIndex (size arr) (toLinearIndex sz ix)) + makeArrayLinear (dComp arr) sz (unsafeIndex arr . fromLinearIndex (size arr)) {-# INLINE unsafeResize #-} -instance Index ix => Extract D ix e where - unsafeExtract !sIx !newSz !arr = - DArray (dComp arr) newSz $ \ !ix -> - unsafeIndex arr (liftIndex2 (+) ix sIx) - {-# INLINE unsafeExtract #-} - - -instance Index ix => Construct D ix e where - setComp c arr = arr { dComp = c } +instance Strategy D where + setComp c arr = arr{dComp = c} {-# INLINE setComp #-} - - makeArray = DArray - {-# INLINE makeArray #-} - + getComp = dComp + {-# INLINE getComp #-} + repr = D -instance Index ix => Source D ix e where - unsafeIndex = INDEX_CHECK("(Source D ix e).unsafeIndex", size, dIndex) +instance Source D e where + unsafeIndex arr = + case dPrefIndex arr of + PrefIndex f -> f + PrefIndexLinear f -> f . toLinearIndex (size arr) {-# INLINE unsafeIndex #-} - -- unsafeLinearSlice ix sz arr = unsafeExtract ix sz (unsafeResize sz arr) - -- {-# INLINE unsafeLinearSlice #-} - - -instance ( Index ix - , Index (Lower ix) - , Elt D ix e ~ Array D (Lower ix) e - ) => - Slice D ix e where - unsafeSlice arr start cut@(SafeSz cutSz) dim = do - newSz <- dropDimM cutSz dim - return $ unsafeResize (SafeSz newSz) (unsafeExtract start cut arr) - {-# INLINE unsafeSlice #-} - - -instance (Elt D ix e ~ Array D (Lower ix) e, Index ix) => OuterSlice D ix e where - - unsafeOuterSlice !arr !i = - DArray (dComp arr) (snd (unconsSz (size arr))) (\ !ix -> unsafeIndex arr (consDim i ix)) + unsafeLinearIndex arr = + case dPrefIndex arr of + PrefIndex f -> f . fromLinearIndex (size arr) + PrefIndexLinear f -> f + {-# INLINE unsafeLinearIndex #-} + unsafePrefIndex = dPrefIndex + {-# INLINE unsafePrefIndex #-} + unsafeOuterSlice !arr !szL !i = + makeArray (dComp arr) szL (unsafeIndex arr . consDim i) {-# INLINE unsafeOuterSlice #-} + unsafeLinearSlice !o !sz arr = + makeArrayLinear (dComp arr) sz $ \ !i -> unsafeLinearIndex arr (i + o) + {-# INLINE unsafeLinearSlice #-} -instance (Elt D ix e ~ Array D (Lower ix) e, Index ix) => InnerSlice D ix e where +-- | /O(1)/ - Extract a portion of an array. Staring index and new size are +-- not validated. +unsafeExtract :: (Source r e, Index ix) => ix -> Sz ix -> Array r ix e -> Array D ix e +unsafeExtract !sIx !newSz !arr = + makeArray (getComp arr) newSz (unsafeIndex arr . liftIndex2 (+) sIx) +{-# INLINE unsafeExtract #-} - unsafeInnerSlice !arr (szL, _) !i = - DArray (dComp arr) szL (\ !ix -> unsafeIndex arr (snocDim ix i)) - {-# INLINE unsafeInnerSlice #-} +-- | /O(1)/ - Take a slice out of an array from within +unsafeSlice + :: (Source r e, Index ix, Index (Lower ix), MonadThrow m) + => Array r ix e + -> ix + -> Sz ix + -> Dim + -> m (Array D (Lower ix) e) +unsafeSlice arr start cut@(SafeSz cutSz) dim = do + newSz <- dropDimM cutSz dim + return $ unsafeResize (SafeSz newSz) (unsafeExtract start cut arr) +{-# INLINE unsafeSlice #-} +-- | /O(1)/ - Take a slice out of an array from the inside +unsafeInnerSlice + :: (Source r e, Index ix) => Array r ix e -> Sz (Lower ix) -> Int -> Array D (Lower ix) e +unsafeInnerSlice !arr szL !i = + DArray (getComp arr) szL $ PrefIndex (unsafeIndex arr . (`snocDim` i)) +{-# INLINE unsafeInnerSlice #-} instance (Eq e, Index ix) => Eq (Array D ix e) where - (==) = eq (==) + (==) = eqArrays (==) {-# INLINE (==) #-} instance (Ord e, Index ix) => Ord (Array D ix e) where - compare = ord compare + compare = compareArrays compare {-# INLINE compare #-} instance Functor (Array D ix) where - fmap f (DArray c sz g) = DArray c sz (f . g) + fmap f (DArray c sz g) = DArray c sz (fmap f g) {-# INLINE fmap #-} - + (<$) e (DArray c sz g) = DArray c sz (e <$ g) + {-# INLINE (<$) #-} instance Index ix => Applicative (Array D ix) where pure = singleton {-# INLINE pure #-} - (<*>) (DArray c1 (SafeSz sz1) uIndex1) (DArray c2 (SafeSz sz2) uIndex2) = - DArray (c1 <> c2) (SafeSz (liftIndex2 min sz1 sz2)) $ \ !ix -> - uIndex1 ix (uIndex2 ix) + (<*>) = liftArray2' id {-# INLINE (<*>) #-} - +#if MIN_VERSION_base(4,10,0) + liftA2 = liftArray2' + {-# INLINE liftA2 #-} +#endif -- | Row-major sequential folding over a Delayed array. instance Index ix => Foldable (Array D ix) where @@ -137,176 +166,176 @@ {-# INLINE null #-} length = totalElem . size {-# INLINE length #-} - toList arr = build (\ c n -> foldrFB c n arr) + elem e = A.any (e ==) + {-# INLINE elem #-} + toList arr = build (\c n -> foldrFB c n arr) {-# INLINE toList #-} - instance Index ix => Load D ix e where - size = dSize - {-# INLINE size #-} - getComp = dComp - {-# INLINE getComp #-} - loadArrayM !scheduler !arr = - splitLinearlyWith_ scheduler (elemsCount arr) (unsafeLinearIndex arr) - {-# INLINE loadArrayM #-} + makeArray comp sz = DArray comp sz . PrefIndex + {-# INLINE makeArray #-} + makeArrayLinear comp sz = DArray comp sz . PrefIndexLinear + {-# INLINE makeArrayLinear #-} + iterArrayLinearST_ !scheduler DArray{..} uWrite = + case dPrefIndex of + PrefIndex f -> + iterTargetFullST_ defRowMajor scheduler 0 dSize $ \ !i -> uWrite i . f + PrefIndexLinear f -> + iterTargetFullST_ defRowMajorLinear scheduler 0 dSize $ \ !i _ -> uWrite i (f i) + {-# INLINE iterArrayLinearST_ #-} -instance Index ix => StrideLoad D ix e +instance Index ix => StrideLoad D ix e where + iterArrayLinearWithStrideST_ !scheduler !stride sz DArray{..} uWrite = + case dPrefIndex of + PrefIndex f -> + iterTargetFullWithStrideST_ defRowMajor scheduler 0 sz stride $ \i -> + uWrite i . f + PrefIndexLinear f -> do + iterTargetFullWithStrideST_ defRowMajor scheduler 0 sz stride $ \i -> + uWrite i . f . toLinearIndex dSize + {-# INLINE iterArrayLinearWithStrideST_ #-} instance Index ix => Stream D ix e where toStream = S.steps {-# INLINE toStream #-} - - -instance (Index ix, Num e) => Num (Array D ix e) where - (+) = unsafeLiftArray2 (+) - {-# INLINE (+) #-} - (-) = unsafeLiftArray2 (-) - {-# INLINE (-) #-} - (*) = unsafeLiftArray2 (*) - {-# INLINE (*) #-} - abs = unsafeLiftArray abs - {-# INLINE abs #-} - signum = unsafeLiftArray signum - {-# INLINE signum #-} - fromInteger = singleton . fromInteger - {-# INLINE fromInteger #-} - -instance (Index ix, Fractional e) => Fractional (Array D ix e) where - (/) = unsafeLiftArray2 (/) - {-# INLINE (/) #-} - fromRational = singleton . fromRational - {-# INLINE fromRational #-} - + toStreamIx = S.steps . imap (,) + {-# INLINE toStreamIx #-} -instance (Index ix, Floating e) => Floating (Array D ix e) where - pi = singleton pi - {-# INLINE pi #-} - exp = unsafeLiftArray exp - {-# INLINE exp #-} - log = unsafeLiftArray log - {-# INLINE log #-} - sin = unsafeLiftArray sin - {-# INLINE sin #-} - cos = unsafeLiftArray cos - {-# INLINE cos #-} - asin = unsafeLiftArray asin - {-# INLINE asin #-} - atan = unsafeLiftArray atan - {-# INLINE atan #-} - acos = unsafeLiftArray acos - {-# INLINE acos #-} - sinh = unsafeLiftArray sinh - {-# INLINE sinh #-} - cosh = unsafeLiftArray cosh - {-# INLINE cosh #-} - asinh = unsafeLiftArray asinh - {-# INLINE asinh #-} - atanh = unsafeLiftArray atanh - {-# INLINE atanh #-} - acosh = unsafeLiftArray acosh - {-# INLINE acosh #-} +-- | Map an index aware function over an array +-- +-- @since 0.1.0 +imap + :: forall r ix e a + . (Index ix, Source r e) + => (ix -> e -> a) + -> Array r ix e + -> Array D ix a +imap f !arr = + case unsafePrefIndex arr of + PrefIndex gix -> DArray (getComp arr) sz $ PrefIndex (\ !ix -> f ix (gix ix)) + PrefIndexLinear gi -> + DArray (getComp arr) sz $ PrefIndex (\ !ix -> f ix (gi (toLinearIndex sz ix))) + where + !sz = size arr +{-# INLINE imap #-} +instance Num e => FoldNumeric D e where + unsafeDotProduct = defaultUnsafeDotProduct + {-# INLINE unsafeDotProduct #-} + powerSumArray = defaultPowerSumArray + {-# INLINE powerSumArray #-} + foldArray = defaultFoldArray + {-# INLINE foldArray #-} instance Num e => Numeric D e where - -- plusScalar arr e = unsafeLiftArray (+ e) arr - -- {-# INLINE plusScalar #-} - -- minusScalar arr e = unsafeLiftArray (subtract e) arr - -- {-# INLINE minusScalar #-} - -- multiplyScalar arr e = unsafeLiftArray (* e) arr - -- {-# INLINE multiplyScalar #-} - -- absPointwise = unsafeLiftArray abs - -- {-# INLINE absPointwise #-} - -- additionPointwise = unsafeLiftArray2 (+) - -- {-# INLINE additionPointwise #-} - -- subtractionPointwise = unsafeLiftArray2 (-) - -- {-# INLINE subtractionPointwise #-} - -- multiplicationPointwise = unsafeLiftArray2 (*) - -- {-# INLINE multiplicationPointwise #-} - -- powerPointwise arr pow = unsafeLiftArray (^ pow) arr - -- {-# INLINE powerPointwise #-} - -- powerSumArray arr = sumArray . powerPointwise arr - -- {-# INLINE powerSumArray #-} - -- unsafeDotProduct a1 a2 = sumArray (multiplicationPointwise a1 a2) - -- {-# INLINE unsafeDotProduct #-} - unsafeLiftArray f arr = arr {dIndex = f . dIndex arr} + unsafeLiftArray f arr = arr{dPrefIndex = f <$> dPrefIndex arr} {-# INLINE unsafeLiftArray #-} - unsafeLiftArray2 f a1 a2 = - DArray (dComp a1 <> dComp a2) (SafeSz (liftIndex2 min (unSz (dSize a1)) (unSz (dSize a2)))) $ \i -> - f (dIndex a1 i) (dIndex a2 i) + unsafeLiftArray2 f a1 a2 = zipWithInternal (size a1) f a1 a2 {-# INLINE unsafeLiftArray2 #-} - -instance Floating e => NumericFloat D e where - -- recipPointwise = liftDArray recip - -- {-# INLINE recipPointwise #-} - -- sqrtPointwise = liftDArray sqrt - -- {-# INLINE sqrtPointwise #-} - -- floorPointwise = liftDArray floor - -- {-# INLINE floorPointwise #-} - -- ceilingPointwise = liftDArray ceiling - -- {-# INLINE ceilingPointwise #-} - -- divisionPointwise = liftDArray2 (/) - -- {-# INLINE divisionPointwise #-} - -- divideScalar arr e = liftDArray (/ e) arr - -- {-# INLINE divideScalar #-} - - +instance Floating e => NumericFloat D e -- | /O(1)/ Conversion from a source array to `D` representation. -delay :: Source r ix e => Array r ix e -> Array D ix e -delay arr = DArray (getComp arr) (size arr) (unsafeIndex arr) +delay :: (Index ix, Source r e) => Array r ix e -> Array D ix e +delay arr = + case unsafePrefIndex arr of + PrefIndex gix -> makeArray (getComp arr) (size arr) gix + PrefIndexLinear gi -> makeArrayLinear (getComp arr) (size arr) gi {-# INLINE [1] delay #-} {-# RULES -"delay" [~1] forall (arr :: Array D ix e) . delay arr = arr - #-} +"delay" [~1] forall (arr :: Array D ix e). delay arr = arr + #-} --- TODO: switch to zipWith --- | /O(min (n1, n2))/ - Compute array equality by applying a comparing function to each element. -eq :: (Source r1 ix e1, Source r2 ix e2) => - (e1 -> e2 -> Bool) -> Array r1 ix e1 -> Array r2 ix e2 -> Bool -eq f arr1 arr2 = - (size arr1 == size arr2) && - F.and - (DArray (getComp arr1 <> getComp arr2) (size arr1) $ \ix -> - f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)) -{-# INLINE eq #-} +-- | Compute array equality by applying a comparing function to each +-- element. Empty arrays are always equal, regardless of their size. +-- +-- @since 0.5.7 +eqArrays + :: (Index ix, Source r1 e1, Source r2 e2) + => (e1 -> e2 -> Bool) + -> Array r1 ix e1 + -> Array r2 ix e2 + -> Bool +eqArrays f arr1 arr2 = + let sz1 = size arr1 + sz2 = size arr2 + in ( sz1 == sz2 + && not + ( A.any + not + ( makeArray @D (getComp arr1 <> getComp arr2) (size arr1) $ \ix -> + f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix) + ) + ) + ) + || (isZeroSz sz1 && isZeroSz sz2) +{-# INLINE eqArrays #-} --- | /O(min (n1, n2))/ - Compute array ordering by applying a comparing function to each element. +-- | Compute array ordering by applying a comparing function to each element. -- The exact ordering is unspecified so this is only intended for use in maps and the like where -- you need an ordering but do not care about which one is used. -ord :: (Source r1 ix e1, Source r2 ix e2) => - (e1 -> e2 -> Ordering) -> Array r1 ix e1 -> Array r2 ix e2 -> Ordering -ord f arr1 arr2 = - compare (size arr1) (size arr2) <> - A.fold - (DArray (getComp arr1 <> getComp arr2) (size arr1) $ \ix -> - f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)) -{-# INLINE ord #-} - - --- -- | The usual map. --- liftArray :: Source r ix b => (b -> e) -> Array r ix b -> Array D ix e --- liftArray f !arr = DArray (getComp arr) (size arr) (f . unsafeIndex arr) --- {-# INLINE liftArray #-} +-- +-- @since 0.5.7 +compareArrays + :: (Index ix, Source r1 e1, Source r2 e2) + => (e1 -> e2 -> Ordering) + -> Array r1 ix e1 + -> Array r2 ix e2 + -> Ordering +compareArrays f arr1 arr2 = + compare (size arr1) (size arr2) + <> A.fold + ( makeArray @D (getComp arr1 <> getComp arr2) (size arr1) $ \ix -> + f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix) + ) +{-# INLINE compareArrays #-} --- -- | Similar to `Data.Massiv.Array.zipWith`, except dimensions of both arrays either have to be the --- -- same, or at least one of the two array must be a singleton array, in which case it will behave as --- -- a `Data.Massiv.Array.map`. --- -- --- -- @since 0.1.4 --- liftArray2 --- :: (Source r1 ix a, Source r2 ix b) --- => (a -> b -> e) -> Array r1 ix a -> Array r2 ix b -> Array D ix e --- liftArray2 f !arr1 !arr2 --- | sz1 == oneSz = liftArray (f (unsafeIndex arr1 zeroIndex)) arr2 --- | sz2 == oneSz = liftArray (`f` unsafeIndex arr2 zeroIndex) arr1 --- | sz1 == sz2 = --- DArray (getComp arr1 <> getComp arr2) sz1 (\ !ix -> f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)) --- | otherwise = throw $ SizeMismatchException (size arr1) (size arr2) --- where --- sz1 = size arr1 --- sz2 = size arr2 --- {-# INLINE liftArray2 #-} +-- | Same as `liftArray2M`, but throws an imprecise exception on mismatched +-- sizes. +-- +-- @since 1.0.0 +liftArray2' + :: (HasCallStack, Index ix, Source r1 a, Source r2 b) + => (a -> b -> e) + -> Array r1 ix a + -> Array r2 ix b + -> Array D ix e +liftArray2' f arr1 arr2 = throwEither $ liftArray2M f arr1 arr2 +{-# INLINE liftArray2' #-} +-- | Similar to `Data.Massiv.Array.zipWith`, except dimensions of both arrays +-- have to be the same, otherwise it throws `SizeMismatchException`. +-- +-- @since 1.0.0 +liftArray2M + :: (Index ix, Source r1 a, Source r2 b, MonadThrow m) + => (a -> b -> e) + -> Array r1 ix a + -> Array r2 ix b + -> m (Array D ix e) +liftArray2M f !arr1 !arr2 + | sz1 == sz2 = pure $ zipWithInternal sz1 f arr1 arr2 + | isZeroSz sz1 && isZeroSz sz2 = pure A.empty + | otherwise = throwM $ SizeMismatchException (size arr1) (size arr2) + where + sz1 = size arr1 + sz2 = size arr2 +{-# INLINE liftArray2M #-} +zipWithInternal + :: (Index ix, Source r1 e1, Source r2 e2) + => Sz ix + -> (e1 -> e2 -> e3) + -> Array r1 ix e1 + -> Array r2 ix e2 + -> Array D ix e3 +zipWithInternal sz f arr1 arr2 = + case unsafePrefIndex arr1 of + PrefIndexLinear gi1 + | PrefIndexLinear gi2 <- unsafePrefIndex arr2 -> + makeArrayLinear comp sz (\ !i -> f (gi1 i) (gi2 i)) + _ -> makeArray comp sz (\ !ix -> f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)) + where + comp = getComp arr1 <> getComp arr2 +{-# INLINE zipWithInternal #-} diff --git a/src/Data/Massiv/Array/Delayed/Push.hs b/src/Data/Massiv/Array/Delayed/Push.hs --- a/src/Data/Massiv/Array/Delayed/Push.hs +++ b/src/Data/Massiv/Array/Delayed/Push.hs @@ -1,122 +1,139 @@ -{-# LANGUAGE LambdaCase #-} {-# LANGUAGE BangPatterns #-} {-# LANGUAGE CPP #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE LambdaCase #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE NamedFieldPuns #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UndecidableInstances #-} + -- | -- Module : Data.Massiv.Array.Delayed.Push --- Copyright : (c) Alexey Kuleshevich 2019 +-- Copyright : (c) Alexey Kuleshevich 2019-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Delayed.Push - ( DL(..) - , Array(..) - , toLoadArray - , makeLoadArrayS - , makeLoadArray - , unsafeMakeLoadArray - , fromStrideLoad - , appendOuterM - , concatOuterM - ) where +module Data.Massiv.Array.Delayed.Push ( + DL (..), + Array (..), + Loader, + toLoadArray, + makeLoadArrayS, + makeLoadArray, + unsafeMakeLoadArray, + unsafeMakeLoadArrayAdjusted, + fromStrideLoad, + appendOuterM, + concatOuterM, +) where import Control.Monad -import Data.Massiv.Core.Common -import Prelude hiding (map, zipWith) import Control.Scheduler as S (traverse_) import Data.Foldable as F +import Data.Massiv.Core.Common +import Prelude hiding (map, zipWith) #include "massiv.h" -- | Delayed load representation. Also known as Push array. -data DL = DL deriving Show +data DL = DL deriving (Show) +type Loader e = + forall s + . Scheduler s () + -- ^ Scheduler that will be used for loading + -> Ix1 + -- ^ Start loading at this linear index + -> (Ix1 -> e -> ST s ()) + -- ^ Linear element writing action + -> (Ix1 -> Sz1 -> e -> ST s ()) + -- ^ Linear region setting action + -> ST s () data instance Array DL ix e = DLArray - { dlComp :: !Comp - , dlSize :: !(Sz ix) - , dlDefault :: !(Maybe e) - , dlLoad :: forall m . Monad m - => Scheduler m () - -> Int -- start loading at this linear index - -> (Int -> e -> m ()) -- linear element writing action - -> m () + { dlComp :: !Comp + , dlSize :: !(Sz ix) + , dlLoad :: Loader e } -instance Index ix => Construct DL ix e where - setComp c arr = arr {dlComp = c} +instance Strategy DL where + getComp = dlComp + {-# INLINE getComp #-} + setComp c arr = arr{dlComp = c} {-# INLINE setComp #-} - makeArrayLinear comp sz f = - DLArray comp sz Nothing $ \scheduler startAt dlWrite -> - splitLinearlyWithStartAtM_ scheduler startAt (totalElem sz) (pure . f) dlWrite - {-# INLINE makeArrayLinear #-} + repr = DL -instance Index ix => Resize DL ix where - unsafeResize !sz arr = arr { dlSize = sz } +instance Index ix => Shape DL ix where + maxLinearSize = Just . SafeSz . elemsCount + {-# INLINE maxLinearSize #-} + +instance Size DL where + size = dlSize + {-# INLINE size #-} + unsafeResize !sz !arr = arr{dlSize = sz} {-# INLINE unsafeResize #-} instance Semigroup (Array DL Ix1 e) where (<>) = mappendDL {-# INLINE (<>) #-} +{- FOURMOLU_DISABLE -} instance Monoid (Array DL Ix1 e) where - mempty = - DLArray - {dlComp = mempty, dlSize = Sz zeroIndex, dlDefault = Nothing, dlLoad = \_ _ _ -> pure ()} + mempty = DLArray {dlComp = mempty, dlSize = zeroSz, dlLoad = \_ _ _ _ -> pure ()} {-# INLINE mempty #-} +#if !MIN_VERSION_base(4,11,0) mappend = mappendDL {-# INLINE mappend #-} +#endif mconcat [] = mempty mconcat [x] = x mconcat [x, y] = x <> y mconcat xs = mconcatDL xs {-# INLINE mconcat #-} +{- FOURMOLU_ENABLE -} -mconcatDL :: forall e . [Array DL Ix1 e] -> Array DL Ix1 e +mconcatDL :: forall e. [Array DL Ix1 e] -> Array DL Ix1 e mconcatDL !arrs = - DLArray {dlComp = foldMap getComp arrs, dlSize = SafeSz k, dlDefault = Nothing, dlLoad = load} + DLArray{dlComp = foldMap getComp arrs, dlSize = SafeSz k, dlLoad = load} where !k = F.foldl' (+) 0 (unSz . size <$> arrs) - load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m () - load scheduler startAt dlWrite = - let loadArr !startAtCur DLArray {dlSize = SafeSz kCur, dlDefault, dlLoad} = do + load + :: forall s + . Scheduler s () + -> Ix1 + -> (Ix1 -> e -> ST s ()) + -> (Ix1 -> Sz1 -> e -> ST s ()) + -> ST s () + load scheduler startAt dlWrite dlSet = + let loadArr !startAtCur DLArray{dlSize = SafeSz kCur, dlLoad} = do let !endAtCur = startAtCur + kCur - scheduleWork_ scheduler $ do - S.traverse_ - (\def -> loopM_ startAtCur (< endAtCur) (+ 1) (`dlWrite` def)) - dlDefault - dlLoad scheduler startAtCur dlWrite + scheduleWork_ scheduler $ dlLoad scheduler startAtCur dlWrite dlSet pure endAtCur {-# INLINE loadArr #-} in foldM_ loadArr startAt arrs {-# INLINE load #-} {-# INLINE mconcatDL #-} - -mappendDL :: forall e . Array DL Ix1 e -> Array DL Ix1 e -> Array DL Ix1 e -mappendDL (DLArray c1 sz1 mDef1 load1) (DLArray c2 sz2 mDef2 load2) = - DLArray {dlComp = c1 <> c2, dlSize = SafeSz (k1 + k2), dlDefault = Nothing, dlLoad = load} +mappendDL :: forall e. Array DL Ix1 e -> Array DL Ix1 e -> Array DL Ix1 e +mappendDL (DLArray c1 sz1 load1) (DLArray c2 sz2 load2) = + DLArray{dlComp = c1 <> c2, dlSize = SafeSz (k1 + k2), dlLoad = load} where !k1 = unSz sz1 !k2 = unSz sz2 - load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m () - load scheduler startAt dlWrite = do - scheduleWork_ scheduler $ do - S.traverse_ (\def1 -> loopM_ startAt (< k1) (+ 1) (`dlWrite` def1)) mDef1 - load1 scheduler startAt dlWrite - scheduleWork_ scheduler $ do - let startAt2 = startAt + k1 - S.traverse_ (\def2 -> loopM_ startAt2 (< startAt2 + k2) (+ 1) (`dlWrite` def2)) mDef2 - load2 scheduler startAt2 dlWrite + load + :: forall s + . Scheduler s () + -> Ix1 + -> (Ix1 -> e -> ST s ()) + -> (Ix1 -> Sz1 -> e -> ST s ()) + -> ST s () + load scheduler !startAt dlWrite dlSet = do + scheduleWork_ scheduler $ load1 scheduler startAt dlWrite dlSet + scheduleWork_ scheduler $ load2 scheduler (startAt + k1) dlWrite dlSet {-# INLINE load #-} {-# INLINE mappendDL #-} @@ -124,67 +141,56 @@ -- agree, otherwise `SizeMismatchException`. -- -- @since 0.4.4 -appendOuterM :: - forall ix e m. (Index ix, MonadThrow m) +appendOuterM + :: forall ix e m + . (Index ix, MonadThrow m) => Array DL ix e -> Array DL ix e -> m (Array DL ix e) -appendOuterM (DLArray c1 sz1 mDef1 load1) (DLArray c2 sz2 mDef2 load2) = do +appendOuterM (DLArray c1 sz1 load1) (DLArray c2 sz2 load2) = do let (!i1, !szl1) = unconsSz sz1 (!i2, !szl2) = unconsSz sz2 unless (szl1 == szl2) $ throwM $ SizeMismatchException sz1 sz2 pure $ - DLArray {dlComp = c1 <> c2, dlSize = consSz (i1 + i2) szl1, dlDefault = Nothing, dlLoad = load} + DLArray{dlComp = c1 <> c2, dlSize = consSz (liftSz2 (+) i1 i2) szl1, dlLoad = load} where - !k1 = totalElem sz1 - !k2 = totalElem sz2 - load :: Monad n => Scheduler n () -> Int -> (Int -> e -> n ()) -> n () - load scheduler !startAt dlWrite = do - scheduleWork_ scheduler $ do - S.traverse_ (\def1 -> loopM_ startAt (< k1) (+ 1) (`dlWrite` def1)) mDef1 - load1 scheduler startAt dlWrite - scheduleWork_ scheduler $ do - let !startAt2 = startAt + k1 - S.traverse_ (\def2 -> loopM_ startAt2 (< startAt2 + k2) (+ 1) (`dlWrite` def2)) mDef2 - load2 scheduler startAt2 dlWrite + load :: Loader e + load scheduler !startAt dlWrite dlSet = do + scheduleWork_ scheduler $ load1 scheduler startAt dlWrite dlSet + scheduleWork_ scheduler $ load2 scheduler (startAt + totalElem sz1) dlWrite dlSet {-# INLINE load #-} {-# INLINE appendOuterM #-} --- | Concat arrays together along the most most dimension. Inner dimensions must agree +-- | Concat arrays together along the outer most dimension. Inner dimensions must agree -- for all arrays in the list, otherwise `SizeMismatchException`. -- -- @since 0.4.4 -concatOuterM :: - forall ix e m. (Index ix, MonadThrow m) +concatOuterM + :: forall ix e m + . (Index ix, MonadThrow m) => [Array DL ix e] -> m (Array DL ix e) concatOuterM = \case [] -> pure empty - (x:xs) -> F.foldlM appendOuterM x xs + (x : xs) -> F.foldlM appendOuterM x xs {-# INLINE concatOuterM #-} - -- | Describe how an array should be loaded into memory sequentially. For parallelizable -- version see `makeLoadArray`. -- -- @since 0.3.1 -makeLoadArrayS :: - Index ix => - Sz ix +makeLoadArrayS + :: forall ix e + . Index ix + => Sz ix -- ^ Size of the resulting array -> e -- ^ Default value to use for all cells that might have been ommitted by the writing function -> (forall m. Monad m => (ix -> e -> m Bool) -> m ()) -- ^ Writing function that described which elements to write into the target array. -> Array DL ix e -makeLoadArrayS sz defVal writer = - DLArray Seq sz (Just defVal) $ \_scheduler !startAt uWrite -> - let safeWrite !ix !e - | isSafeIndex sz ix = uWrite (startAt + toLinearIndex sz ix) e >> pure True - | otherwise = pure False - {-# INLINE safeWrite #-} - in writer safeWrite +makeLoadArrayS sz defVal writer = makeLoadArray Seq sz defVal (const writer) {-# INLINE makeLoadArrayS #-} -- | Specify how an array should be loaded into memory. Unlike `makeLoadArrayS`, loading @@ -192,8 +198,9 @@ -- of this function see `unsafeMakeLoadArray`. -- -- @since 0.4.0 -makeLoadArray :: - Index ix +makeLoadArray + :: forall ix e + . Index ix => Comp -- ^ Computation strategy to use. Directly affects the scheduler that gets created for -- the loading function. @@ -201,19 +208,28 @@ -- ^ Size of the resulting array -> e -- ^ Default value to use for all cells that might have been ommitted by the writing function - -> (forall m. Monad m => - Scheduler m () -> (ix -> e -> m Bool) -> m ()) + -> (forall s. Scheduler s () -> (ix -> e -> ST s Bool) -> ST s ()) -- ^ Writing function that described which elements to write into the target array. It -- accepts a scheduler, that can be used for parallelization, as well as a safe element -- writing function. -> Array DL ix e -makeLoadArray comp sz defVal writer = - DLArray comp sz (Just defVal) $ \scheduler !startAt uWrite -> - let safeWrite !ix !e - | isSafeIndex sz ix = uWrite (startAt + toLinearIndex sz ix) e >> pure True - | otherwise = pure False - {-# INLINE safeWrite #-} - in writer scheduler safeWrite +makeLoadArray comp sz defVal writer = DLArray comp sz load + where + load + :: forall s + . Scheduler s () + -> Ix1 + -> (Ix1 -> e -> ST s ()) + -> (Ix1 -> Sz1 -> e -> ST s ()) + -> ST s () + load scheduler !startAt uWrite uSet = do + uSet startAt (toLinearSz sz) defVal + let safeWrite !ix !e + | isSafeIndex sz ix = True <$ uWrite (startAt + toLinearIndex sz ix) e + | otherwise = pure False + {-# INLINE safeWrite #-} + writer scheduler safeWrite + {-# INLINE load #-} {-# INLINE makeLoadArray #-} -- | Specify how an array can be loaded/computed through creation of a `DL` array. Unlike @@ -222,8 +238,10 @@ -- function does not perform any bounds checking. -- -- @since 0.3.1 -unsafeMakeLoadArray :: - Comp +unsafeMakeLoadArray + :: forall ix e + . Index ix + => Comp -- ^ Computation strategy to use. Directly affects the scheduler that gets created for -- the loading function. -> Sz ix @@ -231,8 +249,7 @@ -> Maybe e -- ^ An element to use for initialization of the mutable array that will be created in -- the future - -> (forall m. Monad m => - Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()) + -> (forall s. Scheduler s () -> Ix1 -> (Ix1 -> e -> ST s ()) -> ST s ()) -- ^ This function accepts: -- -- * A scheduler that can be used for parallelization of loading @@ -242,45 +259,125 @@ -- -- * Linear element writing function -> Array DL ix e -unsafeMakeLoadArray = DLArray +unsafeMakeLoadArray comp sz mDefVal writer = DLArray comp sz load + where + load :: Loader e + load scheduler startAt uWrite uSet = do + S.traverse_ (uSet startAt (toLinearSz sz)) mDefVal + writer scheduler startAt uWrite + {-# INLINE load #-} {-# INLINE unsafeMakeLoadArray #-} +-- | Same as `unsafeMakeLoadArray`, except will ensure that starting index is correctly +-- adjusted. Which means the writing function gets one less argument. +-- +-- @since 0.5.2 +unsafeMakeLoadArrayAdjusted + :: forall ix e + . Index ix + => Comp + -> Sz ix + -> Maybe e + -> (forall s. Scheduler s () -> (Ix1 -> e -> ST s ()) -> ST s ()) + -> Array DL ix e +unsafeMakeLoadArrayAdjusted comp sz mDefVal writer = DLArray comp sz load + where + load + :: forall s + . Scheduler s () + -> Ix1 + -> (Ix1 -> e -> ST s ()) + -> (Ix1 -> Sz1 -> e -> ST s ()) + -> ST s () + load scheduler !startAt uWrite dlSet = do + S.traverse_ (dlSet startAt (toLinearSz sz)) mDefVal + writer scheduler (\i -> uWrite (startAt + i)) + {-# INLINE load #-} +{-# INLINE unsafeMakeLoadArrayAdjusted #-} + -- | Convert any `Load`able array into `DL` representation. -- -- @since 0.3.0 -toLoadArray :: Load r ix e => Array r ix e -> Array DL ix e -toLoadArray arr = - DLArray (getComp arr) (size arr) Nothing $ \scheduler startAt dlWrite -> - loadArrayM scheduler arr (dlWrite . (+ startAt)) -{-# INLINE toLoadArray #-} +toLoadArray + :: forall r ix e + . (Size r, Load r ix e) + => Array r ix e + -> Array DL ix e +toLoadArray arr = DLArray (getComp arr) sz load + where + !sz = size arr + load + :: forall s + . Scheduler s () + -> Ix1 + -> (Ix1 -> e -> ST s ()) + -> (Ix1 -> Sz1 -> e -> ST s ()) + -> ST s () + load scheduler !startAt dlWrite dlSet = + iterArrayLinearWithSetST_ + scheduler + arr + (dlWrite . (+ startAt)) + (\offset -> dlSet (offset + startAt)) + {-# INLINE load #-} +{-# INLINE [1] toLoadArray #-} +{-# RULES "toLoadArray/id" toLoadArray = id #-} + -- | Convert an array that can be loaded with stride into `DL` representation. -- -- @since 0.3.0 fromStrideLoad - :: StrideLoad r ix e => Stride ix -> Array r ix e -> Array DL ix e + :: forall r ix e + . StrideLoad r ix e + => Stride ix + -> Array r ix e + -> Array DL ix e fromStrideLoad stride arr = - DLArray (getComp arr) newsz Nothing $ \scheduler startAt dlWrite -> - loadArrayWithStrideM scheduler stride newsz arr (\ !i -> dlWrite (i + startAt)) + DLArray (getComp arr) newsz load where - newsz = strideSize stride (size arr) + !newsz = strideSize stride (outerSize arr) + load :: Loader e + load scheduler !startAt dlWrite _ = + iterArrayLinearWithStrideST_ scheduler stride newsz arr (\ !i -> dlWrite (i + startAt)) + {-# INLINE load #-} {-# INLINE fromStrideLoad #-} instance Index ix => Load DL ix e where - size = dlSize - {-# INLINE size #-} - getComp = dlComp - {-# INLINE getComp #-} - loadArrayM scheduler DLArray {dlLoad} = dlLoad scheduler 0 - {-# INLINE loadArrayM #-} - defaultElement = dlDefault - {-# INLINE defaultElement #-} + makeArrayLinear comp sz f = DLArray comp sz load + where + load :: Loader e + load scheduler startAt dlWrite _ = + splitLinearlyWithStartAtM_ scheduler startAt (totalElem sz) (pure . f) dlWrite + {-# INLINE load #-} + {-# INLINE makeArrayLinear #-} + replicate comp !sz !e = makeLoadArray comp sz e $ \_ _ -> pure () + {-# INLINE replicate #-} + iterArrayLinearWithSetST_ scheduler DLArray{dlLoad} = dlLoad scheduler 0 + {-# INLINE iterArrayLinearWithSetST_ #-} -instance Functor (Array DL ix) where - fmap f arr = - arr - { dlLoad = - \scheduler startAt uWrite -> dlLoad arr scheduler startAt (\ !i e -> uWrite i (f e)) - , dlDefault = f <$> dlDefault arr - } +instance Index ix => Functor (Array DL ix) where + fmap f arr = arr{dlLoad = loadFunctor arr f} {-# INLINE fmap #-} + (<$) = overwriteFunctor + {-# INLINE (<$) #-} + +overwriteFunctor :: forall ix a b. Index ix => a -> Array DL ix b -> Array DL ix a +overwriteFunctor e arr = arr{dlLoad = load} + where + load :: Loader a + load _ !startAt _ dlSet = dlSet startAt (linearSize arr) e + {-# INLINE load #-} +{-# INLINE overwriteFunctor #-} + +loadFunctor + :: Array DL ix a + -> (a -> b) + -> Scheduler s () + -> Ix1 + -> (Ix1 -> b -> ST s ()) + -> (Ix1 -> Sz1 -> b -> ST s ()) + -> ST s () +loadFunctor arr f scheduler startAt uWrite uSet = + dlLoad arr scheduler startAt (\ !i e -> uWrite i (f e)) (\o sz e -> uSet o sz (f e)) +{-# INLINE loadFunctor #-} diff --git a/src/Data/Massiv/Array/Delayed/Stream.hs b/src/Data/Massiv/Array/Delayed/Stream.hs --- a/src/Data/Massiv/Array/Delayed/Stream.hs +++ b/src/Data/Massiv/Array/Delayed/Stream.hs @@ -3,75 +3,93 @@ {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE TypeFamilies #-} + -- | -- Module : Data.Massiv.Array.Delayed.Stream --- Copyright : (c) Alexey Kuleshevich 2019 +-- Copyright : (c) Alexey Kuleshevich 2019-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Delayed.Stream - ( DS(..) - , Array (..) - , toStreamArray - , toSteps - , fromSteps - , takeS - , dropS - , filterS - , filterM - , mapMaybeS - , mapMaybeM - , catMaybesS - , traverseS - , unfoldr - , unfoldrN - ) where +module Data.Massiv.Array.Delayed.Stream ( + DS (..), + Array (..), + toStreamArray, + toStreamM, + toStreamIxM, + toSteps, + fromSteps, + fromStepsM, +) where import Control.Applicative -import Control.Monad (void) +import Control.Monad.ST import Data.Coerce +import Data.Foldable import Data.Massiv.Array.Delayed.Pull -import qualified Data.Massiv.Array.Manifest.Vector.Stream as S import Data.Massiv.Core.Common +import qualified Data.Massiv.Vector.Stream as S import GHC.Exts -import Prelude hiding (take, drop) -import Data.Vector.Fusion.Bundle.Size (upperBound) +import Prelude hiding (drop, take) --- | Delayed array that will be loaded in an interleaved fashion during parallel --- computation. +-- | Delayed stream array that represents a sequence of values that can be loaded +-- sequentially. Important distinction from other arrays is that its size might no be +-- known until it is computed. data DS = DS newtype instance Array DS Ix1 e = DSArray { dsArray :: S.Steps S.Id e } --- | /O(1)/ - Convert delayed stream arrray into `Steps`. +-- | /O(1)/ - Convert delayed stream array into `Steps`. -- -- @since 0.4.1 -toSteps :: Array DS Ix1 e -> Steps Id e +toSteps :: Vector DS e -> Steps Id e toSteps = coerce {-# INLINE toSteps #-} --- | /O(1)/ - Convert `Steps` into delayed stream arrray +-- | /O(1)/ - Convert `Steps` into delayed stream array -- -- @since 0.4.1 -fromSteps :: Steps Id e -> Array DS Ix1 e +fromSteps :: Steps Id e -> Vector DS e fromSteps = coerce {-# INLINE fromSteps #-} +-- | /O(1)/ - Convert monadic `Steps` into delayed stream array +-- +-- @since 0.5.0 +fromStepsM :: Monad m => Steps m e -> m (Vector DS e) +fromStepsM = fmap DSArray . S.transSteps +{-# INLINE fromStepsM #-} -instance Functor (Array DS Ix1) where +instance Shape DS Ix1 where + linearSizeHint = stepsSize . dsArray + {-# INLINE linearSizeHint #-} - fmap f = coerce . fmap f . dsArray + linearSize = SafeSz . unId . S.length . dsArray + {-# INLINE linearSize #-} + + outerSize = linearSize + {-# INLINE outerSize #-} + + isNull = S.unId . S.null . coerce + {-# INLINE isNull #-} + +-- | For now only `Seq` strategy. +instance Strategy DS where + getComp _ = Seq + setComp _ = id + repr = DS + +instance Functor (Array DS Ix1) where + fmap f = coerce . S.map f . dsArray {-# INLINE fmap #-} + (<$) e = coerce . (e <$) . dsArray + {-# INLINE (<$) #-} instance Applicative (Array DS Ix1) where - pure = fromSteps . S.singleton {-# INLINE pure #-} - (<*>) a1 a2 = fromSteps (S.zipWith ($) (coerce a1) (coerce a2)) {-# INLINE (<*>) #-} @@ -81,101 +99,102 @@ #endif instance Monad (Array DS Ix1) where - - return = fromSteps . S.singleton - {-# INLINE return #-} - (>>=) arr f = coerce (S.concatMap (coerce . f) (dsArray arr)) {-# INLINE (>>=) #-} - instance Foldable (Array DS Ix1) where - - foldr f acc = S.foldr f acc . toSteps + foldr f acc = S.unId . S.foldrLazy f acc . toSteps {-# INLINE foldr #-} - - length = S.length . coerce + foldl f acc = S.unId . S.foldlLazy f acc . toSteps + {-# INLINE foldl #-} + foldl' f acc = S.unId . S.foldl f acc . toSteps + {-# INLINE foldl' #-} + foldr1 f = S.unId . S.foldr1Lazy f . toSteps + {-# INLINE foldr1 #-} + foldl1 f = S.unId . S.foldl1Lazy f . toSteps + {-# INLINE foldl1 #-} + toList = S.toList . coerce + {-# INLINE toList #-} + length = S.unId . S.length . coerce {-# INLINE length #-} - - -- TODO: add more - + null = S.unId . S.null . coerce + {-# INLINE null #-} + sum = S.unId . S.foldl (+) 0 . toSteps + {-# INLINE sum #-} + product = S.unId . S.foldl (*) 1 . toSteps + {-# INLINE product #-} + maximum = S.unId . S.foldl1 max . toSteps + {-# INLINE maximum #-} + minimum = S.unId . S.foldl1 min . toSteps + {-# INLINE minimum #-} instance Semigroup (Array DS Ix1 e) where - (<>) a1 a2 = fromSteps (coerce a1 `S.append` coerce a2) {-# INLINE (<>) #-} - instance Monoid (Array DS Ix1 e) where - mempty = DSArray S.empty {-# INLINE mempty #-} - +#if !MIN_VERSION_base(4,11,0) mappend = (<>) {-# INLINE mappend #-} +#endif instance IsList (Array DS Ix1 e) where type Item (Array DS Ix1 e) = e - - fromList = fromSteps . S.fromList + fromList = fromSteps . fromList {-# INLINE fromList #-} - - fromListN n = fromSteps . S.fromListN n + fromListN n = fromSteps . fromListN n {-# INLINE fromListN #-} - toList = S.toList . coerce {-# INLINE toList #-} - instance S.Stream DS Ix1 e where toStream = coerce {-# INLINE toStream #-} - + toStreamIx = S.indexed . coerce + {-# INLINE toStreamIx #-} -- | Flatten an array into a stream of values. -- -- @since 0.4.1 -toStreamArray :: Source r ix e => Array r ix e -> Array DS Ix1 e +toStreamArray :: (Index ix, Source r e) => Array r ix e -> Vector DS e toStreamArray = DSArray . S.steps -{-# INLINE toStreamArray #-} - -instance Construct DS Ix1 e where - setComp _ arr = arr - {-# INLINE setComp #-} +{-# INLINE [1] toStreamArray #-} - makeArrayLinear _ (Sz k) = fromSteps . S.generate k - {-# INLINE makeArrayLinear #-} +{-# RULES "toStreamArray/id" toStreamArray = id #-} +-- | /O(1)/ - Convert an array into monadic `Steps` +-- +-- @since 0.5.0 +toStreamM :: (Stream r ix e, Monad m) => Array r ix e -> Steps m e +toStreamM = S.transStepsId . toStream +{-# INLINE toStreamM #-} -instance Extract DS Ix1 e where - unsafeExtract sIx newSz = fromSteps . S.slice sIx (unSz newSz) . dsArray - {-# INLINE unsafeExtract #-} +-- | /O(1)/ - Convert an array into monadic `Steps` +-- +-- @since 0.5.0 +toStreamIxM :: (Stream r ix e, Monad m) => Array r ix e -> Steps m (ix, e) +toStreamIxM = S.transStepsId . toStreamIx +{-# INLINE toStreamIxM #-} -- | /O(n)/ - `size` implementation. instance Load DS Ix1 e where - size = coerce . S.length . coerce - {-# INLINE size #-} - - maxSize = coerce . upperBound . stepsSize . dsArray - {-# INLINE maxSize #-} - - getComp _ = Seq - {-# INLINE getComp #-} + makeArrayLinear _ k = fromSteps . S.generate k + {-# INLINE makeArrayLinear #-} + replicate _ k = fromSteps . S.replicate k + {-# INLINE replicate #-} - loadArrayM _scheduler arr uWrite = - case stepsSize (dsArray arr) of - S.Exact _ -> - void $ S.foldlM (\i e -> uWrite i e >> pure (i + 1)) 0 (S.transStepsId (coerce arr)) - _ -> error "Loading Stream array is not supported with loadArrayM" - {-# INLINE loadArrayM #-} + iterArrayLinearST_ _scheduler arr uWrite = + S.mapM_ (uncurry uWrite) $ S.indexed $ S.transStepsId (coerce arr) + {-# INLINE iterArrayLinearST_ #-} - unsafeLoadIntoS marr (DSArray sts) = + unsafeLoadIntoST marr (DSArray sts) = S.unstreamIntoM marr (stepsSize sts) (stepsStream sts) - {-# INLINE unsafeLoadIntoS #-} - - unsafeLoadInto marr arr = liftIO $ unsafeLoadIntoS marr arr - {-# INLINE unsafeLoadInto #-} + {-# INLINE unsafeLoadIntoST #-} + unsafeLoadIntoIO marr arr = stToIO $ unsafeLoadIntoST marr arr + {-# INLINE unsafeLoadIntoIO #-} -- cons :: e -> Array DS Ix1 e -> Array DS Ix1 e -- cons e = coerce . S.cons e . dsArray @@ -189,160 +208,13 @@ -- snoc (DSArray sts) e = DSArray (S.snoc sts e) -- {-# INLINE snoc #-} - -- TODO: skip the stride while loading -- instance StrideLoad DS Ix1 e where --- loadArrayWithStrideM scheduler stride resultSize arr uWrite = +-- iterArrayLinearWithStrideST_ scheduler stride resultSize arr uWrite = -- let strideIx = unStride stride -- DIArray (DArray _ _ f) = arr -- in loopM_ 0 (< numWorkers scheduler) (+ 1) $ \ !start -> -- scheduleWork scheduler $ -- iterLinearM_ resultSize start (totalElem resultSize) (numWorkers scheduler) (<) $ -- \ !i ix -> uWrite i (f (liftIndex2 (*) strideIx ix)) --- {-# INLINE loadArrayWithStrideM #-} - - --- | Right unfolding function. Useful when we do not have any idea ahead of time on how --- many elements the vector will have. --- --- ====__Example__ --- --- >>> import Data.Massiv.Array as A --- >>> unfoldr (\i -> if i < 9 then Just (i*i, i + 1) else Nothing) (0 :: Int) --- Array DS Seq (Sz1 9) --- [ 0, 1, 4, 9, 16, 25, 36, 49, 64 ] --- >>> unfoldr (\i -> if sqrt i < 3 then Just (i * i, i + 1) else Nothing) (0 :: Double) --- Array DS Seq (Sz1 9) --- [ 0.0, 1.0, 4.0, 9.0, 16.0, 25.0, 36.0, 49.0, 64.0 ] --- --- @since 0.4.1 -unfoldr :: (s -> Maybe (e, s)) -> s -> Array DS Ix1 e -unfoldr f = DSArray . S.unfoldr f -{-# INLINE unfoldr #-} - - --- | Right unfolding function with limited number of elements. --- --- ==== __Example__ --- --- >>> import Data.Massiv.Array as A --- >>> unfoldrN 9 (\i -> Just (i*i, i + 1)) (0 :: Int) --- Array DS Seq (Sz1 9) --- [ 0, 1, 4, 9, 16, 25, 36, 49, 64 ] --- --- @since 0.4.1 -unfoldrN :: - Sz1 - -- ^ Maximum number of elements that the vector can have - -> (s -> Maybe (e, s)) - -- ^ Unfolding function. Stops when `Nothing` is reaturned or maximum number of elements - -- is reached. - -> s -- ^ Inititial element. - -> Array DS Ix1 e -unfoldrN n f = DSArray . S.unfoldrN n f -{-# INLINE unfoldrN #-} - --- | Sequentially filter out elements from the array according to the supplied predicate. --- --- ==== __Example__ --- --- >>> import Data.Massiv.Array as A --- >>> arr = makeArrayR D Seq (Sz2 3 4) fromIx2 --- >>> arr --- Array D Seq (Sz (3 :. 4)) --- [ [ (0,0), (0,1), (0,2), (0,3) ] --- , [ (1,0), (1,1), (1,2), (1,3) ] --- , [ (2,0), (2,1), (2,2), (2,3) ] --- ] --- >>> filterS (even . fst) arr --- Array DS Seq (Sz1 8) --- [ (0,0), (0,1), (0,2), (0,3), (2,0), (2,1), (2,2), (2,3) ] --- --- @since 0.4.1 -filterS :: S.Stream r ix e => (e -> Bool) -> Array r ix e -> Array DS Ix1 e -filterS f = DSArray . S.filter f . S.toStream -{-# INLINE filterS #-} - --- | Sequentially filter out elements from the array according to the supplied applicative predicate. --- --- ==== __Example__ --- --- >>> import Data.Massiv.Array as A --- >>> arr = makeArrayR D Seq (Sz2 3 4) fromIx2 --- >>> arr --- Array D Seq (Sz (3 :. 4)) --- [ [ (0,0), (0,1), (0,2), (0,3) ] --- , [ (1,0), (1,1), (1,2), (1,3) ] --- , [ (2,0), (2,1), (2,2), (2,3) ] --- ] --- >>> filterM (Just . odd . fst) arr --- Just (Array DS Seq (Sz1 4) --- [ (1,0), (1,1), (1,2), (1,3) ] --- ) --- >>> filterM (\ix@(_, j) -> print ix >> return (even j)) arr --- (0,0) --- (0,1) --- (0,2) --- (0,3) --- (1,0) --- (1,1) --- (1,2) --- (1,3) --- (2,0) --- (2,1) --- (2,2) --- (2,3) --- Array DS Seq (Sz1 6) --- [ (0,0), (0,2), (1,0), (1,2), (2,0), (2,2) ] --- --- @since 0.4.1 -filterM :: (S.Stream r ix e, Applicative f) => (e -> f Bool) -> Array r ix e -> f (Array DS Ix1 e) -filterM f arr = DSArray <$> S.filterA f (S.toStream arr) -{-# INLINE filterM #-} - - --- | Apply a function to each element of the array, while discarding `Nothing` and --- keeping the `Maybe` result. --- --- @since 0.4.1 -mapMaybeS :: S.Stream r ix a => (a -> Maybe b) -> Array r ix a -> Array DS Ix1 b -mapMaybeS f = DSArray . S.mapMaybe f . S.toStream -{-# INLINE mapMaybeS #-} - --- | Keep all `Maybe`s and discard the `Nothing`s. --- --- @since 0.4.4 -catMaybesS :: S.Stream r ix (Maybe a) => Array r ix (Maybe a) -> Array DS Ix1 a -catMaybesS = mapMaybeS id -{-# INLINE catMaybesS #-} - - --- | Similar to `mapMaybeS`, but with the use of `Applicative` --- --- @since 0.4.1 -mapMaybeM :: - (S.Stream r ix a, Applicative f) => (a -> f (Maybe b)) -> Array r ix a -> f (Array DS Ix1 b) -mapMaybeM f arr = DSArray <$> S.mapMaybeA f (S.toStream arr) -{-# INLINE mapMaybeM #-} - --- | Extract first @n@ elements from the stream vector --- --- @since 0.4.1 -takeS :: Stream r ix e => Sz1 -> Array r ix e -> Array DS Ix1 e -takeS n = fromSteps . S.take (unSz n) . S.toStream -{-# INLINE takeS #-} - --- | Keep all but first @n@ elements from the stream vector. --- --- @since 0.4.1 -dropS :: Stream r ix e => Sz1 -> Array r ix e -> Array DS Ix1 e -dropS n = fromSteps . S.drop (unSz n) . S.toStream -{-# INLINE dropS #-} - - --- | Traverse a stream with an applicative action. --- --- @since 0.4.5 -traverseS :: (S.Stream r ix a, Applicative f) => (a -> f b) -> Array r ix a -> f (Array DS Ix1 b) -traverseS f = fmap fromSteps . S.traverse f . S.toStream -{-# INLINE traverseS #-} +-- {-# INLINE iterArrayLinearWithStrideST_ #-} diff --git a/src/Data/Massiv/Array/Delayed/Windowed.hs b/src/Data/Massiv/Array/Delayed/Windowed.hs --- a/src/Data/Massiv/Array/Delayed/Windowed.hs +++ b/src/Data/Massiv/Array/Delayed/Windowed.hs @@ -8,82 +8,81 @@ {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE TypeOperators #-} {-# LANGUAGE UndecidableInstances #-} + -- | -- Module : Data.Massiv.Array.Delayed.Windowed --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2025 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Delayed.Windowed - ( DW(..) - , Array(..) - , Window(..) - , insertWindow - , getWindow - , dropWindow - , makeWindowedArray - ) where +module Data.Massiv.Array.Delayed.Windowed ( + DW (..), + Array (..), + Window (..), + insertWindow, + getWindow, + dropWindow, + makeWindowedArray, +) where -import Control.Exception (Exception(..)) import Control.Monad (when) +import Control.Scheduler (trivialScheduler_) import Data.Massiv.Array.Delayed.Pull import Data.Massiv.Array.Manifest.Boxed import Data.Massiv.Array.Manifest.Internal import Data.Massiv.Core import Data.Massiv.Core.Common -import Data.Massiv.Core.Index.Internal (Sz(..)) -import Data.Massiv.Core.List (L, showArrayList, showsArrayPrec) +import Data.Massiv.Core.List (showArrayList, showsArrayPrec) import Data.Maybe (fromMaybe) +import GHC.Base (modInt) import GHC.TypeLits -- | Delayed Windowed Array representation. data DW = DW -data Window ix e = Window { windowStart :: !ix - -- ^ Index of where window will start at. - , windowSize :: !(Sz ix) - -- ^ Size of the window - , windowIndex :: ix -> e - -- ^ Indexing function for the window - , windowUnrollIx2 :: !(Maybe Int) - -- ^ Setting this value during stencil application improves cache - -- utilization by unrolling the loop for Ix2 and higher dimensions. - -- Has no affect on arrays with one dimension. - } +data Window ix e = Window + { windowStart :: !ix + -- ^ Index of where window will start at. + , windowSize :: !(Sz ix) + -- ^ Size of the window + , windowIndex :: ix -> e + -- ^ Indexing function for the window + , windowUnrollIx2 :: !(Maybe Int) + -- ^ Setting this value during stencil application improves cache + -- utilization by unrolling the loop for Ix2 and higher dimensions. + -- Has no affect on arrays with one dimension. + } instance Functor (Window ix) where - fmap f arr@Window{windowIndex} = arr { windowIndex = f . windowIndex } + fmap f arr@Window{windowIndex} = arr{windowIndex = f . windowIndex} -data instance Array DW ix e = DWArray { dwArray :: !(Array D ix e) - , dwWindow :: !(Maybe (Window ix e)) - } +data instance Array DW ix e = DWArray + { dwComp :: !Comp + , dwSize :: !(Sz ix) + , dwIndex :: ix -> e + , dwWindow :: !(Maybe (Window ix e)) + } instance (Ragged L ix e, Load DW ix e, Show e) => Show (Array DW ix e) where showsPrec = showsArrayPrec (computeAs B) showList = showArrayList - -instance Index ix => Construct DW ix e where - - setComp c arr = arr { dwArray = (dwArray arr) { dComp = c } } +instance Strategy DW where + setComp c arr = arr{dwComp = c} {-# INLINE setComp #-} - - makeArray c sz f = DWArray (makeArray c sz f) Nothing - {-# INLINE makeArray #-} - - + getComp = dwComp + {-# INLINE getComp #-} + repr = DW instance Functor (Array DW ix) where - fmap f arr@DWArray{dwArray, dwWindow} = + fmap f arr@DWArray{dwIndex, dwWindow} = arr - { dwArray = fmap f dwArray - , dwWindow = fmap f <$> dwWindow - } + { dwIndex = f . dwIndex + , dwWindow = fmap f <$> dwWindow + } {-# INLINE fmap #-} - -- -- -- @since 0.3.0 @@ -124,15 +123,19 @@ -- -- @since 0.1.3 makeWindowedArray - :: Source r ix e - => Array r ix e -- ^ Source array that will have a window inserted into it - -> ix -- ^ Start index for the window - -> Sz ix -- ^ Size of the window - -> (ix -> e) -- ^ Indexing function foto use inside window + :: (Index ix, Source r e) + => Array r ix e + -- ^ Source array that will have a window inserted into it + -> ix + -- ^ Start index for the window + -> Sz ix + -- ^ Size of the window + -> (ix -> e) + -- ^ Indexing function foto use inside window -> Array DW ix e makeWindowedArray !arr wStart wSize wIndex = insertWindow (delay arr) $ - Window {windowStart = wStart, windowSize = wSize, windowIndex = wIndex, windowUnrollIx2 = Nothing} + Window{windowStart = wStart, windowSize = wSize, windowIndex = wIndex, windowUnrollIx2 = Nothing} {-# INLINE makeWindowedArray #-} -- | Inserts a `Window` into a delayed array while scaling the window down if it doesn't fit inside @@ -140,33 +143,37 @@ -- -- @since 0.3.0 insertWindow - :: Source D ix e - => Array D ix e -- ^ Source array that will have a window inserted into it - -> Window ix e -- ^ Window to place inside the delayed array + :: Index ix + => Array D ix e + -- ^ Source array that will have a window inserted into it + -> Window ix e + -- ^ Window to place inside the delayed array -> Array DW ix e insertWindow !arr !window = DWArray - { dwArray = delay arr + { dwComp = getComp arr + , dwSize = arrSize + , dwIndex = unsafeIndex arr , dwWindow = Just $! - Window - { windowStart = wStart' - , windowSize = Sz (liftIndex2 min wSize (liftIndex2 (-) sz wStart')) - , windowIndex = wIndex - , windowUnrollIx2 = wUnrollIx2 - } + Window + { windowStart = wStart' + , windowSize = Sz (liftIndex2 min wSize (liftIndex2 (-) sz wStart')) + , windowIndex = wIndex + , windowUnrollIx2 = wUnrollIx2 + } } where wStart' = unSz (Sz (liftIndex2 min wStart (liftIndex (subtract 1) sz))) - Sz sz = size arr - Window { windowStart = wStart - , windowSize = Sz wSize - , windowIndex = wIndex - , windowUnrollIx2 = wUnrollIx2 - } = window + arrSize@(Sz sz) = size arr + Window + { windowStart = wStart + , windowSize = Sz wSize + , windowIndex = wIndex + , windowUnrollIx2 = wUnrollIx2 + } = window {-# INLINE insertWindow #-} - -- | Get the `Window` from a windowed array. -- -- @since 0.2.1 @@ -178,10 +185,14 @@ -- -- @since 0.3.0 dropWindow :: Array DW ix e -> Array D ix e -dropWindow = dwArray +dropWindow DWArray{..} = + DArray + { dComp = dwComp + , dSize = dwSize + , dPrefIndex = PrefIndex dwIndex + } {-# INLINE dropWindow #-} - zeroWindow :: Index ix => Window ix e zeroWindow = Window zeroIndex zeroSz windowError Nothing {-# INLINE zeroWindow #-} @@ -189,111 +200,112 @@ data EmptyWindowException = EmptyWindowException deriving (Eq, Show) instance Exception EmptyWindowException where - displayException _ = "Index of zero size Window" -windowError :: a -windowError = throwImpossible EmptyWindowException +windowError :: ix -> a +windowError _ = throwImpossible EmptyWindowException {-# NOINLINE windowError #-} - -loadWithIx1 :: - (Monad m) +loadWithIx1 + :: Monad m => (m () -> m ()) -> Array DW Ix1 e -> (Ix1 -> e -> m a) -> m (Ix1 -> Ix1 -> m (), Ix1, Ix1) -loadWithIx1 with (DWArray (DArray _ sz indexB) mWindow) uWrite = do - let Window it wk indexW _ = fromMaybe zeroWindow mWindow +loadWithIx1 with (DWArray _ sz uIndex mWindow) uWrite = do + let Window it wk uwIndex _ = fromMaybe zeroWindow mWindow wEnd = it + unSz wk - with $ iterM_ 0 it 1 (<) $ \ !i -> uWrite i (indexB i) - with $ iterM_ wEnd (unSz sz) 1 (<) $ \ !i -> uWrite i (indexB i) - return (\from to -> with $ iterM_ from to 1 (<) $ \ !i -> uWrite i (indexW i), it, wEnd) + with $ iterA_ 0 it 1 (<) $ \ !i -> uWrite i (uIndex i) + with $ iterA_ wEnd (unSz sz) 1 (<) $ \ !i -> uWrite i (uIndex i) + return (\from to -> with $ iterA_ from to 1 (<) $ \ !i -> uWrite i (uwIndex i), it, wEnd) {-# INLINE loadWithIx1 #-} +instance Index ix => Shape DW ix where + maxLinearSize = Just . linearSize + {-# INLINE maxLinearSize #-} + linearSize = SafeSz . totalElem . dwSize + {-# INLINE linearSize #-} + outerSize = dwSize + {-# INLINE outerSize #-} instance Load DW Ix1 e where - size = dSize . dwArray - {-# INLINE size #-} - getComp = dComp . dwArray - {-# INLINE getComp #-} - loadArrayM scheduler arr uWrite = do + makeArray c sz f = DWArray c sz f Nothing + {-# INLINE makeArray #-} + iterArrayLinearST_ scheduler arr uWrite = do (loadWindow, wStart, wEnd) <- loadWithIx1 (scheduleWork scheduler) arr uWrite let (chunkWidth, slackWidth) = (wEnd - wStart) `quotRem` numWorkers scheduler - loopM_ 0 (< numWorkers scheduler) (+ 1) $ \ !wid -> + loopA_ 0 (< numWorkers scheduler) (+ 1) $ \ !wid -> let !it' = wid * chunkWidth + wStart in loadWindow it' (it' + chunkWidth) when (slackWidth > 0) $ let !itSlack = numWorkers scheduler * chunkWidth + wStart in loadWindow itSlack (itSlack + slackWidth) - {-# INLINE loadArrayM #-} + {-# INLINE iterArrayLinearST_ #-} instance StrideLoad DW Ix1 e where - loadArrayWithStrideM scheduler stride sz arr uWrite = do - (loadWindow, (wStart, wEnd)) <- loadArrayWithIx1 (scheduleWork scheduler) arr stride sz uWrite - let (chunkWidth, slackWidth) = (wEnd - wStart) `quotRem` numWorkers scheduler - loopM_ 0 (< numWorkers scheduler) (+ 1) $ \ !wid -> - let !it' = wid * chunkWidth + wStart - in loadWindow (it', it' + chunkWidth) - when (slackWidth > 0) $ - let !itSlack = numWorkers scheduler * chunkWidth + wStart - in loadWindow (itSlack, itSlack + slackWidth) - {-# INLINE loadArrayWithStrideM #-} + iterArrayLinearWithStrideST_ scheduler stride sz arr uWrite = do + (loadWindow, (wStart, wEnd)) <- loadArrayWithIx1 (scheduleWork scheduler) arr stride sz uWrite + let (chunkWidth, slackWidth) = (wEnd - wStart) `quotRem` numWorkers scheduler + loopA_ 0 (< numWorkers scheduler) (+ 1) $ \ !wid -> + let !it' = wid * chunkWidth + wStart + in loadWindow (it', it' + chunkWidth) + when (slackWidth > 0) $ + let !itSlack = numWorkers scheduler * chunkWidth + wStart + in loadWindow (itSlack, itSlack + slackWidth) + {-# INLINE iterArrayLinearWithStrideST_ #-} -loadArrayWithIx1 :: - (Monad m) +loadArrayWithIx1 + :: Monad m => (m () -> m ()) -> Array DW Ix1 e -> Stride Ix1 -> Sz1 -> (Ix1 -> e -> m a) -> m ((Ix1, Ix1) -> m (), (Ix1, Ix1)) -loadArrayWithIx1 with (DWArray (DArray _ arrSz indexB) mWindow) stride _ uWrite = do - let Window it wk indexW _ = fromMaybe zeroWindow mWindow +loadArrayWithIx1 with (DWArray _ arrSz uIndex mWindow) stride _ uWrite = do + let Window it wk uwIndex _ = fromMaybe zeroWindow mWindow wEnd = it + unSz wk strideIx = unStride stride - with $ iterM_ 0 it strideIx (<) $ \ !i -> uWrite (i `div` strideIx) (indexB i) + with $ iterA_ 0 it strideIx (<) $ \ !i -> uWrite (i `div` strideIx) (uIndex i) with $ - iterM_ (strideStart stride wEnd) (unSz arrSz) strideIx (<) $ \ !i -> - uWrite (i `div` strideIx) (indexB i) + iterA_ (strideStart stride wEnd) (unSz arrSz) strideIx (<) $ \ !i -> + uWrite (i `div` strideIx) (uIndex i) return ( \(from, to) -> with $ - iterM_ (strideStart stride from) to strideIx (<) $ \ !i -> - uWrite (i `div` strideIx) (indexW i) - , (it, wEnd)) + iterA_ (strideStart stride from) to strideIx (<) $ \ !i -> + uWrite (i `div` strideIx) (uwIndex i) + , (it, wEnd) + ) {-# INLINE loadArrayWithIx1 #-} - - -loadWithIx2 :: - Monad m +loadWithIx2 + :: Monad m => (m () -> m ()) -> Array DW Ix2 t1 -> (Int -> t1 -> m ()) -> m (Ix2 -> m (), Ix2) loadWithIx2 with arr uWrite = do - let DWArray (DArray _ (Sz (m :. n)) indexB) window = arr - let Window (it :. jt) (Sz (wm :. wn)) indexW mUnrollHeight = fromMaybe zeroWindow window - let ib :. jb = (wm + it) :. (wn + jt) + let DWArray _ (Sz (m :. n)) uIndex window = arr + Window (it :. jt) (Sz (wm :. wn)) uwIndex mUnrollHeight = fromMaybe zeroWindow window + ib :. jb = (wm + it) :. (wn + jt) !blockHeight = maybe 1 (min 7 . max 1) mUnrollHeight - stride = oneStride - !sz = strideSize stride $ size arr - writeB !ix = uWrite (toLinearIndex sz ix) (indexB ix) + !sz = strideSize oneStride $ outerSize arr + writeB !ix = uWrite (toLinearIndex sz ix) (uIndex ix) {-# INLINE writeB #-} - writeW !ix = uWrite (toLinearIndex sz ix) (indexW ix) + writeW !ix = uWrite (toLinearIndex sz ix) (uwIndex ix) {-# INLINE writeW #-} - with $ iterM_ (0 :. 0) (it :. n) (1 :. 1) (<) writeB - with $ iterM_ (ib :. 0) (m :. n) (1 :. 1) (<) writeB - with $ iterM_ (it :. 0) (ib :. jt) (1 :. 1) (<) writeB - with $ iterM_ (it :. jb) (ib :. n) (1 :. 1) (<) writeB + with $ iterA_ (0 :. 0) (it :. n) (1 :. 1) (<) writeB + with $ iterA_ (ib :. 0) (m :. n) (1 :. 1) (<) writeB + with $ iterA_ (it :. 0) (ib :. jt) (1 :. 1) (<) writeB + with $ iterA_ (it :. jb) (ib :. n) (1 :. 1) (<) writeB let f (it' :. ib') = with $ unrollAndJam blockHeight (it' :. jt) (ib' :. jb) 1 writeW {-# INLINE f #-} return (f, it :. ib) {-# INLINE loadWithIx2 #-} -loadArrayWithIx2 :: - Monad m +loadArrayWithIx2 + :: Monad m => (m () -> m ()) -> Array DW Ix2 e -> Stride Ix2 @@ -301,33 +313,33 @@ -> (Int -> e -> m ()) -> m (Ix2 -> m (), Ix2) loadArrayWithIx2 with arr stride sz uWrite = do - let DWArray (DArray _ (Sz (m :. n)) indexB) window = arr - let Window (it :. jt) (Sz (wm :. wn)) indexW mUnrollHeight = fromMaybe zeroWindow window - let ib :. jb = (wm + it) :. (wn + jt) + let DWArray _ (Sz (m :. n)) uIndex window = arr + Window (it :. jt) (Sz (wm :. wn)) uwIndex mUnrollHeight = fromMaybe zeroWindow window + ib :. jb = (wm + it) :. (wn + jt) !blockHeight = maybe 1 (min 7 . max 1) mUnrollHeight strideIx@(is :. js) = unStride stride - writeB !ix = uWrite (toLinearIndexStride stride sz ix) (indexB ix) + writeB !ix = uWrite (toLinearIndexStride stride sz ix) (uIndex ix) {-# INLINE writeB #-} - writeW !ix = uWrite (toLinearIndexStride stride sz ix) (indexW ix) + writeW !ix = uWrite (toLinearIndexStride stride sz ix) (uwIndex ix) {-# INLINE writeW #-} - with $ iterM_ (0 :. 0) (it :. n) strideIx (<) writeB - with $ iterM_ (strideStart stride (ib :. 0)) (m :. n) strideIx (<) writeB - with $ iterM_ (strideStart stride (it :. 0)) (ib :. jt) strideIx (<) writeB - with $ iterM_ (strideStart stride (it :. jb)) (ib :. n) strideIx (<) writeB - f <- - if is > 1 || blockHeight <= 1 -- Turn off unrolling for vertical strides - then return $ \(it' :. ib') -> - iterM_ (strideStart stride (it' :. jt)) (ib' :. jb) strideIx (<) writeW - else return $ \(it' :. ib') -> - unrollAndJam blockHeight (strideStart stride (it' :. jt)) (ib' :. jb) js writeW - return (f, it :. ib) + with $ iterA_ (0 :. 0) (it :. n) strideIx (<) writeB + with $ iterA_ (strideStart stride (ib :. 0)) (m :. n) strideIx (<) writeB + with $ iterA_ (strideStart stride (it :. 0)) (ib :. jt) strideIx (<) writeB + with $ iterA_ (strideStart stride (it :. jb)) (ib :. n) strideIx (<) writeB + let f (it' :. ib') + | is > 1 || blockHeight <= 1 = + -- Turn off unrolling for vertical strides + iterA_ (strideStart stride (it' :. jt)) (ib' :. jb) strideIx (<) writeW + | otherwise = + unrollAndJam blockHeight (strideStart stride (it' :. jt)) (ib' :. jb) js writeW + {-# INLINE f #-} + return (with . f, it :. ib) {-# INLINE loadArrayWithIx2 #-} - loadWindowIx2 :: Monad m => Int -> (Ix2 -> m ()) -> Ix2 -> m () loadWindowIx2 nWorkers loadWindow (it :. ib) = do let !(chunkHeight, slackHeight) = (ib - it) `quotRem` nWorkers - loopM_ 0 (< nWorkers) (+ 1) $ \ !wid -> + loopA_ 0 (< nWorkers) (+ 1) $ \ !wid -> let !it' = wid * chunkHeight + it in loadWindow (it' :. (it' + chunkHeight)) when (slackHeight > 0) $ @@ -335,225 +347,168 @@ in loadWindow (itSlack :. (itSlack + slackHeight)) {-# INLINE loadWindowIx2 #-} - instance Load DW Ix2 e where - size = dSize . dwArray - {-# INLINE size #-} - getComp = dComp . dwArray - {-# INLINE getComp #-} - loadArrayM scheduler arr uWrite = - loadWithIx2 (scheduleWork scheduler) arr uWrite >>= - uncurry (loadWindowIx2 (numWorkers scheduler)) - {-# INLINE loadArrayM #-} + makeArray c sz f = DWArray c sz f Nothing + {-# INLINE makeArray #-} + iterArrayLinearST_ scheduler arr uWrite = + loadWithIx2 (scheduleWork scheduler) arr uWrite + >>= uncurry (loadWindowIx2 (numWorkers scheduler)) + {-# INLINE iterArrayLinearST_ #-} instance StrideLoad DW Ix2 e where - loadArrayWithStrideM scheduler stride sz arr uWrite = - loadArrayWithIx2 (scheduleWork scheduler) arr stride sz uWrite >>= - uncurry (loadWindowIx2 (numWorkers scheduler)) - {-# INLINE loadArrayWithStrideM #-} - + iterArrayLinearWithStrideST_ scheduler stride sz arr uWrite = + loadArrayWithIx2 (scheduleWork scheduler) arr stride sz uWrite + >>= uncurry (loadWindowIx2 (numWorkers scheduler)) + {-# INLINE iterArrayLinearWithStrideST_ #-} instance (Index (IxN n), Load DW (Ix (n - 1)) e) => Load DW (IxN n) e where - size = dSize . dwArray - {-# INLINE size #-} - getComp = dComp . dwArray - {-# INLINE getComp #-} - loadArrayM = loadWithIxN - {-# INLINE loadArrayM #-} + makeArray c sz f = DWArray c sz f Nothing + {-# INLINE makeArray #-} + iterArrayLinearST_ = loadWithIxN + {-# INLINE iterArrayLinearST_ #-} instance (Index (IxN n), StrideLoad DW (Ix (n - 1)) e) => StrideLoad DW (IxN n) e where - loadArrayWithStrideM = loadArrayWithIxN - {-# INLINE loadArrayWithStrideM #-} + iterArrayLinearWithStrideST_ = loadArrayWithIxN + {-# INLINE iterArrayLinearWithStrideST_ #-} -loadArrayWithIxN :: - (Index ix, Monad m, StrideLoad DW (Lower ix) e) - => Scheduler m () +loadArrayWithIxN + :: (Index ix, StrideLoad DW (Lower ix) e) + => Scheduler s () -> Stride ix -> Sz ix -> Array DW ix e - -> (Int -> e -> m ()) - -> m () + -> (Int -> e -> ST s ()) + -> ST s () loadArrayWithIxN scheduler stride szResult arr uWrite = do - let DWArray darr window = arr - DArray {dSize = szSource, dIndex = indexBorder} = darr - Window {windowStart, windowSize, windowIndex, windowUnrollIx2} = fromMaybe zeroWindow window - !(headSourceSize, lowerSourceSize) = unconsSz szSource + let DWArray _ sz uIndex window = arr + Window{windowStart, windowSize, windowIndex, windowUnrollIx2} = fromMaybe zeroWindow window + !(!headSourceSize, !lowerSourceSize) = unconsSz sz !lowerSize = snd $ unconsSz szResult - !(s, lowerStrideIx) = unconsDim $ unStride stride - !(curWindowStart, lowerWindowStart) = unconsDim windowStart - !(headWindowSz, tailWindowSz) = unconsSz windowSize + !(!s, !lowerStrideIx) = unconsDim $ unStride stride + !(!curWindowStart, lowerWindowStart) = unconsDim windowStart + !(!headWindowSz, tailWindowSz) = unconsSz windowSize !curWindowEnd = curWindowStart + unSz headWindowSz !pageElements = totalElem lowerSize - mkLowerWindow i = + lowerWindow = Window { windowStart = lowerWindowStart , windowSize = tailWindowSz - , windowIndex = windowIndex . consDim i + , windowIndex = \_ -> error "Window index uninitialized" , windowUnrollIx2 = windowUnrollIx2 } - mkLowerArray mw i = - DWArray - {dwArray = DArray Seq lowerSourceSize (indexBorder . consDim i), dwWindow = ($ i) <$> mw} - loadLower mw !i = - loadArrayWithStrideM - scheduler - (Stride lowerStrideIx) - lowerSize - (mkLowerArray mw i) - (\k -> uWrite (k + pageElements * (i `div` s))) - {-# NOINLINE loadLower #-} - loopM_ 0 (< headDim windowStart) (+ s) (loadLower Nothing) - loopM_ + mkLowerWindow !i = + lowerWindow + { windowIndex = windowIndex . consDim i + } + loadLower mkWindow !i = + let !lowerArray = + DWArray + { dwComp = Seq + , dwSize = lowerSourceSize + , dwIndex = uIndex . consDim i + , dwWindow = mkWindow i + } + !innerScheduler = + if numWorkers scheduler <= unSz (strideSize (Stride s) headSourceSize) + then trivialScheduler_ + else scheduler + in scheduleWork_ scheduler $ + iterArrayLinearWithStrideST_ innerScheduler (Stride lowerStrideIx) lowerSize lowerArray $ \k -> + uWrite (k + pageElements * (i `div` s)) + {-# INLINE loadLower #-} + loopA_ 0 (< headDim windowStart) (+ s) (loadLower (const Nothing)) + loopA_ (strideStart (Stride s) curWindowStart) (< curWindowEnd) (+ s) - (loadLower (Just mkLowerWindow)) - loopM_ (strideStart (Stride s) curWindowEnd) (< unSz headSourceSize) (+ s) (loadLower Nothing) + (loadLower (Just . mkLowerWindow)) + loopA_ + (strideStart (Stride s) curWindowEnd) + (< unSz headSourceSize) + (+ s) + (loadLower (const Nothing)) {-# INLINE loadArrayWithIxN #-} - - -loadWithIxN :: - (Index ix, Monad m, Load DW (Lower ix) e) - => Scheduler m () +loadWithIxN + :: (Index ix, Load DW (Lower ix) e) + => Scheduler s () -> Array DW ix e - -> (Int -> e -> m ()) - -> m () + -> (Int -> e -> ST s ()) + -> ST s () loadWithIxN scheduler arr uWrite = do - let DWArray darr window = arr - DArray {dSize = sz, dIndex = indexBorder} = darr - Window {windowStart, windowSize, windowIndex, windowUnrollIx2} = fromMaybe zeroWindow window - !(si, szL) = unconsSz sz + let DWArray _ sz uIndex window = arr + Window{windowStart, windowSize, windowIndex, windowUnrollIx2} = fromMaybe zeroWindow window + !(!si, !szL) = unconsSz sz !windowEnd = liftIndex2 (+) windowStart (unSz windowSize) - !(t, windowStartL) = unconsDim windowStart + !(!t, windowStartL) = unconsDim windowStart !pageElements = totalElem szL - mkLowerWindow i = + lowerWindow = Window { windowStart = windowStartL , windowSize = snd $ unconsSz windowSize - , windowIndex = windowIndex . consDim i + , windowIndex = \_ -> error "Window index uninitialized" , windowUnrollIx2 = windowUnrollIx2 } - mkLowerArray mw i = - DWArray {dwArray = DArray Seq szL (indexBorder . consDim i), dwWindow = ($ i) <$> mw} - loadLower mw !i = - scheduleWork_ scheduler $ - loadArrayM scheduler (mkLowerArray mw i) (\k -> uWrite (k + pageElements * i)) - {-# NOINLINE loadLower #-} - loopM_ 0 (< headDim windowStart) (+ 1) (loadLower Nothing) - loopM_ t (< headDim windowEnd) (+ 1) (loadLower (Just mkLowerWindow)) - loopM_ (headDim windowEnd) (< unSz si) (+ 1) (loadLower Nothing) + mkLowerWindow !i = + lowerWindow + { windowIndex = windowIndex . consDim i + } + loadLower mkWindow !i = + let !lowerArray = + DWArray + { dwComp = Seq + , dwSize = szL + , dwIndex = uIndex . consDim i + , dwWindow = mkWindow i + } + !innerScheduler = + if numWorkers scheduler <= unSz si + then trivialScheduler_ + else scheduler + in scheduleWork_ scheduler $ + iterArrayLinearST_ innerScheduler lowerArray (\k -> uWrite (k + pageElements * i)) + {-# INLINE loadLower #-} + loopA_ 0 (< headDim windowStart) (+ 1) (loadLower (const Nothing)) + loopA_ t (< headDim windowEnd) (+ 1) (loadLower (Just . mkLowerWindow)) + loopA_ (headDim windowEnd) (< unSz si) (+ 1) (loadLower (const Nothing)) {-# INLINE loadWithIxN #-} - - -unrollAndJam :: Monad m => - Int -- ^ Block height - -> Ix2 -- ^ Top corner - -> Ix2 -- ^ Bottom corner - -> Int -- ^ Column Stride - -> (Ix2 -> m ()) -- ^ Writing function - -> m () +unrollAndJam + :: Monad m + => Int + -- ^ Block height. Must not be zero. + -> Ix2 + -- ^ Top corner + -> Ix2 + -- ^ Bottom corner + -> Int + -- ^ Column Stride + -> (Ix2 -> m ()) + -- ^ Writing function + -> m () unrollAndJam !bH (it :. jt) (ib :. jb) js f = do - let f2 (i :. j) = f (i :. j) >> f ((i + 1) :. j) - let f3 (i :. j) = f (i :. j) >> f2 ((i + 1) :. j) - let f4 (i :. j) = f (i :. j) >> f3 ((i + 1) :. j) - let f5 (i :. j) = f (i :. j) >> f4 ((i + 1) :. j) - let f6 (i :. j) = f (i :. j) >> f5 ((i + 1) :. j) - let f7 (i :. j) = f (i :. j) >> f6 ((i + 1) :. j) - let f' = case bH of - 1 -> f - 2 -> f2 - 3 -> f3 - 4 -> f4 - 5 -> f5 - 6 -> f6 - _ -> f7 - let !ibS = ib - ((ib - it) `mod` bH) - loopM_ it (< ibS) (+ bH) $ \ !i -> - loopM_ jt (< jb) (+ js) $ \ !j -> + let + f2 (i :. j) = f (i :. j) >> f ((i + 1) :. j) + f3 (i :. j) = f (i :. j) >> f2 ((i + 1) :. j) + f4 (i :. j) = f (i :. j) >> f3 ((i + 1) :. j) + f5 (i :. j) = f (i :. j) >> f4 ((i + 1) :. j) + f6 (i :. j) = f (i :. j) >> f5 ((i + 1) :. j) + f7 (i :. j) = f (i :. j) >> f6 ((i + 1) :. j) + f' = case bH of + 1 -> f + 2 -> f2 + 3 -> f3 + 4 -> f4 + 5 -> f5 + 6 -> f6 + _ -> f7 + !ibS = ib - ((ib - it) `modInt` bH) + loopA_ it (< ibS) (+ bH) $ \ !i -> + loopA_ jt (< jb) (+ js) $ \ !j -> f' (i :. j) - loopM_ ibS (< ib) (+ 1) $ \ !i -> - loopM_ jt (< jb) (+ js) $ \ !j -> + loopA_ ibS (< ib) (+ 1) $ \ !i -> + loopA_ jt (< jb) (+ js) $ \ !j -> f (i :. j) {-# INLINE unrollAndJam #-} - - -- TODO: Implement Hilbert curve - -toIx2Window :: Window Ix2T e -> Window Ix2 e -toIx2Window Window {..} = - Window - { windowStart = toIx2 windowStart - , windowSize = SafeSz (toIx2 $ unSz windowSize) - , windowIndex = windowIndex . fromIx2 - , windowUnrollIx2 = windowUnrollIx2 - } -{-# INLINE toIx2Window #-} - -toIx2ArrayDW :: Array DW Ix2T e -> Array DW Ix2 e -toIx2ArrayDW DWArray {dwArray, dwWindow} = - DWArray - { dwArray = - dwArray {dIndex = dIndex dwArray . fromIx2, dSize = SafeSz (toIx2 (unSz (dSize dwArray)))} - , dwWindow = fmap toIx2Window dwWindow - } -{-# INLINE toIx2ArrayDW #-} - - -instance Load DW Ix2T e where - size = dSize . dwArray - {-# INLINE size #-} - getComp = dComp . dwArray - {-# INLINE getComp #-} - loadArrayM scheduler arr = - loadArrayWithStrideM scheduler oneStride (size arr) arr - {-# INLINE loadArrayM #-} - -instance StrideLoad DW Ix2T e where - loadArrayWithStrideM scheduler stride sz arr = - loadArrayWithStrideM - scheduler - (Stride $ toIx2 $ unStride stride) - (SafeSz (toIx2 (unSz sz))) - (toIx2ArrayDW arr) - {-# INLINE loadArrayWithStrideM #-} - -instance Load DW Ix3T e where - size = dSize . dwArray - {-# INLINE size #-} - getComp = dComp . dwArray - {-# INLINE getComp #-} - loadArrayM scheduler arr = - loadArrayWithStrideM scheduler oneStride (size arr) arr - {-# INLINE loadArrayM #-} - -instance StrideLoad DW Ix3T e where - loadArrayWithStrideM = loadArrayWithIxN - {-# INLINE loadArrayWithStrideM #-} - - -instance Load DW Ix4T e where - size = dSize . dwArray - {-# INLINE size #-} - getComp = dComp . dwArray - {-# INLINE getComp #-} - loadArrayM scheduler arr = loadArrayWithStrideM scheduler oneStride (size arr) arr - {-# INLINE loadArrayM #-} - -instance StrideLoad DW Ix4T e where - loadArrayWithStrideM = loadArrayWithIxN - {-# INLINE loadArrayWithStrideM #-} - - -instance Load DW Ix5T e where - size = dSize . dwArray - {-# INLINE size #-} - getComp = dComp . dwArray - {-# INLINE getComp #-} - loadArrayM scheduler arr = loadArrayWithStrideM scheduler oneStride (size arr) arr - {-# INLINE loadArrayM #-} -instance StrideLoad DW Ix5T e where - loadArrayWithStrideM = loadArrayWithIxN - {-# INLINE loadArrayWithStrideM #-} diff --git a/src/Data/Massiv/Array/Manifest.hs b/src/Data/Massiv/Array/Manifest.hs --- a/src/Data/Massiv/Array/Manifest.hs +++ b/src/Data/Massiv/Array/Manifest.hs @@ -1,71 +1,161 @@ +{-# LANGUAGE BangPatterns #-} {-# LANGUAGE CPP #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE PatternSynonyms #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} + -- | -- Module : Data.Massiv.Array.Manifest --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Manifest - ( +module Data.Massiv.Array.Manifest ( -- * Manifest - Manifest - , toManifest - , M + Manifest, + + -- ** Generate + generateArray, + generateArrayLinear, + generateArrayS, + generateArrayLinearS, + generateSplitSeedArray, + + -- ** Stateful worker threads + generateArrayWS, + generateArrayLinearWS, + + -- ** Unfold + unfoldrPrimM_, + iunfoldrPrimM_, + unfoldrPrimM, + iunfoldrPrimM, + unfoldlPrimM_, + iunfoldlPrimM_, + unfoldlPrimM, + iunfoldlPrimM, + + -- ** Mapping + forPrimM, + forPrimM_, + iforPrimM, + iforPrimM_, + iforLinearPrimM, + iforLinearPrimM_, + for2PrimM_, + ifor2PrimM_, + -- * Boxed - , B(..) - , N(..) - , Uninitialized(..) + B (..), + BL (..), + BN (..), + N, + pattern N, + Uninitialized (..), + + -- ** Access + findIndex, + -- ** Conversion -- $boxed_conversion_note - , unwrapArray - , evalArray - , unwrapMutableArray - , evalMutableArray - , unwrapNormalFormArray - , evalNormalFormArray - , unwrapNormalFormMutableArray - , evalNormalFormMutableArray + toLazyArray, + evalLazyArray, + forceLazyArray, + unwrapNormalForm, + evalNormalForm, + + -- *** Primitive Boxed Array + unwrapLazyArray, + wrapLazyArray, + unwrapArray, + evalArray, + unwrapMutableArray, + unwrapMutableLazyArray, + evalMutableArray, + unwrapNormalFormArray, + evalNormalFormArray, + unwrapNormalFormMutableArray, + evalNormalFormMutableArray, + + -- *** Boxed Vector + toBoxedVector, + toBoxedMVector, + fromBoxedVector, + fromBoxedMVector, + evalBoxedVector, + evalBoxedMVector, + -- * Primitive - , P(..) - , Prim + P (..), + Prim, + -- ** Conversion - , toByteArray - , fromByteArray - , fromByteArrayM - , toMutableByteArray - , fromMutableByteArray - , fromMutableByteArrayM + + -- *** Primitive ByteArray + toByteArray, + toByteArrayM, + unwrapByteArray, + unwrapByteArrayOffset, + fromByteArray, + fromByteArrayM, + fromByteArrayOffsetM, + toMutableByteArray, + unwrapMutableByteArray, + unwrapMutableByteArrayOffset, + fromMutableByteArray, + fromMutableByteArrayM, + fromMutableByteArrayOffsetM, + + -- *** Primitive Vector + toPrimitiveVector, + toPrimitiveMVector, + fromPrimitiveVector, + fromPrimitiveMVector, + -- * Storable - , S(..) - , Storable + S (..), + Storable, + mallocCompute, + mallocCopy, + -- ** Conversion - , toStorableVector - , toStorableMVector - -- ** Direct Pointer Access - , withPtr + + -- *** Storable Vector + toStorableVector, + toStorableMVector, + fromStorableVector, + fromStorableMVector, + + -- *** Direct Pointer Access + withPtr, + -- * Unboxed - , U(..) - , Unbox + U (..), + Unbox, + -- ** Conversion - , toUnboxedVector - , toUnboxedMVector + + -- *** Unboxed Vector + toUnboxedVector, + toUnboxedMVector, + fromUnboxedVector, + fromUnboxedMVector, + -- * ByteString Conversion - , fromByteString - , castFromByteString - , toByteString - , castToByteString - , toBuilder - ) where + fromByteString, + castFromByteString, + toByteString, + castToByteString, + toBuilder, + castToBuilder, +) where -import Data.ByteString as S +import Control.Monad +import Data.ByteString as S hiding (findIndex) import Data.ByteString.Builder import Data.ByteString.Internal import Data.ByteString.Unsafe as SU @@ -74,23 +164,26 @@ import Data.Massiv.Array.Manifest.Primitive import Data.Massiv.Array.Manifest.Storable import Data.Massiv.Array.Manifest.Unboxed +import Data.Massiv.Array.Mutable import Data.Massiv.Array.Ops.Fold import Data.Massiv.Core.Common -import Data.Massiv.Core.Index.Internal (Sz(..)) import Data.Word (Word8) - --- | /O(1)/ - Convert a strict ByteString into a manifest array. Will return `Nothing` if length +-- | /O(n)/ - Convert a strict ByteString into a manifest array. Will return `Nothing` if length -- doesn't match the total number of elements of new array. -- -- @since 0.2.1 -fromByteString :: - Comp -- ^ Computation strategy - -> ByteString -- ^ Strict ByteString to use as a source. - -> Array M Ix1 Word8 -fromByteString comp bs = MArray comp (SafeSz (S.length bs)) (SU.unsafeIndex bs) +fromByteString + :: Load r Ix1 Word8 + => Comp + -- ^ Computation strategy + -> ByteString + -- ^ Strict ByteString to use as a source. + -> Vector r Word8 +fromByteString comp bs = makeArrayLinear comp (SafeSz (S.length bs)) (SU.unsafeIndex bs) {-# INLINE fromByteString #-} +{- FOURMOLU_DISABLE -} -- | /O(n)/ - Convert any source array into a strict `ByteString`. In case when the source array is -- actually storable, no memory copy will occur. -- @@ -102,31 +195,38 @@ toByteString = castToByteString . #if __GLASGOW_HASKELL__ >= 820 convert - {- For ghc-8.0 `covert` results in "internal error: ARR_WORDS object entered!" -} + {- For ghc-8.0 `convert` results in "internal error: ARR_WORDS object entered!" -} #else compute #endif - --fst $ unfoldrN (totalElem (size arr)) (\ !i -> Just (unsafeLinearIndex arr i, i + 1)) 0 {-# INLINE toByteString #-} +{- FOURMOLU_ENABLE -} --- | /O(n)/ - Conversion of array monoidally into a ByteString Builder. +-- | /O(n)/ - Conversion of array monoidally into a ByteString `Builder`. -- -- @since 0.2.1 -toBuilder :: Source r ix e => (e -> Builder) -> Array r ix e -> Builder +toBuilder :: (Index ix, Source r e) => (e -> Builder) -> Array r ix e -> Builder toBuilder = foldMono {-# INLINE toBuilder #-} +-- | /O(1)/ - Cast a storable array of `Word8` to ByteString `Builder`. +-- +-- @since 0.5.0 +castToBuilder :: Index ix => Array S ix Word8 -> Builder +castToBuilder = byteString . castToByteString +{-# INLINE castToBuilder #-} + -- | /O(1)/ - Cast a `S`torable array into a strict `ByteString` -- -- @since 0.3.0 -castToByteString :: Array S ix Word8 -> ByteString +castToByteString :: Index ix => Array S ix Word8 -> ByteString castToByteString = (\(fp, len) -> PS fp 0 len) . unsafeArrayToForeignPtr {-# INLINE castToByteString #-} -- | /O(1)/ - Cast a strict `ByteString` into a `S`torable array -- -- @since 0.3.0 -castFromByteString :: Comp -> ByteString -> Array S Ix1 Word8 +castFromByteString :: Comp -> ByteString -> Vector S Word8 castFromByteString comp (PS fp offset len) = unsafeArrayFromForeignPtr comp fp offset (Sz len) {-# INLINE castFromByteString #-} @@ -136,3 +236,45 @@ -- `Data.Primitive.Array.Array`, which holds the pointers to values isn't copied around, it is always -- kept as the same array. Conversion to Massiv boxed array will undergo evaluation during which -- computation strategies will be respected. + +-- | /O(n)/ - Perform a row-major search starting at @0@ for an element. Returns the index +-- of the first occurance of an element or `Nothing` if a predicate could not be satisifed +-- after it was applyied to all elements of the array. +-- +-- @since 0.5.5 +findIndex :: (Index ix, Manifest r e) => (e -> Bool) -> Array r ix e -> Maybe ix +findIndex f arr = go 0 + where + !sz = size arr + !k = totalElem sz + go !i = do + guard (i < k) + if f (unsafeLinearIndex arr i) + then Just $ fromLinearIndex sz i + else go (i + 1) +{-# INLINE findIndex #-} + +-- | Very similar to @`computeAs` `S`@ except load the source array into memory allocated +-- with @malloc@ on C heap. It can potentially be useful when iteroperating with some C +-- programs. +-- +-- @since 0.5.9 +mallocCompute + :: forall r ix e. (Size r, Load r ix e, Storable e) => Array r ix e -> IO (Array S ix e) +mallocCompute arr = do + let sz = size arr + marr <- unsafeMallocMArray sz + computeInto marr arr + unsafeFreeze (getComp arr) marr +{-# INLINE mallocCompute #-} + +-- | Allocate memory on C heap with @malloc@ and copy the source array over. +-- +-- @since 0.5.9 +mallocCopy :: forall ix e. (Index ix, Storable e) => Array S ix e -> IO (Array S ix e) +mallocCopy arr = do + let sz = size arr + marr <- unsafeMallocMArray sz + unsafeArrayLinearCopy arr 0 marr 0 (SafeSz (totalElem sz)) + unsafeFreeze (getComp arr) marr +{-# INLINE mallocCopy #-} diff --git a/src/Data/Massiv/Array/Manifest/Boxed.hs b/src/Data/Massiv/Array/Manifest/Boxed.hs --- a/src/Data/Massiv/Array/Manifest/Boxed.hs +++ b/src/Data/Massiv/Array/Manifest/Boxed.hs @@ -1,208 +1,401 @@ -{-# OPTIONS_GHC -fno-warn-orphans #-} {-# LANGUAGE BangPatterns #-} {-# LANGUAGE CPP #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MagicHash #-} {-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE PatternSynonyms #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UndecidableInstances #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} + -- | -- Module : Data.Massiv.Array.Manifest.Boxed --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Manifest.Boxed - ( B(..) - , N(..) - , Array(..) - , unwrapArray - , evalArray - , unwrapMutableArray - , evalMutableArray - , unwrapNormalFormArray - , evalNormalFormArray - , unwrapNormalFormMutableArray - , evalNormalFormMutableArray - , castArrayToVector - , castVectorToArray - , seqArray - , deepseqArray - ) where +module Data.Massiv.Array.Manifest.Boxed ( + B (..), + BL (..), + BN (..), + N, + pattern N, + Array (..), + MArray (..), + wrapLazyArray, + unwrapLazyArray, + unwrapNormalForm, + evalNormalForm, + unwrapArray, + evalArray, + toLazyArray, + evalLazyArray, + forceLazyArray, + unwrapMutableArray, + unwrapMutableLazyArray, + evalMutableArray, + unwrapNormalFormArray, + evalNormalFormArray, + unwrapNormalFormMutableArray, + evalNormalFormMutableArray, + toBoxedVector, + toBoxedMVector, + fromBoxedVector, + fromBoxedMVector, + evalBoxedVector, + evalBoxedMVector, + evalNormalBoxedVector, + evalNormalBoxedMVector, + coerceBoxedArray, + coerceNormalBoxedArray, + seqArray, + deepseqArray, +) where -import Control.DeepSeq (NFData(..), deepseq) +import Control.DeepSeq (NFData (..), deepseq) import Control.Exception import Control.Monad ((>=>)) import Control.Monad.Primitive -import Control.Monad.ST (runST) -import qualified Data.Foldable as F (Foldable(..)) -import Data.Massiv.Array.Delayed.Pull (eq, ord) +import qualified Data.Foldable as F (Foldable (..)) +import Data.Massiv.Array.Delayed.Pull (D) import Data.Massiv.Array.Delayed.Push (DL) import Data.Massiv.Array.Delayed.Stream (DS) -import Data.Massiv.Array.Manifest.Internal (M, computeAs, toManifest) +import Data.Massiv.Array.Manifest.Internal (compute, computeAs) import Data.Massiv.Array.Manifest.List as L -import Data.Massiv.Array.Manifest.Vector.Stream as S (steps) import Data.Massiv.Array.Mutable import Data.Massiv.Array.Ops.Fold import Data.Massiv.Array.Ops.Fold.Internal import Data.Massiv.Array.Ops.Map (traverseA) import Data.Massiv.Core.Common import Data.Massiv.Core.List +import Data.Massiv.Core.Operations +import Data.Massiv.Vector.Stream as S (isteps, steps) import qualified Data.Primitive.Array as A import qualified Data.Vector as VB -import qualified Data.Vector.Mutable as VB -import GHC.Base (build) +import qualified Data.Vector.Mutable as MVB import GHC.Exts as GHC -import Prelude hiding (mapM) import System.IO.Unsafe (unsafePerformIO) +import Prelude hiding (mapM, replicate) +#if !MIN_VERSION_vector(0,13,0) +import Unsafe.Coerce (unsafeCoerce) +#endif #include "massiv.h" -sizeofArray :: A.Array e -> Int -sizeofMutableArray :: A.MutableArray s e -> Int -#if MIN_VERSION_primitive(0,6,2) -sizeofArray = A.sizeofArray -sizeofMutableArray = A.sizeofMutableArray -#else -sizeofArray (A.Array a#) = I# (sizeofArray# a#) -sizeofMutableArray (A.MutableArray ma#) = I# (sizeofMutableArray# ma#) -#endif +---------------- +-- Boxed Lazy -- +---------------- ------------------- --- Boxed Strict -- ------------------- +-- | Array representation for Boxed elements. This data structure is lazy with +-- respect to its elements. +-- +-- ====__Example__ +-- +-- Memoized version of a factorial that relies on laziness. Note that +-- computing memoized factorial of a million would likely overflow memory. +-- +-- >>> import Data.Massiv.Array as A +-- >>> :{ +-- mkMemoFactorial :: Int -> (Int -> Integer) +-- mkMemoFactorial n = +-- let arr = makeVectorR BL Seq (Sz1 n) fact +-- fact i | i == 0 = 1 +-- | otherwise = (arr ! (i - 1)) * toInteger i +-- in (arr !) +-- :} +-- +-- >>> let fact = mkMemoFactorial 1000001 +-- >>> fact 50 +-- 30414093201713378043612608166064768844377641568960512000000000000 +-- >>> length $ show $ fact 5000 +-- 16326 +data BL = BL deriving (Show) --- | Array representation for Boxed elements. This structure is element and --- spine strict, but elements are strict to Weak Head Normal Form (WHNF) only. -data B = B deriving Show +data instance Array BL ix e = BLArray + { blComp :: !Comp + , blSize :: !(Sz ix) + , blOffset :: {-# UNPACK #-} !Int + , blData :: {-# UNPACK #-} !(A.Array e) + } -data instance Array B ix e = BArray { bComp :: !Comp - , bSize :: !(Sz ix) - , bData :: {-# UNPACK #-} !(A.Array e) - } +data instance MArray s BL ix e + = MBLArray !(Sz ix) {-# UNPACK #-} !Int {-# UNPACK #-} !(A.MutableArray s e) -instance (Ragged L ix e, Show e) => Show (Array B ix e) where +instance (Ragged L ix e, Show e) => Show (Array BL ix e) where showsPrec = showsArrayPrec id showList = showArrayList instance (Ragged L ix e, Show e) => Show (Array DL ix e) where - showsPrec = showsArrayPrec (computeAs B) + showsPrec = showsArrayPrec (computeAs BL) showList = showArrayList instance Show e => Show (Array DS Ix1 e) where - showsPrec = showsArrayPrec (computeAs B) + showsPrec = showsArrayPrec (computeAs BL) showList = showArrayList - -instance (Index ix, NFData e) => NFData (Array B ix e) where +instance (Index ix, NFData e) => NFData (Array BL ix e) where rnf = (`deepseqArray` ()) {-# INLINE rnf #-} -instance (Index ix, Eq e) => Eq (Array B ix e) where - (==) = eq (==) +instance (Index ix, Eq e) => Eq (Array BL ix e) where + (==) = eqArrays (==) {-# INLINE (==) #-} -instance (Index ix, Ord e) => Ord (Array B ix e) where - compare = ord compare +instance (Index ix, Ord e) => Ord (Array BL ix e) where + compare = compareArrays compare {-# INLINE compare #-} -instance Index ix => Construct B ix e where - setComp c arr = arr { bComp = c } +instance Strategy BL where + setComp c arr = arr{blComp = c} {-# INLINE setComp #-} - - makeArray !comp !sz f = unsafePerformIO $ generateArray comp sz (\ !ix -> return $! f ix) - {-# INLINE makeArray #-} + getComp = blComp + {-# INLINE getComp #-} + repr = BL -instance Index ix => Source B ix e where - unsafeLinearIndex (BArray _ _ a) = - INDEX_CHECK("(Source B ix e).unsafeLinearIndex", Sz . sizeofArray, A.indexArray) a +instance Source BL e where + unsafeLinearIndex (BLArray _ _sz o a) i = + indexAssert "BL.unsafeLinearIndex" (SafeSz . A.sizeofArray) A.indexArray a (i + o) {-# INLINE unsafeLinearIndex #-} + unsafeOuterSlice (BLArray c _ o a) szL i = BLArray c szL (i * totalElem szL + o) a + {-# INLINE unsafeOuterSlice #-} -instance Index ix => Resize B ix where - unsafeResize !sz !arr = arr { bSize = sz } + unsafeLinearSlice i k (BLArray c _ o a) = BLArray c k (o + i) a + {-# INLINE unsafeLinearSlice #-} + +instance Manifest BL e where + unsafeLinearIndexM (BLArray _ _sz o a) i = + indexAssert "BL.unsafeLinearIndexM" (SafeSz . A.sizeofArray) A.indexArray a (i + o) + {-# INLINE unsafeLinearIndexM #-} + + sizeOfMArray (MBLArray sz _ _) = sz + {-# INLINE sizeOfMArray #-} + + unsafeResizeMArray sz (MBLArray _ off marr) = MBLArray sz off marr + {-# INLINE unsafeResizeMArray #-} + + unsafeLinearSliceMArray i k (MBLArray _ o a) = MBLArray k (i + o) a + {-# INLINE unsafeLinearSliceMArray #-} + + unsafeThaw (BLArray _ sz o a) = MBLArray sz o <$> A.unsafeThawArray a + {-# INLINE unsafeThaw #-} + + unsafeFreeze comp (MBLArray sz o ma) = BLArray comp sz o <$> A.unsafeFreezeArray ma + {-# INLINE unsafeFreeze #-} + + unsafeNew sz = MBLArray sz 0 <$> A.newArray (totalElem sz) uninitialized + {-# INLINE unsafeNew #-} + + initialize _ = return () + {-# INLINE initialize #-} + + newMArray sz e = MBLArray sz 0 <$> A.newArray (totalElem sz) e + {-# INLINE newMArray #-} + + unsafeLinearRead (MBLArray _ o ma) i = + indexAssert "B.unsafeLinearRead" (SafeSz . A.sizeofMutableArray) A.readArray ma (i + o) + {-# INLINE unsafeLinearRead #-} + + unsafeLinearWrite (MBLArray _sz o ma) i = + indexAssert "B.unsafeLinearWrite" (SafeSz . A.sizeofMutableArray) A.writeArray ma (i + o) + {-# INLINE unsafeLinearWrite #-} + +instance Size BL where + size = blSize + {-# INLINE size #-} + unsafeResize !sz !arr = arr{blSize = sz} {-# INLINE unsafeResize #-} -instance Index ix => Extract B ix e where - unsafeExtract !sIx !newSz !arr = unsafeExtract sIx newSz (toManifest arr) - {-# INLINE unsafeExtract #-} +instance Index ix => Shape BL ix where + maxLinearSize = Just . SafeSz . elemsCount + {-# INLINE maxLinearSize #-} +instance Index ix => Load BL ix e where + makeArray comp sz f = compute (makeArray comp sz f :: Array D ix e) + {-# INLINE makeArray #-} -instance ( Index ix - , Index (Lower ix) - , Elt M ix e ~ Array M (Lower ix) e - , Elt B ix e ~ Array M (Lower ix) e - ) => - OuterSlice B ix e where - unsafeOuterSlice arr = unsafeOuterSlice (toManifest arr) - {-# INLINE unsafeOuterSlice #-} + makeArrayLinear !comp !sz f = unsafePerformIO $ generateArrayLinear comp sz (pure . f) + {-# INLINE makeArrayLinear #-} -instance ( Index ix - , Index (Lower ix) - , Elt M ix e ~ Array M (Lower ix) e - , Elt B ix e ~ Array M (Lower ix) e - ) => - InnerSlice B ix e where - unsafeInnerSlice arr = unsafeInnerSlice (toManifest arr) - {-# INLINE unsafeInnerSlice #-} + replicate comp sz e = runST (newMArray sz e >>= unsafeFreeze comp) + {-# INLINE replicate #-} -instance {-# OVERLAPPING #-} Slice B Ix1 e where - unsafeSlice arr i _ _ = pure (unsafeLinearIndex arr i) - {-# INLINE unsafeSlice #-} + iterArrayLinearST_ !scheduler !arr = + splitLinearlyWith_ scheduler (elemsCount arr) (unsafeLinearIndex arr) + {-# INLINE iterArrayLinearST_ #-} +instance Index ix => StrideLoad BL ix e -instance Index ix => Manifest B ix e where +instance Index ix => Stream BL ix e where + toStream = S.steps + {-# INLINE toStream #-} + toStreamIx = S.isteps + {-# INLINE toStreamIx #-} - unsafeLinearIndexM (BArray _ _ a) = - INDEX_CHECK("(Manifest B ix e).unsafeLinearIndexM", Sz . sizeofArray, A.indexArray) a +-- | Row-major sequential folding over a Boxed array. +instance Index ix => Foldable (Array BL ix) where + fold = fold + {-# INLINE fold #-} + foldMap = foldMono + {-# INLINE foldMap #-} + foldl = lazyFoldlS + {-# INLINE foldl #-} + foldl' = foldlS + {-# INLINE foldl' #-} + foldr = foldrFB + {-# INLINE foldr #-} + foldr' = foldrS + {-# INLINE foldr' #-} + null (BLArray _ sz _ _) = totalElem sz == 0 + {-# INLINE null #-} + length = totalElem . size + {-# INLINE length #-} + toList arr = build (\c n -> foldrFB c n arr) + {-# INLINE toList #-} + +instance Index ix => Functor (Array BL ix) where + fmap f arr = makeArrayLinear (blComp arr) (blSize arr) (f . unsafeLinearIndex arr) + {-# INLINE fmap #-} + (<$) e arr = replicate (getComp arr) (size arr) e + {-# INLINE (<$) #-} + +instance Index ix => Traversable (Array BL ix) where + traverse = traverseA + {-# INLINE traverse #-} + +instance (IsList (Array L ix e), Ragged L ix e) => IsList (Array BL ix e) where + type Item (Array BL ix e) = Item (Array L ix e) + fromList = L.fromLists' Seq + {-# INLINE fromList #-} + toList = GHC.toList . toListArray + {-# INLINE toList #-} + +instance Num e => FoldNumeric BL e where + unsafeDotProduct = defaultUnsafeDotProduct + {-# INLINE unsafeDotProduct #-} + powerSumArray = defaultPowerSumArray + {-# INLINE powerSumArray #-} + foldArray = defaultFoldArray + {-# INLINE foldArray #-} + +instance Num e => Numeric BL e where + unsafeLiftArray = defaultUnsafeLiftArray + {-# INLINE unsafeLiftArray #-} + unsafeLiftArray2 = defaultUnsafeLiftArray2 + {-# INLINE unsafeLiftArray2 #-} + +------------------ +-- Boxed Strict -- +------------------ + +-- | Array representation for Boxed elements. Its elements are strict to Weak +-- Head Normal Form (WHNF) only. +data B = B deriving (Show) + +newtype instance Array B ix e = BArray (Array BL ix e) + +newtype instance MArray s B ix e = MBArray (MArray s BL ix e) + +instance (Ragged L ix e, Show e) => Show (Array B ix e) where + showsPrec = showsArrayPrec id + showList = showArrayList + +instance (Index ix, NFData e) => NFData (Array B ix e) where + rnf = (`deepseqArray` ()) . coerce + {-# INLINE rnf #-} + +instance (Index ix, Eq e) => Eq (Array B ix e) where + (==) = eqArrays (==) + {-# INLINE (==) #-} + +instance (Index ix, Ord e) => Ord (Array B ix e) where + compare = compareArrays compare + {-# INLINE compare #-} + +instance Source B e where + unsafeLinearIndex arr = unsafeLinearIndex (toLazyArray arr) + {-# INLINE unsafeLinearIndex #-} + + unsafeLinearSlice i k arr = coerce (unsafeLinearSlice i k (toLazyArray arr)) + {-# INLINE unsafeLinearSlice #-} + + unsafeOuterSlice arr i = coerce (unsafeOuterSlice (toLazyArray arr) i) + {-# INLINE unsafeOuterSlice #-} + +instance Strategy B where + getComp = blComp . coerce + {-# INLINE getComp #-} + setComp c arr = coerceBoxedArray (coerce arr){blComp = c} + {-# INLINE setComp #-} + repr = B + +instance Index ix => Shape B ix where + maxLinearSize = Just . SafeSz . elemsCount + {-# INLINE maxLinearSize #-} + +instance Size B where + size = blSize . coerce + {-# INLINE size #-} + unsafeResize sz = coerce (\arr -> arr{blSize = sz}) + {-# INLINE unsafeResize #-} + +instance Manifest B e where + unsafeLinearIndexM = coerce unsafeLinearIndexM {-# INLINE unsafeLinearIndexM #-} + sizeOfMArray = sizeOfMArray . coerce + {-# INLINE sizeOfMArray #-} -instance Index ix => Mutable B ix e where - data MArray s B ix e = MBArray !(Sz ix) {-# UNPACK #-} !(A.MutableArray s e) + unsafeResizeMArray sz = MBArray . unsafeResizeMArray sz . coerce + {-# INLINE unsafeResizeMArray #-} - msize (MBArray sz _) = sz - {-# INLINE msize #-} + unsafeLinearSliceMArray i k = MBArray . unsafeLinearSliceMArray i k . coerce + {-# INLINE unsafeLinearSliceMArray #-} - unsafeThaw (BArray _ sz a) = MBArray sz <$> A.unsafeThawArray a + unsafeThaw arr = MBArray <$> unsafeThaw (coerce arr) {-# INLINE unsafeThaw #-} - unsafeFreeze comp (MBArray sz ma) = BArray comp sz <$> A.unsafeFreezeArray ma + unsafeFreeze comp marr = BArray <$> unsafeFreeze comp (coerce marr) {-# INLINE unsafeFreeze #-} - unsafeNew sz = MBArray sz <$> A.newArray (totalElem sz) uninitialized + unsafeNew sz = MBArray <$> unsafeNew sz {-# INLINE unsafeNew #-} initialize _ = return () {-# INLINE initialize #-} - unsafeLinearRead (MBArray _ ma) = - INDEX_CHECK("(Mutable B ix e).unsafeLinearRead", Sz . sizeofMutableArray, A.readArray) ma + newMArray sz !e = MBArray <$> newMArray sz e + {-# INLINE newMArray #-} + + unsafeLinearRead ma = unsafeLinearRead (coerce ma) {-# INLINE unsafeLinearRead #-} - unsafeLinearWrite (MBArray _ ma) i e = e `seq` - INDEX_CHECK("(Mutable B ix e).unsafeLinearWrite", Sz . sizeofMutableArray, A.writeArray) ma i e + unsafeLinearWrite ma i e = e `seq` unsafeLinearWrite (coerce ma) i e {-# INLINE unsafeLinearWrite #-} instance Index ix => Load B ix e where - type R B = M - size = bSize - {-# INLINE size #-} - getComp = bComp - {-# INLINE getComp #-} - loadArrayM !scheduler !arr = splitLinearlyWith_ scheduler (elemsCount arr) (unsafeLinearIndex arr) - {-# INLINE loadArrayM #-} + makeArray comp sz f = compute (makeArray comp sz f :: Array D ix e) + {-# INLINE makeArray #-} + makeArrayLinear !comp !sz f = unsafePerformIO $ generateArrayLinear comp sz (pure . f) + {-# INLINE makeArrayLinear #-} + + replicate comp sz e = runST (newMArray sz e >>= unsafeFreeze comp) + {-# INLINE replicate #-} + + iterArrayLinearST_ scheduler = coerce (iterArrayLinearST_ scheduler) + {-# INLINE iterArrayLinearST_ #-} + instance Index ix => StrideLoad B ix e instance Index ix => Stream B ix e where toStream = S.steps {-# INLINE toStream #-} - + toStreamIx = S.isteps + {-# INLINE toStreamIx #-} -- | Row-major sequential folding over a Boxed array. instance Index ix => Foldable (Array B ix) where @@ -218,178 +411,186 @@ {-# INLINE foldr #-} foldr' = foldrS {-# INLINE foldr' #-} - null (BArray _ sz _) = totalElem sz == 0 + null arr = totalElem (size arr) == 0 {-# INLINE null #-} length = totalElem . size {-# INLINE length #-} - toList arr = build (\ c n -> foldrFB c n arr) + toList arr = build (\c n -> foldrFB c n arr) {-# INLINE toList #-} - instance Index ix => Functor (Array B ix) where - fmap f arr = makeArrayLinear (bComp arr) (bSize arr) (f . unsafeLinearIndex arr) + fmap f arr = makeArrayLinear (getComp arr) (size arr) (f . unsafeLinearIndex arr) {-# INLINE fmap #-} + (<$) !e arr = replicate (getComp arr) (size arr) e + {-# INLINE (<$) #-} instance Index ix => Traversable (Array B ix) where traverse = traverseA {-# INLINE traverse #-} -instance ( IsList (Array L ix e) - , Nested LN ix e - , Nested L ix e - , Ragged L ix e - ) => - IsList (Array B ix e) where +instance (IsList (Array L ix e), Ragged L ix e) => IsList (Array B ix e) where type Item (Array B ix e) = Item (Array L ix e) fromList = L.fromLists' Seq {-# INLINE fromList #-} toList = GHC.toList . toListArray {-# INLINE toList #-} +instance Num e => FoldNumeric B e where + unsafeDotProduct = defaultUnsafeDotProduct + {-# INLINE unsafeDotProduct #-} + powerSumArray = defaultPowerSumArray + {-# INLINE powerSumArray #-} + foldArray = defaultFoldArray + {-# INLINE foldArray #-} + +instance Num e => Numeric B e where + unsafeLiftArray = defaultUnsafeLiftArray + {-# INLINE unsafeLiftArray #-} + unsafeLiftArray2 = defaultUnsafeLiftArray2 + {-# INLINE unsafeLiftArray2 #-} + ----------------------- -- Boxed Normal Form -- ----------------------- --- | Array representation for Boxed elements. This structure is element and --- spine strict, and elements are always in Normal Form (NF), therefore `NFData` --- instance is required. -data N = N deriving Show +-- | Array representation for Boxed elements. Its elements are always in Normal +-- Form (NF), therefore `NFData` instance is required. +data BN = BN deriving (Show) -newtype instance Array N ix e = NArray { bArray :: Array B ix e } +-- | Type and pattern `N` have been added for backwards compatibility and will be replaced +-- in the future in favor of `BN`. +-- +-- /Deprecated/ - since 1.0.0 +type N = BN -instance (Ragged L ix e, Show e, NFData e) => Show (Array N ix e) where - showsPrec = showsArrayPrec bArray +pattern N :: N +pattern N = BN + +{-# COMPLETE N #-} + +{-# DEPRECATED N "In favor of more consistently named `BN`" #-} + +newtype instance Array BN ix e = BNArray (Array BL ix e) + +newtype instance MArray s BN ix e = MBNArray (MArray s BL ix e) + +instance (Ragged L ix e, Show e, NFData e) => Show (Array BN ix e) where + showsPrec = showsArrayPrec coerce showList = showArrayList -instance (Index ix, NFData e) => NFData (Array N ix e) where - rnf (NArray barr) = barr `deepseqArray` () +-- | /O(1)/ - `BN` is already in normal form +instance NFData (Array BN ix e) where + rnf = (`seq` ()) {-# INLINE rnf #-} -instance (Index ix, NFData e, Eq e) => Eq (Array N ix e) where - (==) = eq (==) +instance (Index ix, NFData e, Eq e) => Eq (Array BN ix e) where + (==) = eqArrays (==) {-# INLINE (==) #-} -instance (Index ix, NFData e, Ord e) => Ord (Array N ix e) where - compare = ord compare +instance (Index ix, NFData e, Ord e) => Ord (Array BN ix e) where + compare = compareArrays compare {-# INLINE compare #-} - -instance (Index ix, NFData e) => Construct N ix e where - setComp c (NArray arr) = NArray (arr {bComp = c}) +instance Strategy BN where + setComp c = coerce (setComp c) {-# INLINE setComp #-} - makeArray !comp !sz f = - unsafePerformIO $ - generateArray - comp - sz - (\ !ix -> - let res = f ix - in res `deepseq` return res) - {-# INLINE makeArray #-} + getComp = blComp . coerce + {-# INLINE getComp #-} + repr = BN -instance (Index ix, NFData e) => Source N ix e where - unsafeLinearIndex (NArray arr) = - INDEX_CHECK("(Source N ix e).unsafeLinearIndex", Sz . totalElem . size, unsafeLinearIndex) arr +instance NFData e => Source BN e where + unsafeLinearIndex (BNArray arr) = unsafeLinearIndex arr {-# INLINE unsafeLinearIndex #-} - - -instance Index ix => Resize N ix where - unsafeResize !sz = NArray . unsafeResize sz . bArray - {-# INLINE unsafeResize #-} - -instance (Index ix, NFData e) => Extract N ix e where - unsafeExtract !sIx !newSz !arr = unsafeExtract sIx newSz (toManifest arr) - {-# INLINE unsafeExtract #-} - - -instance ( NFData e - , Index ix - , Index (Lower ix) - , Elt M ix e ~ Array M (Lower ix) e - , Elt N ix e ~ Array M (Lower ix) e - ) => - OuterSlice N ix e where - unsafeOuterSlice = unsafeOuterSlice . toManifest + unsafeLinearSlice i k (BNArray a) = coerce (unsafeLinearSlice i k a) + {-# INLINE unsafeLinearSlice #-} + unsafeOuterSlice (BNArray a) i = coerce (unsafeOuterSlice a i) {-# INLINE unsafeOuterSlice #-} -instance ( NFData e - , Index ix - , Index (Lower ix) - , Elt M ix e ~ Array M (Lower ix) e - , Elt N ix e ~ Array M (Lower ix) e - ) => - InnerSlice N ix e where - unsafeInnerSlice = unsafeInnerSlice . toManifest - {-# INLINE unsafeInnerSlice #-} - -instance {-# OVERLAPPING #-} NFData e => Slice N Ix1 e where - unsafeSlice arr i _ _ = pure (unsafeLinearIndex arr i) - {-# INLINE unsafeSlice #-} +instance Index ix => Shape BN ix where + maxLinearSize = Just . SafeSz . elemsCount + {-# INLINE maxLinearSize #-} +instance Size BN where + size = blSize . coerce + {-# INLINE size #-} -instance (Index ix, NFData e) => Manifest N ix e where + unsafeResize !sz = coerce . unsafeResize sz . coerce + {-# INLINE unsafeResize #-} - unsafeLinearIndexM (NArray arr) = - INDEX_CHECK("(Manifest N ix e).unsafeLinearIndexM", Sz . totalElem . size, unsafeLinearIndexM) arr +instance NFData e => Manifest BN e where + unsafeLinearIndexM arr = unsafeLinearIndexM (coerce arr) {-# INLINE unsafeLinearIndexM #-} + sizeOfMArray = sizeOfMArray . coerce + {-# INLINE sizeOfMArray #-} -instance (Index ix, NFData e) => Mutable N ix e where - newtype MArray s N ix e = MNArray { bmArray :: MArray s B ix e } + unsafeResizeMArray sz = coerce . unsafeResizeMArray sz . coerce + {-# INLINE unsafeResizeMArray #-} - msize = msize . bmArray - {-# INLINE msize #-} + unsafeLinearSliceMArray i k = MBNArray . unsafeLinearSliceMArray i k . coerce + {-# INLINE unsafeLinearSliceMArray #-} - unsafeThaw (NArray arr) = MNArray <$> unsafeThaw arr + unsafeThaw arr = MBNArray <$> unsafeThaw (coerce arr) {-# INLINE unsafeThaw #-} - unsafeFreeze comp (MNArray marr) = NArray <$> unsafeFreeze comp marr + unsafeFreeze comp marr = BNArray <$> unsafeFreeze comp (coerce marr) {-# INLINE unsafeFreeze #-} - unsafeNew sz = MNArray <$> unsafeNew sz + unsafeNew sz = MBNArray <$> unsafeNew sz {-# INLINE unsafeNew #-} initialize _ = return () {-# INLINE initialize #-} - unsafeLinearRead (MNArray ma) = - INDEX_CHECK("(Mutable N ix e).unsafeLinearRead", Sz . totalElem . msize, unsafeLinearRead) ma + newMArray sz e = e `deepseq` (MBNArray <$> newMArray sz e) + {-# INLINE newMArray #-} + + unsafeLinearRead ma = unsafeLinearRead (coerce ma) {-# INLINE unsafeLinearRead #-} - unsafeLinearWrite (MNArray ma) i e = e `deepseq` - INDEX_CHECK("(Mutable N ix e).unsafeLinearWrite", Sz . totalElem . msize, unsafeLinearWrite) ma i e + unsafeLinearWrite ma i e = e `deepseq` unsafeLinearWrite (coerce ma) i e {-# INLINE unsafeLinearWrite #-} -instance (Index ix, NFData e) => Load N ix e where - type R N = M - size = bSize . bArray - {-# INLINE size #-} - getComp = bComp . bArray - {-# INLINE getComp #-} - loadArrayM !scheduler !arr = splitLinearlyWith_ scheduler (elemsCount arr) (unsafeLinearIndex arr) - {-# INLINE loadArrayM #-} +instance (Index ix, NFData e) => Load BN ix e where + makeArray comp sz f = compute (makeArray comp sz f :: Array D ix e) + {-# INLINE makeArray #-} + makeArrayLinear !comp !sz f = unsafePerformIO $ generateArrayLinear comp sz (pure . f) + {-# INLINE makeArrayLinear #-} + replicate comp sz e = runST (newMArray sz e >>= unsafeFreeze comp) + {-# INLINE replicate #-} + iterArrayLinearST_ !scheduler !arr = + splitLinearlyWith_ scheduler (elemsCount arr) (unsafeLinearIndex arr) + {-# INLINE iterArrayLinearST_ #-} -instance (Index ix, NFData e) => StrideLoad N ix e +instance (Index ix, NFData e) => StrideLoad BN ix e -instance Index ix => Stream N ix e where +instance (Index ix, NFData e) => Stream BN ix e where toStream = toStream . coerce {-# INLINE toStream #-} - + toStreamIx = toStreamIx . coerce + {-# INLINE toStreamIx #-} -instance ( NFData e - , IsList (Array L ix e) - , Nested LN ix e - , Nested L ix e - , Ragged L ix e - ) => - IsList (Array N ix e) where - type Item (Array N ix e) = Item (Array L ix e) +instance (NFData e, IsList (Array L ix e), Ragged L ix e) => IsList (Array BN ix e) where + type Item (Array BN ix e) = Item (Array L ix e) fromList = L.fromLists' Seq {-# INLINE fromList #-} toList = GHC.toList . toListArray {-# INLINE toList #-} +instance (NFData e, Num e) => FoldNumeric BN e where + unsafeDotProduct = defaultUnsafeDotProduct + {-# INLINE unsafeDotProduct #-} + powerSumArray = defaultPowerSumArray + {-# INLINE powerSumArray #-} + foldArray = defaultFoldArray + {-# INLINE foldArray #-} +instance (NFData e, Num e) => Numeric BN e where + unsafeLiftArray = defaultUnsafeLiftArray + {-# INLINE unsafeLiftArray #-} + unsafeLiftArray2 = defaultUnsafeLiftArray2 + {-# INLINE unsafeLiftArray2 #-} + ---------------------- -- Helper functions -- ---------------------- @@ -397,151 +598,282 @@ uninitialized :: a uninitialized = throw Uninitialized - --- -- | /O(1)/ - Unwrap a fully evaluated boxed array. --- -- --- -- @since 0.2.1 --- unwrapNormalFormArray :: Array N ix e -> Array B ix e --- unwrapNormalFormArray = bArray --- {-# INLINE unwrapNormalFormArray #-} - --- -- | /O(1)/ - Unwrap a fully evaluated mutable boxed array. --- -- --- -- @since 0.2.1 --- unwrapNormalFormMutableArray :: MArray s N ix e -> MArray s B ix e --- unwrapNormalFormMutableArray (MNArray marr) = marr --- {-# INLINE unwrapNormalFormMutableArray #-} - --------------------- -- WHNF conversion -- --------------------- --- | /O(1)/ - Unwrap boxed array. +-- | /O(1)/ - Unwrap boxed array. This will discard any possible slicing that has been +-- applied to the array. -- -- @since 0.2.1 unwrapArray :: Array B ix e -> A.Array e -unwrapArray = bData +unwrapArray = blData . coerce {-# INLINE unwrapArray #-} -- | /O(n)/ - Wrap a boxed array and evaluate all elements to a WHNF. -- -- @since 0.2.1 -evalArray :: - Comp -- ^ Computation strategy - -> A.Array e -- ^ Lazy boxed array from @primitive@ package. - -> Array B Ix1 e -evalArray = fromArraySeq (\a -> a `seqArray` a) +evalArray + :: Comp + -- ^ Computation strategy + -> A.Array e + -- ^ Lazy boxed array from @primitive@ package. + -> Vector B e +evalArray comp a = evalLazyArray $ setComp comp $ wrapLazyArray a {-# INLINE evalArray #-} --- | /O(1)/ - Unwrap mutable boxed array. +-- | /O(1)/ - Unwrap boxed array. This will discard any possible slicing that has been +-- applied to the array. -- +-- @since 0.6.0 +unwrapLazyArray :: Array BL ix e -> A.Array e +unwrapLazyArray = blData +{-# INLINE unwrapLazyArray #-} + +-- | /O(1)/ - Wrap a boxed array. +-- +-- @since 0.6.0 +wrapLazyArray :: A.Array e -> Vector BL e +wrapLazyArray a = BLArray Seq (SafeSz (A.sizeofArray a)) 0 a +{-# INLINE wrapLazyArray #-} + +-- | /O(1)/ - Cast a strict boxed array into a lazy boxed array. +-- +-- @since 0.6.0 +toLazyArray :: Array B ix e -> Array BL ix e +toLazyArray = coerce +{-# INLINE toLazyArray #-} + +-- | /O(n)/ - Evaluate all elements of a boxed lazy array to weak head normal form +-- +-- @since 0.6.0 +evalLazyArray :: Index ix => Array BL ix e -> Array B ix e +evalLazyArray arr = arr `seqArray` BArray arr +{-# INLINE evalLazyArray #-} + +-- | /O(n)/ - Evaluate all elements of a boxed lazy array to normal form +-- +-- @since 0.6.0 +forceLazyArray :: (NFData e, Index ix) => Array BL ix e -> Array N ix e +forceLazyArray arr = arr `deepseqArray` BNArray arr +{-# INLINE forceLazyArray #-} + +-- | /O(1)/ - Unwrap mutable boxed array. This will discard any possible slicing that has been +-- applied to the array. +-- -- @since 0.2.1 unwrapMutableArray :: MArray s B ix e -> A.MutableArray s e -unwrapMutableArray (MBArray _ marr) = marr +unwrapMutableArray (MBArray (MBLArray _ _ marr)) = marr {-# INLINE unwrapMutableArray #-} +-- | /O(1)/ - Unwrap mutable boxed lazy array. This will discard any possible slicing that has been +-- applied to the array. +-- +-- @since 0.6.0 +unwrapMutableLazyArray :: MArray s BL ix e -> A.MutableArray s e +unwrapMutableLazyArray (MBLArray _ _ marr) = marr +{-# INLINE unwrapMutableLazyArray #-} -- | /O(n)/ - Wrap mutable boxed array and evaluate all elements to WHNF. -- -- @since 0.2.1 -evalMutableArray :: - PrimMonad m - => A.MutableArray (PrimState m) e -- ^ Mutable array that will get wrapped +evalMutableArray + :: PrimMonad m + => A.MutableArray (PrimState m) e + -- ^ Mutable array that will get wrapped -> m (MArray (PrimState m) B Ix1 e) -evalMutableArray = fromMutableArraySeq seq +evalMutableArray = fmap MBArray . fromMutableArraySeq seq {-# INLINE evalMutableArray #-} ------------------- -- NF conversion -- ------------------- --- | /O(1)/ - Unwrap a fully evaluated boxed array. +-- | /O(1)/ - Unwrap a fully evaluated boxed array. This will discard any possible slicing +-- that has been applied to the array. -- -- @since 0.2.1 unwrapNormalFormArray :: Array N ix e -> A.Array e -unwrapNormalFormArray = bData . bArray +unwrapNormalFormArray = blData . coerce {-# INLINE unwrapNormalFormArray #-} -- | /O(n)/ - Wrap a boxed array and evaluate all elements to a Normal Form (NF). -- -- @since 0.2.1 -evalNormalFormArray :: - NFData e - => Comp -- ^ Computation strategy - -> A.Array e -- ^ Lazy boxed array +evalNormalFormArray + :: NFData e + => Comp + -- ^ Computation strategy + -> A.Array e + -- ^ Lazy boxed array -> Array N Ix1 e -evalNormalFormArray = fromArraySeq (\a -> a `deepseqArray` NArray a) +evalNormalFormArray comp = forceLazyArray . setComp comp . wrapLazyArray {-# INLINE evalNormalFormArray #-} - --- | /O(1)/ - Unwrap a fully evaluated mutable boxed array. +-- | /O(1)/ - Unwrap a fully evaluated mutable boxed array. This will discard any possible +-- slicing that has been applied to the array. -- -- @since 0.2.1 unwrapNormalFormMutableArray :: MArray s N ix e -> A.MutableArray s e -unwrapNormalFormMutableArray (MNArray (MBArray _ marr)) = marr +unwrapNormalFormMutableArray = unwrapMutableLazyArray . coerce {-# INLINE unwrapNormalFormMutableArray #-} - -- | /O(n)/ - Wrap mutable boxed array and evaluate all elements to NF. -- -- @since 0.2.1 -evalNormalFormMutableArray :: - (PrimMonad m, NFData e) +evalNormalFormMutableArray + :: (PrimMonad m, NFData e) => A.MutableArray (PrimState m) e -> m (MArray (PrimState m) N Ix1 e) -evalNormalFormMutableArray marr = MNArray <$> fromMutableArraySeq deepseq marr +evalNormalFormMutableArray marr = MBNArray <$> fromMutableArraySeq deepseq marr {-# INLINE evalNormalFormMutableArray #-} - ---------------------- -- Helper functions -- ---------------------- -fromMutableArraySeq :: - PrimMonad m +fromMutableArraySeq + :: PrimMonad m => (e -> m () -> m a) -> A.MutableArray (PrimState m) e - -> m (MArray (PrimState m) B Ix1 e) -fromMutableArraySeq with mbarr = do - let !sz = sizeofMutableArray mbarr - loopM_ 0 (< sz) (+ 1) (A.readArray mbarr >=> (`with` return ())) - return $! MBArray (Sz sz) mbarr + -> m (MArray (PrimState m) BL Ix1 e) +fromMutableArraySeq with ma = do + let !sz = A.sizeofMutableArray ma + loopA_ 0 (< sz) (+ 1) (A.readArray ma >=> (`with` return ())) + return $! MBLArray (SafeSz sz) 0 ma {-# INLINE fromMutableArraySeq #-} -fromArraySeq :: - (Array B Ix1 e -> a) - -> Comp - -> A.Array e - -> a -fromArraySeq with comp barr = with (BArray comp (Sz (sizeofArray barr)) barr) -{-# INLINE fromArraySeq #-} - - -seqArray :: Index ix => Array B ix a -> t -> t +seqArray :: Index ix => Array BL ix a -> t -> t seqArray !arr t = foldlInternal (flip seq) () (flip seq) () arr `seq` t {-# INLINE seqArray #-} - -deepseqArray :: (NFData a, Index ix) => Array B ix a -> t -> t +deepseqArray :: (NFData a, Index ix) => Array BL ix a -> t -> t deepseqArray !arr t = foldlInternal (flip deepseq) () (flip seq) () arr `seq` t {-# INLINE deepseqArray #-} +-- | /O(1)/ - Converts array from `N` to `B` representation. +-- +-- @since 0.5.0 +unwrapNormalForm :: Array N ix e -> Array B ix e +unwrapNormalForm = coerce +{-# INLINE unwrapNormalForm #-} --- | Helper function that converts a boxed `A.Array` into a `VB.Vector`. Supplied total number of --- elements is assumed to be the same in the array as provided by the size. -castArrayToVector :: A.Array a -> VB.Vector a -castArrayToVector arr = runST $ do - marr <- A.unsafeThawArray arr - VB.unsafeFreeze $ VB.MVector 0 (sizeofArray arr) marr -{-# INLINE castArrayToVector #-} +-- | /O(n)/ - Compute all elements of a boxed array to NF (normal form) +-- +-- @since 0.5.0 +evalNormalForm :: (Index ix, NFData e) => Array B ix e -> Array N ix e +evalNormalForm (BArray arr) = arr `deepseqArray` BNArray arr +{-# INLINE evalNormalForm #-} +{- FOURMOLU_DISABLE -} +-- | /O(1)/ - Converts a boxed `Array` into a `VB.Vector` without touching any +-- elements. +-- +-- @since 0.5.0 +{-# INLINE toBoxedVector #-} +toBoxedVector :: Index ix => Array BL ix a -> VB.Vector a +toBoxedVector BLArray{blOffset = off, blSize = sz, blData = arr } = +#if MIN_VERSION_vector(0,13,0) + VB.unsafeFromArraySlice arr off (totalElem sz) +#elif MIN_VERSION_vector(0,12,2) + VB.unsafeTake (totalElem sz) (VB.unsafeDrop off (VB.fromArray arr)) +#else + fromVectorCast $ VectorCast off (totalElem sz) arr --- | Cast a Boxed Vector into an Array, but only if it wasn't previously sliced. -castVectorToArray :: VB.Vector a -> Maybe (A.Array a) -castVectorToArray v = +fromVectorCast :: VectorCast a -> VB.Vector a +fromVectorCast = unsafeCoerce +#endif +{- FOURMOLU_ENABLE -} + +-- | /O(1)/ - Converts a boxed `MArray` into a `MVB.MVector`. +-- +-- @since 0.5.0 +toBoxedMVector :: Index ix => MArray s BL ix a -> MVB.MVector s a +toBoxedMVector (MBLArray sz o marr) = MVB.MVector o (totalElem sz) marr +{-# INLINE toBoxedMVector #-} + +-- | /O(n)/ - Convert a boxed vector and evaluate all elements to WHNF. Computation +-- strategy will be respected during evaluation +-- +-- @since 0.5.0 +evalBoxedVector :: Comp -> VB.Vector a -> Array B Ix1 a +evalBoxedVector comp = evalLazyArray . setComp comp . fromBoxedVector +{-# INLINE evalBoxedVector #-} + +-- | /O(n)/ - Convert mutable boxed vector and evaluate all elements to WHNF +-- sequentially. Both keep pointing to the same memory +-- +-- @since 0.5.0 +evalBoxedMVector :: PrimMonad m => MVB.MVector (PrimState m) a -> m (MArray (PrimState m) B Ix1 a) +evalBoxedMVector (MVB.MVector o k ma) = + let marr = MBArray (MBLArray (SafeSz k) o ma) + in marr <$ loopA_ o (< k) (+ 1) (A.readArray ma >=> (`seq` pure ())) +{-# INLINE evalBoxedMVector #-} + +-- | /O(1)/ - Cast a boxed vector without touching any elements. +-- +-- @since 0.6.0 +fromBoxedVector :: VB.Vector a -> Vector BL a +{-# INLINE fromBoxedVector #-} +fromBoxedVector v = + BLArray{blComp = Seq, blSize = SafeSz n, blOffset = offset, blData = arr} + where +#if MIN_VERSION_vector(0,13,0) + (arr, offset, n) = VB.toArraySlice v +#else + VectorCast offset n arr = toVectorCast v + +-- This internal type is needed to get into the internals of a boxed vector, +-- since it is not possible until vector-0.13 version. +data VectorCast a = + VectorCast {-# UNPACK #-}!Int {-# UNPACK #-}!Int {-# UNPACK #-}!(A.Array a) + +toVectorCast :: VB.Vector a -> VectorCast a +toVectorCast = unsafeCoerce +#endif + +-- | /O(1)/ - Convert mutable boxed vector to a lazy mutable boxed array. Both keep +-- pointing to the same memory +-- +-- @since 0.6.0 +fromBoxedMVector :: MVB.MVector s a -> MArray s BL Ix1 a +fromBoxedMVector (MVB.MVector o k ma) = MBLArray (SafeSz k) o ma +{-# INLINE fromBoxedMVector #-} + +-- | /O(1)/ - Cast a boxed lazy array. It is unsafe because it can violate the invariant +-- that all elements of `N` array are in NF. +-- +-- @since 0.6.0 +coerceNormalBoxedArray :: Array BL ix e -> Array N ix e +coerceNormalBoxedArray = coerce +{-# INLINE coerceNormalBoxedArray #-} + +-- | /O(1)/ - Cast a boxed lazy array. It is unsafe because it can violate the invariant +-- that all elements of `B` array are in WHNF. +-- +-- @since 0.6.0 +coerceBoxedArray :: Array BL ix e -> Array B ix e +coerceBoxedArray = coerce +{-# INLINE coerceBoxedArray #-} + +-- | /O(n)/ - Convert mutable boxed vector and evaluate all elements to WHNF +-- sequentially. Both keep pointing to the same memory +-- +-- @since 0.5.0 +evalNormalBoxedMVector + :: (NFData a, PrimMonad m) => MVB.MVector (PrimState m) a -> m (MArray (PrimState m) N Ix1 a) +evalNormalBoxedMVector (MVB.MVector o k ma) = + let marr = MBNArray (MBLArray (SafeSz k) o ma) + in marr <$ loopA_ o (< k) (+ 1) (A.readArray ma >=> pure . rnf) +{-# INLINE evalNormalBoxedMVector #-} + +-- | /O(n)/ - Convert a boxed vector and evaluate all elements to WHNF. Computation +-- strategy will be respected during evaluation +-- +-- @since 0.5.0 +evalNormalBoxedVector :: NFData a => Comp -> VB.Vector a -> Array N Ix1 a +evalNormalBoxedVector comp v = runST $ do - VB.MVector start end marr <- VB.unsafeThaw v - if start == 0 && end == sizeofMutableArray marr - then Just <$> A.unsafeFreezeArray marr - else return Nothing -{-# INLINE castVectorToArray #-} + MVB.MVector o k ma <- VB.unsafeThaw v + forceLazyArray <$> unsafeFreeze comp (MBLArray (SafeSz k) o ma) +{-# INLINE evalNormalBoxedVector #-} 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 @@ -1,211 +1,68 @@ {-# LANGUAGE BangPatterns #-} -{-# LANGUAGE CPP #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MagicHash #-} +{-# LANGUAGE MonoLocalBinds #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE TypeOperators #-} {-# LANGUAGE UndecidableInstances #-} + -- | -- Module : Data.Massiv.Array.Manifest.Internal --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Manifest.Internal - ( M - , Manifest(..) - , Array(..) - , toManifest - , compute - , computeS - , computeIO - , computePrimM - , computeAs - , computeProxy - , computeSource - , computeWithStride - , computeWithStrideAs - , clone - , convert - , convertAs - , convertProxy - , gcastArr - , fromRaggedArrayM - , fromRaggedArray' - , sizeofArray - , sizeofMutableArray - , iterateUntil - , iterateUntilM - ) where +module Data.Massiv.Array.Manifest.Internal ( + Manifest (..), + Array (..), + flattenMArray, + compute, + computeS, + computeP, + computeIO, + computePrimM, + computeAs, + computeProxy, + computeSource, + computeWithStride, + computeWithStrideAs, + clone, + convert, + convertAs, + convertProxy, + gcastArr, + fromRaggedArrayM, + fromRaggedArray', + unsafeLoadIntoS, + unsafeLoadIntoM, + iterateUntil, + iterateUntilM, +) where +import Control.DeepSeq import Control.Exception (try) +import Control.Monad.Primitive import Control.Monad.ST import Control.Scheduler -import qualified Data.Foldable as F (Foldable(..)) import Data.Massiv.Array.Delayed.Pull import Data.Massiv.Array.Mutable -import Data.Massiv.Array.Ops.Fold.Internal -import Data.Massiv.Array.Manifest.Vector.Stream as S (steps) +import Data.Massiv.Array.Mutable.Internal (unsafeCreateArray_) import Data.Massiv.Core.Common import Data.Massiv.Core.List import Data.Maybe (fromMaybe) import Data.Typeable -import GHC.Base hiding (ord) import System.IO.Unsafe (unsafePerformIO) -#if MIN_VERSION_primitive(0,6,2) -import Data.Primitive.Array (sizeofArray, sizeofMutableArray) - -#else -import qualified Data.Primitive.Array as A (Array(..), MutableArray(..)) -import GHC.Exts (sizeofArray#, sizeofMutableArray#) - -sizeofArray :: A.Array a -> Int -sizeofArray (A.Array a) = I# (sizeofArray# a) -{-# INLINE sizeofArray #-} - -sizeofMutableArray :: A.MutableArray s a -> Int -sizeofMutableArray (A.MutableArray ma) = I# (sizeofMutableArray# ma) -{-# INLINE sizeofMutableArray #-} -#endif - - --- | General Manifest representation -data M - -data instance Array M ix e = MArray { mComp :: !Comp - , mSize :: !(Sz ix) - , mLinearIndex :: Int -> e } - -instance (Ragged L ix e, Show e) => Show (Array M ix e) where - showsPrec = showsArrayPrec id - showList = showArrayList - - -instance (Eq e, Index ix) => Eq (Array M ix e) where - (==) = eq (==) - {-# INLINE (==) #-} - -instance (Ord e, Index ix) => Ord (Array M ix e) where - compare = ord compare - {-# INLINE compare #-} - - --- | /O(1)/ - Conversion of `Manifest` arrays to `M` representation. -toManifest :: Manifest r ix e => Array r ix e -> Array M ix e -toManifest !arr = MArray (getComp arr) (size arr) (unsafeLinearIndexM arr) -{-# INLINE toManifest #-} - - --- | Row-major sequentia folding over a Manifest array. -instance Index ix => Foldable (Array M ix) where - fold = fold - {-# INLINE fold #-} - foldMap = foldMono - {-# INLINE foldMap #-} - foldl = lazyFoldlS - {-# INLINE foldl #-} - foldl' = foldlS - {-# INLINE foldl' #-} - foldr = foldrFB - {-# INLINE foldr #-} - foldr' = foldrS - {-# INLINE foldr' #-} - null (MArray _ sz _) = totalElem sz == 0 - {-# INLINE null #-} - length = totalElem . size - {-# INLINE length #-} - toList arr = build (\ c n -> foldrFB c n arr) - {-# INLINE toList #-} - - - -instance Index ix => Source M ix e where - unsafeLinearIndex = mLinearIndex - {-# INLINE unsafeLinearIndex #-} - - -instance Index ix => Manifest M ix e where - - unsafeLinearIndexM = mLinearIndex - {-# INLINE unsafeLinearIndexM #-} - - -instance Index ix => Resize M ix where - unsafeResize !sz !arr = arr { mSize = sz } - {-# INLINE unsafeResize #-} - -instance Index ix => Extract M ix e where - unsafeExtract !sIx !newSz !arr = - MArray (getComp arr) newSz $ \ i -> - unsafeIndex arr (liftIndex2 (+) (fromLinearIndex newSz i) sIx) - {-# INLINE unsafeExtract #-} - - - -instance {-# OVERLAPPING #-} Slice M Ix1 e where - unsafeSlice arr i _ _ = pure (unsafeLinearIndex arr i) - {-# INLINE unsafeSlice #-} - -instance ( Index ix - , Index (Lower ix) - , Elt M ix e ~ Array M (Lower ix) e - ) => - Slice M ix e where - unsafeSlice arr start cutSz dim = do - (_, newSz) <- pullOutSzM cutSz dim - return $ unsafeResize newSz (unsafeExtract start cutSz arr) - {-# INLINE unsafeSlice #-} - -instance {-# OVERLAPPING #-} OuterSlice M Ix1 e where - unsafeOuterSlice !arr = unsafeIndex arr - {-# INLINE unsafeOuterSlice #-} - -instance (Elt M ix e ~ Array M (Lower ix) e, Index ix, Index (Lower ix)) => OuterSlice M ix e where - unsafeOuterSlice !arr !i = - MArray (getComp arr) (snd (unconsSz (size arr))) (unsafeLinearIndex arr . (+ kStart)) - where - !kStart = toLinearIndex (size arr) (consDim i (zeroIndex :: Lower ix)) - {-# INLINE unsafeOuterSlice #-} - -instance {-# OVERLAPPING #-} InnerSlice M Ix1 e where - unsafeInnerSlice !arr _ = unsafeIndex arr - {-# INLINE unsafeInnerSlice #-} - -instance (Elt M ix e ~ Array M (Lower ix) e, Index ix, Index (Lower ix)) => InnerSlice M ix e where - unsafeInnerSlice !arr (szL, m) !i = - MArray (getComp arr) szL (\k -> unsafeLinearIndex arr (k * unSz m + kStart)) - where - !kStart = toLinearIndex (size arr) (snocDim (zeroIndex :: Lower ix) i) - {-# INLINE unsafeInnerSlice #-} - - -instance Index ix => Load M ix e where - size = mSize - {-# INLINE size #-} - getComp = mComp - {-# INLINE getComp #-} - loadArrayM scheduler (MArray _ sz f) = splitLinearlyWith_ scheduler (totalElem sz) f - {-# INLINE loadArrayM #-} - -instance Index ix => StrideLoad M ix e - -instance Index ix => Stream M ix e where - toStream = S.steps - {-# INLINE toStream #-} - - -- | Ensure that Array is computed, i.e. represented with concrete elements in memory, hence is the -- `Mutable` type class restriction. Use `setComp` if you'd like to change computation strategy -- before calling @compute@ -- -- @since 0.1.0 -compute :: forall r ix e r' . (Mutable r ix e, Load r' ix e) => Array r' ix e -> Array r ix e +compute :: forall r ix e r'. (Manifest r e, Load r' ix e) => Array r' ix e -> Array r ix e compute !arr = unsafePerformIO $ computeIO arr {-# INLINE compute #-} @@ -213,18 +70,32 @@ -- the same as `computePrimM`, but executed in `ST`, thus pure. -- -- @since 0.1.0 -computeS :: forall r ix e r' . (Mutable r ix e, Load r' ix e) => Array r' ix e -> Array r ix e +computeS :: forall r ix e r'. (Manifest r e, Load r' ix e) => Array r' ix e -> Array r ix e computeS !arr = runST $ computePrimM arr {-# INLINE computeS #-} +-- | Compute array in parallel using all cores disregarding predefined computation +-- strategy. Computation stategy of the resulting array will match the source, despite +-- that it is diregarded. +-- +-- @since 0.5.4 +computeP + :: forall r ix e r' + . (Manifest r e, Load r' ix e) + => Array r' ix e + -> Array r ix e +computeP arr = setComp (getComp arr) $ compute (setComp Par arr) +{-# INLINE computeP #-} + -- | Very similar to `compute`, but computes an array inside the `IO` monad. Despite being -- deterministic and referentially transparent, because this is an `IO` action it -- can be very useful for enforcing the order of evaluation. Should be a prefered way of -- computing an array during benchmarking. -- -- @since 0.4.5 -computeIO :: - forall r ix e r' m. (Mutable r ix e, Load r' ix e, MonadIO m) +computeIO + :: forall r ix e r' m + . (Manifest r e, Load r' ix e, MonadIO m) => Array r' ix e -> m (Array r ix e) computeIO arr = liftIO (loadArray arr >>= unsafeFreeze (getComp arr)) @@ -234,14 +105,14 @@ -- strategy. -- -- @since 0.4.5 -computePrimM :: - forall r ix e r' m. (Mutable r ix e, Load r' ix e, PrimMonad m) +computePrimM + :: forall r ix e r' m + . (Manifest r e, Load r' ix e, PrimMonad m) => Array r' ix e -> m (Array r ix e) computePrimM arr = loadArrayS arr >>= unsafeFreeze (getComp arr) {-# INLINE computePrimM #-} - -- | Just as `compute`, but let's you supply resulting representation type as an argument. -- -- ====__Examples__ @@ -250,12 +121,10 @@ -- >>> computeAs P $ range Seq (Ix1 0) 10 -- Array P Seq (Sz1 10) -- [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9 ] --- -computeAs :: (Mutable r ix e, Load r' ix e) => r -> Array r' ix e -> Array r ix e +computeAs :: (Manifest r e, Load r' ix e) => r -> Array r' ix e -> Array r ix e computeAs _ = compute {-# INLINE computeAs #-} - -- | Same as `compute` and `computeAs`, but let's you supply resulting representation type as a proxy -- argument. -- @@ -271,122 +140,157 @@ -- [ 0, 1, 4, 9, 16, 25, 36, 49, 64, 81 ] -- -- @since 0.1.1 -computeProxy :: (Mutable r ix e, Load r' ix e) => proxy r -> Array r' ix e -> Array r ix e +computeProxy :: (Manifest r e, Load r' ix e) => proxy r -> Array r' ix e -> Array r ix e computeProxy _ = compute {-# INLINE computeProxy #-} - -- | This is just like `convert`, but restricted to `Source` arrays. Will be a noop if -- resulting type is the same as the input. -- -- @since 0.1.0 -computeSource :: forall r ix e r' . (Mutable r ix e, Source r' ix e) - => Array r' ix e -> Array r ix e -computeSource arr = maybe (compute arr) (\Refl -> arr) (eqT :: Maybe (r' :~: r)) +computeSource + :: forall r ix e r' + . (Manifest r e, Source r' e, Index ix) + => Array r' ix e + -> Array r ix e +computeSource arr = maybe (compute $ delay arr) (\Refl -> arr) (eqT :: Maybe (r' :~: r)) {-# INLINE computeSource #-} - -- | /O(n)/ - Make an exact immutable copy of an Array. -- -- @since 0.1.0 -clone :: Mutable r ix e => Array r ix e -> Array r ix e +clone :: (Manifest r e, Index ix) => Array r ix e -> Array r ix e clone arr = unsafePerformIO $ thaw arr >>= unsafeFreeze (getComp arr) {-# INLINE clone #-} - -- | /O(1)/ - Cast over Array representation -gcastArr :: forall r ix e r' . (Typeable r, Typeable r') - => Array r' ix e -> Maybe (Array r ix e) +gcastArr + :: forall r ix e r' + . (Typeable r, Typeable r') + => Array r' ix e + -> Maybe (Array r ix e) gcastArr arr = fmap (\Refl -> arr) (eqT :: Maybe (r :~: r')) - -- | /O(n)/ - conversion between array types. A full copy will occur, unless when the source and -- result arrays are of the same representation, in which case it is an /O(1)/ operation. -- -- @since 0.1.0 -convert :: forall r ix e r' . (Mutable r ix e, Load r' ix e) - => Array r' ix e -> Array r ix e +convert + :: forall r ix e r' + . (Manifest r e, Load r' ix e) + => Array r' ix e + -> Array r ix e convert arr = fromMaybe (compute arr) (gcastArr arr) {-# INLINE convert #-} -- | Same as `convert`, but let's you supply resulting representation type as an argument. -- -- @since 0.1.0 -convertAs :: (Mutable r ix e, Load r' ix e) - => r -> Array r' ix e -> Array r ix e +convertAs + :: (Manifest r e, Load r' ix e) + => r + -> Array r' ix e + -> Array r ix e convertAs _ = convert {-# INLINE convertAs #-} - -- | Same as `convert` and `convertAs`, but let's you supply resulting representation type as a -- proxy argument. -- -- @since 0.1.1 -convertProxy :: (Mutable r ix e, Load r' ix e) - => proxy r -> Array r' ix e -> Array r ix e +convertProxy + :: (Manifest r e, Load r' ix e) + => proxy r + -> Array r' ix e + -> Array r ix e convertProxy _ = convert {-# INLINE convertProxy #-} - -- | Convert a ragged array into a common array with rectangular shape. Throws `ShapeException` -- whenever supplied ragged array does not have a rectangular shape. -- -- @since 0.4.0 -fromRaggedArrayM :: - forall r ix e r' m . (Mutable r ix e, Ragged r' ix e, Load r' ix e, MonadThrow m) +fromRaggedArrayM + :: forall r ix e r' m + . (Manifest r e, Ragged r' ix e, MonadThrow m) => Array r' ix e -> m (Array r ix e) fromRaggedArrayM arr = - let sz = edgeSize arr + let sz = outerSize arr in either (\(e :: ShapeException) -> throwM e) pure $ - unsafePerformIO $ do - marr <- unsafeNew sz - traverse (\_ -> unsafeFreeze (getComp arr) marr) =<< - try (withScheduler_ (getComp arr) $ \scheduler -> - loadRagged (scheduleWork scheduler) (unsafeLinearWrite marr) 0 (totalElem sz) sz arr) + unsafePerformIO $ do + marr <- unsafeNew sz + traverse (\_ -> unsafeFreeze (getComp arr) marr) + =<< try + ( withMassivScheduler_ (getComp arr) $ \scheduler -> + stToIO $ loadRaggedST scheduler arr (unsafeLinearWrite marr) 0 (totalElem sz) sz + ) {-# INLINE fromRaggedArrayM #-} - --- | Same as `fromRaggedArrayM`, but will throw a pure exception if its shape is not +-- | Same as `fromRaggedArrayM`, but will throw an impure exception if its shape is not -- rectangular. -- -- @since 0.1.1 -fromRaggedArray' :: - forall r ix e r'. (Mutable r ix e, Load r' ix e, Ragged r' ix e) +fromRaggedArray' + :: forall r ix e r' + . (HasCallStack, Manifest r e, Ragged r' ix e) => Array r' ix e -> Array r ix e -fromRaggedArray' arr = either throw id $ fromRaggedArrayM arr +fromRaggedArray' = throwEither . fromRaggedArrayM {-# INLINE fromRaggedArray' #-} - -- | Same as `compute`, but with `Stride`. -- --- /O(n div k)/ - Where @n@ is numer of elements in the source array and @k@ is number of elemts in --- the stride. +-- /O(n div k)/ - Where @n@ is number of elements in the source array and @k@ is number of +-- elements in the stride. -- -- @since 0.3.0 -computeWithStride :: - forall r ix e r'. (Mutable r ix e, StrideLoad r' ix e) +computeWithStride + :: forall r ix e r' + . (Manifest r e, StrideLoad r' ix e) => Stride ix -> Array r' ix e -> Array r ix e computeWithStride stride !arr = unsafePerformIO $ do - let !sz = strideSize stride (size arr) - createArray_ (getComp arr) sz $ \scheduler marr -> - loadArrayWithStrideM scheduler stride sz arr (unsafeLinearWrite marr) + let !sz = strideSize stride (outerSize arr) + unsafeCreateArray_ (getComp arr) sz $ \scheduler marr -> + stToIO $ iterArrayLinearWithStrideST_ scheduler stride sz arr (unsafeLinearWrite marr) {-# INLINE computeWithStride #-} - -- | Same as `computeWithStride`, but with ability to specify resulting array representation. -- -- @since 0.3.0 -computeWithStrideAs :: - (Mutable r ix e, StrideLoad r' ix e) => r -> Stride ix -> Array r' ix e -> Array r ix e +computeWithStrideAs + :: (Manifest r e, StrideLoad r' ix e) => r -> Stride ix -> Array r' ix e -> Array r ix e computeWithStrideAs _ = computeWithStride {-# INLINE computeWithStrideAs #-} +-- | Load into a supplied mutable vector sequentially. Returned array is not +-- necesserally the same vector as the one that was supplied. It will be the +-- same only if it had enough space to load all the elements in. +-- +-- @since 0.5.7 +unsafeLoadIntoS + :: forall r r' ix e m s + . (Load r ix e, Manifest r' e, MonadPrim s m) + => MVector s r' e + -> Array r ix e + -> m (MArray s r' ix e) +unsafeLoadIntoS marr arr = stToPrim $ unsafeLoadIntoS marr arr +{-# INLINE unsafeLoadIntoS #-} +-- | Same as `unsafeLoadIntoS`, but respecting computation strategy. +-- +-- @since 0.5.7 +unsafeLoadIntoM + :: forall r r' ix e m + . (Load r ix e, Manifest r' e, MonadIO m) + => MVector RealWorld r' e + -> Array r ix e + -> m (MArray RealWorld r' ix e) +unsafeLoadIntoM marr arr = liftIO $ unsafeLoadIntoIO marr arr +{-# INLINE unsafeLoadIntoM #-} -- | 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 @@ -397,8 +301,8 @@ -- ====__Example__ -- -- >>> import Data.Massiv.Array --- >>> a = computeAs P $ makeLoadArrayS (Sz2 8 8) (0 :: Int) $ \ w -> w (0 :. 0) 1 >> pure () --- >>> a +-- >>> let arr = computeAs P $ makeLoadArrayS (Sz2 8 8) (0 :: Int) $ \ w -> () <$ w (0 :. 0) 1 +-- >>> arr -- Array P Seq (Sz (8 :. 8)) -- [ [ 1, 0, 0, 0, 0, 0, 0, 0 ] -- , [ 0, 0, 0, 0, 0, 0, 0, 0 ] @@ -409,9 +313,9 @@ -- , [ 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 +-- >>> let nextPascalRow cur above = if cur == 0 then above else cur +-- >>> let 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) arr -- Array P Seq (Sz (8 :. 8)) -- [ [ 1, 0, 0, 0, 0, 0, 0, 0 ] -- , [ 1, 1, 0, 0, 0, 0, 0, 0 ] @@ -424,92 +328,70 @@ -- ] -- -- @since 0.3.6 -iterateUntil :: - (Load r' ix e, Mutable r ix e) +iterateUntil + :: (Load r' ix e, Manifest r e, NFData (Array 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 + -- ^ Initial source array + -> Array r ix e +iterateUntil convergence iteration initArr0 = unsafePerformIO $ do + let loadArr0 = iteration 0 initArr0 + initMVec1 <- unsafeNew (fromMaybe zeroSz (maxLinearSize loadArr0)) + let conv n arr comp marr' = do + arr' <- unsafeFreeze comp marr' + arr' `deepseq` pure (convergence n arr arr', arr') + iterateLoop conv (\n -> pure . iteration n) 0 initArr0 loadArr0 initMVec1 {-# INLINE iterateUntil #-} --- | Monadic version of `iterateUntil` where at each iteration mutable version of an array --- is available. +-- | Monadic version of `iterateUntil` where at each iteration mutable version +-- of an array is available. However it is less efficient then the pure +-- alternative, because an intermediate array must be copied at each +-- iteration. -- -- @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) +iterateUntilM + :: (Load r' ix e, Manifest r e, MonadIO m) + => (Int -> Array r ix e -> MArray RealWorld 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) + -> (Int -> Array r ix e -> m (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 + -- ^ 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 + loadArr0 <- iteration 0 initArr0 + initMVec1 <- liftIO $ unsafeNew (fromMaybe zeroSz (maxLinearSize loadArr0)) + let conv n arr comp marr = (,) <$> convergence n arr marr <*> freeze comp marr + iterateLoop conv iteration 0 initArr0 loadArr0 initMVec1 {-# 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) +iterateLoop + :: (Load r' ix e, Manifest r e, MonadIO m) + => (Int -> Array r ix e -> Comp -> MArray RealWorld r ix e -> m (Bool, Array r ix e)) + -> (Int -> Array r ix e -> m (Array r' ix e)) -> Int -> Array r ix e -> Array r' ix e - -> MArray (PrimState m) r ix e + -> MVector RealWorld r 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' + go n !arr !loadArr !mvec = do + let !comp = getComp loadArr + marr' <- unsafeLoadIntoM mvec loadArr + (shouldStop, arr') <- convergence n arr comp marr' if shouldStop then pure arr' else do - nextMArr <- unsafeThaw arr - go (n + 1) arr' (iteration (n + 1) arr') nextMArr + nextMArr <- liftIO $ unsafeThaw arr + arr'' <- iteration (n + 1) arr' + go (n + 1) arr' arr'' $ flattenMArray nextMArr {-# INLINE iterateLoop #-} diff --git a/src/Data/Massiv/Array/Manifest/List.hs b/src/Data/Massiv/Array/Manifest/List.hs --- a/src/Data/Massiv/Array/Manifest/List.hs +++ b/src/Data/Massiv/Array/Manifest/List.hs @@ -4,47 +4,47 @@ {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} + -- | -- Module : Data.Massiv.Array.Manifest.List --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Manifest.List - ( +module Data.Massiv.Array.Manifest.List ( -- ** List - fromList - , fromListsM - , fromLists' - , toList - , toLists - , toLists2 - , toLists3 - , toLists4 - ) where + fromList, + fromListsM, + fromLists', + toList, + toLists, + toLists2, + toLists3, + toLists4, +) where import Data.Massiv.Array.Manifest.Internal import Data.Massiv.Array.Ops.Fold (foldrInner) import Data.Massiv.Array.Ops.Fold.Internal (foldrFB) import Data.Massiv.Core.Common import Data.Massiv.Core.List -import GHC.Exts (build) +import qualified GHC.Exts as GHC (IsList (..), build) -- | Convert a flat list into a vector -- -- @since 0.1.0 -fromList :: - forall r e. Mutable r Ix1 e - => Comp -- ^ Computation startegy to use - -> [e] -- ^ Flat list - -> Array r Ix1 e +fromList + :: forall r e + . Manifest r e + => Comp + -- ^ Computation startegy to use + -> [e] + -- ^ Flat list + -> Vector r e fromList = fromLists' {-# INLINE fromList #-} - -- | /O(n)/ - Convert a nested list into an array. Nested list must be of a rectangular shape, -- otherwise a runtime error will occur. Also, nestedness must match the rank of resulting array, -- which should be specified through an explicit type signature. @@ -76,27 +76,38 @@ -- , [ [4,5] ] -- ] -- ) --- >>> fromListsM Seq [[[1,2,3]],[[4,5]]] :: Maybe (Array B Ix3 Int) +-- >>> fromListsM Seq [[[1,2,3]],[[4,5]]] :: Maybe (Array B Ix3 Integer) -- Nothing --- >>> fromListsM Seq [[[1,2,3]],[[4,5]]] :: IO (Array B Ix3 Int) --- *** Exception: DimTooShortException: expected (Sz1 3), got (Sz1 2) +-- >>> fromListsM Seq [[[1,2,3]],[[4,5,6],[7,8,9]]] :: IO (Array B Ix3 Integer) +-- *** Exception: DimTooLongException for (Dim 2): expected (Sz1 1), got (Sz1 2) +-- >>> fromListsM Seq [[1,2,3,4],[5,6,7]] :: IO (Matrix B Integer) +-- *** Exception: DimTooShortException for (Dim 1): expected (Sz1 4), got (Sz1 3) -- -- @since 0.3.0 -fromListsM :: forall r ix e m . (Nested LN ix e, Ragged L ix e, Mutable r ix e, MonadThrow m) - => Comp -> [ListItem ix e] -> m (Array r ix e) -fromListsM comp = fromRaggedArrayM . setComp comp . throughNested +fromListsM + :: forall r ix e m + . (Ragged L ix e, Manifest r e, MonadThrow m) + => Comp + -> [ListItem ix e] + -> m (Array r ix e) +fromListsM comp = fromRaggedArrayM . setComp comp . fromListToListArray {-# INLINE fromListsM #-} --- TODO: Figure out QuickCheck properties. Best guess idea so far IMHO is to add it as dependency --- and move Arbitrary instances int the library --- --- prop> fromLists' Seq xs == fromList xs --- --- | Same as `fromListsM`, but will throw a pure error on irregular shaped lists. +fromListToListArray + :: forall ix e + . GHC.IsList (Array L ix e) + => [ListItem ix e] + -> Array L ix e +fromListToListArray = GHC.fromList +{-# INLINE fromListToListArray #-} + +-- | Same as `fromListsM`, but will throw an error on irregular shaped lists. -- --- __Note__: This function is the same as if you would turn on @{-\# LANGUAGE OverloadedLists #-}@ +-- __Note__: This function is the same as if you would turn on @{\-\# LANGUAGE OverloadedLists #-\}@ -- extension. For that reason you can also use `GHC.Exts.fromList`. -- +-- prop> \xs -> fromLists' Seq xs == (fromList Seq xs :: Vector P Int) +-- -- ====__Examples__ -- -- Convert a list of lists into a 2D Array @@ -117,26 +128,18 @@ -- , [ 4, 5, 6 ] -- ] -- --- Example of failure on conversion of an irregular nested list. --- --- >>> fromLists' Seq [[1],[3,4]] :: Array U Ix2 Int --- Array U *** Exception: DimTooLongException --- -- @since 0.1.0 -fromLists' :: forall r ix e . (Nested LN ix e, Ragged L ix e, Mutable r ix e) - => Comp -- ^ Computation startegy to use - -> [ListItem ix e] -- ^ Nested list - -> Array r ix e -fromLists' comp = fromRaggedArray' . setComp comp . throughNested +fromLists' + :: forall r ix e + . (HasCallStack, Ragged L ix e, Manifest r e) + => Comp + -- ^ Computation startegy to use + -> [ListItem ix e] + -- ^ Nested list + -> Array r ix e +fromLists' comp = fromRaggedArray' . setComp comp . fromListToListArray {-# INLINE fromLists' #-} - -throughNested :: forall ix e . Nested LN ix e => [ListItem ix e] -> Array L ix e -throughNested xs = fromNested (fromNested xs :: Array LN ix e) -{-# INLINE throughNested #-} - - - -- | Convert any array to a flat list. -- -- ==== __Examples__ @@ -146,11 +149,10 @@ -- [(0,0),(0,1),(0,2),(1,0),(1,1),(1,2)] -- -- @since 0.1.0 -toList :: Source r ix e => Array r ix e -> [e] -toList !arr = build (\ c n -> foldrFB c n arr) +toList :: (Index ix, Source r e) => Array r ix e -> [e] +toList !arr = GHC.build (\c n -> foldrFB c n arr) {-# INLINE toList #-} - -- | /O(n)/ - Convert an array into a nested list. Number of array dimensions and list nestedness -- will always match, but you can use `toList`, `toLists2`, etc. if flattening of inner dimensions -- is desired. @@ -172,14 +174,14 @@ -- [[[0 :> 0 :. 0,0 :> 0 :. 1,0 :> 0 :. 2]],[[1 :> 0 :. 0,1 :> 0 :. 1,1 :> 0 :. 2]]] -- -- @since 0.1.0 -toLists :: (Nested LN ix e, Construct L ix e, Source r ix e) - => Array r ix e - -> [ListItem ix e] -toLists = toNested . toNested . toListArray +toLists + :: (Ragged L ix e, Shape r ix, Source r e) + => Array r ix e + -- ^ Array to be converted to nested lists + -> [ListItem ix e] +toLists = GHC.toList . toListArray {-# INLINE toLists #-} - - -- | Convert an array with at least 2 dimensions into a list of lists. Inner dimensions will get -- flattened. -- @@ -192,16 +194,16 @@ -- [[(0,0,0),(0,0,1),(0,0,2)],[(1,0,0),(1,0,1),(1,0,2)]] -- -- @since 0.1.0 -toLists2 :: (Source r ix e, Index (Lower ix)) => Array r ix e -> [[e]] +toLists2 :: (Source r e, Index ix, Index (Lower ix)) => Array r ix e -> [[e]] toLists2 = toList . foldrInner (:) [] {-# INLINE toLists2 #-} - -- | Convert an array with at least 3 dimensions into a 3 deep nested list. Inner dimensions will -- get flattened. -- -- @since 0.1.0 -toLists3 :: (Index (Lower (Lower ix)), Index (Lower ix), Source r ix e) => Array r ix e -> [[[e]]] +toLists3 + :: (Source r e, Index ix, Index (Lower ix), Index (Lower (Lower ix))) => Array r ix e -> [[[e]]] toLists3 = toList . foldrInner (:) [] . foldrInner (:) [] {-# INLINE toLists3 #-} @@ -209,13 +211,17 @@ -- get flattened. -- -- @since 0.1.0 -toLists4 :: - ( Index (Lower (Lower (Lower ix))) - , Index (Lower (Lower ix)) +toLists4 + :: ( Source r e + , Index ix , Index (Lower ix) - , Source r ix e + , Index (Lower (Lower ix)) + , Index (Lower (Lower (Lower ix))) ) => Array r ix e -> [[[[e]]]] toLists4 = toList . foldrInner (:) [] . foldrInner (:) [] . foldrInner (:) [] {-# INLINE toLists4 #-} + +-- $setup +-- >>> import Data.Massiv.Array as A 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 @@ -1,369 +1,482 @@ {-# LANGUAGE BangPatterns #-} -{-# LANGUAGE CPP #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MagicHash #-} {-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE NamedFieldPuns #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UnboxedTuples #-} {-# LANGUAGE UndecidableInstances #-} + -- | -- Module : Data.Massiv.Array.Manifest.Primitive --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Manifest.Primitive - ( P(..) - , Array(..) - , Prim - , toByteArray - , fromByteArrayM - , fromByteArray - , toMutableByteArray - , fromMutableByteArrayM - , fromMutableByteArray - , shrinkMutableByteArray - , unsafeAtomicReadIntArray - , unsafeAtomicWriteIntArray - , unsafeCasIntArray - , unsafeAtomicModifyIntArray - , unsafeAtomicAddIntArray - , unsafeAtomicSubIntArray - , unsafeAtomicAndIntArray - , unsafeAtomicNandIntArray - , unsafeAtomicOrIntArray - , unsafeAtomicXorIntArray - ) where +module Data.Massiv.Array.Manifest.Primitive ( + P (..), + Array (..), + MArray (..), + Prim, + toPrimitiveVector, + toPrimitiveMVector, + fromPrimitiveVector, + fromPrimitiveMVector, + toByteArray, + toByteArrayM, + unwrapByteArray, + unwrapByteArrayOffset, + unwrapMutableByteArray, + unwrapMutableByteArrayOffset, + fromByteArray, + fromByteArrayM, + fromByteArrayOffsetM, + toMutableByteArray, + toMutableByteArrayM, + fromMutableByteArrayM, + fromMutableByteArrayOffsetM, + fromMutableByteArray, + shrinkMutableByteArray, + unsafeAtomicReadIntArray, + unsafeAtomicWriteIntArray, + unsafeCasIntArray, + unsafeAtomicModifyIntArray, + unsafeAtomicAddIntArray, + unsafeAtomicSubIntArray, + unsafeAtomicAndIntArray, + unsafeAtomicNandIntArray, + unsafeAtomicOrIntArray, + unsafeAtomicXorIntArray, +) where -import Control.DeepSeq (NFData(..), deepseq) -import Control.Monad.Primitive (PrimMonad(primitive), PrimState, primitive_) -import Data.Massiv.Array.Delayed.Pull (eq, ord) +import Control.DeepSeq (NFData (..), deepseq) +import Control.Monad +import Control.Monad.Primitive (PrimMonad (..), primitive_) +import Data.Massiv.Array.Delayed.Pull -- (eq, ord) import Data.Massiv.Array.Manifest.Internal import Data.Massiv.Array.Manifest.List as A -import Data.Massiv.Array.Manifest.Vector.Stream as S (steps) import Data.Massiv.Array.Mutable import Data.Massiv.Core.Common import Data.Massiv.Core.List -import Data.Primitive (sizeOf) +import Data.Massiv.Core.Operations +import Data.Massiv.Vector.Stream as S (isteps, steps) +import Data.Maybe (fromMaybe) +import Data.Primitive (Prim, sizeOf) import Data.Primitive.ByteArray -import Data.Primitive.Types -import GHC.Base (Int(..)) +import qualified Data.Vector.Primitive as VP +import qualified Data.Vector.Primitive.Mutable as MVP import GHC.Exts as GHC -import Prelude hiding (mapM) import System.IO.Unsafe (unsafePerformIO) - -#include "massiv.h" +import Prelude hiding (mapM) -- | Representation for `Prim`itive elements -data P = P deriving Show +data P = P deriving (Show) -data instance Array P ix e = PArray { pComp :: !Comp - , pSize :: !(Sz ix) - , pData :: {-# UNPACK #-} !ByteArray - } +data instance Array P ix e = PArray + { pComp :: !Comp + , pSize :: !(Sz ix) + , pOffset :: {-# UNPACK #-} !Int + , pData :: {-# UNPACK #-} !ByteArray + } +data instance MArray s P ix e + = MPArray !(Sz ix) {-# UNPACK #-} !Int {-# UNPACK #-} !(MutableByteArray s) + instance (Ragged L ix e, Show e, Prim e) => Show (Array P ix e) where showsPrec = showsArrayPrec id showList = showArrayList instance Index ix => NFData (Array P ix e) where - rnf (PArray c sz a) = c `deepseq` sz `deepseq` a `seq` () + rnf (PArray c sz o a) = c `deepseq` sz `deepseq` o `seq` a `seq` () {-# INLINE rnf #-} +instance NFData ix => NFData (MArray s P ix e) where + rnf (MPArray sz _o _mb) = sz `deepseq` () + {-# INLINE rnf #-} + instance (Prim e, Eq e, Index ix) => Eq (Array P ix e) where - (==) = eq (==) + (==) = eqArrays (==) {-# INLINE (==) #-} instance (Prim e, Ord e, Index ix) => Ord (Array P ix e) where - compare = ord compare + compare = compareArrays compare {-# INLINE compare #-} -instance (Prim e, Index ix) => Construct P ix e where - setComp c arr = arr { pComp = c } +instance Strategy P where + getComp = pComp + {-# INLINE getComp #-} + setComp c arr = arr{pComp = c} {-# INLINE setComp #-} - - makeArray !comp !sz f = unsafePerformIO $ generateArray comp sz (return . f) - {-# INLINE makeArray #-} - -instance (Prim e, Index ix) => Source P ix e where - unsafeLinearIndex _pa@(PArray _ _ a) = - INDEX_CHECK("(Source P ix e).unsafeLinearIndex", - Sz . elemsBA _pa, indexByteArray) a - {-# INLINE unsafeLinearIndex #-} + repr = P +instance Index ix => Shape P ix where + maxLinearSize = Just . SafeSz . elemsCount + {-# INLINE maxLinearSize #-} -instance Index ix => Resize P ix where - unsafeResize !sz !arr = arr { pSize = sz } +instance Size P where + size = pSize + {-# INLINE size #-} + unsafeResize !sz !arr = arr{pSize = sz} {-# INLINE unsafeResize #-} -instance (Prim e, Index ix) => Extract P ix e where - unsafeExtract !sIx !newSz !arr = unsafeExtract sIx newSz (toManifest arr) - {-# INLINE unsafeExtract #-} - - -instance {-# OVERLAPPING #-} Prim e => Slice P Ix1 e where - unsafeSlice arr i _ _ = pure (unsafeLinearIndex arr i) - {-# INLINE unsafeSlice #-} - - -instance ( Prim e - , Index ix - , Index (Lower ix) - , Elt P ix e ~ Elt M ix e - , Elt M ix e ~ Array M (Lower ix) e - ) => - Slice P ix e where - unsafeSlice arr = unsafeSlice (toManifest arr) - {-# INLINE unsafeSlice #-} - -instance {-# OVERLAPPING #-} Prim e => OuterSlice P Ix1 e where - unsafeOuterSlice = unsafeLinearIndex - {-# INLINE unsafeOuterSlice #-} +instance Prim e => Source P e where + unsafeLinearIndex _arr@(PArray _ _ o a) i = + indexAssert "P.unsafeLinearIndex" (SafeSz . elemsBA _arr) indexByteArray a (i + o) + {-# INLINE unsafeLinearIndex #-} -instance ( Prim e - , Index ix - , Index (Lower ix) - , Elt M ix e ~ Array M (Lower ix) e - , Elt P ix e ~ Array M (Lower ix) e - ) => - OuterSlice P ix e where - unsafeOuterSlice arr = unsafeOuterSlice (toManifest arr) + unsafeOuterSlice (PArray c _ o a) szL i = + PArray c szL (i * totalElem szL + o) a {-# INLINE unsafeOuterSlice #-} - -instance {-# OVERLAPPING #-} Prim e => InnerSlice P Ix1 e where - unsafeInnerSlice arr _ = unsafeLinearIndex arr - {-# INLINE unsafeInnerSlice #-} - -instance ( Prim e - , Index ix - , Index (Lower ix) - , Elt M ix e ~ Array M (Lower ix) e - , Elt P ix e ~ Array M (Lower ix) e - ) => - InnerSlice P ix e where - unsafeInnerSlice arr = unsafeInnerSlice (toManifest arr) - {-# INLINE unsafeInnerSlice #-} - -instance (Index ix, Prim e) => Manifest P ix e where + unsafeLinearSlice i k (PArray c _ o a) = PArray c k (i + o) a + {-# INLINE unsafeLinearSlice #-} - unsafeLinearIndexM _pa@(PArray _ _ a) = - INDEX_CHECK("(Manifest P ix e).unsafeLinearIndexM", - Sz . elemsBA _pa, indexByteArray) a +instance Prim e => Manifest P e where + unsafeLinearIndexM _pa@(PArray _ _sz o a) i = + indexAssert "P.unsafeLinearIndexM" (const (Sz (totalElem _sz))) indexByteArray a (i + o) {-# INLINE unsafeLinearIndexM #-} + sizeOfMArray (MPArray sz _ _) = sz + {-# INLINE sizeOfMArray #-} -instance (Index ix, Prim e) => Mutable P ix e where - data MArray s P ix e = MPArray !(Sz ix) {-# UNPACK #-} !(MutableByteArray s) + unsafeResizeMArray sz (MPArray _ off marr) = MPArray sz off marr + {-# INLINE unsafeResizeMArray #-} - msize (MPArray sz _) = sz - {-# INLINE msize #-} + unsafeLinearSliceMArray i k (MPArray _ o a) = MPArray k (i + o) a + {-# INLINE unsafeLinearSliceMArray #-} - unsafeThaw (PArray _ sz a) = MPArray sz <$> unsafeThawByteArray a + unsafeThaw (PArray _ sz o a) = MPArray sz o <$> unsafeThawByteArray a {-# INLINE unsafeThaw #-} - unsafeFreeze comp (MPArray sz a) = PArray comp sz <$> unsafeFreezeByteArray a + unsafeFreeze comp (MPArray sz o a) = PArray comp sz o <$> unsafeFreezeByteArray a {-# INLINE unsafeFreeze #-} unsafeNew sz - | n <= (maxBound :: Int) `div` eSize = MPArray sz <$> newByteArray (n * eSize) + | n <= (maxBound :: Int) `div` eSize = MPArray sz 0 <$> newByteArray (n * eSize) | otherwise = error $ "Array size is too big: " ++ show sz - where !n = totalElem sz - !eSize = sizeOf (undefined :: e) + where + !n = totalElem sz + !eSize = sizeOf (undefined :: e) {-# INLINE unsafeNew #-} - initialize (MPArray sz mba) = - fillByteArray mba 0 (totalElem sz * sizeOf (undefined :: e)) 0 + initialize (MPArray sz o mba) = + let k = totalElem sz * sizeOf (undefined :: e) + in when (k > 0) $ fillByteArray mba o k 0 {-# INLINE initialize #-} - unsafeLinearRead _mpa@(MPArray _ ma) = - INDEX_CHECK("(Mutable P ix e).unsafeLinearRead", - Sz . elemsMBA _mpa, readByteArray) ma + unsafeLinearRead _mpa@(MPArray _sz o ma) i = + indexAssert "P.unsafeLinearRead" (const (Sz (totalElem _sz))) readByteArray ma (i + o) {-# INLINE unsafeLinearRead #-} - unsafeLinearWrite _mpa@(MPArray _ ma) = - INDEX_CHECK("(Mutable P ix e).unsafeLinearWrite", - Sz . elemsMBA _mpa, writeByteArray) ma + unsafeLinearWrite _mpa@(MPArray _sz o ma) i = + indexAssert "P.unsafeLinearWrite" (const (Sz (totalElem _sz))) writeByteArray ma (i + o) {-# INLINE unsafeLinearWrite #-} - unsafeLinearSet (MPArray _ ma) offset (SafeSz sz) = setByteArray ma offset sz + unsafeLinearSet (MPArray _ o ma) offset (SafeSz sz) = setByteArray ma (offset + o) sz {-# 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) + unsafeLinearCopy (MPArray _ oFrom maFrom) iFrom (MPArray _ oTo maTo) iTo (Sz k) = + copyMutableByteArray maTo ((oTo + iTo) * esz) maFrom ((oFrom + 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) + unsafeArrayLinearCopy (PArray _ _ oFrom aFrom) iFrom (MPArray _ oTo maTo) iTo (Sz k) = + copyByteArray maTo ((oTo + iTo) * esz) aFrom ((oFrom + 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 + unsafeLinearShrink (MPArray _ o ma) sz = do + shrinkMutableByteArray ma ((o + totalElem sz) * sizeOf (undefined :: e)) + pure $ MPArray sz o ma {-# INLINE unsafeLinearShrink #-} - unsafeLinearGrow (MPArray _ ma) sz = - MPArray sz <$> resizeMutableByteArrayCompat ma (totalElem sz * sizeOf (undefined :: e)) + unsafeLinearGrow (MPArray _ o ma) sz = + MPArray sz o <$> resizeMutableByteArray ma ((o + totalElem sz) * sizeOf (undefined :: e)) {-# INLINE unsafeLinearGrow #-} - instance (Prim e, Index ix) => Load P ix e where - type R P = M - size = pSize - {-# INLINE size #-} - getComp = pComp - {-# INLINE getComp #-} - loadArrayM !scheduler !arr = + makeArray comp sz f = compute (makeArray comp sz f :: Array D ix e) + {-# INLINE makeArray #-} + makeArrayLinear !comp !sz f = unsafePerformIO $ generateArrayLinear comp sz (pure . f) + {-# INLINE makeArrayLinear #-} + + replicate comp !sz !e = runST (newMArray sz e >>= unsafeFreeze comp) + {-# INLINE replicate #-} + + iterArrayLinearST_ !scheduler !arr = splitLinearlyWith_ scheduler (elemsCount arr) (unsafeLinearIndex arr) - {-# INLINE loadArrayM #-} + {-# INLINE iterArrayLinearST_ #-} instance (Prim e, Index ix) => StrideLoad P ix e instance (Prim e, Index ix) => Stream P ix e where toStream = S.steps {-# INLINE toStream #-} + toStreamIx = S.isteps + {-# INLINE toStreamIx #-} -instance ( Prim e - , IsList (Array L ix e) - , Nested LN ix e - , Nested L ix e - , Ragged L ix e - ) => - IsList (Array P ix e) where +instance (Prim e, Num e) => FoldNumeric P e where + unsafeDotProduct = defaultUnsafeDotProduct + {-# INLINE unsafeDotProduct #-} + powerSumArray = defaultPowerSumArray + {-# INLINE powerSumArray #-} + foldArray = defaultFoldArray + {-# INLINE foldArray #-} + +instance (Prim e, Num e) => Numeric P e where + unsafeLiftArray = defaultUnsafeLiftArray + {-# INLINE unsafeLiftArray #-} + unsafeLiftArray2 = defaultUnsafeLiftArray2 + {-# INLINE unsafeLiftArray2 #-} + +instance (Prim e, Floating e) => NumericFloat P e + +instance (Prim e, IsList (Array L ix e), Ragged L ix e) => IsList (Array P ix e) where type Item (Array P ix e) = Item (Array L ix e) fromList = A.fromLists' Seq {-# INLINE fromList #-} toList = GHC.toList . toListArray {-# INLINE toList #-} - -elemsBA :: forall proxy e . Prim e => proxy e -> ByteArray -> Int +elemsBA :: forall proxy e. Prim e => proxy e -> ByteArray -> Int elemsBA _ a = sizeofByteArray a `div` sizeOf (undefined :: e) {-# INLINE elemsBA #-} - -elemsMBA :: forall proxy e s . Prim e => proxy e -> MutableByteArray s -> Int +elemsMBA :: forall proxy e s. Prim e => proxy e -> MutableByteArray s -> Int elemsMBA _ a = sizeofMutableByteArray a `div` sizeOf (undefined :: e) {-# INLINE elemsMBA #-} - --- | /O(1)/ - Extract the internal `ByteArray`. +-- | /O(n)/ - Ensure that the size matches the internal `ByteArray`. If not make a copy of +-- the slice and return it as `ByteArray` -- -- @since 0.2.1 -toByteArray :: Array P ix e -> ByteArray -toByteArray = pData +toByteArray :: (Index ix, Prim e) => Array P ix e -> ByteArray +toByteArray arr = fromMaybe (unwrapByteArray $ clone arr) $ toByteArrayM arr {-# INLINE toByteArray #-} +-- | /O(1)/ - Extract the internal `ByteArray`. This will ignore any possible slicing that +-- has been applied to the array. Use `toByteArray` in order to preserve slicing or +-- `unwrapByteArrayOffset` to get ahold of the offset +-- +-- @since 0.5.0 +unwrapByteArray :: Array P ix e -> ByteArray +unwrapByteArray = pData +{-# INLINE unwrapByteArray #-} --- | /O(1)/ - Construct a primitive array from the `ByteArray`. Will return `Nothing` if number of --- elements doesn't match. +-- | /O(1)/ - Extract potential linear offset into the underlying `ByteArray`, which can +-- also be extracted with `unwrapByteArray`. -- +-- @since 0.5.9 +unwrapByteArrayOffset :: Array P ix e -> Int +unwrapByteArrayOffset = pOffset +{-# INLINE unwrapByteArrayOffset #-} + +-- | /O(1)/ - Unwrap Ensure that the size matches the internal `ByteArray`. +-- +-- @since 0.5.0 +toByteArrayM :: (Prim e, Index ix, MonadThrow m) => Array P ix e -> m ByteArray +toByteArrayM arr@PArray{pSize, pData} = do + pData <$ guardNumberOfElements pSize (Sz (elemsBA arr pData)) +{-# INLINE toByteArrayM #-} + +-- | /O(1)/ - Construct a primitive array from the `ByteArray`. Will return `Nothing` if +-- number of elements doesn't match. +-- -- @since 0.3.0 fromByteArrayM :: (MonadThrow m, Index ix, Prim e) => Comp -> Sz ix -> ByteArray -> m (Array P ix e) -fromByteArrayM comp sz ba = - guardNumberOfElements sz (Sz (elemsBA arr ba)) >> pure arr - where - arr = PArray comp sz ba +fromByteArrayM comp sz = fromByteArrayOffsetM comp sz 0 {-# INLINE fromByteArrayM #-} +-- | /O(1)/ - Construct a primitive array from the `ByteArray`. Will return `Nothing` if +-- number of elements doesn't match. +-- +-- @since 0.5.9 +fromByteArrayOffsetM + :: (MonadThrow m, Index ix, Prim e) => Comp -> Sz ix -> Int -> ByteArray -> m (Array P ix e) +fromByteArrayOffsetM comp sz off ba = + arr <$ guardNumberOfElements sz (SafeSz (elemsBA arr ba - off)) + where + arr = PArray comp sz off ba +{-# INLINE fromByteArrayOffsetM #-} + -- | /O(1)/ - Construct a flat Array from `ByteArray` -- -- @since 0.4.0 -fromByteArray :: forall e . Prim e => Comp -> ByteArray -> Array P Ix1 e -fromByteArray comp ba = PArray comp (SafeSz (elemsBA (Proxy :: Proxy e) ba)) ba +fromByteArray :: forall e. Prim e => Comp -> ByteArray -> Array P Ix1 e +fromByteArray comp ba = PArray comp (SafeSz (elemsBA (Proxy :: Proxy e) ba)) 0 ba {-# INLINE fromByteArray #-} +-- | /O(1)/ - Extract the internal `MutableByteArray`. This will discard any possible +-- slicing that has been applied to the array. +-- +-- @since 0.5.0 +unwrapMutableByteArray :: MArray s P ix e -> MutableByteArray s +unwrapMutableByteArray (MPArray _ _ mba) = mba +{-# INLINE unwrapMutableByteArray #-} --- | /O(1)/ - Extract the internal `MutableByteArray`. +-- | /O(1)/ - Extract the linear offset into underlying `MutableByteArray`, which can aslo +-- be extracted with `unwrapMutableByteArray`. -- --- @since 0.2.1 -toMutableByteArray :: MArray s P ix e -> MutableByteArray s -toMutableByteArray (MPArray _ mba) = mba +-- @since 0.5.9 +unwrapMutableByteArrayOffset :: MArray s P ix e -> Int +unwrapMutableByteArrayOffset (MPArray _ off _) = off +{-# INLINE unwrapMutableByteArrayOffset #-} + +-- | /O(n)/ - Try to cast a mutable array to `MutableByteArray`, if sizes do not match make +-- a copy. Returns `True` if an array was converted without a copy, in which case it means +-- that the source at the resulting array are still pointing to the same location in memory. +-- +-- @since 0.5.0 +toMutableByteArray + :: forall ix e m + . (Prim e, Index ix, PrimMonad m) + => MArray (PrimState m) P ix e + -> m (Bool, MutableByteArray (PrimState m)) +toMutableByteArray marr@(MPArray sz offset mbas) = + case toMutableByteArrayM marr of + Just mba -> pure (True, mba) + Nothing -> do + let eSize = sizeOf (undefined :: e) + szBytes = totalElem sz * eSize + mbad <- newPinnedByteArray szBytes + copyMutableByteArray mbad 0 mbas (offset * eSize) szBytes + pure (False, mbad) {-# INLINE toMutableByteArray #-} +-- | /O(1)/ - Extract the internal `MutableByteArray`. +-- +-- @since 0.2.1 +toMutableByteArrayM :: (Index ix, Prim e, MonadThrow m) => MArray s P ix e -> m (MutableByteArray s) +toMutableByteArrayM marr@(MPArray sz _ mba) = + mba <$ guardNumberOfElements sz (Sz (elemsMBA marr mba)) +{-# INLINE toMutableByteArrayM #-} -- | /O(1)/ - Construct a primitive mutable array from the `MutableByteArray`. Will throw -- `SizeElementsMismatchException` if number of elements doesn't match. -- -- @since 0.3.0 -fromMutableByteArrayM :: - (MonadThrow m, Index ix, Prim e) => Sz ix -> MutableByteArray s -> m (MArray s P ix e) -fromMutableByteArrayM sz mba = - guardNumberOfElements sz (Sz (elemsMBA marr mba)) >> pure marr - where - marr = MPArray sz mba +fromMutableByteArrayM + :: (MonadThrow m, Index ix, Prim e) => Sz ix -> MutableByteArray s -> m (MArray s P ix e) +fromMutableByteArrayM sz = fromMutableByteArrayOffsetM sz 0 {-# INLINE fromMutableByteArrayM #-} +-- | /O(1)/ - Construct a primitive mutable array from the `MutableByteArray`. Will throw +-- `SizeElementsMismatchException` if number of elements doesn't match. +-- +-- @since 0.5.9 +fromMutableByteArrayOffsetM + :: (MonadThrow m, Index ix, Prim e) => Sz ix -> Ix1 -> MutableByteArray s -> m (MArray s P ix e) +fromMutableByteArrayOffsetM sz off mba = + marr <$ guardNumberOfElements sz (SafeSz (elemsMBA marr mba - off)) + where + marr = MPArray sz off mba +{-# INLINE fromMutableByteArrayOffsetM #-} + -- | /O(1)/ - Construct a flat Array from `MutableByteArray` -- -- @since 0.4.0 -fromMutableByteArray :: forall e s . Prim e => MutableByteArray s -> MArray s P Ix1 e -fromMutableByteArray mba = MPArray (SafeSz (elemsMBA (Proxy :: Proxy e) mba)) mba +fromMutableByteArray :: forall e s. Prim e => MutableByteArray s -> MArray s P Ix1 e +fromMutableByteArray mba = MPArray (SafeSz (elemsMBA (Proxy :: Proxy e) mba)) 0 mba {-# INLINE fromMutableByteArray #-} +-- | /O(1)/ - Cast a primitive array to a primitive vector. +-- +-- @since 0.5.0 +toPrimitiveVector :: Index ix => Array P ix e -> VP.Vector e +toPrimitiveVector PArray{pSize, pOffset, pData} = VP.Vector pOffset (totalElem pSize) pData +{-# INLINE toPrimitiveVector #-} +-- | /O(1)/ - Cast a mutable primitive array to a mutable primitive vector. +-- +-- @since 0.5.0 +toPrimitiveMVector :: Index ix => MArray s P ix e -> MVP.MVector s e +toPrimitiveMVector (MPArray sz offset mba) = MVP.MVector offset (totalElem sz) mba +{-# INLINE toPrimitiveMVector #-} + +-- | /O(1)/ - Cast a primitive vector to a primitive array. +-- +-- @since 0.5.0 +fromPrimitiveVector :: VP.Vector e -> Array P Ix1 e +fromPrimitiveVector (VP.Vector offset len ba) = + PArray{pComp = Seq, pSize = SafeSz len, pOffset = offset, pData = ba} +{-# INLINE fromPrimitiveVector #-} + +-- | /O(1)/ - Cast a mutable primitive vector to a mutable primitive array. +-- +-- @since 0.5.0 +fromPrimitiveMVector :: MVP.MVector s e -> MArray s P Ix1 e +fromPrimitiveMVector (MVP.MVector offset len mba) = MPArray (SafeSz len) offset mba +{-# INLINE fromPrimitiveMVector #-} + -- | Atomically read an `Int` element from the array -- -- @since 0.3.0 -unsafeAtomicReadIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> m Int -unsafeAtomicReadIntArray _mpa@(MPArray sz mba) ix = - INDEX_CHECK( "unsafeAtomicReadIntArray" - , Sz . elemsMBA _mpa - , \(MutableByteArray mba#) (I# i#) -> - primitive $ \s# -> - case atomicReadIntArray# mba# i# s# of - (# s'#, e# #) -> (# s'#, I# e# #)) - mba - (toLinearIndex sz ix) +unsafeAtomicReadIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> m Int +unsafeAtomicReadIntArray _mpa@(MPArray sz o mba) ix = + indexAssert + "P.unsafeAtomicReadIntArray" + (SafeSz . elemsMBA _mpa) + ( \(MutableByteArray mba#) (I# i#) -> + primitive $ \s# -> + case atomicReadIntArray# mba# i# s# of + (# s'#, e# #) -> (# s'#, I# e# #) + ) + mba + (o + toLinearIndex sz ix) {-# INLINE unsafeAtomicReadIntArray #-} -- | Atomically write an `Int` element int the array -- -- @since 0.3.0 -unsafeAtomicWriteIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m () -unsafeAtomicWriteIntArray _mpa@(MPArray sz mba) ix (I# e#) = - INDEX_CHECK( "unsafeAtomicWriteIntArray" - , Sz . elemsMBA _mpa - , \(MutableByteArray mba#) (I# i#) -> - primitive_ (atomicWriteIntArray# mba# i# e#)) - mba - (toLinearIndex sz ix) +unsafeAtomicWriteIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m () +unsafeAtomicWriteIntArray _mpa@(MPArray sz o mba) ix (I# e#) = + indexAssert + "P.unsafeAtomicWriteIntArray" + (SafeSz . elemsMBA _mpa) + ( \(MutableByteArray mba#) (I# i#) -> + primitive_ (atomicWriteIntArray# mba# i# e#) + ) + mba + (o + toLinearIndex sz ix) {-# INLINE unsafeAtomicWriteIntArray #-} -- | Atomically CAS an `Int` in the array. Returns the old value. -- -- @since 0.3.0 -unsafeCasIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> Int -> m Int -unsafeCasIntArray _mpa@(MPArray sz mba) ix (I# e#) (I# n#) = - INDEX_CHECK( "unsafeCasIntArray" - , Sz . elemsMBA _mpa - , \(MutableByteArray mba#) (I# i#) -> - primitive $ \s# -> - case casIntArray# mba# i# e# n# s# of - (# s'#, o# #) -> (# s'#, I# o# #)) - mba - (toLinearIndex sz ix) +unsafeCasIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> Int -> m Int +unsafeCasIntArray _mpa@(MPArray sz o mba) ix (I# e#) (I# n#) = + indexAssert + "P.unsafeCasIntArray" + (SafeSz . elemsMBA _mpa) + ( \(MutableByteArray mba#) (I# i#) -> + primitive $ \s# -> + case casIntArray# mba# i# e# n# s# of + (# s'#, o# #) -> (# s'#, I# o# #) + ) + mba + (o + toLinearIndex sz ix) {-# INLINE unsafeCasIntArray #-} - -- | Atomically modify an `Int` element of the array. Returns the old value. -- -- @since 0.3.0 -unsafeAtomicModifyIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> (Int -> Int) -> m Int -unsafeAtomicModifyIntArray _mpa@(MPArray sz mba) ix f = - INDEX_CHECK("unsafeAtomicModifyIntArray", Sz . elemsMBA _mpa, atomicModify) - mba - (toLinearIndex sz ix) +unsafeAtomicModifyIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> (Int -> Int) -> m Int +unsafeAtomicModifyIntArray _mpa@(MPArray sz o mba) ix f = + indexAssert + "P.unsafeAtomicModifyIntArray" + (SafeSz . elemsMBA _mpa) + atomicModify + mba + (o + toLinearIndex sz ix) where atomicModify (MutableByteArray mba#) (I# i#) = let go s# o# = @@ -372,134 +485,117 @@ (# s'#, o'# #) -> case o# ==# o'# of 0# -> go s# o'# - _ -> (# s'#, I# o# #) + _ -> (# s'#, I# o# #) in primitive $ \s# -> case atomicReadIntArray# mba# i# s# of (# s'#, o# #) -> go s'# o# {-# INLINE atomicModify #-} {-# INLINE unsafeAtomicModifyIntArray #-} - -- | Atomically add to an `Int` element in the array. Returns the old value. -- -- @since 0.3.0 -unsafeAtomicAddIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int -unsafeAtomicAddIntArray _mpa@(MPArray sz mba) ix (I# e#) = - INDEX_CHECK( "unsafeAtomicAddIntArray" - , Sz . elemsMBA _mpa - , \(MutableByteArray mba#) (I# i#) -> - primitive $ \s# -> - case fetchAddIntArray# mba# i# e# s# of - (# s'#, p# #) -> (# s'#, I# p# #)) - mba - (toLinearIndex sz ix) +unsafeAtomicAddIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int +unsafeAtomicAddIntArray _mpa@(MPArray sz o mba) ix (I# e#) = + indexAssert + "P.unsafeAtomicAddIntArray" + (SafeSz . elemsMBA _mpa) + ( \(MutableByteArray mba#) (I# i#) -> + primitive $ \s# -> + case fetchAddIntArray# mba# i# e# s# of + (# s'#, p# #) -> (# s'#, I# p# #) + ) + mba + (o + toLinearIndex sz ix) {-# INLINE unsafeAtomicAddIntArray #-} - -- | Atomically subtract from an `Int` element in the array. Returns the old value. -- -- @since 0.3.0 -unsafeAtomicSubIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int -unsafeAtomicSubIntArray _mpa@(MPArray sz mba) ix (I# e#) = - INDEX_CHECK( "unsafeAtomicSubIntArray" - , Sz . elemsMBA _mpa - , \(MutableByteArray mba#) (I# i#) -> - primitive $ \s# -> - case fetchSubIntArray# mba# i# e# s# of - (# s'#, p# #) -> (# s'#, I# p# #)) - mba - (toLinearIndex sz ix) +unsafeAtomicSubIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int +unsafeAtomicSubIntArray _mpa@(MPArray sz o mba) ix (I# e#) = + indexAssert + "P.unsafeAtomicSubIntArray" + (SafeSz . elemsMBA _mpa) + ( \(MutableByteArray mba#) (I# i#) -> + primitive $ \s# -> + case fetchSubIntArray# mba# i# e# s# of + (# s'#, p# #) -> (# s'#, I# p# #) + ) + mba + (o + toLinearIndex sz ix) {-# INLINE unsafeAtomicSubIntArray #-} - -- | Atomically AND an `Int` element in the array. Returns the old value. -- -- @since 0.3.0 -unsafeAtomicAndIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int -unsafeAtomicAndIntArray _mpa@(MPArray sz mba) ix (I# e#) = - INDEX_CHECK( "unsafeAtomicAndIntArray" - , Sz . elemsMBA _mpa - , \(MutableByteArray mba#) (I# i#) -> - primitive $ \s# -> - case fetchAndIntArray# mba# i# e# s# of - (# s'#, p# #) -> (# s'#, I# p# #)) - mba - (toLinearIndex sz ix) +unsafeAtomicAndIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int +unsafeAtomicAndIntArray _mpa@(MPArray sz o mba) ix (I# e#) = + indexAssert + "P.unsafeAtomicAndIntArray" + (SafeSz . elemsMBA _mpa) + ( \(MutableByteArray mba#) (I# i#) -> + primitive $ \s# -> + case fetchAndIntArray# mba# i# e# s# of + (# s'#, p# #) -> (# s'#, I# p# #) + ) + mba + (o + toLinearIndex sz ix) {-# INLINE unsafeAtomicAndIntArray #-} - -- | Atomically NAND an `Int` element in the array. Returns the old value. -- -- @since 0.3.0 -unsafeAtomicNandIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int -unsafeAtomicNandIntArray _mpa@(MPArray sz mba) ix (I# e#) = - INDEX_CHECK( "unsafeAtomicNandIntArray" - , Sz . elemsMBA _mpa - , \(MutableByteArray mba#) (I# i#) -> - primitive $ \s# -> - case fetchNandIntArray# mba# i# e# s# of - (# s'#, p# #) -> (# s'#, I# p# #)) - mba - (toLinearIndex sz ix) +unsafeAtomicNandIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int +unsafeAtomicNandIntArray _mpa@(MPArray sz o mba) ix (I# e#) = + indexAssert + "P.unsafeAtomicNandIntArray" + (SafeSz . elemsMBA _mpa) + ( \(MutableByteArray mba#) (I# i#) -> + primitive $ \s# -> + case fetchNandIntArray# mba# i# e# s# of + (# s'#, p# #) -> (# s'#, I# p# #) + ) + mba + (o + toLinearIndex sz ix) {-# INLINE unsafeAtomicNandIntArray #-} - -- | Atomically OR an `Int` element in the array. Returns the old value. -- -- @since 0.3.0 -unsafeAtomicOrIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int -unsafeAtomicOrIntArray _mpa@(MPArray sz mba) ix (I# e#) = - INDEX_CHECK( "unsafeAtomicOrIntArray" - , Sz . elemsMBA _mpa - , \(MutableByteArray mba#) (I# i#) -> - primitive $ \s# -> - case fetchOrIntArray# mba# i# e# s# of - (# s'#, p# #) -> (# s'#, I# p# #)) - mba - (toLinearIndex sz ix) +unsafeAtomicOrIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int +unsafeAtomicOrIntArray _mpa@(MPArray sz o mba) ix (I# e#) = + indexAssert + "P.unsafeAtomicOrIntArray" + (SafeSz . elemsMBA _mpa) + ( \(MutableByteArray mba#) (I# i#) -> + primitive $ \s# -> + case fetchOrIntArray# mba# i# e# s# of + (# s'#, p# #) -> (# s'#, I# p# #) + ) + mba + (o + toLinearIndex sz ix) {-# INLINE unsafeAtomicOrIntArray #-} - -- | Atomically XOR an `Int` element in the array. Returns the old value. -- -- @since 0.3.0 -unsafeAtomicXorIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int -unsafeAtomicXorIntArray _mpa@(MPArray sz mba) ix (I# e#) = - INDEX_CHECK( "unsafeAtomicXorIntArray" - , Sz . elemsMBA _mpa - , \(MutableByteArray mba#) (I# i#) -> - primitive $ \s# -> - case fetchXorIntArray# mba# i# e# s# of - (# s'#, p# #) -> (# s'#, I# p# #)) - mba - (toLinearIndex sz ix) +unsafeAtomicXorIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int +unsafeAtomicXorIntArray _mpa@(MPArray sz o mba) ix (I# e#) = + indexAssert + "P.unsafeAtomicXorIntArray" + (SafeSz . elemsMBA _mpa) + ( \(MutableByteArray mba#) (I# i#) -> + primitive $ \s# -> + case fetchXorIntArray# mba# i# e# s# of + (# s'#, p# #) -> (# s'#, I# p# #) + ) + mba + (o + toLinearIndex sz ix) {-# INLINE unsafeAtomicXorIntArray #-} - - -shrinkMutableByteArray :: forall m. (PrimMonad m) - => MutableByteArray (PrimState m) - -> Int -- ^ new size - -> m () -shrinkMutableByteArray (MutableByteArray arr#) (I# n#) - = primitive_ (shrinkMutableByteArray# arr# n#) -{-# INLINE shrinkMutableByteArray #-} - - -resizeMutableByteArrayCompat :: - PrimMonad m => MutableByteArray (PrimState m) -> Int -> m (MutableByteArray (PrimState m)) -#if MIN_VERSION_primitive(0,6,4) -resizeMutableByteArrayCompat = resizeMutableByteArray -#else -resizeMutableByteArrayCompat (MutableByteArray arr#) (I# n#) = - primitive - (\s# -> - case resizeMutableByteArray# arr# n# s# of - (# s'#, arr'# #) -> (# s'#, MutableByteArray arr'# #)) -#endif -{-# INLINE resizeMutableByteArrayCompat #-} 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 @@ -7,170 +7,201 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UndecidableInstances #-} + -- | -- Module : Data.Massiv.Array.Manifest.Storable --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Manifest.Storable - ( S (..) - , Array(..) - , VS.Storable - , toStorableVector - , toStorableMVector - , withPtr - , unsafeWithPtr - , unsafeArrayToForeignPtr - , unsafeMArrayToForeignPtr - , unsafeArrayFromForeignPtr - , unsafeArrayFromForeignPtr0 - , unsafeMArrayFromForeignPtr - , unsafeMArrayFromForeignPtr0 - ) where +module Data.Massiv.Array.Manifest.Storable ( + S (..), + Array (..), + MArray (..), + Storable, + toStorableVector, + toStorableMVector, + fromStorableVector, + fromStorableMVector, + withPtr, + unsafeWithPtr, + unsafeMallocMArray, + unsafeArrayToForeignPtr, + unsafeMArrayToForeignPtr, + unsafeArrayFromForeignPtr, + unsafeArrayFromForeignPtr0, + unsafeMArrayFromForeignPtr, + unsafeMArrayFromForeignPtr0, +) where -import Control.DeepSeq (NFData(..), deepseq) +import Control.DeepSeq (NFData (..), deepseq) +import Control.Exception +import Control.Monad import Control.Monad.IO.Unlift -import Control.Monad.Primitive (unsafePrimToPrim) -import Data.Massiv.Array.Delayed.Pull (eq, ord) +import Control.Monad.Primitive +import Data.Massiv.Array.Delayed.Pull (D, compareArrays, eqArrays) 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.Manifest.Vector.Stream as S (steps) import Data.Massiv.Array.Mutable import Data.Massiv.Core.Common import Data.Massiv.Core.List -import qualified Data.Vector.Generic.Mutable as VGM +import Data.Massiv.Core.Operations +import Data.Massiv.Vector.Stream as S (isteps, steps) +import Data.Primitive.ByteArray +import Data.Primitive.Ptr (setPtr) +import qualified Data.Vector.Generic.Mutable as MVG import qualified Data.Vector.Storable as VS import qualified Data.Vector.Storable.Mutable as MVS +import Data.Word +import Foreign.ForeignPtr +import Foreign.Marshal.Alloc +import Foreign.Marshal.Array (advancePtr, copyArray) 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 GHC.Exts as GHC +import GHC.ForeignPtr import System.IO.Unsafe (unsafePerformIO) - -#include "massiv.h" +import Unsafe.Coerce +import Prelude hiding (mapM) -- | Representation for `Storable` elements -data S = S deriving Show +data S = S deriving (Show) -data instance Array S ix e = SArray { sComp :: !Comp - , sSize :: !(Sz ix) - , sData :: !(VS.Vector e) - } +data instance Array S ix e = SArray + { sComp :: !Comp + , sSize :: !(Sz ix) + , sData :: {-# UNPACK #-} !(ForeignPtr e) + } -instance (Ragged L ix e, Show e, VS.Storable e) => Show (Array S ix e) where +data instance MArray s S ix e = MSArray !(Sz ix) {-# UNPACK #-} !(ForeignPtr e) + +instance (Ragged L ix e, Show e, Storable e) => Show (Array S ix e) where showsPrec = showsArrayPrec id showList = showArrayList instance NFData ix => NFData (Array S ix e) where - rnf (SArray c sz v) = c `deepseq` sz `deepseq` v `deepseq` () + rnf (SArray c sz _v) = c `deepseq` sz `deepseq` () {-# INLINE rnf #-} -instance (VS.Storable e, Eq e, Index ix) => Eq (Array S ix e) where - (==) = eq (==) +instance NFData ix => NFData (MArray s S ix e) where + rnf (MSArray sz _mv) = sz `deepseq` () + {-# INLINE rnf #-} + +instance (Storable e, Eq e, Index ix) => Eq (Array S ix e) where + (==) = eqArrays (==) {-# INLINE (==) #-} -instance (VS.Storable e, Ord e, Index ix) => Ord (Array S ix e) where - compare = ord compare +instance (Storable e, Ord e, Index ix) => Ord (Array S ix e) where + compare = compareArrays compare {-# INLINE compare #-} -instance (VS.Storable e, Index ix) => Construct S ix e where - setComp c arr = arr { sComp = c } +instance Strategy S where + getComp = sComp + {-# INLINE getComp #-} + setComp c arr = arr{sComp = c} {-# INLINE setComp #-} - - makeArray !comp !sz f = unsafePerformIO $ generateArray comp sz (return . f) - {-# INLINE makeArray #-} - - -instance (VS.Storable e, Index ix) => Source S ix e where - unsafeLinearIndex (SArray _ _ v) = - INDEX_CHECK("(Source S ix e).unsafeLinearIndex", Sz . VS.length, VS.unsafeIndex) v - {-# INLINE unsafeLinearIndex #-} + repr = S -instance Index ix => Resize S ix where - unsafeResize !sz !arr = arr { sSize = sz } - {-# INLINE unsafeResize #-} +plusFp :: ForeignPtr a -> Int -> ForeignPtr b +plusFp (ForeignPtr addr c) (I# d) = ForeignPtr (plusAddr# addr d) c -instance (VS.Storable e, Index ix) => Extract S ix e where - unsafeExtract !sIx !newSz !arr = unsafeExtract sIx newSz (toManifest arr) - {-# INLINE unsafeExtract #-} +advanceForeignPtr :: forall e. Storable e => ForeignPtr e -> Int -> ForeignPtr e +advanceForeignPtr fp i = plusFp fp (i * sizeOf (undefined :: e)) +{-# INLINE advanceForeignPtr #-} +indexForeignPtr :: Storable e => ForeignPtr e -> Int -> e +indexForeignPtr fp i = unsafeInlineIO $ unsafeWithForeignPtr fp $ \p -> peekElemOff p i +{-# INLINE indexForeignPtr #-} +instance Storable e => Source S e where + unsafeLinearIndex (SArray _ _sz fp) = + indexAssert "S.unsafeLinearIndex" (const (toLinearSz _sz)) indexForeignPtr fp + {-# INLINE unsafeLinearIndex #-} -instance ( VS.Storable e - , Index ix - , Index (Lower ix) - , Elt M ix e ~ Array M (Lower ix) e - , Elt S ix e ~ Array M (Lower ix) e - ) => - OuterSlice S ix e where - unsafeOuterSlice arr = unsafeOuterSlice (toManifest arr) + unsafeOuterSlice (SArray c _ fp) szL i = + let k = totalElem szL + in SArray c szL $ advanceForeignPtr fp (i * k) {-# INLINE unsafeOuterSlice #-} -instance ( VS.Storable e - , Index ix - , Index (Lower ix) - , Elt M ix e ~ Array M (Lower ix) e - , Elt S ix e ~ Array M (Lower ix) e - ) => - InnerSlice S ix e where - unsafeInnerSlice arr = unsafeInnerSlice (toManifest arr) - {-# INLINE unsafeInnerSlice #-} - -instance {-# OVERLAPPING #-} VS.Storable e => Slice S Ix1 e where - unsafeSlice arr i _ _ = pure (unsafeLinearIndex arr i) - {-# INLINE unsafeSlice #-} + unsafeLinearSlice i k (SArray c _ fp) = + SArray c k $ advanceForeignPtr fp i + {-# INLINE unsafeLinearSlice #-} +instance Index ix => Shape S ix where + maxLinearSize = Just . SafeSz . elemsCount + {-# INLINE maxLinearSize #-} -instance (Index ix, VS.Storable e) => Manifest S ix e where +instance Size S where + size = sSize + {-# INLINE size #-} + unsafeResize !sz !arr = arr{sSize = sz} + {-# INLINE unsafeResize #-} - unsafeLinearIndexM (SArray _ _ v) = - INDEX_CHECK("(Manifest S ix e).unsafeLinearIndexM", Sz . VS.length, VS.unsafeIndex) v +instance Storable e => Manifest S e where + unsafeLinearIndexM (SArray _ _sz fp) = + indexAssert "S.unsafeLinearIndex" (const (toLinearSz _sz)) indexForeignPtr fp {-# INLINE unsafeLinearIndexM #-} + sizeOfMArray (MSArray sz _) = sz + {-# INLINE sizeOfMArray #-} -instance (Index ix, VS.Storable e) => Mutable S ix e where - data MArray s S ix e = MSArray !(Sz ix) !(VS.MVector s e) + unsafeResizeMArray sz (MSArray _ fp) = MSArray sz fp + {-# INLINE unsafeResizeMArray #-} - msize (MSArray sz _) = sz - {-# INLINE msize #-} + unsafeLinearSliceMArray i k (MSArray _ fp) = MSArray k $ advanceForeignPtr fp i + {-# INLINE unsafeLinearSliceMArray #-} - unsafeThaw (SArray _ sz v) = MSArray sz <$> VS.unsafeThaw v + unsafeThaw (SArray _ sz fp) = pure $ MSArray sz fp {-# INLINE unsafeThaw #-} - unsafeFreeze comp (MSArray sz v) = SArray comp sz <$> VS.unsafeFreeze v + unsafeFreeze comp (MSArray sz v) = pure $ SArray comp sz v {-# INLINE unsafeFreeze #-} - unsafeNew sz = MSArray sz <$> MVS.unsafeNew (totalElem sz) + unsafeNew sz = do + let !n = totalElem sz + dummy = undefined :: e + !eSize = sizeOf dummy + when (n > (maxBound :: Int) `div` eSize) $ error $ "Array size is too big: " ++ show sz + unsafeIOToPrim $ do + fp <- mallocPlainForeignPtrAlignedBytes (n * sizeOf dummy) (alignment dummy) + pure $ MSArray sz fp {-# INLINE unsafeNew #-} - initialize (MSArray _ marr) = VGM.basicInitialize marr + initialize (MSArray sz fp) = + unsafeIOToPrim $ + unsafeWithForeignPtr fp $ \p -> + setPtr (castPtr p) (totalElem sz * sizeOf (undefined :: e)) (0 :: Word8) {-# INLINE initialize #-} - unsafeLinearRead (MSArray _ mv) = - INDEX_CHECK("(Mutable S ix e).unsafeLinearRead", Sz . MVS.length, MVS.unsafeRead) mv + unsafeLinearRead (MSArray _sz fp) o = + unsafeIOToPrim $ + indexAssert + "S.unsafeLinearRead" + (const (toLinearSz _sz)) + (\_ _ -> unsafeWithForeignPtr fp (`peekElemOff` o)) + fp + o {-# INLINE unsafeLinearRead #-} - unsafeLinearWrite (MSArray _ mv) = - INDEX_CHECK("(Mutable S ix e).unsafeLinearWrite", Sz . MVS.length, MVS.unsafeWrite) mv + unsafeLinearWrite (MSArray _sz fp) o e = + unsafeIOToPrim $ + indexAssert + "S.unsafeLinearWrite" + (const (toLinearSz _sz)) + (\_ _ -> unsafeWithForeignPtr fp (\p -> pokeElemOff p o e)) + fp + o {-# INLINE unsafeLinearWrite #-} - unsafeLinearSet (MSArray _ mv) i k = VGM.basicSet (MVS.unsafeSlice i (unSz k) mv) + unsafeLinearSet (MSArray _ fp) i k e = + stToPrim (MVG.basicSet (MVS.unsafeFromForeignPtr0 (advanceForeignPtr fp i) (unSz k)) e) {-# INLINE unsafeLinearSet #-} - unsafeLinearCopy marrFrom iFrom marrTo iTo (Sz k) = do - let MSArray _ (MVS.MVector _ fpFrom) = marrFrom - MSArray _ (MVS.MVector _ fpTo) = marrTo + unsafeLinearCopy (MSArray _ fpFrom) iFrom (MSArray _ fpTo) iTo (Sz k) = do unsafePrimToPrim $ - withForeignPtr fpFrom $ \ ptrFrom -> - withForeignPtr fpTo $ \ ptrTo -> do + withForeignPtr fpFrom $ \ptrFrom -> + withForeignPtr fpTo $ \ptrTo -> do let ptrFrom' = advancePtr ptrFrom iFrom ptrTo' = advancePtr ptrTo iTo copyArray ptrTo' ptrFrom' k @@ -181,48 +212,59 @@ unsafeLinearCopy marrFrom iFrom marrTo iTo sz {-# INLINE unsafeArrayLinearCopy #-} - unsafeLinearShrink marr@(MSArray _ mv@(MVS.MVector _ (ForeignPtr _ fpc))) sz = do + unsafeLinearShrink marr@(MSArray _ fp@(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 + pure $ MSArray sz fp PlainPtr mba# -> do unsafePrimToPrim $ shrinkMBA (MutableByteArray mba#) - pure $ MSArray sz mv + pure $ MSArray sz fp _ -> unsafeDefaultLinearShrink marr sz {-# INLINE unsafeLinearShrink #-} - unsafeLinearGrow (MSArray oldSz mv) sz = - MSArray sz <$> MVS.unsafeGrow mv (totalElem sz - totalElem oldSz) - {-# INLINE unsafeLinearGrow #-} +instance (Index ix, Storable e) => Load S ix e where + makeArray comp sz f = compute (makeArray comp sz f :: Array D ix e) + {-# INLINE makeArray #-} + makeArrayLinear !comp !sz f = unsafePerformIO $ generateArrayLinear comp sz (pure . f) + {-# INLINE makeArrayLinear #-} -instance (Index ix, VS.Storable e) => Load S ix e where - type R S = M - size = sSize - {-# INLINE size #-} - getComp = sComp - {-# INLINE getComp #-} - loadArrayM !scheduler !arr = splitLinearlyWith_ scheduler (elemsCount arr) (unsafeLinearIndex arr) - {-# INLINE loadArrayM #-} + replicate comp !sz !e = runST (newMArray sz e >>= unsafeFreeze comp) + {-# INLINE replicate #-} -instance (Index ix, VS.Storable e) => StrideLoad S ix e + iterArrayLinearST_ !scheduler !arr = + splitLinearlyWith_ scheduler (elemsCount arr) (unsafeLinearIndex arr) + {-# INLINE iterArrayLinearST_ #-} -instance (Index ix, VS.Storable e) => Stream S ix e where +instance (Index ix, Storable e) => StrideLoad S ix e + +instance (Index ix, Storable e) => Stream S ix e where toStream = S.steps {-# INLINE toStream #-} + toStreamIx = S.isteps + {-# INLINE toStreamIx #-} +instance (Storable e, Num e) => FoldNumeric S e where + unsafeDotProduct = defaultUnsafeDotProduct + {-# INLINE unsafeDotProduct #-} + powerSumArray = defaultPowerSumArray + {-# INLINE powerSumArray #-} + foldArray = defaultFoldArray + {-# INLINE foldArray #-} -instance ( VS.Storable e - , IsList (Array L ix e) - , Nested LN ix e - , Nested L ix e - , Ragged L ix e - ) => - IsList (Array S ix e) where +instance (Storable e, Num e) => Numeric S e where + unsafeLiftArray = defaultUnsafeLiftArray + {-# INLINE unsafeLiftArray #-} + unsafeLiftArray2 = defaultUnsafeLiftArray2 + {-# INLINE unsafeLiftArray2 #-} + +instance (Storable e, Floating e) => NumericFloat S e + +instance (Storable e, IsList (Array L ix e), Ragged L ix e) => IsList (Array S ix e) where type Item (Array S ix e) = Item (Array L ix e) fromList = A.fromLists' Seq {-# INLINE fromList #-} @@ -233,82 +275,121 @@ -- referential transparency. -- -- @since 0.1.3 -unsafeWithPtr :: (MonadUnliftIO m, VS.Storable a) => Array S ix a -> (Ptr a -> m b) -> m b -unsafeWithPtr arr f = withRunInIO $ \run -> VS.unsafeWith (sData arr) (run . f) +unsafeWithPtr :: MonadUnliftIO m => Array S ix e -> (Ptr e -> m b) -> m b +unsafeWithPtr arr f = withRunInIO $ \run -> unsafeWithForeignPtr (sData arr) (run . f) {-# INLINE unsafeWithPtr #-} - -- | A pointer to the beginning of the mutable array. -- -- @since 0.1.3 -withPtr :: (MonadUnliftIO m, VS.Storable a) => MArray RealWorld S ix a -> (Ptr a -> m b) -> m b -withPtr (MSArray _ mv) f = withRunInIO $ \run -> MVS.unsafeWith mv (run . f) +withPtr :: MonadUnliftIO m => MArray RealWorld S ix e -> (Ptr e -> m b) -> m b +withPtr (MSArray _ fp) f = withRunInIO $ \run -> unsafeWithForeignPtr fp (run . f) {-# INLINE withPtr #-} - -- | /O(1)/ - Unwrap storable array and pull out the underlying storable vector. -- -- @since 0.2.1 -toStorableVector :: Array S ix e -> VS.Vector e -toStorableVector = sData +toStorableVector :: Index ix => Array S ix e -> VS.Vector e +toStorableVector arr = + -- this hack is needed to workaround the redundant Storable constraint + -- see haskell/vector#394 + unsafeCoerce $ + VS.unsafeFromForeignPtr0 (castForeignPtr (sData arr) :: ForeignPtr Word) (totalElem (sSize arr)) {-# INLINE toStorableVector #-} - -- | /O(1)/ - Unwrap storable mutable array and pull out the underlying storable mutable vector. -- -- @since 0.2.1 -toStorableMVector :: MArray s S ix e -> VS.MVector s e -toStorableMVector (MSArray _ mv) = mv +toStorableMVector :: Index ix => MArray s S ix e -> VS.MVector s e +toStorableMVector (MSArray sz fp) = MVS.MVector (totalElem sz) fp {-# INLINE toStorableMVector #-} +-- | /O(1)/ - Cast a storable vector to a storable array. +-- +-- @since 0.5.0 +fromStorableVector :: Comp -> VS.Vector e -> Vector S e +fromStorableVector comp v = + -- unasfeCoerce hack below is needed to workaround the redundant Storable + -- constraint see haskell/vector#394 + case VS.unsafeToForeignPtr0 (unsafeCoerce v :: VS.Vector Word) of + (fp, k) -> SArray{sComp = comp, sSize = SafeSz k, sData = castForeignPtr fp} +{-# INLINE fromStorableVector #-} +-- | /O(1)/ - Cast a mutable storable vector to a mutable storable array. +-- +-- @since 0.5.0 +fromStorableMVector :: MVS.MVector s e -> MVector s S e +fromStorableMVector (MVS.MVector n fp) = MSArray (SafeSz n) fp +{-# INLINE fromStorableMVector #-} + -- | /O(1)/ - Yield the underlying `ForeignPtr` together with its length. -- -- @since 0.3.0 -unsafeArrayToForeignPtr :: VS.Storable e => Array S ix e -> (ForeignPtr e, Int) -unsafeArrayToForeignPtr = VS.unsafeToForeignPtr0 . toStorableVector +unsafeArrayToForeignPtr :: Index ix => Array S ix e -> (ForeignPtr e, Int) +unsafeArrayToForeignPtr (SArray _ sz fp) = (fp, totalElem sz) {-# INLINE unsafeArrayToForeignPtr #-} -- | /O(1)/ - Yield the underlying `ForeignPtr` together with its length. -- -- @since 0.3.0 -unsafeMArrayToForeignPtr :: VS.Storable e => MArray s S ix e -> (ForeignPtr e, Int) -unsafeMArrayToForeignPtr = MVS.unsafeToForeignPtr0 . toStorableMVector +unsafeMArrayToForeignPtr :: Index ix => MArray s S ix e -> (ForeignPtr e, Int) +unsafeMArrayToForeignPtr (MSArray sz fp) = (fp, totalElem sz) {-# INLINE unsafeMArrayToForeignPtr #-} -- | /O(1)/ - Wrap a `ForeignPtr` and it's size into a pure storable array. -- -- @since 0.3.0 -unsafeArrayFromForeignPtr0 :: VS.Storable e => Comp -> ForeignPtr e -> Sz1 -> Array S Ix1 e -unsafeArrayFromForeignPtr0 comp ptr sz = - SArray {sComp = comp, sSize = sz, sData = VS.unsafeFromForeignPtr0 ptr (unSz sz)} +unsafeArrayFromForeignPtr0 :: Comp -> ForeignPtr e -> Sz1 -> Vector S e +unsafeArrayFromForeignPtr0 comp fp sz = SArray{sComp = comp, sSize = sz, sData = fp} {-# INLINE unsafeArrayFromForeignPtr0 #-} -- | /O(1)/ - Wrap a `ForeignPtr`, an offset and it's size into a pure storable array. -- -- @since 0.3.0 -unsafeArrayFromForeignPtr :: VS.Storable e => Comp -> ForeignPtr e -> Int -> Sz1 -> Array S Ix1 e +unsafeArrayFromForeignPtr :: Storable e => Comp -> ForeignPtr e -> Int -> Sz1 -> Array S Ix1 e unsafeArrayFromForeignPtr comp ptr offset sz = - SArray {sComp = comp, sSize = sz, sData = VS.unsafeFromForeignPtr ptr offset (unSz sz)} + SArray{sComp = comp, sSize = sz, sData = advanceForeignPtr ptr offset} {-# INLINE unsafeArrayFromForeignPtr #-} - -- | /O(1)/ - Wrap a `ForeignPtr` and it's size into a mutable storable array. It is still safe to -- modify the pointer, unless the array gets frozen prior to modification. -- -- @since 0.3.0 -unsafeMArrayFromForeignPtr0 :: VS.Storable e => ForeignPtr e -> Sz1 -> MArray s S Ix1 e -unsafeMArrayFromForeignPtr0 fp sz = - MSArray sz (MVS.unsafeFromForeignPtr0 fp (unSz sz)) +unsafeMArrayFromForeignPtr0 :: ForeignPtr e -> Sz1 -> MArray s S Ix1 e +unsafeMArrayFromForeignPtr0 fp sz = MSArray sz fp {-# INLINE unsafeMArrayFromForeignPtr0 #-} - -- | /O(1)/ - Wrap a `ForeignPtr`, an offset and it's size into a mutable storable array. It is -- still safe to modify the pointer, unless the array gets frozen prior to modification. -- -- @since 0.3.0 -unsafeMArrayFromForeignPtr :: VS.Storable e => ForeignPtr e -> Int -> Sz1 -> MArray s S Ix1 e -unsafeMArrayFromForeignPtr fp offset sz = - MSArray sz (MVS.unsafeFromForeignPtr fp offset (unSz sz)) +unsafeMArrayFromForeignPtr :: Storable e => ForeignPtr e -> Int -> Sz1 -> MArray s S Ix1 e +unsafeMArrayFromForeignPtr fp offset sz = MSArray sz (advanceForeignPtr fp offset) {-# INLINE unsafeMArrayFromForeignPtr #-} +-- | Allocate memory using @malloc@ on C heap, instead of on Haskell heap. Memory is left +-- uninitialized +-- +-- @since 0.5.9 +unsafeMallocMArray + :: forall ix e m + . (Index ix, Storable e, PrimMonad m) + => Sz ix + -> m (MArray (PrimState m) S ix e) +unsafeMallocMArray sz = unsafePrimToPrim $ do + let n = totalElem sz + foreignPtr <- mask_ $ do + ptr <- mallocBytes (sizeOf (undefined :: e) * n) + newForeignPtr finalizerFree ptr + pure $ MSArray sz foreignPtr +{-# INLINE unsafeMallocMArray #-} + +#if !MIN_VERSION_base(4,15,0) +-- | A compatibility wrapper for 'GHC.ForeignPtr.unsafeWithForeignPtr' provided +-- by GHC 9.0.1 and later. +-- +-- Only to be used when the continuation is known not to +-- unconditionally diverge lest unsoundness can result. +unsafeWithForeignPtr :: ForeignPtr a -> (Ptr a -> IO b) -> IO b +unsafeWithForeignPtr = withForeignPtr +#endif 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 @@ -1,54 +1,59 @@ {-# LANGUAGE BangPatterns #-} -{-# LANGUAGE CPP #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE NamedFieldPuns #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UndecidableInstances #-} + -- | -- Module : Data.Massiv.Array.Manifest.Unboxed --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Manifest.Unboxed - ( U (..) - , VU.Unbox - , Array(..) - , toUnboxedVector - , toUnboxedMVector - ) where +module Data.Massiv.Array.Manifest.Unboxed ( + U (..), + Unbox, + Array (..), + MArray (..), + toUnboxedVector, + toUnboxedMVector, + fromUnboxedVector, + fromUnboxedMVector, +) where -import Control.DeepSeq (NFData(..), deepseq) -import Data.Massiv.Array.Delayed.Pull (eq, ord) -import Data.Massiv.Array.Manifest.Internal (M, toManifest) +import Control.DeepSeq (NFData (..), deepseq) +import Control.Monad.Primitive (stToPrim) +import Data.Massiv.Array.Delayed.Pull (D, compareArrays, eqArrays) +import Data.Massiv.Array.Manifest.Internal import Data.Massiv.Array.Manifest.List as A -import Data.Massiv.Array.Manifest.Vector.Stream as S (steps) import Data.Massiv.Array.Mutable import Data.Massiv.Core.Common import Data.Massiv.Core.List +import Data.Massiv.Core.Operations +import Data.Massiv.Vector.Stream as S (isteps, steps) import qualified Data.Vector.Generic.Mutable as VGM +import Data.Vector.Unboxed (Unbox) import qualified Data.Vector.Unboxed as VU import qualified Data.Vector.Unboxed.Mutable as MVU -import GHC.Exts as GHC (IsList(..)) -import Prelude hiding (mapM) +import GHC.Exts as GHC (IsList (..)) import System.IO.Unsafe (unsafePerformIO) - -#include "massiv.h" +import Prelude hiding (mapM) -- | Representation for `Unbox`ed elements -data U = U deriving Show +data U = U deriving (Show) -data instance Array U ix e = UArray { uComp :: !Comp - , uSize :: !(Sz ix) - , uData :: !(VU.Vector e) - } +data instance Array U ix e = UArray + { uComp :: !Comp + , uSize :: !(Sz ix) + , uData :: !(VU.Vector e) + } -instance (Ragged L ix e, Show e, VU.Unbox e) => Show (Array U ix e) where +data instance MArray s U ix e = MUArray !(Sz ix) !(VU.MVector s e) + +instance (Ragged L ix e, Show e, Unbox e) => Show (Array U ix e) where showsPrec = showsArrayPrec id showList = showArrayList @@ -56,106 +61,77 @@ rnf (UArray c sz v) = c `deepseq` sz `deepseq` v `deepseq` () {-# INLINE rnf #-} +instance NFData ix => NFData (MArray s U ix e) where + rnf (MUArray sz mv) = sz `deepseq` mv `deepseq` () + {-# INLINE rnf #-} -instance (VU.Unbox e, Index ix) => Construct U ix e where - setComp c arr = arr { uComp = c } +instance Strategy U where + getComp = uComp + {-# INLINE getComp #-} + setComp c arr = arr{uComp = c} {-# INLINE setComp #-} - - makeArray !comp !sz f = unsafePerformIO $ generateArray comp sz (return . f) - {-# INLINE makeArray #-} - + repr = U -instance (VU.Unbox e, Eq e, Index ix) => Eq (Array U ix e) where - (==) = eq (==) +instance (Unbox e, Eq e, Index ix) => Eq (Array U ix e) where + (==) = eqArrays (==) {-# INLINE (==) #-} -instance (VU.Unbox e, Ord e, Index ix) => Ord (Array U ix e) where - compare = ord compare +instance (Unbox e, Ord e, Index ix) => Ord (Array U ix e) where + compare = compareArrays compare {-# INLINE compare #-} - -instance (VU.Unbox e, Index ix) => Source U ix e where +instance Unbox e => Source U e where unsafeLinearIndex (UArray _ _ v) = - INDEX_CHECK("(Source U ix e).unsafeLinearIndex", Sz . VU.length, VU.unsafeIndex) v + indexAssert "U.unsafeLinearIndex" (SafeSz . VU.length) VU.unsafeIndex v {-# INLINE unsafeLinearIndex #-} + unsafeOuterSlice (UArray c _ v) szL i = + let k = totalElem szL + in UArray c szL $ VU.unsafeSlice (i * k) k v + {-# INLINE unsafeOuterSlice #-} -instance Index ix => Resize U ix where - unsafeResize !sz !arr = arr { uSize = sz } - {-# INLINE unsafeResize #-} + unsafeLinearSlice i k (UArray c _ v) = UArray c k $ VU.unsafeSlice i (unSz k) v + {-# INLINE unsafeLinearSlice #-} -instance (VU.Unbox e, Index ix) => Extract U ix e where - unsafeExtract !sIx !newSz !arr = unsafeExtract sIx newSz (toManifest arr) - {-# INLINE unsafeExtract #-} +instance Index ix => Shape U ix where + maxLinearSize = Just . SafeSz . elemsCount + {-# INLINE maxLinearSize #-} -instance (VU.Unbox e, Index ix) => Load U ix e where - type R U = M +instance Size U where size = uSize {-# INLINE size #-} - getComp = uComp - {-# INLINE getComp #-} - loadArrayM !scheduler !arr = splitLinearlyWith_ scheduler (elemsCount arr) (unsafeLinearIndex arr) - {-# INLINE loadArrayM #-} - -instance (VU.Unbox e, Index ix) => StrideLoad U ix e - - -instance {-# OVERLAPPING #-} VU.Unbox e => Slice U Ix1 e where - unsafeSlice arr i _ _ = pure (unsafeLinearIndex arr i) - {-# INLINE unsafeSlice #-} - - -instance ( VU.Unbox e - , Index ix - , Index (Lower ix) - , Elt U ix e ~ Elt M ix e - , Elt M ix e ~ Array M (Lower ix) e - ) => - Slice U ix e where - unsafeSlice arr = unsafeSlice (toManifest arr) - {-# INLINE unsafeSlice #-} - + unsafeResize !sz !arr = arr{uSize = sz} + {-# INLINE unsafeResize #-} -instance {-# OVERLAPPING #-} VU.Unbox e => OuterSlice U Ix1 e where - unsafeOuterSlice = unsafeLinearIndex - {-# INLINE unsafeOuterSlice #-} +instance (Unbox e, Index ix) => Load U ix e where + makeArray comp sz f = compute (makeArray comp sz f :: Array D ix e) + {-# INLINE makeArray #-} -instance ( VU.Unbox e - , Index ix - , Index (Lower ix) - , Elt U ix e ~ Elt M ix e - , Elt M ix e ~ Array M (Lower ix) e - ) => - OuterSlice U ix e where - unsafeOuterSlice arr = unsafeOuterSlice (toManifest arr) - {-# INLINE unsafeOuterSlice #-} + makeArrayLinear !comp !sz f = unsafePerformIO $ generateArrayLinear comp sz (pure . f) + {-# INLINE makeArrayLinear #-} -instance {-# OVERLAPPING #-} VU.Unbox e => InnerSlice U Ix1 e where - unsafeInnerSlice arr _ = unsafeLinearIndex arr - {-# INLINE unsafeInnerSlice #-} + replicate comp !sz !e = runST (newMArray sz e >>= unsafeFreeze comp) + {-# INLINE replicate #-} -instance ( VU.Unbox e - , Index ix - , Index (Lower ix) - , Elt U ix e ~ Elt M ix e - , Elt M ix e ~ Array M (Lower ix) e - ) => - InnerSlice U ix e where - unsafeInnerSlice arr = unsafeInnerSlice (toManifest arr) - {-# INLINE unsafeInnerSlice #-} + iterArrayLinearST_ !scheduler !arr = + splitLinearlyWith_ scheduler (elemsCount arr) (unsafeLinearIndex arr) + {-# INLINE iterArrayLinearST_ #-} -instance (VU.Unbox e, Index ix) => Manifest U ix e where +instance (Unbox e, Index ix) => StrideLoad U ix e +instance Unbox e => Manifest U e where unsafeLinearIndexM (UArray _ _ v) = - INDEX_CHECK("(Manifest U ix e).unsafeLinearIndexM", Sz . VU.length, VU.unsafeIndex) v + indexAssert "S.unsafeLinearIndexM" (SafeSz . VU.length) VU.unsafeIndex v {-# INLINE unsafeLinearIndexM #-} + sizeOfMArray (MUArray sz _) = sz + {-# INLINE sizeOfMArray #-} -instance (VU.Unbox e, Index ix) => Mutable U ix e where - data MArray s U ix e = MUArray !(Sz ix) !(VU.MVector s e) + unsafeResizeMArray sz (MUArray _ mv) = MUArray sz mv + {-# INLINE unsafeResizeMArray #-} - msize (MUArray sz _) = sz - {-# INLINE msize #-} + unsafeLinearSliceMArray i k (MUArray _ mv) = MUArray k $ MVU.unsafeSlice i (unSz k) mv + {-# INLINE unsafeLinearSliceMArray #-} unsafeThaw (UArray _ sz v) = MUArray sz <$> VU.unsafeThaw v {-# INLINE unsafeThaw #-} @@ -166,7 +142,7 @@ unsafeNew sz = MUArray sz <$> MVU.unsafeNew (totalElem sz) {-# INLINE unsafeNew #-} - initialize (MUArray _ marr) = VGM.basicInitialize marr + initialize (MUArray _ marr) = stToPrim (VGM.basicInitialize marr) {-# INLINE initialize #-} unsafeLinearCopy (MUArray _ mvFrom) iFrom (MUArray _ mvTo) iTo (Sz k) = @@ -174,36 +150,43 @@ {-# INLINE unsafeLinearCopy #-} unsafeLinearRead (MUArray _ mv) = - INDEX_CHECK("(Mutable U ix e).unsafeLinearRead", Sz . MVU.length, MVU.unsafeRead) mv + indexAssert "U.unsafeLinearRead" (Sz . MVU.length) MVU.unsafeRead mv {-# INLINE unsafeLinearRead #-} unsafeLinearWrite (MUArray _ mv) = - INDEX_CHECK("(Mutable U ix e).unsafeLinearWrite", Sz . MVU.length, MVU.unsafeWrite) mv + indexAssert "U.unsafeLinearWrite" (Sz . MVU.length) MVU.unsafeWrite mv {-# INLINE unsafeLinearWrite #-} unsafeLinearGrow (MUArray _ mv) sz = MUArray sz <$> MVU.unsafeGrow mv (totalElem sz) {-# INLINE unsafeLinearGrow #-} - -instance (Index ix, VU.Unbox e) => Stream U ix e where +instance (Index ix, Unbox e) => Stream U ix e where toStream = S.steps {-# INLINE toStream #-} - + toStreamIx = S.isteps + {-# INLINE toStreamIx #-} -instance ( VU.Unbox e - , IsList (Array L ix e) - , Nested LN ix e - , Nested L ix e - , Ragged L ix e - ) => - IsList (Array U ix e) where +instance (Unbox e, IsList (Array L ix e), Ragged L ix e) => IsList (Array U ix e) where type Item (Array U ix e) = Item (Array L ix e) fromList = A.fromLists' Seq {-# INLINE fromList #-} toList = GHC.toList . toListArray {-# INLINE toList #-} +instance (VU.Unbox e, Num e) => FoldNumeric U e where + unsafeDotProduct = defaultUnsafeDotProduct + {-# INLINE unsafeDotProduct #-} + powerSumArray = defaultPowerSumArray + {-# INLINE powerSumArray #-} + foldArray = defaultFoldArray + {-# INLINE foldArray #-} +instance (VU.Unbox e, Num e) => Numeric U e where + unsafeLiftArray = defaultUnsafeLiftArray + {-# INLINE unsafeLiftArray #-} + unsafeLiftArray2 = defaultUnsafeLiftArray2 + {-# INLINE unsafeLiftArray2 #-} + -- | /O(1)/ - Unwrap unboxed array and pull out the underlying unboxed vector. -- -- @since 0.2.1 @@ -211,10 +194,23 @@ toUnboxedVector = uData {-# INLINE toUnboxedVector #-} - -- | /O(1)/ - Unwrap unboxed mutable array and pull out the underlying unboxed mutable vector. -- -- @since 0.2.1 toUnboxedMVector :: MArray s U ix e -> VU.MVector s e toUnboxedMVector (MUArray _ mv) = mv {-# INLINE toUnboxedMVector #-} + +-- | /O(1)/ - Wrap an unboxed vector and produce an unboxed flat array. +-- +-- @since 0.6.0 +fromUnboxedVector :: VU.Unbox e => Comp -> VU.Vector e -> Vector U e +fromUnboxedVector comp v = UArray comp (SafeSz (VU.length v)) v +{-# INLINE fromUnboxedVector #-} + +-- | /O(1)/ - Wrap an unboxed mutable vector and produce a mutable unboxed flat array. +-- +-- @since 0.5.0 +fromUnboxedMVector :: Unbox e => VU.MVector s e -> MVector s U e +fromUnboxedMVector mv = MUArray (SafeSz (MVU.length mv)) mv +{-# INLINE fromUnboxedMVector #-} diff --git a/src/Data/Massiv/Array/Manifest/Vector.hs b/src/Data/Massiv/Array/Manifest/Vector.hs --- a/src/Data/Massiv/Array/Manifest/Vector.hs +++ b/src/Data/Massiv/Array/Manifest/Vector.hs @@ -4,31 +4,31 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE TypeOperators #-} + -- | -- Module : Data.Massiv.Array.Manifest.Vector --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Manifest.Vector - ( fromVectorM - , fromVector' - , castFromVector - , toVector - , castToVector - , ARepr - , VRepr - ) where +module Data.Massiv.Array.Manifest.Vector ( + fromVectorM, + fromVector', + castFromVector, + toVector, + castToVector, + ARepr, + VRepr, +) where import Control.Monad (guard, join, msum) +import Data.Kind import Data.Massiv.Array.Manifest.Boxed import Data.Massiv.Array.Manifest.Internal import Data.Massiv.Array.Manifest.Primitive import Data.Massiv.Array.Manifest.Storable import Data.Massiv.Array.Manifest.Unboxed -import Data.Massiv.Array.Mutable import Data.Massiv.Core.Common import Data.Maybe (fromMaybe) import Data.Typeable @@ -39,66 +39,69 @@ import qualified Data.Vector.Unboxed as VU -- | Match vector type to array representation -type family ARepr (v :: * -> *) :: * where +type family ARepr (v :: Type -> Type) :: Type where ARepr VU.Vector = U ARepr VS.Vector = S ARepr VP.Vector = P - ARepr VB.Vector = B + ARepr VB.Vector = BL -- | Match array representation to a vector type -type family VRepr r :: * -> * where +type family VRepr r :: Type -> Type where VRepr U = VU.Vector VRepr S = VS.Vector VRepr P = VP.Vector VRepr B = VB.Vector - VRepr N = VB.Vector - + VRepr BN = VB.Vector + VRepr BL = VB.Vector --- | /O(1)/ - conversion from vector to an array with a corresponding --- representation. Will return `Nothing` if there is a size mismatch, vector has --- been sliced before or if some non-standard vector type is supplied. -castFromVector :: forall v r ix e. (VG.Vector v e, Typeable v, Mutable r ix e, ARepr v ~ r) - => Comp - -> Sz ix -- ^ Size of the result Array - -> v e -- ^ Source Vector - -> Maybe (Array r ix e) +-- | /O(1)/ - conversion from vector to an array with a corresponding representation. Will +-- return `Nothing` if there is a size mismatch or if some non-standard vector type is +-- supplied. Is suppplied is the boxed `Data.Vector.Vector` then it's all elements will be +-- evaluated toWHNF, therefore complexity will be /O(n)/ +castFromVector + :: forall v r ix e + . (VG.Vector v e, Typeable v, Index ix, ARepr v ~ r) + => Comp + -> Sz ix + -- ^ Size of the result Array + -> v e + -- ^ Source Vector + -> Maybe (Array r ix e) castFromVector comp sz vector = do guard (totalElem sz == VG.length vector) msum - [ do Refl <- eqT :: Maybe (v :~: VU.Vector) - uVector <- join $ gcast1 (Just vector) - return $ UArray {uComp = comp, uSize = sz, uData = uVector} - , do Refl <- eqT :: Maybe (v :~: VS.Vector) - sVector <- join $ gcast1 (Just vector) - return $ SArray {sComp = comp, sSize = sz, sData = sVector} - , do Refl <- eqT :: Maybe (v :~: VP.Vector) - VP.Vector 0 _ arr <- join $ gcast1 (Just vector) - return $ PArray {pComp = comp, pSize = sz, pData = arr} - , do Refl <- eqT :: Maybe (v :~: VB.Vector) - bVector <- join $ gcast1 (Just vector) - arr <- castVectorToArray bVector - let barr = BArray {bComp = comp, bSize = sz, bData = arr} - barr `seqArray` return barr + [ do + Refl <- eqT :: Maybe (v :~: VU.Vector) + uVector <- join $ gcast1 (Just vector) + return $ UArray{uComp = comp, uSize = sz, uData = uVector} + , do + Refl <- eqT :: Maybe (v :~: VS.Vector) + sVector <- join $ gcast1 (Just vector) + return $ unsafeResize sz $ fromStorableVector comp sVector + , do + Refl <- eqT :: Maybe (v :~: VP.Vector) + VP.Vector o _ ba <- join $ gcast1 (Just vector) + return $ PArray{pComp = comp, pSize = sz, pOffset = o, pData = ba} + , do + Refl <- eqT :: Maybe (v :~: VB.Vector) + bVector <- join $ gcast1 (Just vector) + pure $ unsafeResize sz $ setComp comp $ fromBoxedVector bVector ] {-# NOINLINE castFromVector #-} - -- | In case when resulting array representation matches the one of vector's it -- will do a /O(1)/ - conversion using `castFromVector`, otherwise Vector elements -- will be copied into a new array. Will throw an error if length of resulting -- array doesn't match the source vector length. -- -- @since 0.3.0 -fromVectorM :: - ( MonadThrow m - , Typeable v - , VG.Vector v a - , Mutable (ARepr v) ix a - , Mutable r ix a - ) +fromVectorM + :: (MonadThrow m, Typeable v, VG.Vector v a, Manifest r a, Load (ARepr v) ix a, Load r ix a) => Comp - -> Sz ix -- ^ Resulting size of the array - -> v a -- ^ Source Vector + -> Sz ix + -- ^ Resulting size of the array + -> v a + -- ^ Source Vector -> m (Array r ix a) fromVectorM comp sz v = case castFromVector comp sz v of @@ -108,45 +111,57 @@ pure (makeArrayLinear comp sz (VG.unsafeIndex v)) {-# NOINLINE fromVectorM #-} - -- | Just like `fromVectorM`, but will throw an exception on a mismatched size. -- -- @since 0.3.0 -fromVector' :: - (Typeable v, VG.Vector v a, Mutable (ARepr v) ix a, Mutable r ix a) +fromVector' + :: (HasCallStack, Typeable v, VG.Vector v a, Load (ARepr v) ix a, Load r ix a, Manifest r a) => Comp - -> Sz ix -- ^ Resulting size of the array - -> v a -- ^ Source Vector + -> Sz ix + -- ^ Resulting size of the array + -> v a + -- ^ Source Vector -> Array r ix a -fromVector' comp sz = either throw id . fromVectorM comp sz +fromVector' comp sz = throwEither . fromVectorM comp sz {-# INLINE fromVector' #-} -- | /O(1)/ - conversion from `Mutable` array to a corresponding vector. Will -- return `Nothing` only if source array representation was not one of `B`, `N`, -- `P`, `S` or `U`. -castToVector :: forall v r ix e . (Mutable r ix e, VRepr r ~ v) - => Array r ix e -> Maybe (v e) +castToVector + :: forall v r ix e + . (Manifest r e, Index ix, VRepr r ~ v) + => Array r ix e + -> Maybe (v e) castToVector arr = msum - [ do Refl <- eqT :: Maybe (r :~: U) - uArr <- gcastArr arr - return $ uData uArr - , do Refl <- eqT :: Maybe (r :~: S) - sArr <- gcastArr arr - return $ sData sArr - , do Refl <- eqT :: Maybe (r :~: P) - pArr <- gcastArr arr - return $ VP.Vector 0 (totalElem (size arr)) $ pData pArr - , do Refl <- eqT :: Maybe (r :~: B) - bArr <- gcastArr arr - return $ castArrayToVector $ bData bArr - , do Refl <- eqT :: Maybe (r :~: N) - bArr <- gcastArr arr - return $ castArrayToVector $ bData $ bArray bArr + [ do + Refl <- eqT :: Maybe (r :~: U) + uArr <- gcastArr arr + return $ uData uArr + , do + Refl <- eqT :: Maybe (r :~: S) + sArr <- gcastArr arr + return $ toStorableVector sArr + , do + Refl <- eqT :: Maybe (r :~: P) + pArr <- gcastArr arr + return $ VP.Vector (pOffset pArr) (totalElem (size arr)) $ pData pArr + , do + Refl <- eqT :: Maybe (r :~: B) + bArr <- gcastArr arr + return $ toBoxedVector $ toLazyArray bArr + , do + Refl <- eqT :: Maybe (r :~: BN) + bArr <- gcastArr arr + return $ toBoxedVector $ toLazyArray $ unwrapNormalForm bArr + , do + Refl <- eqT :: Maybe (r :~: BL) + bArr <- gcastArr arr + return $ toBoxedVector bArr ] {-# NOINLINE castToVector #-} - -- | Convert an array into a vector. Will perform a cast if resulting vector is -- of compatible representation, otherwise memory copy will occur. -- @@ -156,6 +171,7 @@ -- `VS.Vector` in costant time: -- -- >>> import Data.Massiv.Array as A +-- >>> import Data.Massiv.Array.Manifest.Vector (toVector) -- >>> import qualified Data.Vector.Storable as VS -- >>> toVector (makeArrayR S Par (Sz2 5 6) (\(i :. j) -> i + j)) :: VS.Vector Int -- [0,1,2,3,4,5,1,2,3,4,5,6,2,3,4,5,6,7,3,4,5,6,7,8,4,5,6,7,8,9] @@ -167,11 +183,11 @@ -- >>> import qualified Data.Vector.Unboxed as VU -- >>> toVector (makeArrayR S Par (Sz2 5 6) (\(i :. j) -> i + j)) :: VU.Vector Int -- [0,1,2,3,4,5,1,2,3,4,5,6,2,3,4,5,6,7,3,4,5,6,7,8,4,5,6,7,8,9] --- -toVector :: - forall r ix e v. - ( Manifest r ix e - , Mutable (ARepr v) ix e +toVector + :: forall r ix e v + . ( Manifest r e + , Load r ix e + , Manifest (ARepr v) e , VG.Vector v e , VRepr (ARepr v) ~ v ) @@ -182,4 +198,3 @@ (VG.generate (totalElem (size arr)) (unsafeLinearIndex arr)) (castToVector (convert arr :: Array (ARepr v) ix e)) {-# NOINLINE toVector #-} - diff --git a/src/Data/Massiv/Array/Manifest/Vector/Stream.hs b/src/Data/Massiv/Array/Manifest/Vector/Stream.hs deleted file mode 100644 --- a/src/Data/Massiv/Array/Manifest/Vector/Stream.hs +++ /dev/null @@ -1,407 +0,0 @@ -{-# LANGUAGE TupleSections #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} --- | --- Module : Data.Massiv.Array.Manifest.Vector.Stream --- Copyright : (c) Alexey Kuleshevich 2019 --- License : BSD3 --- Maintainer : Alexey Kuleshevich --- Stability : experimental --- Portability : non-portable --- -module Data.Massiv.Array.Manifest.Vector.Stream - ( -- | __Important__ - This module is still experimental, as such it is considered - -- internal and exported for the curious users only. - Steps(..) - , Stream(..) - -- * Conversion - , steps - , isteps - , fromStream - , fromStreamM - , fromStreamExactM - , unstreamExact - , unstreamMax - , unstreamMaxM - , unstreamUnknown - , unstreamUnknownM - , unstreamIntoM - -- * Bundle - , toBundle - , fromBundle - , fromBundleM - -- * Operations on Steps - , length - , empty - , singleton - , generate - , cons - , uncons - , snoc - , drop - , take - , slice - , traverse - , mapM - , concatMap - , append - , zipWith - , zipWithM - -- ** Folding - , foldl - , foldr - , foldlM - , foldrM - -- ** Unfolding - , unfoldr - , unfoldrN - -- * Lists - , toList - , fromList - , fromListN - -- ** Filter - , mapMaybe - , mapMaybeA - , mapMaybeM - , filter - , filterA - , filterM - , transStepsId - -- * Useful re-exports - , module Data.Vector.Fusion.Bundle.Size - , module Data.Vector.Fusion.Util - ) where - -import Data.Maybe (catMaybes) -import qualified Control.Monad as M -import Control.Monad.ST -import Data.Massiv.Core.Common hiding (empty, singleton) -import qualified Data.Traversable as Traversable (traverse) -import qualified Data.Vector.Fusion.Bundle.Monadic as B -import Data.Vector.Fusion.Bundle.Size -import qualified Data.Vector.Fusion.Stream.Monadic as S -import Data.Vector.Fusion.Util -import Prelude hiding (zipWith, mapM, traverse, length, foldl, foldr, filter, concatMap, drop, take) - - --- TODO: benchmark: `fmap snd . isteps` -steps :: forall r ix e m . (Monad m, Source r ix e) => Array r ix e -> Steps m e -steps arr = k `seq` arr `seq` Steps (S.Stream step 0) (Exact k) - where - k = totalElem $ size arr - step i - | i < k = - let e = unsafeLinearIndex arr i - in e `seq` return $ S.Yield e (i + 1) - | otherwise = return S.Done - {-# INLINE step #-} -{-# INLINE steps #-} - - -isteps :: forall r ix e m . (Monad m, Source r ix e) => Array r ix e -> Steps m (ix, e) -isteps arr = k `seq` arr `seq` Steps (S.Stream step 0) (Exact k) - where - sz = size arr - k = totalElem sz - step i - | i < k = - let e = unsafeLinearIndex arr i - in e `seq` return $ S.Yield (fromLinearIndex sz i, e) (i + 1) - | otherwise = return S.Done - {-# INLINE step #-} -{-# INLINE isteps #-} - -toBundle :: (Monad m, Source r ix e) => Array r ix e -> B.Bundle m v e -toBundle arr = - let Steps str k = steps arr - in B.fromStream str k -{-# INLINE toBundle #-} - -fromBundle :: Mutable r Ix1 e => B.Bundle Id v e -> Array r Ix1 e -fromBundle bundle = fromStream (B.sSize bundle) (B.sElems bundle) -{-# INLINE fromBundle #-} - - -fromBundleM :: (Monad m, Mutable r Ix1 e) => B.Bundle m v e -> m (Array r Ix1 e) -fromBundleM bundle = fromStreamM (B.sSize bundle) (B.sElems bundle) -{-# INLINE fromBundleM #-} - - -fromStream :: forall r e . Mutable r Ix1 e => Size -> S.Stream Id e -> Array r Ix1 e -fromStream sz str = - case upperBound sz of - Nothing -> unstreamUnknown str - Just k -> unstreamMax k str -{-# INLINE fromStream #-} - -fromStreamM :: forall r e m. (Monad m, Mutable r Ix1 e) => Size -> S.Stream m e -> m (Array r Ix1 e) -fromStreamM sz str = do - xs <- S.toList str - case upperBound sz of - Nothing -> pure $! unstreamUnknown (S.fromList xs) - Just k -> pure $! unstreamMax k (S.fromList xs) -{-# INLINE fromStreamM #-} - -fromStreamExactM :: - forall r ix e m. (Monad m, Mutable r ix e) - => Sz ix - -> S.Stream m e - -> m (Array r ix e) -fromStreamExactM sz str = do - xs <- S.toList str - pure $! unstreamExact sz (S.fromList xs) -{-# INLINE fromStreamExactM #-} - - -unstreamIntoM :: - (Mutable r Ix1 a, PrimMonad m) - => MArray (PrimState m) r Ix1 a - -> Size - -> S.Stream Id a - -> m (MArray (PrimState m) r Ix1 a) -unstreamIntoM marr sz str = - case sz of - Exact _ -> marr <$ unstreamMaxM marr str - Max _ -> unsafeLinearShrink marr . SafeSz =<< unstreamMaxM marr str - Unknown -> unstreamUnknownM marr str -{-# INLINE unstreamIntoM #-} - - - -unstreamMax :: - forall r e. (Mutable r Ix1 e) - => Int - -> S.Stream Id e - -> Array r Ix1 e -unstreamMax kMax str = - runST $ do - marr <- unsafeNew (SafeSz kMax) - k <- unstreamMaxM marr str - unsafeLinearShrink marr (SafeSz k) >>= unsafeFreeze Seq -{-# INLINE unstreamMax #-} - - -unstreamMaxM :: - (Mutable r ix a, PrimMonad m) => MArray (PrimState m) r ix a -> S.Stream Id a -> m Int -unstreamMaxM marr (S.Stream step s) = stepLoad s 0 - where - stepLoad t i = - case unId (step t) of - S.Yield e' t' -> do - unsafeLinearWrite marr i e' - stepLoad t' (i + 1) - S.Skip t' -> stepLoad t' i - S.Done -> return i - {-# INLINE stepLoad #-} -{-# INLINE unstreamMaxM #-} - - -unstreamUnknown :: Mutable r Ix1 a => S.Stream Id a -> Array r Ix1 a -unstreamUnknown str = - runST $ do - marr <- unsafeNew zeroSz - unstreamUnknownM marr str >>= unsafeFreeze Seq -{-# INLINE unstreamUnknown #-} - - -unstreamUnknownM :: - (Mutable r Ix1 a, PrimMonad m) - => MArray (PrimState m) r Ix1 a - -> S.Stream Id a - -> m (MArray (PrimState m) r Ix1 a) -unstreamUnknownM marrInit (S.Stream step s) = stepLoad s 0 (unSz (msize marrInit)) marrInit - where - stepLoad t i kMax marr - | i < kMax = - case unId (step t) of - S.Yield e' t' -> do - unsafeLinearWrite marr i e' - stepLoad t' (i + 1) kMax marr - S.Skip t' -> stepLoad t' i kMax marr - S.Done -> unsafeLinearShrink marr (SafeSz i) - | otherwise = do - let kMax' = max 1 (kMax * 2) - marr' <- unsafeLinearGrow marr (SafeSz kMax') - stepLoad t i kMax' marr' - {-# INLINE stepLoad #-} -{-# INLINE unstreamUnknownM #-} - - -unstreamExact :: - forall r ix e. (Mutable r ix e) - => Sz ix - -> S.Stream Id e - -> Array r ix e -unstreamExact sz str = - runST $ do - marr <- unsafeNew sz - _ <- unstreamMaxM marr str - unsafeFreeze Seq marr -{-# INLINE unstreamExact #-} - -length :: Steps Id a -> Int -length (Steps str sz) = - case sz of - Exact k -> k - _ -> unId (S.length str) -{-# INLINE length #-} - -empty :: Monad m => Steps m e -empty = Steps S.empty (Exact 0) -{-# INLINE empty #-} - -singleton :: Monad m => e -> Steps m e -singleton e = Steps (S.singleton e) (Exact 1) -{-# INLINE singleton #-} - -generate :: Monad m => Int -> (Int -> e) -> Steps m e -generate k f = Steps (S.generate k f) (Exact k) -{-# INLINE generate #-} - -cons :: Monad m => e -> Steps m e -> Steps m e -cons e (Steps str k) = Steps (S.cons e str) (k + 1) -{-# INLINE cons #-} - -uncons :: Monad m => Steps m e -> m (Maybe (e, Steps m e)) -uncons sts@(Steps str _) = do - mx <- str S.!? 0 - pure $ fmap (, drop 1 sts) mx -{-# INLINE uncons #-} - -snoc :: Monad m => Steps m e -> e -> Steps m e -snoc (Steps str k) e = Steps (S.snoc str e) (k + 1) -{-# INLINE snoc #-} - -traverse :: (Monad m, Applicative f) => (e -> f a) -> Steps Id e -> f (Steps m a) -traverse f (Steps str k) = (`Steps` k) <$> liftListA (Traversable.traverse f) str -{-# INLINE traverse #-} - -append :: Monad m => Steps m e -> Steps m e -> Steps m e -append (Steps str1 k1) (Steps str2 k2) = Steps (str1 S.++ str2) (k1 + k2) -{-# INLINE append #-} - -mapM :: Monad m => (e -> m a) -> Steps m e -> Steps m a -mapM f (Steps str k) = Steps (S.mapM f str) k -{-# INLINE mapM #-} - -zipWith :: Monad m => (a -> b -> e) -> Steps m a -> Steps m b -> Steps m e -zipWith f (Steps str1 k1) (Steps str2 k2) = Steps (S.zipWith f str1 str2) (smaller k1 k2) -{-# INLINE zipWith #-} - -zipWithM :: Monad m => (a -> b -> m c) -> Steps m a -> Steps m b -> Steps m c -zipWithM f (Steps str1 k1) (Steps str2 k2) = Steps (S.zipWithM f str1 str2) (smaller k1 k2) -{-# INLINE zipWithM #-} - -transStepsId :: Monad m => Steps Id e -> Steps m e -transStepsId (Steps sts k) = Steps (S.trans (pure . unId) sts) k -{-# INLINE transStepsId #-} - - -foldr :: (a -> b -> b) -> b -> Steps Id a -> b -foldr f acc sts = unId (S.foldr f acc (stepsStream sts)) -{-# INLINE foldr #-} - - -foldl :: (b -> a -> b) -> b -> Steps Id a -> b -foldl f acc sts = unId (S.foldl f acc (stepsStream sts)) -{-# INLINE foldl #-} - - -foldlM :: Monad m => (a -> b -> m a) -> a -> Steps m b -> m a -foldlM f acc (Steps sts _) = S.foldlM f acc sts -{-# INLINE foldlM #-} - - -foldrM :: Monad m => (b -> a -> m a) -> a -> Steps m b -> m a -foldrM f acc (Steps sts _) = S.foldrM f acc sts -{-# INLINE foldrM #-} - - -mapMaybe :: Monad m => (a -> Maybe e) -> Steps m a -> Steps m e -mapMaybe f (Steps str k) = Steps (S.mapMaybe f str) (toMax k) -{-# INLINE mapMaybe #-} - -concatMap :: Monad m => (a -> Steps m e) -> Steps m a -> Steps m e -concatMap f (Steps str _) = Steps (S.concatMap (stepsStream . f) str) Unknown -{-# INLINE concatMap #-} - - -mapMaybeA :: (Monad m, Applicative f) => (a -> f (Maybe e)) -> Steps Id a -> f (Steps m e) -mapMaybeA f (Steps str k) = (`Steps` toMax k) <$> liftListA (mapMaybeListA f) str -{-# INLINE mapMaybeA #-} - -mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Steps m a -> Steps m b -mapMaybeM f (Steps str k) = Steps (mapMaybeStreamM f str) (toMax k) -{-# INLINE mapMaybeM #-} - -mapMaybeListA :: Applicative f => (a -> f (Maybe b)) -> [a] -> f [b] -mapMaybeListA f = fmap catMaybes . Traversable.traverse f -{-# INLINE mapMaybeListA #-} - -mapMaybeStreamM :: Monad m => (a -> m (Maybe b)) -> S.Stream m a -> S.Stream m b -mapMaybeStreamM f (S.Stream step t) = S.Stream step' t - where - step' s = do - r <- step s - case r of - S.Yield x s' -> do - b <- f x - return $ - case b of - Nothing -> S.Skip s' - Just b' -> S.Yield b' s' - S.Skip s' -> return $ S.Skip s' - S.Done -> return S.Done - {-# INLINE step' #-} -{-# INLINE mapMaybeStreamM #-} - -filter :: Monad m => (a -> Bool) -> Steps m a -> Steps m a -filter f (Steps str k) = Steps (S.filter f str) (toMax k) -{-# INLINE filter #-} - - -filterA :: (Monad m, Applicative f) => (e -> f Bool) -> Steps Id e -> f (Steps m e) -filterA f (Steps str k) = (`Steps` toMax k) <$> liftListA (M.filterM f) str -{-# INLINE filterA #-} - -filterM :: Monad m => (e -> m Bool) -> Steps m e -> Steps m e -filterM f (Steps str k) = Steps (S.filterM f str) (toMax k) -{-# INLINE filterM #-} - -take :: Monad m => Int -> Steps m a -> Steps m a -take n (Steps str _) = Steps (S.take n str) (Max n) -{-# INLINE take #-} - -drop :: Monad m => Int -> Steps m a -> Steps m a -drop n (Steps str k) = Steps (S.drop n str) (k `clampedSubtract` Exact n) -{-# INLINE drop #-} - -slice :: Monad m => Int -> Int -> Steps m a -> Steps m a -slice i k (Steps str _) = Steps (S.slice i k str) (Max k) -{-# INLINE slice #-} - -unfoldr :: Monad m => (s -> Maybe (e, s)) -> s -> Steps m e -unfoldr f e0 = Steps (S.unfoldr f e0) Unknown -{-# INLINE unfoldr #-} - -unfoldrN :: Monad m => Sz1 -> (s -> Maybe (e, s)) -> s -> Steps m e -unfoldrN n f e0 = Steps (S.unfoldrN (unSz n) f e0) (Max (unSz n)) -{-# INLINE unfoldrN #-} - -toList :: Steps Id e -> [e] -toList (Steps str _) = unId (S.toList str) -{-# INLINE toList #-} - -fromList :: Monad m => [e] -> Steps m e -fromList = (`Steps` Unknown) . S.fromList -{-# INLINE fromList #-} - -fromListN :: Monad m => Int -> [e] -> Steps m e -fromListN n = (`Steps` Exact n) . S.fromListN n -{-# INLINE fromListN #-} - -liftListA :: (Monad m, Functor f) => ([a] -> f [b]) -> S.Stream Id a -> f (S.Stream m b) -liftListA f str = S.fromList <$> f (unId (S.toList str)) -{-# INLINE liftListA #-} 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 @@ -1,1095 +1,1457 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE LambdaCase #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} --- | --- Module : Data.Massiv.Array.Mutable --- Copyright : (c) Alexey Kuleshevich 2018-2019 --- License : BSD3 --- Maintainer : Alexey Kuleshevich --- Stability : experimental --- Portability : non-portable --- -module Data.Massiv.Array.Mutable - ( -- ** Size - msize - -- ** Element-wise mutation - , read - , readM - , read' - , write - , write_ - , writeM - , write' - , modify - , modify_ - , modifyM - , modifyM_ - , modify' - , swap - , swap_ - , swapM - , swapM_ - , swap' - -- ** Operations on @MArray@ - -- *** Immutable conversion - , new - , thaw - , thawS - , freeze - , freezeS - -- *** Create mutable - , makeMArray - , makeMArrayLinear - , makeMArrayS - , makeMArrayLinearS - -- *** Create pure - , createArray_ - , createArray - , createArrayS_ - , createArrayS - , createArrayST_ - , createArrayST - -- *** Generate - , generateArray - , generateArrayLinear - , generateArrayS - , generateArrayLinearS - -- *** Stateful worker threads - , generateArrayWS - , generateArrayLinearWS - -- *** Unfold - , unfoldrPrimM_ - , iunfoldrPrimM_ - , unfoldrPrimM - , iunfoldrPrimM - , unfoldlPrimM_ - , iunfoldlPrimM_ - , unfoldlPrimM - , iunfoldlPrimM - -- *** Mapping - , forPrimM - , forPrimM_ - , iforPrimM - , iforPrimM_ - , iforLinearPrimM - , iforLinearPrimM_ - -- *** Modify - , withMArray - , withMArrayS - , withMArrayST - -- *** Initialize - , initialize - , initializeNew - -- ** Computation - , Mutable - , MArray - , RealWorld - , computeInto - , loadArray - , loadArrayS - ) where - --- TODO: add fromListM, et al. - -import Data.Maybe (fromMaybe) -import Control.Monad (void, when, unless, (>=>)) -import Control.Monad.ST -import Control.Scheduler -import Data.Massiv.Core.Common -import Prelude hiding (mapM, read) - --- | /O(n)/ - Initialize a new mutable array. All elements will be set to some default value. For --- boxed arrays in will be a thunk with `Uninitialized` exception, while for others it will be --- simply zeros. --- --- ==== __Examples__ --- --- >>> import Data.Massiv.Array --- >>> marr <- new (Sz2 2 6) :: IO (MArray RealWorld P Ix2 Int) --- >>> freeze Seq marr --- Array P Seq (Sz (2 :. 6)) --- [ [ 0, 0, 0, 0, 0, 0 ] --- , [ 0, 0, 0, 0, 0, 0 ] --- ] --- --- Or using @TypeApplications@: --- --- >>> :set -XTypeApplications --- >>> new @P @Ix2 @Int (Sz2 2 6) >>= freezeS --- Array P Seq (Sz (2 :. 6)) --- [ [ 0, 0, 0, 0, 0, 0 ] --- , [ 0, 0, 0, 0, 0, 0 ] --- ] --- >>> new @B @_ @Int (Sz2 2 6) >>= (`readM` 1) --- *** Exception: Uninitialized --- --- @since 0.1.0 -new :: - forall r ix e m. (Mutable r ix e, PrimMonad m) - => Sz ix - -> m (MArray (PrimState m) r ix e) -new = initializeNew Nothing -{-# INLINE new #-} - --- | /O(n)/ - Make a mutable copy of a pure array. Keep in mind that both `freeze` and `thaw` trigger a --- copy of the full array. --- --- ==== __Example__ --- --- >>> import Data.Massiv.Array --- >>> :set -XTypeApplications --- >>> arr <- fromListsM @U @Ix2 @Double Par [[12,21],[13,31]] --- >>> marr <- thaw arr --- >>> modify marr (pure . (+ 10)) (1 :. 0) --- Just 13.0 --- >>> freeze Par marr --- Array U Par (Sz (2 :. 2)) --- [ [ 12.0, 21.0 ] --- , [ 23.0, 31.0 ] --- ] --- --- @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 $ 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. --- --- ==== __Example__ --- --- >>> import Data.Massiv.Array --- >>> :set -XOverloadedLists --- >>> thawS @P @Ix1 @Double [1..10] --- >>> marr <- thawS @P @Ix1 @Double [1..10] --- >>> writeM marr 5 100 --- >>> freezeS marr --- Array P Seq (Sz1 10) --- [ 1.0, 2.0, 3.0, 4.0, 5.0, 100.0, 7.0, 8.0, 9.0, 10.0 ] --- --- @since 0.3.0 -thawS :: - forall r ix e m. (Mutable r ix e, PrimMonad m) - => Array r ix e - -> m (MArray (PrimState m) r ix e) -thawS arr = do - tmarr <- unsafeNew (size arr) - unsafeArrayLinearCopy arr 0 tmarr 0 (SafeSz (totalElem (size arr))) - pure tmarr -{-# INLINE thawS #-} - - --- | /O(n)/ - Yield an immutable copy of the mutable array. Note that mutable representations --- have to be the same. --- --- ==== __Example__ --- --- >>> import Data.Massiv.Array --- >>> marr <- new @P @_ @Int (Sz2 2 6) --- >>> forM_ (range Seq 0 (Ix2 1 4)) $ \ix -> write marr ix 9 --- >>> freeze Seq marr --- Array P Seq (Sz (2 :. 6)) --- [ [ 9, 9, 9, 9, 0, 0 ] --- , [ 0, 0, 0, 0, 0, 0 ] --- ] --- --- @since 0.1.0 -freeze :: - forall r ix e m. (Mutable r ix e, MonadIO m) - => Comp - -> MArray RealWorld r ix e - -> m (Array r ix e) -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 #-} - - --- | Same as `freeze`, but do the copy of supplied muable array sequentially. Also, unlike `freeze` --- that has to be done in `IO`, `freezeS` can be used with `ST`. --- --- @since 0.3.0 -freezeS :: - forall r ix e m. (Mutable r ix e, PrimMonad m) - => MArray (PrimState m) r ix e - -> m (Array r ix e) -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) -newMaybeInitialized !arr = initializeNew (defaultElement arr) (fromMaybe zeroSz (maxSize arr)) -{-# INLINE newMaybeInitialized #-} - - --- | Load sequentially a pure array into the newly created mutable array. --- --- @since 0.3.0 -loadArrayS :: - forall r ix e r' m. (Load r' ix e, Mutable r ix e, PrimMonad m) - => Array r' ix e - -> m (MArray (PrimState m) r ix e) -loadArrayS arr = do - marr <- newMaybeInitialized arr - unsafeLoadIntoS marr arr -{-# INLINE loadArrayS #-} - - --- | Load a pure array into the newly created mutable array, while respecting computation startegy. --- --- @since 0.3.0 -loadArray :: - forall r ix e r' m. (Load r' ix e, Mutable r ix e, MonadIO m) - => Array r' ix e - -> m (MArray RealWorld r ix e) -loadArray arr = - liftIO $ do - marr <- newMaybeInitialized arr - unsafeLoadInto marr arr -{-# INLINE loadArray #-} - - - --- | Compute an Array while loading the results into the supplied mutable target array. Number of --- elements for arrays must agree, otherwise `SizeElementsMismatchException` exception is thrown. --- --- @since 0.1.3 -computeInto :: - (Load r' ix' e, Mutable r ix e, MonadIO m) - => MArray RealWorld r ix e -- ^ Target Array - -> Array r' ix' e -- ^ Array to load - -> m () -computeInto !mArr !arr = - liftIO $ do - unless (totalElem (msize mArr) == totalElem (size arr)) $ - throwM $ SizeElementsMismatchException (msize mArr) (size arr) - withScheduler_ (getComp arr) $ \scheduler -> loadArrayM scheduler arr (unsafeLinearWrite mArr) -{-# INLINE computeInto #-} - - --- | Create a mutable array using an index aware generating action. --- --- @since 0.3.0 -makeMArrayS :: - forall r ix e m. (Mutable r ix e, PrimMonad m) - => Sz ix -- ^ Size of the create array - -> (ix -> m e) -- ^ Element generating action - -> m (MArray (PrimState m) r ix e) -makeMArrayS sz f = makeMArrayLinearS sz (f . fromLinearIndex sz) -{-# INLINE makeMArrayS #-} - - --- | Same as `makeMArrayS`, but index supplied to the action is row-major linear index. --- --- @since 0.3.0 -makeMArrayLinearS :: - forall r ix e m. (Mutable r ix e, PrimMonad m) - => Sz ix - -> (Int -> m e) - -> m (MArray (PrimState m) r ix e) -makeMArrayLinearS sz f = do - marr <- unsafeNew sz - loopM_ 0 (< totalElem (msize marr)) (+ 1) (\ !i -> f i >>= unsafeLinearWrite marr i) - return marr -{-# INLINE makeMArrayLinearS #-} - --- | Just like `makeMArrayS`, but also accepts computation strategy and runs in `IO`. --- --- @since 0.3.0 -makeMArray :: - forall r ix e m. (PrimMonad m, MonadUnliftIO m, Mutable r ix e) - => Comp - -> Sz ix - -> (ix -> m e) - -> m (MArray (PrimState m) r ix e) -makeMArray comp sz f = makeMArrayLinear comp sz (f . fromLinearIndex sz) -{-# INLINE makeMArray #-} - - --- | Just like `makeMArrayLinearS`, but also accepts computation strategy and runs in `IO`. --- --- @since 0.3.0 -makeMArrayLinear :: - forall r ix e m. (PrimMonad m, MonadUnliftIO m, Mutable r ix e) - => Comp - -> Sz ix - -> (Int -> m e) - -> m (MArray (PrimState m) r ix e) -makeMArrayLinear comp sz f = do - marr <- unsafeNew sz - withScheduler_ comp $ \scheduler -> - splitLinearlyWithM_ scheduler (totalElem sz) f (unsafeLinearWrite marr) - return marr -{-# INLINE makeMArrayLinear #-} - - - - --- | Create a new array by supplying an action that will fill the new blank mutable array. Use --- `createArray` if you'd like to keep the result of the filling function. --- --- ====__Examples__ --- --- >>> :set -XTypeApplications --- >>> import Data.Massiv.Array --- >>> createArray_ @P @_ @Int Seq (Sz1 2) (\ s marr -> scheduleWork s (writeM marr 0 10) >> scheduleWork s (writeM marr 1 11)) --- Array P Seq (Sz1 2) --- [ 10, 11 ] --- --- @since 0.3.0 --- -createArray_ :: - forall r ix e a m. (Mutable r ix e, PrimMonad m, MonadUnliftIO m) - => Comp -- ^ Computation strategy to use after `MArray` gets frozen and onward. - -> Sz ix -- ^ Size of the newly created array - -> (Scheduler m () -> MArray (PrimState m) r ix e -> m a) - -- ^ An action that should fill all elements of the brand new mutable array - -> m (Array r ix e) -createArray_ comp sz action = do - marr <- new sz - withScheduler_ comp (`action` marr) - unsafeFreeze comp marr -{-# INLINE createArray_ #-} - --- | Just like `createArray_`, but together with `Array` it returns results of scheduled filling --- actions. --- --- @since 0.3.0 --- -createArray :: - forall r ix e a m b. (Mutable r ix e, PrimMonad m, MonadUnliftIO m) - => Comp -- ^ Computation strategy to use after `MArray` gets frozen and onward. - -> Sz ix -- ^ Size of the newly created array - -> (Scheduler m a -> MArray (PrimState m) r ix e -> m b) - -- ^ An action that should fill all elements of the brand new mutable array - -> m ([a], Array r ix e) -createArray comp sz action = do - marr <- new sz - a <- withScheduler comp (`action` marr) - arr <- unsafeFreeze comp marr - return (a, arr) -{-# INLINE createArray #-} - - --- | Create a new array by supplying an action that will fill the new blank mutable array. Use --- `createArrayS` if you'd like to keep the result of the filling function. --- --- ====__Examples__ --- --- >>> :set -XTypeApplications --- >>> import Data.Massiv.Array --- >>> createArrayS_ @P @_ @Int Seq (Sz1 2) (\ marr -> write marr 0 10 >> write marr 1 12) --- Array P Seq (Sz1 2) --- [ 10, 12 ] --- --- @since 0.3.0 --- -createArrayS_ :: - forall r ix e a m. (Mutable r ix e, PrimMonad m) - => Comp -- ^ Computation strategy to use after `MArray` gets frozen and onward. - -> Sz ix -- ^ Size of the newly created array - -> (MArray (PrimState m) r ix e -> m a) - -- ^ An action that should fill all elements of the brand new mutable array - -> m (Array r ix e) -createArrayS_ comp sz action = snd <$> createArrayS comp sz action -{-# INLINE createArrayS_ #-} - --- | Just like `createArray_`, but together with `Array` it returns the result of the filling action. --- --- @since 0.3.0 --- -createArrayS :: - forall r ix e a m. (Mutable r ix e, PrimMonad m) - => Comp -- ^ Computation strategy to use after `MArray` gets frozen and onward. - -> Sz ix -- ^ Size of the newly created array - -> (MArray (PrimState m) r ix e -> m a) - -- ^ An action that should fill all elements of the brand new mutable array - -> m (a, Array r ix e) -createArrayS comp sz action = do - marr <- new sz - a <- action marr - arr <- unsafeFreeze comp marr - return (a, arr) -{-# INLINE createArrayS #-} - --- | Just like `createArrayS_`, but restricted to `ST`. --- --- @since 0.3.0 --- -createArrayST_ :: - forall r ix e a. Mutable r ix e - => Comp - -> Sz ix - -> (forall s. MArray s r ix e -> ST s a) - -> Array r ix e -createArrayST_ comp sz action = runST $ createArrayS_ comp sz action -{-# INLINE createArrayST_ #-} - - --- | Just like `createArrayS`, but restricted to `ST`. --- --- @since 0.2.6 --- -createArrayST :: - forall r ix e a. Mutable r ix e - => Comp - -> Sz ix - -> (forall s. MArray s r ix e -> ST s a) - -> (a, Array r ix e) -createArrayST comp sz action = runST $ createArrayS comp sz action -{-# INLINE createArrayST #-} - - --- | Sequentially generate a pure array. Much like `makeArray` creates a pure array this --- function will use `Mutable` interface to generate a pure `Array` in the end, except that --- computation strategy is set to `Seq`. Element producing function no longer has to be pure --- but is a stateful action, becuase it is restricted to `PrimMonad` thus allows for sharing --- the state between computation of each element. --- --- @since 0.2.6 --- --- ====__Examples__ --- --- >>> import Data.Massiv.Array --- >>> import Data.IORef --- >>> ref <- newIORef (0 :: Int) --- >>> generateArrayS (Sz1 6) (\ i -> modifyIORef' ref (+i) >> print i >> pure i) :: IO (Array U Ix1 Int) --- 0 --- 1 --- 2 --- 3 --- 4 --- 5 --- Array U Seq (Sz1 6) --- [ 0, 1, 2, 3, 4, 5 ] --- >>> readIORef ref --- 15 --- -generateArrayS :: - forall r ix e m. (Mutable r ix e, PrimMonad m) - => Sz ix -- ^ Resulting size of the array - -> (ix -> m e) -- ^ Element producing generator - -> m (Array r ix e) -generateArrayS sz gen = generateArrayLinearS sz (gen . fromLinearIndex sz) -{-# INLINE generateArrayS #-} - --- | Same as `generateArray` but with action that accepts row-major linear index. --- --- @since 0.3.0 -generateArrayLinearS :: - forall r ix e m. (Mutable r ix e, PrimMonad m) - => Sz ix -- ^ Resulting size of the array - -> (Int -> m e) -- ^ Element producing generator - -> m (Array r ix e) -generateArrayLinearS sz gen = do - marr <- unsafeNew sz - loopM_ 0 (< totalElem (msize marr)) (+ 1) $ \i -> gen i >>= unsafeLinearWrite marr i - unsafeFreeze Seq marr -{-# INLINE generateArrayLinearS #-} - - --- | Just like `generateArrayS`, except this generator __will__ respect the supplied computation --- strategy, and for that reason it is restricted to `IO`. --- --- @since 0.2.6 -generateArray :: - forall r ix e m. (MonadUnliftIO m, PrimMonad m, Mutable r ix e) - => Comp - -> Sz ix - -> (ix -> m e) - -> m (Array r ix e) -generateArray comp sz f = generateArrayLinear comp sz (f . fromLinearIndex sz) -{-# INLINE generateArray #-} - --- | Just like `generateArrayIO`, but action supplied will receive a row-major linear index. --- --- @since 0.3.0 -generateArrayLinear :: - forall r ix e m. (MonadUnliftIO m, PrimMonad m, Mutable r ix e) - => Comp - -> Sz ix - -> (Int -> m e) - -> m (Array r ix e) -generateArrayLinear comp sz f = makeMArrayLinear comp sz f >>= unsafeFreeze comp -{-# INLINE generateArrayLinear #-} - - --- | Same as `generateArrayWS`, but use linear indexing instead. --- --- @since 0.3.4 -generateArrayLinearWS :: - forall r ix e s m. (Mutable r ix e, MonadUnliftIO m, PrimMonad m) - => WorkerStates s - -> Sz ix - -> (Int -> s -> m e) - -> m (Array r ix e) -generateArrayLinearWS states sz make = do - marr <- unsafeNew sz - withSchedulerWS_ states $ \schedulerWS -> - splitLinearlyWithStatefulM_ - schedulerWS - (totalElem sz) - make - (unsafeLinearWrite marr) - unsafeFreeze (workerStatesComp states) marr -{-# INLINE generateArrayLinearWS #-} - --- | Use per worker thread state while generating elements of the array. Very useful for --- things that are not thread safe. --- --- @since 0.3.4 -generateArrayWS :: - forall r ix e s m. (Mutable r ix e, MonadUnliftIO m, PrimMonad m) - => WorkerStates s - -> Sz ix - -> (ix -> s -> m e) - -> m (Array r ix e) -generateArrayWS states sz make = - generateArrayLinearWS states sz (\ix -> make (fromLinearIndex sz ix)) -{-# INLINE generateArrayWS #-} - - --- | Sequentially unfold an array from the left. --- --- ====__Examples__ --- --- Create an array with Fibonacci numbers while performing and `IO` action on the accumulator for --- each element of the array. --- --- >>> import Data.Massiv.Array --- >>> unfoldrPrimM_ Seq (Sz1 10) (\a@(f0, f1) -> let fn = f0 + f1 in print a >> return (f0, (f1, fn))) (0, 1) :: IO (Array P Ix1 Int) --- (0,1) --- (1,1) --- (1,2) --- (2,3) --- (3,5) --- (5,8) --- (8,13) --- (13,21) --- (21,34) --- (34,55) --- Array P Seq (Sz1 10) --- [ 0, 1, 1, 2, 3, 5, 8, 13, 21, 34 ] --- --- @since 0.3.0 --- -unfoldrPrimM_ :: - forall r ix e a m. (Mutable r ix e, PrimMonad m) - => Comp -- ^ Computation strategy (ignored during initial creation) - -> Sz ix -- ^ Size of the desired array - -> (a -> m (e, a)) -- ^ Unfolding action - -> a -- ^ Initial accumulator - -> m (Array r ix e) -unfoldrPrimM_ comp sz gen acc0 = snd <$> unfoldrPrimM comp sz gen acc0 -{-# INLINE unfoldrPrimM_ #-} - --- | Same as `unfoldrPrimM_` but do the unfolding with index aware function. --- --- @since 0.3.0 --- -iunfoldrPrimM_ :: - forall r ix e a m. (Mutable r ix e, PrimMonad m) - => Comp -- ^ Computation strategy (ignored during initial creation) - -> Sz ix -- ^ Size of the desired array - -> (a -> ix -> m (e, a)) -- ^ Unfolding action - -> a -- ^ Initial accumulator - -> m (Array r ix e) -iunfoldrPrimM_ comp sz gen acc0 = snd <$> iunfoldrPrimM comp sz gen acc0 -{-# INLINE iunfoldrPrimM_ #-} - - --- | Just like `iunfoldrPrimM_`, but also returns the final value of the accumulator. --- --- @since 0.3.0 -iunfoldrPrimM :: - forall r ix e a m. (Mutable r ix e, PrimMonad m) - => Comp -- ^ Computation strategy (ignored during initial creation) - -> Sz ix -- ^ Size of the desired array - -> (a -> ix -> m (e, a)) -- ^ Unfolding action - -> a -- ^ Initial accumulator - -> m (a, Array r ix e) -iunfoldrPrimM comp sz gen acc0 = - createArrayS comp sz $ \marr -> - let sz' = msize marr - in iterLinearM sz' 0 (totalElem sz') 1 (<) acc0 $ \i ix acc -> do - (e, acc') <- gen acc ix - unsafeLinearWrite marr i e - pure $! acc' -{-# INLINE iunfoldrPrimM #-} - --- | Just like `iunfoldrPrimM`, but do the unfolding with index aware function. --- --- @since 0.3.0 -unfoldrPrimM :: - forall r ix e a m. (Mutable r ix e, PrimMonad m) - => Comp -- ^ Computation strategy (ignored during initial creation) - -> Sz ix -- ^ Size of the desired array - -> (a -> m (e, a)) -- ^ Unfolding action - -> a -- ^ Initial accumulator - -> m (a, Array r ix e) -unfoldrPrimM comp sz gen acc0 = - createArrayS comp sz $ \marr -> - let sz' = msize marr - in loopM 0 (< totalElem sz') (+1) acc0 $ \i acc -> do - (e, acc') <- gen acc - unsafeLinearWrite marr i e - pure $! acc' -{-# INLINE unfoldrPrimM #-} - --- | Sequentially unfold an array from the left. --- --- ====__Examples__ --- --- Create an array with Fibonacci numbers starting at the end while performing and `IO` action on --- the accumulator for each element of the array. --- --- >>> import Data.Massiv.Array --- >>> unfoldlPrimM_ Seq (Sz1 10) (\a@(f0, f1) -> let fn = f0 + f1 in print a >> return ((f1, fn), f0)) (0, 1) :: IO (Array P Ix1 Int) --- (0,1) --- (1,1) --- (1,2) --- (2,3) --- (3,5) --- (5,8) --- (8,13) --- (13,21) --- (21,34) --- (34,55) --- Array P Seq (Sz1 10) --- [ 34, 21, 13, 8, 5, 3, 2, 1, 1, 0 ] --- --- @since 0.3.0 --- -unfoldlPrimM_ :: - forall r ix e a m. (Mutable r ix e, PrimMonad m) - => Comp -- ^ Computation strategy (ignored during initial creation) - -> Sz ix -- ^ Size of the desired array - -> (a -> m (a, e)) -- ^ Unfolding action - -> a -- ^ Initial accumulator - -> m (Array r ix e) -unfoldlPrimM_ comp sz gen acc0 = snd <$> unfoldlPrimM comp sz gen acc0 -{-# INLINE unfoldlPrimM_ #-} - --- | Same as `unfoldlPrimM_` but do the unfolding with index aware function. --- --- @since 0.3.0 --- -iunfoldlPrimM_ :: - forall r ix e a m. (Mutable r ix e, PrimMonad m) - => Comp -- ^ Computation strategy (ignored during initial creation) - -> Sz ix -- ^ Size of the desired array - -> (a -> ix -> m (a, e)) -- ^ Unfolding action - -> a -- ^ Initial accumulator - -> m (Array r ix e) -iunfoldlPrimM_ comp sz gen acc0 = snd <$> iunfoldlPrimM comp sz gen acc0 -{-# INLINE iunfoldlPrimM_ #-} - - --- | Just like `iunfoldlPrimM_`, but also returns the final value of the accumulator. --- --- @since 0.3.0 -iunfoldlPrimM :: - forall r ix e a m. (Mutable r ix e, PrimMonad m) - => Comp -- ^ Computation strategy (ignored during initial creation) - -> Sz ix -- ^ Size of the desired array - -> (a -> ix -> m (a, e)) -- ^ Unfolding action - -> a -- ^ Initial accumulator - -> m (a, Array r ix e) -iunfoldlPrimM comp sz gen acc0 = - createArrayS comp sz $ \marr -> - let sz' = msize marr - in iterLinearM sz' (totalElem sz' - 1) 0 (negate 1) (>=) acc0 $ \i ix acc -> do - (acc', e) <- gen acc ix - unsafeLinearWrite marr i e - pure $! acc' -{-# INLINE iunfoldlPrimM #-} - --- | Just like `iunfoldlPrimM`, but do the unfolding with index aware function. --- --- @since 0.3.0 -unfoldlPrimM :: - forall r ix e a m. (Mutable r ix e, PrimMonad m) - => Comp -- ^ Computation strategy (ignored during initial creation) - -> Sz ix -- ^ Size of the desired array - -> (a -> m (a, e)) -- ^ Unfolding action - -> a -- ^ Initial accumulator - -> m (a, Array r ix e) -unfoldlPrimM comp sz gen acc0 = - createArrayS comp sz $ \marr -> - let sz' = msize marr - in loopDeepM 0 (< totalElem sz') (+1) acc0 $ \i acc -> do - (acc', e) <- gen acc - unsafeLinearWrite marr i e - pure $! acc' -{-# INLINE unfoldlPrimM #-} - --- | Sequentially loop over a mutable array while reading each element and applying an --- action to it. There is no mutation to the array, unless the action itself modifies it. --- --- @since 0.4.0 -forPrimM_ :: (Mutable r ix e, PrimMonad m) => MArray (PrimState m) r ix e -> (e -> m ()) -> m () -forPrimM_ marr f = - loopM_ 0 (< totalElem (msize marr)) (+1) (unsafeLinearRead marr >=> f) -{-# INLINE forPrimM_ #-} - --- | Sequentially loop over a mutable array while modifying each element with an action. --- --- @since 0.4.0 -forPrimM :: (Mutable r ix e, PrimMonad m) => MArray (PrimState m) r ix e -> (e -> m e) -> m () -forPrimM marr f = - loopM_ 0 (< totalElem (msize marr)) (+1) (unsafeLinearModify marr f) -{-# INLINE forPrimM #-} - - --- | Sequentially loop over a mutable array while reading each element and applying an --- index aware action to it. There is no mutation to the array, unless the --- action itself modifies it. --- --- @since 0.4.0 -iforPrimM_ :: - (Mutable r ix e, PrimMonad m) => MArray (PrimState m) r ix e -> (ix -> e -> m ()) -> m () -iforPrimM_ marr f = iforLinearPrimM_ marr (f . fromLinearIndex (msize marr)) -{-# INLINE iforPrimM_ #-} - --- | Sequentially loop over a mutable array while modifying each element with an index aware action. --- --- @since 0.4.0 -iforPrimM :: - (Mutable r ix e, PrimMonad m) => MArray (PrimState m) r ix e -> (ix -> e -> m e) -> m () -iforPrimM marr f = iforLinearPrimM marr (f . fromLinearIndex (msize marr)) -{-# INLINE iforPrimM #-} - - --- | Sequentially loop over a mutable array while reading each element and applying a --- linear index aware action to it. There is no mutation to the array, unless the action --- itself modifies it. --- --- @since 0.4.0 -iforLinearPrimM_ :: - (Mutable r ix e, PrimMonad m) => MArray (PrimState m) r ix e -> (Int -> e -> m ()) -> m () -iforLinearPrimM_ marr f = - loopM_ 0 (< totalElem (msize marr)) (+ 1) (\i -> unsafeLinearRead marr i >>= f i) -{-# INLINE iforLinearPrimM_ #-} - --- | Sequentially loop over a mutable array while modifying each element with an index aware action. --- --- @since 0.4.0 -iforLinearPrimM :: - (Mutable r ix e, PrimMonad m) => MArray (PrimState m) r ix e -> (Int -> e -> m e) -> m () -iforLinearPrimM marr f = - loopM_ 0 (< totalElem (msize marr)) (+ 1) (\i -> unsafeLinearModify marr (f i) i) -{-# INLINE iforLinearPrimM #-} - --- | Create a copy of a pure array, mutate it in place and return its frozen version. The big --- difference between `withMArrayS` is that it's not only gonna respect the computation strategy --- supplied to it while making a copy, but it will also pass extra argumens to the action that --- suppose to modify the mutable copy of the source array. These two extra arguments are: --- --- * Number of capabilities derived from the `Comp`utation strategy of the array. --- --- * An action that can be used to schedule arbitrary number of jobs that will be executed in --- parallel. --- --- * And, of course, the mutable array itself. --- --- @since 0.3.0 -withMArray :: - (Mutable r ix e, MonadUnliftIO m) - => Array r ix e - -> (Scheduler m () -> MArray RealWorld r ix e -> m a) - -> m (Array r ix e) -withMArray arr action = do - marr <- thaw arr - withScheduler_ (getComp arr) (`action` marr) - liftIO $ unsafeFreeze (getComp arr) marr -{-# INLINE withMArray #-} - - --- | Create a copy of a pure array, mutate it in place and return its frozen version. The important --- benefit over doing a manual `thawS` followed by a `freezeS` is that an array will be only copied --- once. --- --- @since 0.3.2 -withMArrayS :: - (Mutable r ix e, PrimMonad m) - => Array r ix e - -> (MArray (PrimState m) r ix e -> m a) - -> m (Array r ix e) -withMArrayS arr action = do - marr <- thawS arr - _ <- action marr - unsafeFreeze (getComp arr) marr -{-# INLINE withMArrayS #-} - - --- | Same as `withMArrayS` but in `ST`. This is not only pure, but also the safest way to do --- mutation to the array. --- --- @since 0.2.2 -withMArrayST :: - Mutable r ix e - => Array r ix e - -> (forall s . MArray s r ix e -> ST s a) - -> Array r ix e -withMArrayST arr f = runST $ withMArrayS arr f -{-# INLINE withMArrayST #-} - - --- | /O(1)/ - Lookup an element in the mutable array. Returns `Nothing` when index is out of bounds. --- --- @since 0.1.0 -read :: (Mutable r ix e, PrimMonad m) => - MArray (PrimState m) r ix e -> ix -> m (Maybe e) -read marr ix = - if isSafeIndex (msize marr) ix - then Just <$> unsafeRead marr ix - else return Nothing -{-# INLINE read #-} - - --- | /O(1)/ - Same as `read`, but throws `IndexOutOfBoundsException` on an invalid index. --- --- @since 0.4.0 -readM :: (Mutable r ix e, PrimMonad m, MonadThrow m) => - MArray (PrimState m) r ix e -> ix -> m e -readM marr ix = - read marr ix >>= \case - Just e -> pure e - Nothing -> throwM $ IndexOutOfBoundsException (msize marr) ix -{-# INLINE readM #-} - - --- | /O(1)/ - Same as `read`, but throws `IndexOutOfBoundsException` on an invalid index. --- --- @since 0.1.0 -read' :: (Mutable r ix e, PrimMonad m) => MArray (PrimState m) r ix e -> ix -> m e -read' marr ix = - read marr ix >>= \case - Just e -> pure e - Nothing -> throw $ IndexOutOfBoundsException (msize marr) ix -{-# INLINE read' #-} -{-# DEPRECATED read' "In favor of more general `readM`" #-} - - --- | /O(1)/ - Write an element into the cell of a mutable array. Returns `False` when index is out --- of bounds. --- --- @since 0.1.0 -write :: (Mutable r ix e, PrimMonad m) => MArray (PrimState m) r ix e -> ix -> e -> m Bool -write marr ix e = - if isSafeIndex (msize marr) ix - then unsafeWrite marr ix e >> pure True - else pure False -{-# INLINE write #-} - --- | /O(1)/ - Write an element into the cell of a mutable array. Same as `write` function --- in case of an out of bounds index it is noop, but unlike `write`, there is no --- information is returned about was the writing of element successful or not. In other --- words, just like `writeM`, but doesn't throw an exception. --- --- @since 0.4.4 -write_ :: (Mutable r ix e, PrimMonad m) => MArray (PrimState m) r ix e -> ix -> e -> m () -write_ marr ix = when (isSafeIndex (msize marr) ix) . unsafeWrite marr ix -{-# INLINE write_ #-} - --- | /O(1)/ - Same as `write`, but throws `IndexOutOfBoundsException` on an invalid index. --- --- @since 0.4.0 -writeM :: - (Mutable r ix e, PrimMonad m, MonadThrow m) => MArray (PrimState m) r ix e -> ix -> e -> m () -writeM marr ix e = - write marr ix e >>= (`unless` throwM (IndexOutOfBoundsException (msize marr) ix)) -{-# INLINE writeM #-} - - --- | /O(1)/ - Same as `write`, but lives in IO and throws `IndexOutOfBoundsException` on invalid --- index. --- --- @since 0.1.0 -write' :: - (Mutable r ix e, PrimMonad m) => MArray (PrimState m) r ix e -> ix -> e -> m () -write' marr ix e = write marr ix e >>= (`unless` throw (IndexOutOfBoundsException (msize marr) ix)) -{-# INLINE write' #-} -{-# DEPRECATED write' "In favor of more general `writeM`" #-} - --- | /O(1)/ - Modify an element in the cell of a mutable array with a supplied --- action. Returns the previous value, if index was not out of bounds. --- --- @since 0.1.0 -modify :: - (Mutable r ix e, PrimMonad m) - => MArray (PrimState m) r ix e -- ^ Array to mutate. - -> (e -> m e) -- ^ Monadic action that modifies the element - -> ix -- ^ Index at which to perform modification. - -> m (Maybe e) -modify marr f ix = - if isSafeIndex (msize marr) ix - then Just <$> unsafeModify marr f ix - else return Nothing -{-# INLINE modify #-} - --- | /O(1)/ - Same as `modify`, except that neither the previous value, nor any --- information on whether the modification was successful are returned. In other words, --- just like `modifyM_`, but doesn't throw an exception. --- --- @since 0.4.4 -modify_ :: - (Mutable r ix e, PrimMonad m) - => MArray (PrimState m) r ix e -- ^ Array to mutate. - -> (e -> m e) -- ^ Monadic action that modifies the element - -> ix -- ^ Index at which to perform modification. - -> m () -modify_ marr f ix = when (isSafeIndex (msize marr) ix) $ void $ unsafeModify marr f ix -{-# INLINE modify_ #-} - --- | /O(1)/ - Modify an element in the cell of a mutable array with a supplied --- action. Throws an `IndexOutOfBoundsException` exception for invalid index and returns --- the previous value otherwise. --- --- @since 0.4.0 -modifyM :: - (Mutable r ix e, PrimMonad m, MonadThrow m) - => MArray (PrimState m) r ix e -- ^ Array to mutate. - -> (e -> m e) -- ^ Monadic action that modifies the element - -> ix -- ^ Index at which to perform modification. - -> m e -modifyM marr f ix - | isSafeIndex (msize marr) ix = unsafeModify marr f ix - | otherwise = throwM (IndexOutOfBoundsException (msize marr) ix) -{-# INLINE modifyM #-} - --- | /O(1)/ - Same as `modifyM`, but discard the returned element --- --- ====__Examples__ --- --- >>> :set -XTypeApplications --- >>> import Control.Monad.ST --- >>> import Data.Massiv.Array --- >>> runST $ new @P @Ix1 @Int (Sz1 3) >>= (\ma -> modifyM_ ma (pure . (+10)) 1 >> freezeS ma) --- Array P Seq (Sz1 3) --- [ 0, 10, 0 ] --- --- @since 0.4.0 -modifyM_ :: - (Mutable r ix e, PrimMonad m, MonadThrow m) - => MArray (PrimState m) r ix e -- ^ Array to mutate. - -> (e -> m e) -- ^ Monadic action that modifies the element - -> ix -- ^ Index at which to perform modification. - -> m () -modifyM_ marr f ix = void $ modifyM marr f ix -{-# INLINE modifyM_ #-} - - --- | /O(1)/ - Same as `modify`, but throws an error if index is out of bounds. --- --- @since 0.1.0 -modify' :: (Mutable r ix e, PrimMonad m) => - MArray (PrimState m) r ix e -> (e -> e) -> ix -> m () -modify' marr f ix = - modify marr (pure . f) ix >>= \case - Just _ -> pure () - Nothing -> throw (IndexOutOfBoundsException (msize marr) ix) -{-# INLINE modify' #-} -{-# DEPRECATED modify' "In favor of more general `modifyM`" #-} - - --- | /O(1)/ - Same as `swapM`, but instead of throwing an exception returns `Nothing` when --- either one of the indices is out of bounds and `Just` elements under those indices --- otherwise. --- --- @since 0.1.0 -swap :: (Mutable r ix e, PrimMonad m) => MArray (PrimState m) r ix e -> ix -> ix -> m (Maybe (e, e)) -swap marr ix1 ix2 = - let !sz = msize marr - in if isSafeIndex sz ix1 && isSafeIndex sz ix2 - then Just <$> unsafeSwap marr ix1 ix2 - else pure Nothing -{-# INLINE swap #-} - - --- | /O(1)/ - Same as `swap`, but instead of returning `Nothing` it does nothing. In other --- words, it is similar to `swapM_`, but does not throw any exceptions. --- --- @since 0.4.4 -swap_ :: (Mutable r ix e, PrimMonad m) => MArray (PrimState m) r ix e -> ix -> ix -> m () -swap_ marr ix1 ix2 = - let !sz = msize marr - in when (isSafeIndex sz ix1 && isSafeIndex sz ix2) $ void $ unsafeSwap marr ix1 ix2 -{-# INLINE swap_ #-} - --- | /O(1)/ - Swap two elements in a mutable array under the supplied indices. Throws an --- `IndexOutOfBoundsException` when either one of the indices is out of bounds and --- elements under those indices otherwise. --- --- @since 0.4.0 -swapM :: - (Mutable r ix e, PrimMonad m, MonadThrow m) - => MArray (PrimState m) r ix e - -> ix -- ^ Index for the first element, which will be returned as the first element in the - -- tuple. - -> ix -- ^ Index for the second element, which will be returned as the second element in - -- the tuple. - -> m (e, e) -swapM marr ix1 ix2 - | not (isSafeIndex sz ix1) = throwM $ IndexOutOfBoundsException (msize marr) ix1 - | not (isSafeIndex sz ix2) = throwM $ IndexOutOfBoundsException (msize marr) ix2 - | otherwise = unsafeSwap marr ix1 ix2 - where - !sz = msize marr -{-# INLINE swapM #-} - - --- | /O(1)/ - Same as `swapM`, but discard the returned elements --- --- @since 0.4.0 -swapM_ :: - (Mutable r ix e, PrimMonad m, MonadThrow m) => MArray (PrimState m) r ix e -> ix -> ix -> m () -swapM_ marr ix1 ix2 = void $ swapM marr ix1 ix2 -{-# INLINE swapM_ #-} - - --- | /O(1)/ - Same as `swap`, but throws an `IndexOutOfBoundsException` on invalid indices. --- --- @since 0.1.0 -swap' :: - (Mutable r ix e, PrimMonad m) => MArray (PrimState m) r ix e -> ix -> ix -> m () -swap' marr ix1 ix2 = - swap marr ix1 ix2 >>= \case - Just _ -> pure () - Nothing -> - if isSafeIndex (msize marr) ix1 - then throw $ IndexOutOfBoundsException (msize marr) ix2 - else throw $ IndexOutOfBoundsException (msize marr) ix1 -{-# INLINE swap' #-} -{-# DEPRECATED swap' "In favor of more general `swapM`" #-} +{-# LANGUAGE MonoLocalBinds #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} + +-- | +-- Module : Data.Massiv.Array.Mutable +-- Copyright : (c) Alexey Kuleshevich 2018-2022 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +module Data.Massiv.Array.Mutable ( + -- ** Size + sizeOfMArray, + msize, + resizeMArrayM, + flattenMArray, + outerSliceMArrayM, + outerSlicesMArray, + + -- ** Element-wise mutation + read, + readM, + write, + write_, + writeM, + modify, + modify_, + modifyM, + modifyM_, + swap, + swap_, + swapM, + swapM_, + zipSwapM_, + + -- ** Operations on @MArray@ + + -- *** Immutable conversion + thaw, + thawS, + freeze, + freezeS, + + -- *** Create mutable + newMArray, + newMArray', + makeMArray, + makeMArrayLinear, + makeMArrayS, + makeMArrayLinearS, + + -- *** Create pure + createArray_, + createArray, + createArrayS_, + createArrayS, + createArrayST_, + createArrayST, + + -- *** Generate + generateArray, + generateArrayLinear, + generateArrayS, + generateArrayLinearS, + generateSplitSeedArray, + + -- *** Stateful worker threads + generateArrayWS, + generateArrayLinearWS, + + -- *** Unfold + unfoldrPrimM_, + iunfoldrPrimM_, + unfoldrPrimM, + iunfoldrPrimM, + unfoldlPrimM_, + iunfoldlPrimM_, + unfoldlPrimM, + iunfoldlPrimM, + + -- *** Mapping + forPrimM, + forPrimM_, + iforPrimM, + iforPrimM_, + iforLinearPrimM, + iforLinearPrimM_, + for2PrimM_, + ifor2PrimM_, + + -- *** Modify + withMArray, + withMArray_, + withLoadMArray_, + withMArrayS, + withLoadMArrayS, + withMArrayS_, + withLoadMArrayS_, + withMArrayST, + withLoadMArrayST, + withMArrayST_, + withLoadMArrayST_, + + -- *** Initialize + initialize, + initializeNew, + + -- ** Computation + Manifest, + MArray, + RealWorld, + computeInto, + loadArray, + loadArrayS, +) where + +-- TODO: add fromListM, et al. + +import Control.Monad (unless, void, when, (>=>)) +import Control.Monad.Primitive +import Control.Monad.ST +import Control.Scheduler +import Data.IORef +import Data.Massiv.Array.Delayed.Pull (D) +import Data.Massiv.Array.Mutable.Internal +import Data.Massiv.Core.Common +import Data.Maybe (fromMaybe) +import System.IO.Unsafe (unsafePerformIO) +import Prelude hiding (mapM, read) + +-- | /O(1)/ - Change the size of a mutable array. Throws +-- `SizeElementsMismatchException` if total number of elements does not match +-- the supplied array. +-- +-- @since 1.0.0 +resizeMArrayM + :: (Manifest r e, Index ix', Index ix, MonadThrow m) + => Sz ix' + -> MArray s r ix e + -> m (MArray s r ix' e) +resizeMArrayM sz marr = + unsafeResizeMArray sz marr <$ guardNumberOfElements (sizeOfMArray marr) sz +{-# INLINE resizeMArrayM #-} + +-- | /O(1)/ - Change a mutable array to a mutable vector. +-- +-- @since 1.0.0 +flattenMArray :: (Manifest r e, Index ix) => MArray s r ix e -> MVector s r e +flattenMArray marr = unsafeResizeMArray (toLinearSz (sizeOfMArray marr)) marr +{-# INLINE flattenMArray #-} + +-- | /O(1)/ - Slice a mutable array from the outside, while reducing its +-- dimensionality by one. Same as `Data.Massiv.Array.!?>` operator, but for +-- mutable arrays. +-- +-- @since 1.0.0 +outerSliceMArrayM + :: forall r ix e m s + . (MonadThrow m, Index (Lower ix), Index ix, Manifest r e) + => MArray s r ix e + -> Ix1 + -> m (MArray s r (Lower ix) e) +outerSliceMArrayM !marr !i = do + let (k, szL) = unconsSz (sizeOfMArray marr) + unless (isSafeIndex k i) $ throwM $ IndexOutOfBoundsException k i + pure $ unsafeResizeMArray szL $ unsafeLinearSliceMArray (i * totalElem szL) (toLinearSz szL) marr +{-# INLINE outerSliceMArrayM #-} + +-- | /O(1)/ - Take all outer slices of a mutable array and construct a delayed +-- vector out of them. In other words it applies `outerSliceMArrayM` to each +-- outer index. Same as `Data.Massiv.Array.outerSlices` function, but for +-- mutable arrays. +-- +-- ====__Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> arr <- resizeM (Sz2 4 7) $ makeArrayR P Seq (Sz1 28) (+10) +-- >>> arr +-- Array P Seq (Sz (4 :. 7)) +-- [ [ 10, 11, 12, 13, 14, 15, 16 ] +-- , [ 17, 18, 19, 20, 21, 22, 23 ] +-- , [ 24, 25, 26, 27, 28, 29, 30 ] +-- , [ 31, 32, 33, 34, 35, 36, 37 ] +-- ] +-- +-- Here we can see we can get individual rows from a mutable matrix +-- +-- >>> marr <- thawS arr +-- >>> import Control.Monad ((<=<)) +-- >>> mapIO_ (print <=< freezeS) $ outerSlicesMArray Seq marr +-- Array P Seq (Sz1 7) +-- [ 10, 11, 12, 13, 14, 15, 16 ] +-- Array P Seq (Sz1 7) +-- [ 17, 18, 19, 20, 21, 22, 23 ] +-- Array P Seq (Sz1 7) +-- [ 24, 25, 26, 27, 28, 29, 30 ] +-- Array P Seq (Sz1 7) +-- [ 31, 32, 33, 34, 35, 36, 37 ] +-- +-- For the sake of example what if our goal was to mutate array in such a way +-- that rows from the top half were swapped with the bottom half: +-- +-- >>> (top, bottom) <- splitAtM 1 2 $ outerSlicesMArray Seq marr +-- >>> mapIO_ (print <=< freezeS) top +-- Array P Seq (Sz1 7) +-- [ 10, 11, 12, 13, 14, 15, 16 ] +-- Array P Seq (Sz1 7) +-- [ 17, 18, 19, 20, 21, 22, 23 ] +-- >>> mapIO_ (print <=< freezeS) bottom +-- Array P Seq (Sz1 7) +-- [ 24, 25, 26, 27, 28, 29, 30 ] +-- Array P Seq (Sz1 7) +-- [ 31, 32, 33, 34, 35, 36, 37 ] +-- >>> szipWithM_ (zipSwapM_ 0) top bottom +-- >>> freezeS marr +-- Array P Seq (Sz (4 :. 7)) +-- [ [ 24, 25, 26, 27, 28, 29, 30 ] +-- , [ 31, 32, 33, 34, 35, 36, 37 ] +-- , [ 10, 11, 12, 13, 14, 15, 16 ] +-- , [ 17, 18, 19, 20, 21, 22, 23 ] +-- ] +-- +-- @since 1.0.0 +outerSlicesMArray + :: forall r ix e s + . (Index (Lower ix), Index ix, Manifest r e) + => Comp + -> MArray s r ix e + -> Vector D (MArray s r (Lower ix) e) +outerSlicesMArray comp marr = + makeArray comp k (\i -> unsafeResizeMArray szL $ unsafeLinearSliceMArray (i * unSz kL) kL marr) + where + kL = toLinearSz szL + (k, szL) = unconsSz $ sizeOfMArray marr +{-# INLINE outerSlicesMArray #-} + +-- | /O(n)/ - Initialize a new mutable array. All elements will be set to some default value. For +-- boxed arrays it will be a thunk with `Uninitialized` exception, while for others it will be +-- simply zeros. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array +-- >>> marr <- newMArray' (Sz2 2 6) :: IO (MArray RealWorld P Ix2 Int) +-- >>> freeze Seq marr +-- Array P Seq (Sz (2 :. 6)) +-- [ [ 0, 0, 0, 0, 0, 0 ] +-- , [ 0, 0, 0, 0, 0, 0 ] +-- ] +-- +-- Or using @TypeApplications@: +-- +-- >>> :set -XTypeApplications +-- >>> newMArray' @P @Ix2 @Int (Sz2 2 6) >>= freezeS +-- Array P Seq (Sz (2 :. 6)) +-- [ [ 0, 0, 0, 0, 0, 0 ] +-- , [ 0, 0, 0, 0, 0, 0 ] +-- ] +-- >>> newMArray' @B @_ @Int (Sz2 2 6) >>= freezeS +-- *** Exception: Uninitialized +-- +-- @since 0.6.0 +newMArray' + :: forall r ix e m + . (Manifest r e, Index ix, PrimMonad m) + => Sz ix + -> m (MArray (PrimState m) r ix e) +newMArray' sz = unsafeNew sz >>= \ma -> ma <$ initialize ma +{-# INLINE newMArray' #-} + +-- | /O(n)/ - Make a mutable copy of a pure array. Keep in mind that both `freeze` and `thaw` trigger a +-- copy of the full array. +-- +-- ==== __Example__ +-- +-- >>> import Data.Massiv.Array +-- >>> :set -XTypeApplications +-- >>> arr <- fromListsM @U @Ix2 @Double Par [[12,21],[13,31]] +-- >>> marr <- thaw arr +-- >>> modify marr (pure . (+ 10)) (1 :. 0) +-- Just 13.0 +-- >>> freeze Par marr +-- Array U Par (Sz (2 :. 2)) +-- [ [ 12.0, 21.0 ] +-- , [ 23.0, 31.0 ] +-- ] +-- +-- @since 0.1.0 +thaw + :: forall r ix e m. (Manifest r e, Index ix, MonadIO m) => Array r ix e -> m (MArray RealWorld r ix e) +thaw arr = + liftIO $ do + let sz = size arr + totalLength = totalElem sz + marr <- unsafeNew sz + withMassivScheduler_ (getComp arr) $ \scheduler -> + splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do + loopA_ 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. +-- +-- ==== __Example__ +-- +-- >>> import Data.Massiv.Array +-- >>> :set -XOverloadedLists +-- >>> thawS @P @Ix1 @Double [1..10] +-- >>> marr <- thawS @P @Ix1 @Double [1..10] +-- >>> writeM marr 5 100 +-- >>> freezeS marr +-- Array P Seq (Sz1 10) +-- [ 1.0, 2.0, 3.0, 4.0, 5.0, 100.0, 7.0, 8.0, 9.0, 10.0 ] +-- +-- @since 0.3.0 +thawS + :: forall r ix e m + . (Manifest r e, Index ix, PrimMonad m) + => Array r ix e + -> m (MArray (PrimState m) r ix e) +thawS arr = do + tmarr <- unsafeNew (size arr) + unsafeArrayLinearCopy arr 0 tmarr 0 (SafeSz (totalElem (size arr))) + pure tmarr +{-# INLINE thawS #-} + +-- | /O(n)/ - Yield an immutable copy of the mutable array. Note that mutable representations +-- have to be the same. +-- +-- ==== __Example__ +-- +-- >>> import Data.Massiv.Array +-- >>> marr <- newMArray @P (Sz2 2 6) (0 :: Int) +-- >>> forM_ (range Seq 0 (Ix2 1 4)) $ \ix -> write marr ix 9 +-- >>> freeze Seq marr +-- Array P Seq (Sz (2 :. 6)) +-- [ [ 9, 9, 9, 9, 0, 0 ] +-- , [ 0, 0, 0, 0, 0, 0 ] +-- ] +-- +-- @since 0.1.0 +freeze + :: forall r ix e m + . (Manifest r e, Index ix, MonadIO m) + => Comp + -> MArray RealWorld r ix e + -> m (Array r ix e) +freeze comp smarr = + liftIO $ do + let sz = sizeOfMArray smarr + totalLength = totalElem sz + tmarr <- unsafeNew sz + withMassivScheduler_ comp $ \scheduler -> + splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do + loopA_ 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 #-} + +-- | Same as `freeze`, but do the copy of supplied muable array sequentially. Also, unlike `freeze` +-- that has to be done in `IO`, `freezeS` can be used with `ST`. +-- +-- @since 0.3.0 +freezeS + :: forall r ix e m + . (Manifest r e, Index ix, PrimMonad m) + => MArray (PrimState m) r ix e + -> m (Array r ix e) +freezeS smarr = do + let sz = sizeOfMArray smarr + tmarr <- unsafeNew sz + unsafeLinearCopy smarr 0 tmarr 0 (SafeSz (totalElem sz)) + unsafeFreeze Seq tmarr +{-# INLINE freezeS #-} + +unsafeNewUpper + :: (Load r' ix e, Manifest r e, PrimMonad m) => Array r' ix e -> m (MArray (PrimState m) r Ix1 e) +unsafeNewUpper !arr = unsafeNew (fromMaybe zeroSz (maxLinearSize arr)) +{-# INLINE unsafeNewUpper #-} + +-- | Load sequentially a pure array into the newly created mutable array. +-- +-- @since 0.3.0 +loadArrayS + :: forall r ix e r' m + . (Load r' ix e, Manifest r e, PrimMonad m) + => Array r' ix e + -> m (MArray (PrimState m) r ix e) +loadArrayS arr = do + marr <- unsafeNewUpper arr + stToPrim $ unsafeLoadIntoST marr arr +{-# INLINE loadArrayS #-} + +-- | Load a pure array into the newly created mutable array, while respecting computation startegy. +-- +-- @since 0.3.0 +loadArray + :: forall r ix e r' m + . (Load r' ix e, Manifest r e, MonadIO m) + => Array r' ix e + -> m (MArray RealWorld r ix e) +loadArray arr = + liftIO $ do + marr <- unsafeNewUpper arr + unsafeLoadIntoIO marr arr +{-# INLINE loadArray #-} + +-- | Compute an Array while loading the results into the supplied mutable target array. Number of +-- elements for arrays must agree, otherwise `SizeElementsMismatchException` exception is thrown. +-- +-- @since 0.1.3 +computeInto + :: (Load r' ix' e, Manifest r e, Index ix, MonadIO m) + => MArray RealWorld r ix e + -- ^ Target Array + -> Array r' ix' e + -- ^ Array to load + -> m () +computeInto !mArr !arr = + liftIO $ do + let sz = outerSize arr + unless (totalElem (sizeOfMArray mArr) == totalElem sz) $ + throwM $ + SizeElementsMismatchException (sizeOfMArray mArr) sz + withMassivScheduler_ (getComp arr) $ \scheduler -> + stToPrim $ iterArrayLinearST_ scheduler arr (unsafeLinearWrite mArr) +{-# INLINE computeInto #-} + +-- | Create a mutable array using an index aware generating action. +-- +-- @since 0.3.0 +makeMArrayS + :: forall r ix e m + . (Manifest r e, Index ix, PrimMonad m) + => Sz ix + -- ^ Size of the create array + -> (ix -> m e) + -- ^ Element generating action + -> m (MArray (PrimState m) r ix e) +makeMArrayS sz f = makeMArrayLinearS sz (f . fromLinearIndex sz) +{-# INLINE makeMArrayS #-} + +-- | Same as `makeMArrayS`, but index supplied to the action is row-major linear index. +-- +-- @since 0.3.0 +makeMArrayLinearS + :: forall r ix e m + . (Manifest r e, Index ix, PrimMonad m) + => Sz ix + -> (Int -> m e) + -> m (MArray (PrimState m) r ix e) +makeMArrayLinearS sz f = do + marr <- unsafeNew sz + loopA_ 0 (< totalElem (sizeOfMArray marr)) (+ 1) (\ !i -> f i >>= unsafeLinearWrite marr i) + return marr +{-# INLINE makeMArrayLinearS #-} + +-- | Just like `makeMArrayS`, but also accepts computation strategy and runs in `IO`. +-- +-- @since 0.3.0 +makeMArray + :: forall r ix e m + . (MonadUnliftIO m, Manifest r e, Index ix) + => Comp + -> Sz ix + -> (ix -> m e) + -> m (MArray RealWorld r ix e) +makeMArray comp sz f = makeMArrayLinear comp sz (f . fromLinearIndex sz) +{-# INLINE makeMArray #-} + +-- | Just like `makeMArrayLinearS`, but also accepts computation strategy and runs in `IO`. +-- +-- @since 0.3.0 +makeMArrayLinear + :: forall r ix e m + . (MonadUnliftIO m, Manifest r e, Index ix) + => Comp + -> Sz ix + -> (Int -> m e) + -> m (MArray RealWorld r ix e) +makeMArrayLinear comp sz f = do + marr <- liftIO $ unsafeNew sz + withScheduler_ comp $ \scheduler -> + withRunInIO $ \run -> + splitLinearlyWithM_ scheduler (totalElem sz) (run . f) (unsafeLinearWrite marr) + return marr +{-# INLINE makeMArrayLinear #-} + +-- | Create a new array by supplying an action that will fill the new blank mutable array. Use +-- `createArray` if you'd like to keep the result of the filling function. +-- +-- ====__Examples__ +-- +-- >>> :set -XTypeApplications +-- >>> import Data.Massiv.Array +-- >>> createArray_ @P @_ @Int Seq (Sz1 2) (\ s marr -> scheduleWork s (writeM marr 0 10) >> scheduleWork s (writeM marr 1 11)) +-- Array P Seq (Sz1 2) +-- [ 10, 11 ] +-- +-- @since 0.3.0 +createArray_ + :: forall r ix e a m + . (Manifest r e, Index ix, MonadUnliftIO m) + => Comp + -- ^ Computation strategy to use after `MArray` gets frozen and onward. + -> Sz ix + -- ^ Size of the newly created array + -> (Scheduler RealWorld () -> MArray RealWorld r ix e -> m a) + -- ^ An action that should fill all elements of the brand new mutable array + -> m (Array r ix e) +createArray_ comp sz action = do + marr <- liftIO $ newMArray' sz + withScheduler_ comp (`action` marr) + liftIO $ unsafeFreeze comp marr +{-# INLINE createArray_ #-} + +-- | Just like `createArray_`, but together with `Array` it returns results of scheduled filling +-- actions. +-- +-- @since 0.3.0 +createArray + :: forall r ix e a m b + . (Manifest r e, Index ix, MonadUnliftIO m) + => Comp + -- ^ Computation strategy to use after `MArray` gets frozen and onward. + -> Sz ix + -- ^ Size of the newly created array + -> (Scheduler RealWorld a -> MArray RealWorld r ix e -> m b) + -- ^ An action that should fill all elements of the brand new mutable array + -> m ([a], Array r ix e) +createArray comp sz action = do + marr <- liftIO $ newMArray' sz + a <- withScheduler comp (`action` marr) + arr <- liftIO $ unsafeFreeze comp marr + return (a, arr) +{-# INLINE createArray #-} + +-- | Create a new array by supplying an action that will fill the new blank mutable array. Use +-- `createArrayS` if you'd like to keep the result of the filling function. +-- +-- ====__Examples__ +-- +-- >>> :set -XTypeApplications +-- >>> import Data.Massiv.Array +-- >>> createArrayS_ @P @_ @Int (Sz1 2) (\ marr -> write marr 0 10 >> write marr 1 12) +-- Array P Seq (Sz1 2) +-- [ 10, 12 ] +-- +-- @since 0.3.0 +createArrayS_ + :: forall r ix e a m + . (Manifest r e, Index ix, PrimMonad m) + => Sz ix + -- ^ Size of the newly created array + -> (MArray (PrimState m) r ix e -> m a) + -- ^ An action that should fill all elements of the brand new mutable array + -> m (Array r ix e) +createArrayS_ sz action = snd <$> createArrayS sz action +{-# INLINE createArrayS_ #-} + +-- | Just like `createArray_`, but together with `Array` it returns the result of the filling action. +-- +-- @since 0.3.0 +createArrayS + :: forall r ix e a m + . (Manifest r e, Index ix, PrimMonad m) + => Sz ix + -- ^ Size of the newly created array + -> (MArray (PrimState m) r ix e -> m a) + -- ^ An action that should fill all elements of the brand new mutable array + -> m (a, Array r ix e) +createArrayS sz action = do + marr <- newMArray' sz + a <- action marr + arr <- unsafeFreeze Seq marr + return (a, arr) +{-# INLINE createArrayS #-} + +-- | Just like `createArrayS_`, but restricted to `ST`. +-- +-- @since 0.3.0 +createArrayST_ + :: forall r ix e a + . (Manifest r e, Index ix) + => Sz ix + -> (forall s. MArray s r ix e -> ST s a) + -> Array r ix e +createArrayST_ sz action = runST $ createArrayS_ sz action +{-# INLINE createArrayST_ #-} + +-- | Just like `createArrayS`, but restricted to `ST`. +-- +-- @since 0.2.6 +createArrayST + :: forall r ix e a + . (Manifest r e, Index ix) + => Sz ix + -> (forall s. MArray s r ix e -> ST s a) + -> (a, Array r ix e) +createArrayST sz action = runST $ createArrayS sz action +{-# INLINE createArrayST #-} + +-- | Sequentially generate a pure array. Much like `makeArray` creates a pure array this +-- function will use `Manifest` interface to generate a pure `Array` in the end, except that +-- computation strategy is set to `Seq`. Element producing function no longer has to be pure +-- but is a stateful action, becuase it is restricted to `PrimMonad` thus allows for sharing +-- the state between computation of each element. +-- +-- ====__Examples__ +-- +-- >>> import Data.Massiv.Array +-- >>> import Data.IORef +-- >>> ref <- newIORef (0 :: Int) +-- >>> generateArrayS (Sz1 6) (\ i -> modifyIORef' ref (+i) >> print i >> pure i) :: IO (Array U Ix1 Int) +-- 0 +-- 1 +-- 2 +-- 3 +-- 4 +-- 5 +-- Array U Seq (Sz1 6) +-- [ 0, 1, 2, 3, 4, 5 ] +-- >>> readIORef ref +-- 15 +-- +-- @since 0.2.6 +generateArrayS + :: forall r ix e m + . (Manifest r e, Index ix, PrimMonad m) + => Sz ix + -- ^ Size of the array + -> (ix -> m e) + -- ^ Element producing action + -> m (Array r ix e) +generateArrayS sz gen = generateArrayLinearS sz (gen . fromLinearIndex sz) +{-# INLINE generateArrayS #-} + +-- | Same as `generateArray` but with action that accepts row-major linear index. +-- +-- @since 0.3.0 +generateArrayLinearS + :: forall r ix e m + . (Manifest r e, Index ix, PrimMonad m) + => Sz ix + -- ^ Resulting size of the array + -> (Int -> m e) + -- ^ Element producing generator + -> m (Array r ix e) +generateArrayLinearS sz gen = do + marr <- unsafeNew sz + loopA_ 0 (< totalElem (sizeOfMArray marr)) (+ 1) $ \i -> gen i >>= unsafeLinearWrite marr i + unsafeFreeze Seq marr +{-# INLINE generateArrayLinearS #-} + +-- | Just like `generateArrayS`, except this generator __will__ respect the supplied computation +-- strategy, and for that reason it is restricted to `IO`. +-- +-- @since 0.2.6 +generateArray + :: forall r ix e m + . (MonadUnliftIO m, Manifest r e, Index ix) + => Comp + -> Sz ix + -> (ix -> m e) + -> m (Array r ix e) +generateArray comp sz f = generateArrayLinear comp sz (f . fromLinearIndex sz) +{-# INLINE generateArray #-} + +-- | Just like `generateArray`, except generating action will receive a row-major linear +-- index. +-- +-- @since 0.3.0 +generateArrayLinear + :: forall r ix e m + . (MonadUnliftIO m, Manifest r e, Index ix) + => Comp + -> Sz ix + -> (Ix1 -> m e) + -> m (Array r ix e) +generateArrayLinear comp sz f = makeMArrayLinear comp sz f >>= liftIO . unsafeFreeze comp +{-# INLINE generateArrayLinear #-} + +-- | Similar to `Data.Massiv.Array.makeSplitSeedArray`, except it will produce a +-- Manifest array and will return back the last unused seed together with all +-- final seeds produced by each scheduled job. This function can be thought of +-- as an unfolding done in parallel while iterating in a customizable manner. +-- +-- @since 1.0.2 +generateSplitSeedArray + :: forall r ix e g it + . (Iterator it, Manifest r e, Index ix) + => it + -- ^ Iterator + -> g + -- ^ Initial seed + -> (forall s. g -> ST s (g, g)) + -- ^ An ST action that can split a seed into two independent seeds. It will + -- be called the same number of times as the number of jobs that will get + -- scheduled during parallelization. Eg. only once for the sequential case. + -> Comp + -- ^ Computation strategy. + -> Sz ix + -- ^ Resulting size of the array. + -> (forall s. Ix1 -> ix -> g -> ST s (e, g)) + -- ^ An ST action that produces a value and the next seed. It takes both + -- versions of the index, in linear and in multi-dimensional forms, as well + -- as the current seeding value. Returns the element for the array cell + -- together with the new seed that will be used for the next element + -- generation + -> (g, [g], Array r ix e) + -- ^ Returned values are: + -- + -- * The final split of the supplied seed. + -- + -- * Results of scheduled jobs in the same order that they where scheduled + -- + -- * Final array that was fully filled using the supplied action and iterator. +generateSplitSeedArray it seed splitSeed comp sz genFunc = + unsafePerformIO $ do + marr <- unsafeNew sz + ref <- newIORef Nothing + res <- withSchedulerR comp $ \scheduler -> do + fin <- stToIO $ + iterTargetFullAccST it scheduler 0 sz seed splitSeed $ \ !i ix !g -> + genFunc i ix g >>= \(x, g') -> g' <$ unsafeLinearWrite marr i x + writeIORef ref $ Just fin + mFin <- readIORef ref + case res of + Finished gs + | Just fin <- mFin -> do + arr <- unsafeFreeze comp marr + pure (fin, gs, arr) + -- This case does not make much sence for array filling and can only + -- happen with a custom 'Iterator' defined outside massiv, therefore it is + -- ok to not support it. + _ -> + error $ + "Parallelized array filling finished prematurely. " + ++ "This feature is not supported by the 'generateSplitSeedArray' function." +{-# INLINE generateSplitSeedArray #-} + +-- | Same as `generateArrayWS`, but use linear indexing instead. +-- +-- @since 0.3.4 +generateArrayLinearWS + :: forall r ix e s m + . (Manifest r e, Index ix, MonadUnliftIO m, PrimMonad m) + => WorkerStates s + -> Sz ix + -> (Int -> s -> m e) + -> m (Array r ix e) +generateArrayLinearWS states sz make = do + marr <- unsafeNew sz + withSchedulerWS_ states $ \schedulerWS -> + splitLinearlyWithStatefulM_ + schedulerWS + (totalElem sz) + make + (unsafeLinearWrite marr) + unsafeFreeze (workerStatesComp states) marr +{-# INLINE generateArrayLinearWS #-} + +-- | Use per worker thread state while generating elements of the array. Very useful for +-- things that are not thread safe. +-- +-- @since 0.3.4 +generateArrayWS + :: forall r ix e s m + . (Manifest r e, Index ix, MonadUnliftIO m, PrimMonad m) + => WorkerStates s + -> Sz ix + -> (ix -> s -> m e) + -> m (Array r ix e) +generateArrayWS states sz make = generateArrayLinearWS states sz (make . fromLinearIndex sz) +{-# INLINE generateArrayWS #-} + +-- | Sequentially unfold an array from the left. +-- +-- ====__Examples__ +-- +-- Create an array with Fibonacci numbers while performing an `IO` action at each iteration. +-- +-- >>> import Data.Massiv.Array +-- >>> unfoldrPrimM_ (Sz1 10) (\(f0, f1) -> (f0, (f1, f0 + f1)) <$ print f1) (0, 1) :: IO (Array P Ix1 Int) +-- 1 +-- 1 +-- 2 +-- 3 +-- 5 +-- 8 +-- 13 +-- 21 +-- 34 +-- 55 +-- Array P Seq (Sz1 10) +-- [ 0, 1, 1, 2, 3, 5, 8, 13, 21, 34 ] +-- +-- @since 0.3.0 +unfoldrPrimM_ + :: forall r ix e a m + . (Manifest r e, Index ix, PrimMonad m) + => Sz ix + -- ^ Size of the desired array + -> (a -> m (e, a)) + -- ^ Unfolding action + -> a + -- ^ Initial accumulator + -> m (Array r ix e) +unfoldrPrimM_ sz gen acc0 = snd <$> unfoldrPrimM sz gen acc0 +{-# INLINE unfoldrPrimM_ #-} + +-- | Same as `unfoldrPrimM_` but do the unfolding with index aware function. +-- +-- @since 0.3.0 +iunfoldrPrimM_ + :: forall r ix e a m + . (Manifest r e, Index ix, PrimMonad m) + => Sz ix + -- ^ Size of the desired array + -> (a -> ix -> m (e, a)) + -- ^ Unfolding action + -> a + -- ^ Initial accumulator + -> m (Array r ix e) +iunfoldrPrimM_ sz gen acc0 = snd <$> iunfoldrPrimM sz gen acc0 +{-# INLINE iunfoldrPrimM_ #-} + +-- | Just like `iunfoldrPrimM_`, but also returns the final value of the accumulator. +-- +-- @since 0.3.0 +iunfoldrPrimM + :: forall r ix e a m + . (Manifest r e, Index ix, PrimMonad m) + => Sz ix + -- ^ Size of the desired array + -> (a -> ix -> m (e, a)) + -- ^ Unfolding action + -> a + -- ^ Initial accumulator + -> m (a, Array r ix e) +iunfoldrPrimM sz gen acc0 = + unsafeCreateArrayS sz $ \marr -> + let sz' = sizeOfMArray marr + in iterLinearM sz' 0 (totalElem sz') 1 (<) acc0 $ \ !i ix !acc -> do + (e, acc') <- gen acc ix + unsafeLinearWrite marr i e + pure acc' +{-# INLINE iunfoldrPrimM #-} + +-- | Just like `iunfoldrPrimM`, but do the unfolding with index aware function. +-- +-- @since 0.3.0 +unfoldrPrimM + :: forall r ix e a m + . (Manifest r e, Index ix, PrimMonad m) + => Sz ix + -- ^ Size of the desired array + -> (a -> m (e, a)) + -- ^ Unfolding action + -> a + -- ^ Initial accumulator + -> m (a, Array r ix e) +unfoldrPrimM sz gen acc0 = + unsafeCreateArrayS sz $ \marr -> + let sz' = sizeOfMArray marr + in loopM 0 (< totalElem sz') (+ 1) acc0 $ \ !i !acc -> do + (e, acc') <- gen acc + unsafeLinearWrite marr i e + pure acc' +{-# INLINE unfoldrPrimM #-} + +-- | Sequentially unfold an array from the left. +-- +-- ====__Examples__ +-- +-- Create an array with Fibonacci numbers starting at the end while performing and `IO` action on +-- the accumulator for each element of the array. +-- +-- >>> import Data.Massiv.Array +-- >>> unfoldlPrimM_ (Sz1 10) (\a@(f0, f1) -> let fn = f0 + f1 in print a >> return ((f1, fn), f0)) (0, 1) :: IO (Array P Ix1 Int) +-- (0,1) +-- (1,1) +-- (1,2) +-- (2,3) +-- (3,5) +-- (5,8) +-- (8,13) +-- (13,21) +-- (21,34) +-- (34,55) +-- Array P Seq (Sz1 10) +-- [ 34, 21, 13, 8, 5, 3, 2, 1, 1, 0 ] +-- +-- @since 0.3.0 +unfoldlPrimM_ + :: forall r ix e a m + . (Manifest r e, Index ix, PrimMonad m) + => Sz ix + -- ^ Size of the desired array + -> (a -> m (a, e)) + -- ^ Unfolding action + -> a + -- ^ Initial accumulator + -> m (Array r ix e) +unfoldlPrimM_ sz gen acc0 = snd <$> unfoldlPrimM sz gen acc0 +{-# INLINE unfoldlPrimM_ #-} + +-- | Same as `unfoldlPrimM_` but do the unfolding with index aware function. +-- +-- @since 0.3.0 +iunfoldlPrimM_ + :: forall r ix e a m + . (Manifest r e, Index ix, PrimMonad m) + => Sz ix + -- ^ Size of the desired array + -> (a -> ix -> m (a, e)) + -- ^ Unfolding action + -> a + -- ^ Initial accumulator + -> m (Array r ix e) +iunfoldlPrimM_ sz gen acc0 = snd <$> iunfoldlPrimM sz gen acc0 +{-# INLINE iunfoldlPrimM_ #-} + +-- | Just like `iunfoldlPrimM_`, but also returns the final value of the accumulator. +-- +-- @since 0.3.0 +iunfoldlPrimM + :: forall r ix e a m + . (Manifest r e, Index ix, PrimMonad m) + => Sz ix + -- ^ Size of the desired array + -> (a -> ix -> m (a, e)) + -- ^ Unfolding action + -> a + -- ^ Initial accumulator + -> m (a, Array r ix e) +iunfoldlPrimM sz gen acc0 = + unsafeCreateArrayS sz $ \marr -> + let sz' = sizeOfMArray marr + in iterLinearM sz' (totalElem sz' - 1) 0 (negate 1) (>=) acc0 $ \ !i ix !acc -> do + (acc', e) <- gen acc ix + unsafeLinearWrite marr i e + pure acc' +{-# INLINE iunfoldlPrimM #-} + +-- | Just like `iunfoldlPrimM`, but do the unfolding with index aware function. +-- +-- @since 0.3.0 +unfoldlPrimM + :: forall r ix e a m + . (Manifest r e, Index ix, PrimMonad m) + => Sz ix + -- ^ Size of the desired array + -> (a -> m (a, e)) + -- ^ Unfolding action + -> a + -- ^ Initial accumulator + -> m (a, Array r ix e) +unfoldlPrimM sz gen acc0 = + unsafeCreateArrayS sz $ \marr -> + let sz' = sizeOfMArray marr + in loopDeepM 0 (< totalElem sz') (+ 1) acc0 $ \ !i !acc -> do + (acc', e) <- gen acc + unsafeLinearWrite marr i e + pure acc' +{-# INLINE unfoldlPrimM #-} + +-- | Sequentially loop over a mutable array while reading each element and applying an +-- action to it. There is no mutation to the array, unless the action itself modifies it. +-- +-- @since 0.4.0 +forPrimM_ + :: (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> (e -> m ()) -> m () +forPrimM_ marr f = + loopA_ 0 (< totalElem (sizeOfMArray marr)) (+ 1) (unsafeLinearRead marr >=> f) +{-# INLINE forPrimM_ #-} + +-- | Sequentially loop over a mutable array while modifying each element with an action. +-- +-- @since 0.4.0 +forPrimM + :: (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> (e -> m e) -> m () +forPrimM marr f = + loopA_ 0 (< totalElem (sizeOfMArray marr)) (+ 1) (unsafeLinearModify marr f) +{-# INLINE forPrimM #-} + +-- | Sequentially loop over a mutable array while reading each element and applying an +-- index aware action to it. There is no mutation to the array, unless the +-- action itself modifies it. +-- +-- @since 0.4.0 +iforPrimM_ + :: (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> (ix -> e -> m ()) -> m () +iforPrimM_ marr f = iforLinearPrimM_ marr (f . fromLinearIndex (sizeOfMArray marr)) +{-# INLINE iforPrimM_ #-} + +-- | Sequentially loop over a mutable array while modifying each element with an index aware action. +-- +-- @since 0.4.0 +iforPrimM + :: (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> (ix -> e -> m e) -> m () +iforPrimM marr f = iforLinearPrimM marr (f . fromLinearIndex (sizeOfMArray marr)) +{-# INLINE iforPrimM #-} + +-- | Sequentially loop over a mutable array while reading each element and applying a +-- linear index aware action to it. There is no mutation to the array, unless the action +-- itself modifies it. +-- +-- @since 0.4.0 +iforLinearPrimM_ + :: (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> (Int -> e -> m ()) -> m () +iforLinearPrimM_ marr f = + loopA_ 0 (< totalElem (sizeOfMArray marr)) (+ 1) (\i -> unsafeLinearRead marr i >>= f i) +{-# INLINE iforLinearPrimM_ #-} + +-- | Sequentially loop over a mutable array while modifying each element with an index aware action. +-- +-- @since 0.4.0 +iforLinearPrimM + :: (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> (Int -> e -> m e) -> m () +iforLinearPrimM marr f = + loopA_ 0 (< totalElem (sizeOfMArray marr)) (+ 1) (\i -> unsafeLinearModify marr (f i) i) +{-# INLINE iforLinearPrimM #-} + +-- | Sequentially loop over the intersection of two mutable arrays while reading +-- elements from both and applying an action to it. There is no mutation to the +-- actual arrays, unless the action itself modifies either one of them. +-- +-- @since 1.0.0 +for2PrimM_ + :: forall r1 r2 e1 e2 ix m + . (PrimMonad m, Index ix, Manifest r1 e1, Manifest r2 e2) + => MArray (PrimState m) r1 ix e1 + -> MArray (PrimState m) r2 ix e2 + -> (e1 -> e2 -> m ()) + -> m () +for2PrimM_ m1 m2 f = ifor2PrimM_ m1 m2 (const f) +{-# INLINE for2PrimM_ #-} + +-- | Same as `for2PrimM_`, but with index aware action. +-- +-- @since 1.0.0 +ifor2PrimM_ + :: forall r1 r2 e1 e2 ix m + . (PrimMonad m, Index ix, Manifest r1 e1, Manifest r2 e2) + => MArray (PrimState m) r1 ix e1 + -> MArray (PrimState m) r2 ix e2 + -> (ix -> e1 -> e2 -> m ()) + -> m () +ifor2PrimM_ m1 m2 f = do + let sz = liftIndex2 min (unSz (sizeOfMArray m1)) (unSz (sizeOfMArray m2)) + iterA_ zeroIndex sz oneIndex (<) $ \ix -> do + e1 <- unsafeRead m1 ix + e2 <- unsafeRead m2 ix + f ix e1 e2 +{-# INLINE ifor2PrimM_ #-} + +-- | Same as `withMArray_`, but allows to keep artifacts of scheduled tasks. +-- +-- @since 0.5.0 +withMArray + :: (Manifest r e, Index ix, MonadUnliftIO m) + => Array r ix e + -> (Scheduler RealWorld a -> MArray RealWorld r ix e -> m b) + -> m ([a], Array r ix e) +withMArray arr action = do + marr <- thaw arr + xs <- withScheduler (getComp arr) (`action` marr) + liftIO ((,) xs <$> unsafeFreeze (getComp arr) marr) +{-# INLINE withMArray #-} + +-- | Create a copy of a pure array, mutate it in place and return its frozen version. The big +-- difference between `withMArrayS` is that it's not only gonna respect the computation strategy +-- supplied to it while making a copy, but it will also pass extra argumens to the action that +-- suppose to modify the mutable copy of the source array. These two extra arguments are: +-- +-- * Number of capabilities derived from the `Comp`utation strategy of the array. +-- +-- * An action that can be used to schedule arbitrary number of jobs that will be executed in +-- parallel. +-- +-- * And, of course, the mutable array itself. +-- +-- @since 0.5.0 +withMArray_ + :: (Manifest r e, Index ix, MonadUnliftIO m) + => Array r ix e + -> (Scheduler RealWorld () -> MArray RealWorld r ix e -> m a) + -> m (Array r ix e) +withMArray_ arr action = do + marr <- thaw arr + withScheduler_ (getComp arr) (`action` marr) + liftIO $ unsafeFreeze (getComp arr) marr +{-# INLINE withMArray_ #-} + +-- | Same as `withMArray_`, but the array supplied to this function can be any loadable +-- array. For that reason it will be faster if supplied array is delayed. +-- +-- @since 0.6.1 +withLoadMArray_ + :: forall r ix e r' m b + . (Load r' ix e, Manifest r e, MonadUnliftIO m) + => Array r' ix e + -> (Scheduler RealWorld () -> MArray RealWorld r ix e -> m b) + -> m (Array r ix e) +withLoadMArray_ arr action = do + marr <- loadArray arr + withScheduler_ (getComp arr) (`action` marr) + liftIO $ unsafeFreeze (getComp arr) marr +{-# INLINE [2] withLoadMArray_ #-} + +{-# RULES +"withLoadMArray_/withMArray_" [~2] withLoadMArray_ = withMArray_ +"withLoadMArrayS/withMArrayS" [~2] withLoadMArrayS = withMArrayS +"withLoadMArrayS_/withMArrayS_" [~2] withLoadMArrayS_ = withMArrayS_ + #-} + +-- | Create a copy of a pure array, mutate it in place and return its frozen version. The important +-- benefit over doing a manual `thawS` followed by a `freezeS` is that an array will only be copied +-- once. +-- +-- @since 0.5.0 +withMArrayS + :: (Manifest r e, Index ix, PrimMonad m) + => Array r ix e + -> (MArray (PrimState m) r ix e -> m a) + -> m (a, Array r ix e) +withMArrayS arr action = do + marr <- thawS arr + a <- action marr + (,) a <$> unsafeFreeze (getComp arr) marr +{-# INLINE withMArrayS #-} + +-- | Same as `withMArrayS`, except it discards the value produced by the supplied action +-- +-- @since 0.5.0 +withMArrayS_ + :: (Manifest r e, Index ix, PrimMonad m) + => Array r ix e + -> (MArray (PrimState m) r ix e -> m a) + -> m (Array r ix e) +withMArrayS_ arr action = snd <$> withMArrayS arr action +{-# INLINE withMArrayS_ #-} + +-- | Same as `withMArrayS`, but will work with any loadable array. +-- +-- @since 0.6.1 +withLoadMArrayS + :: forall r ix e r' m a + . (Load r' ix e, Manifest r e, PrimMonad m) + => Array r' ix e + -> (MArray (PrimState m) r ix e -> m a) + -> m (a, Array r ix e) +withLoadMArrayS arr action = do + marr <- loadArrayS arr + a <- action marr + (,) a <$> unsafeFreeze (getComp arr) marr +{-# INLINE [2] withLoadMArrayS #-} + +-- | Same as `withMArrayS_`, but will work with any loadable array. +-- +-- @since 0.6.1 +withLoadMArrayS_ + :: forall r ix e r' m a + . (Load r' ix e, Manifest r e, PrimMonad m) + => Array r' ix e + -> (MArray (PrimState m) r ix e -> m a) + -> m (Array r ix e) +withLoadMArrayS_ arr action = snd <$> withLoadMArrayS arr action +{-# INLINE [2] withLoadMArrayS_ #-} + +-- | Same as `withMArrayS` but in `ST`. This is not only pure, but also the safest way to do +-- mutation to the array. +-- +-- @since 0.5.0 +withMArrayST + :: (Manifest r e, Index ix) + => Array r ix e + -> (forall s. MArray s r ix e -> ST s a) + -> (a, Array r ix e) +withMArrayST arr f = runST $ withMArrayS arr f +{-# INLINE withMArrayST #-} + +-- | Same as `withMArrayS` but in `ST`. This is not only pure, but also the safest way to do +-- mutation to the array. +-- +-- @since 0.5.0 +withMArrayST_ + :: (Manifest r e, Index ix) => Array r ix e -> (forall s. MArray s r ix e -> ST s a) -> Array r ix e +withMArrayST_ arr f = runST $ withMArrayS_ arr f +{-# INLINE withMArrayST_ #-} + +-- | Same as `withMArrayST`, but works with any loadable array. +-- +-- @since 0.6.1 +withLoadMArrayST + :: forall r ix e r' a + . (Load r' ix e, Manifest r e) + => Array r' ix e + -> (forall s. MArray s r ix e -> ST s a) + -> (a, Array r ix e) +withLoadMArrayST arr f = runST $ withLoadMArrayS arr f +{-# INLINE [2] withLoadMArrayST #-} + +-- | Same as `withMArrayST_`, but works with any loadable array. +-- +-- @since 0.6.1 +withLoadMArrayST_ + :: forall r ix e r' a + . (Load r' ix e, Manifest r e) + => Array r' ix e + -> (forall s. MArray s r ix e -> ST s a) + -> Array r ix e +withLoadMArrayST_ arr f = runST $ withLoadMArrayS_ arr f +{-# INLINE [2] withLoadMArrayST_ #-} + +-- | /O(1)/ - Lookup an element in the mutable array. Returns `Nothing` when index is out of bounds. +-- +-- @since 0.1.0 +read + :: (Manifest r e, Index ix, PrimMonad m) + => MArray (PrimState m) r ix e + -> ix + -> m (Maybe e) +read marr ix = + if isSafeIndex (sizeOfMArray marr) ix + then Just <$> unsafeRead marr ix + else return Nothing +{-# INLINE read #-} + +-- | /O(1)/ - Same as `read`, but throws `IndexOutOfBoundsException` on an invalid index. +-- +-- @since 0.4.0 +readM + :: (Manifest r e, Index ix, PrimMonad m, MonadThrow m) + => MArray (PrimState m) r ix e + -> ix + -> m e +readM marr ix = + read marr ix >>= \case + Just e -> pure e + Nothing -> throwM $ IndexOutOfBoundsException (sizeOfMArray marr) ix +{-# INLINE readM #-} + +-- | /O(1)/ - Write an element into the cell of a mutable array. Returns `False` when index is out +-- of bounds. +-- +-- @since 0.1.0 +write :: (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> ix -> e -> m Bool +write marr ix e = + if isSafeIndex (sizeOfMArray marr) ix + then unsafeWrite marr ix e >> pure True + else pure False +{-# INLINE write #-} + +-- | /O(1)/ - Write an element into the cell of a mutable array. Same as `write` function +-- in case of an out of bounds index it is noop, but unlike `write`, there is no +-- information is returned about was the writing of element successful or not. In other +-- words, just like `writeM`, but doesn't throw an exception. +-- +-- @since 0.4.4 +write_ :: (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> ix -> e -> m () +write_ marr ix = when (isSafeIndex (sizeOfMArray marr) ix) . unsafeWrite marr ix +{-# INLINE write_ #-} + +-- | /O(1)/ - Same as `write`, but throws `IndexOutOfBoundsException` on an invalid index. +-- +-- @since 0.4.0 +writeM + :: (Manifest r e, Index ix, PrimMonad m, MonadThrow m) + => MArray (PrimState m) r ix e -> ix -> e -> m () +writeM marr ix e = + write marr ix e >>= (`unless` throwM (IndexOutOfBoundsException (sizeOfMArray marr) ix)) +{-# INLINE writeM #-} + +-- | /O(1)/ - Modify an element in the cell of a mutable array with a supplied +-- action. Returns the previous value, if index was not out of bounds. +-- +-- @since 0.1.0 +modify + :: (Manifest r e, Index ix, PrimMonad m) + => MArray (PrimState m) r ix e + -- ^ Array to mutate. + -> (e -> m e) + -- ^ Monadic action that modifies the element + -> ix + -- ^ Index at which to perform modification. + -> m (Maybe e) +modify marr f ix = + if isSafeIndex (sizeOfMArray marr) ix + then Just <$> unsafeModify marr f ix + else return Nothing +{-# INLINE modify #-} + +-- | /O(1)/ - Same as `modify`, except that neither the previous value, nor any +-- information on whether the modification was successful are returned. In other words, +-- just like `modifyM_`, but doesn't throw an exception. +-- +-- @since 0.4.4 +modify_ + :: (Manifest r e, Index ix, PrimMonad m) + => MArray (PrimState m) r ix e + -- ^ Array to mutate. + -> (e -> m e) + -- ^ Monadic action that modifies the element + -> ix + -- ^ Index at which to perform modification. + -> m () +modify_ marr f ix = when (isSafeIndex (sizeOfMArray marr) ix) $ void $ unsafeModify marr f ix +{-# INLINE modify_ #-} + +-- | /O(1)/ - Modify an element in the cell of a mutable array with a supplied +-- action. Throws an `IndexOutOfBoundsException` exception for invalid index and returns +-- the previous value otherwise. +-- +-- @since 0.4.0 +modifyM + :: (Manifest r e, Index ix, PrimMonad m, MonadThrow m) + => MArray (PrimState m) r ix e + -- ^ Array to mutate. + -> (e -> m e) + -- ^ Monadic action that modifies the element + -> ix + -- ^ Index at which to perform modification. + -> m e +modifyM marr f ix + | isSafeIndex (sizeOfMArray marr) ix = unsafeModify marr f ix + | otherwise = throwM (IndexOutOfBoundsException (sizeOfMArray marr) ix) +{-# INLINE modifyM #-} + +-- | /O(1)/ - Same as `modifyM`, but discard the returned element +-- +-- ====__Examples__ +-- +-- >>> :set -XTypeApplications +-- >>> import Control.Monad.ST +-- >>> import Data.Massiv.Array +-- >>> runST $ newMArray' @P @Ix1 @Int (Sz1 3) >>= (\ma -> modifyM_ ma (pure . (+10)) 1 >> freezeS ma) +-- Array P Seq (Sz1 3) +-- [ 0, 10, 0 ] +-- +-- @since 0.4.0 +modifyM_ + :: (Manifest r e, Index ix, PrimMonad m, MonadThrow m) + => MArray (PrimState m) r ix e + -- ^ Array to mutate. + -> (e -> m e) + -- ^ Monadic action that modifies the element + -> ix + -- ^ Index at which to perform modification. + -> m () +modifyM_ marr f ix = void $ modifyM marr f ix +{-# INLINE modifyM_ #-} + +-- | /O(1)/ - Same as `swapM`, but instead of throwing an exception returns `Nothing` when +-- either one of the indices is out of bounds and `Just` elements under those indices +-- otherwise. +-- +-- @since 0.1.0 +swap + :: (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> ix -> ix -> m (Maybe (e, e)) +swap marr ix1 ix2 = + let !sz = sizeOfMArray marr + in if isSafeIndex sz ix1 && isSafeIndex sz ix2 + then Just <$> unsafeSwap marr ix1 ix2 + else pure Nothing +{-# INLINE swap #-} + +-- | /O(1)/ - Same as `swap`, but instead of returning `Nothing` it does nothing. In other +-- words, it is similar to `swapM_`, but does not throw any exceptions. +-- +-- @since 0.4.4 +swap_ :: (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> ix -> ix -> m () +swap_ marr ix1 ix2 = + let !sz = sizeOfMArray marr + in when (isSafeIndex sz ix1 && isSafeIndex sz ix2) $ void $ unsafeSwap marr ix1 ix2 +{-# INLINE swap_ #-} + +-- | /O(1)/ - Swap two elements in a mutable array under the supplied indices. Throws an +-- `IndexOutOfBoundsException` when either one of the indices is out of bounds and +-- elements under those indices otherwise. +-- +-- @since 0.4.0 +swapM + :: (Manifest r e, Index ix, PrimMonad m, MonadThrow m) + => MArray (PrimState m) r ix e + -> ix + -- ^ Index for the first element, which will be returned as the first element in the + -- tuple. + -> ix + -- ^ Index for the second element, which will be returned as the second element in + -- the tuple. + -> m (e, e) +swapM marr ix1 ix2 + | not (isSafeIndex sz ix1) = throwM $ IndexOutOfBoundsException (sizeOfMArray marr) ix1 + | not (isSafeIndex sz ix2) = throwM $ IndexOutOfBoundsException (sizeOfMArray marr) ix2 + | otherwise = unsafeSwap marr ix1 ix2 + where + !sz = sizeOfMArray marr +{-# INLINE swapM #-} + +-- | /O(1)/ - Same as `swapM`, but discard the returned elements +-- +-- @since 0.4.0 +swapM_ + :: (Manifest r e, Index ix, PrimMonad m, MonadThrow m) + => MArray (PrimState m) r ix e + -> ix + -> ix + -> m () +swapM_ marr ix1 ix2 = void $ swapM marr ix1 ix2 +{-# INLINE swapM_ #-} + +-- | Swap elements in the intersection of two mutable arrays starting at the +-- initial index. +-- +-- @since 1.0.0 +zipSwapM_ + :: forall r1 r2 ix e m s + . (MonadPrim s m, Manifest r2 e, Manifest r1 e, Index ix) + => ix + -> MArray s r1 ix e + -> MArray s r2 ix e + -> m () +zipSwapM_ startIx m1 m2 = do + let sz1 = sizeOfMArray m1 + sz2 = sizeOfMArray m2 + sz = liftIndex2 min (unSz sz1) (unSz sz2) + iterA_ startIx sz oneIndex (<) $ \ix -> do + let i1 = toLinearIndex sz1 ix + i2 = toLinearIndex sz2 ix + e1 <- unsafeLinearRead m1 i1 + e2 <- unsafeLinearRead m2 i2 + unsafeLinearWrite m2 i2 e1 + unsafeLinearWrite m1 i1 e2 +{-# INLINE zipSwapM_ #-} + +-- | Get the size of a mutable array. +-- +-- @since 0.1.0 +msize :: (Manifest r e, Index ix) => MArray s r ix e -> Sz ix +msize = sizeOfMArray +{-# DEPRECATED msize "In favor of `sizeOfMArray`" #-} 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 @@ -1,41 +1,44 @@ {-# LANGUAGE ExplicitForAll #-} {-# LANGUAGE FlexibleContexts #-} + -- | -- Module : Data.Massiv.Array.Mutable.Algorithms --- Copyright : (c) Alexey Kuleshevich 2019 +-- Copyright : (c) Alexey Kuleshevich 2019-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Mutable.Algorithms - ( quicksortM_ - , unstablePartitionM - , iterateUntilM - ) where +module Data.Massiv.Array.Mutable.Algorithms ( + quicksortM_, + quicksortByM_, + unstablePartitionM, + iterateUntilM, +) where -import Data.Massiv.Array.Ops.Sort import Data.Massiv.Array.Manifest.Internal (iterateUntilM) +import Data.Massiv.Array.Ops.Sort import Data.Massiv.Core.Common - -- | Partition elements of the supplied mutable vector according to the predicate. -- -- ==== __Example__ -- -- >>> import Data.Massiv.Array as A -- >>> import Data.Massiv.Array.Mutable.Algorithms +-- >>> :set -XOverloadedLists -- >>> m <- thaw ([2,1,50,10,20,8] :: Array P Ix1 Int) --- >>> unstablePartitionM m (<= 10) +-- >>> unstablePartitionM m (pure . (<= 10)) -- 4 -- >>> freeze Seq m -- Array P Seq (Sz1 6) -- [ 2, 1, 8, 10, 20, 50 ] -- --- @since 0.3.2 -unstablePartitionM :: - forall r e m. (Mutable r Ix1 e, PrimMonad m) - => MArray (PrimState m) r Ix1 e - -> (e -> Bool) -- ^ Predicate +-- @since 1.0.0 +unstablePartitionM + :: forall r e m + . (Manifest r e, PrimMonad m) + => MVector (PrimState m) r e + -> (e -> m Bool) + -- ^ Predicate -> m Ix1 -unstablePartitionM marr f = unsafeUnstablePartitionRegionM marr f 0 (unSz (msize marr) - 1) +unstablePartitionM marr f = unsafeUnstablePartitionRegionM marr f 0 (unSz (sizeOfMArray marr) - 1) diff --git a/src/Data/Massiv/Array/Mutable/Atomic.hs b/src/Data/Massiv/Array/Mutable/Atomic.hs --- a/src/Data/Massiv/Array/Mutable/Atomic.hs +++ b/src/Data/Massiv/Array/Mutable/Atomic.hs @@ -1,30 +1,29 @@ -{-# LANGUAGE BangPatterns #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} + -- | -- Module : Data.Massiv.Array.Mutable.Atomic --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Mutable.Atomic - ( -- * Atomic element-wise mutation - atomicReadIntArray - , atomicWriteIntArray - , atomicModifyIntArray - , atomicAddIntArray - , atomicSubIntArray - , atomicAndIntArray - , atomicNandIntArray - , atomicOrIntArray - , atomicXorIntArray - , casIntArray - ) where +module Data.Massiv.Array.Mutable.Atomic ( + -- * Atomic element-wise mutation + atomicReadIntArray, + atomicWriteIntArray, + atomicModifyIntArray, + atomicAddIntArray, + atomicSubIntArray, + atomicAndIntArray, + atomicNandIntArray, + atomicOrIntArray, + atomicXorIntArray, + casIntArray, +) where import Control.Monad.Primitive import Data.Massiv.Array.Manifest.Primitive @@ -35,115 +34,110 @@ -- | Atomically read an `Int` element from the array -- -- @since 0.3.0 -atomicReadIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> m (Maybe Int) +atomicReadIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> m (Maybe Int) atomicReadIntArray marr ix - | isSafeIndex (msize marr) ix = Just <$> unsafeAtomicReadIntArray marr ix + | isSafeIndex (sizeOfMArray marr) ix = Just <$> unsafeAtomicReadIntArray marr ix | otherwise = pure Nothing {-# INLINE atomicReadIntArray #-} - -- | Atomically write an `Int` element int the array. Returns `True` if supplied index was correct -- and write was successfull. -- -- @since 0.3.0 -atomicWriteIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Bool +atomicWriteIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Bool atomicWriteIntArray marr ix f - | isSafeIndex (msize marr) ix = unsafeAtomicWriteIntArray marr ix f >> pure True + | isSafeIndex (sizeOfMArray marr) ix = unsafeAtomicWriteIntArray marr ix f >> pure True | otherwise = pure False {-# INLINE atomicWriteIntArray #-} - -- | Atomically CAS (Compare-and-Swap) an `Int` in the array. Returns the old value. -- -- @since 0.3.0 -casIntArray :: - (Index ix, PrimMonad m) - => MArray (PrimState m) P ix Int -- ^ Array to mutate - -> ix -- ^ Index at which to mutate - -> Int -- ^ Expected value - -> Int -- ^ New value +casIntArray + :: (Index ix, PrimMonad m) + => MArray (PrimState m) P ix Int + -- ^ Array to mutate + -> ix + -- ^ Index at which to mutate + -> Int + -- ^ Expected value + -> Int + -- ^ New value -> m (Maybe Int) casIntArray marr ix e n - | isSafeIndex (msize marr) ix = Just <$> unsafeCasIntArray marr ix e n + | isSafeIndex (sizeOfMArray marr) ix = Just <$> unsafeCasIntArray marr ix e n | otherwise = pure Nothing {-# INLINE casIntArray #-} - -- | Atomically modify an `Int` element of the array. Returns the old value, unless the -- supplied index was out of bounds. -- -- @since 0.3.0 -atomicModifyIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> (Int -> Int) -> m (Maybe Int) +atomicModifyIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> (Int -> Int) -> m (Maybe Int) atomicModifyIntArray marr ix f - | isSafeIndex (msize marr) ix = Just <$> unsafeAtomicModifyIntArray marr ix f + | isSafeIndex (sizeOfMArray marr) ix = Just <$> unsafeAtomicModifyIntArray marr ix f | otherwise = pure Nothing {-# INLINE atomicModifyIntArray #-} - -- | Atomically add to an `Int` element in the array. Returns the old value. -- -- @since 0.3.0 -atomicAddIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m (Maybe Int) +atomicAddIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m (Maybe Int) atomicAddIntArray marr ix e - | isSafeIndex (msize marr) ix = Just <$> unsafeAtomicAddIntArray marr ix e + | isSafeIndex (sizeOfMArray marr) ix = Just <$> unsafeAtomicAddIntArray marr ix e | otherwise = pure Nothing {-# INLINE atomicAddIntArray #-} - -- | Atomically subtract from an `Int` element in the array. Returns the old value. -- -- @since 0.3.0 -atomicSubIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m (Maybe Int) +atomicSubIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m (Maybe Int) atomicSubIntArray marr ix e - | isSafeIndex (msize marr) ix = Just <$> unsafeAtomicSubIntArray marr ix e + | isSafeIndex (sizeOfMArray marr) ix = Just <$> unsafeAtomicSubIntArray marr ix e | otherwise = pure Nothing {-# INLINE atomicSubIntArray #-} - -- | Atomically AND an `Int` element in the array. Returns the old value. -- -- @since 0.3.0 -atomicAndIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m (Maybe Int) +atomicAndIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m (Maybe Int) atomicAndIntArray marr ix e - | isSafeIndex (msize marr) ix = Just <$> unsafeAtomicAndIntArray marr ix e + | isSafeIndex (sizeOfMArray marr) ix = Just <$> unsafeAtomicAndIntArray marr ix e | otherwise = pure Nothing {-# INLINE atomicAndIntArray #-} - -- | Atomically NAND an `Int` element in the array. Returns the old value. -- -- @since 0.3.0 -atomicNandIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m (Maybe Int) +atomicNandIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m (Maybe Int) atomicNandIntArray marr ix e - | isSafeIndex (msize marr) ix = Just <$> unsafeAtomicNandIntArray marr ix e + | isSafeIndex (sizeOfMArray marr) ix = Just <$> unsafeAtomicNandIntArray marr ix e | otherwise = pure Nothing {-# INLINE atomicNandIntArray #-} - -- | Atomically OR an `Int` element in the array. Returns the old value. -- -- @since 0.3.0 -atomicOrIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m (Maybe Int) +atomicOrIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m (Maybe Int) atomicOrIntArray marr ix e - | isSafeIndex (msize marr) ix = Just <$> unsafeAtomicOrIntArray marr ix e + | isSafeIndex (sizeOfMArray marr) ix = Just <$> unsafeAtomicOrIntArray marr ix e | otherwise = pure Nothing {-# INLINE atomicOrIntArray #-} - -- | Atomically XOR an `Int` element in the array. Returns the old value. -- -- @since 0.3.0 -atomicXorIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m (Maybe Int) +atomicXorIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m (Maybe Int) atomicXorIntArray marr ix e - | isSafeIndex (msize marr) ix = Just <$> unsafeAtomicXorIntArray marr ix e + | isSafeIndex (sizeOfMArray marr) ix = Just <$> unsafeAtomicXorIntArray marr ix e | otherwise = pure Nothing {-# INLINE atomicXorIntArray #-} diff --git a/src/Data/Massiv/Array/Mutable/Internal.hs b/src/Data/Massiv/Array/Mutable/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Array/Mutable/Internal.hs @@ -0,0 +1,78 @@ +{-# LANGUAGE ExplicitForAll #-} + +-- | +-- Module : Data.Massiv.Array.Mutable.Internal +-- Copyright : (c) Alexey Kuleshevich 2018-2022 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +module Data.Massiv.Array.Mutable.Internal ( + unsafeCreateArray, + unsafeCreateArray_, + unsafeCreateArrayS, +) where + +import Control.Scheduler +import Data.Massiv.Core.Common + +-- | Same as `Data.Massiv.Array.Mutable.createArrayS`, but memory will not be initialized +-- and for unboxed types might contain garbage. +-- +-- @since 0.5.0 +unsafeCreateArrayS + :: forall r ix e a m + . (Manifest r e, Index ix, PrimMonad m) + => Sz ix + -- ^ Size of the newly created array + -> (MArray (PrimState m) r ix e -> m a) + -- ^ An action that should fill all elements of the brand new mutable array + -> m (a, Array r ix e) +unsafeCreateArrayS sz action = do + marr <- unsafeNew sz + a <- action marr + arr <- unsafeFreeze Seq marr + return (a, arr) +{-# INLINE unsafeCreateArrayS #-} + +-- | Same as `Data.Massiv.Array.Mutable.createArray`, but memory will not be initialized +-- and for unboxed types might contain garbage. +-- +-- @since 0.5.0 +unsafeCreateArray + :: forall r ix e a m b + . (Manifest r e, Index ix, MonadUnliftIO m) + => Comp + -- ^ Computation strategy to use after `MArray` gets frozen and onward. + -> Sz ix + -- ^ Size of the newly created array + -> (Scheduler RealWorld a -> MArray RealWorld r ix e -> m b) + -- ^ An action that should fill all elements of the brand new mutable array + -> m ([a], Array r ix e) +unsafeCreateArray comp sz action = do + marr <- liftIO $ unsafeNew sz + a <- withScheduler comp (`action` marr) + arr <- liftIO $ unsafeFreeze comp marr + return (a, arr) +{-# INLINE unsafeCreateArray #-} + +-- | Same as `Data.Massiv.Array.Mutable.createArray_`, but memory will not be initialized +-- and for unboxed types might contain garbage. +-- +-- @since 0.5.0 +unsafeCreateArray_ + :: forall r ix e a m b + . (Manifest r e, Index ix, MonadUnliftIO m) + => Comp + -- ^ Computation strategy to use after `MArray` gets frozen and onward. + -> Sz ix + -- ^ Size of the newly created array + -> (Scheduler RealWorld a -> MArray RealWorld r ix e -> m b) + -- ^ An action that should fill all elements of the brand new mutable array + -> m (Array r ix e) +unsafeCreateArray_ comp sz action = do + marr <- liftIO $ unsafeNew sz + withScheduler_ comp (`action` marr) + arr <- liftIO $ unsafeFreeze comp marr + return arr +{-# INLINE unsafeCreateArray_ #-} 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 @@ -3,477 +3,1388 @@ {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} --- | --- Module : Data.Massiv.Array.Numeric --- Copyright : (c) Alexey Kuleshevich 2018-2019 --- License : BSD3 --- Maintainer : Alexey Kuleshevich --- Stability : experimental --- Portability : non-portable --- -module Data.Massiv.Array.Numeric - ( -- * Num - (.+.) - , (.+) - , (+.) - , (.-.) - , (.-) - , (-.) - , (.*.) - , (.*) - , (*.) - , (.^) - , (|*|) - , multiplyTransposed - , identityMatrix - , negateA - , absA - , signumA - , fromIntegerA - -- * Integral - , quotA - , remA - , divA - , modA - , quotRemA - , divModA - -- * Fractional - , (./.) - , (./) - , (.^^) - , recipA - , fromRationalA - -- * Floating - , piA - , expA - , logA - , sqrtA - , (.**) - , logBaseA - , sinA - , cosA - , tanA - , asinA - , acosA - , atanA - , sinhA - , coshA - , tanhA - , asinhA - , acoshA - , atanhA - -- * RealFrac - , truncateA - , roundA - , ceilingA - , floorA - -- * RealFloat - , atan2A - ) 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 -import Data.Massiv.Array.Ops.Transform as A -import Data.Massiv.Core -import Data.Massiv.Core.Common -import Data.Massiv.Core.Operations -import Data.Massiv.Core.Index.Internal (Sz(SafeSz)) -import Prelude as P - - -infixr 8 .^, .^^ -infixl 7 .*., .*, *., ./., ./, `quotA`, `remA`, `divA`, `modA` -infixl 6 .+., .+, +., .-., .-, -. - -liftArray2Matching - :: (Source r1 ix a, Source r2 ix b) - => (a -> b -> e) -> Array r1 ix a -> Array r2 ix b -> Array D ix e -liftArray2Matching f !arr1 !arr2 - | sz1 == sz2 = - makeArray - (getComp arr1 <> getComp arr2) - sz1 - (\ !ix -> f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)) - | otherwise = throw $ SizeMismatchException (size arr1) (size arr2) - where - sz1 = size arr1 - sz2 = size arr2 -{-# INLINE liftArray2Matching #-} - -liftArray2M :: - (Load r ix e, Numeric r e, MonadThrow m) - => (e -> e -> e) - -> Array r ix e - -> Array r ix e - -> m (Array r ix e) -liftArray2M f a1 a2 - | size a1 == size a2 = pure $ unsafeLiftArray2 f a1 a2 - | otherwise = throwM $ SizeMismatchException (size a1) (size a2) -{-# INLINE liftArray2M #-} - - -liftNumericArray2M :: - (Load r ix e, MonadThrow m) - => (Array r ix e -> Array r ix e -> Array r ix e) - -> Array r ix e - -> Array r ix e - -> m (Array r ix e) -liftNumericArray2M f a1 a2 - | size a1 == size a2 = pure $ f a1 a2 - | otherwise = throwM $ SizeMismatchException (size a1) (size a2) -{-# INLINE liftNumericArray2M #-} - - --- | Add two arrays together pointwise. Throws `SizeMismatchException` if arrays sizes do --- not match. --- --- @since 0.4.0 -(.+.) :: - (Load r ix e, Numeric r e, MonadThrow m) => Array r ix e -> Array r ix e -> m (Array r ix e) -(.+.) = liftNumericArray2M additionPointwise -{-# INLINE (.+.) #-} - --- | Add a scalar to each element of the array. Array is on the left. --- --- @since 0.1.0 -(.+) :: (Index ix, Numeric r e) => Array r ix e -> e -> Array r ix e -(.+) = plusScalar -{-# INLINE (.+) #-} - --- | Add a scalar to each element of the array. Array is on the right. --- --- @since 0.4.0 -(+.) :: (Index ix, Numeric r e) => e -> Array r ix e -> Array r ix e -(+.) = flip plusScalar -{-# INLINE (+.) #-} - --- | Subtract two arrays pointwise. Throws `SizeMismatchException` if arrays sizes do not --- match. --- --- @since 0.4.0 -(.-.) :: - (Load r ix e, Numeric r e, MonadThrow m) => Array r ix e -> Array r ix e -> m (Array r ix e) -(.-.) = liftNumericArray2M subtractionPointwise -{-# INLINE (.-.) #-} - - --- | Subtract a scalar from each element of the array. Array is on the left. --- --- @since 0.1.0 -(.-) :: (Index ix, Numeric r e) => Array r ix e -> e -> Array r ix e -(.-) = minusScalar -{-# INLINE (.-) #-} - --- | Subtract a scalar from each element of the array. Array is on the right. --- --- @since 0.4.0 -(-.) :: (Index ix, Numeric r e) => e -> Array r ix e -> Array r ix e -(-.) = flip minusScalar -{-# INLINE (-.) #-} - - --- | Multiply two arrays together pointwise. --- --- @since 0.4.0 -(.*.) :: - (Load r ix e, Numeric r e, MonadThrow m) => Array r ix e -> Array r ix e -> m (Array r ix e) -(.*.) = liftNumericArray2M multiplicationPointwise -{-# INLINE (.*.) #-} - -(.*) :: (Index ix, Numeric r e) => Array r ix e -> e -> Array r ix e -(.*) = multiplyScalar -{-# INLINE (.*) #-} - -(*.) :: (Index ix, Numeric r e) => e -> Array r ix e -> Array r ix e -(*.) = flip multiplyScalar -{-# INLINE (*.) #-} - -(.^) :: (Index ix, Numeric r e) => Array r ix e -> Int -> Array r ix e -(.^) = powerPointwise -{-# INLINE (.^) #-} - --- | Perform matrix multiplication. Inner dimensions must agree, otherwise `SizeMismatchException`. -(|*|) :: - (Mutable r Ix2 e, Source r' Ix2 e, OuterSlice r Ix2 e, Source (R r) Ix1 e, Num e, MonadThrow m) - => Array r Ix2 e - -> Array r' Ix2 e - -> m (Array r Ix2 e) -(|*|) a1 a2 = compute <$> multArrs a1 a2 -{-# INLINE [1] (|*|) #-} - -{-# RULES -"multDoubleTranspose" [~1] forall arr1 arr2 . arr1 |*| transpose arr2 = - multiplyTransposedFused arr1 (convert arr2) - #-} - -multiplyTransposedFused :: - ( Mutable r Ix2 e - , OuterSlice r Ix2 e - , Source (R r) Ix1 e - , Num e - , MonadThrow m - ) - => Array r Ix2 e - -> Array r Ix2 e - -> m (Array r Ix2 e) -multiplyTransposedFused arr1 arr2 = compute <$> multiplyTransposed arr1 arr2 -{-# INLINE multiplyTransposedFused #-} - - -multArrs :: forall r r' e m. - ( Mutable r Ix2 e - , Source r' Ix2 e - , OuterSlice r Ix2 e - , Source (R r) Ix1 e - , Num e - , MonadThrow m - ) - => Array r Ix2 e -> Array r' Ix2 e -> m (Array D Ix2 e) -multArrs arr1 arr2 = multiplyTransposed arr1 arr2' - where - arr2' :: Array r Ix2 e - arr2' = compute $ transpose arr2 -{-# INLINE multArrs #-} - --- | It is quite often that second matrix gets transposed before multiplication (eg. A * A'), but --- due to layout of data in memory it is more efficient to transpose the second array again. -multiplyTransposed :: - ( Manifest r Ix2 e - , OuterSlice r Ix2 e - , Source (R r) Ix1 e - , Num e - , MonadThrow m - ) - => Array r Ix2 e - -> Array r Ix2 e - -> m (Array D Ix2 e) -multiplyTransposed arr1 arr2 - | n1 /= m2 = throwM $ SizeMismatchException (size arr1) (size arr2) - | otherwise = - pure $ - DArray (getComp arr1 <> getComp arr2) (SafeSz (m1 :. n2)) $ \(i :. j) -> - A.foldlS (+) 0 (A.zipWith (*) (unsafeOuterSlice arr1 i) (unsafeOuterSlice arr2 j)) - where - 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 :: Sz1 -> Array DL Ix2 Int -identityMatrix (Sz n) = makeLoadArrayS (Sz2 n n) 0 $ \ w -> loopM_ 0 (< n) (+1) $ \ i -> w (i :. i) 1 -{-# INLINE identityMatrix #-} - - -negateA :: (Index ix, Numeric r e) => Array r ix e -> Array r ix e -negateA = unsafeLiftArray negate -{-# INLINE negateA #-} - -absA :: (Index ix, Numeric r e) => Array r ix e -> Array r ix e -absA = absPointwise -{-# INLINE absA #-} - -signumA :: (Index ix, Numeric r e) => Array r ix e -> Array r ix e -signumA = unsafeLiftArray signum -{-# INLINE signumA #-} - -fromIntegerA :: (Index ix, Num e) => Integer -> Array D ix e -fromIntegerA = singleton . fromInteger -{-# INLINE fromIntegerA #-} - -(./.) :: - (Load r ix e, NumericFloat r e, MonadThrow m) - => Array r ix e - -> Array r ix e - -> m (Array r ix e) -(./.) = liftNumericArray2M divisionPointwise -{-# INLINE (./.) #-} - -(./) ::(Index ix, NumericFloat r e) => Array r ix e -> e -> Array r ix e -(./) = divideScalar -{-# INLINE (./) #-} - -(.^^) - :: (Index ix, Numeric r e, Fractional e, Integral b) - => Array r ix e -> b -> Array r ix e -(.^^) arr n = unsafeLiftArray (^^ n) arr -{-# INLINE (.^^) #-} - -recipA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e -recipA = recipPointwise -{-# INLINE recipA #-} - - -fromRationalA - :: (Index ix, Fractional e) - => Rational -> Array D ix e -fromRationalA = singleton . fromRational -{-# INLINE fromRationalA #-} - -piA - :: (Index ix, Floating e) - => Array D ix e -piA = singleton pi -{-# INLINE piA #-} - -expA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e -expA = unsafeLiftArray exp -{-# INLINE expA #-} - -sqrtA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e -sqrtA = unsafeLiftArray sqrt -{-# INLINE sqrtA #-} - -logA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e -logA = unsafeLiftArray log -{-# INLINE logA #-} - -logBaseA - :: (Source r1 ix e, Source r2 ix e, Floating e) - => Array r1 ix e -> Array r2 ix e -> Array D ix e -logBaseA = liftArray2Matching logBase -{-# INLINE logBaseA #-} - -(.**) - :: (Source r1 ix e, Source r2 ix e, Floating e) - => Array r1 ix e -> Array r2 ix e -> Array D ix e -(.**) = liftArray2Matching (**) -{-# INLINE (.**) #-} - - - -sinA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e -sinA = unsafeLiftArray sin -{-# INLINE sinA #-} - -cosA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e -cosA = unsafeLiftArray cos -{-# INLINE cosA #-} - -tanA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e -tanA = unsafeLiftArray cos -{-# INLINE tanA #-} - -asinA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e -asinA = unsafeLiftArray asin -{-# INLINE asinA #-} - -atanA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e -atanA = unsafeLiftArray atan -{-# INLINE atanA #-} - -acosA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e -acosA = unsafeLiftArray acos -{-# INLINE acosA #-} - -sinhA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e -sinhA = unsafeLiftArray sinh -{-# INLINE sinhA #-} - -tanhA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e -tanhA = unsafeLiftArray cos -{-# INLINE tanhA #-} - -coshA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e -coshA = unsafeLiftArray cosh -{-# INLINE coshA #-} - -asinhA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e -asinhA = unsafeLiftArray asinh -{-# INLINE asinhA #-} - -acoshA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e -acoshA = unsafeLiftArray acosh -{-# INLINE acoshA #-} - -atanhA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e -atanhA = unsafeLiftArray atanh -{-# INLINE atanhA #-} - - -quotA - :: (Source r1 ix e, Source r2 ix e, Integral e) - => Array r1 ix e -> Array r2 ix e -> Array D ix e -quotA = liftArray2Matching quot -{-# INLINE quotA #-} - - -remA - :: (Source r1 ix e, Source r2 ix e, Integral e) - => Array r1 ix e -> Array r2 ix e -> Array D ix e -remA = liftArray2Matching rem -{-# INLINE remA #-} - -divA - :: (Source r1 ix e, Source r2 ix e, Integral e) - => Array r1 ix e -> Array r2 ix e -> Array D ix e -divA = liftArray2Matching div -{-# INLINE divA #-} - -modA - :: (Source r1 ix e, Source r2 ix e, Integral e) - => Array r1 ix e -> Array r2 ix e -> Array D ix e -modA = liftArray2Matching mod -{-# INLINE modA #-} - - - -quotRemA - :: (Source r1 ix e, Source r2 ix e, Integral e) - => Array r1 ix e -> Array r2 ix e -> (Array D ix e, Array D ix e) -quotRemA arr1 = A.unzip . liftArray2Matching (quotRem) arr1 -{-# INLINE quotRemA #-} - - -divModA - :: (Source r1 ix e, Source r2 ix e, Integral e) - => Array r1 ix e -> Array r2 ix e -> (Array D ix e, Array D ix e) -divModA arr1 = A.unzip . liftArray2Matching (divMod) arr1 -{-# INLINE divModA #-} - - - -truncateA - :: (Index ix, Numeric r e, RealFrac a, Integral e) - => Array r ix a -> Array r ix e -truncateA = unsafeLiftArray truncate -{-# INLINE truncateA #-} - - -roundA :: (Index ix, Numeric r e, RealFrac a, Integral e) => Array r ix a -> Array r ix e -roundA = unsafeLiftArray round -{-# INLINE roundA #-} - - -ceilingA :: (Index ix, Numeric r e, RealFrac a, Integral e) => Array r ix a -> Array r ix e -ceilingA = unsafeLiftArray ceiling -{-# INLINE ceilingA #-} - - -floorA :: (Index ix, Numeric r e, RealFrac a, Integral e) => Array r ix a -> Array r ix e -floorA = unsafeLiftArray floor -{-# INLINE floorA #-} - -atan2A :: - (Load r ix e, Numeric r e, RealFloat e, MonadThrow m) - => Array r ix e - -> Array r ix e - -> m (Array r ix e) -atan2A = liftArray2M atan2 -{-# INLINE atan2A #-} + +-- | +-- Module : Data.Massiv.Array.Numeric +-- Copyright : (c) Alexey Kuleshevich 2018-2022 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +module Data.Massiv.Array.Numeric ( + -- * Numeric + Numeric, + NumericFloat, + liftNumArray2M, + + -- ** Pointwise addition + (.+), + (+.), + (.+.), + (!+!), + sumArraysM, + sumArrays', + + -- ** Pointwise subtraction + (.-), + (-.), + (.-.), + (!-!), + + -- ** Pointwise multiplication + (.*), + (*.), + (.*.), + (!*!), + (.^), + productArraysM, + productArrays', + + -- ** Dot product + (!.!), + dotM, + + -- ** Matrix multiplication + (.><), + (!><), + multiplyMatrixByVector, + (><.), + (><.), + (!> (e -> e -> e) + -> Array r ix e + -> Array r ix e + -> m (Array r ix e) +liftNumArray2M f a1 a2 + | size a1 == size a2 = pure $ unsafeLiftArray2 f a1 a2 + | isZeroSz sz1 && isZeroSz sz2 = pure $ unsafeResize zeroSz a1 + | otherwise = throwM $ SizeMismatchException sz1 sz2 + where + !sz1 = size a1 + !sz2 = size a2 +{-# INLINE liftNumArray2M #-} + +applyExactSize2M + :: (Index ix, Size r, MonadThrow m) + => (Array r ix e -> Array r ix e -> Array r ix e) + -> Array r ix e + -> Array r ix e + -> m (Array r ix e) +applyExactSize2M f a1 a2 + | size a1 == size a2 = pure $! f a1 a2 + | isZeroSz sz1 && isZeroSz sz2 = pure $! unsafeResize zeroSz a1 + | otherwise = throwM $! SizeMismatchException sz1 sz2 + where + !sz1 = size a1 + !sz2 = size a2 +{-# INLINE applyExactSize2M #-} + +-- | Add two arrays together pointwise. Same as `!+!` but produces monadic computation +-- that allows for handling failure. +-- +-- /__Throws Exception__/: `SizeMismatchException` when array sizes do not match. +-- +-- @since 0.4.0 +(.+.) :: (Index ix, Numeric r e, MonadThrow m) => Array r ix e -> Array r ix e -> m (Array r ix e) +(.+.) = applyExactSize2M additionPointwise +{-# INLINE (.+.) #-} + +-- | Add two arrays together pointwise. Prefer to use monadic version of this function +-- `.+.` whenever possible, because it is better to avoid partial functions. +-- +-- [Partial] Mismatched array sizes will result in an impure exception being thrown. +-- +-- ====__Example__ +-- +-- >>> let a1 = Ix1 0 ... 10 +-- >>> let a2 = Ix1 20 ... 30 +-- >>> a1 !+! a2 +-- Array D Seq (Sz1 11) +-- [ 20, 22, 24, 26, 28, 30, 32, 34, 36, 38, 40 ] +-- +-- @since 0.5.6 +(!+!) :: (HasCallStack, Index ix, Numeric r e) => Array r ix e -> Array r ix e -> Array r ix e +(!+!) a1 a2 = throwEither (a1 .+. a2) +{-# INLINE (!+!) #-} + +-- | Add a scalar to each element of the array. Array is on the left. +-- +-- @since 0.1.0 +(.+) :: (Index ix, Numeric r e) => Array r ix e -> e -> Array r ix e +(.+) = plusScalar +{-# INLINE (.+) #-} + +-- | Add a scalar to each element of the array. Array is on the right. +-- +-- @since 0.4.0 +(+.) :: (Index ix, Numeric r e) => e -> Array r ix e -> Array r ix e +(+.) = flip plusScalar +{-# INLINE (+.) #-} + +-- | Subtract two arrays pointwise. Same as `!-!` but produces monadic computation that +-- allows for handling failure. +-- +-- /__Throws Exception__/: `SizeMismatchException` when array sizes do not match. +-- +-- @since 0.4.0 +(.-.) + :: (Index ix, Numeric r e, MonadThrow m) => Array r ix e -> Array r ix e -> m (Array r ix e) +(.-.) = applyExactSize2M subtractionPointwise +{-# INLINE (.-.) #-} + +-- | Subtract one array from another pointwise. Prefer to use monadic version of this +-- function `.-.` whenever possible, because it is better to avoid partial functions. +-- +-- [Partial] Mismatched array sizes will result in an impure exception being thrown. +-- +-- ====__Example__ +-- +-- >>> let a1 = Ix1 0 ... 10 +-- >>> let a2 = Ix1 20 ... 30 +-- >>> a1 !-! a2 +-- Array D Seq (Sz1 11) +-- [ -20, -20, -20, -20, -20, -20, -20, -20, -20, -20, -20 ] +-- +-- @since 0.5.6 +(!-!) :: (Index ix, Numeric r e) => Array r ix e -> Array r ix e -> Array r ix e +(!-!) a1 a2 = throwEither (a1 .-. a2) +{-# INLINE (!-!) #-} + +-- | Subtract a scalar from each element of the array. Array is on the left. +-- +-- @since 0.4.0 +(.-) :: (Index ix, Numeric r e) => Array r ix e -> e -> Array r ix e +(.-) = minusScalar +{-# INLINE (.-) #-} + +-- | Subtract each element of the array from a scalar. Array is on the right. +-- +-- @since 0.5.6 +(-.) :: (Index ix, Numeric r e) => e -> Array r ix e -> Array r ix e +(-.) = scalarMinus +{-# INLINE (-.) #-} + +-- | Multiply two arrays together pointwise. Same as `!*!` but produces monadic +-- computation that allows for handling failure. +-- +-- /__Throws Exception__/: `SizeMismatchException` when array sizes do not match. +-- +-- @since 0.4.0 +(.*.) + :: (Index ix, Numeric r e, MonadThrow m) => Array r ix e -> Array r ix e -> m (Array r ix e) +(.*.) = applyExactSize2M multiplicationPointwise +{-# INLINE (.*.) #-} + +-- | Multiplication of two arrays pointwise, +-- i.e. . +-- Prefer to use monadic version of this function `.*.` whenever possible, +-- because it is better to avoid partial functions. +-- +-- [Partial] Mismatched array sizes will result in an impure exception being thrown. +-- +-- ====__Example__ +-- +-- >>> let a1 = Ix1 0 ... 10 +-- >>> let a2 = Ix1 20 ... 30 +-- >>> a1 !*! a2 +-- Array D Seq (Sz1 11) +-- [ 0, 21, 44, 69, 96, 125, 156, 189, 224, 261, 300 ] +-- +-- @since 0.5.6 +(!*!) :: (Index ix, Numeric r e) => Array r ix e -> Array r ix e -> Array r ix e +(!*!) a1 a2 = throwEither (a1 .*. a2) +{-# INLINE (!*!) #-} + +-- | Multiply each element of the array by a scalar value. Scalar is on the right. +-- +-- ====__Example__ +-- +-- >>> let arr = Ix1 20 ..: 25 +-- >>> arr +-- Array D Seq (Sz1 5) +-- [ 20, 21, 22, 23, 24 ] +-- >>> arr .* 10 +-- Array D Seq (Sz1 5) +-- [ 200, 210, 220, 230, 240 ] +-- +-- @since 0.4.0 +(.*) :: (Index ix, Numeric r e) => Array r ix e -> e -> Array r ix e +(.*) = multiplyScalar +{-# INLINE (.*) #-} + +-- | Multiply each element of the array by a scalar value. Scalar is on the left. +-- +-- ====__Example__ +-- +-- >>> let arr = Ix1 20 ..: 25 +-- >>> arr +-- Array D Seq (Sz1 5) +-- [ 20, 21, 22, 23, 24 ] +-- >>> 10 *. arr +-- Array D Seq (Sz1 5) +-- [ 200, 210, 220, 230, 240 ] +-- +-- @since 0.4.0 +(*.) :: (Index ix, Numeric r e) => e -> Array r ix e -> Array r ix e +(*.) = flip multiplyScalar +{-# INLINE (*.) #-} + +-- | Raise each element of the array to a power. +-- +-- ====__Example__ +-- +-- >>> let arr = Ix1 20 ..: 25 +-- >>> arr +-- Array D Seq (Sz1 5) +-- [ 20, 21, 22, 23, 24 ] +-- >>> arr .^ 3 +-- Array D Seq (Sz1 5) +-- [ 8000, 9261, 10648, 12167, 13824 ] +-- +-- @since 0.4.0 +(.^) :: (Index ix, Numeric r e) => Array r ix e -> Int -> Array r ix e +(.^) = powerPointwise +{-# INLINE (.^) #-} + +-- | Dot product of two vectors. +-- +-- [Partial] Throws an impure exception when lengths of vectors do not match +-- +-- @since 0.5.6 +(!.!) :: (Numeric r e, Source r e) => Vector r e -> Vector r e -> e +(!.!) v1 v2 = throwEither $ dotM v1 v2 +{-# INLINE (!.!) #-} + +-- | Dot product of two vectors. +-- +-- /__Throws Exception__/: `SizeMismatchException` when lengths of vectors do not match +-- +-- @since 0.5.6 +dotM :: (FoldNumeric r e, Source r e, MonadThrow m) => Vector r e -> Vector r e -> m e +dotM v1 v2 + | size v1 /= size v2 = throwM $ SizeMismatchException (size v1) (size v2) + | comp == Seq = pure $! unsafeDotProduct v1 v2 + | otherwise = pure $! unsafePerformIO $ unsafeDotProductIO v1 v2 + where + comp = getComp v1 <> getComp v2 +{-# INLINE dotM #-} + +unsafeDotProductIO + :: (MonadUnliftIO m, Index ix, FoldNumeric r b, Source r b) + => Array r ix b + -> Array r ix b + -> m b +unsafeDotProductIO v1 v2 = do + results <- + withScheduler comp $ \scheduler -> + splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> liftIO $ do + let n = SafeSz chunkLength + loopA_ 0 (< slackStart) (+ chunkLength) $ \ !start -> + scheduleWork scheduler $ + pure $! + unsafeDotProduct (unsafeLinearSlice start n v1) (unsafeLinearSlice start n v2) + when (slackStart < totalLength) $ do + let k = SafeSz (totalLength - slackStart) + scheduleWork scheduler $ + pure $! + unsafeDotProduct (unsafeLinearSlice slackStart k v1) (unsafeLinearSlice slackStart k v2) + pure $! F.foldl' (+) 0 results + where + totalLength = totalElem (size v1) + comp = getComp v1 <> getComp v2 +{-# INLINE unsafeDotProductIO #-} + +-- | Compute L2 norm of an array. +-- +-- @since 0.5.6 +normL2 :: (FoldNumeric r e, Source r e, Index ix, Floating e) => Array r ix e -> e +normL2 v + | getComp v == Seq = sqrt $! powerSumArray v 2 + | otherwise = sqrt $! unsafePerformIO $ powerSumArrayIO v 2 +{-# INLINE normL2 #-} + +powerSumArrayIO + :: (MonadUnliftIO m, Index ix, FoldNumeric r b, Source r b) + => Array r ix b + -> Int + -> m b +powerSumArrayIO v p = do + results <- + withScheduler (getComp v) $ \scheduler -> + splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> liftIO $ do + let n = SafeSz chunkLength + loopA_ 0 (< slackStart) (+ chunkLength) $ \ !start -> + scheduleWork scheduler $ pure $! powerSumArray (unsafeLinearSlice start n v) p + when (slackStart < totalLength) $ do + let k = SafeSz (totalLength - slackStart) + scheduleWork scheduler $ pure $! powerSumArray (unsafeLinearSlice slackStart k v) p + pure $! F.foldl' (+) 0 results + where + totalLength = totalElem (size v) +{-# INLINE powerSumArrayIO #-} + +-- | Multiply a matrix by a column vector. Same as `!><` but produces monadic +-- computation that allows for handling failure. +-- +-- /__Throws Exception__/: `SizeMismatchException` when inner dimensions of arrays do not match. +-- +-- @since 0.5.6 +(.><) + :: (MonadThrow m, FoldNumeric r e, Source r e) + => Matrix r e + -- ^ Matrix + -> Vector r e + -- ^ Column vector (Used many times, so make sure it is computed) + -> m (Vector D e) +(.><) mm v + | mCols /= n = throwM $ SizeMismatchException (size mm) (Sz2 n 1) + | mRows == 0 || mCols == 0 = pure $ setComp comp empty + | otherwise = pure $ makeArray comp (Sz1 mRows) $ \i -> + unsafeDotProduct (unsafeLinearSlice (i * n) sz mm) v + where + comp = getComp mm <> getComp v + Sz2 mRows mCols = size mm + sz@(Sz1 n) = size v +{-# INLINE (.><) #-} + +-- | Multiply matrix by a column vector. Same as `.><` but returns computed version of a vector +-- +-- /__Throws Exception__/: `SizeMismatchException` when inner dimensions of arrays do not match. +-- +-- @since 0.5.7 +multiplyMatrixByVector + :: (MonadThrow m, Numeric r e, Manifest r e) + => Matrix r e + -- ^ Matrix + -> Vector r e + -- ^ Column vector (Used many times, so make sure it is computed) + -> m (Vector r e) +multiplyMatrixByVector mm v = compute <$> mm .>< v +{-# INLINE multiplyMatrixByVector #-} + +-- | Multiply a matrix by a column vector +-- +-- [Partial] Throws impure exception when inner dimensions do not agree +-- +-- @since 0.5.6 +(!><) + :: (Numeric r e, Source r e) + => Matrix r e + -- ^ Matrix + -> Vector r e + -- ^ Column vector (Used many times, so make sure it is computed) + -> Vector D e +(!><) mm v = throwEither (mm .>< v) +{-# INLINE (!><) #-} + +-- | Multiply a row vector by a matrix. Same as `><.) + :: (MonadThrow m, Numeric r e, Manifest r e) + => Vector r e + -- ^ Row vector + -> Matrix r e + -- ^ Matrix + -> m (Vector r e) +(><.) = multiplyVectorByMatrix +{-# INLINE (><.) #-} + +-- | Multiply a row vector by a matrix. Same as `><.` but returns computed vector instead of +-- a delayed one. +-- +-- /__Throws Exception__/: `SizeMismatchException` when inner dimensions of arrays do not match. +-- +-- @since 0.5.7 +multiplyVectorByMatrix + :: (MonadThrow m, Numeric r e, Manifest r e) + => Vector r e + -- ^ Row vector + -> Matrix r e + -- ^ Matrix + -> m (Vector r e) +multiplyVectorByMatrix v mm + | mRows /= n = throwM $ SizeMismatchException (Sz2 1 n) (size mm) + | mRows == 0 || mCols == 0 = pure $ runST (unsafeFreeze comp =<< unsafeNew zeroSz) + | otherwise = + pure $! + unsafePerformIO $ do + mv <- newMArray (Sz mCols) 0 + withMassivScheduler_ comp $ \scheduler -> do + let loopCols x ivto = + fix $ \go im iv -> + when (iv < ivto) $ do + _ <- unsafeLinearModify mv (\a -> pure $ a + unsafeLinearIndex mm im * x) iv + go (im + 1) (iv + 1) + loopRows i0 from to = + flip fix i0 $ \go i -> + when (i < mRows) $ do + loopCols (unsafeLinearIndex v i) to (i * mCols + from) from + go (i + 1) + splitLinearlyM_ scheduler mCols (loopRows 0) + unsafeFreeze comp mv + where + comp = getComp mm <> getComp v + Sz2 mRows mCols = size mm + Sz1 n = size v +{-# INLINE multiplyVectorByMatrix #-} + +-- | Multiply a row vector by a matrix. +-- +-- [Partial] Throws impure exception when inner dimensions do not agree +-- +-- @since 0.5.6 +(> Vector r e + -- ^ Row vector (Used many times, so make sure it is computed) + -> Matrix r e + -- ^ Matrix + -> Vector r e +(><. mm) +{-# INLINE (>>> import Data.Massiv.Array +-- >>> a1 = makeArrayR P Seq (Sz2 5 6) $ \(i :. j) -> i + j +-- >>> a2 = makeArrayR P Seq (Sz2 6 5) $ \(i :. j) -> i - j +-- >>> a1 !> Matrix r e -> Matrix r e -> Matrix r e +(!><.) :: (Numeric r e, Manifest r e, MonadThrow m) => Matrix r e -> Matrix r e -> m (Matrix r e) +(.><.) = multiplyMatrices +{-# INLINE (.><.) #-} + +-- | Synonym for `.><.` +-- +-- @since 0.5.6 +multiplyMatrices + :: (Numeric r e, Manifest r e, MonadThrow m) => Matrix r e -> Matrix r e -> m (Matrix r e) +multiplyMatrices arrA arrB + -- mA == 1 = -- TODO: call multiplyVectorByMatrix + -- nA == 1 = -- TODO: call multiplyMatrixByVector + | nA /= mB = throwM $ SizeMismatchException (size arrA) (size arrB) + | isEmpty arrA || isEmpty arrB = pure $ runST (unsafeFreeze comp =<< unsafeNew zeroSz) + | otherwise = pure $! unsafePerformIO $ do + marrC <- newMArray (SafeSz (mA :. nB)) 0 + withScheduler_ comp $ \scheduler -> do + let withC00 iA jB f = + let !ixC00 = iA * nB + jB + in f ixC00 =<< unsafeLinearRead marrC ixC00 + withC01 ixC00 f = + let !ixC01 = ixC00 + 1 + in f ixC01 =<< unsafeLinearRead marrC ixC01 + withC10 ixC00 f = + let !ixC10 = ixC00 + nB + in f ixC10 =<< unsafeLinearRead marrC ixC10 + withC11 ixC01 f = + let !ixC11 = ixC01 + nB + in f ixC11 =<< unsafeLinearRead marrC ixC11 + withB00 iB jB f = + let !ixB00 = iB * nB + jB + in f ixB00 $! unsafeLinearIndex arrB ixB00 + withB00B10 iB jB f = + withB00 iB jB $ \ixB00 b00 -> + let !ixB10 = ixB00 + nB + in f ixB00 b00 ixB10 $! unsafeLinearIndex arrB ixB10 + withA00 iA jA f = + let !ixA00 = iA * nA + jA + in f ixA00 $! unsafeLinearIndex arrA ixA00 + withA00A10 iA jA f = + withA00 iA jA $ \ixA00 a00 -> + let !ixA10 = ixA00 + nA + in f ixA00 a00 ixA10 $! unsafeLinearIndex arrA ixA10 + let loopColsB_UnRowBColA_UnRowA a00 a01 a10 a11 iA iB jB + | jB < n2B = do + withB00B10 iB jB $ \ixB00 b00 ixB10 b10 -> do + let !b01 = unsafeLinearIndex arrB (ixB00 + 1) + !b11 = unsafeLinearIndex arrB (ixB10 + 1) + withC00 iA jB $ \ixC00 c00 -> do + unsafeLinearWrite marrC ixC00 (c00 + a00 * b00 + a01 * b10) + withC01 ixC00 $ \ixC01 c01 -> do + unsafeLinearWrite marrC ixC01 (c01 + a00 * b01 + a01 * b11) + withC10 ixC00 $ \ixC10 c10 -> + unsafeLinearWrite marrC ixC10 (c10 + a10 * b00 + a11 * b10) + withC11 ixC01 $ \ixC11 c11 -> + unsafeLinearWrite marrC ixC11 (c11 + a10 * b01 + a11 * b11) + loopColsB_UnRowBColA_UnRowA a00 a01 a10 a11 iA iB (jB + 2) + | jB < nB = withB00B10 iB jB $ \_ b00 _ b10 -> + withC00 iA jB $ \ixC00 c00 -> do + unsafeLinearWrite marrC ixC00 (c00 + a00 * b00 + a01 * b10) + withC10 ixC00 $ \ixC10 c10 -> + unsafeLinearWrite marrC ixC10 (c10 + a10 * b00 + a11 * b10) + | otherwise = pure () + + loopColsB_UnRowBColA_RowA a00 a01 iA iB jB + | jB < n2B = do + withB00B10 iB jB $ \ixB00 b00 ixB10 b10 -> do + let !b01 = unsafeLinearIndex arrB (ixB00 + 1) + !b11 = unsafeLinearIndex arrB (ixB10 + 1) + withC00 iA jB $ \ixC00 c00 -> do + unsafeLinearWrite marrC ixC00 (c00 + a00 * b00 + a01 * b10) + withC01 ixC00 $ \ixC01 c01 -> + unsafeLinearWrite marrC ixC01 (c01 + a00 * b01 + a01 * b11) + loopColsB_UnRowBColA_RowA a00 a01 iA iB (jB + 2) + | jB < nB = withB00B10 iB jB $ \_ b00 _ b10 -> + withC00 iA jB $ \ixC00 c00 -> + unsafeLinearWrite marrC ixC00 (c00 + a00 * b00 + a01 * b10) + | otherwise = pure () + + loopColsB_RowBColA_UnRowA a00 a10 iA iB jB + | jB < n2B = do + withB00 iB jB $ \ixB00 b00 -> do + let !b01 = unsafeLinearIndex arrB (ixB00 + 1) + withC00 iA jB $ \ixC00 c00 -> do + unsafeLinearWrite marrC ixC00 (c00 + a00 * b00) + withC01 ixC00 $ \ixC01 c01 -> do + unsafeLinearWrite marrC ixC01 (c01 + a00 * b01) + withC10 ixC00 $ \ixC10 c10 -> + unsafeLinearWrite marrC ixC10 (c10 + a10 * b00) + withC11 ixC01 $ \ixC11 c11 -> + unsafeLinearWrite marrC ixC11 (c11 + a10 * b01) + loopColsB_RowBColA_UnRowA a00 a10 iA iB (jB + 2) + | jB < nB = withB00 iB jB $ \_ b00 -> + withC00 iA jB $ \ixC00 c00 -> do + unsafeLinearWrite marrC ixC00 (c00 + a00 * b00) + withC10 ixC00 $ \ixC10 c10 -> + unsafeLinearWrite marrC ixC10 (c10 + a10 * b00) + | otherwise = pure () + + loopColsB_RowBColA_RowA a00 iA iB jB + | jB < n2B = do + withB00 iB jB $ \ixB00 b00 -> do + let !b01 = unsafeLinearIndex arrB (ixB00 + 1) + withC00 iA jB $ \ixC00 c00 -> do + unsafeLinearWrite marrC ixC00 (c00 + a00 * b00) + withC01 ixC00 $ \ixC01 c01 -> do + unsafeLinearWrite marrC ixC01 (c01 + a00 * b01) + loopColsB_RowBColA_RowA a00 iA iB (jB + 2) + | jB < nB = withB00 iB jB $ \_ b00 -> + withC00 iA jB $ \ixC00 c00 -> + unsafeLinearWrite marrC ixC00 (c00 + a00 * b00) + | otherwise = pure () + + loopRowsB_UnRowA iA iB + | iB < m2B = do + withA00A10 iA iB $ \ixA00 a00 ixA10 a10 -> do + let !a01 = unsafeLinearIndex arrA (ixA00 + 1) + !a11 = unsafeLinearIndex arrA (ixA10 + 1) + loopColsB_UnRowBColA_UnRowA a00 a01 a10 a11 iA iB 0 + loopRowsB_UnRowA iA (iB + 2) + | iB < mB = + withA00A10 iA iB $ \_ a00 _ a10 -> loopColsB_RowBColA_UnRowA a00 a10 iA iB 0 + | otherwise = pure () + + loopRowsB_RowA iA iB + | iB < m2B = do + withA00 iA iB $ \ixA00 a00 -> do + let !a01 = unsafeLinearIndex arrA (ixA00 + 1) + loopColsB_UnRowBColA_RowA a00 a01 iA iB 0 + loopRowsB_RowA iA (iB + 2) + | iB < mB = withA00 iA iB $ \_ a00 -> loopColsB_RowBColA_RowA a00 iA iB 0 + | otherwise = pure () + + loopRowsA iA + | iA < m2A = do + scheduleWork_ scheduler $ loopRowsB_UnRowA iA 0 + loopRowsA (iA + 2) + | iA < mA = scheduleWork_ scheduler $ loopRowsB_RowA iA 0 + | otherwise = pure () + loopRowsA 0 + + unsafeFreeze comp marrC + where + comp = getComp arrA <> getComp arrB + m2A = mA - mA `rem` 2 + m2B = mB - mB `rem` 2 + n2B = nB - nB `rem` 2 + Sz (mA :. nA) = size arrA + Sz (mB :. nB) = size arrB +{-# INLINEABLE multiplyMatrices #-} + +-- | Computes the matrix-matrix multiplication where second matrix is transposed (i.e. M +-- x N') +-- +-- > m1 .><. transpose m2 == multiplyMatricesTransposed m1 m2 +-- +-- @since 0.5.6 +multiplyMatricesTransposed + :: (Numeric r e, Manifest r e, MonadThrow m) + => Matrix r e + -> Matrix r e + -> m (Matrix D e) +multiplyMatricesTransposed arr1 arr2 + | n1 /= m2 = throwM $ SizeMismatchException (size arr1) (Sz2 m2 n2) + | isEmpty arr1 || isEmpty arr2 = pure $ setComp comp empty + | otherwise = + pure $ + makeArray comp (SafeSz (m1 :. n2)) $ \(i :. j) -> + unsafeDotProduct (unsafeLinearSlice (i * n1) n arr1) (unsafeLinearSlice (j * n1) n arr2) + where + comp = getComp arr1 <> getComp arr2 + n = SafeSz n1 + SafeSz (m1 :. n1) = size arr1 + SafeSz (n2 :. m2) = size arr2 +{-# INLINE multiplyMatricesTransposed #-} + +-- | 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 :: Num e => Sz1 -> Matrix DL e +identityMatrix (Sz n) = + makeLoadArrayS (Sz2 n n) 0 $ \w -> loopA_ 0 (< n) (+ 1) $ \i -> w (i :. i) 1 +{-# INLINE identityMatrix #-} + +-- | Create a lower triangular (L in LU decomposition) matrix of size @NxN@ +-- +-- ==== __Example__ +-- +-- >>> import Data.Massiv.Array +-- >>> lowerTriangular Seq 5 (\(i :. j) -> i + j) +-- Array DL Seq (Sz (5 :. 5)) +-- [ [ 0, 0, 0, 0, 0 ] +-- , [ 1, 2, 0, 0, 0 ] +-- , [ 2, 3, 4, 0, 0 ] +-- , [ 3, 4, 5, 6, 0 ] +-- , [ 4, 5, 6, 7, 8 ] +-- ] +-- +-- @since 0.5.2 +lowerTriangular :: forall e. Num e => Comp -> Sz1 -> (Ix2 -> e) -> Matrix DL e +lowerTriangular comp (Sz1 n) f = DLArray comp (SafeSz (n :. n)) load + where + load :: Loader e + load scheduler startAt uWrite uSet = do + forM_ (0 ..: n) $ \i -> do + let !k = startAt + i * n + scheduleWork_ scheduler $ do + forM_ (0 ... i) $ \j -> uWrite (k + j) (f (i :. j)) + uSet (k + i + 1) (Sz (n - i - 1)) 0 +{-# INLINE lowerTriangular #-} + +-- | Create an upper triangular (U in LU decomposition) matrix of size @NxN@ +-- +-- ==== __Example__ +-- +-- >>> import Data.Massiv.Array +-- >>> upperTriangular Par 5 (\(i :. j) -> i + j) +-- Array DL Par (Sz (5 :. 5)) +-- [ [ 0, 1, 2, 3, 4 ] +-- , [ 0, 2, 3, 4, 5 ] +-- , [ 0, 0, 4, 5, 6 ] +-- , [ 0, 0, 0, 6, 7 ] +-- , [ 0, 0, 0, 0, 8 ] +-- ] +-- +-- @since 0.5.2 +upperTriangular :: forall e. Num e => Comp -> Sz1 -> (Ix2 -> e) -> Matrix DL e +upperTriangular comp (Sz1 n) f = DLArray comp (SafeSz (n :. n)) load + where + load :: Loader e + load scheduler startAt uWrite uSet = do + forM_ (0 ..: n) $ \i -> do + let !k = startAt + i * n + scheduleWork_ scheduler $ do + uSet k (SafeSz i) 0 + forM_ (i ..: n) $ \j -> uWrite (k + j) (f (i :. j)) +{-# INLINE upperTriangular #-} + +-- | Negate each element of the array +-- +-- @since 0.4.0 +negateA :: (Index ix, Numeric r e) => Array r ix e -> Array r ix e +negateA = unsafeLiftArray negate +{-# INLINE negateA #-} + +-- | Apply `abs` to each element of the array +-- +-- @since 0.4.0 +absA :: (Index ix, Numeric r e) => Array r ix e -> Array r ix e +absA = absPointwise +{-# INLINE absA #-} + +-- | Apply `signum` to each element of the array +-- +-- @since 0.4.0 +signumA :: (Index ix, Numeric r e) => Array r ix e -> Array r ix e +signumA = unsafeLiftArray signum +{-# INLINE signumA #-} + +-- | Divide each element of one array by another pointwise. Same as `!/!` but produces +-- monadic computation that allows for handling failure. +-- +-- /__Throws Exception__/: `SizeMismatchException` when array sizes do not match. +-- +-- @since 0.4.0 +(./.) + :: (Index ix, NumericFloat r e, MonadThrow m) + => Array r ix e + -> Array r ix e + -> m (Array r ix e) +(./.) = applyExactSize2M divisionPointwise +{-# INLINE (./.) #-} + +-- | Divide two arrays pointwise. Prefer to use monadic version of this function `./.` +-- whenever possible, because it is better to avoid partial functions. +-- +-- [Partial] Mismatched array sizes will result in an impure exception being thrown. +-- +-- ====__Example__ +-- +-- >>> let arr1 = fromIntegral <$> (Ix1 20 ..: 25) :: Array D Ix1 Float +-- >>> let arr2 = fromIntegral <$> (Ix1 100 ..: 105) :: Array D Ix1 Float +-- >>> arr1 !/! arr2 +-- Array D Seq (Sz1 5) +-- [ 0.2, 0.20792079, 0.21568628, 0.22330096, 0.23076923 ] +-- +-- @since 0.5.6 +(!/!) :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e -> Array r ix e +(!/!) a1 a2 = throwEither (a1 ./. a2) +{-# INLINE (!/!) #-} + +-- | Divide a scalar value by each element of the array. +-- +-- > e /. arr == e *. recipA arr +-- +-- ====__Example__ +-- +-- >>> let arr = fromIntegral <$> (Ix1 20 ..: 25) :: Array D Ix1 Float +-- >>> arr +-- Array D Seq (Sz1 5) +-- [ 20.0, 21.0, 22.0, 23.0, 24.0 ] +-- >>> 100 /. arr +-- Array D Seq (Sz1 5) +-- [ 5.0, 4.7619047, 4.5454545, 4.347826, 4.1666665 ] +-- +-- @since 0.5.6 +(/.) :: (Index ix, NumericFloat r e) => e -> Array r ix e -> Array r ix e +(/.) = scalarDivide +{-# INLINE (/.) #-} + +-- | Divide each element of the array by a scalar value. +-- +-- ====__Example__ +-- +-- >>> let arr = fromIntegral <$> (Ix1 20 ..: 25) :: Array D Ix1 Float +-- >>> arr +-- Array D Seq (Sz1 5) +-- [ 20.0, 21.0, 22.0, 23.0, 24.0 ] +-- >>> arr ./ 100 +-- Array D Seq (Sz1 5) +-- [ 0.2, 0.21, 0.22, 0.23, 0.24 ] +-- +-- @since 0.4.0 +(./) :: (Index ix, NumericFloat r e) => Array r ix e -> e -> Array r ix e +(./) = divideScalar +{-# INLINE (./) #-} + +(.^^) + :: (Index ix, Numeric r e, Fractional e, Integral b) + => Array r ix e + -> b + -> Array r ix e +(.^^) arr n = unsafeLiftArray (^^ n) arr +{-# INLINE (.^^) #-} + +-- | Apply reciprocal to each element of the array. +-- +-- > recipA arr == 1 /. arr +-- +-- @since 0.4.0 +recipA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e +recipA = recipPointwise +{-# INLINE recipA #-} + +-- | Apply exponent to each element of the array. +-- +-- > expA arr == map exp arr +-- +-- @since 0.4.0 +expA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e +expA = unsafeLiftArray exp +{-# INLINE expA #-} + +-- | Apply square root to each element of the array. +-- +-- > sqrtA arr == map sqrt arr +-- +-- @since 0.4.0 +sqrtA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e +sqrtA = unsafeLiftArray sqrt +{-# INLINE sqrtA #-} + +-- | Apply logarithm to each element of the array. +-- +-- > logA arr == map log arr +-- +-- @since 0.4.0 +logA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e +logA = unsafeLiftArray log +{-# INLINE logA #-} + +-- | Apply logarithm to each element of the array where the base is in the same cell in +-- the second array. +-- +-- > logBaseA arr1 arr2 == zipWith logBase arr1 arr2 +-- +-- [Partial] Throws an error when arrays do not have matching sizes +-- +-- @since 0.4.0 +logBaseA + :: (Index ix, Source r1 e, Source r2 e, Floating e) + => Array r1 ix e + -> Array r2 ix e + -> Array D ix e +logBaseA = liftArray2' logBase +{-# INLINE logBaseA #-} + +-- TODO: siwtch to +-- (breaking) logBaseA :: Array r ix e -> e -> Array D ix e +-- logBasesM :: Array r ix e -> Array r ix e -> m (Array D ix e) + +-- | Apply power to each element of the array where the power value is in the same cell +-- in the second array. +-- +-- > arr1 .** arr2 == zipWith (**) arr1 arr2 +-- +-- [Partial] Throws an error when arrays do not have matching sizes +-- +-- @since 0.4.0 +(.**) + :: (Index ix, Source r1 e, Source r2 e, Floating e) + => Array r1 ix e + -> Array r2 ix e + -> Array D ix e +(.**) = liftArray2' (**) +{-# INLINE (.**) #-} + +-- TODO: +-- !**! :: Array r1 ix e -> Array r2 ix e -> Array D ix e +-- .**. :: Array r1 ix e -> Array r2 ix e -> m (Array D ix e) +-- (breaking) .** :: Array r1 ix e -> e -> Array D ix e + +-- | Apply sine function to each element of the array. +-- +-- > sinA arr == map sin arr +-- +-- @since 0.4.0 +sinA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e +sinA = unsafeLiftArray sin +{-# INLINE sinA #-} + +-- | Apply cosine function to each element of the array. +-- +-- > cosA arr == map cos arr +-- +-- @since 0.4.0 +cosA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e +cosA = unsafeLiftArray cos +{-# INLINE cosA #-} + +-- | Apply tangent function to each element of the array. +-- +-- > tanA arr == map tan arr +-- +-- @since 0.4.0 +tanA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e +tanA = unsafeLiftArray tan +{-# INLINE tanA #-} + +-- | Apply arcsine function to each element of the array. +-- +-- > asinA arr == map asin arr +-- +-- @since 0.4.0 +asinA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e +asinA = unsafeLiftArray asin +{-# INLINE asinA #-} + +-- | Apply arctangent function to each element of the array. +-- +-- > atanA arr == map atan arr +-- +-- @since 0.4.0 +atanA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e +atanA = unsafeLiftArray atan +{-# INLINE atanA #-} + +-- | Apply arccosine function to each element of the array. +-- +-- > acosA arr == map acos arr +-- +-- @since 0.4.0 +acosA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e +acosA = unsafeLiftArray acos +{-# INLINE acosA #-} + +-- | Apply hyperbolic sine function to each element of the array. +-- +-- > sinhA arr == map sinh arr +-- +-- @since 0.4.0 +sinhA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e +sinhA = unsafeLiftArray sinh +{-# INLINE sinhA #-} + +-- | Apply hyperbolic tangent function to each element of the array. +-- +-- > tanhA arr == map tanh arr +-- +-- @since 0.4.0 +tanhA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e +tanhA = unsafeLiftArray tanh +{-# INLINE tanhA #-} + +-- | Apply hyperbolic cosine function to each element of the array. +-- +-- > coshA arr == map cosh arr +-- +-- @since 0.4.0 +coshA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e +coshA = unsafeLiftArray cosh +{-# INLINE coshA #-} + +-- | Apply inverse hyperbolic sine function to each element of the array. +-- +-- > asinhA arr == map asinh arr +-- +-- @since 0.4.0 +asinhA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e +asinhA = unsafeLiftArray asinh +{-# INLINE asinhA #-} + +-- | Apply inverse hyperbolic cosine function to each element of the array. +-- +-- > acoshA arr == map acosh arr +-- +-- @since 0.4.0 +acoshA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e +acoshA = unsafeLiftArray acosh +{-# INLINE acoshA #-} + +-- | Apply inverse hyperbolic tangent function to each element of the array. +-- +-- > atanhA arr == map atanh arr +-- +-- @since 0.4.0 +atanhA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e +atanhA = unsafeLiftArray atanh +{-# INLINE atanhA #-} + +-- | Perform a pointwise quotient where first array contains numerators and the second +-- one denominators +-- +-- > quotA arr1 arr2 == zipWith quot arr1 arr2 +-- +-- [Partial] Mismatched array sizes will result in an impure exception being thrown. +-- +-- @since 0.1.0 +quotA + :: (HasCallStack, Index ix, Source r1 e, Source r2 e, Integral e) + => Array r1 ix e + -> Array r2 ix e + -> Array D ix e +quotA = liftArray2' quot +{-# INLINE quotA #-} + +-- | Perform a pointwise remainder computation +-- +-- > remA arr1 arr2 == zipWith rem arr1 arr2 +-- +-- [Partial] Mismatched array sizes will result in an impure exception being thrown. +-- +-- @since 0.1.0 +remA + :: (HasCallStack, Index ix, Source r1 e, Source r2 e, Integral e) + => Array r1 ix e + -> Array r2 ix e + -> Array D ix e +remA = liftArray2' rem +{-# INLINE remA #-} + +-- | Perform a pointwise integer division where first array contains numerators and the +-- second one denominators +-- +-- > divA arr1 arr2 == zipWith div arr1 arr2 +-- +-- [Partial] Mismatched array sizes will result in an impure exception being thrown. +-- +-- @since 0.1.0 +divA + :: (HasCallStack, Index ix, Source r1 e, Source r2 e, Integral e) + => Array r1 ix e + -> Array r2 ix e + -> Array D ix e +divA = liftArray2' div +{-# INLINE divA #-} + +-- TODO: +-- * Array r ix e -> Array r ix e -> m (Array r ix e) +-- * Array r ix e -> e -> Array r ix e +-- * e -> Array r ix e -> Array r ix e + +-- | Perform a pointwise modulo computation +-- +-- > modA arr1 arr2 == zipWith mod arr1 arr2 +-- +-- [Partial] Mismatched array sizes will result in an impure exception being thrown. +-- +-- @since 0.1.0 +modA + :: (HasCallStack, Index ix, Source r1 e, Source r2 e, Integral e) + => Array r1 ix e + -> Array r2 ix e + -> Array D ix e +modA = liftArray2' mod +{-# INLINE modA #-} + +-- | Perform a pointwise quotient with remainder where first array contains numerators +-- and the second one denominators +-- +-- > quotRemA arr1 arr2 == zipWith quotRem arr1 arr2 +-- +-- [Partial] Mismatched array sizes will result in an impure exception being thrown. +-- +-- @since 0.1.0 +quotRemA + :: (HasCallStack, Index ix, Source r1 e, Source r2 e, Integral e) + => Array r1 ix e + -> Array r2 ix e + -> (Array D ix e, Array D ix e) +quotRemA arr1 = A.unzip . liftArray2' quotRem arr1 +{-# INLINE quotRemA #-} + +-- | Perform a pointwise integer division with modulo where first array contains +-- numerators and the second one denominators +-- +-- > divModA arr1 arr2 == zipWith divMod arr1 arr2 +-- +-- [Partial] Mismatched array sizes will result in an impure exception being thrown. +-- +-- @since 0.1.0 +divModA + :: (HasCallStack, Index ix, Source r1 e, Source r2 e, Integral e) + => Array r1 ix e + -> Array r2 ix e + -> (Array D ix e, Array D ix e) +divModA arr1 = A.unzip . liftArray2' divMod arr1 +{-# INLINE divModA #-} + +-- | Truncate each element of the array. +-- +-- > truncateA arr == map truncate arr +-- +-- @since 0.1.0 +truncateA :: (Index ix, Source r a, RealFrac a, Integral e) => Array r ix a -> Array D ix e +truncateA = A.map truncate +{-# INLINE truncateA #-} + +-- | Round each element of the array. +-- +-- > truncateA arr == map truncate arr +-- +-- @since 0.1.0 +roundA :: (Index ix, Source r a, RealFrac a, Integral e) => Array r ix a -> Array D ix e +roundA = A.map round +{-# INLINE roundA #-} + +-- | Ceiling of each element of the array. +-- +-- > truncateA arr == map truncate arr +-- +-- @since 0.1.0 +ceilingA :: (Index ix, Source r a, RealFrac a, Integral e) => Array r ix a -> Array D ix e +ceilingA = A.map ceiling +{-# INLINE ceilingA #-} + +-- | Floor each element of the array. +-- +-- > truncateA arr == map truncate arr +-- +-- @since 0.1.0 +floorA :: (Index ix, Source r a, RealFrac a, Integral e) => Array r ix a -> Array D ix e +floorA = A.map floor +{-# INLINE floorA #-} + +-- | Perform atan2 pointwise +-- +-- > atan2A arr1 arr2 == zipWith atan2 arr1 arr2 +-- +-- /__Throws Exception__/: `SizeMismatchException` when array sizes do not match. +-- +-- @since 0.1.0 +atan2A + :: (Index ix, Numeric r e, RealFloat e, MonadThrow m) + => Array r ix e + -> Array r ix e + -> m (Array r ix e) +atan2A = liftNumArray2M atan2 +{-# INLINE atan2A #-} + +-- | Same as `sumArraysM`, compute sum of arrays pointwise. All arrays must have the same +-- size, otherwise it will result in an error. +-- +-- @since 1.0.0 +sumArrays' + :: (HasCallStack, Foldable t, Load r ix e, Numeric r e) => t (Array r ix e) -> Array r ix e +sumArrays' = throwEither . sumArraysM +{-# INLINE sumArrays' #-} + +-- | Compute sum of arrays pointwise. All arrays must have the same size. +-- +-- ====__Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> sumArraysM [] :: IO (Array P Ix3 Int) +-- Array P Seq (Sz (0 :> 0 :. 0)) +-- [ ] +-- >>> arr = A.makeArrayR P Seq (Sz3 4 5 6) $ \(i :> j :. k) -> i + j * k +-- >>> arr +-- Array P Seq (Sz (4 :> 5 :. 6)) +-- [ [ [ 0, 0, 0, 0, 0, 0 ] +-- , [ 0, 1, 2, 3, 4, 5 ] +-- , [ 0, 2, 4, 6, 8, 10 ] +-- , [ 0, 3, 6, 9, 12, 15 ] +-- , [ 0, 4, 8, 12, 16, 20 ] +-- ] +-- , [ [ 1, 1, 1, 1, 1, 1 ] +-- , [ 1, 2, 3, 4, 5, 6 ] +-- , [ 1, 3, 5, 7, 9, 11 ] +-- , [ 1, 4, 7, 10, 13, 16 ] +-- , [ 1, 5, 9, 13, 17, 21 ] +-- ] +-- , [ [ 2, 2, 2, 2, 2, 2 ] +-- , [ 2, 3, 4, 5, 6, 7 ] +-- , [ 2, 4, 6, 8, 10, 12 ] +-- , [ 2, 5, 8, 11, 14, 17 ] +-- , [ 2, 6, 10, 14, 18, 22 ] +-- ] +-- , [ [ 3, 3, 3, 3, 3, 3 ] +-- , [ 3, 4, 5, 6, 7, 8 ] +-- , [ 3, 5, 7, 9, 11, 13 ] +-- , [ 3, 6, 9, 12, 15, 18 ] +-- , [ 3, 7, 11, 15, 19, 23 ] +-- ] +-- ] +-- >>> sumArraysM $ outerSlices arr +-- Array P Seq (Sz (5 :. 6)) +-- [ [ 6, 6, 6, 6, 6, 6 ] +-- , [ 6, 10, 14, 18, 22, 26 ] +-- , [ 6, 14, 22, 30, 38, 46 ] +-- , [ 6, 18, 30, 42, 54, 66 ] +-- , [ 6, 22, 38, 54, 70, 86 ] +-- ] +-- >>> sumArraysM $ innerSlices arr +-- Array D Seq (Sz (4 :. 5)) +-- [ [ 0, 15, 30, 45, 60 ] +-- , [ 6, 21, 36, 51, 66 ] +-- , [ 12, 27, 42, 57, 72 ] +-- , [ 18, 33, 48, 63, 78 ] +-- ] +-- +-- @since 1.0.0 +sumArraysM + :: (Foldable t, Load r ix e, Numeric r e, MonadThrow m) => t (Array r ix e) -> m (Array r ix e) +sumArraysM as = + case F.toList as of + [] -> pure empty + (x : xs) -> F.foldlM (.+.) x xs +{-# INLINE sumArraysM #-} + +-- OPTIMIZE: Allocate a single result array and write sums into it incrementally. + +-- | Same as `productArraysM`. Compute product of arrays pointwise. All arrays must have +-- the same size, otherwise it +-- will result in an error. +-- +-- @since 1.0.0 +productArrays' + :: (HasCallStack, Foldable t, Load r ix e, Numeric r e) => t (Array r ix e) -> Array r ix e +productArrays' = throwEither . productArraysM +{-# INLINE productArrays' #-} + +-- | Compute product of arrays pointwise. All arrays must have the same size. +-- +-- ====__Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> productArraysM [] :: IO (Array P Ix3 Int) +-- Array P Seq (Sz (0 :> 0 :. 0)) +-- [ ] +-- >>> arr = A.makeArrayR P Seq (Sz3 4 5 6) $ \(i :> j :. k) -> i + j * k +-- >>> arr +-- Array P Seq (Sz (4 :> 5 :. 6)) +-- [ [ [ 0, 0, 0, 0, 0, 0 ] +-- , [ 0, 1, 2, 3, 4, 5 ] +-- , [ 0, 2, 4, 6, 8, 10 ] +-- , [ 0, 3, 6, 9, 12, 15 ] +-- , [ 0, 4, 8, 12, 16, 20 ] +-- ] +-- , [ [ 1, 1, 1, 1, 1, 1 ] +-- , [ 1, 2, 3, 4, 5, 6 ] +-- , [ 1, 3, 5, 7, 9, 11 ] +-- , [ 1, 4, 7, 10, 13, 16 ] +-- , [ 1, 5, 9, 13, 17, 21 ] +-- ] +-- , [ [ 2, 2, 2, 2, 2, 2 ] +-- , [ 2, 3, 4, 5, 6, 7 ] +-- , [ 2, 4, 6, 8, 10, 12 ] +-- , [ 2, 5, 8, 11, 14, 17 ] +-- , [ 2, 6, 10, 14, 18, 22 ] +-- ] +-- , [ [ 3, 3, 3, 3, 3, 3 ] +-- , [ 3, 4, 5, 6, 7, 8 ] +-- , [ 3, 5, 7, 9, 11, 13 ] +-- , [ 3, 6, 9, 12, 15, 18 ] +-- , [ 3, 7, 11, 15, 19, 23 ] +-- ] +-- ] +-- >>> productArraysM $ outerSlices arr +-- Array P Seq (Sz (5 :. 6)) +-- [ [ 0, 0, 0, 0, 0, 0 ] +-- , [ 0, 24, 120, 360, 840, 1680 ] +-- , [ 0, 120, 840, 3024, 7920, 17160 ] +-- , [ 0, 360, 3024, 11880, 32760, 73440 ] +-- , [ 0, 840, 7920, 32760, 93024, 212520 ] +-- ] +-- >>> productArraysM $ innerSlices arr +-- Array D Seq (Sz (4 :. 5)) +-- [ [ 0, 0, 0, 0, 0 ] +-- , [ 1, 720, 10395, 58240, 208845 ] +-- , [ 64, 5040, 46080, 209440, 665280 ] +-- , [ 729, 20160, 135135, 524880, 1514205 ] +-- ] +-- +-- @since 1.0.0 +productArraysM + :: (Foldable t, Load r ix e, Numeric r e, MonadThrow m) => t (Array r ix e) -> m (Array r ix e) +productArraysM as = + case F.toList as of + [] -> pure empty + (x : xs) -> F.foldlM (.*.) x xs +{-# INLINE productArraysM #-} diff --git a/src/Data/Massiv/Array/Numeric/Integral.hs b/src/Data/Massiv/Array/Numeric/Integral.hs --- a/src/Data/Massiv/Array/Numeric/Integral.hs +++ b/src/Data/Massiv/Array/Numeric/Integral.hs @@ -1,46 +1,48 @@ -{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE BangPatterns #-} {-# LANGUAGE FlexibleContexts #-} + -- | -- Module : Data.Massiv.Array.Numeric.Integral --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Numeric.Integral - ( +module Data.Massiv.Array.Numeric.Integral ( -- $integral_intro - -- - -- * Integral Approximation - -- ** Midpoint Rule - midpointRule - , midpointStencil + midpointRule, + midpointStencil, + -- ** Trapezoid Rule - , trapezoidRule - , trapezoidStencil + trapezoidRule, + trapezoidStencil, + -- ** Simpson's Rule - , simpsonsRule - , simpsonsStencil + simpsonsRule, + simpsonsStencil, + -- * General Integral approximation - , integrateWith - , integralApprox + integrateWith, + integralApprox, + -- * From functions + -- ** Sampled at the edge - , fromFunction - -- ** Sampled at the midpoint - , fromFunctionMidpoint - ) where + fromFunction, -import Data.Coerce -import Data.Massiv.Array.Delayed.Pull (D) -import Data.Massiv.Array.Delayed.Windowed (DW) -import Data.Massiv.Array.Manifest.Internal -import Data.Massiv.Array.Ops.Construct (rangeInclusive) -import Data.Massiv.Array.Ops.Transform (extract') -import Data.Massiv.Array.Stencil -import Data.Massiv.Core.Common + -- ** Sampled at the midpoint + fromFunctionMidpoint, +) where +import Data.Coerce +import Data.Massiv.Array.Delayed.Pull (D) +import Data.Massiv.Array.Delayed.Windowed (DW) +import Data.Massiv.Array.Manifest.Internal +import Data.Massiv.Array.Ops.Construct (rangeInclusive) +import Data.Massiv.Array.Ops.Transform (extract') +import Data.Massiv.Array.Stencil +import Data.Massiv.Array.Unsafe +import Data.Massiv.Core.Common -- | -- @@ -49,18 +51,20 @@ -- \[ -- \int_{{\,a}}^{{\,b}}{{f\left( x \right)\,dx}} \approx \Delta x \cdot \,f\left( {x_1 + \frac{\Delta x}{2}} \right) + \Delta x \cdot \,f\left( {x_2 + \frac{\Delta x}{2}} \right) + \cdots + \Delta x \cdot \,f\left( {x_n + \frac{\Delta x}{2}} \right) -- \] -midpointStencil :: - (Fractional e, Index ix) - => e -- ^ @Δx@ - distance between sample points - -> Dim -- ^ Dimension along which to integrate - -> Int -- ^ @n@ - number of sample points. +midpointStencil + :: (Fractional e, Index ix) + => e + -- ^ @Δx@ - distance between sample points + -> Dim + -- ^ Dimension along which to integrate + -> Int + -- ^ @n@ - number of sample points. -> Stencil ix e e midpointStencil dx dim k = - makeStencilDef 0 (Sz (setDim' (pureIndex 1) dim k)) zeroIndex $ \g -> - pure dx * loop 0 (< k) (+ 1) 0 (\i -> (+ g (setDim' zeroIndex dim i))) + makeUnsafeStencil (Sz (setDim' (pureIndex 1) dim k)) zeroIndex $ \_ g -> + dx * loop 0 (< k) (+ 1) 0 (\i -> (+ g (setDim' zeroIndex dim i))) {-# INLINE midpointStencil #-} - -- | -- -- __Trapezoid Rule__ @@ -68,20 +72,23 @@ -- \[ -- \int_{{\,a}}^{{\,b}}{{f\left( x \right)\,dx}} \approx \frac{{\Delta x}}{2}\cdot\left( {f\left( {{x_0}} \right) + f\left( {{x_1}} \right)} \right) + \frac{{\Delta x}}{2}\cdot\left( {f\left( {{x_1}} \right) + f\left( {{x_2}} \right)} \right) + \cdots + \frac{{\Delta x}}{2}\cdot\left( {f\left( {{x_{n - 1}}} \right) + f\left( {{x_n}} \right)} \right) -- \] -trapezoidStencil :: - (Fractional e, Index ix) - => e -- ^ @Δx@ - distance between sample points - -> Dim -- ^ Dimension along which to integrate - -> Int -- ^ @n@ - number of sample points. +trapezoidStencil + :: (Fractional e, Index ix) + => e + -- ^ @Δx@ - distance between sample points + -> Dim + -- ^ Dimension along which to integrate + -> Int + -- ^ @n@ - number of sample points. -> Stencil ix e e trapezoidStencil dx dim n = - makeStencilDef 0 (Sz (setDim' (pureIndex 1) dim (n + 1))) zeroIndex $ \g -> - pure dx / 2 * - (loop 1 (< n) (+ 1) (g zeroIndex) (\i -> (+ 2 * g (setDim' zeroIndex dim i))) + - g (setDim' zeroIndex dim n)) + makeUnsafeStencil (Sz (setDim' (pureIndex 1) dim (n + 1))) zeroIndex $ \_ g -> + (dx / 2) + * ( loop 1 (< n) (+ 1) (g zeroIndex) (\i -> (+ 2 * g (setDim' zeroIndex dim i))) + + g (setDim' zeroIndex dim n) + ) {-# INLINE trapezoidStencil #-} - -- | -- -- __Simpson's Rule__ @@ -89,31 +96,36 @@ -- \[ -- \int_{{\,a}}^{{\,b}}{{f\left( x \right)\,dx}} \approx \frac{{\Delta x}}{3}\cdot\left( {f\left( {{x_0}} \right) + 4\cdot f\left( {{x_1}} \right) + f\left( {{x_2}} \right)} \right) + \frac{{\Delta x}}{3}\cdot\left( {f\left( {{x_2}} \right) + 4\cdot f\left( {{x_3}} \right) + f\left( {{x_4}} \right)} \right) + \cdots + \frac{{\Delta x}}{3}\cdot\left( {f\left( {{x_{n - 2}}} \right) + 4\cdot f\left( {{x_{n - 1}}} \right) + f\left( {{x_n}} \right)} \right) -- \] -simpsonsStencil :: - (Fractional e, Index ix) - => e -- ^ @Δx@ - distance between sample points - -> Dim -- ^ Dimension along which to integrate - -> Int -- ^ @n@ - Number of sample points. This value should be even, otherwise error. +simpsonsStencil + :: (Fractional e, Index ix) + => e + -- ^ @Δx@ - distance between sample points + -> Dim + -- ^ Dimension along which to integrate + -> Int + -- ^ @n@ - Number of sample points. This value should be even, otherwise error. -> Stencil ix e e simpsonsStencil dx dim n | odd n = - error $ - "Number of sample points for Simpson's rule stencil should be even, but received: " ++ show n + error $ + "Number of sample points for Simpson's rule stencil should be even, but received: " ++ show n | otherwise = - makeStencilDef 0 (Sz (setDim' (pureIndex 1) dim (n + 1))) zeroIndex $ \g -> - let simAcc i (prev, acc) = - let !fx3 = g (setDim' zeroIndex dim (i + 2)) - !newAcc = acc + prev + 4 * g (setDim' zeroIndex dim (i + 1)) + fx3 - in (fx3, newAcc) - in pure dx / 3 * snd (loop 2 (< n - 1) (+ 2) (simAcc 0 (g zeroIndex, 0)) simAcc) + makeUnsafeStencil (Sz (setDim' (pureIndex 1) dim (n + 1))) zeroIndex $ \_ g -> + let simAcc i (prev, acc) = + let !fx3 = g (setDim' zeroIndex dim (i + 2)) + !newAcc = acc + prev + 4 * g (setDim' zeroIndex dim (i + 1)) + fx3 + in (fx3, newAcc) + in dx / 3 * snd (loop 2 (< n - 1) (+ 2) (simAcc 0 (g zeroIndex, 0)) simAcc) {-# INLINE simpsonsStencil #-} -- | Integrate with a stencil along a particular dimension. -integrateWith :: - (Fractional e, StrideLoad DW ix e, Mutable r ix e) +integrateWith + :: (Fractional e, StrideLoad DW ix e, Manifest r e) => (Dim -> Int -> Stencil ix e e) - -> Dim -- ^ Dimension along which integration should be estimated. - -> Int -- ^ @n@ - Number of samples + -> Dim + -- ^ Dimension along which integration should be estimated. + -> Int + -- ^ @n@ - Number of samples -> Array r ix e -> Array r ix e integrateWith stencil dim n arr = @@ -122,73 +134,95 @@ !nsz = setDim' (pureIndex 1) dim n {-# INLINE integrateWith #-} - -- | Compute an approximation of integral using a supplied rule in a form of `Stencil`. -integralApprox :: - (Fractional e, StrideLoad DW ix e, Mutable r ix e) - => (e -> Dim -> Int -> Stencil ix e e) -- ^ Integration Stencil - -> e -- ^ @d@ - Length of interval per cell - -> Sz ix -- ^ @sz@ - Result size of the matrix - -> Int -- ^ @n@ - Number of samples - -> Array r ix e -- ^ Array with values of @f(x,y,..)@ that will be used as source for integration. - -> Array M ix e +integralApprox + :: (Fractional e, StrideLoad DW ix e, Manifest r e) + => (e -> Dim -> Int -> Stencil ix e e) + -- ^ Integration Stencil + -> e + -- ^ @d@ - Length of interval per cell + -> Sz ix + -- ^ @sz@ - Result size of the matrix + -> Int + -- ^ @n@ - Number of samples + -> Array r ix e + -- ^ Array with values of @f(x,y,..)@ that will be used as source for integration. + -> Array D ix e integralApprox stencil d sz n arr = - extract' zeroIndex sz $ toManifest $ loop 1 (<= coerce (dimensions sz)) (+ 1) arr integrateAlong + extract' zeroIndex sz $ loop 1 (<= coerce (dimensions sz)) (+ 1) arr integrateAlong where !dx = d / fromIntegral n integrateAlong dim = integrateWith (stencil dx) (Dim dim) n {-# INLINE integrateAlong #-} {-# INLINE integralApprox #-} - -- | Use midpoint rule to approximate an integral. -midpointRule :: - (Fractional e, StrideLoad DW ix e, Mutable r ix e) - => Comp -- ^ Computation strategy. - -> r -- ^ Intermediate array representation. - -> ((Int -> e) -> ix -> e) -- ^ @f(x,y,...)@ - Function to integrate - -> e -- ^ @a@ - Starting value point. - -> e -- ^ @d@ - Distance per matrix cell. - -> Sz ix -- ^ @sz@ - Result matrix size. - -> Int -- ^ @n@ - Number of sample points per cell in each direction. - -> Array M ix e +midpointRule + :: (Fractional e, StrideLoad DW ix e, Manifest r e) + => Comp + -- ^ Computation strategy. + -> r + -- ^ Intermediate array representation. + -> ((Int -> e) -> ix -> e) + -- ^ @f(x,y,...)@ - Function to integrate + -> e + -- ^ @a@ - Starting value point. + -> e + -- ^ @d@ - Distance per matrix cell. + -> Sz ix + -- ^ @sz@ - Result matrix size. + -> Int + -- ^ @n@ - Number of sample points per cell in each direction. + -> Array D ix e midpointRule comp r f a d sz n = integralApprox midpointStencil d sz n $ computeAs r $ fromFunctionMidpoint comp f a d sz n {-# INLINE midpointRule #-} - -- | Use trapezoid rule to approximate an integral. -trapezoidRule :: - (Fractional e, StrideLoad DW ix e, Mutable r ix e) - => Comp -- ^ Computation strategy - -> r -- ^ Intermediate array representation - -> ((Int -> e) -> ix -> e) -- ^ @f(x,y,...)@ - function to integrate - -> e -- ^ @a@ - Starting value point. - -> e -- ^ @d@ - Distance per matrix cell. - -> Sz ix -- ^ @sz@ - Result matrix size. - -> Int -- ^ @n@ - Number of sample points per cell in each direction. - -> Array M ix e +trapezoidRule + :: (Fractional e, StrideLoad DW ix e, Manifest r e) + => Comp + -- ^ Computation strategy + -> r + -- ^ Intermediate array representation + -> ((Int -> e) -> ix -> e) + -- ^ @f(x,y,...)@ - function to integrate + -> e + -- ^ @a@ - Starting value point. + -> e + -- ^ @d@ - Distance per matrix cell. + -> Sz ix + -- ^ @sz@ - Result matrix size. + -> Int + -- ^ @n@ - Number of sample points per cell in each direction. + -> Array D ix e trapezoidRule comp r f a d sz n = integralApprox trapezoidStencil d sz n $ computeAs r $ fromFunction comp f a d sz n {-# INLINE trapezoidRule #-} -- | Use Simpson's rule to approximate an integral. -simpsonsRule :: - (Fractional e, StrideLoad DW ix e, Mutable r ix e) - => Comp -- ^ Computation strategy - -> r -- ^ Intermediate array representation - -> ((Int -> e) -> ix -> e) -- ^ @f(x,y,...)@ - Function to integrate - -> e -- ^ @a@ - Starting value point. - -> e -- ^ @d@ - Distance per matrix cell. - -> Sz ix -- ^ @sz@ - Result matrix size. - -> Int -- ^ @n@ - Number of sample points per cell in each direction. This value must be even, - -- otherwise error. - -> Array M ix e +simpsonsRule + :: (Fractional e, StrideLoad DW ix e, Manifest r e) + => Comp + -- ^ Computation strategy + -> r + -- ^ Intermediate array representation + -> ((Int -> e) -> ix -> e) + -- ^ @f(x,y,...)@ - Function to integrate + -> e + -- ^ @a@ - Starting value point. + -> e + -- ^ @d@ - Distance per matrix cell. + -> Sz ix + -- ^ @sz@ - Result matrix size. + -> Int + -- ^ @n@ - Number of sample points per cell in each direction. This value must be even, + -- otherwise error. + -> Array D ix e simpsonsRule comp r f a d sz n = integralApprox simpsonsStencil d sz n $ computeAs r $ fromFunction comp f a d sz n {-# INLINE simpsonsRule #-} - -- | Create an array from a function with sample points at the edges -- -- >>> fromFunction Seq (\ scale (i :. j) -> scale i + scale j :: Double) (-2) 1 (Sz 4) 2 @@ -203,17 +237,21 @@ -- , [ -0.5, 0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5 ] -- , [ 0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0 ] -- ] --- -fromFunction :: - (Index ix, Fractional a) - => Comp -- ^ Computation strategy +fromFunction + :: (Index ix, Fractional a) + => Comp + -- ^ Computation strategy -> ((Int -> a) -> ix -> e) -- ^ A function that will produce elements of scaled up array. First argument is a scaling -- function that should be applied to individual indicies. - -> a -- ^ @a@ - Starting point - -> a -- ^ @d@ - Distance per cell - -> Sz ix -- ^ @sz@ - Size of the desired array - -> Int -- ^ @n@ - Scaling factor, i.e. number of sample points per cell. + -> a + -- ^ @a@ - Starting point + -> a + -- ^ @d@ - Distance per cell + -> Sz ix + -- ^ @sz@ - Size of the desired array + -> Int + -- ^ @n@ - Scaling factor, i.e. number of sample points per cell. -> Array D ix e fromFunction comp f a d (Sz sz) n = f scale <$> rangeInclusive comp zeroIndex (liftIndex (n *) sz) @@ -223,7 +261,6 @@ {-# INLINE scale #-} {-# INLINE fromFunction #-} - -- | Similar to `fromFunction`, but will create an array from a function with sample points in the -- middle of cells. -- @@ -238,10 +275,15 @@ -- , [ -0.5, 0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0 ] -- , [ 0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5 ] -- ] --- fromFunctionMidpoint - :: (Index ix, Fractional a) => - Comp -> ((Int -> a) -> ix -> e) -> a -> a -> Sz ix -> Int -> Array D ix e + :: (Index ix, Fractional a) + => Comp + -> ((Int -> a) -> ix -> e) + -> a + -> a + -> Sz ix + -> Int + -> Array D ix e fromFunctionMidpoint comp f a d (Sz sz) n = f scale <$> rangeInclusive comp zeroIndex (liftIndex (\i -> n * i - 1) sz) where @@ -251,7 +293,6 @@ {-# INLINE scale #-} {-# INLINE fromFunctionMidpoint #-} - -- $integral_intro -- -- Inspiration for the code in this module was taken from [Paul Dawkins Online @@ -259,8 +300,9 @@ -- Approximation](http://tutorial.math.lamar.edu/Classes/CalcII/ApproximatingDefIntegrals.aspx), -- so if you need to brush up on some theory it is a great place to start. -- --- Implementation-wise, integral approximation here relies heavily on stencils with stride, as such --- computation is fast and is automatically parallelizable. +-- Implementation-wise, integral approximation here relies heavily on stencils +-- with stride, because such computation is fast and is automatically +-- parallelizable. -- -- Here are some examples of where this can be useful: -- @@ -281,7 +323,7 @@ -- stencils to compute an integral, but there are already functions that will do both steps for you: -- -- >>> simpsonsRule Seq U (\ scale x -> f (scale x)) 0 2 (Sz1 1) 4 --- Array M Seq (Sz1 1) +-- Array D Seq (Sz1 1) -- [ 17.353626 ] -- -- @scale@ is the function that will change an array index into equally spaced and @@ -304,7 +346,7 @@ -- The problem with above example is that computed values do not accurately represent the total -- value contained within each vector cell. For that reason if your were to later use it for example -- as convolution stencil, approximation would be very poor. The way to solve it is to approximate --- an integral across each cell of vector by drastically blowing up the `xArr` and then reducing it +-- an integral across each cell of vector by drastically blowing up the @xArr@ and then reducing it -- to a smaller array by using one of the approximation rules: -- -- >>> startValue = -2 :: Float @@ -317,14 +359,14 @@ -- [ -2.0, -1.75, -1.5, -1.25, -1.0, -0.75, -0.5, -0.25, 0.0, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0 ] -- >>> yArrX4 = computeAs U $ fmap f xArrX4 -- >>> integralApprox trapezoidStencil distPerCell desiredSize numSamples yArrX4 --- Array M Seq (Sz1 4) +-- Array D Seq (Sz1 4) -- [ 16.074406, 1.4906789, 1.4906789, 16.074408 ] -- -- We can clearly see the difference is huge, but it doesn't mean it is much better than our -- previous estimate. In order to get more accurate results we can use a better Simpson's rule for --- approximation and many more sample points. There is no need to create individual arrays `xArr` --- and `yArr`, there are functions like `simpsonRule` that will take care it for you: +-- approximation and many more sample points. There is no need to create individual arrays @xArrX4@ +-- and @yArrX4@, there are functions like `simpsonsRule` that will take care of it for us: -- -- >>> simpsonsRule Seq U (\ scale i -> f (scale i)) startValue distPerCell desiredSize 128 --- Array M Seq (Sz1 4) +-- Array D Seq (Sz1 4) -- [ 14.989977, 1.4626511, 1.4626517, 14.989977 ] diff --git a/src/Data/Massiv/Array/Ops/Construct.hs b/src/Data/Massiv/Array/Ops/Construct.hs --- a/src/Data/Massiv/Array/Ops/Construct.hs +++ b/src/Data/Massiv/Array/Ops/Construct.hs @@ -1,79 +1,86 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE DataKinds #-} -{-# LANGUAGE ExplicitForAll #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} + -- | -- Module : Data.Massiv.Array.Ops.Construct --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Ops.Construct - ( -- ** With constant value - empty - , singleton - , replicate - -- ** With a function - , makeArray - , makeArrayLinear - , makeArrayR - , makeArrayLinearR - , makeVectorR - -- *** Iterating - , iterateN - , iiterateN - -- *** Unfolding - , unfoldr - , unfoldrN - , unfoldlS_ - , iunfoldlS_ - , unfoldrS_ - , iunfoldrS_ - -- *** Random - , randomArray - , randomArrayS - , randomArrayWS - -- *** Applicative - , makeArrayA - , makeArrayAR - , makeArrayLinearA - -- ** Enumeration - , (...) - , (..:) - , range - , rangeStepM - , rangeStep' - , rangeInclusive - , rangeStepInclusiveM - , rangeStepInclusive' - , rangeSize - , rangeStepSize - , enumFromN - , enumFromStepN - -- ** Expansion - , expandWithin - , expandWithinM - , expandWithin' - , expandOuter - , expandInner - ) where +module Data.Massiv.Array.Ops.Construct ( + -- ** With constant value + empty, + singleton, + replicate, + -- ** With a function + makeArray, + makeArrayLinear, + makeArrayR, + makeArrayLinearR, + makeVectorR, + + -- *** Iterating + iterateN, + iiterateN, + + -- *** Unfolding + unfoldlS_, + iunfoldlS_, + unfoldrS_, + iunfoldrS_, + makeSplitSeedArray, + + -- *** Random + uniformArray, + uniformRangeArray, + randomArray, + randomArrayS, + randomArrayWS, + + -- *** Applicative + makeArrayA, + makeArrayAR, + makeArrayLinearA, + + -- ** Enumeration + (...), + (..:), + range, + rangeStepM, + rangeStep', + rangeInclusive, + rangeStepInclusiveM, + rangeStepInclusive', + rangeSize, + rangeStepSize, + enumFromN, + enumFromStepN, + + -- ** Expansion + expandWithin, + expandWithinM, + expandWithin', + expandOuter, + expandInner, +) where + import Control.Applicative hiding (empty) -import Control.Monad (when, void) +import Control.Monad (void, when) import Control.Monad.ST import Data.Massiv.Array.Delayed.Pull import Data.Massiv.Array.Delayed.Push -import Data.Massiv.Array.Delayed.Stream (unfoldr, unfoldrN) +-- import Data.Massiv.Array.Delayed.Stream (unfoldr, unfoldrN) import Data.Massiv.Array.Mutable import Data.Massiv.Core.Common -import Prelude as P hiding (enumFromTo, replicate) +import System.Random.Stateful +import Prelude hiding (enumFromTo, replicate) -- | Just like `makeArray` but with ability to specify the result representation as an -- argument. Note the `Data.Massiv.Array.U`nboxed type constructor in the below example. @@ -94,38 +101,29 @@ -- ] -- -- @since 0.1.0 -makeArrayR :: Construct r ix e => r -> Comp -> Sz ix -> (ix -> e) -> Array r ix e +makeArrayR :: Load r ix e => r -> Comp -> Sz ix -> (ix -> e) -> Array r ix e makeArrayR _ = makeArray {-# INLINE makeArrayR #-} -- | Same as `makeArrayLinear`, but with ability to supply resulting representation -- -- @since 0.3.0 -makeArrayLinearR :: Construct r ix e => r -> Comp -> Sz ix -> (Int -> e) -> Array r ix e +makeArrayLinearR :: Load r ix e => r -> Comp -> Sz ix -> (Int -> e) -> Array r ix e makeArrayLinearR _ = makeArrayLinear {-# INLINE makeArrayLinearR #-} -- | Same as `makeArrayR`, but restricted to 1-dimensional arrays. -- -- @since 0.1.0 -makeVectorR :: Construct r Ix1 e => r -> Comp -> Sz1 -> (Ix1 -> e) -> Array r Ix1 e +makeVectorR :: Load r Ix1 e => r -> Comp -> Sz1 -> (Ix1 -> e) -> Vector r e makeVectorR _ = makeArray {-# INLINE makeVectorR #-} - --- | Replicate the same element --- --- @since 0.3.0 -replicate :: forall r ix e . Construct r ix e => Comp -> Sz ix -> e -> Array r ix e -replicate comp sz e = makeArray comp sz (const e) -{-# INLINE replicate #-} - - newtype STA r ix a = STA {_runSTA :: forall s. MArray s r ix a -> ST s (Array r ix a)} -runSTA :: Mutable r ix e => Sz ix -> STA r ix e -> Array r ix e +runSTA :: (Manifest r e, Index ix) => Sz ix -> STA r ix e -> Array r ix e runSTA !sz (STA m) = runST (unsafeNew sz >>= m) -{-# INLINE runSTA #-} +{-# INLINE runSTA #-} -- | Similar to `makeArray`, but construct the array sequentially using an `Applicative` interface. -- @@ -134,46 +132,39 @@ -- -- -- @since 0.2.6 -makeArrayA :: - forall r ix e f. (Mutable r ix e, Applicative f) +makeArrayA + :: forall r ix e f + . (Manifest r e, Index ix, Applicative f) => Sz ix -> (ix -> f e) -> f (Array r ix e) -makeArrayA !sz f = - let n = totalElem sz - go !i - | i < n = - liftA2 - (\e (STA st) -> STA (\ma -> unsafeLinearWrite ma i e >> st ma)) - (f (fromLinearIndex sz i)) - (go (i + 1)) - | otherwise = pure (STA (unsafeFreeze Seq)) - in runSTA sz <$> go 0 -{-# INLINE makeArrayA #-} +makeArrayA sz@(Sz n) f = + fmap (runSTA sz) $ + iterF zeroIndex n oneIndex (<) (pure (STA (unsafeFreeze Seq))) $ \ix g -> + liftA2 (\e (STA st) -> STA (\ma -> unsafeWrite ma ix e >> st ma)) (f ix) g +{-# INLINE makeArrayA #-} -- | Same as `makeArrayA`, but with linear index. -- -- @since 0.4.5 -makeArrayLinearA :: - forall r ix e f. (Mutable r ix e, Applicative f) +makeArrayLinearA + :: forall r ix e f + . (Manifest r e, Index ix, Applicative f) => Sz ix -> (Int -> f e) -> f (Array r ix e) makeArrayLinearA !sz f = - let n = totalElem sz - go !i - | i < n = - liftA2 (\e (STA st) -> STA (\ma -> unsafeLinearWrite ma i e >> st ma)) (f i) (go (i + 1)) - | otherwise = pure (STA (unsafeFreeze Seq)) - in runSTA sz <$> go 0 -{-# INLINE makeArrayLinearA #-} - + fmap (runSTA sz) $ + loopF 0 (< totalElem sz) (+ 1) (pure (STA (unsafeFreeze Seq))) $ \i -> + liftA2 (\e (STA st) -> STA (\ma -> unsafeLinearWrite ma i e >> st ma)) (f i) +{-# INLINE makeArrayLinearA #-} -- | Same as `makeArrayA`, but with ability to supply result array representation. -- -- @since 0.2.6 -makeArrayAR :: - forall r ix e f. (Mutable r ix e, Applicative f) +makeArrayAR + :: forall r ix e f + . (Manifest r e, Index ix, Applicative f) => r -> Sz ix -> (ix -> f e) @@ -181,7 +172,6 @@ makeArrayAR _ = makeArrayA {-# INLINE makeArrayAR #-} - -- | Sequentially iterate over each cell in the array in the row-major order while continuously -- aplying the accumulator at each step. -- @@ -195,29 +185,34 @@ -- ] -- -- @since 0.3.0 -iterateN :: forall ix e . Index ix => Sz ix -> (e -> e) -> e -> Array DL ix e +iterateN :: forall ix e. Index ix => Sz ix -> (e -> e) -> e -> Array DL ix e iterateN sz f = unfoldrS_ sz $ \a -> let !a' = f a in (a', a') {-# INLINE iterateN #-} -- | Same as `iterateN`, but with index aware function. -- -- @since 0.3.0 -iiterateN :: forall ix e . Index ix => Sz ix -> (e -> ix -> e) -> e -> Array DL ix e +iiterateN :: forall ix e. Index ix => Sz ix -> (e -> ix -> e) -> e -> Array DL ix e iiterateN sz f = iunfoldrS_ sz $ \a ix -> let !a' = f a ix in (a', a') {-# INLINE iiterateN #-} - --- | Right unfold of a delayed load array. For the inverse direction use `unfoldlS_`. +-- | Right unfold into a delayed load array. For the opposite direction use `unfoldlS_`. -- -- ==== __Examples__ -- -- >>> import Data.Massiv.Array --- >>> unfoldrS_ (Sz1 10) (\xs -> (head xs, tail xs)) ([10 ..] :: [Int]) +-- >>> unfoldrS_ (Sz1 10) (\xs -> (Prelude.head xs, Prelude.tail xs)) ([10 ..] :: [Int]) -- Array DL Seq (Sz1 10) -- [ 10, 11, 12, 13, 14, 15, 16, 17, 18, 19 ] -- -- @since 0.3.0 -unfoldrS_ :: forall ix e a . Construct DL ix e => Sz ix -> (a -> (e, a)) -> a -> Array DL ix e +unfoldrS_ + :: forall ix e a + . Index ix + => Sz ix + -> (a -> (e, a)) + -> a + -> Array DL ix e unfoldrS_ sz f = iunfoldrS_ sz (\a _ -> f a) {-# INLINE unfoldrS_ #-} @@ -225,29 +220,29 @@ -- -- @since 0.3.0 iunfoldrS_ - :: Construct DL ix e => Sz ix -> (a -> ix -> (e, a)) -> a -> Array DL ix e -iunfoldrS_ sz f acc0 = - DLArray - { dlComp = Seq - , dlSize = sz - , dlDefault = Nothing - , dlLoad = - \_ startAt dlWrite -> - void $ - loopM startAt (< (totalElem sz + startAt)) (+ 1) acc0 $ \ !i !acc -> do - let (e, acc') = f acc $ fromLinearIndex sz (i - startAt) - dlWrite i e - pure acc' - } + :: forall ix e a + . Index ix + => Sz ix + -> (a -> ix -> (e, a)) + -> a + -> Array DL ix e +iunfoldrS_ sz f acc0 = DLArray{dlComp = Seq, dlSize = sz, dlLoad = load} + where + load :: Loader e + load _ startAt dlWrite _ = + void $ + iterTargetM defRowMajor startAt sz zeroIndex oneStride acc0 $ \ !i !ix !acc -> + case f acc ix of + (e, !acc') -> acc' <$ dlWrite i e + {-# INLINE load #-} {-# INLINE iunfoldrS_ #-} - -- | Unfold sequentially from the end. There is no way to save the accumulator after -- unfolding is done, since resulting array is delayed, but it's possible to use --- `Data.Massiv.Array.Mutable.unfoldlPrimM` to achive such effect. +-- `Data.Massiv.Array.Mutable.unfoldlPrimM` to achieve such effect. -- -- @since 0.3.0 -unfoldlS_ :: Construct DL ix e => Sz ix -> (a -> (a, e)) -> a -> Array DL ix e +unfoldlS_ :: Index ix => Sz ix -> (a -> (a, e)) -> a -> Array DL ix e unfoldlS_ sz f = iunfoldlS_ sz (const f) {-# INLINE unfoldlS_ #-} @@ -255,22 +250,23 @@ -- -- @since 0.3.0 iunfoldlS_ - :: Construct DL ix e => Sz ix -> (ix -> a -> (a, e)) -> a -> Array DL ix e -iunfoldlS_ sz f acc0 = - DLArray - { dlComp = Seq - , dlSize = sz - , dlDefault = Nothing - , dlLoad = - \ _ startAt dlWrite -> - void $ loopDeepM startAt (< (totalElem sz + startAt)) (+ 1) acc0 $ \ !i !acc -> do - let (acc', e) = f (fromLinearIndex sz (i - startAt)) acc - dlWrite i e - pure acc' - } + :: forall ix e a + . Index ix + => Sz ix + -> (ix -> a -> (a, e)) + -> a + -> Array DL ix e +iunfoldlS_ sz f acc0 = DLArray{dlComp = Seq, dlSize = sz, dlLoad = load} + where + load :: Loader e + load _ startAt dlWrite _ = + void $ + loopDeepM startAt (< totalElem sz + startAt) (+ 1) acc0 $ \ !i !acc -> + let (acc', e) = f (fromLinearIndex sz (i - startAt)) acc + in acc' <$ dlWrite i e + {-# INLINE load #-} {-# INLINE iunfoldlS_ #-} - -- | Create an array with random values by using a pure splittable random number generator -- such as one provided by either [splitmix](https://www.stackage.org/package/splitmix) or -- [random](https://www.stackage.org/package/random) packages. If you don't have a @@ -280,6 +276,9 @@ -- to parallelize the random value generation, but also guarantee that it will be -- deterministic, granted none of the arguments have changed. -- +-- __Note__: Starting with massiv-1.1.0 this function will be deprecated in +-- favor of a more general `genSplitArray` +-- -- ==== __Examples__ -- -- >>> import Data.Massiv.Array @@ -292,49 +291,131 @@ -- ] -- -- >>> import Data.Massiv.Array --- >>> import System.Random as System --- >>> gen = System.mkStdGen 217 --- >>> randomArray gen System.split System.random (ParN 2) (Sz2 2 3) :: Array DL Ix2 Double +-- >>> import System.Random as Random +-- >>> gen = Random.mkStdGen 217 +-- >>> randomArray gen Random.split Random.random (ParN 2) (Sz2 2 3) :: Array DL Ix2 Double -- Array DL (ParN 2) (Sz (2 :. 3)) --- [ [ 0.15191527341922206, 0.2045537167404079, 0.9635356052820256 ] --- , [ 9.308278528094238e-2, 0.7200934018606843, 0.23173694193083583 ] +-- [ [ 0.2616843941380331, 0.600959468331641, 0.4382415961606372 ] +-- , [ 0.27812817813217605, 0.2993277194932741, 0.2774105268603957 ] -- ] -- --- @since 0.3.3 -randomArray :: - forall ix e g. Index ix - => g -- ^ Initial random value generator +-- @since 1.0.0 +randomArray + :: forall ix e g + . Index ix + => g + -- ^ Initial random value generator -> (g -> (g, g)) - -- ^ A function that can split a generator in two independent generators + -- ^ A function that can split a generator into two independent + -- generators. It will only be called if supplied computation strategy + -- needs more than one worker threads. -> (g -> (e, g)) - -- ^ A function that produces a random value and the next generator - -> Comp -- ^ Computation strategy. - -> Sz ix -- ^ Resulting size of the array. + -- ^ A function that produces a random value and the next generator + -> Comp + -- ^ Computation strategy. + -> Sz ix + -- ^ Resulting size of the array. -> Array DL ix e -randomArray gen splitGen nextRandom comp sz = - unsafeMakeLoadArray comp sz Nothing $ \scheduler startAt writeAt -> - splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do - let slackStartAt = slackStart + startAt - writeRandom k genII = do - let (e, genII') = nextRandom genII - writeAt k e - pure genII' - genForSlack <- - loopM startAt (< slackStartAt) (+ chunkLength) gen $ \start genI -> do - let (genI0, genI1) = - if numWorkers scheduler == 1 - then (genI, genI) - else splitGen genI - scheduleWork_ scheduler $ - void $ loopM start (< (start + chunkLength)) (+ 1) genI0 writeRandom - pure genI1 - when (slackStartAt < totalLength + startAt) $ - scheduleWork_ scheduler $ - void $ loopM slackStartAt (< totalLength + startAt) (+ 1) genForSlack writeRandom +randomArray gen splitGen' nextRandom comp sz = unsafeMakeLoadArray comp sz Nothing load where !totalLength = totalElem sz + load :: forall s. Scheduler s () -> Ix1 -> (Ix1 -> e -> ST s ()) -> ST s () + load scheduler startAt writeAt = + splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do + let slackStartAt = slackStart + startAt + writeRandom k genII = + let (e, genII') = nextRandom genII + in genII' <$ writeAt k e + genForSlack <- + loopM startAt (< slackStartAt) (+ chunkLength) gen $ \start genI -> do + let (genI0, genI1) = + if numWorkers scheduler == 1 + then (genI, genI) + else splitGen' genI + scheduleWork_ scheduler $ + void $ + loopM start (< start + chunkLength) (+ 1) genI0 writeRandom + pure genI1 + when (slackStart < totalLength) $ + scheduleWork_ scheduler $ + void $ + loopM slackStartAt (< totalLength + startAt) (+ 1) genForSlack writeRandom {-# INLINE randomArray #-} +-- | Create a delayed array with an initial seed and a splitting function. It is +-- somewhat similar to `iunfoldlS_` function, but it is capable of parallelizing +-- computation and iterating over the array accoriding to the supplied +-- `Iterator`. Upon parallelization every job will get the second part of the +-- result produced by the split function, while the first part will be used for +-- subsequent splits. This function is similar to +-- `Data.Massiv.Array.Manifest.generateSplitSeedArray` +-- +-- @since 1.0.2 +makeSplitSeedArray + :: forall ix e g it + . (Iterator it, Index ix) + => it + -- ^ Iterator + -> g + -- ^ Initial seed + -> (g -> (g, g)) + -- ^ A function that can split a seed into two independent seeds. It will + -- be called the same number of times as the number of jobs that will get + -- scheduled during parallelization. Eg. only once for the sequential case. + -> Comp + -- ^ Computation strategy. + -> Sz ix + -- ^ Resulting size of the array. + -> (Ix1 -> ix -> g -> (e, g)) + -- ^ A function that produces a value and the next seed. It takes both + -- versions of the index, in linear and in multi-dimensional forms, as well as + -- the current seeding value. + -> Array DL ix e +makeSplitSeedArray it seed splitSeed comp sz genFunc = + DLArray{dlComp = comp, dlSize = sz, dlLoad = load} + where + load :: Loader e + load scheduler startAt writeAt _ = + iterTargetFullAccST_ it scheduler startAt sz seed (pure . splitSeed) $ \i ix g -> + case genFunc (i - startAt) ix g of + (x, g') -> g' <$ writeAt i x + {-# INLINE load #-} +{-# INLINE makeSplitSeedArray #-} + +-- | Generate a random array where all elements are sampled from a uniform distribution. +-- +-- @since 1.0.0 +uniformArray + :: forall ix e g + . (Index ix, RandomGen g, Uniform e) + => g + -- ^ Initial random value generator. + -> Comp + -- ^ Computation strategy. + -> Sz ix + -- ^ Resulting size of the array. + -> Array DL ix e +uniformArray gen = randomArray gen split uniform +{-# INLINE uniformArray #-} + +-- | Same as `uniformArray`, but will generate values in a supplied range. +-- +-- @since 1.0.0 +uniformRangeArray + :: forall ix e g + . (Index ix, RandomGen g, UniformRange e) + => g + -- ^ Initial random value generator. + -> (e, e) + -- ^ Inclusive range in which values will be generated in. + -> Comp + -- ^ Computation strategy. + -> Sz ix + -- ^ Resulting size of the array. + -> Array DL ix e +uniformRangeArray gen r = randomArray gen split (uniformR r) +{-# INLINE uniformRangeArray #-} + -- | Similar to `randomArray` but performs generation sequentially, which means it doesn't -- require splitability property. Another consequence is that it returns the new generator -- together with /manifest/ array of random values. @@ -364,20 +445,23 @@ -- >>> gen = System.mkStdGen 217 -- >>> snd $ randomArrayS gen (Sz2 2 3) System.random :: Array P Ix2 Double -- Array P Seq (Sz (2 :. 3)) --- [ [ 0.7972230393466304, 0.4485860543300083, 0.257773196880671 ] --- , [ 0.19115043859955794, 0.33784788936970034, 3.479381605706322e-2 ] +-- [ [ 0.11217260506402493, 0.8870919238985904, 0.2616843941380331 ] +-- , [ 0.600959468331641, 0.4382415961606372, 0.8375162573397977 ] -- ] -- -- @since 0.3.4 -randomArrayS :: - forall r ix e g. Mutable r ix e - => g -- ^ Initial random value generator - -> Sz ix -- ^ Resulting size of the array. +randomArrayS + :: forall r ix e g + . (Manifest r e, Index ix) + => g + -- ^ Initial random value generator + -> Sz ix + -- ^ Resulting size of the array. -> (g -> (e, g)) - -- ^ A function that produces a random value and the next generator + -- ^ A function that produces a random value and the next generator -> (g, Array r ix e) randomArrayS gen sz nextRandom = - runST $ unfoldrPrimM Seq sz (pure . nextRandom) gen + runST $ unfoldrPrimM sz (pure . nextRandom) gen {-# INLINE randomArrayS #-} -- | This is a stateful approach of generating random values. If your generator is pure @@ -391,36 +475,44 @@ -- -- ==== __Examples__ -- --- In the example below we take a stateful random generator from +-- In the example below we take a stateful random number generator from -- [wmc-random](https://www.stackage.org/package/mwc-random), which is not thread safe, --- and safely parallelize it by giving each thread it's own generator: +-- and safely parallelize it by giving each thread it's own generator. There is a caveat +-- of course, statistical independence will depend on the entropy in your initial seeds, +-- so do not use the example below verbatim, since initial seeds are sequential numbers. -- --- > λ> import Data.Massiv.Array --- > λ> import System.Random.MWC (createSystemRandom, uniformR) --- > λ> import System.Random.MWC.Distributions (standard) --- > λ> gens <- initWorkerStates Par (\_ -> createSystemRandom) --- > λ> randomArrayWS gens (Sz2 2 3) standard :: IO (Array P Ix2 Double) --- > Array P Par (Sz (2 :. 3)) --- > [ [ -0.9066144845415213, 0.5264323240310042, -1.320943607597422 ] --- > , [ -0.6837929005619592, -0.3041255565826211, 6.53353089112833e-2 ] --- > ] --- > λ> randomArrayWS gens (Sz1 10) (uniformR (0, 9)) :: IO (Array P Ix1 Int) --- > Array P Par (Sz1 10) --- > [ 3, 6, 1, 2, 1, 7, 6, 0, 8, 8 ] +-- >>> import Data.Massiv.Array as A +-- >>> import System.Random.MWC as MWC (initialize) +-- >>> import System.Random.Stateful (uniformRM) +-- >>> import Control.Scheduler (initWorkerStates, getWorkerId) +-- >>> :set -XTypeApplications +-- >>> gens <- initWorkerStates Par (MWC.initialize . A.toPrimitiveVector . A.singleton @P @Ix1 . fromIntegral . getWorkerId) +-- >>> randomArrayWS gens (Sz2 2 3) (uniformRM (0, 9)) :: IO (Matrix P Double) +-- Array P Par (Sz (2 :. 3)) +-- [ [ 8.999240522095299, 6.832223390653755, 3.065728078741671 ] +-- , [ 7.242581103346686, 2.4565807301968623, 0.4514262066689775 ] +-- ] +-- >>> randomArrayWS gens (Sz1 6) (uniformRM (0, 9)) :: IO (Vector P Int) +-- Array P Par (Sz1 6) +-- [ 8, 8, 7, 1, 1, 2 ] -- -- @since 0.3.4 -randomArrayWS :: - forall r ix e g m. (Mutable r ix e, MonadUnliftIO m, PrimMonad m) - => WorkerStates g -- ^ Use `initWorkerStates` to initialize you per thread generators - -> Sz ix -- ^ Resulting size of the array - -> (g -> m e) -- ^ Generate the value using the per thread generator. +randomArrayWS + :: forall r ix e g m + . (Manifest r e, Index ix, MonadUnliftIO m, PrimMonad m) + => WorkerStates g + -- ^ Use `Control.Scheduler.initWorkerStates` to initialize you per thread generators + -> Sz ix + -- ^ Resulting size of the array + -> (g -> m e) + -- ^ Generate the value using the per thread generator. -> m (Array r ix e) randomArrayWS states sz genRandom = generateArrayLinearWS states sz (const genRandom) {-# INLINE randomArrayWS #-} infix 4 ..., ..: --- | Handy synonym for `rangeInclusive` `Seq` +-- | Handy synonym for @`rangeInclusive` `Seq`@. Similar to @..@ for list. -- -- >>> Ix1 4 ... 10 -- Array D Seq (Sz1 7) @@ -431,7 +523,7 @@ (...) = rangeInclusive Seq {-# INLINE (...) #-} --- | Handy synonym for `range` `Seq` +-- | Handy synonym for @`range` `Seq`@ -- -- >>> Ix1 4 ..: 10 -- Array D Seq (Sz1 6) @@ -442,9 +534,9 @@ (..:) = range Seq {-# INLINE (..:) #-} - -- prop> range comp from to == rangeStep comp from 1 to -- + -- | Create an array of indices with a range from start to finish (not-including), where indices are -- incremeted by one. -- @@ -468,6 +560,8 @@ -- | Same as `range`, but with a custom step. -- +-- /__Throws Exceptions__/: `IndexZeroException` +-- -- ==== __Examples__ -- -- >>> import Data.Massiv.Array @@ -478,19 +572,25 @@ -- *** Exception: IndexZeroException: 0 -- -- @since 0.3.0 -rangeStepM :: (Index ix, MonadThrow m) => - Comp -- ^ Computation strategy - -> ix -- ^ Start - -> ix -- ^ Step (Can't have zeros) - -> ix -- ^ End - -> m (Array D ix ix) +rangeStepM + :: forall ix m + . (Index ix, MonadThrow m) + => Comp + -- ^ Computation strategy + -> ix + -- ^ Start + -> ix + -- ^ Step. Negative and positive values are ok, but can't have zeros + -> ix + -- ^ End + -> m (Array D ix ix) rangeStepM comp !from !step !to | foldlIndex (\acc i -> acc || i == 0) False step = throwM $ IndexZeroException step | otherwise = - let dist = liftIndex2 (-) to from - sz = liftIndex2 div dist step - r = liftIndex signum $ liftIndex2 mod dist step - in pure $ rangeStepSize comp from step (Sz (liftIndex2 (+) sz r)) + let dist = liftIndex2 (-) to from + sz = liftIndex2 div dist step + r = liftIndex signum $ liftIndex2 mod dist step + in pure $ rangeStepSize comp from step (Sz (liftIndex2 (+) sz r)) {-# INLINE rangeStepM #-} -- | Same as `rangeStepM`, but will throw an error whenever @step@ contains zeros. @@ -503,8 +603,8 @@ -- [ 1, 3, 5 ] -- -- @since 0.3.0 -rangeStep' :: Index ix => Comp -> ix -> ix -> ix -> Array D ix ix -rangeStep' comp from step = either throw id . rangeStepM comp from step +rangeStep' :: (HasCallStack, Index ix) => Comp -> ix -> ix -> ix -> Array D ix ix +rangeStep' comp from step = throwEither . rangeStepM comp from step {-# INLINE rangeStep' #-} -- | Just like `range`, except the finish index is included. @@ -515,8 +615,7 @@ rangeSize comp ixFrom (Sz (liftIndex2 (-) (liftIndex (+ 1) ixTo) ixFrom)) {-# INLINE rangeInclusive #-} - --- | Just like `rangeStep`, except the finish index is included. +-- | Just like `rangeStepM`, except the finish index is included. -- -- @since 0.3.0 rangeStepInclusiveM :: (MonadThrow m, Index ix) => Comp -> ix -> ix -> ix -> m (Array D ix ix) @@ -526,39 +625,49 @@ -- | Just like `range`, except the finish index is included. -- -- @since 0.3.1 -rangeStepInclusive' :: Index ix => Comp -> ix -> ix -> ix -> Array D ix ix -rangeStepInclusive' comp ixFrom step = either throw id . rangeStepInclusiveM comp ixFrom step +rangeStepInclusive' :: (HasCallStack, Index ix) => Comp -> ix -> ix -> ix -> Array D ix ix +rangeStepInclusive' comp ixFrom step = throwEither . rangeStepInclusiveM comp ixFrom step {-# INLINE rangeStepInclusive' #-} - -- | Create an array of specified size with indices starting with some index at position @0@ and -- incremented by @1@ until the end of the array is reached -- -- @since 0.3.0 -rangeSize :: Index ix => - Comp - -> ix -- ^ @x@ - start value - -> Sz ix -- ^ @sz@ - Size of resulting array - -> Array D ix ix +rangeSize + :: Index ix + => Comp + -- ^ Computation strategy + -> ix + -- ^ @x@ - start value + -> Sz ix + -- ^ @sz@ - Size of resulting array + -> Array D ix ix rangeSize comp !from !sz = makeArray comp sz (liftIndex2 (+) from) {-# INLINE rangeSize #-} -- | Same as `rangeSize`, but with ability to specify the step. -- -- @since 0.3.0 -rangeStepSize :: Index ix => - Comp - -> ix -- ^ @x@ - start value - -> ix -- ^ @delta@ - step value - -> Sz ix -- ^ @sz@ - Size of resulting array - -> Array D ix ix +rangeStepSize + :: Index ix + => Comp + -- ^ Computation strategy + -> ix + -- ^ @x@ - start value + -> ix + -- ^ @delta@ - step value + -> Sz ix + -- ^ @sz@ - Size of resulting array + -> Array D ix ix rangeStepSize comp !from !step !sz = makeArray comp sz (liftIndex2 (+) from . liftIndex2 (*) step) {-# INLINE rangeStepSize #-} - --- | Same as `enumFromStepN` with step @delta = 1@. +-- | Same as `enumFromStepN` with step @dx = 1@. -- +-- /Related/: `Data.Massiv.Vector.senumFromN`, `Data.Massiv.Vector.senumFromStepN`, +-- `enumFromStepN`, `rangeSize`, `rangeStepSize`, `range` +-- -- ==== __Examples__ -- -- >>> import Data.Massiv.Array @@ -566,39 +675,66 @@ -- Array D Seq (Sz1 3) -- [ 5.0, 6.0, 7.0 ] -- +-- __/Similar/__: +-- +-- [@Prelude.`Prelude.enumFromTo`@] Very similar to @[i .. i + n - 1]@, except that +-- `enumFromN` is faster, but it only works for `Num` and not for `Enum` elements +-- +-- [@Data.Vector.Generic.`Data.Vector.Generic.enumFromN`@] +-- -- @since 0.1.0 -enumFromN :: Num e => - Comp - -> e -- ^ @x@ - start value - -> Sz1 -- ^ @n@ - length of resulting vector. - -> Array D Ix1 e -enumFromN comp !from !sz = makeArray comp sz $ \ i -> fromIntegral i + from +enumFromN + :: Num e + => Comp + -> e + -- ^ @x@ - start value + -> Sz1 + -- ^ @n@ - length of resulting vector. + -> Vector D e +enumFromN comp !from !sz = makeArrayLinear comp sz $ \i -> from + fromIntegral i {-# INLINE enumFromN #-} - --- | Create a vector with length @n@ that has it's 0th value set to @x@ and gradually increasing --- with @step@ delta until the end. Similar to: @`Data.Massiv.Array.fromList'` `Seq` $ `take` n [x, --- x + delta ..]@. Major difference is that `fromList` constructs an `Array` with manifest --- representation, while `enumFromStepN` is delayed. +-- | Enumerate from a starting number @x@ exactly @n@ times with a custom step value +-- @dx@. Unlike `Data.Massiv.Vector.senumFromStepN`, there is no dependency on neigboring +-- elements therefore `enumFromStepN` is parallelizable. -- +-- /Related/: `Data.Massiv.Vector.senumFromN`, `Data.Massiv.Vector.senumFromStepN`, +-- `enumFromN`, `rangeSize`, `rangeStepSize`, `range`, `rangeStepM` +-- -- ==== __Examples__ -- -- >>> import Data.Massiv.Array -- >>> enumFromStepN Seq 1 (0.1 :: Double) 5 -- Array D Seq (Sz1 5) -- [ 1.0, 1.1, 1.2, 1.3, 1.4 ] +-- >>> enumFromStepN Seq (-pi :: Float) (pi/4) 9 +-- Array D Seq (Sz1 9) +-- [ -3.1415927, -2.3561945, -1.5707964, -0.78539824, 0.0, 0.78539824, 1.5707963, 2.3561947, 3.1415927 ] -- +-- __/Similar/__: +-- +-- [@Prelude.`Prelude.enumFrom`@] Similar to @take n [x, x + dx ..]@, except that +-- `enumFromStepN` is parallelizable and it only works for `Num` and not for `Enum` +-- elements. Floating point value will be slightly different as well. +-- +-- [@Data.Vector.Generic.`Data.Vector.Generic.enumFromStepN`@] Similar in the +-- outcome, but very different in the way it works. +-- +-- -- @since 0.1.0 -enumFromStepN :: Num e => - Comp - -> e -- ^ @x@ - start value - -> e -- ^ @delta@ - step value - -> Sz1 -- ^ @n@ - length of resulting vector - -> Array D Ix1 e -enumFromStepN comp !from !step !sz = makeArray comp sz $ \ i -> from + fromIntegral i * step +enumFromStepN + :: Num e + => Comp + -> e + -- ^ @x@ - start number + -> e + -- ^ @dx@ - step number + -> Sz1 + -- ^ @n@ - length of resulting vector + -> Vector D e +enumFromStepN comp !from !step !sz = makeArrayLinear comp sz $ \i -> from + fromIntegral i * step {-# INLINE enumFromStepN #-} - -- | Function that expands an array to one with a higher dimension. -- -- This is useful for constructing arrays where there is shared computation @@ -640,8 +776,9 @@ -- ] -- -- @since 0.2.6 -expandWithin :: - forall ix e r n a. (IsIndexDimension ix n, Manifest r (Lower ix) a) +expandWithin + :: forall n ix e r a + . (IsIndexDimension ix n, Index (Lower ix), Manifest r a) => Dimension n -> Sz1 -> (a -> Ix1 -> e) @@ -661,13 +798,14 @@ -- -- @since 0.2.6 expandWithin' - :: (Index ix, Manifest r (Lower ix) a) + :: forall r ix a b + . (HasCallStack, Index ix, Index (Lower ix), Manifest r a) => Dim -> Sz1 -> (a -> Ix1 -> b) -> Array r (Lower ix) a -> Array D ix b -expandWithin' dim k f arr = either throw id $ expandWithinM dim k f arr +expandWithin' dim k f = throwEither . expandWithinM dim k f {-# INLINE expandWithin' #-} -- | Similar to `expandWithin`, except that dimension is specified at a value level, which means it @@ -675,7 +813,8 @@ -- -- @since 0.4.0 expandWithinM - :: (Index ix, Manifest r (Lower ix) a, MonadThrow m) + :: forall r ix a b m + . (Index ix, Index (Lower ix), Manifest r a, MonadThrow m) => Dim -> Sz1 -> (a -> Ix1 -> b) @@ -693,7 +832,8 @@ -- -- @since 0.2.6 expandOuter - :: (Index ix, Manifest r (Lower ix) a) + :: forall r ix a b + . (Index ix, Index (Lower ix), Manifest r a) => Sz1 -> (a -> Ix1 -> b) -> Array r (Lower ix) a @@ -711,7 +851,8 @@ -- -- @since 0.2.6 expandInner - :: (Index ix, Manifest r (Lower ix) a) + :: forall r ix a b + . (Index ix, Index (Lower ix), Manifest r a) => Sz1 -> (a -> Ix1 -> b) -> Array r (Lower ix) a diff --git a/src/Data/Massiv/Array/Ops/Fold.hs b/src/Data/Massiv/Array/Ops/Fold.hs --- a/src/Data/Massiv/Array/Ops/Fold.hs +++ b/src/Data/Massiv/Array/Ops/Fold.hs @@ -3,154 +3,164 @@ {-# LANGUAGE GADTs #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE ScopedTypeVariables #-} + -- | -- Module : Data.Massiv.Array.Ops.Fold --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Ops.Fold - ( +module Data.Massiv.Array.Ops.Fold ( -- ** Unstructured folds - -- $unstruct_folds - fold - , ifoldMono - , foldMono - , ifoldSemi - , foldSemi - , foldOuterSlice - , ifoldOuterSlice - , foldInnerSlice - , ifoldInnerSlice - , minimumM - , minimum' - , maximumM - , maximum' - , sum - , product - , and - , or - , all - , any + fold, + ifoldMono, + foldMono, + ifoldSemi, + foldSemi, + foldOuterSlice, + ifoldOuterSlice, + foldInnerSlice, + ifoldInnerSlice, + minimumM, + minimum', + maximumM, + maximum', + sum, + product, + and, + or, + all, + any, + elem, + eqArrays, + compareArrays, -- ** Single dimension folds + -- *** Safe inner most + -- -- Folding along the inner most dimension will always be faster when compared to doing the same -- operation along any other dimension, this is due to the fact that inner most folds follow the -- memory layout of data. - , ifoldlInner - , foldlInner - , ifoldrInner - , foldrInner - , foldInner + ifoldlInner, + foldlInner, + ifoldrInner, + foldrInner, + foldInner, + -- *** Type safe within - , ifoldlWithin - , foldlWithin - , ifoldrWithin - , foldrWithin - , foldWithin + ifoldlWithin, + foldlWithin, + ifoldrWithin, + foldrWithin, + foldWithin, + -- *** Partial within - , ifoldlWithin' - , foldlWithin' - , ifoldrWithin' - , foldrWithin' - , foldWithin' + ifoldlWithin', + foldlWithin', + ifoldrWithin', + foldrWithin', + foldWithin', -- ** Sequential folds - -- $seq_folds - - , foldlS - , foldrS - , ifoldlS - , ifoldrS + foldlS, + foldrS, + ifoldlS, + ifoldrS, -- *** Monadic - , foldlM - , foldrM - , foldlM_ - , foldrM_ - , ifoldlM - , ifoldrM - , ifoldlM_ - , ifoldrM_ + foldlM, + foldrM, + foldlM_, + foldrM_, + ifoldlM, + ifoldrM, + ifoldlM_, + ifoldrM_, -- *** Special folds - , foldrFB - , lazyFoldlS - , lazyFoldrS + foldrFB, + lazyFoldlS, + lazyFoldrS, -- ** Parallel folds - -- $par_folds - - , foldlP - , foldrP - , ifoldlP - , ifoldrP - , ifoldlIO - , ifoldrIO + foldlP, + foldrP, + ifoldlP, + ifoldrP, + ifoldlIO, + ifoldrIO, -- , splitReduce - ) where +) where import Data.Massiv.Array.Delayed.Pull import Data.Massiv.Array.Ops.Construct import Data.Massiv.Array.Ops.Fold.Internal import Data.Massiv.Core import Data.Massiv.Core.Common -import Prelude hiding (all, and, any, foldl, foldr, map, maximum, minimum, or, product, sum) +import Prelude hiding (all, and, any, elem, foldl, foldr, map, maximum, minimum, or, product, sum) -- | /O(n)/ - Monoidal fold over an array with an index aware function. Also known as reduce. -- -- @since 0.2.4 -ifoldMono :: - (Source r ix e, Monoid m) - => (ix -> e -> m) -- ^ Convert each element of an array to an appropriate `Monoid`. - -> Array r ix e -- ^ Source array +ifoldMono + :: (Index ix, Source r e, Monoid m) + => (ix -> e -> m) + -- ^ Convert each element of an array to an appropriate `Monoid`. + -> Array r ix e + -- ^ Source array -> m ifoldMono f = ifoldlInternal (\a ix e -> a `mappend` f ix e) mempty mappend mempty {-# INLINE ifoldMono #-} - -- | /O(n)/ - Semigroup fold over an array with an index aware function. -- -- @since 0.2.4 -ifoldSemi :: - (Source r ix e, Semigroup m) - => (ix -> e -> m) -- ^ Convert each element of an array to an appropriate `Semigroup`. - -> m -- ^ Initial element that must be neutral to the (`<>`) function. - -> Array r ix e -- ^ Source array +ifoldSemi + :: (Index ix, Source r e, Semigroup m) + => (ix -> e -> m) + -- ^ Convert each element of an array to an appropriate `Semigroup`. -> m + -- ^ Initial element that must be neutral to the (`<>`) function. + -> Array r ix e + -- ^ Source array + -> m ifoldSemi f m = ifoldlInternal (\a ix e -> a <> f ix e) m (<>) m {-# INLINE ifoldSemi #-} - -- | /O(n)/ - Semigroup fold over an array. -- -- @since 0.1.6 -foldSemi :: - (Source r ix e, Semigroup m) - => (e -> m) -- ^ Convert each element of an array to an appropriate `Semigroup`. - -> m -- ^ Initial element that must be neutral to the (`<>`) function. - -> Array r ix e -- ^ Source array +foldSemi + :: (Index ix, Source r e, Semigroup m) + => (e -> m) + -- ^ Convert each element of an array to an appropriate `Semigroup`. -> m + -- ^ Initial element that must be neutral to the (`<>`) function. + -> Array r ix e + -- ^ Source array + -> m foldSemi f m = foldlInternal (\a e -> a <> f e) m (<>) m {-# INLINE foldSemi #-} - -- | Left fold along a specified dimension with an index aware function. -- -- @since 0.2.4 -ifoldlWithin :: (Index (Lower ix), IsIndexDimension ix n, Source r ix e) => - Dimension n -> (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a +ifoldlWithin + :: (Index (Lower ix), IsIndexDimension ix n, Source r e) + => Dimension n + -> (ix -> a -> e -> a) + -> a + -> Array r ix e + -> Array D (Lower ix) a ifoldlWithin dim = ifoldlWithin' (fromDimension dim) {-# INLINE ifoldlWithin #-} - -- | Left fold along a specified dimension. -- -- ====__Example__ @@ -171,36 +181,53 @@ -- [ [5,0], [6,1], [7,2], [8,3], [9,4] ] -- -- @since 0.2.4 -foldlWithin :: (Index (Lower ix), IsIndexDimension ix n, Source r ix e) => - Dimension n -> (a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a +foldlWithin + :: (Index (Lower ix), IsIndexDimension ix n, Source r e) + => Dimension n + -> (a -> e -> a) + -> a + -> Array r ix e + -> Array D (Lower ix) a foldlWithin dim f = ifoldlWithin dim (const f) {-# INLINE foldlWithin #-} - -- | Right fold along a specified dimension with an index aware function. -- -- @since 0.2.4 -ifoldrWithin :: (Index (Lower ix), IsIndexDimension ix n, Source r ix e) => - Dimension n -> (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a +ifoldrWithin + :: (Index (Lower ix), IsIndexDimension ix n, Source r e) + => Dimension n + -> (ix -> e -> a -> a) + -> a + -> Array r ix e + -> Array D (Lower ix) a ifoldrWithin dim = ifoldrWithin' (fromDimension dim) {-# INLINE ifoldrWithin #-} - -- | Right fold along a specified dimension. -- -- @since 0.2.4 -foldrWithin :: (Index (Lower ix), IsIndexDimension ix n, Source r ix e) => - Dimension n -> (e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a +foldrWithin + :: (Index (Lower ix), IsIndexDimension ix n, Source r e) + => Dimension n + -> (e -> a -> a) + -> a + -> Array r ix e + -> Array D (Lower ix) a foldrWithin dim f = ifoldrWithin dim (const f) {-# INLINE foldrWithin #-} - -- | Similar to `ifoldlWithin`, except that dimension is specified at a value level, which means it -- will throw an exception on an invalid dimension. -- -- @since 0.2.4 -ifoldlWithin' :: (Index (Lower ix), Source r ix e) => - Dim -> (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a +ifoldlWithin' + :: (HasCallStack, Index (Lower ix), Index ix, Source r e) + => Dim + -> (ix -> a -> e -> a) + -> a + -> Array r ix e + -> Array D (Lower ix) a ifoldlWithin' dim f acc0 arr = makeArray (getComp arr) (SafeSz szl) $ \ixl -> iter @@ -215,24 +242,32 @@ (k, szl) = pullOutDim' sz dim {-# INLINE ifoldlWithin' #-} - -- | Similar to `foldlWithin`, except that dimension is specified at a value level, which means it will -- throw an exception on an invalid dimension. -- -- @since 0.2.4 -foldlWithin' :: (Index (Lower ix), Source r ix e) => - Dim -> (a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a +foldlWithin' + :: (HasCallStack, Index (Lower ix), Index ix, Source r e) + => Dim + -> (a -> e -> a) + -> a + -> Array r ix e + -> Array D (Lower ix) a foldlWithin' dim f = ifoldlWithin' dim (const f) {-# INLINE foldlWithin' #-} - -- | Similar to `ifoldrWithin`, except that dimension is specified at a value level, which means it -- will throw an exception on an invalid dimension. -- -- -- @since 0.2.4 -ifoldrWithin' :: (Index (Lower ix), Source r ix e) => - Dim -> (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a +ifoldrWithin' + :: (HasCallStack, Index (Lower ix), Index ix, Source r e) + => Dim + -> (ix -> e -> a -> a) + -> a + -> Array r ix e + -> Array D (Lower ix) a ifoldrWithin' dim f acc0 arr = makeArray (getComp arr) (SafeSz szl) $ \ixl -> iter @@ -251,56 +286,77 @@ -- will throw an exception on an invalid dimension. -- -- @since 0.2.4 -foldrWithin' :: (Index (Lower ix), Source r ix e) => - Dim -> (e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a +foldrWithin' + :: (HasCallStack, Index (Lower ix), Index ix, Source r e) + => Dim + -> (e -> a -> a) + -> a + -> Array r ix e + -> Array D (Lower ix) a foldrWithin' dim f = ifoldrWithin' dim (const f) {-# INLINE foldrWithin' #-} - -- | Left fold over the inner most dimension with index aware function. -- -- @since 0.2.4 -ifoldlInner :: (Index (Lower ix), Source r ix e) => - (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a +ifoldlInner + :: (Index (Lower ix), Index ix, Source r e) + => (ix -> a -> e -> a) + -> a + -> Array r ix e + -> Array D (Lower ix) a ifoldlInner = ifoldlWithin' 1 {-# INLINE ifoldlInner #-} -- | Left fold over the inner most dimension. -- -- @since 0.2.4 -foldlInner :: (Index (Lower ix), Source r ix e) => - (a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a +foldlInner + :: (Index (Lower ix), Index ix, Source r e) + => (a -> e -> a) + -> a + -> Array r ix e + -> Array D (Lower ix) a foldlInner = foldlWithin' 1 {-# INLINE foldlInner #-} -- | Right fold over the inner most dimension with index aware function. -- -- @since 0.2.4 -ifoldrInner :: (Index (Lower ix), Source r ix e) => - (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a +ifoldrInner + :: (Index (Lower ix), Index ix, Source r e) + => (ix -> e -> a -> a) + -> a + -> Array r ix e + -> Array D (Lower ix) a ifoldrInner = ifoldrWithin' 1 {-# INLINE ifoldrInner #-} -- | Right fold over the inner most dimension. -- -- @since 0.2.4 -foldrInner :: (Index (Lower ix), Source r ix e) => - (e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a +foldrInner + :: (Index (Lower ix), Index ix, Source r e) + => (e -> a -> a) + -> a + -> Array r ix e + -> Array D (Lower ix) a foldrInner = foldrWithin' 1 {-# INLINE foldrInner #-} -- | Monoidal fold over the inner most dimension. -- -- @since 0.4.3 -foldInner :: (Monoid e, Index (Lower ix), Source r ix e) => Array r ix e -> Array D (Lower ix) e +foldInner + :: (Monoid e, Index (Lower ix), Index ix, Source r e) => Array r ix e -> Array D (Lower ix) e foldInner = foldlInner mappend mempty {-# INLINE foldInner #-} -- | Monoidal fold over some internal dimension. -- -- @since 0.4.3 -foldWithin :: - (Source r ix a, Monoid a, Index (Lower ix), IsIndexDimension ix n) +foldWithin + :: (Source r a, Monoid a, Index (Lower ix), IsIndexDimension ix n) => Dimension n -> Array r ix a -> Array D (Lower ix) a @@ -311,15 +367,14 @@ -- result in `IndexDimensionException` if supplied dimension is invalid. -- -- @since 0.4.3 -foldWithin' :: - (Source r ix a, Monoid a, Index (Lower ix)) +foldWithin' + :: (HasCallStack, Index ix, Source r a, Monoid a, Index (Lower ix)) => Dim -> Array r ix a -> Array D (Lower ix) a foldWithin' dim = foldlWithin' dim mappend mempty {-# INLINE foldWithin' #-} - -- | Reduce each outer slice into a monoid and mappend results together -- -- ==== __Example__ @@ -338,23 +393,30 @@ -- 1620 -- -- @since 0.4.3 -foldOuterSlice :: (OuterSlice r ix e, Monoid m) => (Elt r ix e -> m) -> Array r ix e -> m +foldOuterSlice + :: (Index ix, Index (Lower ix), Source r e, Monoid m) + => (Array r (Lower ix) e -> m) + -> Array r ix e + -> m foldOuterSlice f = ifoldOuterSlice (const f) {-# INLINE foldOuterSlice #-} - -- | Reduce each outer slice into a monoid with an index aware function and mappend results -- together -- -- @since 0.4.3 -ifoldOuterSlice :: (OuterSlice r ix e, Monoid m) => (Ix1 -> Elt r ix e -> m) -> Array r ix e -> m -ifoldOuterSlice f arr = foldMono g $ range (getComp arr) 0 (headDim (unSz (size arr))) +ifoldOuterSlice + :: (Index ix, Index (Lower ix), Source r e, Monoid m) + => (Ix1 -> Array r (Lower ix) e -> m) + -> Array r ix e + -> m +ifoldOuterSlice f arr = foldMono g $ range (getComp arr) 0 k where - g i = f i (unsafeOuterSlice arr i) + (Sz1 k, szL) = unconsSz $ size arr + g i = f i (unsafeOuterSlice arr szL i) {-# INLINE g #-} {-# INLINE ifoldOuterSlice #-} - -- | Reduce each inner slice into a monoid and mappend results together -- -- ==== __Example__ @@ -373,67 +435,71 @@ -- 19575 -- -- @since 0.4.3 -foldInnerSlice :: (InnerSlice r ix e, Monoid m) => (Elt r ix e -> m) -> Array r ix e -> m +foldInnerSlice + :: (Source r e, Index ix, Monoid m) => (Array D (Lower ix) e -> m) -> Array r ix e -> m foldInnerSlice f = ifoldInnerSlice (const f) {-# INLINE foldInnerSlice #-} - -- | Reduce each inner slice into a monoid with an index aware function and mappend -- results together -- -- @since 0.4.3 -ifoldInnerSlice :: (InnerSlice r ix e, Monoid m) => (Ix1 -> Elt r ix e -> m) -> Array r ix e -> m +ifoldInnerSlice + :: (Source r e, Index ix, Monoid m) => (Ix1 -> Array D (Lower ix) e -> m) -> Array r ix e -> m ifoldInnerSlice f arr = foldMono g $ range (getComp arr) 0 (unSz k) where - szs@(_, !k) = unsnocSz (size arr) - g i = f i (unsafeInnerSlice arr szs i) + (szL, !k) = unsnocSz (size arr) + g i = f i (unsafeInnerSlice arr szL i) {-# INLINE g #-} {-# INLINE ifoldInnerSlice #-} -- | /O(n)/ - Compute maximum of all elements. -- -- @since 0.3.0 -maximumM :: (MonadThrow m, Source r ix e, Ord e) => Array r ix e -> m e +maximumM :: (MonadThrow m, Shape r ix, Source r e, Ord e) => Array r ix e -> m e maximumM arr = - if isEmpty arr - then throwM (SizeEmptyException (size arr)) - else let !e0 = unsafeIndex arr zeroIndex - in pure $ foldlInternal max e0 max e0 arr + if isNull arr + then throwM (SizeEmptyException (size arr)) + else + let !e0 = unsafeIndex arr zeroIndex + in pure $ foldlInternal max e0 max e0 arr {-# INLINE maximumM #-} - -- | /O(n)/ - Compute maximum of all elements. -- -- @since 0.3.0 -maximum' :: (Source r ix e, Ord e) => Array r ix e -> e -maximum' = either throw id . maximumM +maximum' + :: forall r ix e + . (HasCallStack, Shape r ix, Source r e, Ord e) + => Array r ix e + -> e +maximum' = throwEither . maximumM {-# INLINE maximum' #-} - -- | /O(n)/ - Compute minimum of all elements. -- -- @since 0.3.0 -minimumM :: (MonadThrow m, Source r ix e, Ord e) => Array r ix e -> m e +minimumM :: (MonadThrow m, Shape r ix, Source r e, Ord e) => Array r ix e -> m e minimumM arr = - if isEmpty arr - then throwM (SizeEmptyException (size arr)) - else let !e0 = unsafeIndex arr zeroIndex - in pure $ foldlInternal min e0 min e0 arr + if isNull arr + then throwM (SizeEmptyException (size arr)) + else + let !e0 = unsafeIndex arr zeroIndex + in pure $ foldlInternal min e0 min e0 arr {-# INLINE minimumM #-} -- | /O(n)/ - Compute minimum of all elements. -- -- @since 0.3.0 -minimum' :: (Source r ix e, Ord e) => Array r ix e -> e -minimum' = either throw id . minimumM +minimum' :: forall r ix e. (HasCallStack, Shape r ix, Source r e, Ord e) => Array r ix e -> e +minimum' = throwEither . minimumM {-# INLINE minimum' #-} - -- -- | /O(n)/ - Compute sum of all elements. -- -- -- -- @since 0.1.0 -- sum' :: --- forall r ix e. (Source r ix e, Numeric r e) +-- forall r ix e. (Index ix, Source r e, Numeric r e) -- => Array r ix e -- -> IO e -- sum' = splitReduce (\_ -> pure . sumArray) (\x y -> pure (x + y)) 0 @@ -442,79 +508,67 @@ -- | /O(n)/ - Compute sum of all elements. -- -- @since 0.1.0 -sum :: (Source r ix e, Num e) => Array r ix e -> e +sum :: (Index ix, Source r e, Num e) => Array r ix e -> e sum = foldlInternal (+) 0 (+) 0 {-# INLINE sum #-} - -- | /O(n)/ - Compute product of all elements. -- -- @since 0.1.0 -product :: (Source r ix e, Num e) => Array r ix e -> e +product :: (Index ix, Source r e, Num e) => Array r ix e -> e product = foldlInternal (*) 1 (*) 1 {-# INLINE product #-} - -- | /O(n)/ - Compute conjunction of all elements. -- -- @since 0.1.0 -and :: Source r ix Bool => Array r ix Bool -> Bool -and = foldlInternal (&&) True (&&) True +and :: (Index ix, Source r Bool) => Array r ix Bool -> Bool +and = all id {-# INLINE and #-} - -- | /O(n)/ - Compute disjunction of all elements. -- -- @since 0.1.0 -or :: Source r ix Bool => Array r ix Bool -> Bool -or = foldlInternal (||) False (||) False +or :: (Index ix, Source r Bool) => Array r ix Bool -> Bool +or = any id {-# INLINE or #-} - --- | /O(n)/ - Determines whether all element of the array satisfy the predicate. +-- | /O(n)/ - Determines whether all elements of the array satisfy a predicate. -- -- @since 0.1.0 -all :: Source r ix e => (e -> Bool) -> Array r ix e -> Bool -all f = foldlInternal (\acc e -> acc && f e) True (&&) True +all :: (Index ix, Source r e) => (e -> Bool) -> Array r ix e -> Bool +all f = not . any (not . f) {-# INLINE all #-} --- | /O(n)/ - Determines whether any element of the array satisfies the predicate. +-- | /O(n)/ - Determines whether an element is present in the array. -- --- @since 0.1.0 -any :: Source r ix e => (e -> Bool) -> Array r ix e -> Bool -any f = foldlInternal (\acc e -> acc || f e) False (||) False -{-# INLINE any #-} - - -{- $unstruct_folds - -Functions in this section will fold any `Source` array with respect to the inner -`Comp`utation strategy setting. - --} - - -{- $seq_folds - -Functions in this section will fold any `Source` array sequentially, regardless of the inner -`Comp`utation strategy setting. - --} - - -{- $par_folds - -__Note__ It is important to compile with @-threaded -with-rtsopts=-N@ flags, otherwise there will be -no parallelization. +-- @since 0.5.5 +elem :: (Eq e, Index ix, Source r e) => e -> Array r ix e -> Bool +elem e = any (e ==) +{-# INLINE elem #-} -Functions in this section will fold any `Source` array in parallel, regardless of the inner -`Comp`utation strategy setting. All of the parallel structured folds are performed inside `IO` -monad, because referential transparency can't generally be preserved and results will depend on the -number of cores/capabilities that computation is being performed on. +-- $unstruct_folds +-- +-- Functions in this section will fold any `Source` array with respect to the inner +-- `Comp`utation strategy setting. -In contrast to sequential folds, each parallel folding function accepts two functions and two -initial elements as arguments. This is necessary because an array is first split into chunks, which -folded individually on separate cores with the first function, and the results of those folds are -further folded with the second function. +-- $seq_folds +-- +-- Functions in this section will fold any `Source` array sequentially, regardless of the inner +-- `Comp`utation strategy setting. --} +-- $par_folds +-- +-- __Note__ It is important to compile with @-threaded -with-rtsopts=-N@ flags, otherwise +-- there will be no parallelization. +-- +-- Functions in this section will fold any `Source` array in parallel, regardless of the +-- inner `Comp`utation strategy setting. All of the parallel structured folds are +-- performed inside `IO` monad, because referential transparency can't generally be +-- preserved and results will depend on the number of cores/capabilities that computation +-- is being performed on. +-- +-- In contrast to sequential folds, each parallel folding function accepts two functions +-- and two initial elements as arguments. This is necessary because an array is first +-- split into chunks, which folded individually on separate cores with the first function, +-- and the results of those folds are further folded with the second function. diff --git a/src/Data/Massiv/Array/Ops/Fold/Internal.hs b/src/Data/Massiv/Array/Ops/Fold/Internal.hs --- a/src/Data/Massiv/Array/Ops/Fold/Internal.hs +++ b/src/Data/Massiv/Array/Ops/Fold/Internal.hs @@ -2,216 +2,228 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UndecidableInstances #-} + -- | -- Module : Data.Massiv.Array.Ops.Fold.Internal --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Ops.Fold.Internal - ( - foldlS - , foldrS - , ifoldlS - , ifoldrS - --Monadic - , foldlM - , foldrM - , foldlM_ - , foldrM_ - , ifoldlM - , ifoldrM - , ifoldlM_ - , ifoldrM_ - --Special folds - , fold - , foldMono - , foldlInternal - , ifoldlInternal - , foldrFB - , lazyFoldlS - , lazyFoldrS +module Data.Massiv.Array.Ops.Fold.Internal ( + foldlS, + foldrS, + ifoldlS, + ifoldrS, + -- Monadic + foldlM, + foldrM, + foldlM_, + foldrM_, + ifoldlM, + ifoldrM, + ifoldlM_, + ifoldrM_, + -- Special folds + fold, + foldMono, + foldlInternal, + ifoldlInternal, + foldrFB, + lazyFoldlS, + lazyFoldrS, -- Parallel folds - , foldlP - , foldrP - , ifoldlP - , ifoldrP - , ifoldlIO - , ifoldrIO - -- , splitReduce - ) where + foldlP, + foldrP, + ifoldlP, + ifoldrP, + foldlIO, + ifoldlIO, + ifoldrIO, + splitReduce, + any, + anySu, + anyPu, +) where import Control.Monad (void, when) +import Control.Monad.Primitive import Control.Scheduler import qualified Data.Foldable as F import Data.Functor.Identity (runIdentity) import Data.Massiv.Core.Common -import Prelude hiding (foldl, foldr) import System.IO.Unsafe (unsafePerformIO) - +import Prelude hiding (any, foldl, foldr) -- | /O(n)/ - Unstructured fold of an array. -- -- @since 0.3.0 -fold :: - (Monoid e, Source r ix e) - => Array r ix e -- ^ Source array +fold + :: (Monoid e, Index ix, Source r e) + => Array r ix e + -- ^ Source array -> e fold = foldlInternal mappend mempty mappend mempty {-# INLINE fold #-} - -- | /O(n)/ - This is exactly like `Data.Foldable.foldMap`, but for arrays. Fold over an array, -- while converting each element into a `Monoid`. Also known as map-reduce. If elements of the array -- are already a `Monoid` you can use `fold` instead. -- -- @since 0.1.4 -foldMono :: - (Source r ix e, Monoid m) - => (e -> m) -- ^ Convert each element of an array to an appropriate `Monoid`. - -> Array r ix e -- ^ Source array +foldMono + :: (Index ix, Source r e, Monoid m) + => (e -> m) + -- ^ Convert each element of an array to an appropriate `Monoid`. + -> Array r ix e + -- ^ Source array -> m foldMono f = foldlInternal (\a e -> a `mappend` f e) mempty mappend mempty {-# INLINE foldMono #-} - -- | /O(n)/ - Monadic left fold. -- -- @since 0.1.0 -foldlM :: (Source r ix e, Monad m) => (a -> e -> m a) -> a -> Array r ix e -> m a -foldlM f = ifoldlM (\ a _ b -> f a b) +foldlM :: (Index ix, Source r e, Monad m) => (a -> e -> m a) -> a -> Array r ix e -> m a +foldlM f acc arr = + case unsafePrefIndex arr of + PrefIndex gix -> + iterM zeroIndex (unSz sz) (pureIndex 1) (<) acc $ \ !ix !a -> f a (gix ix) + PrefIndexLinear gi -> + loopM 0 (< totalElem sz) (+ 1) acc $ \ !i !a -> f a (gi i) + where + sz = size arr {-# INLINE foldlM #-} - -- | /O(n)/ - Monadic left fold, that discards the result. -- -- @since 0.1.0 -foldlM_ :: (Source r ix e, Monad m) => (a -> e -> m a) -> a -> Array r ix e -> m () -foldlM_ f = ifoldlM_ (\ a _ b -> f a b) +foldlM_ :: (Index ix, Source r e, Monad m) => (a -> e -> m a) -> a -> Array r ix e -> m () +foldlM_ f acc = void . foldlM f acc {-# INLINE foldlM_ #-} - -- | /O(n)/ - Monadic left fold with an index aware function. -- -- @since 0.1.0 -ifoldlM :: (Source r ix e, Monad m) => (a -> ix -> e -> m a) -> a -> Array r ix e -> m a +ifoldlM :: (Index ix, Source r e, Monad m) => (a -> ix -> e -> m a) -> a -> Array r ix e -> m a ifoldlM f !acc !arr = - iterM zeroIndex (unSz (size arr)) (pureIndex 1) (<) acc $ \ !ix !a -> f a ix (unsafeIndex arr ix) + case unsafePrefIndex arr of + PrefIndex gix -> + iterM zeroIndex (unSz (size arr)) (pureIndex 1) (<) acc $ \ !ix !a -> f a ix (gix ix) + PrefIndexLinear gi -> + iterTargetM defRowMajor 0 (size arr) zeroIndex oneStride acc $ \i ix !a -> f a ix (gi i) {-# INLINE ifoldlM #-} - -- | /O(n)/ - Monadic left fold with an index aware function, that discards the result. -- -- @since 0.1.0 -ifoldlM_ :: (Source r ix e, Monad m) => (a -> ix -> e -> m a) -> a -> Array r ix e -> m () +ifoldlM_ :: (Index ix, Source r e, Monad m) => (a -> ix -> e -> m a) -> a -> Array r ix e -> m () ifoldlM_ f acc = void . ifoldlM f acc {-# INLINE ifoldlM_ #-} - -- | /O(n)/ - Monadic right fold. -- -- @since 0.1.0 -foldrM :: (Source r ix e, Monad m) => (e -> a -> m a) -> a -> Array r ix e -> m a -foldrM f = ifoldrM (\_ e a -> f e a) +foldrM :: (Index ix, Source r e, Monad m) => (e -> a -> m a) -> a -> Array r ix e -> m a +foldrM f acc arr = + case unsafePrefIndex arr of + PrefIndex gix -> + iterM (liftIndex (subtract 1) (unSz sz)) zeroIndex (pureIndex (-1)) (>=) acc (f . gix) + PrefIndexLinear gi -> + loopM (totalElem sz - 1) (>= 0) (subtract 1) acc (f . gi) + where + !sz = size arr {-# INLINE foldrM #-} - -- | /O(n)/ - Monadic right fold, that discards the result. -- -- @since 0.1.0 -foldrM_ :: (Source r ix e, Monad m) => (e -> a -> m a) -> a -> Array r ix e -> m () +foldrM_ :: (Index ix, Source r e, Monad m) => (e -> a -> m a) -> a -> Array r ix e -> m () foldrM_ f = ifoldrM_ (\_ e a -> f e a) {-# INLINE foldrM_ #-} - -- | /O(n)/ - Monadic right fold with an index aware function. -- -- @since 0.1.0 -ifoldrM :: (Source r ix e, Monad m) => (ix -> e -> a -> m a) -> a -> Array r ix e -> m a +ifoldrM :: (Index ix, Source r e, Monad m) => (ix -> e -> a -> m a) -> a -> Array r ix e -> m a ifoldrM f !acc !arr = - iterM (liftIndex (subtract 1) (unSz (size arr))) zeroIndex (pureIndex (-1)) (>=) acc $ \ !ix !acc0 -> - f ix (unsafeIndex arr ix) acc0 + iterM (liftIndex (subtract 1) (unSz (size arr))) zeroIndex (pureIndex (-1)) (>=) acc $ \ !ix -> + f ix (unsafeIndex arr ix) {-# INLINE ifoldrM #-} - -- | /O(n)/ - Monadic right fold with an index aware function, that discards the result. -- -- @since 0.1.0 -ifoldrM_ :: (Source r ix e, Monad m) => (ix -> e -> a -> m a) -> a -> Array r ix e -> m () +ifoldrM_ :: (Index ix, Source r e, Monad m) => (ix -> e -> a -> m a) -> a -> Array r ix e -> m () ifoldrM_ f !acc !arr = void $ ifoldrM f acc arr {-# INLINE ifoldrM_ #-} - - -- | /O(n)/ - Left fold, computed sequentially with lazy accumulator. -- -- @since 0.1.0 -lazyFoldlS :: Source r ix e => (a -> e -> a) -> a -> Array r ix e -> a -lazyFoldlS f initAcc arr = go initAcc 0 where +lazyFoldlS :: (Index ix, Source r e) => (a -> e -> a) -> a -> Array r ix e -> a +lazyFoldlS f initAcc arr = go initAcc 0 + where len = totalElem (size arr) - go acc k | k < len = go (f acc (unsafeLinearIndex arr k)) (k + 1) - | otherwise = acc + go acc !k + | k < len = go (f acc (unsafeLinearIndex arr k)) (k + 1) + | otherwise = acc {-# INLINE lazyFoldlS #-} - -- | /O(n)/ - Right fold, computed sequentially with lazy accumulator. -- -- @since 0.1.0 -lazyFoldrS :: Source r ix e => (e -> a -> a) -> a -> Array r ix e -> a +lazyFoldrS :: (Index ix, Source r e) => (e -> a -> a) -> a -> Array r ix e -> a lazyFoldrS = foldrFB {-# INLINE lazyFoldrS #-} - -- | /O(n)/ - Left fold, computed sequentially. -- -- @since 0.1.0 -foldlS :: Source r ix e => (a -> e -> a) -> a -> Array r ix e -> a -foldlS f = ifoldlS (\ a _ e -> f a e) +foldlS :: (Index ix, Source r e) => (a -> e -> a) -> a -> Array r ix e -> a +foldlS f acc = runIdentity . foldlM (\a e -> pure $! f a e) acc {-# INLINE foldlS #-} - -- | /O(n)/ - Left fold with an index aware function, computed sequentially. -- -- @since 0.1.0 -ifoldlS :: Source r ix e - => (a -> ix -> e -> a) -> a -> Array r ix e -> a -ifoldlS f acc = runIdentity . ifoldlM (\ a ix e -> return $ f a ix e) acc +ifoldlS + :: (Index ix, Source r e) + => (a -> ix -> e -> a) + -> a + -> Array r ix e + -> a +ifoldlS f acc = runIdentity . ifoldlM (\a ix e -> pure $! f a ix e) acc {-# INLINE ifoldlS #-} - -- | /O(n)/ - Right fold, computed sequentially. -- -- @since 0.1.0 -foldrS :: Source r ix e => (e -> a -> a) -> a -> Array r ix e -> a -foldrS f = ifoldrS (\_ e a -> f e a) +foldrS :: (Index ix, Source r e) => (e -> a -> a) -> a -> Array r ix e -> a +foldrS f acc = runIdentity . foldrM (\e a -> pure $! f e a) acc {-# INLINE foldrS #-} - -- | /O(n)/ - Right fold with an index aware function, computed sequentially. -- -- @since 0.1.0 -ifoldrS :: Source r ix e => (ix -> e -> a -> a) -> a -> Array r ix e -> a -ifoldrS f acc = runIdentity . ifoldrM (\ ix e a -> return $ f ix e a) acc +ifoldrS :: (Index ix, Source r e) => (ix -> e -> a -> a) -> a -> Array r ix e -> a +ifoldrS f acc = runIdentity . ifoldrM (\ix e a -> pure $! f ix e a) acc {-# INLINE ifoldrS #-} - -- | Version of foldr that supports @foldr/build@ list fusion implemented by GHC. -- -- @since 0.1.0 -foldrFB :: Source r ix e => (e -> b -> b) -> b -> Array r ix e -> b +foldrFB :: (Index ix, Source r e) => (e -> b -> b) -> b -> Array r ix e -> b foldrFB c n arr = go 0 where !k = totalElem (size arr) go !i | i == k = n - | otherwise = let !v = unsafeLinearIndex arr i in v `c` go (i + 1) + | otherwise = let v = unsafeLinearIndex arr i in v `c` go (i + 1) {-# INLINE [0] foldrFB #-} - - -- | /O(n)/ - Left fold, computed with respect of array's computation strategy. Because we do -- potentially split the folding among many threads, we also need a combining function and an -- accumulator for the results. Depending on the number of threads being used, results can be @@ -230,13 +242,20 @@ -- [1,0,3,2,5,4] -- -- @since 0.1.0 -foldlP :: (MonadIO m, Source r ix e) => - (a -> e -> a) -- ^ Folding function @g@. - -> a -- ^ Accumulator. Will be applied to @g@ multiple times, thus must be neutral. - -> (b -> a -> b) -- ^ Chunk results folding function @f@. - -> b -- ^ Accumulator for results of chunks folding. - -> Array r ix e -> m b -foldlP f fAcc g gAcc = liftIO . ifoldlP (\ x _ -> f x) fAcc g gAcc +foldlP + :: (MonadIO m, Index ix, Source r e) + => (a -> e -> a) + -- ^ Folding function @g@. + -> a + -- ^ Accumulator. Will be applied to @g@ multiple times, thus must be neutral. + -> (b -> a -> b) + -- ^ Chunk results folding function @f@. + -> b + -- ^ Accumulator for results of chunks folding. + -> Array r ix e + -> m b +foldlP f fAcc g gAcc = + liftIO . foldlIO (\acc -> pure . f acc) fAcc (\acc -> pure . g acc) gAcc {-# INLINE foldlP #-} -- | /O(n)/ - Left fold with an index aware function, computed in parallel. Just @@ -244,13 +263,18 @@ -- element it is being applied to. -- -- @since 0.1.0 -ifoldlP :: (MonadIO m, Source r ix e) => - (a -> ix -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> m b +ifoldlP + :: (MonadIO m, Index ix, Source r e) + => (a -> ix -> e -> a) + -> a + -> (b -> a -> b) + -> b + -> Array r ix e + -> m b ifoldlP f fAcc g gAcc = - liftIO . ifoldlIO (\acc ix -> return . f acc ix) fAcc (\acc -> return . g acc) gAcc + liftIO . ifoldlIO (\acc ix -> pure . f acc ix) fAcc (\acc -> pure . g acc) gAcc {-# INLINE ifoldlP #-} - -- | /O(n)/ - Right fold, computed with respect to computation strategy. Same as `foldlP`, except -- directed from the last element in the array towards beginning. -- @@ -265,19 +289,24 @@ -- [[0,1],[2,3],[4,5]] -- -- @since 0.1.0 -foldrP :: (MonadIO m, Source r ix e) => - (e -> a -> a) -> a -> (a -> b -> b) -> b -> Array r ix e -> m b +foldrP + :: (MonadIO m, Index ix, Source r e) + => (e -> a -> a) + -> a + -> (a -> b -> b) + -> b + -> Array r ix e + -> m b foldrP f fAcc g gAcc = liftIO . ifoldrP (const f) fAcc g gAcc {-# INLINE foldrP #-} - -- | /O(n)/ - Right fold with an index aware function, while respecting the computation strategy. -- Same as `ifoldlP`, except directed from the last element in the array towards -- beginning, but also row-major. -- -- @since 0.1.0 -ifoldrP :: - (MonadIO m, Source r ix e) +ifoldrP + :: (MonadIO m, Index ix, Source r e) => (ix -> e -> a -> a) -> a -> (a -> b -> b) @@ -287,93 +316,230 @@ ifoldrP f fAcc g gAcc = liftIO . ifoldrIO (\ix e -> pure . f ix e) fAcc (\e -> pure . g e) gAcc {-# INLINE ifoldrP #-} - -- | This folding function breaks referential transparency on some functions -- @f@, therefore it is kept here for internal use only. -foldlInternal :: Source r ix e => (a -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> b +foldlInternal + :: (Index ix, Source r e) => (a -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> b foldlInternal g initAcc f resAcc = unsafePerformIO . foldlP g initAcc f resAcc {-# INLINE foldlInternal #-} - -ifoldlInternal :: Source r ix e => (a -> ix -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> b +ifoldlInternal + :: (Index ix, Source r e) => (a -> ix -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> b ifoldlInternal g initAcc f resAcc = unsafePerformIO . ifoldlP g initAcc f resAcc {-# INLINE ifoldlInternal #-} +-- | Similar to `foldlP`, except that folding functions themselves do live in IO +-- +-- @since 0.1.0 +foldlIO + :: (MonadUnliftIO m, Index ix, Source r e) + => (a -> e -> m a) + -- ^ Index aware folding IO action + -> a + -- ^ Accumulator + -> (b -> a -> m b) + -- ^ Folding action that is applied to the results of a parallel fold + -> b + -- ^ Accumulator for chunks folding + -> Array r ix e + -> m b +foldlIO f !initAcc g !tAcc !arr + | getComp arr == Seq = foldlM f initAcc arr >>= g tAcc + | otherwise = do + let splitAcc _ = pure (initAcc, initAcc) + !sz = size arr + results <- + withScheduler (getComp arr) $ \scheduler -> + withRunInIO $ \run -> + stToPrim $ + case unsafePrefIndex arr of + PrefIndex gix -> + iterFullAccST defRowMajor scheduler zeroIndex sz initAcc splitAcc $ \ !ix !acc -> + ioToPrim (run (f acc (gix ix))) + PrefIndexLinear gi -> + iterFullAccST defRowMajor scheduler 0 (toLinearSz sz) initAcc splitAcc $ \ !i !acc -> + ioToPrim (run (f acc (gi i))) + F.foldlM g tAcc results +{-# INLINE foldlIO #-} -- | Similar to `ifoldlP`, except that folding functions themselves do live in IO -- -- @since 0.1.0 -ifoldlIO :: - (MonadUnliftIO m, Source r ix e) - => (a -> ix -> e -> m a) -- ^ Index aware folding IO action - -> a -- ^ Accumulator - -> (b -> a -> m b) -- ^ Folding action that is applied to the results of a parallel fold - -> b -- ^ Accumulator for chunks folding +ifoldlIO + :: (MonadUnliftIO m, Index ix, Source r e) + => (a -> ix -> e -> m a) + -- ^ Index aware folding IO action + -> a + -- ^ Accumulator + -> (b -> a -> m b) + -- ^ Folding action that is applied to the results of a parallel fold + -> b + -- ^ Accumulator for chunks folding -> Array r ix e -> m b -ifoldlIO f !initAcc g !tAcc !arr = do +ifoldlIO f !initAcc g !tAcc !arr + | getComp arr == Seq = ifoldlM f initAcc arr >>= g tAcc + | otherwise = do + let !sz = size arr + splitAcc _ = pure (initAcc, initAcc) + results <- + withScheduler (getComp arr) $ \scheduler -> + withRunInIO $ \run -> + stToPrim $ + case unsafePrefIndex arr of + PrefIndex gix -> + iterFullAccST defRowMajor scheduler zeroIndex sz initAcc splitAcc $ \ !ix !acc -> + ioToPrim (run (f acc ix (gix ix))) + PrefIndexLinear gi -> + iterTargetFullAccST defRowMajor scheduler 0 sz initAcc splitAcc $ \ !i !ix !acc -> + ioToPrim (run (f acc ix (gi i))) + F.foldlM g tAcc results +{-# INLINE ifoldlIO #-} + +-- | Slice an array into linear row-major vector chunks and apply an action to each of +-- them. Number of chunks will depend on the computation strategy. Results of each action +-- will be combined with a folding function. +-- +-- @since 1.0.0 +splitReduce + :: (MonadUnliftIO m, Index ix, Source r e) + => (Scheduler RealWorld a -> Vector r e -> m a) + -> (b -> a -> m b) + -- ^ Folding action that is applied to the results of a parallel fold + -> b + -- ^ Accumulator for chunks folding + -> Array r ix e + -> m b +splitReduce f g !tAcc !arr = do let !sz = size arr !totalLength = totalElem sz results <- - withScheduler (getComp arr) $ \scheduler -> - splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do - loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start -> - scheduleWork scheduler $ - iterLinearM sz start (start + chunkLength) 1 (<) initAcc $ \ !i ix !acc -> - f acc ix (unsafeLinearIndex arr i) - when (slackStart < totalLength) $ - scheduleWork scheduler $ - iterLinearM sz slackStart totalLength 1 (<) initAcc $ \ !i ix !acc -> - f acc ix (unsafeLinearIndex arr i) + withScheduler (getComp arr) $ \scheduler -> do + withRunInIO $ \run -> do + splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do + loopA_ 0 (< slackStart) (+ chunkLength) $ \ !start -> + scheduleWork scheduler $ + run $ + f scheduler $ + unsafeLinearSlice start (SafeSz chunkLength) arr + when (slackStart < totalLength) $ + scheduleWork scheduler $ + run $ + f scheduler $ + unsafeLinearSlice slackStart (SafeSz (totalLength - slackStart)) arr F.foldlM g tAcc results -{-# INLINE ifoldlIO #-} - --- -- | Split an array into linear row-major vector chunks and apply an action to each of --- -- them. Number of chunks will depend on the computation strategy. Results of each action --- -- will be combined with a folding function. --- -- --- -- @since 0.4.1 --- splitReduce :: --- (MonadUnliftIO m, Source r ix e) --- => (Scheduler m a -> Array r Ix1 e -> m a) --- -> (b -> a -> m b) -- ^ Folding action that is applied to the results of a parallel fold --- -> b -- ^ Accumulator for chunks folding --- -> Array r ix e --- -> m b --- splitReduce f g !tAcc !arr = do --- let !sz = size arr --- !totalLength = totalElem sz --- results <- --- withScheduler (getComp arr) $ \scheduler -> --- splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do --- loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start -> --- scheduleWork scheduler $ f scheduler $ unsafeLinearSlice start (SafeSz chunkLength) arr --- when (slackStart < totalLength) $ --- scheduleWork scheduler $ --- f scheduler $ unsafeLinearSlice slackStart (SafeSz (totalLength - slackStart)) arr --- F.foldlM g tAcc results --- {-# INLINE splitReduce #-} - - +{-# INLINE splitReduce #-} -- | Similar to `ifoldrP`, except that folding functions themselves do live in IO -- -- @since 0.1.0 -ifoldrIO :: (MonadUnliftIO m, Source r ix e) => - (ix -> e -> a -> m a) -> a -> (a -> b -> m b) -> b -> Array r ix e -> m b -ifoldrIO f !initAcc g !tAcc !arr = do +ifoldrIO + :: (MonadUnliftIO m, Index ix, Source r e) + => (ix -> e -> a -> m a) + -> a + -> (a -> b -> m b) + -> b + -> Array r ix e + -> m b +ifoldrIO f !initAcc g !tAcc !arr + | getComp arr == Seq = ifoldrM f initAcc arr >>= (`g` tAcc) + | otherwise = do + let !sz = size arr + !totalLength = totalElem sz + results <- + withRunInIO $ \run -> do + withScheduler (getComp arr) $ \scheduler -> + splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do + when (slackStart < totalLength) $ + scheduleWork scheduler $ + run $ + iterLinearM sz (totalLength - 1) slackStart (-1) (>=) initAcc $ \ !i ix -> + f ix (unsafeLinearIndex arr i) + loopA_ slackStart (> 0) (subtract chunkLength) $ \ !start -> + scheduleWork scheduler $ + run $ + iterLinearM sz (start - 1) (start - chunkLength) (-1) (>=) initAcc $ \ !i ix -> + f ix (unsafeLinearIndex arr i) + F.foldlM (flip g) tAcc results +{-# INLINE ifoldrIO #-} + +-- | Sequential implementation of `any` with unrolling +anySu :: (Index ix, Source r e) => (e -> Bool) -> Array r ix e -> Bool +anySu f arr = go 0 + where + !k = elemsCount arr + !k4 = k - (k `rem` 4) + go !i + | i < k4 = + f (unsafeLinearIndex arr i) + || f (unsafeLinearIndex arr (i + 1)) + || f (unsafeLinearIndex arr (i + 2)) + || f (unsafeLinearIndex arr (i + 3)) + || go (i + 4) + | i < k = f (unsafeLinearIndex arr i) || go (i + 1) + | otherwise = False +{-# INLINE anySu #-} + +-- | Implementaton of `any` on a slice of an array with short-circuiting using batch cancellation. +anySliceSuM + :: (Index ix, Source r e) + => Batch RealWorld Bool + -> Ix1 + -> Sz1 + -> (e -> Bool) + -> Array r ix e + -> IO Bool +anySliceSuM batch ix0 (Sz1 k) f arr = go ix0 + where + !k' = k - ix0 + !k4 = ix0 + (k' - (k' `rem` 4)) + go !i + | i < k4 = do + let r = + f (unsafeLinearIndex arr i) + || f (unsafeLinearIndex arr (i + 1)) + || f (unsafeLinearIndex arr (i + 2)) + || f (unsafeLinearIndex arr (i + 3)) + in if r + then cancelBatchWith batch True + else do + done <- hasBatchFinished batch + if done + then pure True + else go (i + 4) + | i < k = + if f (unsafeLinearIndex arr i) + then cancelBatchWith batch True + else go (i + 1) + | otherwise = pure False +{-# INLINE anySliceSuM #-} + +-- | Parallelizable implementation of `any` with unrolling +anyPu :: (Index ix, Source r e) => (e -> Bool) -> Array r ix e -> IO Bool +-- TODO: switch to splitReduce +-- anyPu f arr = +-- splitReduce anySu (\r acc -> pure (r || acc)) False +anyPu f arr = do let !sz = size arr !totalLength = totalElem sz results <- - withScheduler (getComp arr) $ \ scheduler -> - splitLinearly (numWorkers scheduler) totalLength $ \ chunkLength slackStart -> do + withScheduler (getComp arr) $ \scheduler -> do + batch <- getCurrentBatch scheduler + splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do + loopA_ 0 (< slackStart) (+ chunkLength) $ \ !start -> + scheduleWork scheduler $ anySliceSuM batch start (Sz (start + chunkLength)) f arr when (slackStart < totalLength) $ scheduleWork scheduler $ - iterLinearM sz (totalLength - 1) slackStart (-1) (>=) initAcc $ \ !i ix !acc -> - f ix (unsafeLinearIndex arr i) acc - loopM_ slackStart (> 0) (subtract chunkLength) $ \ !start -> - scheduleWork scheduler $ - iterLinearM sz (start - 1) (start - chunkLength) (-1) (>=) initAcc $ \ !i ix !acc -> - f ix (unsafeLinearIndex arr i) acc - F.foldlM (flip g) tAcc results -{-# INLINE ifoldrIO #-} + anySliceSuM batch slackStart (Sz totalLength) f arr + pure $ F.foldl' (||) False results +{-# INLINE anyPu #-} + +-- | /O(n)/ - Determines whether any element of the array satisfies a predicate. +-- +-- @since 0.1.0 +any :: (Index ix, Source r e) => (e -> Bool) -> Array r ix e -> Bool +any f arr = + case getComp arr of + Seq -> anySu f arr + _ -> unsafePerformIO $ anyPu f arr +{-# INLINE any #-} diff --git a/src/Data/Massiv/Array/Ops/Map.hs b/src/Data/Massiv/Array/Ops/Map.hs --- a/src/Data/Massiv/Array/Ops/Map.hs +++ b/src/Data/Massiv/Array/Ops/Map.hs @@ -1,167 +1,256 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MonoLocalBinds #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE ScopedTypeVariables #-} + -- | -- Module : Data.Massiv.Array.Ops.Map --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Ops.Map - ( map - , imap +module Data.Massiv.Array.Ops.Map ( + map, + imap, + -- ** Traversing + -- *** Applicative - , traverseA - , traverseA_ - , itraverseA - , itraverseA_ - , traverseAR - , itraverseAR - , sequenceA - , sequenceA_ - , traverseS + traverseA, + traverseA_, + itraverseA, + itraverseA_, + sequenceA, + sequenceA_, + -- *** PrimMonad - , traversePrim - , itraversePrim - , traversePrimR - , itraversePrimR + traversePrim, + itraversePrim, + -- ** Monadic mapping + -- *** Sequential - , mapM - , mapMR - , forM - , forMR - , imapM - , imapMR - , iforM - , iforMR - , mapM_ - , forM_ - , imapM_ - , iforM_ + mapM, + forM, + imapM, + iforM, + mapM_, + forM_, + imapM_, + iforM_, + -- *** Parallelizable - , mapIO - , mapWS - , mapIO_ - , imapIO - , imapWS - , imapIO_ - , forIO - , forWS - , forIO_ - , iforIO - , iforWS - , iforIO_ - , imapSchedulerM_ - , iforSchedulerM_ + mapIO, + mapWS, + mapIO_, + imapIO, + imapWS, + imapIO_, + forIO, + forWS, + forIO_, + iforIO, + iforWS, + iforIO_, + imapSchedulerM_, + iforSchedulerM_, + iterArrayLinearM_, + iterArrayLinearWithSetM_, + iterArrayLinearWithStrideM_, + -- ** Zipping - , zip - , zip3 - , unzip - , unzip3 - , zipWith - , zipWith3 - , izipWith - , izipWith3 - , liftArray2 + zip, + zip3, + zip4, + unzip, + unzip3, + unzip4, + zipWith, + zipWith3, + zipWith4, + izipWith, + izipWith3, + izipWith4, + -- *** Applicative - , zipWithA - , izipWithA - , zipWith3A - , izipWith3A - ) where + zipWithA, + izipWithA, + zipWith3A, + izipWith3A, +) where import Control.Monad (void) -import Control.Monad.Primitive (PrimMonad) +import Control.Monad.Primitive import Control.Scheduler import Data.Coerce import Data.Massiv.Array.Delayed.Pull -import Data.Massiv.Array.Delayed.Stream +import Data.Massiv.Array.Manifest.List import Data.Massiv.Array.Mutable import Data.Massiv.Array.Ops.Construct (makeArrayA, makeArrayLinearA) import Data.Massiv.Core.Common -import Data.Massiv.Core.Index.Internal (Sz(..)) -import Prelude hiding (map, mapM, mapM_, sequenceA, traverse, unzip, unzip3, - zip, zip3, zipWith, zipWith3) +import Data.Traversable (traverse) +import Prelude hiding ( + map, + mapM, + mapM_, + sequenceA, + traverse, + unzip, + unzip3, + zip, + zip3, + zipWith, + zipWith3, + ) -------------------------------------------------------------------------------- -- map ------------------------------------------------------------------------- -------------------------------------------------------------------------------- -- | Map a function over an array -map :: Source r ix e' => (e' -> e) -> Array r ix e' -> Array D ix e -map f = imap (const f) +-- +-- @since 0.1.0 +map :: (Index ix, Source r e') => (e' -> e) -> Array r ix e' -> Array D ix e +map f = fmap f . delay {-# INLINE map #-} --- | Map an index aware function over an array -imap :: Source r ix e' => (ix -> e' -> e) -> Array r ix e' -> Array D ix e -imap f !arr = DArray (getComp arr) (size arr) (\ !ix -> f ix (unsafeIndex arr ix)) -{-# INLINE imap #-} - -------------------------------------------------------------------------------- -- zip ------------------------------------------------------------------------- -------------------------------------------------------------------------------- -- | Zip two arrays -zip :: (Source r1 ix e1, Source r2 ix e2) - => Array r1 ix e1 -> Array r2 ix e2 -> Array D ix (e1, e2) +-- +-- @since 0.1.0 +zip + :: (Index ix, Source r1 e1, Source r2 e2) + => Array r1 ix e1 + -> Array r2 ix e2 + -> Array D ix (e1, e2) zip = zipWith (,) {-# INLINE zip #-} -- | Zip three arrays -zip3 :: (Source r1 ix e1, Source r2 ix e2, Source r3 ix e3) - => Array r1 ix e1 -> Array r2 ix e2 -> Array r3 ix e3 -> Array D ix (e1, e2, e3) +-- +-- @since 0.1.0 +zip3 + :: (Index ix, Source r1 e1, Source r2 e2, Source r3 e3) + => Array r1 ix e1 + -> Array r2 ix e2 + -> Array r3 ix e3 + -> Array D ix (e1, e2, e3) zip3 = zipWith3 (,,) {-# INLINE zip3 #-} +-- | Zip four arrays +-- +-- @since 0.5.4 +zip4 + :: (Index ix, Source r1 e1, Source r2 e2, Source r3 e3, Source r4 e4) + => Array r1 ix e1 + -> Array r2 ix e2 + -> Array r3 ix e3 + -> Array r4 ix e4 + -> Array D ix (e1, e2, e3, e4) +zip4 = zipWith4 (,,,) +{-# INLINE zip4 #-} + -- | Unzip two arrays -unzip :: Source r ix (e1, e2) => Array r ix (e1, e2) -> (Array D ix e1, Array D ix e2) +-- +-- @since 0.1.0 +unzip :: (Index ix, Source r (e1, e2)) => Array r ix (e1, e2) -> (Array D ix e1, Array D ix e2) unzip arr = (map fst arr, map snd arr) {-# INLINE unzip #-} -- | Unzip three arrays -unzip3 :: Source r ix (e1, e2, e3) - => Array r ix (e1, e2, e3) -> (Array D ix e1, Array D ix e2, Array D ix e3) -unzip3 arr = (map (\ (e, _, _) -> e) arr, map (\ (_, e, _) -> e) arr, map (\ (_, _, e) -> e) arr) +-- +-- @since 0.1.0 +unzip3 + :: (Index ix, Source r (e1, e2, e3)) + => Array r ix (e1, e2, e3) + -> (Array D ix e1, Array D ix e2, Array D ix e3) +unzip3 arr = (map (\(e, _, _) -> e) arr, map (\(_, e, _) -> e) arr, map (\(_, _, e) -> e) arr) {-# INLINE unzip3 #-} +-- | Unzip four arrays +-- +-- @since 0.5.4 +unzip4 + :: (Index ix, Source r (e1, e2, e3, e4)) + => Array r ix (e1, e2, e3, e4) + -> (Array D ix e1, Array D ix e2, Array D ix e3, Array D ix e4) +unzip4 arr = + ( map (\(e, _, _, _) -> e) arr + , map (\(_, e, _, _) -> e) arr + , map (\(_, _, e, _) -> e) arr + , map (\(_, _, _, e) -> e) arr + ) +{-# INLINE unzip4 #-} + -------------------------------------------------------------------------------- -- zipWith --------------------------------------------------------------------- -------------------------------------------------------------------------------- -- | Zip two arrays with a function. Resulting array will be an intersection of -- source arrays in case their dimensions do not match. -zipWith :: (Source r1 ix e1, Source r2 ix e2) - => (e1 -> e2 -> e) -> Array r1 ix e1 -> Array r2 ix e2 -> Array D ix e -zipWith f = izipWith (\ _ e1 e2 -> f e1 e2) +zipWith + :: (Index ix, Source r1 e1, Source r2 e2) + => (e1 -> e2 -> e) + -> Array r1 ix e1 + -> Array r2 ix e2 + -> Array D ix e +zipWith f arr1 arr2 = DArray comp sz prefIndex + where + sz = SafeSz (liftIndex2 min (coerce (size arr1)) (coerce (size arr2))) + comp = getComp arr1 <> getComp arr2 + prefIndex = PrefIndex (\ix -> f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)) +-- Somehow checking for size equality destroys performance +-- | PrefIndexLinear gi1 <- unsafePrefIndex arr1, +-- PrefIndexLinear gi2 <- unsafePrefIndex arr2, +-- size arr1 == size arr2 = +-- PrefIndexLinear (\i -> f (gi1 i) (gi2 i)) +-- | otherwise = PrefIndex (\ix -> f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)) {-# INLINE zipWith #-} - -- | Just like `zipWith`, except with an index aware function. -izipWith :: (Source r1 ix e1, Source r2 ix e2) - => (ix -> e1 -> e2 -> e) -> Array r1 ix e1 -> Array r2 ix e2 -> Array D ix e +izipWith + :: (Index ix, Source r1 e1, Source r2 e2) + => (ix -> e1 -> e2 -> e) + -> Array r1 ix e1 + -> Array r2 ix e2 + -> Array D ix e izipWith f arr1 arr2 = DArray (getComp arr1 <> getComp arr2) - (SafeSz (liftIndex2 min (coerce (size arr1)) (coerce (size arr2)))) $ \ !ix -> - f ix (unsafeIndex arr1 ix) (unsafeIndex arr2 ix) + (SafeSz (liftIndex2 min (coerce (size arr1)) (coerce (size arr2)))) + (PrefIndex (\ix -> f ix (unsafeIndex arr1 ix) (unsafeIndex arr2 ix))) {-# INLINE izipWith #-} - -- | Just like `zipWith`, except zip three arrays with a function. -zipWith3 :: (Source r1 ix e1, Source r2 ix e2, Source r3 ix e3) - => (e1 -> e2 -> e3 -> e) -> Array r1 ix e1 -> Array r2 ix e2 -> Array r3 ix e3 -> Array D ix e -zipWith3 f = izipWith3 (\ _ e1 e2 e3 -> f e1 e2 e3) +zipWith3 + :: (Index ix, Source r1 e1, Source r2 e2, Source r3 e3) + => (e1 -> e2 -> e3 -> e) + -> Array r1 ix e1 + -> Array r2 ix e2 + -> Array r3 ix e3 + -> Array D ix e +zipWith3 f arr1 arr2 arr3 = izipWith3 (\_ e1 e2 e3 -> f e1 e2 e3) arr1 arr2 arr3 +-- See note on zipWith +-- | sz1 == size arr2 && sz1 == size arr3 +-- , PrefIndexLinear gi1 <- unsafePrefIndex arr1 +-- , PrefIndexLinear gi2 <- unsafePrefIndex arr2 +-- , PrefIndexLinear gi3 <- unsafePrefIndex arr3 = +-- makeArrayLinear comp sz1 (\ !i -> f (gi1 i) (gi2 i) (gi3 i)) +-- | otherwise = izipWith3 (\_ e1 e2 e3 -> f e1 e2 e3) arr1 arr2 arr3 +-- where +-- comp = getComp arr1 <> getComp arr2 <> getComp arr3 +-- sz1 = size arr1 {-# INLINE zipWith3 #-} - -- | Just like `zipWith3`, except with an index aware function. izipWith3 - :: (Source r1 ix e1, Source r2 ix e2, Source r3 ix e3) + :: (Index ix, Source r1 e1, Source r2 e2, Source r3 e3) => (ix -> e1 -> e2 -> e3 -> e) -> Array r1 ix e1 -> Array r2 ix e2 @@ -170,50 +259,114 @@ izipWith3 f arr1 arr2 arr3 = DArray (getComp arr1 <> getComp arr2 <> getComp arr3) - (SafeSz - (liftIndex2 - min - (liftIndex2 min (coerce (size arr1)) (coerce (size arr2))) - (coerce (size arr3)))) $ \ !ix -> - f ix (unsafeIndex arr1 ix) (unsafeIndex arr2 ix) (unsafeIndex arr3 ix) + ( SafeSz + ( liftIndex2 + min + (liftIndex2 min (coerce (size arr1)) (coerce (size arr2))) + (coerce (size arr3)) + ) + ) + (PrefIndex $ \ !ix -> f ix (unsafeIndex arr1 ix) (unsafeIndex arr2 ix) (unsafeIndex arr3 ix)) {-# INLINE izipWith3 #-} +-- | Just like `zipWith`, except zip four arrays with a function. +-- +-- @since 0.5.4 +zipWith4 + :: (Index ix, Source r1 e1, Source r2 e2, Source r3 e3, Source r4 e4) + => (e1 -> e2 -> e3 -> e4 -> e) + -> Array r1 ix e1 + -> Array r2 ix e2 + -> Array r3 ix e3 + -> Array r4 ix e4 + -> Array D ix e +zipWith4 f arr1 arr2 arr3 arr4 = + izipWith4 (\_ e1 e2 e3 e4 -> f e1 e2 e3 e4) arr1 arr2 arr3 arr4 +-- See note on zipWith +-- | sz1 == size arr2 && sz1 == size arr3 && sz1 == size arr4 +-- , PrefIndexLinear gi1 <- unsafePrefIndex arr1 +-- , PrefIndexLinear gi2 <- unsafePrefIndex arr2 +-- , PrefIndexLinear gi3 <- unsafePrefIndex arr3 +-- , PrefIndexLinear gi4 <- unsafePrefIndex arr4 = +-- makeArrayLinear comp sz1 (\ !i -> f (gi1 i) (gi2 i) (gi3 i) (gi4 i)) +-- | otherwise = izipWith4 (\ _ e1 e2 e3 e4 -> f e1 e2 e3 e4) arr1 arr2 arr3 arr4 +-- where +-- comp = getComp arr1 <> getComp arr2 <> getComp arr3 <> getComp arr4 +-- sz1 = size arr1 +{-# INLINE zipWith4 #-} --- | Similar to `zipWith`, except does it sequentiall and using the `Applicative`. Note that --- resulting array has Mutable representation. +-- | Just like `zipWith4`, except with an index aware function. -- +-- @since 0.5.4 +izipWith4 + :: (Index ix, Source r1 e1, Source r2 e2, Source r3 e3, Source r4 e4) + => (ix -> e1 -> e2 -> e3 -> e4 -> e) + -> Array r1 ix e1 + -> Array r2 ix e2 + -> Array r3 ix e3 + -> Array r4 ix e4 + -> Array D ix e +izipWith4 f arr1 arr2 arr3 arr4 = + makeArray + (getComp arr1 <> getComp arr2 <> getComp arr3 <> getComp arr4) + ( SafeSz + ( liftIndex2 + min + ( liftIndex2 + min + (liftIndex2 min (coerce (size arr1)) (coerce (size arr2))) + (coerce (size arr3)) + ) + (coerce (size arr4)) + ) + ) + ( \ !ix -> + f ix (unsafeIndex arr1 ix) (unsafeIndex arr2 ix) (unsafeIndex arr3 ix) (unsafeIndex arr4 ix) + ) +{-# INLINE izipWith4 #-} + +-- | Similar to `zipWith`, except does it sequentially and using the `Applicative`. Note that +-- resulting array has Manifest representation. +-- -- @since 0.3.0 -zipWithA :: - (Source r1 ix e1, Source r2 ix e2, Applicative f, Mutable r ix e) +zipWithA + :: (Source r1 e1, Source r2 e2, Applicative f, Manifest r e, Index ix) => (e1 -> e2 -> f e) -> Array r1 ix e1 -> Array r2 ix e2 -> f (Array r ix e) -zipWithA f = izipWithA (const f) +zipWithA f arr1 arr2 + | sz1 == size arr2 + , PrefIndexLinear gi1 <- unsafePrefIndex arr1 + , PrefIndexLinear gi2 <- unsafePrefIndex arr2 = + setComp (getComp arr1 <> getComp arr2) <$> makeArrayLinearA sz1 (\ !i -> f (gi1 i) (gi2 i)) + | otherwise = izipWithA (const f) arr1 arr2 + where + !sz1 = size arr1 {-# INLINE zipWithA #-} --- | Similar to `zipWith`, except does it sequentiall and using the `Applicative`. Note that --- resulting array has Mutable representation. +-- | Similar to `zipWith`, except does it sequentially and using the `Applicative`. Note that +-- resulting array has Manifest representation. -- -- @since 0.3.0 -izipWithA :: - (Source r1 ix e1, Source r2 ix e2, Applicative f, Mutable r ix e) +izipWithA + :: (Source r1 e1, Source r2 e2, Applicative f, Manifest r e, Index ix) => (ix -> e1 -> e2 -> f e) -> Array r1 ix e1 -> Array r2 ix e2 -> f (Array r ix e) izipWithA f arr1 arr2 = - setComp (getComp arr1 <> getComp arr2) <$> - makeArrayA - (SafeSz (liftIndex2 min (coerce (size arr1)) (coerce (size arr2)))) - (\ !ix -> f ix (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)) + setComp (getComp arr1 <> getComp arr2) + <$> makeArrayA + (SafeSz (liftIndex2 min (coerce (size arr1)) (coerce (size arr2)))) + (\ !ix -> f ix (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)) {-# INLINE izipWithA #-} -- | Same as `zipWithA`, but for three arrays. -- -- @since 0.3.0 -zipWith3A :: - (Source r1 ix e1, Source r2 ix e2, Source r3 ix e3, Applicative f, Mutable r ix e) +zipWith3A + :: (Source r1 e1, Source r2 e2, Source r3 e3, Applicative f, Manifest r e, Index ix) => (e1 -> e2 -> e3 -> f e) -> Array r1 ix e1 -> Array r2 ix e2 @@ -225,76 +378,65 @@ -- | Same as `izipWithA`, but for three arrays. -- -- @since 0.3.0 -izipWith3A :: - (Source r1 ix e1, Source r2 ix e2, Source r3 ix e3, Applicative f, Mutable r ix e) +izipWith3A + :: (Source r1 e1, Source r2 e2, Source r3 e3, Applicative f, Manifest r e, Index ix) => (ix -> e1 -> e2 -> e3 -> f e) -> Array r1 ix e1 -> Array r2 ix e2 -> Array r3 ix e3 -> f (Array r ix e) izipWith3A f arr1 arr2 arr3 = - setComp (getComp arr1 <> getComp arr2 <> getComp arr3) <$> - makeArrayA sz (\ !ix -> f ix (unsafeIndex arr1 ix) (unsafeIndex arr2 ix) (unsafeIndex arr3 ix)) + setComp (getComp arr1 <> getComp arr2 <> getComp arr3) + <$> makeArrayA sz (\ !ix -> f ix (unsafeIndex arr1 ix) (unsafeIndex arr2 ix) (unsafeIndex arr3 ix)) where sz = SafeSz $ - liftIndex2 min (liftIndex2 min (coerce (size arr1)) (coerce (size arr2))) (coerce (size arr3)) + liftIndex2 min (liftIndex2 min (coerce (size arr1)) (coerce (size arr2))) (coerce (size arr3)) {-# INLINE izipWith3A #-} - - --- | Similar to `Data.Massiv.Array.zipWith`, except dimensions of both arrays either have to be the --- same, or at least one of the two array must be a singleton array, in which case it will behave as --- a `Data.Massiv.Array.map`. --- --- @since 0.1.4 -liftArray2 - :: (Source r1 ix a, Source r2 ix b) - => (a -> b -> e) -> Array r1 ix a -> Array r2 ix b -> Array D ix e -liftArray2 f !arr1 !arr2 - | sz1 == oneSz = map (f (unsafeIndex arr1 zeroIndex)) arr2 - | sz2 == oneSz = map (`f` unsafeIndex arr2 zeroIndex) arr1 - | sz1 == sz2 = - DArray (getComp arr1 <> getComp arr2) sz1 (\ !ix -> f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)) - | otherwise = throw $ SizeMismatchException (size arr1) (size arr2) - where - sz1 = size arr1 - sz2 = size arr2 -{-# INLINE liftArray2 #-} - - -------------------------------------------------------------------------------- -- traverse -------------------------------------------------------------------- -------------------------------------------------------------------------------- -- | Traverse with an `Applicative` action over an array sequentially. -- --- /Note/ - using `traversePrim` will always be faster, althought not always possible. +-- /Note/ - using `traversePrim` instead will always be significantly faster, roughly +-- about 30 times faster in practice. -- -- @since 0.2.6 --- -traverseA :: - forall r ix e r' a f . (Source r' ix a, Mutable r ix e, Applicative f) +traverseA + :: forall r ix e r' a f + . (Source r' a, Manifest r e, Index ix, Applicative f) => (a -> f e) -> Array r' ix a -> f (Array r ix e) -traverseA f arr = makeArrayLinearA (size arr) (f . unsafeLinearIndex arr) +traverseA f arr = + unsafeResize (size arr) . fromList (getComp arr) <$> traverse f (toList arr) {-# INLINE traverseA #-} -- | Traverse sequentially over a source array, while discarding the result. -- -- @since 0.3.0 --- -traverseA_ :: forall r ix e a f . (Source r ix e, Applicative f) => (e -> f a) -> Array r ix e -> f () -traverseA_ f arr = loopA_ 0 (< totalElem (size arr)) (+ 1) (f . unsafeLinearIndex arr) +traverseA_ + :: forall r ix e a f + . (Index ix, Source r e, Applicative f) + => (e -> f a) + -> Array r ix e + -> f () +traverseA_ f arr = + case unsafePrefIndex arr of + PrefIndex gix -> iterA_ zeroIndex (unSz sz) oneIndex (<) (f . gix) + PrefIndexLinear gi -> loopA_ 0 (< totalElem sz) (+ 1) (f . gi) + where + sz = size arr {-# INLINE traverseA_ #-} -- | Sequence actions in a source array. -- -- @since 0.3.0 --- -sequenceA :: - forall r ix e r' f. (Source r' ix (f e), Mutable r ix e, Applicative f) +sequenceA + :: forall r ix e r' f + . (Source r' (f e), Manifest r e, Index ix, Applicative f) => Array r' ix (f e) -> f (Array r ix e) sequenceA = traverseA id @@ -303,18 +445,20 @@ -- | Sequence actions in a source array, while discarding the result. -- -- @since 0.3.0 --- -sequenceA_ :: forall r ix e f . (Source r ix (f e), Applicative f) => Array r ix (f e) -> f () +sequenceA_ + :: forall r ix e f + . (Index ix, Source r (f e), Applicative f) + => Array r ix (f e) + -> f () sequenceA_ = traverseA_ id {-# INLINE sequenceA_ #-} - -- | Traverse with an `Applicative` index aware action over an array sequentially. -- -- @since 0.2.6 --- -itraverseA :: - forall r ix e r' a f . (Source r' ix a, Mutable r ix e, Applicative f) +itraverseA + :: forall r ix e r' a f + . (Source r' a, Manifest r e, Index ix, Applicative f) => (ix -> a -> f e) -> Array r' ix a -> f (Array r ix e) @@ -322,114 +466,69 @@ setComp (getComp arr) <$> makeArrayA (size arr) (\ !ix -> f ix (unsafeIndex arr ix)) {-# INLINE itraverseA #-} - -- | Traverse with an `Applicative` index aware action over an array sequentially. -- -- @since 0.2.6 --- -itraverseA_ :: - forall r ix e a f. (Source r ix a, Applicative f) +itraverseA_ + :: forall r ix e a f + . (Source r a, Index ix, Applicative f) => (ix -> a -> f e) -> Array r ix a -> f () itraverseA_ f arr = - loopA_ 0 (< totalElem sz) (+ 1) (\ !i -> f (fromLinearIndex sz i) (unsafeLinearIndex arr i)) + case unsafePrefIndex arr of + PrefIndex gix -> + iterA_ zeroIndex (unSz sz) oneIndex (<) (\ !ix -> f ix (gix ix)) + PrefIndexLinear gi -> + iterTargetA_ defRowMajor 0 sz zeroIndex oneStride $ \i ix -> f ix (gi i) where sz = size arr {-# INLINE itraverseA_ #-} - - --- | Same as `traverseA`, except with ability to specify representation. --- --- @since 0.2.6 --- -traverseAR :: - (Source r' ix a, Mutable r ix b, Applicative f) - => r - -> (a -> f b) - -> Array r' ix a - -> f (Array r ix b) -traverseAR _ = traverseA -{-# INLINE traverseAR #-} -{-# DEPRECATED traverseAR "In favor of `traverseA`" #-} - --- | Same as `itraverseA`, except with ability to specify representation. --- --- @since 0.2.6 --- -itraverseAR :: - (Source r' ix a, Mutable r ix b, Applicative f) - => r - -> (ix -> a -> f b) - -> Array r' ix a - -> f (Array r ix b) -itraverseAR _ = itraverseA -{-# INLINE itraverseAR #-} -{-# DEPRECATED itraverseAR "In favor of `itraverseA`" #-} - - - -- | Traverse sequentially within `PrimMonad` over an array with an action. -- -- @since 0.3.0 --- -traversePrim :: - forall r ix b r' a m . (Source r' ix a, Mutable r ix b, PrimMonad m) +traversePrim + :: forall r ix b r' a m + . (Source r' a, Manifest r b, Index ix, PrimMonad m) => (a -> m b) -> Array r' ix a -> m (Array r ix b) -traversePrim f = itraversePrim (const f) +traversePrim f arr = do + let sz = size arr + marr <- unsafeNew sz + case unsafePrefIndex arr of + PrefIndex gix -> + iterTargetA_ defRowMajor 0 sz zeroIndex oneStride $ \i ix -> + f (gix ix) >>= unsafeLinearWrite marr i + PrefIndexLinear gi -> + loopA_ 0 (< totalElem sz) (+ 1) $ \i -> + f (gi i) >>= unsafeLinearWrite marr i + unsafeFreeze (getComp arr) marr {-# INLINE traversePrim #-} -- | Same as `traversePrim`, but traverse with index aware action. -- -- @since 0.3.0 --- -itraversePrim :: - forall r ix b r' a m . (Source r' ix a, Mutable r ix b, PrimMonad m) +itraversePrim + :: forall r ix b r' a m + . (Source r' a, Manifest r b, Index ix, PrimMonad m) => (ix -> a -> m b) -> Array r' ix a -> m (Array r ix b) -itraversePrim f arr = - setComp (getComp arr) <$> - generateArrayLinearS - (size arr) - (\ !i -> - let ix = fromLinearIndex (size arr) i - in f ix (unsafeLinearIndex arr i)) +itraversePrim f arr = do + let sz = size arr + marr <- unsafeNew sz + case unsafePrefIndex arr of + PrefIndex gix -> + iterTargetA_ defRowMajor 0 sz zeroIndex oneStride $ \i ix -> + f ix (gix ix) >>= unsafeLinearWrite marr i + PrefIndexLinear gi -> + iterTargetA_ defRowMajor 0 sz zeroIndex oneStride $ \i ix -> + f ix (gi i) >>= unsafeLinearWrite marr i + unsafeFreeze (getComp arr) marr {-# INLINE itraversePrim #-} - --- | Same as `traversePrim`, but with ability to specify the desired representation. --- --- @since 0.3.0 --- -traversePrimR :: - (Source r' ix a, Mutable r ix b, PrimMonad m) - => r - -> (a -> m b) - -> Array r' ix a - -> m (Array r ix b) -traversePrimR _ = traversePrim -{-# INLINE traversePrimR #-} -{-# DEPRECATED traversePrimR "In favor of `traversePrim`" #-} - --- | Same as `itraversePrim`, but with ability to specify the desired representation. --- --- @since 0.3.0 --- -itraversePrimR :: - (Source r' ix a, Mutable r ix b, PrimMonad m) - => r - -> (ix -> a -> m b) - -> Array r' ix a - -> m (Array r ix b) -itraversePrimR _ = itraversePrim -{-# INLINE itraversePrimR #-} -{-# DEPRECATED itraversePrimR "In favor of `itraversePrim`" #-} - - -------------------------------------------------------------------------------- -- mapM ------------------------------------------------------------------------ -------------------------------------------------------------------------------- @@ -437,106 +536,53 @@ -- | Map a monadic action over an array sequentially. -- -- @since 0.2.6 -mapM :: - forall r ix b r' a m. (Source r' ix a, Mutable r ix b, Monad m) - => (a -> m b) -- ^ Mapping action - -> Array r' ix a -- ^ Source array +mapM + :: forall r ix b r' a m + . (Source r' a, Manifest r b, Index ix, Monad m) + => (a -> m b) + -- ^ Mapping action + -> Array r' ix a + -- ^ Source array -> m (Array r ix b) mapM = traverseA {-# INLINE mapM #-} - --- | Same as `mapM`, except with ability to specify result representation. --- --- @since 0.2.6 -mapMR :: - forall r ix b r' a m. (Source r' ix a, Mutable r ix b, Monad m) - => r - -> (a -> m b) - -> Array r' ix a - -> m (Array r ix b) -mapMR _ = traverseA -{-# INLINE mapMR #-} - - -- | Same as `mapM` except with arguments flipped. -- -- @since 0.2.6 -forM :: - forall r ix b r' a m. (Source r' ix a, Mutable r ix b, Monad m) +forM + :: forall r ix b r' a m + . (Source r' a, Manifest r b, Index ix, Monad m) => Array r' ix a -> (a -> m b) -> m (Array r ix b) forM = flip traverseA {-# INLINE forM #-} - --- | Same as `forM`, except with ability to specify result representation. --- --- @since 0.2.6 -forMR :: - forall r ix b r' a m. (Source r' ix a, Mutable r ix b, Monad m) - => r - -> Array r' ix a - -> (a -> m b) - -> m (Array r ix b) -forMR _ = flip traverseA -{-# INLINE forMR #-} - - - --- | Map a monadic action over an array sequentially. +-- | Map an index aware monadic action over an array sequentially. -- -- @since 0.2.6 -imapM :: - forall r ix b r' a m. (Source r' ix a, Mutable r ix b, Monad m) +imapM + :: forall r ix b r' a m + . (Source r' a, Manifest r b, Index ix, Monad m) => (ix -> a -> m b) -> Array r' ix a -> m (Array r ix b) imapM = itraverseA {-# INLINE imapM #-} - --- | Same as `imapM`, except with ability to specify result representation. +-- | Same as `forM`, except with an index aware action. -- --- @since 0.2.6 -imapMR :: - forall r ix b r' a m. (Source r' ix a, Mutable r ix b, Monad m) - => r +-- @since 0.5.1 +iforM + :: forall r ix b r' a m + . (Source r' a, Manifest r b, Index ix, Monad m) + => Array r' ix a -> (ix -> a -> m b) - -> Array r' ix a -> m (Array r ix b) -imapMR _ = itraverseA -{-# INLINE imapMR #-} - - - --- | Same as `forM`, except map an index aware action. --- --- @since 0.2.6 -iforM :: - forall r ix b r' a m. (Source r' ix a, Mutable r ix b, Monad m) - => (ix -> a -> m b) - -> Array r' ix a - -> m (Array r ix b) -iforM = itraverseA +iforM = flip itraverseA {-# INLINE iforM #-} - --- | Same as `iforM`, except with ability to specify result representation. --- --- @since 0.2.6 --- -iforMR :: - forall r ix b r' a m. (Source r' ix a, Mutable r ix b, Monad m) - => r - -> (ix -> a -> m b) - -> Array r' ix a - -> m (Array r ix b) -iforMR _ = itraverseA -{-# INLINE iforMR #-} - - -- | Map a monadic function over an array sequentially, while discarding the result. -- -- ==== __Examples__ @@ -550,11 +596,10 @@ -- 58 -- -- @since 0.1.0 -mapM_ :: (Source r ix a, Monad m) => (a -> m b) -> Array r ix a -> m () -mapM_ f !arr = iterM_ zeroIndex (unSz (size arr)) (pureIndex 1) (<) (f . unsafeIndex arr) +mapM_ :: (Source r a, Index ix, Monad m) => (a -> m b) -> Array r ix a -> m () +mapM_ = traverseA_ {-# INLINE mapM_ #-} - -- | Just like `mapM_`, except with flipped arguments. -- -- ==== __Examples__ @@ -568,102 +613,122 @@ -- >>> A.forM_ (range Seq (Ix1 0) 1000) $ \ i -> modifyIORef' ref (+i) -- >>> readIORef ref -- 499500 --- -forM_ :: (Source r ix a, Monad m) => Array r ix a -> (a -> m b) -> m () -forM_ = flip mapM_ +forM_ :: (Source r a, Index ix, Monad m) => Array r ix a -> (a -> m b) -> m () +forM_ = flip traverseA_ {-# INLINE forM_ #-} +-- | Map a monadic index aware function over an array sequentially, while discarding the result. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array +-- >>> imapM_ (curry print) $ range Seq (Ix1 10) 15 +-- (0,10) +-- (1,11) +-- (2,12) +-- (3,13) +-- (4,14) +-- +-- @since 0.1.0 +imapM_ :: (Index ix, Source r a, Monad m) => (ix -> a -> m b) -> Array r ix a -> m () +imapM_ = itraverseA_ +{-# INLINE imapM_ #-} -- | Just like `imapM_`, except with flipped arguments. -iforM_ :: (Source r ix a, Monad m) => Array r ix a -> (ix -> a -> m b) -> m () -iforM_ = flip imapM_ +iforM_ :: (Source r a, Index ix, Monad m) => Array r ix a -> (ix -> a -> m b) -> m () +iforM_ = flip itraverseA_ {-# INLINE iforM_ #-} - -- | Map an `IO` action over an `Array`. Underlying computation strategy is respected and will be -- parallelized when requested. Unfortunately no fusion is possible and new array will be create -- upon each call. -- -- @since 0.2.6 -mapIO :: - forall r ix b r' a m. (Source r' ix a, Mutable r ix b, MonadUnliftIO m, PrimMonad m) +mapIO + :: forall r ix b r' a m + . (Size r', Load r' ix a, Manifest r b, MonadUnliftIO m) => (a -> m b) -> Array r' ix a -> m (Array r ix b) mapIO action = imapIO (const action) {-# INLINE mapIO #-} --- | Similar to `mapIO`, but ignores the result of mapping action and does not create a resulting --- array, therefore it is faster. Use this instead of `mapIO` when result is irrelevant. +-- | Similar to `mapIO`, but ignores the result of mapping action and does not +-- create a resulting array, therefore it is faster. Use this instead of `mapIO` +-- when result is irrelevant. Most importantly it will follow the iteration +-- logic outlined by the supplied array. -- -- @since 0.2.6 -mapIO_ :: (Source r b e, MonadUnliftIO m) => (e -> m a) -> Array r b e -> m () -mapIO_ action = imapIO_ (const action) +mapIO_ + :: forall r ix e a m + . (Load r ix e, MonadUnliftIO m) + => (e -> m a) + -> Array r ix e + -> m () +mapIO_ action arr = + withRunInIO $ \run -> + withMassivScheduler_ (getComp arr) $ \scheduler -> + iterArrayLinearM_ scheduler arr (\_ -> void . run . action) {-# INLINE mapIO_ #-} -- | Same as `mapIO_`, but map an index aware action instead. -- -- @since 0.2.6 -imapIO_ :: (Source r ix e, MonadUnliftIO m) => (ix -> e -> m a) -> Array r ix e -> m () +imapIO_ + :: forall r ix e a m + . (Load r ix e, MonadUnliftIO m) + => (ix -> e -> m a) + -> Array r ix e + -> m () imapIO_ action arr = - withScheduler_ (getComp arr) $ \scheduler -> imapSchedulerM_ scheduler action arr + withRunInIO $ \run -> + withMassivScheduler_ (getComp arr) $ \scheduler -> + let sz = outerSize arr + in -- It is ok to use outerSize in context of DS and L. Former is 1-dim, + -- so sz is never evaluated and for the latter outerSize has to be + -- called regardless how this function is implemented. + iterArrayLinearM_ scheduler arr (\i -> void . run . action (fromLinearIndex sz i)) {-# INLINE imapIO_ #-} --- | Same as `imapM_`, but will use the supplied scheduler. --- --- @since 0.3.1 -imapSchedulerM_ :: - (Source r ix e, Monad m) => Scheduler m () -> (ix -> e -> m a) -> Array r ix e -> m () -imapSchedulerM_ scheduler action arr = do - let sz = size arr - splitLinearlyWith_ - scheduler - (totalElem sz) - (unsafeLinearIndex arr) - (\i -> void . action (fromLinearIndex sz i)) -{-# INLINE imapSchedulerM_ #-} - - --- | Same as `imapM_`, but will use the supplied scheduler. --- --- @since 0.3.1 -iforSchedulerM_ :: - (Source r ix e, Monad m) => Scheduler m () -> Array r ix e -> (ix -> e -> m a) -> m () -iforSchedulerM_ scheduler arr action = imapSchedulerM_ scheduler action arr -{-# INLINE iforSchedulerM_ #-} - - --- | Same as `mapIO` but map an index aware action instead. +-- | Same as `mapIO` but map an index aware action instead. Respects computation strategy. -- -- @since 0.2.6 -imapIO :: - forall r ix b r' a m. (Source r' ix a, Mutable r ix b, MonadUnliftIO m, PrimMonad m) +imapIO + :: forall r ix b r' a m + . (Size r', Load r' ix a, Manifest r b, MonadUnliftIO m) => (ix -> a -> m b) -> Array r' ix a -> m (Array r ix b) -imapIO action arr = generateArray (getComp arr) (size arr) $ \ix -> action ix (unsafeIndex arr ix) +imapIO action arr = do + let sz = size arr + withRunInIO $ \run -> do + marr <- unsafeNew sz + withMassivScheduler_ (getComp arr) $ \scheduler -> + iterArrayLinearM_ scheduler arr $ \ !i e -> + run (action (fromLinearIndex sz i) e) >>= unsafeLinearWrite marr i + unsafeFreeze (getComp arr) marr {-# INLINE imapIO #-} -- | Same as `mapIO` but with arguments flipped. -- -- @since 0.2.6 -forIO :: - forall r ix b r' a m. (Source r' ix a, Mutable r ix b, MonadUnliftIO m, PrimMonad m) +forIO + :: forall r ix b r' a m + . (Size r', Load r' ix a, Manifest r b, MonadUnliftIO m) => Array r' ix a -> (a -> m b) -> m (Array r ix b) forIO = flip mapIO {-# INLINE forIO #-} - - --- | Same as `imapIO`, but ignores the inner computation strategy and uses stateful --- workers during computation instead. Use `initWorkerStates` for the `WorkerStates` --- initialization. +-- | Same as `imapIO`, but ignores the inner computation strategy and uses +-- stateful workers during computation instead. Use +-- `Control.Scheduler.initWorkerStates` for the `WorkerStates` initialization. -- -- @since 0.3.4 -imapWS :: - forall r ix b r' a s m. (Source r' ix a, Mutable r ix b, MonadUnliftIO m, PrimMonad m) +imapWS + :: forall r ix b r' a s m + . (Source r' a, Manifest r b, Index ix, MonadUnliftIO m, PrimMonad m) => WorkerStates s -> (ix -> a -> s -> m b) -> Array r' ix a @@ -674,21 +739,22 @@ -- | Same as `imapWS`, but without the index. -- -- @since 0.3.4 -mapWS :: - forall r ix b r' a s m. (Source r' ix a, Mutable r ix b, MonadUnliftIO m, PrimMonad m) +mapWS + :: forall r ix b r' a s m + . (Source r' a, Manifest r b, Index ix, MonadUnliftIO m, PrimMonad m) => WorkerStates s -> (a -> s -> m b) -> Array r' ix a -> m (Array r ix b) -mapWS states f = imapWS states (\ _ -> f) +mapWS states f = imapWS states (\_ -> f) {-# INLINE mapWS #-} - -- | Same as `imapWS`, but with source array and mapping action arguments flipped. -- -- @since 0.3.4 -iforWS :: - forall r ix b r' a s m. (Source r' ix a, Mutable r ix b, MonadUnliftIO m, PrimMonad m) +iforWS + :: forall r ix b r' a s m + . (Source r' a, Manifest r b, Index ix, MonadUnliftIO m, PrimMonad m) => WorkerStates s -> Array r' ix a -> (ix -> a -> s -> m b) @@ -699,17 +765,16 @@ -- | Same as `iforWS`, but without the index. -- -- @since 0.3.4 -forWS :: - forall r ix b r' a s m. (Source r' ix a, Mutable r ix b, MonadUnliftIO m, PrimMonad m) +forWS + :: forall r ix b r' a s m + . (Source r' a, Manifest r b, Index ix, MonadUnliftIO m, PrimMonad m) => WorkerStates s -> Array r' ix a -> (a -> s -> m b) -> m (Array r ix b) -forWS states arr f = imapWS states (\ _ -> f) arr +forWS states arr f = imapWS states (\_ -> f) arr {-# INLINE forWS #-} - - -- | Same as `mapIO_` but with arguments flipped. -- -- ==== __Example__ @@ -725,15 +790,16 @@ -- 499500 -- -- @since 0.2.6 -forIO_ :: (Source r ix e, MonadUnliftIO m) => Array r ix e -> (e -> m a) -> m () +forIO_ :: (Load r ix e, MonadUnliftIO m) => Array r ix e -> (e -> m a) -> m () forIO_ = flip mapIO_ {-# INLINE forIO_ #-} -- | Same as `imapIO` but with arguments flipped. -- -- @since 0.2.6 -iforIO :: - forall r ix b r' a m. (Source r' ix a, Mutable r ix b, MonadUnliftIO m, PrimMonad m) +iforIO + :: forall r ix b r' a m + . (Size r', Load r' ix a, Manifest r b, MonadUnliftIO m) => Array r' ix a -> (ix -> a -> m b) -> m (Array r ix b) @@ -743,6 +809,133 @@ -- | Same as `imapIO_` but with arguments flipped. -- -- @since 0.2.6 -iforIO_ :: (Source r ix a, MonadUnliftIO m) => Array r ix a -> (ix -> a -> m b) -> m () +iforIO_ + :: forall r ix e a m + . (Load r ix e, MonadUnliftIO m) + => Array r ix e + -> (ix -> e -> m a) + -> m () iforIO_ = flip imapIO_ {-# INLINE iforIO_ #-} + +iterArrayLinearM_ + :: forall r ix e m s + . (Load r ix e, MonadPrimBase s m) + => Scheduler s () + -> Array r ix e + -- ^ Array that is being loaded + -> (Int -> e -> m ()) + -- ^ Function that writes an element into target array + -> m () +iterArrayLinearM_ scheduler arr f = + stToPrim $ iterArrayLinearST_ scheduler arr (\i -> primToPrim . f i) +{-# INLINE iterArrayLinearM_ #-} + +iterArrayLinearWithSetM_ + :: forall r ix e m s + . (Load r ix e, MonadPrimBase s m) + => Scheduler s () + -> Array r ix e + -- ^ Array that is being loaded + -> (Int -> e -> m ()) + -- ^ Function that writes an element into target array + -> (Ix1 -> Sz1 -> e -> m ()) + -- ^ Function that efficiently sets a region of an array + -- to the supplied value target array + -> m () +iterArrayLinearWithSetM_ scheduler arr f set = + stToPrim $ + iterArrayLinearWithSetST_ scheduler arr (\i -> primToPrim . f i) (\i n -> primToPrim . set i n) +{-# INLINE iterArrayLinearWithSetM_ #-} + +iterArrayLinearWithStrideM_ + :: forall r ix e m s + . (StrideLoad r ix e, MonadPrimBase s m) + => Scheduler s () + -> Stride ix + -- ^ Stride to use + -> Sz ix + -- ^ Size of the target array affected by the stride. + -> Array r ix e + -- ^ Array that is being loaded + -> (Int -> e -> m ()) + -- ^ Function that writes an element into target array + -> m () +iterArrayLinearWithStrideM_ scheduler stride sz arr f = + stToPrim $ iterArrayLinearWithStrideST_ scheduler stride sz arr (\i -> primToPrim . f i) +{-# INLINE iterArrayLinearWithStrideM_ #-} + +-- iterArrayM_ :: +-- Scheduler s () +-- -> Array r ix e -- ^ Array that is being loaded +-- -> (Int -> e -> ST s ()) -- ^ Function that writes an element into target array +-- -> ST s () +-- iterArrayM_ scheduler arr uWrite + +-- Deprecated + +-- | Same as `imapM_`, but will use the supplied scheduler. +-- +-- @since 0.3.1 +imapSchedulerM_ + :: (Index ix, Source r e, MonadPrimBase s m) + => Scheduler s () + -> (ix -> e -> m a) + -> Array r ix e + -> m () +imapSchedulerM_ scheduler action arr = do + let sz = size arr + splitLinearlyWith_ + scheduler + (totalElem sz) + (unsafeLinearIndex arr) + (\i -> void . action (fromLinearIndex sz i)) +{-# INLINE imapSchedulerM_ #-} + +-- | Same as `imapM_`, but will use the supplied scheduler. +-- +-- @since 0.3.1 +iforSchedulerM_ + :: (Index ix, Source r e, MonadPrimBase s m) + => Scheduler s () + -> Array r ix e + -> (ix -> e -> m a) + -> m () +iforSchedulerM_ scheduler arr action = imapSchedulerM_ scheduler action arr +{-# INLINE iforSchedulerM_ #-} + +-- -- | Load an array into memory. +-- -- +-- -- @since 0.3.0 +-- loadArrayM +-- :: Scheduler s () +-- -> Array r ix e -- ^ Array that is being loaded +-- -> (Int -> e -> ST s ()) -- ^ Function that writes an element into target array +-- -> ST s () +-- loadArrayM scheduler arr uWrite = +-- loadArrayWithSetM scheduler arr uWrite $ \offset sz e -> +-- loopM_ offset (< (offset + unSz sz)) (+1) (`uWrite` e) +-- {-# INLINE loadArrayM #-} + +-- -- | Load an array into memory, just like `loadArrayM`. Except it also accepts a +-- -- function that is potentially optimized for setting many cells in a region to the same +-- -- value +-- -- +-- -- @since 0.5.8 +-- loadArrayWithSetM +-- :: Scheduler s () +-- -> Array r ix e -- ^ Array that is being loaded +-- -> (Ix1 -> e -> ST s ()) -- ^ Function that writes an element into target array +-- -> (Ix1 -> Sz1 -> e -> ST s ()) -- ^ Function that efficiently sets a region of an array +-- -- to the supplied value target array +-- -> ST s () +-- loadArrayWithSetM scheduler arr uWrite _ = loadArrayM scheduler arr uWrite +-- {-# INLINE loadArrayWithSetM #-} + +-- iterArrayLinearWithStrideST +-- :: Scheduler s () +-- -> Stride ix -- ^ Stride to use +-- -> Sz ix -- ^ Size of the target array affected by the stride. +-- -> Array r ix e -- ^ Array that is being loaded +-- -> (Int -> e -> ST s ()) -- ^ Function that writes an element into target array +-- -> ST s () diff --git a/src/Data/Massiv/Array/Ops/Slice.hs b/src/Data/Massiv/Array/Ops/Slice.hs --- a/src/Data/Massiv/Array/Ops/Slice.hs +++ b/src/Data/Massiv/Array/Ops/Slice.hs @@ -1,35 +1,44 @@ {-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} + -- | -- Module : Data.Massiv.Array.Ops.Slice --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Ops.Slice - ( +module Data.Massiv.Array.Ops.Slice ( -- ** From the outside - (!>) - , (!?>) - , (??>) + (!>), + (!?>), + (??>), + -- ** From the inside - , () - , () - , () - ) where + (), + (), + (), + -- ** Many slices + outerSlices, + innerSlices, + withinSlices, + withinSlicesM, +) where + import Control.Monad (unless) +import Data.Massiv.Array.Delayed.Pull import Data.Massiv.Core.Common - infixl 4 !>, !?>, ??>, , , - -- | /O(1)/ - Slices the array from the outside. For 2-dimensional array this will -- be equivalent of taking a row. Throws an error when index is out of bounds. -- @@ -54,41 +63,47 @@ -- ] -- ] -- >>> arr !> 2 --- Array M Seq (Sz (2 :. 4)) +-- Array U Seq (Sz (2 :. 4)) -- [ [ (2,0,0), (2,0,1), (2,0,2), (2,0,3) ] -- , [ (2,1,0), (2,1,1), (2,1,2), (2,1,3) ] -- ] -- --- There is nothing wrong with chaining, mixing and matching slicing operators, or even using them --- to index arrays: +-- There is nothing wrong with chaining, mixing and matching slicing operators: -- --- >>> arr !> 2 !> 0 !> 3 +-- >>> arr !> 2 !> 0 ! 3 -- (2,0,3) --- >>> arr !> 2 >> evaluateM (arr !> 2 >> (arr !> 2 !> 0 !> 3) == (arr ! 2 :> 0 :. 3) +-- >>> (arr !> 2 !> 0 ! 3) == (arr ! 2 :> 0 :. 3) -- True -- -- -- @since 0.1.0 -(!>) :: OuterSlice r ix e => Array r ix e -> Int -> Elt r ix e -(!>) !arr !ix = either throw id (arr !?> ix) +(!>) + :: forall r ix e + . (HasCallStack, Index ix, Index (Lower ix), Source r e) + => Array r ix e + -> Int + -> Array r (Lower ix) e +(!>) !arr !ix = throwEither (arr !?> ix) {-# INLINE (!>) #-} - -- | /O(1)/ - Just like `!>` slices the array from the outside, but returns -- `Nothing` when index is out of bounds. -- -- @since 0.1.0 -(!?>) :: (MonadThrow m, OuterSlice r ix e) => Array r ix e -> Int -> m (Elt r ix e) -(!?>) !arr !i - | isSafeIndex sz i = pure $ unsafeOuterSlice arr i - | otherwise = throwM $ IndexOutOfBoundsException sz i - where - !sz = fst (unconsSz (size arr)) +(!?>) + :: forall r ix e m + . (MonadThrow m, Index ix, Index (Lower ix), Source r e) + => Array r ix e + -> Int + -> m (Array r (Lower ix) e) +(!?>) !arr !i = do + let (k, szL) = unconsSz (size arr) + unless (isSafeIndex k i) $ throwM $ IndexOutOfBoundsException k i + pure $ unsafeOuterSlice arr szL i {-# INLINE (!?>) #-} - -- | /O(1)/ - Safe slicing continuation from the outside. Similarly to (`!>`) slices the array from -- the outside, but takes `Maybe` array as input and returns `Nothing` when index is out of bounds. -- @@ -96,81 +111,234 @@ -- -- >>> import Data.Massiv.Array -- >>> arr = makeArrayR U Seq (Sz (3 :> 2 :. 4)) fromIx3 --- >>> arr !?> 2 ??> 0 ??> 3 :: Maybe Ix3T +-- >>> arr !?> 2 ??> 0 ?? 3 :: Maybe Ix3T -- Just (2,0,3) --- >>> arr !?> 2 ??> 0 ??> -1 :: Maybe Ix3T +-- >>> arr !?> 2 ??> 0 ?? -1 :: Maybe Ix3T -- Nothing -- >>> arr !?> 2 ??> -10 ?? 1 -- *** Exception: IndexOutOfBoundsException: -10 is not safe for (Sz1 2) -- -- @since 0.1.0 -(??>) :: (MonadThrow m, OuterSlice r ix e) => m (Array r ix e) -> Int -> m (Elt r ix e) +(??>) + :: forall r ix e m + . (MonadThrow m, Index ix, Index (Lower ix), Source r e) + => m (Array r ix e) + -> Int + -> m (Array r (Lower ix) e) (??>) marr !ix = marr >>= (!?> ix) {-# INLINE (??>) #-} - -- | /O(1)/ - Safe slice from the inside -- -- @since 0.1.0 -( Array r ix e -> Int -> m (Elt r ix e) -( Array r ix e + -> Int + -> m (Array D (Lower ix) e) +(`) slice an array from an opposite direction. -- -- @since 0.1.0 -( Array r ix e -> Int -> Elt r ix e -( res - Left exc -> throw exc +( Array r ix e + -> Int + -> Array D (Lower ix) e +( m (Array r ix e) -> Int -> m (Elt r ix e) +( m (Array r ix e) + -> Int + -> m (Array D (Lower ix) e) (>= (`), but fails gracefully with a `Nothing`, instead of an error -- -- @since 0.1.0 -() :: (MonadThrow m, Slice r ix e) => Array r ix e -> (Dim, Int) -> m (Elt r ix e) +() + :: forall r ix e m + . (MonadThrow m, Index ix, Index (Lower ix), Source r e) + => Array r ix e + -> (Dim, Int) + -> m (Array D (Lower ix) e) () !arr (dim, i) = do (m, szl) <- pullOutSzM (size arr) dim unless (isSafeIndex m i) $ throwM $ IndexOutOfBoundsException m i - start <- setDimM zeroIndex dim i cutSz <- insertSzM szl dim oneSz - unsafeSlice arr start cutSz dim + internalInnerSlice dim cutSz arr i {-# INLINE () #-} +internalInnerSlice + :: (MonadThrow m, Index ix, Index (Lower ix), Source r e) + => Dim + -> Sz ix + -> Array r ix e + -> Ix1 + -> m (Array D (Lower ix) e) +internalInnerSlice dim cutSz arr i = do + start <- setDimM zeroIndex dim i + unsafeSlice arr start cutSz dim +{-# INLINE internalInnerSlice #-} -- prop> arr !> i == arr (dimensions (size arr), i) -- prop> arr (1,i) -- + -- | /O(1)/ - Slices the array in any available dimension. Throws an error when -- index is out of bounds or dimensions is invalid. -- -- @since 0.1.0 -() :: Slice r ix e => Array r ix e -> (Dim, Int) -> Elt r ix e -() !arr !dix = - case arr dix of - Right res -> res - Left exc -> throw exc +() + :: forall r ix e + . (HasCallStack, Index ix, Index (Lower ix), Source r e) + => Array r ix e + -> (Dim, Int) + -> Array D (Lower ix) e +() !arr !dix = throwEither (arr dix) {-# INLINE () #-} - -- | /O(1)/ - Safe slicing continuation from within. -- -- @since 0.1.0 -() :: (MonadThrow m, Slice r ix e) => m (Array r ix e) -> (Dim, Int) -> m (Elt r ix e) +() + :: forall r ix e m + . (MonadThrow m, Index ix, Index (Lower ix), Source r e) + => m (Array r ix e) + -> (Dim, Int) + -> m (Array D (Lower ix) e) () !marr !ix = marr >>= ( ix) {-# INLINE () #-} + +-- | Create a delayed array of outer slices. +-- +-- ====__Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> A.mapM_ print $ outerSlices (0 ..: (3 :. 2)) +-- Array D Seq (Sz1 2) +-- [ 0 :. 0, 0 :. 1 ] +-- Array D Seq (Sz1 2) +-- [ 1 :. 0, 1 :. 1 ] +-- Array D Seq (Sz1 2) +-- [ 2 :. 0, 2 :. 1 ] +-- +-- @since 0.5.4 +outerSlices + :: forall r ix e + . (Index ix, Index (Lower ix), Source r e) + => Array r ix e + -> Array D Ix1 (Array r (Lower ix) e) +outerSlices arr = makeArray (getComp arr) k (unsafeOuterSlice (setComp Seq arr) szL) + where + (k, szL) = unconsSz $ size arr +{-# INLINE outerSlices #-} + +-- | Create a delayed array of inner slices. +-- +-- ====__Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> A.mapM_ print $ innerSlices (0 ..: (3 :. 2)) +-- Array D Seq (Sz1 3) +-- [ 0 :. 0, 1 :. 0, 2 :. 0 ] +-- Array D Seq (Sz1 3) +-- [ 0 :. 1, 1 :. 1, 2 :. 1 ] +-- +-- @since 0.5.4 +innerSlices + :: forall r ix e + . (Index ix, Source r e) + => Array r ix e + -> Array D Ix1 (Array D (Lower ix) e) +innerSlices arr = makeArray (getComp arr) k (unsafeInnerSlice (setComp Seq arr) szL) + where + (szL, k) = unsnocSz $ size arr +{-# INLINE innerSlices #-} + +-- | Create a delayed array of slices from within. Checks dimension at compile time. +-- +-- ====__Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> arr = fromIx3 <$> (0 ..: (4 :> 3 :. 2)) +-- >>> print arr +-- Array D Seq (Sz (4 :> 3 :. 2)) +-- [ [ [ (0,0,0), (0,0,1) ] +-- , [ (0,1,0), (0,1,1) ] +-- , [ (0,2,0), (0,2,1) ] +-- ] +-- , [ [ (1,0,0), (1,0,1) ] +-- , [ (1,1,0), (1,1,1) ] +-- , [ (1,2,0), (1,2,1) ] +-- ] +-- , [ [ (2,0,0), (2,0,1) ] +-- , [ (2,1,0), (2,1,1) ] +-- , [ (2,2,0), (2,2,1) ] +-- ] +-- , [ [ (3,0,0), (3,0,1) ] +-- , [ (3,1,0), (3,1,1) ] +-- , [ (3,2,0), (3,2,1) ] +-- ] +-- ] +-- >>> A.mapM_ print $ withinSlices Dim2 arr +-- Array D Seq (Sz (4 :. 2)) +-- [ [ (0,0,0), (0,0,1) ] +-- , [ (1,0,0), (1,0,1) ] +-- , [ (2,0,0), (2,0,1) ] +-- , [ (3,0,0), (3,0,1) ] +-- ] +-- Array D Seq (Sz (4 :. 2)) +-- [ [ (0,1,0), (0,1,1) ] +-- , [ (1,1,0), (1,1,1) ] +-- , [ (2,1,0), (2,1,1) ] +-- , [ (3,1,0), (3,1,1) ] +-- ] +-- Array D Seq (Sz (4 :. 2)) +-- [ [ (0,2,0), (0,2,1) ] +-- , [ (1,2,0), (1,2,1) ] +-- , [ (2,2,0), (2,2,1) ] +-- , [ (3,2,0), (3,2,1) ] +-- ] +-- +-- @since 0.5.4 +withinSlices + :: forall n r ix e + . (IsIndexDimension ix n, Index (Lower ix), Source r e) + => Dimension n + -> Array r ix e + -> Array D Ix1 (Array D (Lower ix) e) +withinSlices dim = either throwImpossible id . withinSlicesM (fromDimension dim) +{-# INLINE withinSlices #-} + +-- | Create a delayed array of slices from within. Same as `withinSlices`, but throws an +-- error on invalid dimension. +-- +-- /__Throws Exceptions__/: `IndexDimensionException` +-- +-- @since 0.5.4 +withinSlicesM + :: forall r ix e m + . (MonadThrow m, Index ix, Index (Lower ix), Source r e) + => Dim + -> Array r ix e + -> m (Array D Ix1 (Array D (Lower ix) e)) +withinSlicesM dim arr = do + (k, szl) <- pullOutSzM (size arr) dim + cutSz <- insertSzM szl dim oneSz + pure $ makeArray Seq k (either throwImpossible id . internalInnerSlice dim cutSz arr) +{-# INLINE withinSlicesM #-} diff --git a/src/Data/Massiv/Array/Ops/Sort.hs b/src/Data/Massiv/Array/Ops/Sort.hs --- a/src/Data/Massiv/Array/Ops/Sort.hs +++ b/src/Data/Massiv/Array/Ops/Sort.hs @@ -1,30 +1,42 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE ExplicitForAll #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MonoLocalBinds #-} + -- | -- Module : Data.Massiv.Array.Ops.Sort --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Ops.Sort - ( tally - , quicksort - , quicksortM_ - , unsafeUnstablePartitionRegionM - ) where +module Data.Massiv.Array.Ops.Sort ( + tally, + quicksort, + quicksortBy, + quicksortByM, + quicksortAs, + quicksortAsBy, + quicksortAsByM, + quicksortM_, + quicksortByM_, + unsafeUnstablePartitionRegionM, +) where import Control.Monad (when) +import Control.Monad.IO.Unlift +import Control.Monad.Primitive import Control.Scheduler +import Data.Bits (countLeadingZeros) import Data.Massiv.Array.Delayed.Stream import Data.Massiv.Array.Mutable import Data.Massiv.Array.Ops.Transform import Data.Massiv.Core.Common +import Data.Massiv.Vector (scatMaybes, sunfoldrN) +import Data.Word (Word64) import System.IO.Unsafe --- | Count how many occurance of each element there is in the array. Results will be +-- | Count number of occurrences of each element in the array. Results will be -- sorted in ascending order of the element. -- -- ==== __Example__ @@ -39,109 +51,196 @@ -- [ (1,1), (2,3), (3,1), (4,2), (5,1) ] -- -- @since 0.4.4 -tally :: (Mutable r Ix1 e, Resize r ix, Load r ix e, Ord e) => Array r ix e -> Array DS Ix1 (e, Int) +tally :: (Manifest r e, Load r ix e, Ord e) => Array r ix e -> Vector DS (e, Int) tally arr | isEmpty arr = setComp (getComp arr) empty - | otherwise = catMaybesS $ unfoldrN (sz + 1) count (0, 0, sorted ! 0) + | otherwise = scatMaybes $ sunfoldrN (liftSz2 (+) sz oneSz) count (0, 0, sorted ! 0) where sz@(Sz k) = size sorted count (!i, !n, !prev) | i < k = - let !e' = unsafeLinearIndex sorted i - in if prev == e' - then Just (Nothing, (i + 1, n + 1, prev)) - else Just (Just (prev, n), (i + 1, 1, e')) + let !e' = unsafeLinearIndex sorted i + in if prev == e' + then Just (Nothing, (i + 1, n + 1, prev)) + else Just (Just (prev, n), (i + 1, 1, e')) | otherwise = Just (Just (prev, n), (i + 1, n, prev)) {-# INLINE count #-} sorted = quicksort $ flatten arr {-# INLINE tally #-} - - -- | Partition a segment of a vector. Starting and ending indices are unchecked. -- --- @since 0.3.2 -unsafeUnstablePartitionRegionM :: - forall r e m. (Mutable r Ix1 e, PrimMonad m) - => MArray (PrimState m) r Ix1 e - -> (e -> Bool) - -> Ix1 -- ^ Start index of the region - -> Ix1 -- ^ End index of the region +-- @since 1.0.0 +unsafeUnstablePartitionRegionM + :: forall r e m + . (Manifest r e, PrimMonad m) + => MVector (PrimState m) r e + -> (e -> m Bool) + -> Ix1 + -- ^ Start index of the region + -> Ix1 + -- ^ End index of the region -> m Ix1 unsafeUnstablePartitionRegionM marr f start end = fromLeft start (end + 1) where fromLeft i j | i == j = pure i | otherwise = do - x <- unsafeRead marr i - if f x - then fromLeft (i + 1) j - else fromRight i (j - 1) + e <- f =<< unsafeLinearRead marr i + if e + then fromLeft (i + 1) j + else fromRight i (j - 1) fromRight i j | i == j = pure i | otherwise = do - x <- unsafeRead marr j - if f x - then do - unsafeWrite marr j =<< unsafeRead marr i - unsafeWrite marr i x - fromLeft (i + 1) j - else fromRight i (j - 1) + x <- unsafeLinearRead marr j + e <- f x + if e + then do + unsafeLinearWrite marr j =<< unsafeLinearRead marr i + unsafeLinearWrite marr i x + fromLeft (i + 1) j + else fromRight i (j - 1) {-# INLINE unsafeUnstablePartitionRegionM #-} +-- | Same as `quicksort` except it accepts any array that is computable. +-- +-- @since 1.0.2 +quicksortAs + :: (Load r Ix1 e, Manifest r' e, Ord e) => r' -> Vector r e -> Vector r' e +quicksortAs _ arr = unsafePerformIO $ withLoadMArray_ arr quicksortM_ +{-# INLINE quicksortAs #-} --- | This is an implementation of [Quicksort](https://en.wikipedia.org/wiki/Quicksort), which is an --- efficient, but unstable sort that uses Median-of-three for pivot choosing, as such it performs --- very well not only for random values, but also for common edge cases like already sorted, --- reversed sorted and arrays with many duplicate elements. It will also respect the computation --- strategy and will result in a nice speed up for systems with multiple CPUs. +-- | Same as `quicksortBy` except it accepts any array that is computable. -- +-- @since 1.0.2 +quicksortAsBy + :: (Load r Ix1 e, Manifest r' e) => r' -> (e -> e -> Ordering) -> Vector r e -> Vector r' e +quicksortAsBy _ f arr = + unsafePerformIO $ withLoadMArray_ arr (quicksortByM_ (\x y -> pure $ f x y)) +{-# INLINE quicksortAsBy #-} + +-- | Same as `quicksortByM` except it accepts any array that is computable. +-- +-- @since 1.0.2 +quicksortAsByM + :: (Load r Ix1 e, Manifest r' e, MonadUnliftIO m) + => r' + -> (e -> e -> m Ordering) + -> Vector r e + -> m (Vector r' e) +quicksortAsByM _ f arr = + withRunInIO $ \run -> withLoadMArray_ arr (quicksortByM_ (\x y -> run (f x y))) +{-# INLINE quicksortAsByM #-} + +-- | This is an implementation of +-- [Quicksort](https://en.wikipedia.org/wiki/Quicksort), which is an efficient, +-- but unstable sort. This implementation uses Median-of-three for pivot +-- choosing, as such it performs very well not only for random values, but also +-- for common edge cases like already sorted, reversed sorted and arrays with +-- many duplicate elements. It will also respect the computation strategy and +-- will result in a nice speed up for systems with multiple CPUs. +-- -- @since 0.3.2 -quicksort :: - (Mutable r Ix1 e, Ord e) => Array r Ix1 e -> Array r Ix1 e -quicksort arr = unsafePerformIO $ withMArray arr quicksortM_ +quicksort + :: (Manifest r e, Ord e) => Vector r e -> Vector r e +quicksort arr = unsafePerformIO $ withMArray_ arr quicksortM_ {-# INLINE quicksort #-} +-- | Same as `quicksortBy`, but instead of `Ord` constraint expects a custom `Ordering`. +-- +-- @since 0.6.1 +quicksortByM + :: (Manifest r e, MonadUnliftIO m) => (e -> e -> m Ordering) -> Vector r e -> m (Vector r e) +quicksortByM f arr = withRunInIO $ \run -> withMArray_ arr (quicksortByM_ (\x y -> run (f x y))) +{-# INLINE quicksortByM #-} +-- | Same as `quicksortBy`, but instead of `Ord` constraint expects a custom `Ordering`. +-- +-- @since 0.6.1 +quicksortBy :: Manifest r e => (e -> e -> Ordering) -> Vector r e -> Vector r e +quicksortBy f arr = + unsafePerformIO $ withMArray_ arr (quicksortByM_ (\x y -> pure $ f x y)) +{-# INLINE quicksortBy #-} --- | Mutable version of `quicksort` +-- | Manifest version of `quicksort` -- -- @since 0.3.2 -quicksortM_ :: - (Ord e, Mutable r Ix1 e, PrimMonad m) - => Scheduler m () - -> MArray (PrimState m) r Ix1 e +quicksortM_ + :: (Ord e, Manifest r e, MonadPrimBase s m) + => Scheduler s () + -> MVector s r e -> m () -quicksortM_ scheduler marr = - scheduleWork scheduler $ qsort (numWorkers scheduler) 0 (unSz (msize marr) - 1) +quicksortM_ = quicksortInternalM_ (\e1 e2 -> pure $ e1 < e2) (\e1 e2 -> pure $ e1 == e2) +{-# INLINE quicksortM_ #-} + +-- | Same as `quicksortM_`, but instead of `Ord` constraint expects a custom `Ordering`. +-- +-- @since 0.6.1 +quicksortByM_ + :: (Manifest r e, MonadPrimBase s m) + => (e -> e -> m Ordering) + -> Scheduler s () + -> MVector s r e + -> m () +quicksortByM_ compareM = + quicksortInternalM_ (\x y -> (LT ==) <$> compareM x y) (\x y -> (EQ ==) <$> compareM x y) +{-# INLINE quicksortByM_ #-} + +quicksortInternalM_ + :: (Manifest r e, MonadPrimBase s m) + => (e -> e -> m Bool) + -> (e -> e -> m Bool) + -> Scheduler s () + -> MVector s r e + -> m () +quicksortInternalM_ fLT fEQ scheduler marr + | numWorkers scheduler < 2 || depthPar <= 0 = qsortSeq 0 (k - 1) + | otherwise = qsortPar depthPar 0 (k - 1) where - leSwap i j = do - ei <- unsafeRead marr i - ej <- unsafeRead marr j - if ei < ej + -- How deep into the search tree should we continue scheduling jobs. Constants below + -- were discovered empirically: + depthPar = min (logNumWorkers + 4) (logSize - 10) + k = unSz (sizeOfMArray marr) + -- We must use log becuase decinding into a tree creates an exponential number of jobs + logNumWorkers = 63 - countLeadingZeros (fromIntegral (numWorkers scheduler) :: Word64) + -- Using many cores on small vectors only makes things slower + logSize = 63 - countLeadingZeros (fromIntegral k :: Word64) + ltSwap i j = do + ei <- unsafeLinearRead marr i + ej <- unsafeLinearRead marr j + lt <- fLT ei ej + if lt then do - unsafeWrite marr i ej - unsafeWrite marr j ei + unsafeLinearWrite marr i ej + unsafeLinearWrite marr j ei pure ei else pure ej - {-# INLINE leSwap #-} + {-# INLINE ltSwap #-} getPivot lo hi = do let !mid = (hi + lo) `div` 2 - _ <- leSwap mid lo - _ <- leSwap hi lo - leSwap mid hi + _ <- ltSwap mid lo + _ <- ltSwap hi lo + ltSwap mid hi {-# INLINE getPivot #-} - qsort !n !lo !hi = + qsortPar !n !lo !hi = when (lo < hi) $ do p <- getPivot lo hi - l <- unsafeUnstablePartitionRegionM marr (< p) lo (hi - 1) - h <- unsafeUnstablePartitionRegionM marr (== p) l hi + l <- unsafeUnstablePartitionRegionM marr (`fLT` p) lo (hi - 1) + h <- unsafeUnstablePartitionRegionM marr (`fEQ` p) l hi if n > 0 then do let !n' = n - 1 - scheduleWork scheduler $ qsort n' lo (l - 1) - scheduleWork scheduler $ qsort n' h hi + scheduleWork scheduler $ qsortPar n' lo (l - 1) + scheduleWork scheduler $ qsortPar n' h hi else do - qsort n lo (l - 1) - qsort n h hi -{-# INLINE quicksortM_ #-} + qsortSeq lo (l - 1) + qsortSeq h hi + qsortSeq !lo !hi = + when (lo < hi) $ do + p <- getPivot lo hi + l <- unsafeUnstablePartitionRegionM marr (`fLT` p) lo (hi - 1) + h <- unsafeUnstablePartitionRegionM marr (`fEQ` p) l hi + qsortSeq lo (l - 1) + qsortSeq h hi +{-# INLINE quicksortInternalM_ #-} diff --git a/src/Data/Massiv/Array/Ops/Transform.hs b/src/Data/Massiv/Array/Ops/Transform.hs --- a/src/Data/Massiv/Array/Ops/Transform.hs +++ b/src/Data/Massiv/Array/Ops/Transform.hs @@ -1,980 +1,1356 @@ {-# LANGUAGE BangPatterns #-} -{-# LANGUAGE ExplicitForAll #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE TypeFamilies #-} -{-# OPTIONS_GHC -fno-warn-redundant-constraints #-} --- | --- Module : Data.Massiv.Array.Ops.Transform --- Copyright : (c) Alexey Kuleshevich 2018-2019 --- License : BSD3 --- Maintainer : Alexey Kuleshevich --- Stability : experimental --- Portability : non-portable --- -module Data.Massiv.Array.Ops.Transform - ( -- ** Transpose - transpose - , transposeInner - , transposeOuter - -- ** Reverse - , reverse - , reverse' - , reverseM - -- ** Backpermute - , backpermuteM - , backpermute' - -- ** Resize - , resizeM - , resize' - , flatten - -- ** Extract - , extractM - , extract' - , extractFromToM - , extractFromTo' - , deleteRowsM - , deleteColumnsM - , deleteRegionM - -- ** Append/Split - , cons - , unconsM - -- , headM - -- , head' - , snoc - , unsnocM - -- , lastM - -- , last' - , appendOuterM - , appendM - , append' - , concatOuterM - , concatM - , concat' - , splitAtM - , splitAt' - , splitExtractM - , takeS - , dropS - -- ** Upsample/Downsample - , upsample - , downsample - -- ** Zoom - , zoom - , zoomWithGrid - -- ** Transform - , transformM - , transform' - , transform2M - , transform2' - ) where - -import Control.Scheduler (traverse_) -import Control.Monad as M (foldM_, unless, forM_) -import Data.Bifunctor (bimap) -import Data.Foldable as F (foldl', foldrM, toList) -import qualified Data.List as L (uncons) -import Data.Massiv.Array.Delayed.Pull -import Data.Massiv.Array.Delayed.Push -import Data.Massiv.Array.Delayed.Stream -import Data.Massiv.Array.Mutable -import Data.Massiv.Array.Ops.Construct -import Data.Massiv.Array.Ops.Map -import Data.Massiv.Core.Common -import Data.Massiv.Core.Index.Internal (Sz(SafeSz)) -import Prelude as P hiding (concat, splitAt, traverse, mapM_, reverse, take, drop) - - --- | Extract a sub-array from within a larger source array. Array that is being extracted must be --- fully encapsulated in a source array, otherwise `SizeSubregionException` will be thrown. -extractM :: (MonadThrow m, Extract r ix e) - => ix -- ^ Starting index - -> Sz ix -- ^ Size of the resulting array - -> Array r ix e -- ^ Source array - -> m (Array (R r) ix e) -extractM !sIx !newSz !arr - | isSafeIndex sz1 sIx && isSafeIndex eIx1 sIx && isSafeIndex sz1 eIx = - pure $ unsafeExtract sIx newSz arr - | otherwise = throwM $ SizeSubregionException (size arr) sIx newSz - where - sz1 = Sz (liftIndex (+1) (unSz (size arr))) - eIx1 = Sz (liftIndex (+1) eIx) - eIx = liftIndex2 (+) sIx $ unSz newSz -{-# INLINE extractM #-} - --- | Same as `extractM`, but will throw a runtime exception from pure code if supplied dimensions --- are incorrect. --- --- @since 0.1.0 -extract' :: Extract r ix e - => ix -- ^ Starting index - -> Sz ix -- ^ Size of the resulting array - -> Array r ix e -- ^ Source array - -> Array (R r) ix e -extract' sIx newSz = either throw id . extractM sIx newSz -{-# INLINE extract' #-} - - --- | Similar to `extractM`, except it takes starting and ending index. Result array will not include --- the ending index. --- --- @since 0.3.0 -extractFromToM :: (MonadThrow m, Extract r ix e) => - ix -- ^ Starting index - -> ix -- ^ Index up to which elements should be extracted. - -> Array r ix e -- ^ Source array. - -> m (Array (R r) ix e) -extractFromToM sIx eIx = extractM sIx (Sz (liftIndex2 (-) eIx sIx)) -{-# INLINE extractFromToM #-} - --- | Same as `extractFromTo`, but throws an error on invalid indices. --- --- @since 0.2.4 -extractFromTo' :: Extract r ix e => - ix -- ^ Starting index - -> ix -- ^ Index up to which elmenets should be extracted. - -> Array r ix e -- ^ Source array. - -> Array (R r) ix e -extractFromTo' sIx eIx = extract' sIx $ Sz (liftIndex2 (-) eIx sIx) -{-# INLINE extractFromTo' #-} - - --- | /O(1)/ - Changes the shape of an array. Returns `Nothing` if total --- number of elements does not match the source array. --- --- @since 0.3.0 -resizeM :: - (MonadThrow m, Index ix', Load r ix e, Resize r ix) - => Sz ix' - -> Array r ix e - -> m (Array r ix' e) -resizeM sz arr = guardNumberOfElements (size arr) sz >> pure (unsafeResize sz arr) -{-# INLINE resizeM #-} - --- | Same as `resizeM`, but will throw an error if supplied dimensions are incorrect. --- --- @since 0.1.0 -resize' :: (Index ix', Load r ix e, Resize r ix) => Sz ix' -> Array r ix e -> Array r ix' e -resize' sz = either throw id . resizeM sz -{-# INLINE resize' #-} - --- | /O(1)/ - Reduce a multi-dimensional array into a flat vector --- --- @since 0.3.1 -flatten :: (Load r ix e, Resize r ix) => Array r ix e -> Array r Ix1 e -flatten arr = unsafeResize (SafeSz (totalElem (size arr))) arr -{-# INLINE flatten #-} - - --- | Transpose a 2-dimensional array --- --- ==== __Examples__ --- --- >>> import Data.Massiv.Array --- >>> arr = makeArrayLinearR D Seq (Sz (2 :. 3)) id --- >>> arr --- Array D Seq (Sz (2 :. 3)) --- [ [ 0, 1, 2 ] --- , [ 3, 4, 5 ] --- ] --- >>> transpose arr --- Array D Seq (Sz (3 :. 2)) --- [ [ 0, 3 ] --- , [ 1, 4 ] --- , [ 2, 5 ] --- ] --- --- @since 0.1.0 -transpose :: Source r Ix2 e => Array r Ix2 e -> Array D Ix2 e -transpose = transposeInner -{-# INLINE [1] transpose #-} - -{-# RULES -"transpose . transpose" [~1] forall arr . transpose (transpose arr) = delay arr -"transposeInner . transposeInner" [~1] forall arr . transposeInner (transposeInner arr) = delay arr -"transposeOuter . transposeOuter" [~1] forall arr . transposeOuter (transposeOuter arr) = delay arr - #-} - - --- | Transpose inner two dimensions of at least rank-2 array. --- --- ===__Examples__ --- --- >>> import Data.Massiv.Array --- >>> arr = makeArrayLinearR U Seq (Sz (2 :> 3 :. 4)) id --- >>> arr --- Array U Seq (Sz (2 :> 3 :. 4)) --- [ [ [ 0, 1, 2, 3 ] --- , [ 4, 5, 6, 7 ] --- , [ 8, 9, 10, 11 ] --- ] --- , [ [ 12, 13, 14, 15 ] --- , [ 16, 17, 18, 19 ] --- , [ 20, 21, 22, 23 ] --- ] --- ] --- >>> transposeInner arr --- Array D Seq (Sz (3 :> 2 :. 4)) --- [ [ [ 0, 1, 2, 3 ] --- , [ 12, 13, 14, 15 ] --- ] --- , [ [ 4, 5, 6, 7 ] --- , [ 16, 17, 18, 19 ] --- ] --- , [ [ 8, 9, 10, 11 ] --- , [ 20, 21, 22, 23 ] --- ] --- ] --- --- @since 0.1.0 -transposeInner :: (Index (Lower ix), Source r' ix e) - => Array r' ix e -> Array D ix e -transposeInner !arr = makeArray (getComp arr) newsz newVal - where - transInner !ix = - either throwImpossible id $ do - n <- getDimM ix dix - m <- getDimM ix (dix - 1) - ix' <- setDimM ix dix m - setDimM ix' (dix - 1) n - {-# INLINE transInner #-} - newVal = unsafeIndex arr . transInner - {-# INLINE newVal #-} - !newsz = Sz (transInner (unSz (size arr))) - !dix = dimensions newsz -{-# INLINE [1] transposeInner #-} - --- | Transpose outer two dimensions of at least rank-2 array. --- --- ====__Examples__ --- --- >>> import Data.Massiv.Array --- >>> :set -XTypeApplications --- >>> arr = makeArrayLinear @U Seq (Sz (2 :> 3 :. 4)) id --- >>> arr --- Array U Seq (Sz (2 :> 3 :. 4)) --- [ [ [ 0, 1, 2, 3 ] --- , [ 4, 5, 6, 7 ] --- , [ 8, 9, 10, 11 ] --- ] --- , [ [ 12, 13, 14, 15 ] --- , [ 16, 17, 18, 19 ] --- , [ 20, 21, 22, 23 ] --- ] --- ] --- >>> transposeOuter arr --- Array D Seq (Sz (2 :> 4 :. 3)) --- [ [ [ 0, 4, 8 ] --- , [ 1, 5, 9 ] --- , [ 2, 6, 10 ] --- , [ 3, 7, 11 ] --- ] --- , [ [ 12, 16, 20 ] --- , [ 13, 17, 21 ] --- , [ 14, 18, 22 ] --- , [ 15, 19, 23 ] --- ] --- ] --- --- --- @since 0.1.0 -transposeOuter :: (Index (Lower ix), Source r' ix e) - => Array r' ix e -> Array D ix e -transposeOuter !arr = makeArray (getComp arr) newsz newVal - where - transOuter !ix = - either throwImpossible id $ do - n <- getDimM ix 1 - m <- getDimM ix 2 - ix' <- setDimM ix 1 m - setDimM ix' 2 n - {-# INLINE transOuter #-} - newVal = unsafeIndex arr . transOuter - {-# INLINE newVal #-} - !newsz = Sz (transOuter (unSz (size arr))) -{-# INLINE [1] transposeOuter #-} - --- | Reverse an array along some dimension. Dimension supplied is checked at compile time. --- --- ==== __Example__ --- --- >>> import Data.Massiv.Array as A --- >>> arr = makeArrayLinear Seq (Sz2 4 5) (+10) :: Array D Ix2 Int --- >>> arr --- Array D Seq (Sz (4 :. 5)) --- [ [ 10, 11, 12, 13, 14 ] --- , [ 15, 16, 17, 18, 19 ] --- , [ 20, 21, 22, 23, 24 ] --- , [ 25, 26, 27, 28, 29 ] --- ] --- >>> A.reverse Dim1 arr --- Array D Seq (Sz (4 :. 5)) --- [ [ 14, 13, 12, 11, 10 ] --- , [ 19, 18, 17, 16, 15 ] --- , [ 24, 23, 22, 21, 20 ] --- , [ 29, 28, 27, 26, 25 ] --- ] --- >>> A.reverse Dim2 arr --- Array D Seq (Sz (4 :. 5)) --- [ [ 25, 26, 27, 28, 29 ] --- , [ 20, 21, 22, 23, 24 ] --- , [ 15, 16, 17, 18, 19 ] --- , [ 10, 11, 12, 13, 14 ] --- ] --- --- @since 0.4.1 -reverse :: (IsIndexDimension ix n, Source r ix e) => Dimension n -> Array r ix e -> Array D ix e -reverse dim = reverse' (fromDimension dim) -{-# INLINE reverse #-} - --- | Similarly to `reverse`, flip an array along a particular dimension, but throws --- `IndexDimensionException` for an incorrect dimension. --- --- @since 0.4.1 -reverseM :: (MonadThrow m, Source r ix e) => Dim -> Array r ix e -> m (Array D ix e) -reverseM dim arr = do - let sz = size arr - k <- getDimM (unSz sz) dim - pure $ makeArray (getComp arr) sz $ \ ix -> - unsafeIndex arr (snd $ modifyDim' ix dim (\i -> k - i - 1)) -{-# INLINE reverseM #-} - --- | Reverse an array along some dimension. Same as `reverseM`, but throws the --- `IndexDimensionException` from pure code. --- --- @since 0.4.1 -reverse' :: Source r ix e => Dim -> Array r ix e -> Array D ix e -reverse' dim = either throw id . reverseM dim -{-# INLINE reverse' #-} - --- | Rearrange elements of an array into a new one by using a function that maps indices of the --- newly created one into the old one. This function can throw `IndexOutOfBoundsException`. --- --- ===__Examples__ --- --- >>> import Data.Massiv.Array --- >>> :set -XTypeApplications --- >>> arr = makeArrayLinear @D Seq (Sz (2 :> 3 :. 4)) id --- >>> arr --- Array D Seq (Sz (2 :> 3 :. 4)) --- [ [ [ 0, 1, 2, 3 ] --- , [ 4, 5, 6, 7 ] --- , [ 8, 9, 10, 11 ] --- ] --- , [ [ 12, 13, 14, 15 ] --- , [ 16, 17, 18, 19 ] --- , [ 20, 21, 22, 23 ] --- ] --- ] --- >>> backpermuteM @U (Sz (4 :. 2)) (\(i :. j) -> j :> j :. i) arr --- Array U Seq (Sz (4 :. 2)) --- [ [ 0, 16 ] --- , [ 1, 17 ] --- , [ 2, 18 ] --- , [ 3, 19 ] --- ] --- --- @since 0.3.0 -backpermuteM :: - forall r ix e r' ix' m. - (Mutable r ix e, Source r' ix' e, MonadUnliftIO m, PrimMonad m, MonadThrow m) - => Sz ix -- ^ Size of the result array - -> (ix -> ix') -- ^ A function that maps indices of the new array into the source one. - -> Array r' ix' e -- ^ Source array. - -> m (Array r ix e) -backpermuteM sz ixF !arr = generateArray (getComp arr) sz (evaluateM arr . ixF) -{-# INLINE backpermuteM #-} - --- | Similar to `backpermuteM`, with a few notable differences: --- --- * Creates a delayed array, instead of manifest, therefore it can be fused --- * Respects computation strategy, so it can be parallelized --- * Throws a runtime `IndexOutOfBoundsException` from pure code. --- --- @since 0.3.0 -backpermute' :: (Source r' ix' e, Index ix) => - Sz ix -- ^ Size of the result array - -> (ix -> ix') -- ^ A function that maps indices of the new array into the source one. - -> Array r' ix' e -- ^ Source array. - -> Array D ix e -backpermute' sz ixF !arr = makeArray (getComp arr) sz (evaluate' arr . ixF) -{-# INLINE backpermute' #-} - - --- | /O(1)/ - Add an element to the vector from the left side --- --- @since 0.3.0 -cons :: e -> Array DL Ix1 e -> Array DL Ix1 e -cons e arr = - arr - { dlSize = SafeSz (1 + unSz (dlSize arr)) - , dlLoad = - \scheduler startAt uWrite -> - uWrite startAt e >> dlLoad arr scheduler (startAt + 1) uWrite - } -{-# INLINE cons #-} - - --- -- | /O(1)/ - Take the first element off the vector from the left side. --- -- --- -- @since 0.4.3 --- headM :: (MonadThrow m, Source r Ix1 e) => Array r Ix1 e -> m e --- headM = fmap fst . unconsM --- {-# INLINE headM #-} - --- -- | /O(1)/ - Take the first element off the vector from the left side. Throws --- -- `SizeEmptyException` --- -- --- -- @since 0.4.3 --- head' :: Source r Ix1 e => Array r Ix1 e -> e --- head' = either throw id . headM --- {-# INLINE head' #-} - - --- -- | /O(1)/ - Take the last element off the vector from the right side. --- -- --- -- @since 0.4.3 --- lastM :: (MonadThrow m, Source r Ix1 e) => Array r Ix1 e -> m e --- lastM = fmap snd . unsnocM --- {-# INLINE lastM #-} - --- -- | /O(1)/ - Take the last element off the vector from the right side. Throws --- -- `SizeEmptyException` --- -- --- -- @since 0.4.3 --- last' :: Source r Ix1 e => Array r Ix1 e -> e --- last' = either throw id . lastM --- {-# INLINE last' #-} - - - --- | /O(1)/ - Take one element off the vector from the left side. --- --- @since 0.3.0 -unconsM :: (MonadThrow m, Source r Ix1 e) => Array r Ix1 e -> m (e, Array D Ix1 e) -unconsM arr - | 0 == totalElem sz = throwM $ SizeEmptyException sz - | otherwise = - pure - ( unsafeLinearIndex arr 0 - , makeArray (getComp arr) (SafeSz (unSz sz - 1)) (\ !i -> unsafeLinearIndex arr (i + 1))) - where - !sz = size arr -{-# INLINE unconsM #-} - --- | /O(1)/ - Add an element to the vector from the right side --- --- @since 0.3.0 -snoc :: Array DL Ix1 e -> e -> Array DL Ix1 e -snoc arr e = - arr - { dlSize = SafeSz (1 + k) - , dlLoad = - \scheduler startAt uWrite -> dlLoad arr scheduler startAt uWrite >> uWrite (k + startAt) e - } - where - !k = unSz (size arr) -{-# INLINE snoc #-} - - --- | /O(1)/ - Take one element off the vector from the right side. --- --- @since 0.3.0 -unsnocM :: (MonadThrow m, Source r Ix1 e) => Array r Ix1 e -> m (Array D Ix1 e, e) -unsnocM arr - | k < 0 = throwM $ SizeEmptyException sz - | otherwise = - pure (makeArray (getComp arr) (SafeSz k) (unsafeLinearIndex arr), unsafeLinearIndex arr k) - where - !sz = size arr - !k = unSz sz - 1 -{-# INLINE unsnocM #-} - - - --- | Append two arrays together along a particular dimension. Sizes of both arrays must match, with --- an allowed exception of the dimension they are being appended along, otherwise `Nothing` is --- returned. --- --- ===__Examples__ --- --- Append two 2D arrays along both dimensions. Note that they do agree on inner dimensions. --- --- >>> import Data.Massiv.Array --- >>> arrA = makeArrayR U Seq (Sz2 2 3) (\(i :. j) -> ('A', i, j)) --- >>> arrB = makeArrayR U Seq (Sz2 2 3) (\(i :. j) -> ('B', i, j)) --- >>> appendM 1 arrA arrB --- Array DL Seq (Sz (2 :. 6)) --- [ [ ('A',0,0), ('A',0,1), ('A',0,2), ('B',0,0), ('B',0,1), ('B',0,2) ] --- , [ ('A',1,0), ('A',1,1), ('A',1,2), ('B',1,0), ('B',1,1), ('B',1,2) ] --- ] --- >>> appendM 2 arrA arrB --- Array DL Seq (Sz (4 :. 3)) --- [ [ ('A',0,0), ('A',0,1), ('A',0,2) ] --- , [ ('A',1,0), ('A',1,1), ('A',1,2) ] --- , [ ('B',0,0), ('B',0,1), ('B',0,2) ] --- , [ ('B',1,0), ('B',1,1), ('B',1,2) ] --- ] --- --- Now appending arrays with different sizes: --- --- >>> arrC = makeArrayR U Seq (Sz (2 :. 4)) (\(i :. j) -> ('C', i, j)) --- >>> appendM 1 arrA arrC --- Array DL Seq (Sz (2 :. 7)) --- [ [ ('A',0,0), ('A',0,1), ('A',0,2), ('C',0,0), ('C',0,1), ('C',0,2), ('C',0,3) ] --- , [ ('A',1,0), ('A',1,1), ('A',1,2), ('C',1,0), ('C',1,1), ('C',1,2), ('C',1,3) ] --- ] --- >>> appendM 2 arrA arrC --- *** Exception: SizeMismatchException: (Sz (2 :. 3)) vs (Sz (2 :. 4)) --- --- @since 0.3.0 -appendM :: (MonadThrow m, Source r1 ix e, Source r2 ix e) => - Dim -> Array r1 ix e -> Array r2 ix e -> m (Array DL ix e) -appendM n !arr1 !arr2 = do - let !sz1 = size arr1 - !sz2 = size arr2 - (k1, szl1) <- pullOutSzM sz1 n - (k2, szl2) <- pullOutSzM sz2 n - unless (szl1 == szl2) $ throwM $ SizeMismatchException sz1 sz2 - let k1' = unSz k1 - newSz <- insertSzM szl1 n (SafeSz (k1' + unSz k2)) - return $ - DLArray - { dlComp = getComp arr1 <> getComp arr2 - , dlSize = newSz - , dlDefault = Nothing - , dlLoad = - \scheduler startAt dlWrite -> do - scheduleWork scheduler $ - iterM_ zeroIndex (unSz sz1) (pureIndex 1) (<) $ \ix -> - dlWrite (startAt + toLinearIndex newSz ix) (unsafeIndex arr1 ix) - scheduleWork scheduler $ - iterM_ zeroIndex (unSz sz2) (pureIndex 1) (<) $ \ix -> - let i = getDim' ix n - ix' = setDim' ix n (i + k1') - in dlWrite (startAt + toLinearIndex newSz ix') (unsafeIndex arr2 ix) - } -{-# INLINE appendM #-} - - --- | Same as `appendM`, but will throw an exception in pure code on mismatched sizes. --- --- @since 0.3.0 -append' :: (Source r1 ix e, Source r2 ix e) => - Dim -> Array r1 ix e -> Array r2 ix e -> Array DL ix e -append' dim arr1 arr2 = either throw id $ appendM dim arr1 arr2 -{-# INLINE append' #-} - --- | Concat many arrays together along some dimension. --- --- @since 0.3.0 -concat' :: (Foldable f, Source r ix e) => Dim -> f (Array r ix e) -> Array DL ix e -concat' n arrs = either throw id $ concatM n arrs -{-# INLINE concat' #-} - --- | Concatenate many arrays together along some dimension. It is important that all sizes are --- equal, with an exception of the dimensions along which concatenation happens, otherwise it doues --- result in a `SizeMismatchException` exception. --- --- @since 0.3.0 -concatM :: - (MonadThrow m, Foldable f, Source r ix e) => Dim -> f (Array r ix e) -> m (Array DL ix e) -concatM n !arrsF = - case L.uncons (F.toList arrsF) of - Nothing -> pure empty - Just (a, arrs) -> do - let sz = unSz (size a) - szs = P.map (unSz . size) arrs - (k, szl) <- pullOutDimM sz n - -- / remove the dimension out of all sizes along which concatenation will happen - (ks, szls) <- - F.foldrM (\ !csz (ks, szls) -> bimap (: ks) (: szls) <$> pullOutDimM csz n) ([], []) szs - -- / make sure to fail as soon as at least one of the arrays has a mismatching inner size - traverse_ - (\(sz', _) -> throwM (SizeMismatchException (SafeSz sz) (SafeSz sz'))) - (dropWhile ((== szl) . snd) $ P.zip szs szls) - let kTotal = SafeSz $ F.foldl' (+) k ks - newSz <- insertSzM (SafeSz szl) n kTotal - return $ - DLArray - { dlComp = mconcat $ P.map getComp arrs - , dlSize = newSz - , dlDefault = Nothing - , dlLoad = - \scheduler startAt dlWrite -> - let arrayLoader !kAcc (kCur, arr) = do - scheduleWork scheduler $ - iterM_ zeroIndex (unSz (size arr)) (pureIndex 1) (<) $ \ix -> - let i = getDim' ix n - ix' = setDim' ix n (i + kAcc) - in dlWrite (startAt + toLinearIndex newSz ix') (unsafeIndex arr ix) - pure (kAcc + kCur) - in M.foldM_ arrayLoader 0 $ (k, a) : P.zip ks arrs - } -{-# INLINE concatM #-} - - --- | /O(1)/ - Split an array at an index along a specified dimension. --- --- @since 0.3.0 -splitAtM :: - (MonadThrow m, Extract r ix e) - => Dim -- ^ Dimension along which to split - -> Int -- ^ Index along the dimension to split at - -> Array r ix e -- ^ Source array - -> m (Array (R r) ix e, Array (R r) ix e) -splitAtM dim i arr = do - let Sz sz = size arr - eIx <- setDimM sz dim i - sIx <- setDimM zeroIndex dim i - arr1 <- extractFromToM zeroIndex eIx arr - arr2 <- extractFromToM sIx sz arr - return (arr1, arr2) -{-# INLINE splitAtM #-} - --- | Same as `splitAt`, but will throw an error instead of returning `Nothing` on wrong dimension --- and index out of bounds. --- --- @since 0.1.0 -splitAt' :: Extract r ix e => - Dim -> Int -> Array r ix e -> (Array (R r) ix e, Array (R r) ix e) -splitAt' dim i arr = either throw id $ splitAtM dim i arr -{-# INLINE splitAt' #-} - - --- | Split an array in three parts across some dimension --- --- @since 0.3.5 -splitExtractM :: - (MonadThrow m, Extract r ix e, Source (R r) ix e) - => Dim -- ^ Dimension along which to do the extraction - -> Ix1 -- ^ Start index along the dimension that needs to be extracted - -> Sz Ix1 -- ^ Size of the extracted array along the dimension that it will be extracted - -> Array r ix e - -> m (Array (R r) ix e, Array (R r) ix e, Array (R r) ix e) -splitExtractM dim startIx1 (Sz extractSzIx1) arr = do - let Sz szIx = size arr - midStartIx <- setDimM zeroIndex dim startIx1 - midExtractSzIx <- setDimM szIx dim extractSzIx1 - midArr <- extractM midStartIx (Sz midExtractSzIx) arr - leftArrSzIx <- setDimM szIx dim startIx1 - leftArr <- extractM zeroIndex (Sz leftArrSzIx) arr - rightArrStartIx <- setDimM zeroIndex dim (startIx1 + extractSzIx1) - rightArr <- extractFromToM rightArrStartIx szIx arr - pure (leftArr, midArr, rightArr) -{-# INLINE splitExtractM #-} - --- | Delete a region from an array along the specified dimension. --- --- ==== __Examples__ --- --- >>> import Data.Massiv.Array --- >>> arr = fromIx3 <$> (0 :> 0 :. 0 ..: 3 :> 2 :. 6) --- >>> deleteRegionM 1 2 3 arr --- Array DL Seq (Sz (3 :> 2 :. 3)) --- [ [ [ (0,0,0), (0,0,1), (0,0,5) ] --- , [ (0,1,0), (0,1,1), (0,1,5) ] --- ] --- , [ [ (1,0,0), (1,0,1), (1,0,5) ] --- , [ (1,1,0), (1,1,1), (1,1,5) ] --- ] --- , [ [ (2,0,0), (2,0,1), (2,0,5) ] --- , [ (2,1,0), (2,1,1), (2,1,5) ] --- ] --- ] --- >>> v = Ix1 0 ... 10 --- >>> deleteRegionM 1 3 5 v --- Array DL Seq (Sz1 6) --- [ 0, 1, 2, 8, 9, 10 ] --- --- @since 0.3.5 -deleteRegionM :: - (MonadThrow m, Extract r ix e, Source (R r) ix e) - => Dim -- ^ Along which axis should the removal happen - -> Ix1 -- ^ At which index to start dropping slices - -> Sz Ix1 -- ^ Number of slices to drop - -> Array r ix e -- ^ Array that will have it's subarray removed - -> m (Array DL ix e) -deleteRegionM dim ix sz arr = do - (leftArr, _, rightArr) <- splitExtractM dim ix sz arr - appendM dim leftArr rightArr -{-# INLINE deleteRegionM #-} - --- | Similar to `deleteRegionM`, but drop a specified number of rows from an array that --- has at least 2 dimensions. --- --- ====__Example__ --- --- >>> import Data.Massiv.Array --- >>> arr = fromIx2 <$> (0 :. 0 ..: 3 :. 6) --- >>> arr --- Array D Seq (Sz (3 :. 6)) --- [ [ (0,0), (0,1), (0,2), (0,3), (0,4), (0,5) ] --- , [ (1,0), (1,1), (1,2), (1,3), (1,4), (1,5) ] --- , [ (2,0), (2,1), (2,2), (2,3), (2,4), (2,5) ] --- ] --- >>> deleteRowsM 1 1 arr --- Array DL Seq (Sz (2 :. 6)) --- [ [ (0,0), (0,1), (0,2), (0,3), (0,4), (0,5) ] --- , [ (2,0), (2,1), (2,2), (2,3), (2,4), (2,5) ] --- ] --- --- @since 0.3.5 -deleteRowsM :: - (MonadThrow m, Extract r ix e, Source (R r) ix e, Index (Lower ix)) - => Ix1 - -> Sz Ix1 - -> Array r ix e - -> m (Array DL ix e) -deleteRowsM = deleteRegionM 2 -{-# INLINE deleteRowsM #-} - --- | Similar to `deleteRegionM`, but drop a specified number of columns an array. --- --- ====__Example__ --- --- >>> import Data.Massiv.Array --- >>> arr = fromIx2 <$> (0 :. 0 ..: 3 :. 6) --- >>> arr --- Array D Seq (Sz (3 :. 6)) --- [ [ (0,0), (0,1), (0,2), (0,3), (0,4), (0,5) ] --- , [ (1,0), (1,1), (1,2), (1,3), (1,4), (1,5) ] --- , [ (2,0), (2,1), (2,2), (2,3), (2,4), (2,5) ] --- ] --- >>> deleteColumnsM 2 3 arr --- Array DL Seq (Sz (3 :. 3)) --- [ [ (0,0), (0,1), (0,5) ] --- , [ (1,0), (1,1), (1,5) ] --- , [ (2,0), (2,1), (2,5) ] --- ] --- --- @since 0.3.5 -deleteColumnsM :: - (MonadThrow m, Extract r ix e, Source (R r) ix e) - => Ix1 - -> Sz Ix1 - -> Array r ix e - -> m (Array DL ix e) -deleteColumnsM = deleteRegionM 1 -{-# INLINE deleteColumnsM #-} - - --- | Discard elements from the source array according to the stride. --- --- @since 0.3.0 -downsample :: Source r ix e => Stride ix -> Array r ix e -> Array DL ix e -downsample stride arr = - DLArray - { dlComp = getComp arr - , dlSize = resultSize - , dlDefault = defaultElement arr - , dlLoad = - \scheduler startAt dlWrite -> - splitLinearlyWithStartAtM_ - scheduler - startAt - (totalElem resultSize) - (pure . unsafeLinearWriteWithStride) - dlWrite - } - where - resultSize = strideSize stride (size arr) - strideIx = unStride stride - unsafeLinearWriteWithStride = - unsafeIndex arr . liftIndex2 (*) strideIx . fromLinearIndex resultSize - {-# INLINE unsafeLinearWriteWithStride #-} -{-# INLINE downsample #-} - - --- | Insert the same element into a `Load`able array according to the stride. --- --- @since 0.3.0 -upsample - :: Load r ix e => e -> Stride ix -> Array r ix e -> Array DL ix e -upsample !fillWith safeStride arr = - DLArray - { dlComp = getComp arr - , dlSize = newsz - , dlDefault = Just fillWith - , dlLoad = - \scheduler startAt dlWrite -> do - M.forM_ (defaultElement arr) $ \prevFillWith -> - loopM_ - startAt - (< totalElem sz) - (+ 1) - (\i -> dlWrite (adjustLinearStride (i + startAt)) prevFillWith) - loadArrayM scheduler arr (\i -> dlWrite (adjustLinearStride (i + startAt))) - } - where - adjustLinearStride = toLinearIndex newsz . timesStride . fromLinearIndex sz - {-# INLINE adjustLinearStride #-} - timesStride !ix = liftIndex2 (*) stride ix - {-# INLINE timesStride #-} - !stride = unStride safeStride - !sz = size arr - !newsz = SafeSz (timesStride $ unSz sz) -{-# INLINE upsample #-} - - --- | General array transformation, that forces computation and produces a manifest array. --- --- @since 0.3.0 -transformM :: - forall r ix e r' ix' e' a m. - (Mutable r ix e, Source r' ix' e', MonadUnliftIO m, PrimMonad m, MonadThrow m) - => (Sz ix' -> m (Sz ix, a)) - -> (a -> (ix' -> m e') -> ix -> m e) - -> Array r' ix' e' - -> m (Array r ix e) -transformM getSzM getM arr = do - (sz, a) <- getSzM (size arr) - generateArray (getComp arr) sz (getM a (evaluateM arr)) -{-# INLINE transformM #-} - - --- | General array transformation --- --- @since 0.3.0 -transform' :: - (Source r' ix' e', Index ix) - => (Sz ix' -> (Sz ix, a)) - -> (a -> (ix' -> e') -> ix -> e) - -> Array r' ix' e' - -> Array D ix e -transform' getSz get arr = makeArray (getComp arr) sz (get a (evaluate' arr)) - where - (sz, a) = getSz (size arr) -{-# INLINE transform' #-} - --- | Same as `transformM`, but operates on two arrays --- --- @since 0.3.0 -transform2M :: - (Mutable r ix e, Source r1 ix1 e1, Source r2 ix2 e2, MonadUnliftIO m, PrimMonad m, MonadThrow m) - => (Sz ix1 -> Sz ix2 -> m (Sz ix, a)) - -> (a -> (ix1 -> m e1) -> (ix2 -> m e2) -> ix -> m e) - -> Array r1 ix1 e1 - -> Array r2 ix2 e2 - -> m (Array r ix e) -transform2M getSzM getM arr1 arr2 = do - (sz, a) <- getSzM (size arr1) (size arr2) - generateArray (getComp arr1 <> getComp arr2) sz (getM a (evaluateM arr1) (evaluateM arr2)) -{-# INLINE transform2M #-} - - --- | Same as `transform'`, but operates on two arrays --- --- @since 0.3.0 -transform2' :: - (Source r1 ix1 e1, Source r2 ix2 e2, Index ix) - => (Sz ix1 -> Sz ix2 -> (Sz ix, a)) - -> (a -> (ix1 -> e1) -> (ix2 -> e2) -> ix -> e) - -> Array r1 ix1 e1 - -> Array r2 ix2 e2 - -> Array D ix e -transform2' getSz get arr1 arr2 = - makeArray (getComp arr1 <> getComp arr2) sz (get a (evaluate' arr1) (evaluate' arr2)) - where - (sz, a) = getSz (size arr1) (size arr2) -{-# INLINE transform2' #-} - - - --- | Replicate each element of the array by a factor in stride along each dimension and surround each --- such group with a box of supplied grid value. It will essentially zoom up an array and create a --- grid around each element from the original array. Very useful for zooming up images to inspect --- individual pixels. --- --- ==== __Example__ --- --- >>> import Data.Massiv.Array as A --- >>> arr = resize' (Sz2 3 2) (Ix1 1 ... 6) --- >>> arr --- Array D Seq (Sz (3 :. 2)) --- [ [ 1, 2 ] --- , [ 3, 4 ] --- , [ 5, 6 ] --- ] --- >>> zoomWithGrid 0 (Stride (2 :. 3)) arr --- Array DL Seq (Sz (10 :. 9)) --- [ [ 0, 0, 0, 0, 0, 0, 0, 0, 0 ] --- , [ 0, 1, 1, 1, 0, 2, 2, 2, 0 ] --- , [ 0, 1, 1, 1, 0, 2, 2, 2, 0 ] --- , [ 0, 0, 0, 0, 0, 0, 0, 0, 0 ] --- , [ 0, 3, 3, 3, 0, 4, 4, 4, 0 ] --- , [ 0, 3, 3, 3, 0, 4, 4, 4, 0 ] --- , [ 0, 0, 0, 0, 0, 0, 0, 0, 0 ] --- , [ 0, 5, 5, 5, 0, 6, 6, 6, 0 ] --- , [ 0, 5, 5, 5, 0, 6, 6, 6, 0 ] --- , [ 0, 0, 0, 0, 0, 0, 0, 0, 0 ] --- ] --- --- @since 0.3.1 -zoomWithGrid :: - Source r ix e - => e -- ^ Value to use for the grid - -> Stride ix -- ^ Scaling factor - -> Array r ix e -- ^ Source array - -> Array DL ix e -zoomWithGrid gridVal (Stride zoomFactor) arr = - unsafeMakeLoadArray Seq newSz (Just gridVal) $ \scheduler _ writeElement -> - iforSchedulerM_ scheduler arr $ \ !ix !e -> do - let !kix = liftIndex2 (*) ix kx - mapM_ (\ !ix' -> writeElement (toLinearIndex newSz ix') e) $ - range Seq (liftIndex (+1) kix) (liftIndex2 (+) kix kx) - where - !kx = liftIndex (+1) zoomFactor - !lastNewIx = liftIndex2 (*) kx $ unSz (size arr) - !newSz = Sz (liftIndex (+1) lastNewIx) -{-# INLINE zoomWithGrid #-} - --- | Increaze the size of the array accoridng to the stride multiplier while replicating --- the same element to fill the neighbors. It is exactly the same as `zoomWithGrid`, but --- without the grid. --- --- ==== __Example__ --- --- >>> import Data.Massiv.Array as A --- >>> arr = resize' (Sz3 1 3 2) (Ix1 1 ... 6) --- >>> arr --- Array D Seq (Sz (1 :> 3 :. 2)) --- [ [ [ 1, 2 ] --- , [ 3, 4 ] --- , [ 5, 6 ] --- ] --- ] --- >>> zoom (Stride (2 :> 2 :. 3)) arr --- Array DL Seq (Sz (2 :> 6 :. 6)) --- [ [ [ 1, 1, 1, 2, 2, 2 ] --- , [ 1, 1, 1, 2, 2, 2 ] --- , [ 3, 3, 3, 4, 4, 4 ] --- , [ 3, 3, 3, 4, 4, 4 ] --- , [ 5, 5, 5, 6, 6, 6 ] --- , [ 5, 5, 5, 6, 6, 6 ] --- ] --- , [ [ 1, 1, 1, 2, 2, 2 ] --- , [ 1, 1, 1, 2, 2, 2 ] --- , [ 3, 3, 3, 4, 4, 4 ] --- , [ 3, 3, 3, 4, 4, 4 ] --- , [ 5, 5, 5, 6, 6, 6 ] --- , [ 5, 5, 5, 6, 6, 6 ] --- ] --- ] --- --- @since 0.4.4 -zoom :: - Source r ix e - => Stride ix -- ^ Scaling factor - -> Array r ix e -- ^ Source array - -> Array DL ix e -zoom (Stride zoomFactor) arr = - unsafeMakeLoadArray Seq newSz Nothing $ \scheduler _ writeElement -> - iforSchedulerM_ scheduler arr $ \ !ix !e -> do - let !kix = liftIndex2 (*) ix zoomFactor - mapM_ (\ !ix' -> writeElement (toLinearIndex newSz ix') e) $ - range Seq kix (liftIndex2 (+) kix zoomFactor) - where - !lastNewIx = liftIndex2 (*) zoomFactor $ unSz (size arr) - !newSz = Sz lastNewIx +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# OPTIONS_GHC -fno-warn-redundant-constraints #-} + +-- | +-- Module : Data.Massiv.Array.Ops.Transform +-- Copyright : (c) Alexey Kuleshevich 2018-2022 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +module Data.Massiv.Array.Ops.Transform ( + -- ** Transpose + transpose, + transposeInner, + transposeOuter, + + -- ** Reverse + reverse, + reverse', + reverseM, + + -- ** Backpermute + backpermuteM, + backpermute', + + -- ** Resize + resizeM, + resize', + flatten, + + -- ** Extract + extractM, + extract', + extractFromToM, + extractFromTo', + deleteRowsM, + deleteColumnsM, + deleteRegionM, + + -- ** Append/Split + appendOuterM, + appendM, + append', + concatOuterM, + concatM, + concat', + stackSlicesM, + stackOuterSlicesM, + stackInnerSlicesM, + splitAtM, + splitAt', + splitExtractM, + replaceSlice, + replaceOuterSlice, + + -- ** Upsample/Downsample + upsample, + downsample, + + -- ** Zoom + zoom, + zoomWithGrid, + + -- ** Transform + transformM, + transform', + transform2M, + transform2', +) where + +import Control.Monad as M (foldM_, forM_, unless) +import Control.Monad.ST +import Control.Scheduler (traverse_) +import Data.Bifunctor (bimap) +import Data.Foldable as F (foldl', foldrM, length, toList) +import qualified Data.List as L (uncons) +import Data.Massiv.Array.Delayed.Pull +import Data.Massiv.Array.Delayed.Push +import Data.Massiv.Array.Mutable +import Data.Massiv.Array.Ops.Construct +import Data.Massiv.Array.Ops.Map +import Data.Massiv.Core +import Data.Massiv.Core.Common -- (size, unsafeIndex, unsafeResize, evaluate', evaluateM) +import Prelude as P hiding ( + concat, + drop, + mapM_, + reverse, + splitAt, + take, + traverse, + ) + +-- | Extract a sub-array from within a larger source array. Array that is being extracted must be +-- fully encapsulated in a source array, otherwise `SizeSubregionException` will be thrown. +-- +-- ====__Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> m <- resizeM (Sz (3 :. 3)) $ Ix1 1 ... 9 +-- >>> m +-- Array D Seq (Sz (3 :. 3)) +-- [ [ 1, 2, 3 ] +-- , [ 4, 5, 6 ] +-- , [ 7, 8, 9 ] +-- ] +-- >>> extractM (0 :. 1) (Sz (2 :. 2)) m +-- Array D Seq (Sz (2 :. 2)) +-- [ [ 2, 3 ] +-- , [ 5, 6 ] +-- ] +-- >>> a <- resizeM (Sz (3 :> 2 :. 4)) $ Ix1 11 ... 34 +-- >>> a +-- Array D Seq (Sz (3 :> 2 :. 4)) +-- [ [ [ 11, 12, 13, 14 ] +-- , [ 15, 16, 17, 18 ] +-- ] +-- , [ [ 19, 20, 21, 22 ] +-- , [ 23, 24, 25, 26 ] +-- ] +-- , [ [ 27, 28, 29, 30 ] +-- , [ 31, 32, 33, 34 ] +-- ] +-- ] +-- >>> extractM (0 :> 1 :. 1) (Sz (3 :> 1 :. 2)) a +-- Array D Seq (Sz (3 :> 1 :. 2)) +-- [ [ [ 16, 17 ] +-- ] +-- , [ [ 24, 25 ] +-- ] +-- , [ [ 32, 33 ] +-- ] +-- ] +-- +-- @since 0.3.0 +extractM + :: forall r ix e m + . (MonadThrow m, Index ix, Source r e) + => ix + -- ^ Starting index + -> Sz ix + -- ^ Size of the resulting array + -> Array r ix e + -- ^ Source array + -> m (Array D ix e) +extractM !sIx !newSz !arr + | isSafeIndex sz1 sIx && isSafeIndex eIx1 sIx && isSafeIndex sz1 eIx = + pure $ unsafeExtract sIx newSz arr + | otherwise = throwM $ SizeSubregionException (size arr) sIx newSz + where + sz1 = Sz (liftIndex (+ 1) (unSz (size arr))) + eIx1 = Sz (liftIndex (+ 1) eIx) + eIx = liftIndex2 (+) sIx $ unSz newSz +{-# INLINE extractM #-} + +-- | Same as `extractM`, but will throw a runtime exception from pure code if supplied dimensions +-- are incorrect. +-- +-- @since 0.1.0 +extract' + :: forall r ix e + . (HasCallStack, Index ix, Source r e) + => ix + -- ^ Starting index + -> Sz ix + -- ^ Size of the resulting array + -> Array r ix e + -- ^ Source array + -> Array D ix e +extract' sIx newSz = throwEither . extractM sIx newSz +{-# INLINE extract' #-} + +-- | Similar to `extractM`, except it takes starting and ending index. Result array will not include +-- the ending index. +-- +-- ====__Examples__ +-- +-- >>> a <- resizeM (Sz (3 :> 2 :. 4)) $ Ix1 11 ... 34 +-- >>> a +-- Array D Seq (Sz (3 :> 2 :. 4)) +-- [ [ [ 11, 12, 13, 14 ] +-- , [ 15, 16, 17, 18 ] +-- ] +-- , [ [ 19, 20, 21, 22 ] +-- , [ 23, 24, 25, 26 ] +-- ] +-- , [ [ 27, 28, 29, 30 ] +-- , [ 31, 32, 33, 34 ] +-- ] +-- ] +-- >>> extractFromToM (1 :> 0 :. 1) (3 :> 2 :. 4) a +-- Array D Seq (Sz (2 :> 2 :. 3)) +-- [ [ [ 20, 21, 22 ] +-- , [ 24, 25, 26 ] +-- ] +-- , [ [ 28, 29, 30 ] +-- , [ 32, 33, 34 ] +-- ] +-- ] +-- +-- @since 0.3.0 +extractFromToM + :: forall r ix e m + . (MonadThrow m, Index ix, Source r e) + => ix + -- ^ Starting index + -> ix + -- ^ Index up to which elements should be extracted. + -> Array r ix e + -- ^ Source array. + -> m (Array D ix e) +extractFromToM sIx eIx = extractM sIx (Sz (liftIndex2 (-) eIx sIx)) +{-# INLINE extractFromToM #-} + +-- | Same as `extractFromToM`, but throws an error on invalid indices. +-- +-- @since 0.2.4 +extractFromTo' + :: forall r ix e + . (HasCallStack, Index ix, Source r e) + => ix + -- ^ Starting index + -> ix + -- ^ Index up to which elmenets should be extracted. + -> Array r ix e + -- ^ Source array. + -> Array D ix e +extractFromTo' sIx eIx = extract' sIx $ Sz (liftIndex2 (-) eIx sIx) +{-# INLINE extractFromTo' #-} + +-- | /O(1)/ - Change the size of an array. Throws +-- `SizeElementsMismatchException` if total number of elements does not match +-- the supplied array. +-- +-- @since 0.3.0 +resizeM + :: forall r ix ix' e m + . (MonadThrow m, Index ix', Index ix, Size r) + => Sz ix' + -> Array r ix e + -> m (Array r ix' e) +resizeM sz arr = guardNumberOfElements (size arr) sz >> pure (unsafeResize sz arr) +{-# INLINE resizeM #-} + +-- | Same as `resizeM`, but will throw an error if supplied dimensions are incorrect. +-- +-- @since 0.1.0 +resize' + :: forall r ix ix' e + . (HasCallStack, Index ix', Index ix, Size r) + => Sz ix' + -> Array r ix e + -> Array r ix' e +resize' sz = throwEither . resizeM sz +{-# INLINE resize' #-} + +-- | /O(1)/ - Reduce a multi-dimensional array into a flat vector +-- +-- @since 0.3.1 +flatten :: forall r ix e. (Index ix, Size r) => Array r ix e -> Vector r e +flatten arr = unsafeResize (SafeSz (totalElem (size arr))) arr +{-# INLINE flatten #-} + +-- | Transpose a 2-dimensional array +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array +-- >>> arr = makeArrayLinearR D Seq (Sz (2 :. 3)) id +-- >>> arr +-- Array D Seq (Sz (2 :. 3)) +-- [ [ 0, 1, 2 ] +-- , [ 3, 4, 5 ] +-- ] +-- >>> transpose arr +-- Array D Seq (Sz (3 :. 2)) +-- [ [ 0, 3 ] +-- , [ 1, 4 ] +-- , [ 2, 5 ] +-- ] +-- +-- @since 0.1.0 +transpose :: forall r e. Source r e => Matrix r e -> Matrix D e +transpose = transposeInner +{-# INLINE [1] transpose #-} + +{-# RULES +"transpose . transpose" [~1] forall arr. transpose (transpose arr) = delay arr +"transposeInner . transposeInner" [~1] forall arr. transposeInner (transposeInner arr) = delay arr +"transposeOuter . transposeOuter" [~1] forall arr. transposeOuter (transposeOuter arr) = delay arr + #-} + +-- | Transpose inner two dimensions of at least rank-2 array. +-- +-- ===__Examples__ +-- +-- >>> import Data.Massiv.Array +-- >>> arr = makeArrayLinearR U Seq (Sz (2 :> 3 :. 4)) id +-- >>> arr +-- Array U Seq (Sz (2 :> 3 :. 4)) +-- [ [ [ 0, 1, 2, 3 ] +-- , [ 4, 5, 6, 7 ] +-- , [ 8, 9, 10, 11 ] +-- ] +-- , [ [ 12, 13, 14, 15 ] +-- , [ 16, 17, 18, 19 ] +-- , [ 20, 21, 22, 23 ] +-- ] +-- ] +-- >>> transposeInner arr +-- Array D Seq (Sz (3 :> 2 :. 4)) +-- [ [ [ 0, 1, 2, 3 ] +-- , [ 12, 13, 14, 15 ] +-- ] +-- , [ [ 4, 5, 6, 7 ] +-- , [ 16, 17, 18, 19 ] +-- ] +-- , [ [ 8, 9, 10, 11 ] +-- , [ 20, 21, 22, 23 ] +-- ] +-- ] +-- +-- @since 0.1.0 +transposeInner + :: forall r ix e + . (Index (Lower ix), Index ix, Source r e) + => Array r ix e + -> Array D ix e +transposeInner !arr = makeArray (getComp arr) newsz newVal + where + transInner !ix = + either throwImpossible id $ do + n <- getDimM ix dix + m <- getDimM ix (dix - 1) + ix' <- setDimM ix dix m + setDimM ix' (dix - 1) n + {-# INLINE transInner #-} + newVal = unsafeIndex arr . transInner + {-# INLINE newVal #-} + !newsz = Sz (transInner (unSz (size arr))) + !dix = dimensions newsz +{-# INLINE [1] transposeInner #-} + +-- | Transpose outer two dimensions of at least rank-2 array. +-- +-- ====__Examples__ +-- +-- >>> import Data.Massiv.Array +-- >>> :set -XTypeApplications +-- >>> arr = makeArrayLinear @U Seq (Sz (2 :> 3 :. 4)) id +-- >>> arr +-- Array U Seq (Sz (2 :> 3 :. 4)) +-- [ [ [ 0, 1, 2, 3 ] +-- , [ 4, 5, 6, 7 ] +-- , [ 8, 9, 10, 11 ] +-- ] +-- , [ [ 12, 13, 14, 15 ] +-- , [ 16, 17, 18, 19 ] +-- , [ 20, 21, 22, 23 ] +-- ] +-- ] +-- >>> transposeOuter arr +-- Array D Seq (Sz (2 :> 4 :. 3)) +-- [ [ [ 0, 4, 8 ] +-- , [ 1, 5, 9 ] +-- , [ 2, 6, 10 ] +-- , [ 3, 7, 11 ] +-- ] +-- , [ [ 12, 16, 20 ] +-- , [ 13, 17, 21 ] +-- , [ 14, 18, 22 ] +-- , [ 15, 19, 23 ] +-- ] +-- ] +-- +-- +-- @since 0.1.0 +transposeOuter + :: forall r ix e + . (Index (Lower ix), Index ix, Source r e) + => Array r ix e + -> Array D ix e +transposeOuter !arr = makeArray (getComp arr) newsz newVal + where + transOuter !ix = + either throwImpossible id $ do + n <- getDimM ix 1 + m <- getDimM ix 2 + ix' <- setDimM ix 1 m + setDimM ix' 2 n + {-# INLINE transOuter #-} + newVal = unsafeIndex arr . transOuter + {-# INLINE newVal #-} + !newsz = Sz (transOuter (unSz (size arr))) +{-# INLINE [1] transposeOuter #-} + +-- | Reverse an array along some dimension. Dimension supplied is checked at compile time. +-- +-- ==== __Example__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> arr = makeArrayLinear Seq (Sz2 4 5) (+10) :: Array D Ix2 Int +-- >>> arr +-- Array D Seq (Sz (4 :. 5)) +-- [ [ 10, 11, 12, 13, 14 ] +-- , [ 15, 16, 17, 18, 19 ] +-- , [ 20, 21, 22, 23, 24 ] +-- , [ 25, 26, 27, 28, 29 ] +-- ] +-- >>> A.reverse Dim1 arr +-- Array D Seq (Sz (4 :. 5)) +-- [ [ 14, 13, 12, 11, 10 ] +-- , [ 19, 18, 17, 16, 15 ] +-- , [ 24, 23, 22, 21, 20 ] +-- , [ 29, 28, 27, 26, 25 ] +-- ] +-- >>> A.reverse Dim2 arr +-- Array D Seq (Sz (4 :. 5)) +-- [ [ 25, 26, 27, 28, 29 ] +-- , [ 20, 21, 22, 23, 24 ] +-- , [ 15, 16, 17, 18, 19 ] +-- , [ 10, 11, 12, 13, 14 ] +-- ] +-- +-- @since 0.4.1 +reverse + :: forall n r ix e + . (IsIndexDimension ix n, Index ix, Source r e) + => Dimension n + -> Array r ix e + -> Array D ix e +reverse dim = reverse' (fromDimension dim) +{-# INLINE reverse #-} + +-- | Similarly to `reverse`, flip an array along a particular dimension, but throws +-- `IndexDimensionException` for an incorrect dimension. +-- +-- @since 0.4.1 +reverseM + :: forall r ix e m + . (MonadThrow m, Index ix, Source r e) + => Dim + -> Array r ix e + -> m (Array D ix e) +reverseM dim arr = do + let sz = size arr + k <- getDimM (unSz sz) dim + pure $ makeArray (getComp arr) sz $ \ix -> + unsafeIndex arr (snd $ modifyDim' ix dim (\i -> k - i - 1)) +{-# INLINE reverseM #-} + +-- | Reverse an array along some dimension. Same as `reverseM`, but throws the +-- `IndexDimensionException` from pure code. +-- +-- @since 0.4.1 +reverse' + :: forall r ix e + . (HasCallStack, Index ix, Source r e) + => Dim + -> Array r ix e + -> Array D ix e +reverse' dim = throwEither . reverseM dim +{-# INLINE reverse' #-} + +-- | Rearrange elements of an array into a new one by using a function that maps indices of the +-- newly created one into the old one. This function can throw `IndexOutOfBoundsException`. +-- +-- ===__Examples__ +-- +-- >>> import Data.Massiv.Array +-- >>> :set -XTypeApplications +-- >>> arr = makeArrayLinear @D Seq (Sz (2 :> 3 :. 4)) id +-- >>> arr +-- Array D Seq (Sz (2 :> 3 :. 4)) +-- [ [ [ 0, 1, 2, 3 ] +-- , [ 4, 5, 6, 7 ] +-- , [ 8, 9, 10, 11 ] +-- ] +-- , [ [ 12, 13, 14, 15 ] +-- , [ 16, 17, 18, 19 ] +-- , [ 20, 21, 22, 23 ] +-- ] +-- ] +-- >>> backpermuteM @U (Sz (4 :. 2)) (\(i :. j) -> j :> j :. i) arr +-- Array U Seq (Sz (4 :. 2)) +-- [ [ 0, 16 ] +-- , [ 1, 17 ] +-- , [ 2, 18 ] +-- , [ 3, 19 ] +-- ] +-- +-- @since 0.3.0 +backpermuteM + :: forall r ix e r' ix' m + . (Manifest r e, Index ix, Source r' e, Index ix', MonadUnliftIO m, PrimMonad m, MonadThrow m) + => Sz ix + -- ^ Size of the result array + -> (ix -> ix') + -- ^ A function that maps indices of the new array into the source one. + -> Array r' ix' e + -- ^ Source array. + -> m (Array r ix e) +backpermuteM sz ixF !arr = generateArray (getComp arr) sz (evaluateM arr . ixF) +{-# INLINE backpermuteM #-} + +-- | Similar to `backpermuteM`, with a few notable differences: +-- +-- * Creates a delayed array, instead of manifest, therefore it can be fused +-- * Respects computation strategy, so it can be parallelized +-- * Throws a runtime `IndexOutOfBoundsException` from pure code. +-- +-- @since 0.3.0 +backpermute' + :: forall r ix ix' e + . (HasCallStack, Source r e, Index ix, Index ix') + => Sz ix' + -- ^ Size of the result array + -> (ix' -> ix) + -- ^ A function that maps indices of the new array into the source one. + -> Array r ix e + -- ^ Source array. + -> Array D ix' e +backpermute' sz ixF !arr = makeArray (getComp arr) sz (evaluate' arr . ixF) +{-# INLINE backpermute' #-} + +-- | Append two arrays together along a particular dimension. Sizes of both arrays must match, with +-- an allowed exception of the dimension they are being appended along, otherwise `Nothing` is +-- returned. +-- +-- ====__Examples__ +-- +-- Append two 2D arrays along both dimensions. Note that they do agree on inner dimensions. +-- +-- >>> import Data.Massiv.Array +-- >>> arrA = makeArrayR U Seq (Sz2 2 3) (\(i :. j) -> ('A', i, j)) +-- >>> arrB = makeArrayR U Seq (Sz2 2 3) (\(i :. j) -> ('B', i, j)) +-- >>> appendM 1 arrA arrB +-- Array DL Seq (Sz (2 :. 6)) +-- [ [ ('A',0,0), ('A',0,1), ('A',0,2), ('B',0,0), ('B',0,1), ('B',0,2) ] +-- , [ ('A',1,0), ('A',1,1), ('A',1,2), ('B',1,0), ('B',1,1), ('B',1,2) ] +-- ] +-- >>> appendM 2 arrA arrB +-- Array DL Seq (Sz (4 :. 3)) +-- [ [ ('A',0,0), ('A',0,1), ('A',0,2) ] +-- , [ ('A',1,0), ('A',1,1), ('A',1,2) ] +-- , [ ('B',0,0), ('B',0,1), ('B',0,2) ] +-- , [ ('B',1,0), ('B',1,1), ('B',1,2) ] +-- ] +-- +-- Now appending arrays with different sizes: +-- +-- >>> arrC = makeArrayR U Seq (Sz (2 :. 4)) (\(i :. j) -> ('C', i, j)) +-- >>> appendM 1 arrA arrC +-- Array DL Seq (Sz (2 :. 7)) +-- [ [ ('A',0,0), ('A',0,1), ('A',0,2), ('C',0,0), ('C',0,1), ('C',0,2), ('C',0,3) ] +-- , [ ('A',1,0), ('A',1,1), ('A',1,2), ('C',1,0), ('C',1,1), ('C',1,2), ('C',1,3) ] +-- ] +-- >>> appendM 2 arrA arrC +-- *** Exception: SizeMismatchException: (Sz (2 :. 3)) vs (Sz (2 :. 4)) +-- +-- @since 0.3.0 +appendM + :: forall r1 r2 ix e m + . (MonadThrow m, Index ix, Source r1 e, Source r2 e) + => Dim + -> Array r1 ix e + -> Array r2 ix e + -> m (Array DL ix e) +appendM n !arr1 !arr2 = do + let !sz1 = size arr1 + !sz2 = size arr2 + (k1, szl1) <- pullOutSzM sz1 n + (k2, szl2) <- pullOutSzM sz2 n + unless (szl1 == szl2) $ throwM $ SizeMismatchException sz1 sz2 + let !k1' = unSz k1 + newSz <- insertSzM szl1 n (SafeSz (k1' + unSz k2)) + let load :: Loader e + load scheduler !startAt dlWrite _dlSet = do + scheduleWork scheduler $ + iterA_ zeroIndex (unSz sz1) (pureIndex 1) (<) $ \ix -> + dlWrite (startAt + toLinearIndex newSz ix) (unsafeIndex arr1 ix) + scheduleWork scheduler $ + iterA_ zeroIndex (unSz sz2) (pureIndex 1) (<) $ \ix -> + let i = getDim' ix n + ix' = setDim' ix n (i + k1') + in dlWrite (startAt + toLinearIndex newSz ix') (unsafeIndex arr2 ix) + {-# INLINE load #-} + return $ + DLArray + { dlComp = getComp arr1 <> getComp arr2 + , dlSize = newSz + , dlLoad = load + } +{-# INLINE appendM #-} + +-- | Same as `appendM`, but will throw an exception in pure code on mismatched sizes. +-- +-- @since 0.3.0 +append' + :: forall r1 r2 ix e + . (HasCallStack, Index ix, Source r1 e, Source r2 e) + => Dim + -> Array r1 ix e + -> Array r2 ix e + -> Array DL ix e +append' dim arr1 arr2 = throwEither $ appendM dim arr1 arr2 +{-# INLINE append' #-} + +-- | Concat many arrays together along some dimension. +-- +-- @since 0.3.0 +concat' + :: forall f r ix e + . (HasCallStack, Foldable f, Index ix, Source r e) + => Dim + -> f (Array r ix e) + -> Array DL ix e +concat' n = throwEither . concatM n +{-# INLINE concat' #-} + +-- | Concatenate many arrays together along some dimension. It is important that all sizes are +-- equal, with an exception of the dimensions along which concatenation happens. +-- +-- /__Exceptions__/: `IndexDimensionException`, `SizeMismatchException` +-- +-- @since 0.3.0 +concatM + :: forall r ix e f m + . (MonadThrow m, Foldable f, Index ix, Source r e) + => Dim + -> f (Array r ix e) + -> m (Array DL ix e) +concatM n arrsF = + case L.uncons (F.toList arrsF) of + Nothing -> pure empty + Just (a, arrs) -> do + let sz = unSz (size a) + szs = unSz . size <$> arrs + (k, szl) <- pullOutDimM sz n + -- / remove the dimension out of all sizes along which concatenation will happen + (ks, szls) <- + F.foldrM (\ !csz (ks, szls) -> bimap (: ks) (: szls) <$> pullOutDimM csz n) ([], []) szs + -- / make sure to fail as soon as at least one of the arrays has a mismatching inner size + traverse_ + (\(sz', _) -> throwM (SizeMismatchException (SafeSz sz) (SafeSz sz'))) + (dropWhile ((== szl) . snd) $ P.zip szs szls) + let kTotal = SafeSz $ F.foldl' (+) k ks + newSz <- insertSzM (SafeSz szl) n kTotal + let load :: Loader e + load scheduler startAt dlWrite _dlSet = + let arrayLoader !kAcc (!kCur, arr) = do + scheduleWork scheduler $ + iforM_ arr $ \ix e -> do + i <- getDimM ix n + ix' <- setDimM ix n (i + kAcc) + dlWrite (startAt + toLinearIndex newSz ix') e + pure $! kAcc + kCur + {-# INLINE arrayLoader #-} + in M.foldM_ arrayLoader 0 $ (k, a) : P.zip ks arrs + {-# INLINE load #-} + return $ + DLArray{dlComp = getComp a <> foldMap getComp arrs, dlSize = newSz, dlLoad = load} +{-# INLINE concatM #-} + +-- | Stack slices on top of each other along the specified dimension. +-- +-- /__Exceptions__/: `IndexDimensionException`, `SizeMismatchException` +-- +-- ====__Examples__ +-- +-- Here are the three different ways to stack up two 2D Matrix pages into a 3D array. +-- +-- >>> import Data.Massiv.Array as A +-- >>> x = compute (iterateN 3 succ 0) :: Matrix P Int +-- >>> y = compute (iterateN 3 succ 9) :: Matrix P Int +-- >>> x +-- Array P Seq (Sz (3 :. 3)) +-- [ [ 1, 2, 3 ] +-- , [ 4, 5, 6 ] +-- , [ 7, 8, 9 ] +-- ] +-- >>> y +-- Array P Seq (Sz (3 :. 3)) +-- [ [ 10, 11, 12 ] +-- , [ 13, 14, 15 ] +-- , [ 16, 17, 18 ] +-- ] +-- >>> stackSlicesM 1 [x, y] :: IO (Array DL Ix3 Int) +-- Array DL Seq (Sz (3 :> 3 :. 2)) +-- [ [ [ 1, 10 ] +-- , [ 2, 11 ] +-- , [ 3, 12 ] +-- ] +-- , [ [ 4, 13 ] +-- , [ 5, 14 ] +-- , [ 6, 15 ] +-- ] +-- , [ [ 7, 16 ] +-- , [ 8, 17 ] +-- , [ 9, 18 ] +-- ] +-- ] +-- >>> stackSlicesM 2 [x, y] :: IO (Array DL Ix3 Int) +-- Array DL Seq (Sz (3 :> 2 :. 3)) +-- [ [ [ 1, 2, 3 ] +-- , [ 10, 11, 12 ] +-- ] +-- , [ [ 4, 5, 6 ] +-- , [ 13, 14, 15 ] +-- ] +-- , [ [ 7, 8, 9 ] +-- , [ 16, 17, 18 ] +-- ] +-- ] +-- >>> stackSlicesM 3 [x, y] :: IO (Array DL Ix3 Int) +-- Array DL Seq (Sz (2 :> 3 :. 3)) +-- [ [ [ 1, 2, 3 ] +-- , [ 4, 5, 6 ] +-- , [ 7, 8, 9 ] +-- ] +-- , [ [ 10, 11, 12 ] +-- , [ 13, 14, 15 ] +-- , [ 16, 17, 18 ] +-- ] +-- ] +-- +-- @since 0.5.4 +stackSlicesM + :: forall r ix e f m + . (Foldable f, MonadThrow m, Index (Lower ix), Source r e, Index ix) + => Dim + -> f (Array r (Lower ix) e) + -> m (Array DL ix e) +stackSlicesM dim !arrsF = do + case L.uncons (F.toList arrsF) of + Nothing -> pure empty + Just (a, arrs) -> do + let sz = size a + len = SafeSz (F.length arrsF) + -- / make sure all arrays have the same size + M.forM_ arrsF $ \arr -> + unless (sz == size arr) $ throwM (SizeMismatchException sz (size arr)) + newSz <- insertSzM sz dim len + let load :: Loader e + load scheduler startAt dlWrite _dlSet = + let loadIndex k ix = dlWrite (toLinearIndex newSz (insertDim' ix dim k) + startAt) + arrayLoader !k arr = (k + 1) <$ scheduleWork scheduler (imapM_ (loadIndex k) arr) + {-# INLINE arrayLoader #-} + in M.foldM_ arrayLoader 0 arrsF + {-# INLINE load #-} + return $ + DLArray{dlComp = foldMap getComp arrs, dlSize = newSz, dlLoad = load} +{-# INLINE stackSlicesM #-} + +-- | Specialized `stackSlicesM` to handling stacking from the outside. It is the inverse of +-- `Data.Massiv.Array.outerSlices`. +-- +-- /__Exceptions__/: `SizeMismatchException` +-- +-- ====__Examples__ +-- +-- In this example we stack vectors as row of a matrix from top to bottom: +-- +-- >>> import Data.Massiv.Array as A +-- >>> x = compute (iterateN 3 succ 0) :: Matrix P Int +-- >>> x +-- Array P Seq (Sz (3 :. 3)) +-- [ [ 1, 2, 3 ] +-- , [ 4, 5, 6 ] +-- , [ 7, 8, 9 ] +-- ] +-- >>> rows = outerSlices x +-- >>> A.mapM_ print rows +-- Array P Seq (Sz1 3) +-- [ 1, 2, 3 ] +-- Array P Seq (Sz1 3) +-- [ 4, 5, 6 ] +-- Array P Seq (Sz1 3) +-- [ 7, 8, 9 ] +-- >>> stackOuterSlicesM rows :: IO (Matrix DL Int) +-- Array DL Seq (Sz (3 :. 3)) +-- [ [ 1, 2, 3 ] +-- , [ 4, 5, 6 ] +-- , [ 7, 8, 9 ] +-- ] +-- +-- @since 0.5.4 +stackOuterSlicesM + :: forall r ix e f m + . (Foldable f, MonadThrow m, Index (Lower ix), Source r e, Index ix) + => f (Array r (Lower ix) e) + -> m (Array DL ix e) +stackOuterSlicesM = stackSlicesM (dimensions (Proxy :: Proxy ix)) +{-# INLINE stackOuterSlicesM #-} + +-- | Specialized `stackSlicesM` to handling stacking from the inside. It is the inverse of +-- `Data.Massiv.Array.innerSlices`. +-- +-- /__Exceptions__/: `SizeMismatchException` +-- +-- ====__Examples__ +-- +-- In this example we stack vectors as columns of a matrix from left to right: +-- +-- >>> import Data.Massiv.Array as A +-- >>> x = compute (iterateN 3 succ 0) :: Matrix P Int +-- >>> x +-- Array P Seq (Sz (3 :. 3)) +-- [ [ 1, 2, 3 ] +-- , [ 4, 5, 6 ] +-- , [ 7, 8, 9 ] +-- ] +-- >>> columns = innerSlices x +-- >>> A.mapM_ print columns +-- Array D Seq (Sz1 3) +-- [ 1, 4, 7 ] +-- Array D Seq (Sz1 3) +-- [ 2, 5, 8 ] +-- Array D Seq (Sz1 3) +-- [ 3, 6, 9 ] +-- >>> stackInnerSlicesM columns :: IO (Matrix DL Int) +-- Array DL Seq (Sz (3 :. 3)) +-- [ [ 1, 2, 3 ] +-- , [ 4, 5, 6 ] +-- , [ 7, 8, 9 ] +-- ] +-- +-- @since 0.5.4 +stackInnerSlicesM + :: forall r ix e f m + . (Foldable f, MonadThrow m, Index (Lower ix), Source r e, Index ix) + => f (Array r (Lower ix) e) + -> m (Array DL ix e) +stackInnerSlicesM = stackSlicesM 1 +{-# INLINE stackInnerSlicesM #-} + +-- | /O(1)/ - Split an array into two at an index along a specified dimension. +-- +-- /Related/: 'splitAt'', `splitExtractM`, 'Data.Massiv.Vector.sliceAt'', `Data.Massiv.Vector.sliceAtM` +-- +-- /__Exceptions__/: `IndexDimensionException`, `SizeSubregionException` +-- +-- @since 0.3.0 +splitAtM + :: forall r ix e m + . (MonadThrow m, Index ix, Source r e) + => Dim + -- ^ Dimension along which to split + -> Int + -- ^ Index along the dimension to split at + -> Array r ix e + -- ^ Source array + -> m (Array D ix e, Array D ix e) +splitAtM dim i arr = do + let Sz sz = size arr + eIx <- setDimM sz dim i + sIx <- setDimM zeroIndex dim i + arr1 <- extractFromToM zeroIndex eIx arr + arr2 <- extractFromToM sIx sz arr + return (arr1, arr2) +{-# INLINE splitAtM #-} + +-- | /O(1)/ - Split an array into two at an index along a specified dimension. Throws an +-- error for a wrong dimension or incorrect indices. +-- +-- /Related/: `splitAtM`, `splitExtractM`, 'Data.Massiv.Vector.sliceAt'', `Data.Massiv.Vector.sliceAtM` +-- +-- ==== __Examples__ +-- +-- +-- @since 0.1.0 +splitAt' + :: forall r ix e + . (HasCallStack, Index ix, Source r e) + => Dim + -> Int + -> Array r ix e + -> (Array D ix e, Array D ix e) +splitAt' dim i = throwEither . splitAtM dim i +{-# INLINE splitAt' #-} + +-- | Split an array in three parts across some dimension +-- +-- @since 0.3.5 +splitExtractM + :: forall r ix e m + . (MonadThrow m, Index ix, Source r e) + => Dim + -- ^ Dimension along which to do the extraction + -> Ix1 + -- ^ Start index along the dimension that needs to be extracted + -> Sz Ix1 + -- ^ Size of the extracted array along the dimension that it will be extracted + -> Array r ix e + -> m (Array D ix e, Array D ix e, Array D ix e) +splitExtractM dim startIx1 (Sz extractSzIx1) arr = do + let Sz szIx = size arr + midStartIx <- setDimM zeroIndex dim startIx1 + midExtractSzIx <- setDimM szIx dim extractSzIx1 + midArr <- extractM midStartIx (Sz midExtractSzIx) arr + leftArrSzIx <- setDimM szIx dim startIx1 + leftArr <- extractM zeroIndex (Sz leftArrSzIx) arr + rightArrStartIx <- setDimM zeroIndex dim (startIx1 + extractSzIx1) + rightArr <- extractFromToM rightArrStartIx szIx arr + pure (leftArr, midArr, rightArr) +{-# INLINE splitExtractM #-} + +-- | Replace a slice of an array with another one +-- +-- ==== __Example__ +-- +-- >>> import Data.Massiv.Array +-- >>> arr = makeArrayR U Seq (Sz3 3 4 5) fromIx3 +-- >>> arr' = makeArrayR U Seq (Sz3 3 4 5) (fromIx3 . liftIndex (* 100)) +-- >>> replaceSlice 2 1 (arr' (2, 3)) arr +-- Array DL Seq (Sz (3 :> 4 :. 5)) +-- [ [ [ (0,0,0), (0,0,1), (0,0,2), (0,0,3), (0,0,4) ] +-- , [ (0,300,0), (0,300,100), (0,300,200), (0,300,300), (0,300,400) ] +-- , [ (0,2,0), (0,2,1), (0,2,2), (0,2,3), (0,2,4) ] +-- , [ (0,3,0), (0,3,1), (0,3,2), (0,3,3), (0,3,4) ] +-- ] +-- , [ [ (1,0,0), (1,0,1), (1,0,2), (1,0,3), (1,0,4) ] +-- , [ (100,300,0), (100,300,100), (100,300,200), (100,300,300), (100,300,400) ] +-- , [ (1,2,0), (1,2,1), (1,2,2), (1,2,3), (1,2,4) ] +-- , [ (1,3,0), (1,3,1), (1,3,2), (1,3,3), (1,3,4) ] +-- ] +-- , [ [ (2,0,0), (2,0,1), (2,0,2), (2,0,3), (2,0,4) ] +-- , [ (200,300,0), (200,300,100), (200,300,200), (200,300,300), (200,300,400) ] +-- , [ (2,2,0), (2,2,1), (2,2,2), (2,2,3), (2,2,4) ] +-- , [ (2,3,0), (2,3,1), (2,3,2), (2,3,3), (2,3,4) ] +-- ] +-- ] +-- +-- @since 0.6.1 +replaceSlice + :: forall r r' ix e m + . (MonadThrow m, Source r e, Source r' e, Index ix, Index (Lower ix)) + => Dim + -> Ix1 + -> Array r' (Lower ix) e + -> Array r ix e + -> m (Array DL ix e) +replaceSlice dim i sl arr = do + (l, m, r) <- splitExtractM dim i (SafeSz 1) arr + m' <- resizeM (size m) sl + concatM dim [l, delay m', r] +{-# INLINE replaceSlice #-} + +-- | Replace an outer slice of an array with another one +-- +-- ==== __Example__ +-- +-- >>> import Data.Massiv.Array +-- >>> arr = makeArrayR U Seq (Sz3 3 4 5) fromIx3 +-- >>> arr' = makeArrayR U Seq (Sz3 3 4 5) (fromIx3 . liftIndex (* 100)) +-- >>> replaceOuterSlice 1 (arr' !> 2) arr +-- Array DL Seq (Sz (3 :> 4 :. 5)) +-- [ [ [ (0,0,0), (0,0,1), (0,0,2), (0,0,3), (0,0,4) ] +-- , [ (0,1,0), (0,1,1), (0,1,2), (0,1,3), (0,1,4) ] +-- , [ (0,2,0), (0,2,1), (0,2,2), (0,2,3), (0,2,4) ] +-- , [ (0,3,0), (0,3,1), (0,3,2), (0,3,3), (0,3,4) ] +-- ] +-- , [ [ (200,0,0), (200,0,100), (200,0,200), (200,0,300), (200,0,400) ] +-- , [ (200,100,0), (200,100,100), (200,100,200), (200,100,300), (200,100,400) ] +-- , [ (200,200,0), (200,200,100), (200,200,200), (200,200,300), (200,200,400) ] +-- , [ (200,300,0), (200,300,100), (200,300,200), (200,300,300), (200,300,400) ] +-- ] +-- , [ [ (2,0,0), (2,0,1), (2,0,2), (2,0,3), (2,0,4) ] +-- , [ (2,1,0), (2,1,1), (2,1,2), (2,1,3), (2,1,4) ] +-- , [ (2,2,0), (2,2,1), (2,2,2), (2,2,3), (2,2,4) ] +-- , [ (2,3,0), (2,3,1), (2,3,2), (2,3,3), (2,3,4) ] +-- ] +-- ] +-- +-- @since 0.6.1 +replaceOuterSlice + :: forall r ix e m + . (MonadThrow m, Index ix, Source r e, Load r (Lower ix) e) + => Ix1 + -> Array r (Lower ix) e + -> Array r ix e + -> m (Array DL ix e) +replaceOuterSlice i sl arr = replaceSlice (dimensions (size arr)) i sl arr +{-# INLINE replaceOuterSlice #-} + +-- | Delete a region from an array along the specified dimension. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array +-- >>> arr = fromIx3 <$> (0 :> 0 :. 0 ..: 3 :> 2 :. 6) +-- >>> deleteRegionM 1 2 3 arr +-- Array DL Seq (Sz (3 :> 2 :. 3)) +-- [ [ [ (0,0,0), (0,0,1), (0,0,5) ] +-- , [ (0,1,0), (0,1,1), (0,1,5) ] +-- ] +-- , [ [ (1,0,0), (1,0,1), (1,0,5) ] +-- , [ (1,1,0), (1,1,1), (1,1,5) ] +-- ] +-- , [ [ (2,0,0), (2,0,1), (2,0,5) ] +-- , [ (2,1,0), (2,1,1), (2,1,5) ] +-- ] +-- ] +-- >>> v = Ix1 0 ... 10 +-- >>> deleteRegionM 1 3 5 v +-- Array DL Seq (Sz1 6) +-- [ 0, 1, 2, 8, 9, 10 ] +-- +-- @since 0.3.5 +deleteRegionM + :: forall r ix e m + . (MonadThrow m, Index ix, Source r e) + => Dim + -- ^ Along which axis should the removal happen + -> Ix1 + -- ^ At which index to start dropping slices + -> Sz Ix1 + -- ^ Number of slices to drop + -> Array r ix e + -- ^ Array that will have it's subarray removed + -> m (Array DL ix e) +deleteRegionM dim ix sz arr = do + (leftArr, _, rightArr) <- splitExtractM dim ix sz arr + appendM dim leftArr rightArr +{-# INLINE deleteRegionM #-} + +-- | Similar to `deleteRegionM`, but drop a specified number of rows from an array that +-- has at least 2 dimensions. +-- +-- ====__Example__ +-- +-- >>> import Data.Massiv.Array +-- >>> arr = fromIx2 <$> (0 :. 0 ..: 3 :. 6) +-- >>> arr +-- Array D Seq (Sz (3 :. 6)) +-- [ [ (0,0), (0,1), (0,2), (0,3), (0,4), (0,5) ] +-- , [ (1,0), (1,1), (1,2), (1,3), (1,4), (1,5) ] +-- , [ (2,0), (2,1), (2,2), (2,3), (2,4), (2,5) ] +-- ] +-- >>> deleteRowsM 1 1 arr +-- Array DL Seq (Sz (2 :. 6)) +-- [ [ (0,0), (0,1), (0,2), (0,3), (0,4), (0,5) ] +-- , [ (2,0), (2,1), (2,2), (2,3), (2,4), (2,5) ] +-- ] +-- +-- @since 0.3.5 +deleteRowsM + :: forall r ix e m + . (MonadThrow m, Index ix, Index (Lower ix), Source r e) + => Ix1 + -> Sz Ix1 + -> Array r ix e + -> m (Array DL ix e) +deleteRowsM = deleteRegionM 2 +{-# INLINE deleteRowsM #-} + +-- | Similar to `deleteRegionM`, but drop a specified number of columns an array. +-- +-- ====__Example__ +-- +-- >>> import Data.Massiv.Array +-- >>> arr = fromIx2 <$> (0 :. 0 ..: 3 :. 6) +-- >>> arr +-- Array D Seq (Sz (3 :. 6)) +-- [ [ (0,0), (0,1), (0,2), (0,3), (0,4), (0,5) ] +-- , [ (1,0), (1,1), (1,2), (1,3), (1,4), (1,5) ] +-- , [ (2,0), (2,1), (2,2), (2,3), (2,4), (2,5) ] +-- ] +-- >>> deleteColumnsM 2 3 arr +-- Array DL Seq (Sz (3 :. 3)) +-- [ [ (0,0), (0,1), (0,5) ] +-- , [ (1,0), (1,1), (1,5) ] +-- , [ (2,0), (2,1), (2,5) ] +-- ] +-- +-- @since 0.3.5 +deleteColumnsM + :: forall r ix e m + . (MonadThrow m, Index ix, Source r e) + => Ix1 + -> Sz Ix1 + -> Array r ix e + -> m (Array DL ix e) +deleteColumnsM = deleteRegionM 1 +{-# INLINE deleteColumnsM #-} + +-- | Discard elements from the source array according to the stride. +-- +-- @since 0.3.0 +downsample + :: forall r ix e + . (Source r e, Load r ix e) + => Stride ix + -> Array r ix e + -> Array DL ix e +downsample stride arr = + DLArray{dlComp = getComp arr, dlSize = resultSize, dlLoad = load} + where + resultSize = strideSize stride (size arr) + strideIx = unStride stride + unsafeLinearWriteWithStride = + unsafeIndex arr . liftIndex2 (*) strideIx . fromLinearIndex resultSize + {-# INLINE unsafeLinearWriteWithStride #-} + load :: Loader e + load scheduler startAt dlWrite _ = + splitLinearlyWithStartAtM_ + scheduler + startAt + (totalElem resultSize) + (pure . unsafeLinearWriteWithStride) + dlWrite + {-# INLINE load #-} +{-# INLINE downsample #-} + +-- | Insert the same element into a `Load`able array according to the supplied stride. +-- +-- ====__Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> arr = iterateN (Sz2 3 2) succ (0 :: Int) +-- >>> arr +-- Array DL Seq (Sz (3 :. 2)) +-- [ [ 1, 2 ] +-- , [ 3, 4 ] +-- , [ 5, 6 ] +-- ] +-- >>> upsample 0 (Stride (2 :. 3)) arr +-- Array DL Seq (Sz (6 :. 6)) +-- [ [ 1, 0, 0, 2, 0, 0 ] +-- , [ 0, 0, 0, 0, 0, 0 ] +-- , [ 3, 0, 0, 4, 0, 0 ] +-- , [ 0, 0, 0, 0, 0, 0 ] +-- , [ 5, 0, 0, 6, 0, 0 ] +-- , [ 0, 0, 0, 0, 0, 0 ] +-- ] +-- >>> upsample 9 (Stride (1 :. 2)) arr +-- Array DL Seq (Sz (3 :. 4)) +-- [ [ 1, 9, 2, 9 ] +-- , [ 3, 9, 4, 9 ] +-- , [ 5, 9, 6, 9 ] +-- ] +-- +-- @since 0.3.0 +upsample + :: forall r ix e + . Load r ix e + => e + -- ^ Element to use for filling the newly added cells + -> Stride ix + -- ^ Fill cells according to this stride + -> Array r ix e + -- ^ Array that will have cells added to + -> Array DL ix e +upsample !fillWith safeStride arr = + DLArray + { dlComp = getComp arr + , dlSize = newsz + , dlLoad = load + } + where + load :: Loader e + load scheduler startAt uWrite uSet = do + uSet startAt (toLinearSz newsz) fillWith + iterArrayLinearST_ scheduler arr (\i -> uWrite (adjustLinearStride (i + startAt))) + {-# INLINE load #-} + adjustLinearStride = toLinearIndex newsz . timesStride . fromLinearIndex sz + {-# INLINE adjustLinearStride #-} + timesStride !ix = liftIndex2 (*) stride ix + {-# INLINE timesStride #-} + !stride = unStride safeStride + ~sz = outerSize arr -- intentionally lazy in case it is used with DS + !newsz = SafeSz (timesStride $ unSz sz) +{-# INLINE upsample #-} + +-- | General array transformation, that forces computation and produces a manifest array. +-- +-- @since 0.3.0 +transformM + :: forall r ix e r' ix' e' a m + . (Manifest r e, Index ix, Source r' e', Index ix', MonadUnliftIO m, PrimMonad m, MonadThrow m) + => (Sz ix' -> m (Sz ix, a)) + -> (a -> (ix' -> m e') -> ix -> m e) + -> Array r' ix' e' + -> m (Array r ix e) +transformM getSzM getM arr = do + (sz, a) <- getSzM (size arr) + generateArray (getComp arr) sz (getM a (evaluateM arr)) +{-# INLINE transformM #-} + +-- | General array transformation +-- +-- @since 0.3.0 +transform' + :: forall ix e r' ix' e' a + . (HasCallStack, Source r' e', Index ix', Index ix) + => (Sz ix' -> (Sz ix, a)) + -> (a -> (ix' -> e') -> ix -> e) + -> Array r' ix' e' + -> Array D ix e +transform' getSz get arr = makeArray (getComp arr) sz (get a (evaluate' arr)) + where + (sz, a) = getSz (size arr) +{-# INLINE transform' #-} + +-- | Same as `transformM`, but operates on two arrays +-- +-- @since 0.3.0 +transform2M + :: ( Manifest r e + , Index ix + , Source r1 e1 + , Source r2 e2 + , Index ix1 + , Index ix2 + , MonadUnliftIO m + , PrimMonad m + , MonadThrow m + ) + => (Sz ix1 -> Sz ix2 -> m (Sz ix, a)) + -> (a -> (ix1 -> m e1) -> (ix2 -> m e2) -> ix -> m e) + -> Array r1 ix1 e1 + -> Array r2 ix2 e2 + -> m (Array r ix e) +transform2M getSzM getM arr1 arr2 = do + (sz, a) <- getSzM (size arr1) (size arr2) + generateArray (getComp arr1 <> getComp arr2) sz (getM a (evaluateM arr1) (evaluateM arr2)) +{-# INLINE transform2M #-} + +-- | Same as 'transform'', but operates on two arrays +-- +-- @since 0.3.0 +transform2' + :: (HasCallStack, Source r1 e1, Source r2 e2, Index ix, Index ix1, Index ix2) + => (Sz ix1 -> Sz ix2 -> (Sz ix, a)) + -> (a -> (ix1 -> e1) -> (ix2 -> e2) -> ix -> e) + -> Array r1 ix1 e1 + -> Array r2 ix2 e2 + -> Array D ix e +transform2' getSz get arr1 arr2 = + makeArray (getComp arr1 <> getComp arr2) sz (get a (evaluate' arr1) (evaluate' arr2)) + where + (sz, a) = getSz (size arr1) (size arr2) +{-# INLINE transform2' #-} + +-- | Replicate each element of the array by a factor in stride along each dimension and surround each +-- such group with a box of supplied grid value. It will essentially zoom up an array and create a +-- grid around each element from the original array. Very useful for zooming up images to inspect +-- individual pixels. +-- +-- ==== __Example__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> arr = resize' (Sz2 3 2) (Ix1 1 ... 6) +-- >>> arr +-- Array D Seq (Sz (3 :. 2)) +-- [ [ 1, 2 ] +-- , [ 3, 4 ] +-- , [ 5, 6 ] +-- ] +-- >>> zoomWithGrid 0 (Stride (2 :. 3)) arr +-- Array DL Seq (Sz (10 :. 9)) +-- [ [ 0, 0, 0, 0, 0, 0, 0, 0, 0 ] +-- , [ 0, 1, 1, 1, 0, 2, 2, 2, 0 ] +-- , [ 0, 1, 1, 1, 0, 2, 2, 2, 0 ] +-- , [ 0, 0, 0, 0, 0, 0, 0, 0, 0 ] +-- , [ 0, 3, 3, 3, 0, 4, 4, 4, 0 ] +-- , [ 0, 3, 3, 3, 0, 4, 4, 4, 0 ] +-- , [ 0, 0, 0, 0, 0, 0, 0, 0, 0 ] +-- , [ 0, 5, 5, 5, 0, 6, 6, 6, 0 ] +-- , [ 0, 5, 5, 5, 0, 6, 6, 6, 0 ] +-- , [ 0, 0, 0, 0, 0, 0, 0, 0, 0 ] +-- ] +-- +-- @since 0.3.1 +zoomWithGrid + :: forall r ix e + . (Index ix, Source r e) + => e + -- ^ Value to use for the grid + -> Stride ix + -- ^ Scaling factor + -> Array r ix e + -- ^ Source array + -> Array DL ix e +zoomWithGrid gridVal (Stride zoomFactor) arr = unsafeMakeLoadArray Seq newSz (Just gridVal) load + where + !kx = liftIndex (+ 1) zoomFactor + !lastNewIx = liftIndex2 (*) kx $ unSz (size arr) + !newSz = Sz (liftIndex (+ 1) lastNewIx) + load :: forall s. Scheduler s () -> Ix1 -> (Ix1 -> e -> ST s ()) -> ST s () + load scheduler _ writeElement = + iforSchedulerM_ scheduler arr $ \ !ix !e -> + let !kix = liftIndex2 (*) ix kx + in mapM_ (\ !ix' -> writeElement (toLinearIndex newSz ix') e) $ + range Seq (liftIndex (+ 1) kix) (liftIndex2 (+) kix kx) + {-# INLINE load #-} +{-# INLINE zoomWithGrid #-} + +-- | Increaze the size of the array accoridng to the stride multiplier while replicating +-- the same element to fill the neighbors. It is exactly the same as `zoomWithGrid`, but +-- without the grid. +-- +-- ==== __Example__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> arr = resize' (Sz3 1 3 2) (Ix1 1 ... 6) +-- >>> arr +-- Array D Seq (Sz (1 :> 3 :. 2)) +-- [ [ [ 1, 2 ] +-- , [ 3, 4 ] +-- , [ 5, 6 ] +-- ] +-- ] +-- >>> zoom (Stride (2 :> 2 :. 3)) arr +-- Array DL Seq (Sz (2 :> 6 :. 6)) +-- [ [ [ 1, 1, 1, 2, 2, 2 ] +-- , [ 1, 1, 1, 2, 2, 2 ] +-- , [ 3, 3, 3, 4, 4, 4 ] +-- , [ 3, 3, 3, 4, 4, 4 ] +-- , [ 5, 5, 5, 6, 6, 6 ] +-- , [ 5, 5, 5, 6, 6, 6 ] +-- ] +-- , [ [ 1, 1, 1, 2, 2, 2 ] +-- , [ 1, 1, 1, 2, 2, 2 ] +-- , [ 3, 3, 3, 4, 4, 4 ] +-- , [ 3, 3, 3, 4, 4, 4 ] +-- , [ 5, 5, 5, 6, 6, 6 ] +-- , [ 5, 5, 5, 6, 6, 6 ] +-- ] +-- ] +-- +-- @since 0.4.4 +zoom + :: forall r ix e + . (Index ix, Source r e) + => Stride ix + -- ^ Scaling factor + -> Array r ix e + -- ^ Source array + -> Array DL ix e +zoom (Stride zoomFactor) arr = unsafeMakeLoadArray Seq newSz Nothing load + where + !lastNewIx = liftIndex2 (*) zoomFactor $ unSz (size arr) + !newSz = Sz lastNewIx + load :: forall s. Scheduler s () -> Ix1 -> (Ix1 -> e -> ST s ()) -> ST s () + load scheduler _ writeElement = + iforSchedulerM_ scheduler arr $ \ !ix !e -> + let !kix = liftIndex2 (*) ix zoomFactor + in mapM_ (\ !ix' -> writeElement (toLinearIndex newSz ix') e) $ + range Seq kix (liftIndex2 (+) kix zoomFactor) + {-# INLINE load #-} {-# INLINE zoom #-} 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,52 +1,53 @@ {-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE TypeFamilies #-} + -- | -- Module : Data.Massiv.Array.Stencil --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Stencil - ( -- * Stencil - Stencil - , Value - , makeStencil - , makeStencilDef - , getStencilSize - , getStencilCenter +module Data.Massiv.Array.Stencil ( + -- * Stencil + Stencil, + makeStencil, + getStencilSize, + getStencilCenter, + -- ** Padding - , Padding(..) - , noPadding - , samePadding + Padding (..), + noPadding, + samePadding, + -- ** Application - , mapStencil - , applyStencil + mapStencil, + applyStencil, + -- ** Common stencils - , idStencil - , sumStencil - , productStencil - , avgStencil - , maxStencil - , minStencil - , foldlStencil - , foldrStencil - , foldStencil + idStencil, + sumStencil, + productStencil, + avgStencil, + maxStencil, + minStencil, + foldlStencil, + foldrStencil, + foldStencil, + -- ** Profunctor - , dimapStencil - , lmapStencil - , rmapStencil + dimapStencil, + lmapStencil, + rmapStencil, + -- * Convolution - , module Data.Massiv.Array.Stencil.Convolution - -- * Re-export - , Default(def) - ) where + module Data.Massiv.Array.Stencil.Convolution, +) where import Data.Coerce -import Data.Default.Class (Default(def)) import Data.Massiv.Array.Delayed.Windowed import Data.Massiv.Array.Manifest import Data.Massiv.Array.Stencil.Convolution @@ -72,16 +73,18 @@ -- `Data.Massiv.Array.compute`d in order to be useful. -- -- @since 0.1.0 -mapStencil :: - (Source r ix e, Manifest r ix e) - => Border e -- ^ Border resolution technique - -> Stencil ix e a -- ^ Stencil to map over the array - -> Array r ix e -- ^ Source array +mapStencil + :: (Index ix, Manifest r e) + => Border e + -- ^ Border resolution technique + -> Stencil ix e a + -- ^ Stencil to map over the array + -> Array r ix e + -- ^ Source array -> Array DW ix a mapStencil b stencil = applyStencil (samePadding stencil b) stencil {-# INLINE mapStencil #-} - -- | Padding of the source array before stencil application. -- -- ==== __Examples__ @@ -147,11 +150,12 @@ -- -- @since 0.4.3 data Padding ix e = Padding - { paddingFromOrigin :: !(Sz ix) - , paddingFromBottom :: !(Sz ix) + { paddingFromOrigin :: !(Sz ix) + , paddingFromBottom :: !(Sz ix) , paddingWithElement :: !(Border e) -- ^ Element to do padding with - } deriving (Eq, Show) + } + deriving (Eq, Show) -- | Also known as "valid" padding. When stencil is applied to an array, that array will -- shrink, unless the stencil is of size 1. @@ -170,7 +174,7 @@ samePadding (Stencil (Sz sSz) sCenter _) border = Padding { paddingFromOrigin = Sz sCenter - , paddingFromBottom = Sz (liftIndex2 (-) sSz (liftIndex (+1) sCenter)) + , paddingFromBottom = Sz (liftIndex2 (-) sSz (liftIndex (+ 1) sCenter)) , paddingWithElement = border } @@ -180,13 +184,15 @@ -- depend on the padding. -- -- @since 0.4.3 -applyStencil :: - (Source r ix e, Manifest r ix e) +applyStencil + :: (Index ix, Manifest r e) => Padding ix e -- ^ Padding to be applied to the source array. This will dictate the resulting size of -- the array. No padding will cause it to shrink by the size of the stencil - -> Stencil ix e a -- ^ Stencil to apply to the array - -> Array r ix e -- ^ Source array + -> Stencil ix e a + -- ^ Stencil to apply to the array + -> Array r ix e + -- ^ Source array -> Array DW ix a applyStencil (Padding (Sz po) (Sz pb) border) (Stencil sSz sCenter stencilF) !arr = insertWindow warr window @@ -196,7 +202,7 @@ DArray (getComp arr) sz - (unValue . stencilF (Value . borderIndex border arr) . liftIndex2 (+) offset) + (PrefIndex (stencilF (borderIndex border arr) (borderIndex border arr) . liftIndex2 (+) offset)) -- Size by which the resulting array will shrink (not accounting for padding) !shrinkSz = Sz (liftIndex (subtract 1) (unSz sSz)) !sz = liftSz2 (-) (SafeSz (liftIndex2 (+) po (liftIndex2 (+) pb (unSz (size arr))))) shrinkSz @@ -205,18 +211,20 @@ Window { windowStart = po , windowSize = wsz - , windowIndex = unValue . stencilF (Value . unsafeIndex arr) . liftIndex2 (+) offset + , windowIndex = stencilF (unsafeIndex arr) (index' arr) . liftIndex2 (+) offset , windowUnrollIx2 = unSz . fst <$> pullOutSzM sSz 2 } {-# INLINE applyStencil #-} - -- | Construct a stencil from a function, which describes how to calculate the -- value at a point while having access to neighboring elements with a function -- that accepts idices relative to the center of stencil. Trying to index -- outside the stencil box will result in a runtime error upon stencil -- creation. -- +-- /Note/ - Once correctness of stencil is verified then switching to `makeUnsafeStencil` +-- is recommended in order to get the most performance out of the `Stencil` +-- -- ==== __Example__ -- -- Below is an example of creating a `Stencil`, which, when mapped over a @@ -225,19 +233,23 @@ -- -- /Note/ - Make sure to add an @INLINE@ pragma, otherwise performance will be terrible. -- --- > average3x3Stencil :: (Default a, Fractional a) => Stencil Ix2 a a --- > average3x3Stencil = makeStencil (Sz (3 :. 3)) (1 :. 1) $ \ get -> --- > ( get (-1 :. -1) + get (-1 :. 0) + get (-1 :. 1) + --- > get ( 0 :. -1) + get ( 0 :. 0) + get ( 0 :. 1) + --- > get ( 1 :. -1) + get ( 1 :. 0) + get ( 1 :. 1) ) / 9 --- > {-# INLINE average3x3Stencil #-} +-- @ +-- average3x3Stencil :: Fractional a => Stencil Ix2 a a +-- average3x3Stencil = makeStencil (Sz (3 :. 3)) (1 :. 1) $ \ get -> +-- ( get (-1 :. -1) + get (-1 :. 0) + get (-1 :. 1) + +-- get ( 0 :. -1) + get ( 0 :. 0) + get ( 0 :. 1) + +-- get ( 1 :. -1) + get ( 1 :. 0) + get ( 1 :. 1) ) / 9 +-- {\-# INLINE average3x3Stencil #-\} +-- @ -- -- @since 0.1.0 makeStencil - :: (Index ix, Default e) - => Sz ix -- ^ Size of the stencil - -> ix -- ^ Center of the stencil - -> ((ix -> Value e) -> Value a) + :: Index ix + => Sz ix + -- ^ Size of the stencil + -> ix + -- ^ Center of the stencil + -> ((ix -> e) -> a) -- ^ Stencil function that receives a "get" function as it's argument that can -- retrieve values of cells in the source array with respect to the center of -- the stencil. Stencil function must return a value that will be assigned to @@ -245,33 +257,17 @@ -- cannot go outside the boundaries of the stencil, otherwise an error will be -- raised during stencil creation. -> Stencil ix e a -makeStencil = makeStencilDef def -{-# INLINE makeStencil #-} - --- | Same as `makeStencil`, but with ability to specify default value for stencil validation. --- --- @since 0.2.3 -makeStencilDef - :: Index ix - => e -- ^ Default element that will be used for stencil validation only. - -> Sz ix -- ^ Size of the stencil - -> ix -- ^ Center of the stencil - -> ((ix -> Value e) -> Value a) - -- ^ Stencil function. - -> Stencil ix e a -makeStencilDef defVal !sSz !sCenter relStencil = - validateStencil defVal $ Stencil sSz sCenter stencil +makeStencil !sSz !sCenter relStencil = Stencil sSz sCenter stencil where - stencil getVal !ix = - inline relStencil $ \ !ixD -> getVal (liftIndex2 (+) ix ixD) + stencil _ getVal !ix = inline (relStencil (getVal . liftIndex2 (+) ix)) {-# INLINE stencil #-} -{-# INLINE makeStencilDef #-} +{-# INLINE makeStencil #-} -- | Identity stencil that does not change the elements of the source array. -- -- @since 0.4.3 idStencil :: Index ix => Stencil ix e e -idStencil = makeUnsafeStencil oneSz zeroIndex $ \ _ get -> get zeroIndex +idStencil = makeUnsafeStencil oneSz zeroIndex $ \_ get -> get zeroIndex {-# INLINE idStencil #-} -- | Stencil that does a left fold in a row-major order. Regardless of the supplied size @@ -333,7 +329,9 @@ iter ixStart zeroIndex (pureIndex (-1)) (>=) acc0 $ \ix -> f (get ix) {-# INLINE foldrStencil #-} - +-- | Create a stencil that will fold all elements in the region monoidally. +-- +-- @since 0.4.3 foldStencil :: (Monoid e, Index ix) => Sz ix -> Stencil ix e e foldStencil = foldlStencil mappend mempty {-# INLINE foldStencil #-} @@ -370,7 +368,6 @@ maxStencil :: (Bounded e, Ord e, Index ix) => Sz ix -> Stencil ix e e maxStencil = dimapStencil coerce getMax . foldStencil {-# INLINE maxStencil #-} - -- | Create a stencil centered at 0 that will extract the maximum value in the region of -- supplied size. diff --git a/src/Data/Massiv/Array/Stencil/Convolution.hs b/src/Data/Massiv/Array/Stencil/Convolution.hs --- a/src/Data/Massiv/Array/Stencil/Convolution.hs +++ b/src/Data/Massiv/Array/Stencil/Convolution.hs @@ -1,19 +1,19 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE FlexibleContexts #-} + -- | -- Module : Data.Massiv.Array.Stencil.Convolution --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Stencil.Convolution - ( makeConvolutionStencil - , makeConvolutionStencilFromKernel - , makeCorrelationStencil - , makeCorrelationStencilFromKernel - ) where +module Data.Massiv.Array.Stencil.Convolution ( + makeConvolutionStencil, + makeConvolutionStencilFromKernel, + makeCorrelationStencil, + makeCorrelationStencilFromKernel, +) where import Data.Massiv.Array.Ops.Fold (ifoldlS) import Data.Massiv.Array.Stencil.Internal @@ -23,41 +23,46 @@ -- | Create a convolution stencil by specifying border resolution technique and -- an accumulator function. -- +-- /Note/ - Using `Data.Massiv.Array.Stencil.Unsafe.makeUnsafeConvolutionStencil` will be +-- slightly faster, therefore it is recommended to switch from this function, after manual +-- verification that the created stencil behaves as expected. +-- -- ==== __Examples__ -- -- Here is how to create a 2D horizontal Sobel Stencil: -- --- > sobelX :: Num e => Stencil Ix2 e e --- > sobelX = makeConvolutionStencil (Sz2 3 3) (1 :. 1) $ --- > \f -> f (-1 :. -1) (-1) . f (-1 :. 1) 1 . --- > f ( 0 :. -1) (-2) . f ( 0 :. 1) 2 . --- > f ( 1 :. -1) (-1) . f ( 1 :. 1) 1 --- > {-# INLINE sobelX #-} +-- @ +-- sobelX :: Num e => Stencil Ix2 e e +-- sobelX = makeConvolutionStencil (Sz2 3 3) (1 :. 1) +-- $ \f -> f (-1 :. -1) (-1) . f (-1 :. 1) 1 . +-- f ( 0 :. -1) (-2) . f ( 0 :. 1) 2 . +-- f ( 1 :. -1) (-1) . f ( 1 :. 1) 1 +-- {\-# INLINE sobelX #-\} +-- @ -- -- @since 0.1.0 makeConvolutionStencil :: (Index ix, Num e) => Sz ix -> ix - -> ((ix -> Value e -> Value e -> Value e) -> Value e -> Value e) + -> ((ix -> e -> e -> e) -> e -> e) -> Stencil ix e e makeConvolutionStencil !sz !sCenter relStencil = - validateStencil 0 $ Stencil sz sInvertCenter stencil + Stencil sz sInvertCenter stencil where !sInvertCenter = liftIndex2 (-) (liftIndex (subtract 1) (unSz sz)) sCenter - stencil getVal !ix = + stencil _ getVal !ix = (inline relStencil $ \ !ixD !kVal !acc -> getVal (liftIndex2 (-) ix ixD) * kVal + acc) 0 {-# INLINE stencil #-} {-# INLINE makeConvolutionStencil #-} - -- | Make a stencil out of a Kernel Array. This `Stencil` will be slower than if -- `makeConvolutionStencil` is used, but sometimes we just really don't know the -- kernel at compile time. -- -- @since 0.1.0 makeConvolutionStencilFromKernel - :: (Manifest r ix e, Num e) + :: (Manifest r e, Index ix, Num e) => Array r ix e -> Stencil ix e e makeConvolutionStencilFromKernel kArr = Stencil sz sInvertCenter stencil @@ -66,27 +71,30 @@ !szi1 = liftIndex (subtract 1) szi !sInvertCenter = liftIndex2 (-) szi1 sCenter !sCenter = liftIndex (`quot` 2) szi - stencil getVal !ix = Value (ifoldlS accum 0 kArr) where - !ixOff = liftIndex2 (+) ix sCenter - accum !acc !kIx !kVal = - unValue (getVal (liftIndex2 (-) ixOff kIx)) * kVal + acc - {-# INLINE accum #-} + stencil uget _ !ix = ifoldlS accum 0 kArr + where + !ixOff = liftIndex2 (+) ix sCenter + accum !acc !kIx !kVal = uget (liftIndex2 (-) ixOff kIx) * kVal + acc + {-# INLINE accum #-} {-# INLINE stencil #-} {-# INLINE makeConvolutionStencilFromKernel #-} - -- | Make a stencil -- +-- /Note/ - Using `Data.Massiv.Array.Stencil.Unsafe.makeUnsafeCorrelationStencil` will be +-- much faster, therefore it is recommended to switch from this function, after manual +-- verification that the created stencil behaves as expected. +-- -- @since 0.1.5 makeCorrelationStencil :: (Index ix, Num e) => Sz ix -> ix - -> ((ix -> Value e -> Value e -> Value e) -> Value e -> Value e) + -> ((ix -> e -> e -> e) -> e -> e) -> Stencil ix e e -makeCorrelationStencil !sSz !sCenter relStencil = validateStencil 0 $ Stencil sSz sCenter stencil +makeCorrelationStencil !sSz !sCenter relStencil = Stencil sSz sCenter stencil where - stencil getVal !ix = + stencil _ getVal !ix = (inline relStencil $ \ !ixD !kVal !acc -> getVal (liftIndex2 (+) ix ixD) * kVal + acc) 0 {-# INLINE stencil #-} {-# INLINE makeCorrelationStencil #-} @@ -97,17 +105,17 @@ -- -- @since 0.1.5 makeCorrelationStencilFromKernel - :: (Manifest r ix e, Num e) + :: (Manifest r e, Index ix, Num e) => Array r ix e -> Stencil ix e e makeCorrelationStencilFromKernel kArr = Stencil sz sCenter stencil where !sz = size kArr !sCenter = liftIndex (`div` 2) $ unSz sz - stencil getVal !ix = Value (ifoldlS accum 0 kArr) where - !ixOff = liftIndex2 (-) ix sCenter - accum !acc !kIx !kVal = - unValue (getVal (liftIndex2 (+) ixOff kIx)) * kVal + acc - {-# INLINE accum #-} + stencil uget _ !ix = ifoldlS accum 0 kArr + where + !ixOff = liftIndex2 (-) ix sCenter + accum !acc !kIx !kVal = uget (liftIndex2 (+) ixOff kIx) * kVal + acc + {-# INLINE accum #-} {-# INLINE stencil #-} {-# INLINE makeCorrelationStencilFromKernel #-} 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 @@ -1,147 +1,54 @@ {-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE RecordWildCards #-} {-# LANGUAGE ScopedTypeVariables #-} + -- | -- Module : Data.Massiv.Array.Stencil.Internal --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Stencil.Internal - ( Stencil(..) - , Value(..) - , dimapStencil - , lmapStencil - , rmapStencil - , validateStencil - ) where +module Data.Massiv.Array.Stencil.Internal ( + Stencil (..), + dimapStencil, + lmapStencil, + rmapStencil, +) where import Control.Applicative import Control.DeepSeq -import Data.Massiv.Array.Delayed.Pull import Data.Massiv.Core.Common --- | 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 --- `Data.Array.makeConvolutionStencil` in order to create a stencil. +-- | 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.Massiv.Array.makeStencil` and `Data.Massiv.Array.makeConvolutionStencil` in order +-- to create a stencil. data Stencil ix e a = Stencil - { stencilSize :: !(Sz ix) + { stencilSize :: !(Sz ix) , stencilCenter :: !ix - , stencilFunc :: (ix -> Value e) -> ix -> Value a + , stencilFunc :: (ix -> e) -> (ix -> e) -> ix -> a } instance Index ix => NFData (Stencil ix e a) where rnf (Stencil sz ix f) = sz `deepseq` ix `deepseq` f `seq` () --- | This is a simple wrapper for value of an array cell. It is used in order to improve safety of --- `Stencil` mapping. Using various class instances, such as `Num` and `Functor` for example, make --- it possible to manipulate the value, without having direct access to it. -newtype Value e = Value { unValue :: e } deriving (Show, Bounded) - -instance Functor Value where - fmap f (Value e) = Value (f e) - {-# INLINE fmap #-} - -instance Applicative Value where - pure = Value - {-# INLINE pure #-} - (<*>) (Value f) (Value e) = Value (f e) - {-# INLINE (<*>) #-} - --- | @since 0.1.5 -instance Semigroup a => Semigroup (Value a) where - Value a <> Value b = Value (a <> b) - {-# INLINE (<>) #-} - --- | @since 0.1.5 -instance Monoid a => Monoid (Value a) where - mempty = Value mempty - {-# INLINE mempty #-} - Value a `mappend` Value b = Value (a `mappend` b) - {-# INLINE mappend #-} - -instance Num e => Num (Value e) where - (+) = liftA2 (+) - {-# INLINE (+) #-} - (*) = liftA2 (*) - {-# INLINE (*) #-} - negate = fmap negate - {-# INLINE negate #-} - abs = fmap abs - {-# INLINE abs #-} - signum = fmap signum - {-# INLINE signum #-} - fromInteger = Value . fromInteger - {-# INLINE fromInteger #-} - -instance Fractional e => Fractional (Value e) where - (/) = liftA2 (/) - {-# INLINE (/) #-} - recip = fmap recip - {-# INLINE recip #-} - fromRational = pure . fromRational - {-# INLINE fromRational #-} - -instance Floating e => Floating (Value e) where - pi = pure pi - {-# INLINE pi #-} - exp = fmap exp - {-# INLINE exp #-} - log = fmap log - {-# INLINE log #-} - sqrt = fmap sqrt - {-# INLINE sqrt #-} - (**) = liftA2 (**) - {-# INLINE (**) #-} - logBase = liftA2 logBase - {-# INLINE logBase #-} - sin = fmap sin - {-# INLINE sin #-} - cos = fmap cos - {-# INLINE cos #-} - tan = fmap tan - {-# INLINE tan #-} - asin = fmap asin - {-# INLINE asin #-} - acos = fmap acos - {-# INLINE acos #-} - atan = fmap atan - {-# INLINE atan #-} - sinh = fmap sinh - {-# INLINE sinh #-} - cosh = fmap cosh - {-# INLINE cosh #-} - tanh = fmap tanh - {-# INLINE tanh #-} - asinh = fmap asinh - {-# INLINE asinh #-} - acosh = fmap acosh - {-# INLINE acosh #-} - atanh = fmap atanh - {-# INLINE atanh #-} - - - - instance Functor (Stencil ix e) where fmap = rmapStencil {-# INLINE fmap #-} - -- Profunctor -- | A Profunctor dimap. Same caviat applies as in `lmapStencil` -- -- @since 0.2.3 dimapStencil :: (c -> d) -> (a -> b) -> Stencil ix d a -> Stencil ix c b -dimapStencil f g stencil@Stencil {stencilFunc = sf} = stencil {stencilFunc = sf'} +dimapStencil f g stencil@Stencil{stencilFunc = sf} = stencil{stencilFunc = sf'} where - sf' s = Value . g . unValue . sf (Value . f . unValue . s) + sf' us s = g . sf (f . us) (f . s) {-# INLINE sf' #-} {-# INLINE dimapStencil #-} @@ -154,48 +61,56 @@ -- -- @since 0.2.3 lmapStencil :: (c -> d) -> Stencil ix d a -> Stencil ix c a -lmapStencil f stencil@Stencil {stencilFunc = sf} = stencil {stencilFunc = sf'} +lmapStencil f stencil@Stencil{stencilFunc = sf} = stencil{stencilFunc = sf'} where - sf' s = sf (Value . f . unValue . s) + sf' us s = sf (f . us) (f . s) {-# INLINE sf' #-} {-# INLINE lmapStencil #-} --- | A covariant map over the right most type argument. In other words a usual Functor `fmap`: +-- | A covariant map over the right most type argument. In other words the usual `fmap` +-- from `Functor`: -- -- > fmap == rmapStencil -- -- @since 0.2.3 rmapStencil :: (a -> b) -> Stencil ix e a -> Stencil ix e b -rmapStencil f stencil@Stencil {stencilFunc = sf} = stencil {stencilFunc = sf'} +rmapStencil f stencil@Stencil{stencilFunc = sf} = stencil{stencilFunc = sf'} where - sf' s = Value . f . unValue . sf s + sf' us s = f . sf us s {-# INLINE sf' #-} {-# INLINE rmapStencil #-} +unionStencilCenters :: Index ix => Stencil ix e1 a1 -> Stencil ix e2 a2 -> ix +unionStencilCenters (Stencil _ sC1 _) (Stencil _ sC2 _) = liftIndex2 max sC1 sC2 +{-# INLINE unionStencilCenters #-} +unionStencilSizes :: Index ix => ix -> Stencil ix e1 a1 -> Stencil ix e2 a2 -> Sz ix +unionStencilSizes maxCenter (Stencil (SafeSz sSz1) sC1 _) (Stencil (SafeSz sSz2) sC2 _) = + Sz $ liftIndex2 (+) maxCenter $ liftIndex2 max (liftIndex2 (-) sSz1 sC1) (liftIndex2 (-) sSz2 sC2) +{-# INLINE unionStencilSizes #-} --- TODO: Figure out interchange law (u <*> pure y = pure ($ y) <*> u) and issue --- with discarding size and center. Best idea so far is to increase stencil size to --- the maximum one and shift the center of the other stencil so that they both match --- up. This approach would also remove requirement to validate the result --- Stencil - both stencils are trusted, increasing the size will not affect the --- safety. +-- TODO: Test interchange law (u <*> pure y = pure ($ y) <*> u) instance Index ix => Applicative (Stencil ix e) where - pure a = Stencil oneSz zeroIndex (const (const (Value a))) + pure a = Stencil oneSz zeroIndex (\_ _ _ -> a) {-# INLINE pure #-} - (<*>) (Stencil (SafeSz sSz1) sC1 f1) (Stencil (SafeSz sSz2) sC2 f2) = Stencil newSz maxCenter stF + (<*>) s1@(Stencil _ _ f1) s2@(Stencil _ _ f2) = Stencil newSz maxCenter stF where - stF gV !ix = Value (unValue (f1 gV ix) (unValue (f2 gV ix))) + stF ug gV !ix = f1 ug gV ix (f2 ug gV ix) {-# INLINE stF #-} - !newSz = - Sz - (liftIndex2 - (+) - maxCenter - (liftIndex2 max (liftIndex2 (-) sSz1 sC1) (liftIndex2 (-) sSz2 sC2))) - !maxCenter = liftIndex2 max sC1 sC2 + !newSz = unionStencilSizes maxCenter s1 s2 + !maxCenter = unionStencilCenters s1 s2 {-# INLINE (<*>) #-} +#if MIN_VERSION_base(4,10,0) + liftA2 f s1@(Stencil _ _ f1) s2@(Stencil _ _ f2) = Stencil newSz maxCenter stF + where + stF ug gV !ix = f (f1 ug gV ix) (f2 ug gV ix) + {-# INLINE stF #-} + !newSz = unionStencilSizes maxCenter s1 s2 + !maxCenter = unionStencilCenters s1 s2 + {-# INLINE liftA2 #-} +#endif + instance (Index ix, Num a) => Num (Stencil ix e a) where (+) = liftA2 (+) {-# INLINE (+) #-} @@ -257,21 +172,3 @@ {-# INLINE acosh #-} atanh = fmap atanh {-# INLINE atanh #-} - - -safeStencilIndex :: Index ix => Array D ix e -> ix -> e -safeStencilIndex DArray {..} ix - | isSafeIndex dSize ix = dIndex ix - | otherwise = throw $ IndexOutOfBoundsException dSize ix - - --- | Make sure constructed stencil doesn't index outside the allowed stencil size boundary. -validateStencil - :: Index ix - => e -> Stencil ix e a -> Stencil ix e a -validateStencil d s@(Stencil sSz sCenter stencil) - | isSafeIndex sSz sCenter = - let valArr = DArray Seq sSz (const d) - in stencil (Value . safeStencilIndex valArr) sCenter `seq` s - | otherwise = throw $ IndexOutOfBoundsException sSz sCenter -{-# INLINE validateStencil #-} diff --git a/src/Data/Massiv/Array/Stencil/Unsafe.hs b/src/Data/Massiv/Array/Stencil/Unsafe.hs --- a/src/Data/Massiv/Array/Stencil/Unsafe.hs +++ b/src/Data/Massiv/Array/Stencil/Unsafe.hs @@ -2,92 +2,27 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE RecordWildCards #-} + -- | -- Module : Data.Massiv.Array.Stencil.Unsafe --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Stencil.Unsafe - ( -- * Stencil - makeUnsafeStencil - , mapStencilUnsafe - , forStencilUnsafe - ) where +module Data.Massiv.Array.Stencil.Unsafe ( + -- * Stencil + makeUnsafeStencil, + makeUnsafeConvolutionStencil, + makeUnsafeCorrelationStencil, + unsafeTransformStencil, +) where -import Data.Massiv.Array.Delayed.Windowed (Array(..), DW, Window(..), - insertWindow) import Data.Massiv.Array.Stencil.Internal import Data.Massiv.Core.Common import GHC.Exts (inline) - --- | Just as `mapStencilUnsafe` this is an unsafe version of the stencil --- mapping. Arguments are in slightly different order and the indexing function returns --- `Nothing` for elements outside the border. --- --- @since 0.1.7 -forStencilUnsafe :: - (Source r ix e, Manifest r ix e) - => Array r ix e - -> Sz ix -- ^ Size of the stencil - -> ix -- ^ Center of the stencil - -> ((ix -> Maybe e) -> a) - -- ^ Stencil function that receives a "get" function as it's argument that can - -- retrieve values of cells in the source array with respect to the center of - -- the stencil. Stencil function must return a value that will be assigned to - -- the cell in the result array. Offset supplied to the "get" function - -- cannot go outside the boundaries of the stencil. - -> Array DW ix a -forStencilUnsafe !arr !sSz !sCenter relStencil = - insertWindow (DArray (getComp arr) sz (stencil (index arr))) window - where - !window = - Window - { windowStart = sCenter - , windowSize = windowSz - , windowIndex = stencil (Just . unsafeIndex arr) - , windowUnrollIx2 = unSz . fst <$> pullOutSzM windowSz 2 - } - !sz = size arr - !windowSz = Sz (liftIndex2 (-) (unSz sz) (liftIndex (subtract 1) (unSz sSz))) - stencil getVal !ix = inline relStencil $ \ !ixD -> getVal (liftIndex2 (+) ix ixD) - {-# INLINE stencil #-} -{-# INLINE forStencilUnsafe #-} - - --- | This is an unsafe version of `Data.Massiv.Array.Stencil.mapStencil`, that does no --- stencil validation. There is no performance difference between the two, but the unsafe --- version has an advantage of not requiring to deal with `Value` wrapper. --- --- @since 0.4.3 -mapStencilUnsafe :: - Manifest r ix e - => Border e - -> Sz ix - -> ix - -> ((ix -> e) -> a) - -> Array r ix e - -> Array DW ix a -mapStencilUnsafe b sSz sCenter stencilF !arr = insertWindow warr window - where - !warr = DArray (getComp arr) sz (stencil (borderIndex b arr)) - !window = - Window - { windowStart = sCenter - , windowSize = windowSz - , windowIndex = stencil (unsafeIndex arr) - , windowUnrollIx2 = unSz . fst <$> pullOutSzM sSz 2 - } - !sz = size arr - !windowSz = Sz (liftIndex2 (-) (unSz sz) (liftIndex (subtract 1) (unSz sSz))) - stencil getVal !ix = inline stencilF $ \ !ixD -> getVal (liftIndex2 (+) ix ixD) - {-# INLINE stencil #-} -{-# INLINE mapStencilUnsafe #-} - - -- | Similar to `Data.Massiv.Array.Stencil.makeStencil`, but there are no guarantees that the -- stencil will not read out of bounds memory. This stencil is also a bit more powerful in sense it -- gets an extra peice of information, namely the exact index for the element it is constructing. @@ -95,14 +30,130 @@ -- @since 0.3.0 makeUnsafeStencil :: Index ix - => Sz ix -- ^ Size of the stencil - -> ix -- ^ Center of the stencil + => Sz ix + -- ^ Size of the stencil + -> ix + -- ^ Center of the stencil -> (ix -> (ix -> e) -> a) -- ^ Stencil function. -> Stencil ix e a makeUnsafeStencil !sSz !sCenter relStencil = Stencil sSz sCenter stencil where - stencil getVal !ix = - Value $ inline $ relStencil ix (unValue . getVal . liftIndex2 (+) ix) + stencil unsafeGetVal _getVal !ix = + inline (relStencil ix (unsafeGetVal . liftIndex2 (+) ix)) {-# INLINE stencil #-} {-# INLINE makeUnsafeStencil #-} + +-- | Same as `Data.Massiv.Array.Stencil.makeConvolutionStencil`, but will result in +-- reading memory out of bounds and potential segfaults if supplied arguments are not valid. +-- +-- @since 0.6.0 +makeUnsafeConvolutionStencil + :: (Index ix, Num e) + => Sz ix + -> ix + -> ((ix -> e -> e -> e) -> e -> e) + -> Stencil ix e e +makeUnsafeConvolutionStencil !sz !sCenter relStencil = + Stencil sz sInvertCenter stencil + where + !sInvertCenter = liftIndex2 (-) (liftIndex (subtract 1) (unSz sz)) sCenter + stencil uget _ !ix = + (inline relStencil $ \ !ixD !kVal !acc -> uget (liftIndex2 (-) ix ixD) * kVal + acc) 0 + {-# INLINE stencil #-} +{-# INLINE makeUnsafeConvolutionStencil #-} + +-- | Same as `Data.Massiv.Array.Stencil.makeCorrelationStencil`, but will result in +-- reading memory out of bounds and potential segfaults if supplied arguments are not +-- valid. +-- +-- @since 0.6.0 +makeUnsafeCorrelationStencil + :: (Index ix, Num e) + => Sz ix + -> ix + -> ((ix -> e -> e -> e) -> e -> e) + -> Stencil ix e e +makeUnsafeCorrelationStencil !sSz !sCenter relStencil = Stencil sSz sCenter stencil + where + stencil _ getVal !ix = + (inline relStencil $ \ !ixD !kVal !acc -> getVal (liftIndex2 (+) ix ixD) * kVal + acc) 0 + {-# INLINE stencil #-} +{-# INLINE makeUnsafeCorrelationStencil #-} + +-- | Perform an arbitrary transformation of a stencil. This stencil modifier can be used for +-- example to turn a vector stencil into a matrix stencil implement, or transpose a matrix +-- stencil. It is really easy to get this wrong, so be extremely careful. +-- +-- ====__Examples__ +-- +-- Convert a 1D stencil into a row or column 2D stencil: +-- +-- >>> import Data.Massiv.Array +-- >>> import Data.Massiv.Array.Unsafe +-- >>> let arr = compute $ iterateN 3 succ 0 :: Array P Ix2 Int +-- >>> arr +-- Array P Seq (Sz (3 :. 3)) +-- [ [ 1, 2, 3 ] +-- , [ 4, 5, 6 ] +-- , [ 7, 8, 9 ] +-- ] +-- >>> let rowStencil = unsafeTransformStencil (\(Sz n) -> Sz (1 :. n)) (0 :.) $ \ f uget getVal (i :. j) -> f (uget . (i :.)) (getVal . (i :.)) j +-- >>> applyStencil noPadding (rowStencil (sumStencil (Sz1 3))) arr +-- Array DW Seq (Sz (3 :. 1)) +-- [ [ 6 ] +-- , [ 15 ] +-- , [ 24 ] +-- ] +-- >>> let columnStencil = unsafeTransformStencil (\(Sz n) -> Sz (n :. 1)) (:. 0) $ \ f uget getVal (i :. j) -> f (uget . (:. j)) (getVal . (:. j)) i +-- >>> applyStencil noPadding (columnStencil (sumStencil (Sz1 3))) arr +-- Array DW Seq (Sz (1 :. 3)) +-- [ [ 12, 15, 18 ] +-- ] +-- +-- @since 0.5.4 +unsafeTransformStencil + :: (Sz ix' -> Sz ix) + -- ^ Forward modifier for the size + -> (ix' -> ix) + -- ^ Forward index modifier + -> ( ((ix' -> e) -> (ix' -> e) -> ix' -> a) + -> (ix -> e) + -> (ix -> e) + -> ix + -> a + ) + -- ^ Inverse stencil function modifier + -> Stencil ix' e a + -- ^ Original stencil. + -> Stencil ix e a +unsafeTransformStencil transformSize transformIndex transformFunc Stencil{..} = + Stencil + { stencilSize = transformSize stencilSize + , stencilCenter = transformIndex stencilCenter + , stencilFunc = transformFunc stencilFunc + } +{-# INLINE unsafeTransformStencil #-} + +{- + +Invalid stencil transformer function. + +TODO: figure out if there is a safe way to do stencil index trnasformation. + +transformStencil :: + (Default e, Index ix) + => (Sz ix' -> Sz ix) + -- ^ Forward modifier for the size + -> (ix' -> ix) + -- ^ Forward index modifier + -> (ix -> ix') + -- ^ Inverse index modifier + -> Stencil ix' e a + -- ^ Original stencil. + -> Stencil ix e a +transformStencil transformSize transformIndex transformIndex' stencil = + validateStencil def $! unsafeTransformStencil transformSize transformIndex transformIndex' stencil +{-# INLINE transformStencil #-} + +-} 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 @@ -1,103 +1,156 @@ {-# LANGUAGE BangPatterns #-} +{-# LANGUAGE ExplicitForAll #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE PatternSynonyms #-} + -- | -- Module : Data.Massiv.Array.Unsafe --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Unsafe - ( -- * Creation - unsafeMakeLoadArray - -- * Indexing - , Sz(SafeSz) - , Stride(SafeStride) - , unsafeIndex - , unsafeLinearIndex - , unsafeLinearIndexM - -- * Manipulations - , unsafeBackpermute - , unsafeResize - , unsafeExtract - , unsafeTransform - , unsafeTransform2 - -- * Slicing - , unsafeSlice - , unsafeOuterSlice - , unsafeInnerSlice - -- , unsafeLinearSlice - -- * Mutable interface - , unsafeThaw - , unsafeFreeze - , unsafeNew - -- ** Read - , unsafeRead - , unsafeLinearRead - -- ** Write - , unsafeWrite - , unsafeLinearWrite - -- ** Modify - , unsafeModify - , unsafeLinearModify - -- ** Swap - , unsafeSwap - , unsafeLinearSwap - -- ** Range modification - , unsafeLinearSet - , unsafeLinearCopy - , unsafeArrayLinearCopy - -- ** Resizing - , unsafeLinearShrink - , unsafeLinearGrow - -- * Pointer access - , unsafeWithPtr - , unsafeArrayToForeignPtr - , unsafeMArrayToForeignPtr - , unsafeArrayFromForeignPtr - , unsafeArrayFromForeignPtr0 - , unsafeMArrayFromForeignPtr - , unsafeMArrayFromForeignPtr0 - -- ** Atomic Operations - , unsafeAtomicReadIntArray - , unsafeAtomicWriteIntArray - , unsafeAtomicModifyIntArray - , unsafeAtomicAddIntArray - , unsafeAtomicSubIntArray - , unsafeAtomicAndIntArray - , unsafeAtomicNandIntArray - , unsafeAtomicOrIntArray - , unsafeAtomicXorIntArray - , unsafeCasIntArray - -- ** Other operations - , unsafeUnstablePartitionRegionM - ) where +module Data.Massiv.Array.Unsafe ( + -- * Creation + unsafeMakeLoadArray, + unsafeMakeLoadArrayAdjusted, -import Data.Massiv.Array.Delayed.Pull (D) -import Data.Massiv.Array.Delayed.Push (unsafeMakeLoadArray) + -- * Indexing + Sz (SafeSz), + Stride (SafeStride), + unsafeIndex, + unsafePrefIndex, + unsafeLinearIndex, + unsafeLinearIndexM, + + -- * Manipulations + unsafeBackpermute, + unsafeResize, + unsafeExtract, + unsafeTransform, + unsafeTransform2, + + -- * Slicing + unsafeSlice, + unsafeOuterSlice, + unsafeInnerSlice, + unsafeLinearSlice, + + -- * Mutable interface + unsafeResizeMArray, + unsafeLinearSliceMArray, + unsafeThaw, + unsafeFreeze, + unsafeNew, + unsafeLoadIntoST, + unsafeLoadIntoIO, + unsafeLoadIntoS, + unsafeLoadIntoM, + unsafeCreateArray, + unsafeCreateArray_, + unsafeCreateArrayS, + + -- ** Read + unsafeRead, + unsafeLinearRead, + + -- ** Write + unsafeWrite, + unsafeLinearWrite, + + -- ** Modify + unsafeModify, + unsafeLinearModify, + + -- ** Swap + unsafeSwap, + unsafeLinearSwap, + + -- ** Range modification + unsafeLinearSet, + unsafeLinearCopy, + unsafeArrayLinearCopy, + + -- ** Resizing + unsafeLinearShrink, + unsafeLinearGrow, + + -- * Pointer access + unsafeMallocMArray, + unsafeWithPtr, + unsafeArrayToForeignPtr, + unsafeMArrayToForeignPtr, + unsafeArrayFromForeignPtr, + unsafeArrayFromForeignPtr0, + unsafeMArrayFromForeignPtr, + unsafeMArrayFromForeignPtr0, + + -- ** Atomic Operations + unsafeAtomicReadIntArray, + unsafeAtomicWriteIntArray, + unsafeAtomicModifyIntArray, + unsafeAtomicAddIntArray, + unsafeAtomicSubIntArray, + unsafeAtomicAndIntArray, + unsafeAtomicNandIntArray, + unsafeAtomicOrIntArray, + unsafeAtomicXorIntArray, + unsafeCasIntArray, + + -- ** Other operations + coerceBoxedArray, + coerceNormalBoxedArray, + unsafeUnstablePartitionRegionM, + module Data.Massiv.Vector.Unsafe, + module Data.Massiv.Array.Stencil.Unsafe, + + -- * Constructors + Array (PArray, SArray, UArray, BArray, BLArray, BNArray, DArray, DLArray, DSArray, DIArray, DWArray), + MArray (MPArray, MSArray, MUArray, MBArray, MBLArray, MBNArray), +) where + +import Data.Massiv.Array.Delayed.Interleaved (Array (DIArray)) +import Data.Massiv.Array.Delayed.Pull ( + D, + unsafeExtract, + unsafeInnerSlice, + unsafeSlice, + ) +import Data.Massiv.Array.Delayed.Push ( + Array (DLArray), + unsafeMakeLoadArray, + unsafeMakeLoadArrayAdjusted, + ) +import Data.Massiv.Array.Delayed.Stream (Array (DSArray)) +import Data.Massiv.Array.Delayed.Windowed (Array (DWArray)) +import Data.Massiv.Array.Manifest.Boxed +import Data.Massiv.Array.Manifest.Internal import Data.Massiv.Array.Manifest.Primitive import Data.Massiv.Array.Manifest.Storable -import Data.Massiv.Core.Common -import Data.Massiv.Core.Index.Internal (Sz(SafeSz)) -import Data.Massiv.Core.Index.Stride (Stride(SafeStride)) +import Data.Massiv.Array.Manifest.Unboxed +import Data.Massiv.Array.Mutable.Internal import Data.Massiv.Array.Ops.Sort (unsafeUnstablePartitionRegionM) +import Data.Massiv.Array.Stencil.Unsafe +import Data.Massiv.Core.Common +import Data.Massiv.Core.Index.Stride (Stride (SafeStride)) +import Data.Massiv.Vector.Unsafe -unsafeBackpermute :: (Source r' ix' e, Index ix) => - Sz ix -> (ix -> ix') -> Array r' ix' e -> Array D ix e -unsafeBackpermute !sz ixF !arr = - makeArray (getComp arr) sz $ \ !ix -> unsafeIndex arr (ixF ix) +unsafeBackpermute + :: (Index ix', Source r' e, Index ix) + => Sz ix + -> (ix -> ix') + -> Array r' ix' e + -> Array D ix e +unsafeBackpermute !sz ixF !arr = makeArray (getComp arr) sz (unsafeIndex arr . ixF) {-# INLINE unsafeBackpermute #-} --- | Same `Data.Array.transform'`, except no bounds checking is performed, thus making it faster, +-- | Same 'Data.Array.transform'', except no bounds checking is performed, thus making it faster, -- but unsafe. -- -- @since 0.3.0 -unsafeTransform :: - (Source r' ix' e', Index ix) +unsafeTransform + :: (Index ix', Source r' e', Index ix) => (Sz ix' -> (Sz ix, a)) -> (a -> (ix' -> e') -> ix -> e) -> Array r' ix' e' @@ -107,12 +160,12 @@ (sz, a) = getSz (size arr) {-# INLINE unsafeTransform #-} --- | Same `Data.Array.transform2'`, except no bounds checking is performed, thus making it faster, +-- | Same 'Data.Array.transform2'', except no bounds checking is performed, thus making it faster, -- but unsafe. -- -- @since 0.3.0 -unsafeTransform2 :: - (Source r1 ix1 e1, Source r2 ix2 e2, Index ix) +unsafeTransform2 + :: (Index ix1, Source r1 e1, Index ix2, Source r2 e2, Index ix) => (Sz ix1 -> Sz ix2 -> (Sz ix, a)) -> (a -> (ix1 -> e1) -> (ix2 -> e2) -> ix -> e) -> Array r1 ix1 e1 diff --git a/src/Data/Massiv/Core.hs b/src/Data/Massiv/Core.hs --- a/src/Data/Massiv/Core.hs +++ b/src/Data/Massiv/Core.hs @@ -1,55 +1,71 @@ -- | -- Module : Data.Massiv.Core --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Core - ( Array(List, unList) - , Elt - , Construct - , Load(R, loadArrayM, defaultElement) - , Stream(..) - , Source - , Resize - , Extract - , StrideLoad(..) - , Slice - , OuterSlice - , InnerSlice - , Manifest - , Mutable - , Ragged - , Nested(..) - , NestedStruct - , L(..) - , LN - , ListItem - , Comp(Seq, Par, Par', ParOn, ParN) - , WorkerStates - , initWorkerStates - , module Data.Massiv.Core.Index +module Data.Massiv.Core ( + Array (LArray), + List (..), + Vector, + MVector, + Matrix, + MMatrix, + Load (iterArrayLinearST_, iterArrayLinearWithSetST_), + Stream (..), + Source, + PrefIndex (..), + Size, + Shape (..), + LengthHint (..), + StrideLoad (..), + Manifest, + Mutable, + Ragged, + L (..), + ListItem, + Scheduler, + SchedulerWS, + Strategy, + Comp (Seq, Par, Par', ParOn, ParN), + getComp, + setComp, + appComp, + WorkerStates, + initWorkerStates, + scheduleWork, + scheduleWork_, + module Data.Massiv.Core.Index, + + -- * Numeric + FoldNumeric, + Numeric, + NumericFloat, + -- * Exceptions - , MonadThrow(..) - , throw - , Exception(..) - , SomeException - , IndexException(..) - , SizeException(..) - , ShapeException(..) - , module Data.Massiv.Core.Exception + MonadThrow (..), + IndexException (..), + SizeException (..), + ShapeException (..), + module Data.Massiv.Core.Exception, + -- * Stateful Monads - , MonadUnliftIO - , MonadIO(liftIO) - , PrimMonad(PrimState) - ) where + MonadUnliftIO, + MonadIO (liftIO), + PrimMonad (PrimState), +) where -import Control.Exception (Exception(..), SomeException) -import Control.Scheduler (initWorkerStates) +import Control.Scheduler (SchedulerWS, initWorkerStates) import Data.Massiv.Core.Common +import Data.Massiv.Core.Exception import Data.Massiv.Core.Index import Data.Massiv.Core.List -import Data.Massiv.Core.Exception +import Data.Massiv.Core.Operations (FoldNumeric, Numeric, NumericFloat) +-- | Append computation strategy using `Comp`'s `Monoid` instance. +-- +-- @since 0.6.0 +appComp :: Strategy r => Comp -> Array r ix e -> Array r ix e +appComp comp arr = setComp (comp <> getComp arr) arr +{-# INLINEABLE appComp #-} 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 @@ -1,864 +1,1145 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE CPP #-} -{-# LANGUAGE DefaultSignatures #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE UndecidableInstances #-} --- | --- Module : Data.Massiv.Core.Common --- Copyright : (c) Alexey Kuleshevich 2018-2019 --- License : BSD3 --- Maintainer : Alexey Kuleshevich --- Stability : experimental --- Portability : non-portable -module Data.Massiv.Core.Common - ( Array - , Elt - , Steps(..) - , Stream(..) - , Construct(..) - , Source(..) - , Load(..) - , StrideLoad(..) - , Resize(..) - , Extract(..) - , Slice(..) - , OuterSlice(..) - , InnerSlice(..) - , Manifest(..) - , Mutable(..) - , Comp(..) - , Scheduler - , numWorkers - , scheduleWork - , scheduleWork_ - , WorkerStates - , unsafeRead - , unsafeWrite - , unsafeModify - , unsafeLinearModify - , unsafeSwap - , unsafeLinearSwap - , unsafeDefaultLinearShrink - , Ragged(..) - , Nested(..) - , NestedStruct - , empty - , singleton - -- * Size - , elemsCount - , isEmpty - , Sz(SafeSz) - , Size(..) - -- * Indexing - , (!?) - , index - , indexM - , (!) - , index' - , (??) - , defaultIndex - , borderIndex - , evaluateM - , evaluate' - , module Data.Massiv.Core.Index - -- * Common Operations - , imapM_ - , Semigroup((<>)) - -- * Exceptions - , MonadThrow(..) - , throw - , IndexException(..) - , SizeException(..) - , ShapeException(..) - , module Data.Massiv.Core.Exception - , Proxy(..) - , Id(..) - -- * Stateful Monads - , MonadUnliftIO - , MonadIO(liftIO) - , PrimMonad(PrimState) - ) where - -#if !MIN_VERSION_base(4,11,0) -import Data.Semigroup -#endif -import Control.Exception (throw) -import Control.Monad.Catch (MonadThrow(..)) -import Control.Monad.IO.Unlift (MonadIO(liftIO), MonadUnliftIO) -import Control.Monad.Primitive -import Control.Scheduler (Comp(..), Scheduler, WorkerStates, numWorkers, - scheduleWork, scheduleWork_, withScheduler_, trivialScheduler_) -import Data.Massiv.Core.Exception -import Data.Massiv.Core.Index -import Data.Massiv.Core.Index.Internal (Sz(SafeSz)) -import Data.Typeable -import Data.Vector.Fusion.Bundle.Size -import qualified Data.Vector.Fusion.Stream.Monadic as S -import Data.Vector.Fusion.Util - -#include "massiv.h" - --- | The array family. Representations @r@ describes how data is arranged or computed. All arrays --- have a common property that each index @ix@ always maps to the same unique element, even if that --- element does not exist in memory and has to be computed upon lookup. Data is always arranged in a --- nested fashion, depth of which is controlled by @`Rank` ix@. -data family Array r ix e :: * - -type family Elt r ix e :: * where - Elt r Ix1 e = e - Elt r ix e = Array (R r) (Lower ix) e - -type family NestedStruct r ix e :: * - - - -class Stream r ix e where - toStream :: Array r ix e -> Steps Id e - -data Steps m e = Steps - { stepsStream :: S.Stream m e - , stepsSize :: Size - } - -instance Monad m => Functor (Steps m) where - fmap f s = s { stepsStream = S.map f (stepsStream s) } - {-# INLINE fmap #-} - - --- | Array types that can be constructed. -class (Typeable r, Index ix) => Construct r ix e where - {-# MINIMAL setComp,(makeArray|makeArrayLinear) #-} - - -- | Set computation strategy for this array - -- - -- ==== __Example__ - -- - -- >>> :set -XTypeApplications - -- >>> import Data.Massiv.Array - -- >>> a = singleton @DL @Ix1 @Int 0 - -- >>> a - -- Array DL Seq (Sz1 1) - -- [ 0 ] - -- >>> setComp (ParN 6) a -- use 6 capabilities - -- Array DL (ParN 6) (Sz1 1) - -- [ 0 ] - -- - setComp :: Comp -> Array r ix e -> Array r ix e - - -- | Construct an Array. Resulting type either has to be unambiguously inferred or restricted - -- manually, like in the example below. Use "Data.Massiv.Array.makeArrayR" if you'd like to - -- specify representation as an argument. - -- - -- >>> import Data.Massiv.Array - -- >>> makeArray Seq (Sz (3 :. 4)) (\ (i :. j) -> if i == j then i else 0) :: Array D Ix2 Int - -- Array D Seq (Sz (3 :. 4)) - -- [ [ 0, 0, 0, 0 ] - -- , [ 0, 1, 0, 0 ] - -- , [ 0, 0, 2, 0 ] - -- ] - -- - -- Instead of restricting the full type manually we can use `TypeApplications` as convenience: - -- - -- >>> :set -XTypeApplications - -- >>> makeArray @P @_ @Double Seq (Sz2 3 4) $ \(i :. j) -> logBase (fromIntegral i) (fromIntegral j) - -- Array P Seq (Sz (3 :. 4)) - -- [ [ NaN, -0.0, -0.0, -0.0 ] - -- , [ -Infinity, NaN, Infinity, Infinity ] - -- , [ -Infinity, 0.0, 1.0, 1.5849625007211563 ] - -- ] - -- - -- @since 0.1.0 - makeArray :: - Comp -- ^ Computation strategy. Useful constructors are `Seq` and `Par` - -> Sz ix -- ^ Size of the result array. - -> (ix -> e) -- ^ Function to generate elements at a particular index - -> Array r ix e - makeArray comp sz f = makeArrayLinear comp sz (f . fromLinearIndex sz) - {-# INLINE makeArray #-} - - -- | Same as `makeArray`, but produce elements using linear row-major index. - -- - -- >>> import Data.Massiv.Array - -- >>> makeArrayLinear Seq (Sz (2 :. 4)) id :: Array D Ix2 Int - -- Array D Seq (Sz (2 :. 4)) - -- [ [ 0, 1, 2, 3 ] - -- , [ 4, 5, 6, 7 ] - -- ] - -- - -- @since 0.3.0 - makeArrayLinear :: Comp -> Sz ix -> (Int -> e) -> Array r ix e - makeArrayLinear comp sz f = makeArray comp sz (f . toLinearIndex sz) - {-# INLINE makeArrayLinear #-} - - - -class Index ix => Resize r ix where - -- | /O(1)/ - Change the size of an array. Total number of elements should be the same, but it is - -- not validated. - unsafeResize :: Index ix' => Sz ix' -> Array r ix e -> Array r ix' e - - -class Load r ix e => Extract r ix e where - -- | /O(1)/ - Extract a portion of an array. Staring index and new size are - -- not validated. - unsafeExtract :: ix -> Sz ix -> Array r ix e -> Array (R r) ix e - - --- | Arrays that can be used as source to practically any manipulation function. -class Load r ix e => Source r ix e where - {-# MINIMAL (unsafeIndex|unsafeLinearIndex) #-} - - -- | Lookup element in the array. No bounds check is performed and access of - -- arbitrary memory is possible when invalid index is supplied. - -- - -- @since 0.1.0 - unsafeIndex :: Array r ix e -> ix -> e - unsafeIndex = - INDEX_CHECK("(Source r ix e).unsafeIndex", - size, \ !arr -> unsafeLinearIndex arr . toLinearIndex (size arr)) - {-# INLINE unsafeIndex #-} - - -- | Lookup element in the array using flat index in a row-major fashion. No - -- bounds check is performed - -- - -- @since 0.1.0 - unsafeLinearIndex :: Array r ix e -> Int -> e - unsafeLinearIndex !arr = unsafeIndex arr . fromLinearIndex (size arr) - {-# INLINE unsafeLinearIndex #-} - - -- -- | Source arrays also give us ability to look at their linear slices - -- -- - -- -- @since 0.4.0 - -- unsafeLinearSlice :: Ix1 -> Sz1 -> Array r ix e -> Array r Ix1 e - --- | Any array that can be computed and loaded into memory -class (Typeable r, Index ix) => Load r ix e where - type family R r :: * - type instance R r = r - - -- | Get computation strategy of this array - -- - -- @since 0.1.0 - getComp :: Array r ix e -> Comp - - -- | Get the size of an immutabe array - -- - -- @since 0.1.0 - size :: Array r ix e -> Sz ix - - - -- | Load an array into memory. - -- - -- @since 0.3.0 - loadArrayM - :: Monad m => - Scheduler m () - -> Array r ix e -- ^ Array that is being loaded - -> (Int -> e -> m ()) -- ^ Function that writes an element into target array - -> m () - - defaultElement :: Array r ix e -> Maybe e - defaultElement _ = Nothing - {-# INLINE defaultElement #-} - - -- | /O(1)/ - Get the possible maximum size of an immutabe array. If the lookup of size - -- in constant time is not possible, `Nothing` should be returned. This value will be - -- used as the initial size of the mutable array in which loading will happen. - -- - -- @since 0.4.1 - maxSize :: Array r ix e -> Maybe (Sz ix) - maxSize = Just . size - {-# INLINE maxSize #-} - - -- | Load into a supplied mutable array sequentially. Returned array does npt have to be - -- the same - -- - -- @since 0.4.1 - unsafeLoadIntoS :: - (Mutable r' ix e, PrimMonad m) - => MArray (PrimState m) r' ix e - -> Array r ix e - -> m (MArray (PrimState m) r' ix e) - unsafeLoadIntoS marr arr = do - loadArrayM trivialScheduler_ arr (unsafeLinearWrite marr) - pure marr - {-# INLINE unsafeLoadIntoS #-} - - -- | Same as `unsafeLoadIntoS`, but with respect of computation startegy. - -- - -- @since 0.4.1 - unsafeLoadInto :: - (Mutable r' ix e, MonadIO m) - => MArray RealWorld r' ix e - -> Array r ix e - -> m (MArray RealWorld r' ix e) - unsafeLoadInto marr arr = do - liftIO $ withScheduler_ (getComp arr) $ \scheduler -> - loadArrayM scheduler arr (unsafeLinearWrite marr) - pure marr - {-# INLINE unsafeLoadInto #-} - - -class Load r ix e => StrideLoad r ix e where - -- | Load an array into memory with stride. Default implementation requires an instance of - -- `Source`. - loadArrayWithStrideM - :: Monad m => - Scheduler m () - -> Stride ix -- ^ Stride to use - -> Sz ix -- ^ Size of the target array affected by the stride. - -> Array r ix e -- ^ Array that is being loaded - -> (Int -> e -> m ()) -- ^ Function that writes an element into target array - -> m () - default loadArrayWithStrideM - :: (Source r ix e, Monad m) => - Scheduler m () - -> Stride ix - -> Sz ix - -> Array r ix e - -> (Int -> e -> m ()) - -> m () - loadArrayWithStrideM scheduler stride resultSize arr = - splitLinearlyWith_ scheduler (totalElem resultSize) unsafeLinearWriteWithStride - where - !strideIx = unStride stride - unsafeLinearWriteWithStride = - unsafeIndex arr . liftIndex2 (*) strideIx . fromLinearIndex resultSize - {-# INLINE unsafeLinearWriteWithStride #-} - {-# INLINE loadArrayWithStrideM #-} - - -class Load r ix e => OuterSlice r ix e where - -- | /O(1)/ - Take a slice out of an array from the outside - unsafeOuterSlice :: Array r ix e -> Int -> Elt r ix e - -class Load r ix e => InnerSlice r ix e where - unsafeInnerSlice :: Array r ix e -> (Sz (Lower ix), Sz Int) -> Int -> Elt r ix e - -class Load r ix e => Slice r ix e where - unsafeSlice :: MonadThrow m => Array r ix e -> ix -> Sz ix -> Dim -> m (Elt r ix e) - - --- | Manifest arrays are backed by actual memory and values are looked up versus --- computed as it is with delayed arrays. Because of this fact indexing functions --- @(`!`)@, @(`!?`)@, etc. are constrained to manifest arrays only. -class (Load r ix e, Source r ix e) => Manifest r ix e where - - unsafeLinearIndexM :: Array r ix e -> Int -> e - - -class (Construct r ix e, Manifest r ix e) => Mutable r ix e where - data MArray s r ix e :: * - - -- | Get the size of a mutable array. - -- - -- @since 0.1.0 - msize :: MArray s r ix e -> Sz ix - - -- | Convert immutable array into a mutable array without copy. - -- - -- @since 0.1.0 - unsafeThaw :: PrimMonad m => Array r ix e -> m (MArray (PrimState m) r ix e) - - -- | Convert mutable array into an immutable array without copy. - -- - -- @since 0.1.0 - unsafeFreeze :: PrimMonad m => Comp -> MArray (PrimState m) r ix e -> m (Array r ix e) - - -- | Create new mutable array, leaving it's elements uninitialized. Size isn't validated either. - -- - -- @since 0.1.0 - unsafeNew :: PrimMonad m => Sz ix -> m (MArray (PrimState m) r ix e) - - -- | Read an element at linear row-major index - -- - -- @since 0.1.0 - unsafeLinearRead :: PrimMonad m => MArray (PrimState m) r ix e -> Int -> m e - - -- | Write an element into mutable array with linear row-major index - -- - -- @since 0.1.0 - unsafeLinearWrite :: PrimMonad m => MArray (PrimState m) r ix e -> Int -> e -> m () - - -- | Initialize mutable array to some default value. - -- - -- @since 0.3.0 - initialize :: PrimMonad m => MArray (PrimState m) r ix e -> m () - - -- | Create new mutable array while initializing all elements to some default value. - -- - -- @since 0.3.0 - initializeNew :: PrimMonad m => Maybe e -> Sz ix -> m (MArray (PrimState m) r ix e) - initializeNew mdef sz = do - marr <- unsafeNew sz - case mdef of - Just val -> unsafeLinearSet marr 0 (SafeSz (totalElem sz)) val - Nothing -> initialize marr - 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 -> Ix1 -> Sz1 -> e -> m () - unsafeLinearSet marr offset len e = - loopM_ offset (< (offset + unSz len)) (+1) (\i -> unsafeLinearWrite marr i e) - {-# INLINE unsafeLinearSet #-} - - -- | Copy part of one mutable array into another - -- - -- @since 0.3.6 - unsafeLinearCopy :: (Mutable r ix' e, 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 :: (Mutable r ix' e, 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 or - -- equal. There is no guarantee that the original array is left unchanged, so it should - -- no longer be used. - -- - -- @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 - -- or equal. There is no guarantee that the original array is left unchanged, so it - -- should no longer be used. - -- - -- @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 --- --- @since 0.1.0 -unsafeRead :: (Mutable r ix e, PrimMonad m) => - MArray (PrimState m) r ix e -> ix -> m e -unsafeRead !marr !ix = unsafeLinearRead marr (toLinearIndex (msize marr) ix) -{-# INLINE unsafeRead #-} - --- | Write an element into array --- --- @since 0.1.0 -unsafeWrite :: (Mutable r ix e, PrimMonad m) => - MArray (PrimState m) r ix e -> ix -> e -> m () -unsafeWrite !marr !ix = unsafeLinearWrite marr (toLinearIndex (msize marr) ix) -{-# INLINE unsafeWrite #-} - - --- | Modify an element in the array with a monadic action. Returns the previous value. --- --- @since 0.4.0 -unsafeLinearModify :: (Mutable r ix e, PrimMonad m) => - MArray (PrimState m) r ix e -> (e -> m e) -> Int -> m e -unsafeLinearModify !marr f !i = do - v <- unsafeLinearRead marr i - v' <- f v - unsafeLinearWrite marr i v' - pure v -{-# INLINE unsafeLinearModify #-} - --- | Modify an element in the array with a monadic action. Returns the previous value. --- --- @since 0.4.0 -unsafeModify :: (Mutable r ix e, PrimMonad m) => - MArray (PrimState m) r ix e -> (e -> m e) -> ix -> m e -unsafeModify marr f ix = unsafeLinearModify marr f (toLinearIndex (msize marr) ix) -{-# INLINE unsafeModify #-} - --- | Swap two elements in a mutable array under the supplied indices. Returns the previous --- values. --- --- @since 0.4.0 -unsafeSwap :: (Mutable r ix e, PrimMonad m) => - MArray (PrimState m) r ix e -> ix -> ix -> m (e, e) -unsafeSwap !marr !ix1 !ix2 = unsafeLinearSwap marr (toLinearIndex sz ix1) (toLinearIndex sz ix2) - where sz = msize marr -{-# INLINE unsafeSwap #-} - - --- | Swap two elements in a mutable array under the supplied linear indices. Returns the --- previous values. --- --- @since 0.4.0 -unsafeLinearSwap :: (Mutable r ix e, PrimMonad m) => - MArray (PrimState m) r ix e -> Int -> Int -> m (e, e) -unsafeLinearSwap !marr !i1 !i2 = do - val1 <- unsafeLinearRead marr i1 - val2 <- unsafeLinearRead marr i2 - unsafeLinearWrite marr i1 val2 - unsafeLinearWrite marr i2 val1 - return (val1, val2) -{-# INLINE unsafeLinearSwap #-} - - -class Nested r ix e where - fromNested :: NestedStruct r ix e -> Array r ix e - - toNested :: Array r ix e -> NestedStruct r ix e - -class Construct r ix e => Ragged r ix e where - - emptyR :: Comp -> Array r ix e - - isNull :: Array r ix e -> Bool - - consR :: Elt r ix e -> Array r ix e -> Array r ix e - - unconsR :: Array r ix e -> Maybe (Elt r ix e, Array r ix e) - - generateRaggedM :: Monad m => Comp -> Sz ix -> (ix -> m e) -> m (Array r ix e) - - edgeSize :: Array r ix e -> Sz ix - - flattenRagged :: Array r ix e -> Array r Ix1 e - - loadRagged :: - Monad m => (m () -> m ()) -> (Int -> e -> m a) -> Int -> Int -> Sz ix -> Array r ix e -> m () - - -- TODO: test property: - -- (read $ raggedFormat show "\n" (ls :: Array L (IxN n) Int)) == ls - raggedFormat :: (e -> String) -> String -> Array r ix e -> String - - - --- | Create an Array with no elements. By itself it is not particularly useful, but it serves as a --- nice base for constructing larger arrays. --- --- ==== __Examples__ --- --- >>> import Data.Massiv.Array as A --- >>> :set -XTypeApplications --- >>> xs = empty @DL @Ix1 @Double --- >>> snoc (cons 4 (cons 5 xs)) 22 --- Array DL Seq (Sz1 3) --- [ 4.0, 5.0, 22.0 ] --- --- @since 0.3.0 -empty :: - forall r ix e. Construct r ix e - => Array r ix e -empty = makeArray Seq zeroSz (const (throwImpossible Uninitialized)) -{-# INLINE empty #-} - --- | Create an Array with a single element. --- --- ==== __Examples__ --- --- >>> import Data.Massiv.Array as A --- >>> singleton 7 :: Array D Ix4 Double --- Array D Seq (Sz (1 :> 1 :> 1 :. 1)) --- [ [ [ [ 7.0 ] --- ] --- ] --- ] --- --- Instead of specifying type signature we could use @TypeApplications@ --- --- >>> :set -XTypeApplications --- >>> singleton @U @Ix4 @Double 7 --- Array U Seq (Sz (1 :> 1 :> 1 :. 1)) --- [ [ [ [ 7.0 ] --- ] --- ] --- ] --- --- @since 0.1.0 -singleton :: - forall r ix e. Construct r ix e - => e -- ^ The only element - -> Array r ix e -singleton = makeArray Seq oneSz . const -{-# INLINE singleton #-} - - -infixl 4 !, !?, ?? - --- | Infix version of `index'`. --- --- ==== __Examples__ --- --- >>> import Data.Massiv.Array as A --- >>> a = computeAs U $ iterateN (Sz (2 :. 3)) succ (0 :: Int) --- >>> a --- Array U Seq (Sz (2 :. 3)) --- [ [ 1, 2, 3 ] --- , [ 4, 5, 6 ] --- ] --- >>> a ! 0 :. 2 --- 3 --- >>> a ! 0 :. 3 --- *** Exception: IndexOutOfBoundsException: (0 :. 3) is not safe for (Sz (2 :. 3)) --- --- @since 0.1.0 -(!) :: Manifest r ix e => Array r ix e -> ix -> e -(!) = index' -{-# INLINE (!) #-} - - --- | Infix version of `indexM`. --- --- ==== __Examples__ --- --- >>> import Data.Massiv.Array as A --- >>> :set -XTypeApplications --- >>> a <- fromListsM @U @Ix2 @Int Seq [[1,2,3],[4,5,6]] --- >>> a --- Array U Seq (Sz (2 :. 3)) --- [ [ 1, 2, 3 ] --- , [ 4, 5, 6 ] --- ] --- >>> a !? 0 :. 2 --- 3 --- >>> a !? 0 :. 3 --- *** Exception: IndexOutOfBoundsException: (0 :. 3) is not safe for (Sz (2 :. 3)) --- >>> a !? 0 :. 3 :: Maybe Int --- Nothing --- --- @since 0.1.0 -(!?) :: (Manifest r ix e, MonadThrow m) => Array r ix e -> ix -> m e -(!?) = indexM -{-# INLINE (!?) #-} - - --- | /O(1)/ - Lookup an element in the array, where array itself is wrapped with --- `MonadThrow`. This operator is useful when used together with slicing or other --- functions that can fail. --- --- ==== __Examples__ --- --- >>> import Data.Massiv.Array as A --- >>> :set -XTypeApplications --- >>> ma = fromListsM @U @Ix3 @Int @Maybe Seq [[[1,2,3]],[[4,5,6]]] --- >>> ma --- Just (Array U Seq (Sz (2 :> 1 :. 3)) --- [ [ [ 1, 2, 3 ] --- ] --- , [ [ 4, 5, 6 ] --- ] --- ] --- ) --- >>> ma ??> 1 --- Just (Array M Seq (Sz (1 :. 3)) --- [ [ 4, 5, 6 ] --- ] --- ) --- >>> ma ??> 1 ?? 0 :. 2 --- Just 6 --- >>> ma ?? 1 :> 0 :. 2 --- Just 6 --- --- @since 0.1.0 -(??) :: (Manifest r ix e, MonadThrow m) => m (Array r ix e) -> ix -> m e -(??) marr ix = marr >>= (!? ix) -{-# INLINE (??) #-} - --- | /O(1)/ - Lookup an element in the array. Returns `Nothing`, when index is out of bounds and --- returns the element at the supplied index otherwise. Use `indexM` instead, since it is more --- generaland can just as well be used with `Maybe`. --- --- @since 0.1.0 -index :: Manifest r ix e => Array r ix e -> ix -> Maybe e -index = indexM -{-# INLINE index #-} - --- | /O(1)/ - Lookup an element in the array. Throws `IndexOutOfBoundsException`, when index is out --- of bounds and returns the element at the supplied index otherwise. --- --- @since 0.3.0 -indexM :: (Manifest r ix e, MonadThrow m) => Array r ix e -> ix -> m e -indexM = evaluateM -{-# INLINE indexM #-} - --- | /O(1)/ - Lookup an element in the array, while using default element when index is out of --- bounds. --- --- ==== __Examples__ --- --- >>> import Data.Massiv.Array --- >>> :set -XOverloadedLists --- >>> xs = [0..100] :: Array P Ix1 Int --- >>> defaultIndex 999 xs 100 --- 100 --- >>> defaultIndex 999 xs 101 --- 999 --- --- @since 0.1.0 -defaultIndex :: Manifest r ix e => e -> Array r ix e -> ix -> e -defaultIndex defVal = borderIndex (Fill defVal) -{-# INLINE defaultIndex #-} - --- | /O(1)/ - Lookup an element in the array. Use a border resolution technique --- when index is out of bounds. --- --- ==== __Examples__ --- --- >>> import Data.Massiv.Array as A --- >>> :set -XOverloadedLists --- >>> xs = [0..100] :: Array U Ix1 Int --- >>> borderIndex Wrap xs <$> range Seq 99 104 --- Array D Seq (Sz1 5) --- [ 99, 100, 0, 1, 2 ] --- --- @since 0.1.0 -borderIndex :: Manifest r ix e => Border e -> Array r ix e -> ix -> e -borderIndex border arr = handleBorderIndex border (size arr) (unsafeIndex arr) -{-# INLINE borderIndex #-} - --- | /O(1)/ - Lookup an element in the array. This is a partial function and it can throw --- `IndexOutOfBoundsException` inside pure code. It is safer to use `index` instead. --- --- ==== __Examples__ --- --- >>> import Data.Massiv.Array --- >>> :set -XOverloadedLists --- >>> xs = [0..100] :: Array U Ix1 Int --- >>> index' xs 50 --- 50 --- >>> index' xs 150 --- *** Exception: IndexOutOfBoundsException: 150 is not safe for (Sz1 101) --- --- @since 0.1.0 -index' :: Manifest r ix e => Array r ix e -> ix -> e -index' = evaluate' -{-# INLINE index' #-} - --- | This is just like `indexM` function, but it allows getting values from --- delayed arrays as well as `Manifest`. As the name suggests, indexing into a --- delayed array at the same index multiple times will cause evaluation of the --- value each time and can destroy the performace if used without care. --- --- ==== __Examples__ --- --- >>> import Control.Exception --- >>> import Data.Massiv.Array --- >>> evaluateM (range Seq (Ix2 10 20) (100 :. 210)) 50 :: Either SomeException Ix2 --- Right (60 :. 70) --- >>> evaluateM (range Seq (Ix2 10 20) (100 :. 210)) 150 :: Either SomeException Ix2 --- Left (IndexOutOfBoundsException: (150 :. 150) is not safe for (Sz (90 :. 190))) --- --- @since 0.3.0 -evaluateM :: (Source r ix e, MonadThrow m) => Array r ix e -> ix -> m e -evaluateM arr ix = - handleBorderIndex - (Fill (throwM (IndexOutOfBoundsException (size arr) ix))) - (size arr) - (pure . unsafeIndex arr) - ix -{-# INLINE evaluateM #-} - --- | Similar to `evaluateM`, but will throw an exception in pure code. --- --- ==== __Examples__ --- --- >>> import Data.Massiv.Array --- >>> evaluate' (range Seq (Ix2 10 20) (100 :. 210)) 50 --- 60 :. 70 --- >>> evaluate' (range Seq (Ix2 10 20) (100 :. 210)) 150 --- *** Exception: IndexOutOfBoundsException: (150 :. 150) is not safe for (Sz (90 :. 190)) --- --- @since 0.3.0 -evaluate' :: Source r ix e => Array r ix e -> ix -> e -evaluate' arr ix = - handleBorderIndex - (Fill (throw (IndexOutOfBoundsException (size arr) ix))) - (size arr) - (unsafeIndex arr) - ix -{-# INLINE evaluate' #-} - - --- | Map a monadic index aware function over an array sequentially, while discarding the result. --- --- ==== __Examples__ --- --- >>> import Data.Massiv.Array --- >>> imapM_ (curry print) $ range Seq (Ix1 10) 15 --- (0,10) --- (1,11) --- (2,12) --- (3,13) --- (4,14) --- --- @since 0.1.0 -imapM_ :: (Source r ix a, Monad m) => (ix -> a -> m b) -> Array r ix a -> m () -imapM_ f !arr = - iterM_ zeroIndex (unSz (size arr)) (pureIndex 1) (<) $ \ !ix -> f ix (unsafeIndex arr ix) -{-# INLINE imapM_ #-} - - --- | /O(1)/ - Get the number of elements in the array --- --- ==== __Examples__ --- --- >>> import Data.Massiv.Array --- >>> elemsCount $ range Seq (Ix1 10) 15 --- 5 --- --- @since 0.1.0 -elemsCount :: Load r ix e => Array r ix e -> Int -elemsCount = totalElem . size -{-# INLINE elemsCount #-} - --- | /O(1)/ - Check if array has no elements. --- --- ==== __Examples__ --- --- >>> import Data.Massiv.Array --- >>> isEmpty $ range Seq (Ix2 10 20) (11 :. 21) --- False --- >>> isEmpty $ range Seq (Ix2 10 20) (10 :. 21) --- True --- --- @since 0.1.0 -isEmpty :: Load r ix e => Array r ix e -> Bool -isEmpty !arr = 0 == elemsCount arr -{-# INLINE isEmpty #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE DefaultSignatures #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | +-- Module : Data.Massiv.Core.Common +-- Copyright : (c) Alexey Kuleshevich 2018-2022 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +module Data.Massiv.Core.Common ( + Array, + Vector, + Matrix, + MArray, + MVector, + MMatrix, + Steps (..), + Stream (..), + Strategy (..), + Source (..), + PrefIndex (..), + Load (..), + StrideLoad (..), + Size (..), + Shape (..), + Manifest (..), + Mutable, + Comp (..), + Scheduler, + numWorkers, + scheduleWork, + scheduleWork_, + WorkerStates, + unsafeRead, + unsafeWrite, + unsafeModify, + unsafeLinearModify, + unsafeSwap, + unsafeLinearSwap, + unsafeDefaultLinearShrink, + Ragged (..), + empty, + singleton, + + -- * Size + elemsCount, + isNotNull, + isEmpty, + isNotEmpty, + Sz (SafeSz), + LengthHint (..), + + -- * Indexing + (!?), + index, + indexM, + (!), + index', + (??), + defaultIndex, + borderIndex, + evaluateM, + evaluate', + inline0, + inline1, + inline2, + module Data.Massiv.Core.Index, + + -- * Common Operations + Semigroup ((<>)), + + -- * Exceptions + MonadThrow (..), + IndexException (..), + SizeException (..), + ShapeException (..), + module Data.Massiv.Core.Exception, + Proxy (..), + Id (..), + + -- * Stateful Monads + runST, + ST, + MonadUnliftIO (..), + MonadIO (liftIO), + PrimMonad (PrimState), + RealWorld, +) where + +#if !MIN_VERSION_base(4,11,0) +import Data.Semigroup (Semigroup((<>))) +#endif +import Control.Monad.Catch (MonadThrow (..)) +import Control.Monad.IO.Unlift (MonadIO (liftIO), MonadUnliftIO (..)) +import Control.Monad.Primitive +import Control.Monad.ST +import Control.Scheduler ( + Comp (..), + Scheduler, + WorkerStates, + numWorkers, + scheduleWork, + scheduleWork_, + trivialScheduler_, + ) +import Data.Kind +import Data.Massiv.Core.Exception +import Data.Massiv.Core.Index +import Data.Massiv.Core.Index.Internal (Sz (SafeSz)) +import qualified Data.Stream.Monadic as S (Stream) +import Data.Typeable +import Data.Vector.Fusion.Util +import GHC.Exts (IsList) + +#include "massiv.h" + +-- | The array family. Representations @r@ describe how data is arranged or computed. All +-- arrays have a common property that each index @ix@ always maps to the same unique +-- element @e@, even if that element does not yet exist in memory and the array has to be +-- computed in order to get the value of that element. Data is always arranged in a nested +-- row-major fashion. Rank of an array is specified by @`Dimensions` ix@. +-- +-- @since 0.1.0 +data family Array r ix e :: Type + +-- | Type synonym for a single dimension array, or simply a flat vector. +-- +-- @since 0.5.0 +type Vector r e = Array r Ix1 e + +-- | Type synonym for a two-dimentsional array, or simply a matrix. +-- +-- @since 0.5.0 +type Matrix r e = Array r Ix2 e + +-- | Mutable version of a `Manifest` `Array`. The extra type argument @s@ is for +-- the state token used by `IO` and `ST`. +-- +-- @since 0.1.0 +data family MArray s r ix e :: Type + +-- | Type synonym for a single dimension mutable array, or simply a flat mutable vector. +-- +-- @since 0.5.0 +type MVector s r e = MArray s r Ix1 e + +-- | Type synonym for a two-dimentsional mutable array, or simply a mutable matrix. +-- +-- @since 0.5.0 +type MMatrix s r e = MArray s r Ix2 e + +class Load r ix e => Stream r ix e where + toStream :: Array r ix e -> Steps Id e + + toStreamIx :: Array r ix e -> Steps Id (ix, e) + +data Steps m e = Steps + { stepsStream :: S.Stream m e + , stepsSize :: LengthHint + } + +class Typeable r => Strategy r where + -- | Set computation strategy for this array + -- + -- ==== __Example__ + -- + -- >>> :set -XTypeApplications + -- >>> import Data.Massiv.Array + -- >>> a = singleton @DL @Ix1 @Int 0 + -- >>> a + -- Array DL Seq (Sz1 1) + -- [ 0 ] + -- >>> setComp (ParN 6) a -- use 6 capabilities + -- Array DL (ParN 6) (Sz1 1) + -- [ 0 ] + setComp :: Comp -> Array r ix e -> Array r ix e + + -- | Get computation strategy of this array + -- + -- @since 0.1.0 + getComp :: Array r ix e -> Comp + + -- | Array representation. Representation is never evaluated in @massiv@, + -- therefore default implementation is bottom. However, it is recommended to + -- supply a constructor that doesn't result in an error when evaluated. + -- + -- @since 1.0.2 + repr :: r + repr = + error $ + "Array representation should never be evaluated: " + ++ show (typeRep (Proxy :: Proxy r)) + +-- | Size hint +-- +-- @since 1.0.0 +data LengthHint + = -- | Exact known size + LengthExact Sz1 + | -- | Upper bound on the size + LengthMax Sz1 + | -- | Unknown size + LengthUnknown + deriving (Eq, Show) + +-- | The shape of an array. It is different from `Size` in that it can be applicable to +-- non-square matrices and might not be available in constant time. +-- +-- @since 1.0.0 +class Index ix => Shape r ix where + -- | /O(1)/ - Check what do we know about the number of elements without doing any work + -- + -- @since 1.0.0 + linearSizeHint :: Array r ix e -> LengthHint + linearSizeHint = LengthExact . linearSize + {-# INLINE linearSizeHint #-} + + -- | /O(n)/ - possibly iterate over the whole array before producing the answer + -- + -- @since 0.5.8 + linearSize :: Array r ix e -> Sz1 + default linearSize :: Size r => Array r ix e -> Sz1 + linearSize = SafeSz . elemsCount + {-# INLINE linearSize #-} + + -- | /O(n)/ - Rectangular size of an array that is inferred from looking at the first row in + -- each dimensions. For rectangular arrays this is the same as `size` + -- + -- @since 1.0.0 + outerSize :: Array r ix e -> Sz ix + default outerSize :: Size r => Array r ix e -> Sz ix + outerSize = size + {-# INLINE outerSize #-} + + -- | /O(1)/ - Get the possible maximum linear size of an immutabe array. If the lookup + -- of size in constant time is not possible, `Nothing` will be returned. This value + -- will be used as the initial size of the mutable array into which the loading will + -- happen. + -- + -- @since 1.0.0 + maxLinearSize :: Array r ix e -> Maybe Sz1 + maxLinearSize = lengthHintUpperBound . linearSizeHint + {-# INLINE maxLinearSize #-} + + -- | /O(1)/ - Check whether an array is empty or not. + -- + -- ==== __Examples__ + -- + -- >>> import Data.Massiv.Array + -- >>> isNull $ range Seq (Ix2 10 20) (11 :. 21) + -- False + -- >>> isNull $ range Seq (Ix2 10 20) (10 :. 21) + -- True + -- >>> isNull (empty :: Array D Ix5 Int) + -- True + -- >>> isNull $ sfromList [] + -- True + -- + -- @since 1.0.0 + isNull :: Array r ix e -> Bool + isNull = (zeroSz ==) . linearSize + {-# INLINE isNull #-} + +lengthHintUpperBound :: LengthHint -> Maybe Sz1 +lengthHintUpperBound = \case + LengthExact sz -> Just sz + LengthMax sz -> Just sz + LengthUnknown -> Nothing +{-# INLINE lengthHintUpperBound #-} + +-- | Arrays that have information about their size availible in constant +-- time. +class Size r where + -- | /O(1)/ - Get the exact size of an immutabe array. Most of the time will + -- produce the size in constant time, except for `Data.Massiv.Array.DS` + -- representation, which could result in evaluation of the whole stream. See + -- `maxLinearSize` and `Data.Massiv.Vector.slength` for more info. + -- + -- @since 0.1.0 + size :: Array r ix e -> Sz ix + + -- | /O(1)/ - Change the size of an array. Total number of elements should be the same, but it is + -- not validated. + -- + -- @since 0.1.0 + unsafeResize :: (Index ix, Index ix') => Sz ix' -> Array r ix e -> Array r ix' e + +-- | Prefered indexing function. +data PrefIndex ix e + = PrefIndex (ix -> e) + | PrefIndexLinear (Int -> e) + +instance Functor (PrefIndex ix) where + fmap f = \case + PrefIndex ig -> PrefIndex (f . ig) + PrefIndexLinear ig -> PrefIndexLinear (f . ig) + {-# INLINE fmap #-} + (<$) e _ = PrefIndexLinear (const e) + {-# INLINE (<$) #-} + +-- | Arrays that can be used as source to practically any manipulation function. +class (Strategy r, Size r) => Source r e where + {-# MINIMAL (unsafeIndex | unsafeLinearIndex), unsafeLinearSlice #-} + + -- | Lookup element in the array. No bounds check is performed and access of + -- arbitrary memory is possible when invalid index is supplied. + -- + -- @since 0.1.0 + unsafeIndex :: HAS_CALL_STACK => Index ix => Array r ix e -> ix -> e + unsafeIndex !arr = unsafeLinearIndex arr . toLinearIndex (size arr) + {-# INLINE unsafeIndex #-} + + -- | Lookup element in the array using flat index in a row-major fashion. No + -- bounds check is performed + -- + -- @since 0.1.0 + unsafeLinearIndex :: HAS_CALL_STACK => Index ix => Array r ix e -> Int -> e + unsafeLinearIndex !arr = unsafeIndex arr . fromLinearIndex (size arr) + {-# INLINE unsafeLinearIndex #-} + + -- | Alternative indexing function that can choose an index that is most + -- efficient for underlying representation + -- + -- @since 1.0.2 + unsafePrefIndex :: HAS_CALL_STACK => Index ix => Array r ix e -> PrefIndex ix e + unsafePrefIndex !arr = PrefIndexLinear (unsafeLinearIndex arr) + {-# INLINE unsafePrefIndex #-} + + -- | /O(1)/ - Take a slice out of an array from the outside + -- + -- @since 0.1.0 + unsafeOuterSlice + :: HAS_CALL_STACK + => (Index ix, Index (Lower ix)) + => Array r ix e + -> Sz (Lower ix) + -> Int + -> Array r (Lower ix) e + unsafeOuterSlice arr sz i = unsafeResize sz $ unsafeLinearSlice i (toLinearSz sz) arr + {-# INLINE unsafeOuterSlice #-} + + -- | /O(1)/ - Source arrays also give us ability to look at their linear slices in + -- constant time + -- + -- @since 0.5.0 + unsafeLinearSlice :: Index ix => Ix1 -> Sz1 -> Array r ix e -> Array r Ix1 e + +-- | Any array that can be computed and loaded into memory +class (Strategy r, Shape r ix) => Load r ix e where + {-# MINIMAL (makeArray | makeArrayLinear), (iterArrayLinearST_ | iterArrayLinearWithSetST_) #-} + + -- | Construct an Array. Resulting type either has to be unambiguously inferred or restricted + -- manually, like in the example below. Use "Data.Massiv.Array.makeArrayR" if you'd like to + -- specify representation as an argument. + -- + -- >>> import Data.Massiv.Array + -- >>> makeArray Seq (Sz (3 :. 4)) (\ (i :. j) -> if i == j then i else 0) :: Array D Ix2 Int + -- Array D Seq (Sz (3 :. 4)) + -- [ [ 0, 0, 0, 0 ] + -- , [ 0, 1, 0, 0 ] + -- , [ 0, 0, 2, 0 ] + -- ] + -- + -- Instead of restricting the full type manually we can use @TypeApplications@ as convenience: + -- + -- >>> :set -XTypeApplications + -- >>> makeArray @P @_ @Double Seq (Sz2 3 4) $ \(i :. j) -> logBase (fromIntegral i) (fromIntegral j) + -- Array P Seq (Sz (3 :. 4)) + -- [ [ NaN, -0.0, -0.0, -0.0 ] + -- , [ -Infinity, NaN, Infinity, Infinity ] + -- , [ -Infinity, 0.0, 1.0, 1.5849625007211563 ] + -- ] + -- + -- @since 0.1.0 + makeArray + :: Comp + -- ^ Computation strategy. Useful constructors are `Seq` and `Par` + -> Sz ix + -- ^ Size of the result array. + -> (ix -> e) + -- ^ Function to generate elements at a particular index + -> Array r ix e + makeArray comp sz f = makeArrayLinear comp sz (f . fromLinearIndex sz) + {-# INLINE makeArray #-} + + -- | Same as `makeArray`, but produce elements using linear row-major index. + -- + -- >>> import Data.Massiv.Array + -- >>> makeArrayLinear Seq (Sz (2 :. 4)) id :: Array D Ix2 Int + -- Array D Seq (Sz (2 :. 4)) + -- [ [ 0, 1, 2, 3 ] + -- , [ 4, 5, 6, 7 ] + -- ] + -- + -- @since 0.3.0 + makeArrayLinear :: Comp -> Sz ix -> (Int -> e) -> Array r ix e + makeArrayLinear comp sz f = makeArray comp sz (f . toLinearIndex sz) + {-# INLINE makeArrayLinear #-} + + -- | Construct an array of the specified size that contains the same element in all of + -- the cells. + -- + -- @since 0.3.0 + replicate :: Comp -> Sz ix -> e -> Array r ix e + replicate comp sz !e = makeArrayLinear comp sz (const e) + {-# INLINE replicate #-} + + -- | Iterate over an array with a ST action that is applied to each element and its index. + -- + -- @since 1.0.0 + iterArrayLinearST_ + :: Scheduler s () + -> Array r ix e + -- ^ Array that is being loaded + -> (Int -> e -> ST s ()) + -- ^ Function that writes an element into target array + -> ST s () + iterArrayLinearST_ scheduler arr uWrite = + iterArrayLinearWithSetST_ scheduler arr uWrite $ \offset sz e -> + loopA_ offset (< (offset + unSz sz)) (+ 1) (`uWrite` e) + {-# INLINE iterArrayLinearST_ #-} + + -- | Similar to `iterArrayLinearST_`. Except it also accepts a function that is + -- potentially optimized for setting many cells in a region to the same + -- value. + -- + -- @since 1.0.0 + iterArrayLinearWithSetST_ + :: Scheduler s () + -> Array r ix e + -- ^ Array that is being loaded + -> (Ix1 -> e -> ST s ()) + -- ^ Function that writes an element into target array + -> (Ix1 -> Sz1 -> e -> ST s ()) + -- ^ Function that efficiently sets a region of an array + -- to the supplied value target array + -> ST s () + iterArrayLinearWithSetST_ scheduler arr uWrite _ = iterArrayLinearST_ scheduler arr uWrite + {-# INLINE iterArrayLinearWithSetST_ #-} + + -- | Load into a supplied mutable array sequentially. Returned array does not have to be + -- the same. + -- + -- @since 1.0.0 + unsafeLoadIntoST + :: Manifest r' e + => MVector s r' e + -> Array r ix e + -> ST s (MArray s r' ix e) + unsafeLoadIntoST mvec arr = do + let sz = outerSize arr + mvec' <- resizeMVector mvec $ toLinearSz sz + iterArrayLinearWithSetST_ trivialScheduler_ arr (unsafeLinearWrite mvec') (unsafeLinearSet mvec') + pure $ unsafeResizeMArray sz mvec' + {-# INLINE unsafeLoadIntoST #-} + + -- | Same as `unsafeLoadIntoST`, but respecting computation strategy. + -- + -- @since 1.0.0 + unsafeLoadIntoIO + :: Manifest r' e + => MVector RealWorld r' e + -> Array r ix e + -> IO (MArray RealWorld r' ix e) + unsafeLoadIntoIO mvec arr = do + let sz = outerSize arr + mvec' <- resizeMVector mvec $ toLinearSz sz + withMassivScheduler_ (getComp arr) $ \scheduler -> + stToIO $ + iterArrayLinearWithSetST_ scheduler arr (unsafeLinearWrite mvec') (unsafeLinearSet mvec') + pure $ unsafeResizeMArray sz mvec' + {-# INLINE unsafeLoadIntoIO #-} + +resizeMVector + :: (Manifest r e, PrimMonad f) + => MVector (PrimState f) r e + -> Sz1 + -> f (MVector (PrimState f) r e) +resizeMVector mvec k = + let mk = sizeOfMArray mvec + in if k == mk + then pure mvec + else + if k < mk + then unsafeLinearShrink mvec k + else unsafeLinearGrow mvec k +{-# INLINE resizeMVector #-} + +class Load r ix e => StrideLoad r ix e where + -- | Load an array into memory with stride. Default implementation requires an instance of + -- `Source`. + iterArrayLinearWithStrideST_ + :: Scheduler s () + -> Stride ix + -- ^ Stride to use + -> Sz ix + -- ^ Size of the target array affected by the stride. + -> Array r ix e + -- ^ Array that is being loaded + -> (Int -> e -> ST s ()) + -- ^ Function that writes an element into target array + -> ST s () + default iterArrayLinearWithStrideST_ + :: Source r e + => Scheduler s () + -> Stride ix + -> Sz ix + -> Array r ix e + -> (Int -> e -> ST s ()) + -> ST s () + iterArrayLinearWithStrideST_ scheduler stride resultSize arr = + splitLinearlyWith_ scheduler (totalElem resultSize) unsafeLinearIndexWithStride + where + !strideIx = unStride stride + unsafeLinearIndexWithStride = + unsafeIndex arr . liftIndex2 (*) strideIx . fromLinearIndex resultSize + {-# INLINE unsafeLinearIndexWithStride #-} + {-# INLINE iterArrayLinearWithStrideST_ #-} + +-- class (Load r ix e) => StrideLoad r ix e where +-- class (Size r, StrideLoad r ix e) => StrideLoadP r ix e where +-- +-- unsafeLoadIntoWithStrideST :: -- TODO: this would remove Size constraint and allow DS and LN instances for vectors. +-- Manifest r' ix e +-- => Array r ix e +-- -> Stride ix -- ^ Stride to use +-- -> MArray RealWorld r' ix e +-- -> m (MArray RealWorld r' ix e) + +-- | Starting with massiv-1.0 `Mutable` and `Manifest` are synonymous. Since massiv-1.1 +-- it is deprecated and will be removed in massiv-1.2 +type Mutable r e = Manifest r e + +{-# DEPRECATED Mutable "In favor of `Manifest`" #-} + +-- | Manifest arrays are backed by actual memory and values are looked up versus +-- computed as it is with delayed arrays. Because manifest arrays are located in +-- memory their contents can be mutated once thawed into `MArray`. The process +-- of changed a mutable `MArray` back into an immutable `Array` is called +-- freezing. +class Source r e => Manifest r e where + unsafeLinearIndexM :: Index ix => Array r ix e -> Int -> e + + -- | /O(1)/ - Get the size of a mutable array. + -- + -- @since 1.0.0 + sizeOfMArray :: Index ix => MArray s r ix e -> Sz ix + + -- | /O(1)/ - Change the size of a mutable array. The actual number of + -- elements should stay the same. + -- + -- @since 1.0.0 + unsafeResizeMArray :: (Index ix', Index ix) => Sz ix' -> MArray s r ix e -> MArray s r ix' e + + -- | /O(1)/ - Take a linear slice out of a mutable array. + -- + -- @since 1.0.0 + unsafeLinearSliceMArray :: Index ix => Ix1 -> Sz1 -> MArray s r ix e -> MVector s r e + + -- | Convert immutable array into a mutable array without copy. + -- + -- @since 0.1.0 + unsafeThaw :: (Index ix, PrimMonad m) => Array r ix e -> m (MArray (PrimState m) r ix e) + + -- | Convert mutable array into an immutable array without copy. + -- + -- @since 0.1.0 + unsafeFreeze :: (Index ix, PrimMonad m) => Comp -> MArray (PrimState m) r ix e -> m (Array r ix e) + + -- | Create new mutable array, leaving it's elements uninitialized. Size isn't validated either. + -- + -- @since 0.1.0 + unsafeNew :: (Index ix, PrimMonad m) => Sz ix -> m (MArray (PrimState m) r ix e) + + -- | Read an element at linear row-major index + -- + -- @since 0.1.0 + unsafeLinearRead + :: HAS_CALL_STACK + => (Index ix, PrimMonad m) + => MArray (PrimState m) r ix e + -> Int + -> m e + + -- | Write an element into mutable array with linear row-major index + -- + -- @since 0.1.0 + unsafeLinearWrite + :: HAS_CALL_STACK + => (Index ix, PrimMonad m) + => MArray (PrimState m) r ix e + -> Int + -> e + -> m () + + -- | Initialize mutable array to some default value. + -- + -- @since 0.3.0 + initialize :: (Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> m () + + -- | Create new mutable array while initializing all elements to some default value. + -- + -- @since 0.3.0 + initializeNew :: (Index ix, PrimMonad m) => Maybe e -> Sz ix -> m (MArray (PrimState m) r ix e) + initializeNew Nothing sz = unsafeNew sz >>= \ma -> ma <$ initialize ma + initializeNew (Just e) sz = newMArray sz e + {-# INLINE initializeNew #-} + + -- | Create new mutable array while initializing all elements to the specified value. + -- + -- @since 0.6.0 + newMArray :: (Index ix, PrimMonad m) => Sz ix -> e -> m (MArray (PrimState m) r ix e) + newMArray sz e = do + marr <- unsafeNew sz + marr <$ unsafeLinearSet marr 0 (SafeSz (totalElem sz)) e + {-# INLINE newMArray #-} + + -- | Set all cells in the mutable array within the range to a specified value. + -- + -- @since 0.3.0 + unsafeLinearSet + :: HAS_CALL_STACK + => (Index ix, PrimMonad m) + => MArray (PrimState m) r ix e + -> Ix1 + -> Sz1 + -> e + -> m () + unsafeLinearSet marr offset len e = + loopA_ offset (< (offset + unSz len)) (+ 1) (\i -> unsafeLinearWrite marr i e) + {-# INLINE unsafeLinearSet #-} + + -- | Copy part of one mutable array into another + -- + -- @since 0.3.6 + unsafeLinearCopy + :: HAS_CALL_STACK + => (Index ix', Index ix, 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 + loopA_ 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 + :: HAS_CALL_STACK + => (Index ix', Index ix, 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 + loopA_ 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 or + -- equal. There is no guarantee that the original array is left unchanged, so it should + -- no longer be used. + -- + -- @since 0.3.6 + unsafeLinearShrink + :: HAS_CALL_STACK + => (Index ix, 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 + -- or equal. There is no guarantee that the original array is left unchanged, so it + -- should no longer be used. + -- + -- @since 0.3.6 + unsafeLinearGrow + :: HAS_CALL_STACK + => (Index ix, 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 (sizeOfMArray marr)) + pure marr' + {-# INLINE unsafeLinearGrow #-} + +unsafeDefaultLinearShrink + :: HAS_CALL_STACK + => (Manifest r e, Index ix, 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 +-- +-- @since 0.1.0 +unsafeRead + :: HAS_CALL_STACK + => (Manifest r e, Index ix, PrimMonad m) + => MArray (PrimState m) r ix e + -> ix + -> m e +unsafeRead marr = unsafeLinearRead marr . toLinearIndex (sizeOfMArray marr) +{-# INLINE unsafeRead #-} + +-- | Write an element into array +-- +-- @since 0.1.0 +unsafeWrite + :: HAS_CALL_STACK + => (Manifest r e, Index ix, PrimMonad m) + => MArray (PrimState m) r ix e + -> ix + -> e + -> m () +unsafeWrite marr = unsafeLinearWrite marr . toLinearIndex (sizeOfMArray marr) +{-# INLINE unsafeWrite #-} + +-- | Modify an element in the array with a monadic action. Returns the previous value. +-- +-- @since 0.4.0 +unsafeLinearModify + :: HAS_CALL_STACK + => (Manifest r e, Index ix, PrimMonad m) + => MArray (PrimState m) r ix e + -> (e -> m e) + -> Int + -> m e +unsafeLinearModify !marr f !i = do + v <- unsafeLinearRead marr i + v' <- f v + unsafeLinearWrite marr i v' + pure v +{-# INLINE unsafeLinearModify #-} + +-- | Modify an element in the array with a monadic action. Returns the previous value. +-- +-- @since 0.4.0 +unsafeModify + :: HAS_CALL_STACK + => (Manifest r e, Index ix, PrimMonad m) + => MArray (PrimState m) r ix e + -> (e -> m e) + -> ix + -> m e +unsafeModify marr f ix = unsafeLinearModify marr f (toLinearIndex (sizeOfMArray marr) ix) +{-# INLINE unsafeModify #-} + +-- | Swap two elements in a mutable array under the supplied indices. Returns the previous +-- values. +-- +-- @since 0.4.0 +unsafeSwap + :: HAS_CALL_STACK + => (Manifest r e, Index ix, PrimMonad m) + => MArray (PrimState m) r ix e + -> ix + -> ix + -> m (e, e) +unsafeSwap !marr !ix1 !ix2 = unsafeLinearSwap marr (toLinearIndex sz ix1) (toLinearIndex sz ix2) + where + sz = sizeOfMArray marr +{-# INLINE unsafeSwap #-} + +-- | Swap two elements in a mutable array under the supplied linear indices. Returns the +-- previous values. +-- +-- @since 0.4.0 +unsafeLinearSwap + :: HAS_CALL_STACK + => (Manifest r e, Index ix, PrimMonad m) + => MArray (PrimState m) r ix e + -> Int + -> Int + -> m (e, e) +unsafeLinearSwap !marr !i1 !i2 = do + val1 <- unsafeLinearRead marr i1 + val2 <- unsafeLinearRead marr i2 + unsafeLinearWrite marr i1 val2 + unsafeLinearWrite marr i2 val1 + return (val1, val2) +{-# INLINE unsafeLinearSwap #-} + +class (IsList (Array r ix e), Load r ix e) => Ragged r ix e where + generateRaggedM :: Monad m => Comp -> Sz ix -> (ix -> m e) -> m (Array r ix e) + + flattenRagged :: Array r ix e -> Vector r e + + loadRaggedST + :: Scheduler s () -> Array r ix e -> (Ix1 -> e -> ST s ()) -> Ix1 -> Ix1 -> Sz ix -> ST s () + + raggedFormat :: (e -> String) -> String -> Array r ix e -> String + +-- | Create an Array with no elements. By itself it is not particularly useful, but it serves as a +-- nice base for constructing larger arrays. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> :set -XTypeApplications +-- >>> xs = empty @DL @Ix1 @Double +-- >>> snoc (cons 4 (cons 5 xs)) 22 +-- Array DL Seq (Sz1 3) +-- [ 4.0, 5.0, 22.0 ] +-- +-- @since 0.3.0 +empty + :: forall r ix e + . Load r ix e + => Array r ix e +empty = makeArray Seq zeroSz (const (throwImpossible Uninitialized)) +{-# INLINE empty #-} + +-- | Create an Array with a single element. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> singleton 7 :: Array D Ix4 Double +-- Array D Seq (Sz (1 :> 1 :> 1 :. 1)) +-- [ [ [ [ 7.0 ] +-- ] +-- ] +-- ] +-- +-- Instead of specifying type signature we could use @TypeApplications@ +-- +-- >>> :set -XTypeApplications +-- >>> singleton @U @Ix4 @Double 7 +-- Array U Seq (Sz (1 :> 1 :> 1 :. 1)) +-- [ [ [ [ 7.0 ] +-- ] +-- ] +-- ] +-- +-- @since 0.1.0 +singleton + :: forall r ix e + . Load r ix e + => e + -- ^ The only element + -> Array r ix e +singleton = makeArray Seq oneSz . const +{-# INLINE singleton #-} + +infixl 4 !, !?, ?? + +-- | /O(1)/ - Infix version of 'index''. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> a = computeAs U $ iterateN (Sz (2 :. 3)) succ (0 :: Int) +-- >>> a +-- Array U Seq (Sz (2 :. 3)) +-- [ [ 1, 2, 3 ] +-- , [ 4, 5, 6 ] +-- ] +-- >>> a ! 0 :. 2 +-- 3 +-- +-- @since 0.1.0 +(!) + :: forall r ix e + . (HasCallStack, Manifest r e, Index ix) + => Array r ix e + -> ix + -> e +(!) arr = throwEither . evaluateM arr +{-# INLINE (!) #-} + +-- | /O(1)/ - Infix version of `indexM`. +-- +-- /__Exceptions__/: `IndexOutOfBoundsException` +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> :set -XTypeApplications +-- >>> a <- fromListsM @U @Ix2 @Int Seq [[1,2,3],[4,5,6]] +-- >>> a +-- Array U Seq (Sz (2 :. 3)) +-- [ [ 1, 2, 3 ] +-- , [ 4, 5, 6 ] +-- ] +-- >>> a !? 0 :. 2 +-- 3 +-- >>> a !? 0 :. 3 +-- *** Exception: IndexOutOfBoundsException: (0 :. 3) is not safe for (Sz (2 :. 3)) +-- >>> a !? 0 :. 3 :: Maybe Int +-- Nothing +-- +-- @since 0.1.0 +(!?) + :: forall r ix e m + . (Index ix, Manifest r e, MonadThrow m) + => Array r ix e + -> ix + -> m e +(!?) = indexM +{-# INLINE (!?) #-} + +-- | /O(1)/ - Lookup an element in the array, where array itself is wrapped with +-- `MonadThrow`. This operator is useful when used together with slicing or other +-- functions that can fail. +-- +-- /__Exceptions__/: `IndexOutOfBoundsException` +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> :set -XTypeApplications +-- >>> ma = fromListsM @U @Ix3 @Int @Maybe Seq [[[1,2,3]],[[4,5,6]]] +-- >>> ma +-- Just (Array U Seq (Sz (2 :> 1 :. 3)) +-- [ [ [ 1, 2, 3 ] +-- ] +-- , [ [ 4, 5, 6 ] +-- ] +-- ] +-- ) +-- >>> ma ??> 1 +-- Just (Array U Seq (Sz (1 :. 3)) +-- [ [ 4, 5, 6 ] +-- ] +-- ) +-- >>> ma ??> 1 ?? 0 :. 2 +-- Just 6 +-- >>> ma ?? 1 :> 0 :. 2 +-- Just 6 +-- +-- @since 0.1.0 +(??) :: (Index ix, Manifest r e, MonadThrow m) => m (Array r ix e) -> ix -> m e +(??) marr ix = marr >>= (!? ix) +{-# INLINE (??) #-} + +-- | /O(1)/ - Lookup an element in the array. Returns `Nothing`, when index is out of bounds and +-- returns the element at the supplied index otherwise. Use `indexM` instead, since it is more +-- general and it can just as well be used with `Maybe`. +-- +-- @since 0.1.0 +index :: (Index ix, Manifest r e) => Array r ix e -> ix -> Maybe e +index = indexM +{-# INLINE index #-} + +-- | /O(1)/ - Lookup an element in the array. +-- +-- /__Exceptions__/: `IndexOutOfBoundsException` +-- +-- @since 0.3.0 +indexM :: (Index ix, Manifest r e, MonadThrow m) => Array r ix e -> ix -> m e +indexM = evaluateM +{-# INLINE indexM #-} + +-- | /O(1)/ - Lookup an element in the array, while using default element when index is out of +-- bounds. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array +-- >>> :set -XOverloadedLists +-- >>> xs = [0..100] :: Array P Ix1 Int +-- >>> defaultIndex 999 xs 100 +-- 100 +-- >>> defaultIndex 999 xs 101 +-- 999 +-- +-- @since 0.1.0 +defaultIndex :: (Index ix, Manifest r e) => e -> Array r ix e -> ix -> e +defaultIndex defVal = borderIndex (Fill defVal) +{-# INLINE defaultIndex #-} + +-- | /O(1)/ - Lookup an element in the array. Use a border resolution technique +-- when index is out of bounds. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> :set -XOverloadedLists +-- >>> xs = [0..100] :: Array U Ix1 Int +-- >>> borderIndex Wrap xs <$> range Seq 99 104 +-- Array D Seq (Sz1 5) +-- [ 99, 100, 0, 1, 2 ] +-- +-- @since 0.1.0 +borderIndex :: (Index ix, Manifest r e) => Border e -> Array r ix e -> ix -> e +borderIndex border arr = handleBorderIndex border (size arr) (unsafeIndex arr) +{-# INLINE borderIndex #-} + +-- | /O(1)/ - Lookup an element in the array. This is a partial function and it will throw +-- an error when index is out of bounds. It is safer to use `indexM` instead. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array +-- >>> :set -XOverloadedLists +-- >>> xs = [0..100] :: Array U Ix1 Int +-- >>> index' xs 50 +-- 50 +-- +-- @since 0.1.0 +index' :: (HasCallStack, Index ix, Manifest r e) => Array r ix e -> ix -> e +index' arr ix = throwEither (evaluateM arr ix) +{-# INLINE index' #-} + +-- | This is just like `indexM` function, but it allows getting values from +-- delayed arrays as well as `Manifest`. As the name suggests, indexing into a +-- delayed array at the same index multiple times will cause evaluation of the +-- value each time and can destroy the performace if used without care. +-- +-- ==== __Examples__ +-- +-- >>> import Control.Exception +-- >>> import Data.Massiv.Array +-- >>> evaluateM (range Seq (Ix2 10 20) (100 :. 210)) 50 :: Either SomeException Ix2 +-- Right (60 :. 70) +-- >>> evaluateM (range Seq (Ix2 10 20) (100 :. 210)) 150 :: Either SomeException Ix2 +-- Left (IndexOutOfBoundsException: (150 :. 150) is not safe for (Sz (90 :. 190))) +-- +-- @since 0.3.0 +evaluateM :: (Index ix, Source r e, MonadThrow m) => Array r ix e -> ix -> m e +evaluateM arr ix + | isSafeIndex (size arr) ix = pure (unsafeIndex arr ix) + | otherwise = throwM (IndexOutOfBoundsException (size arr) ix) +{-# INLINE evaluateM #-} + +-- | Similar to `evaluateM`, but will throw an error on out of bounds indices. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array +-- >>> evaluate' (range Seq (Ix2 10 20) (100 :. 210)) 50 +-- 60 :. 70 +-- +-- @since 0.3.0 +evaluate' :: (HasCallStack, Index ix, Source r e) => Array r ix e -> ix -> e +evaluate' arr ix = throwEither (evaluateM arr ix) +{-# INLINE evaluate' #-} + +-- | /O(1)/ - Check if array has elements. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array +-- >>> isNotNull (singleton 1 :: Array D Ix2 Int) +-- True +-- >>> isNotNull (empty :: Array D Ix2 Int) +-- False +-- +-- @since 0.5.1 +isNotNull :: Shape r ix => Array r ix e -> Bool +isNotNull = not . isNull +{-# INLINE isNotNull #-} + +-- | /O(1)/ - Check if array has elements. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array +-- >>> isEmpty (singleton 1 :: Array D Ix2 Int) +-- False +-- >>> isEmpty (empty :: Array D Ix2 Int) +-- True +-- +-- @since 1.0.0 +isEmpty :: (Index ix, Size r) => Array r ix e -> Bool +isEmpty = (== 0) . elemsCount +{-# INLINE isEmpty #-} + +-- | /O(1)/ - Check if array has elements. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array +-- >>> isNotEmpty (singleton 1 :: Array D Ix2 Int) +-- True +-- >>> isNotEmpty (empty :: Array D Ix2 Int) +-- False +-- +-- @since 1.0.0 +isNotEmpty :: (Index ix, Size r) => Array r ix e -> Bool +isNotEmpty = not . isEmpty +{-# INLINE isNotEmpty #-} + +-- | /O(1)/ - Get the number of elements in the array. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array +-- >>> elemsCount $ range Seq (Ix1 10) 15 +-- 5 +-- +-- @since 0.1.0 +elemsCount :: (Index ix, Size r) => Array r ix e -> Int +elemsCount = totalElem . size +{-# INLINE elemsCount #-} + +inline0 :: (a -> b) -> a -> b +inline0 f = f +{-# INLINE [0] inline0 #-} + +inline1 :: (a -> b) -> a -> b +inline1 f = f +{-# INLINE [1] inline1 #-} + +inline2 :: (a -> b) -> a -> b +inline2 f = f +{-# INLINE [2] inline2 #-} diff --git a/src/Data/Massiv/Core/Exception.hs b/src/Data/Massiv/Core/Exception.hs --- a/src/Data/Massiv/Core/Exception.hs +++ b/src/Data/Massiv/Core/Exception.hs @@ -1,17 +1,32 @@ {-# LANGUAGE CPP #-} {-# LANGUAGE GADTs #-} +{-# LANGUAGE ImplicitParams #-} +{-# LANGUAGE LambdaCase #-} {-# OPTIONS_GHC -fno-warn-orphans #-} -module Data.Massiv.Core.Exception - ( ImpossibleException(..) - , throwImpossible - , Uninitialized(..) - , guardNumberOfElements - ) where +-- | +-- Module : Data.Massiv.Core.Exception +-- Copyright : (c) Alexey Kuleshevich 2019-2022 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +module Data.Massiv.Core.Exception ( + throwImpossible, + throwEither, + Uninitialized (..), + guardNumberOfElements, + Exception (..), + SomeException, + HasCallStack, +) where + import Control.Exception import Control.Monad import Control.Monad.Catch import Data.Massiv.Core.Index.Internal +import GHC.Exception +import GHC.Stack #if !MIN_VERSION_exceptions(0, 10, 3) import Control.Monad.ST (ST) @@ -22,30 +37,38 @@ throwM = unsafeIOToST . throwIO #endif - -newtype ImpossibleException = - ImpossibleException SomeException - deriving (Show) - -throwImpossible :: Exception e => e -> a -throwImpossible = throw . ImpossibleException . toException +-- | Throw an impossible error. +-- +-- @since 0.5.6 +throwImpossible :: HasCallStack => Exception e => e -> a +throwImpossible exc = throw (errorCallWithCallStackException msg ?callStack) + where + msg = + " ImpossibleException (" + ++ displayException exc + ++ "): Either one of the unsafe functions was used or it is a bug in the library. " + ++ "In latter case please report this error." {-# NOINLINE throwImpossible #-} -instance Exception ImpossibleException where - displayException (ImpossibleException exc) = - " ImpossibleException (" ++ - displayException exc ++ - "): Either one of the unsafe functions was used or it is a bug in the library. " ++ - "In latter case please report this error." +-- | Throw an error on `Left` or produce the result on `Right`. Exception type is lost, so +-- do not expect to be able to catch it as such. Stick to `IO` if you need exception control +-- flow. +-- +-- @since 0.5.6 +throwEither :: HasCallStack => Either SomeException a -> a +throwEither = + \case + Left exc -> throw (errorCallWithCallStackException (displayException exc) ?callStack) + Right res -> res +{-# INLINE throwEither #-} -- | An error that gets thrown when an unitialized element of a boxed array gets accessed. Can only --- happen when array was constructed with `unsafeNew`. -data Uninitialized = Uninitialized deriving Show +-- happen when array was constructed with `Data.Massiv.Array.Unsafe.unsafeNew`. +data Uninitialized = Uninitialized deriving (Show) instance Exception Uninitialized where displayException Uninitialized = "Array element is uninitialized" - -- | Throw `SizeElementsMismatchException` whenever number of elements in both sizes do -- not match. -- @@ -54,4 +77,3 @@ guardNumberOfElements sz sz' = unless (totalElem sz == totalElem sz') $ throwM $ SizeElementsMismatchException sz sz' {-# INLINE guardNumberOfElements #-} - diff --git a/src/Data/Massiv/Core/Index.hs b/src/Data/Massiv/Core/Index.hs --- a/src/Data/Massiv/Core/Index.hs +++ b/src/Data/Massiv/Core/Index.hs @@ -1,161 +1,214 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE GADTs #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE LambdaCase #-} {-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE TypeOperators #-} + -- | -- Module : Data.Massiv.Core.Index --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Core.Index - ( Ix0(..) - , type Ix1 - , pattern Ix1 - , type Ix2(Ix2, (:.)) - , IxN((:>), Ix3, Ix4, Ix5) - , type Ix3 - , type Ix4 - , type Ix5 - , Ix +module Data.Massiv.Core.Index ( + Ix0 (..), + type Ix1, + pattern Ix1, + type Ix2 (Ix2, (:.)), + IxN ((:>), Ix3, Ix4, Ix5), + HighIxN, + type Ix3, + type Ix4, + type Ix5, + Ix, + -- ** Size - , type Sz1 - , type Sz2 - , type Sz3 - , type Sz4 - , type Sz5 - , Sz(Sz, Sz1, Sz2, Sz3, Sz4, Sz5) - , unSz - , zeroSz - , oneSz - , liftSz - , liftSz2 - , consSz - , unconsSz - , snocSz - , unsnocSz - , setSzM - , insertSzM - , pullOutSzM + type Sz1, + type Sz2, + type Sz3, + type Sz4, + type Sz5, + Sz (Sz, Sz1, Sz2, Sz3, Sz4, Sz5), + unSz, + zeroSz, + oneSz, + liftSz, + liftSz2, + consSz, + unconsSz, + snocSz, + unsnocSz, + setSzM, + insertSzM, + pullOutSzM, + toLinearSz, + mkSzM, + -- ** Dimension - , Dim(..) - , Dimension(Dim1, Dim2, Dim3, Dim4, Dim5, DimN) - , IsIndexDimension - , IsDimValid - , ReportInvalidDim + Dim (..), + Dimension (Dim1, Dim2, Dim3, Dim4, Dim5, DimN), + IsIndexDimension, + IsDimValid, + ReportInvalidDim, + -- ** Stride - , Stride(Stride) - , unStride - , toLinearIndexStride - , strideStart - , strideSize - , oneStride + Stride (Stride), + unStride, + toLinearIndexStride, + strideStart, + strideSize, + oneStride, + -- ** Border - , Border(..) - , handleBorderIndex + Border (..), + handleBorderIndex, + -- ** Index functions - , Lower - , Index(..) - , zeroIndex - , oneIndex - , isNonEmpty - , headDim - , tailDim - , lastDim - , initDim - , getDim' - , setDim' - , modifyDim' - , dropDimM - , dropDim' - , pullOutDim' - , insertDim' - , fromDimension - , getDimension - , setDimension - , modifyDimension - , dropDimension - , pullOutDimension - , insertDimension + Lower, + Index (..), + zeroIndex, + oneIndex, + isZeroSz, + isNotZeroSz, + headDim, + tailDim, + lastDim, + initDim, + getDim', + setDim', + modifyDim', + dropDimM, + dropDim', + pullOutDim', + insertDim', + fromDimension, + getDimension, + setDimension, + modifyDimension, + dropDimension, + pullOutDimension, + insertDimension, + -- * Iterators - , iter - , iterLinearM - , iterLinearM_ - , module Data.Massiv.Core.Iterator - , module Data.Massiv.Core.Index.Tuple + iter, + iterA_, + iterM_, + iterLinearM, + iterLinearM_, + module Data.Massiv.Core.Loop, + module Data.Massiv.Core.Index.Iterator, + module Data.Massiv.Core.Index.Tuple, + -- * Exceptions - , IndexException(..) - , SizeException(..) - , ShapeException(..) - , guardNumberOfElements - , indexWith - ) where + IndexException (..), + SizeException (..), + ShapeException (..), + guardNumberOfElements, + indexAssert, + indexWith, +) where import Control.DeepSeq -import Control.Exception (throw) -import Control.Monad.Catch (MonadThrow(..)) +import Control.Monad.Catch (MonadThrow (..)) +import Data.Coerce import Data.Functor.Identity (runIdentity) -import Data.Massiv.Core.Exception (guardNumberOfElements) +import Data.Massiv.Core.Exception import Data.Massiv.Core.Index.Internal +import Data.Massiv.Core.Index.Iterator import Data.Massiv.Core.Index.Ix import Data.Massiv.Core.Index.Stride import Data.Massiv.Core.Index.Tuple -import Data.Massiv.Core.Iterator +import Data.Massiv.Core.Loop +import GHC.Base (modInt) import GHC.TypeLits +#include "massiv.h" + +-- | 1-dimensional type synonym for size. +-- +-- @since 0.3.0 +type Sz1 = Sz Ix1 + +-- | 2-dimensional size type synonym. +-- +-- @since 0.3.0 +type Sz2 = Sz Ix2 + +-- | 3-dimensional size type synonym. +-- +-- @since 0.3.0 +type Sz3 = Sz Ix3 + +-- | 4-dimensional size type synonym. +-- +-- @since 0.3.0 +type Sz4 = Sz Ix4 + +-- | 5-dimensional size type synonym. +-- +-- @since 0.3.0 +type Sz5 = Sz Ix5 + -- | Approach to be used near the borders during various transformations. -- Whenever a function needs information not only about an element of interest, but -- also about it's neighbors, it will go out of bounds near the array edges, -- hence is this set of approaches that specify how to handle such situation. -data Border e = - Fill e -- ^ Fill in a constant element. - -- - -- @ - -- outside | Array | outside - -- ('Fill' 0) : 0 0 0 0 | 1 2 3 4 | 0 0 0 0 - -- @ - -- - | Wrap -- ^ Wrap around from the opposite border of the array. - -- - -- @ - -- outside | Array | outside - -- 'Wrap' : 1 2 3 4 | 1 2 3 4 | 1 2 3 4 - -- @ - -- - | Edge -- ^ Replicate the element at the edge. - -- - -- @ - -- outside | Array | outside - -- 'Edge' : 1 1 1 1 | 1 2 3 4 | 4 4 4 4 - -- @ - -- - | Reflect -- ^ Mirror like reflection. - -- - -- @ - -- outside | Array | outside - -- 'Reflect' : 4 3 2 1 | 1 2 3 4 | 4 3 2 1 - -- @ - -- - | Continue -- ^ Also mirror like reflection, but without repeating the edge element. - -- - -- @ - -- outside | Array | outside - -- 'Continue' : 1 4 3 2 | 1 2 3 4 | 3 2 1 4 - -- @ - -- +data Border e + = -- | Fill in a constant element. + -- + -- @ + -- outside | Array | outside + -- ('Fill' 0) : 0 0 0 0 | 1 2 3 4 | 0 0 0 0 + -- @ + Fill e + | -- | Wrap around from the opposite border of the array. + -- + -- @ + -- outside | Array | outside + -- 'Wrap' : 1 2 3 4 | 1 2 3 4 | 1 2 3 4 + -- @ + Wrap + | -- | Replicate the element at the edge. + -- + -- @ + -- outside | Array | outside + -- 'Edge' : 1 1 1 1 | 1 2 3 4 | 4 4 4 4 + -- @ + Edge + | -- | Mirror like reflection. + -- + -- @ + -- outside | Array | outside + -- 'Reflect' : 4 3 2 1 | 1 2 3 4 | 4 3 2 1 + -- @ + Reflect + | -- | Also mirror like reflection, but without repeating the edge element. + -- + -- @ + -- outside | Array | outside + -- 'Continue' : 1 4 3 2 | 1 2 3 4 | 3 2 1 4 + -- @ + Continue deriving (Eq, Show) instance NFData e => NFData (Border e) where - rnf b = case b of - Fill e -> rnf e - Wrap -> () - Edge -> () - Reflect -> () - Continue -> () + rnf = \case + Fill e -> rnf e + Wrap -> () + Edge -> () + Reflect -> () + Continue -> () +instance Functor Border where + fmap f = \case + Fill e -> Fill (f e) + Wrap -> Wrap + Edge -> Edge + Reflect -> Reflect + Continue -> Continue -- | Apply a border resolution technique to an index -- @@ -169,25 +222,34 @@ -- 1 :. 2 -- -- @since 0.1.0 -handleBorderIndex :: - Index ix - => Border e -- ^ Broder resolution technique - -> Sz ix -- ^ Size - -> (ix -> e) -- ^ Index function that produces an element - -> ix -- ^ Index +handleBorderIndex + :: Index ix + => Border e + -- ^ Broder resolution technique + -> Sz ix + -- ^ Size + -> (ix -> e) + -- ^ Index function that produces an element + -> ix + -- ^ Index -> e handleBorderIndex border !sz getVal !ix = case border of - Fill val -> if isSafeIndex sz ix then getVal ix else val - Wrap -> getVal (repairIndex sz ix wrap wrap) - Edge -> getVal (repairIndex sz ix (const (const 0)) (\ (SafeSz k) _ -> k - 1)) - Reflect -> getVal (repairIndex sz ix (\ (SafeSz k) !i -> (abs i - 1) `mod` k) - (\ (SafeSz k) !i -> (-i - 1) `mod` k)) - Continue -> getVal (repairIndex sz ix (\ (SafeSz k) !i -> abs i `mod` k) - (\ (SafeSz k) !i -> (-i - 2) `mod` k)) - - where wrap (SafeSz k) i = i `mod` k - {-# INLINE [1] wrap #-} + Fill val + | isSafeIndex sz ix -> getVal ix + | otherwise -> val + Wrap -> + getVal $ + repairIndex sz ix (\(SafeSz k) i -> i `modInt` k) (\(SafeSz k) i -> i `modInt` k) + Edge -> + getVal $ + repairIndex sz ix (const (const 0)) (\(SafeSz k) _ -> k - 1) + Reflect -> + getVal $ + repairIndex sz ix (\(SafeSz k) i -> (-i - 1) `modInt` k) (\(SafeSz k) i -> (-i - 1) `modInt` k) + Continue -> + getVal $ + repairIndex sz ix (\(SafeSz k) i -> negate i `modInt` k) (\(SafeSz k) i -> (-i - 2) `modInt` k) {-# INLINE [1] handleBorderIndex #-} -- | Index with all zeros @@ -209,22 +271,42 @@ oneIndex = pureIndex 1 {-# INLINE [1] oneIndex #-} --- | Checks whether array with this size can hold at least one element. +-- | Checks whether size can hold at least one element. -- -- ==== __Examples__ -- --- >>> isNonEmpty (Sz3 1 0 2) +-- >>> isNotZeroSz (Sz3 1 0 2) -- False -- --- @since 0.1.0 -isNonEmpty :: Index ix => Sz ix -> Bool -isNonEmpty !sz = isSafeIndex sz zeroIndex -{-# INLINE [1] isNonEmpty #-} +-- @since 1.0.0 +isNotZeroSz :: Index ix => Sz ix -> Bool +isNotZeroSz !sz = isSafeIndex sz zeroIndex +{-# INLINE [1] isNotZeroSz #-} + -- TODO: benchmark against (also adjust `isEmpty` with fastest): -- - foldlIndex (*) 1 (unSz sz) /= 0 -- - foldlIndex (\a x -> a && x /= 0) True (unSz sz) -- - totalElem sz == 0 +-- | Checks whether size can hold at least one element. +-- +-- ==== __Examples__ +-- +-- >>> isZeroSz (Sz3 1 0 2) +-- True +-- +-- @since 1.0.0 +isZeroSz :: Index ix => Sz ix -> Bool +isZeroSz = not . isNotZeroSz +{-# INLINE [1] isZeroSz #-} + +-- | Convert a size to a linear size. +-- +-- @since 0.5.8 +toLinearSz :: Index ix => Sz ix -> Sz1 +toLinearSz = coerce . totalElem +{-# INLINE [1] toLinearSz #-} + -- | Get the outmost dimension of the index. -- -- ==== __Examples__ @@ -273,8 +355,8 @@ initDim = fst . unsnocDim {-# INLINE [1] initDim #-} --- | Change the value of a specific dimension within the index. Throws `IndexException`. See --- `setDimM` for a safer version and `setDimension` for a type safe version. +-- | Change the value of a specific dimension within the index. See `setDimM` for a safer +-- version and `setDimension` for a type safe version. -- -- ==== __Examples__ -- @@ -282,26 +364,24 @@ -- 2 :> 10 :> 4 :. 5 -- -- @since 0.2.4 -setDim' :: Index ix => ix -> Dim -> Int -> ix -setDim' ix dim = either throw id . setDimM ix dim +setDim' :: (HasCallStack, Index ix) => ix -> Dim -> Int -> ix +setDim' ix dim = throwEither . setDimM ix dim {-# INLINE [1] setDim' #-} --- | Change the value from a specific dimension within the index. Throws `IndexException`. See +-- | Change the value from a specific dimension within the index. See -- `getDimM` for a safer version and `getDimension` for a type safe version. -- -- ==== __Examples__ -- -- >>> getDim' (2 :> 3 :> 4 :. 5) 3 -- 3 --- >>> getDim' (2 :> 3 :> 4 :. 5) 0 --- *** Exception: IndexDimensionException: (Dim 0) for (2 :> 3 :> 4 :. 5) -- -- @since 0.2.4 -getDim' :: Index ix => ix -> Dim -> Int -getDim' ix = either throw id . getDimM ix +getDim' :: (HasCallStack, Index ix) => ix -> Dim -> Int +getDim' ix = throwEither . getDimM ix {-# INLINE [1] getDim' #-} --- | Update the value of a specific dimension within the index. Throws `IndexException`. See +-- | Update the value of a specific dimension within the index. See -- `modifyDimM` for a safer version and `modifyDimension` for a type safe version. -- -- ==== __Examples__ @@ -310,17 +390,17 @@ -- (4,2 :> 3 :> 14 :. 5) -- -- @since 0.4.1 -modifyDim' :: Index ix => ix -> Dim -> (Int -> Int) -> (Int, ix) -modifyDim' ix dim = either throw id . modifyDimM ix dim +modifyDim' :: (HasCallStack, Index ix) => ix -> Dim -> (Int -> Int) -> (Int, ix) +modifyDim' ix dim = throwEither . modifyDimM ix dim {-# INLINE [1] modifyDim' #-} -- | Remove a dimension from the index. -- -- ==== __Examples__ -- --- λ> dropDimM (2 :> 3 :> 4 :. 5) 3 :: Maybe Ix3 +-- >>> dropDimM (2 :> 3 :> 4 :. 5) 3 :: Maybe Ix3 -- Just (2 :> 4 :. 5) --- λ> dropDimM (2 :> 3 :> 4 :. 5) 6 :: Maybe Ix3 +-- >>> dropDimM (2 :> 3 :> 4 :. 5) 6 :: Maybe Ix3 -- Nothing -- -- @since 0.3.0 @@ -334,41 +414,36 @@ -- -- >>> dropDim' (2 :> 3 :> 4 :. 5) 3 -- 2 :> 4 :. 5 --- >>> dropDim' (2 :> 3 :> 4 :. 5) 6 --- *** Exception: IndexDimensionException: (Dim 6) for (2 :> 3 :> 4 :. 5) -- -- @since 0.2.4 -dropDim' :: Index ix => ix -> Dim -> Lower ix -dropDim' ix = either throw id . dropDimM ix +dropDim' :: (HasCallStack, Index ix) => ix -> Dim -> Lower ix +dropDim' ix = throwEither . dropDimM ix {-# INLINE [1] dropDim' #-} --- | Lower the dimension of the index by pulling the specified dimension. Throws `IndexException`. See +-- | Lower the dimension of the index by pulling the specified dimension. See -- `pullOutDimM` for a safer version and `pullOutDimension` for a type safe version. -- -- ==== __Examples__ -- --- λ> pullOutDim' (2 :> 3 :> 4 :. 5) 3 +-- >>> pullOutDim' (2 :> 3 :> 4 :. 5) 3 -- (3,2 :> 4 :. 5) -- -- @since 0.2.4 -pullOutDim' :: Index ix => ix -> Dim -> (Int, Lower ix) -pullOutDim' ix = either throw id . pullOutDimM ix +pullOutDim' :: (HasCallStack, Index ix) => ix -> Dim -> (Int, Lower ix) +pullOutDim' ix = throwEither . pullOutDimM ix {-# INLINE [1] pullOutDim' #-} --- | Raise the dimension of the index by inserting one in the specified dimension. Throws --- `IndexException`. See `insertDimM` for a safer version and `insertDimension` for a type safe --- version. +-- | Raise the dimension of the index by inserting one in the specified dimension. See +-- `insertDimM` for a safer version and `insertDimension` for a type safe version. -- -- ==== __Examples__ -- -- >>> insertDim' (2 :> 3 :> 4 :. 5) 3 10 :: Ix5 -- 2 :> 3 :> 10 :> 4 :. 5 --- >>> insertDim' (2 :> 3 :> 4 :. 5) 11 10 :: Ix5 --- *** Exception: IndexDimensionException: (Dim 11) for (2 :> 3 :> 4 :. 5) -- -- @since 0.2.4 -insertDim' :: Index ix => Lower ix -> Dim -> Int -> ix -insertDim' ix dim = either throw id . insertDimM ix dim +insertDim' :: (HasCallStack, Index ix) => Lower ix -> Dim -> Int -> ix +insertDim' ix dim = throwEither . insertDimM ix dim {-# INLINE [1] insertDim' #-} -- | Get the value level `Dim` from the type level equivalent. @@ -422,7 +497,6 @@ getDimension ix = getDim' ix . fromDimension {-# INLINE [1] getDimension #-} - -- | Type safe way of dropping a particular dimension, thus lowering index -- dimensionality. -- @@ -479,41 +553,54 @@ -- 3615 -- -- @since 0.1.0 -iter :: Index ix - => ix -- ^ Start index - -> ix -- ^ End index - -> ix -- ^ Increment - -> (Int -> Int -> Bool) -- ^ Continuation condition - -> a -- ^ Accumulator - -> (ix -> a -> a) -- ^ Iterating function +iter + :: Index ix + => ix + -- ^ Start index + -> ix + -- ^ End index + -> ix + -- ^ Increment + -> (Int -> Int -> Bool) + -- ^ Continuation condition -> a + -- ^ Accumulator + -> (ix -> a -> a) + -- ^ Iterating function + -> a iter sIx eIx incIx cond acc f = runIdentity $ iterM sIx eIx incIx cond acc (\ix -> return . f ix) {-# INLINE iter #-} - -- | Iterate over N-dimensional space linearly from start to end in row-major fashion with an -- accumulator -- -- ==== __Examples__ -- -- >>> sz = Sz2 3 4 --- >>> iterLinearM sz 0 3 1 (<) 100 $ \ k ix acc -> print (fromLinearIndex sz k == ix) >> pure (acc + k) +-- >>> iterLinearM sz 0 3 1 (<) 100 $ \ k ix acc -> (acc + k) <$ print (fromLinearIndex sz k == ix) -- True -- True -- True -- 103 -- -- @since 0.1.0 -iterLinearM :: (Index ix, Monad m) - => Sz ix -- ^ Size - -> Int -- ^ Linear start (must be non-negative) - -> Int -- ^ Linear end (must be less than or equal to @`totalElem` sz@) - -> Int -- ^ Increment (must not be zero) - -> (Int -> Int -> Bool) -- ^ Continuation condition (continue if @True@) - -> a -- ^ Accumulator - -> (Int -> ix -> a -> m a) - -> m a +iterLinearM + :: (Index ix, Monad m) + => Sz ix + -- ^ Size + -> Int + -- ^ Linear start (must be non-negative) + -> Int + -- ^ Linear end (must be less than or equal to @`totalElem` sz@) + -> Int + -- ^ Increment (must not be zero) + -> (Int -> Int -> Bool) + -- ^ Continuation condition (continue if @True@) + -> a + -- ^ Accumulator + -> (Int -> ix -> a -> m a) + -> m a iterLinearM !sz !k0 !k1 !inc cond !acc f = loopM k0 (`cond` k1) (+ inc) acc $ \ !i !acc0 -> f i (fromLinearIndex sz i) acc0 {-# INLINE iterLinearM #-} @@ -529,44 +616,76 @@ -- True -- -- @since 0.1.0 -iterLinearM_ :: (Index ix, Monad m) => - Sz ix -- ^ Size - -> Int -- ^ Start (must be non-negative) - -> Int -- ^ End - -> Int -- ^ Increment (must not be zero) - -> (Int -> Int -> Bool) -- ^ Continuation condition (continue if @True@) - -> (Int -> ix -> m ()) -- ^ Monadic action that takes index in both forms - -> m () +iterLinearM_ + :: (Index ix, Monad m) + => Sz ix + -- ^ Size + -> Int + -- ^ Start (must be non-negative) + -> Int + -- ^ End + -> Int + -- ^ Increment (must not be zero) + -> (Int -> Int -> Bool) + -- ^ Continuation condition (continue if @True@) + -> (Int -> ix -> m ()) + -- ^ Monadic action that takes index in both forms + -> m () iterLinearM_ sz !k0 !k1 !inc cond f = - loopM_ k0 (`cond` k1) (+ inc) $ \ !i -> f i (fromLinearIndex sz i) + loopA_ k0 (`cond` k1) (+ inc) $ \ !i -> f i (fromLinearIndex sz i) {-# INLINE iterLinearM_ #-} +-- | This is used by the @unsafe-checks@ cabal flag. +-- +-- @since 1.1.0 +#ifdef MASSIV_UNSAFE_CHECKS +indexAssert :: (HasCallStack, Index ix) => String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e +indexAssert funName getSize f arr ix + | isSafeIndex sz ix = f arr ix + | otherwise = _errorIx ("<" ++ funName ++ ">") sz ix + where + sz = getSize arr +#else +indexAssert :: String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e +indexAssert _funName _getSize f arr ix = f arr ix +#endif +{-# INLINE indexAssert #-} -- | This is used by @INDEX_CHECK@ macro and thus used whenever the @unsafe-checks@ cabal -- flag is on. -- -- @since 0.4.0 -indexWith :: - Index ix - => String -- ^ Source file name, eg. __FILE__ - -> Int -- ^ Line number in th source file, eg. __LINE__ +indexWith + :: Index ix + => String + -- ^ Source file name, eg. __FILE__ + -> Int + -- ^ Line number in th source file, eg. __LINE__ -> String - -> (arr -> Sz ix) -- ^ Get size of the array - -> (arr -> ix -> e) -- ^ Indexing function - -> arr -- ^ Array - -> ix -- ^ Index + -> (arr -> Sz ix) + -- ^ Get size of the array + -> (arr -> ix -> e) + -- ^ Indexing function + -> arr + -- ^ Array + -> ix + -- ^ Index -> e indexWith fileName lineNo funName getSize f arr ix | isSafeIndex sz ix = f arr ix - | otherwise = errorIx ("<" ++ fileName ++ ":" ++ show lineNo ++ "> " ++ funName) sz ix + | otherwise = _errorIx ("<" ++ fileName ++ ":" ++ show lineNo ++ "> " ++ funName) sz ix where sz = getSize arr +{-# DEPRECATED indexWith "In favor of `indexAssert` that uses HasCallStack" #-} --- | Helper function for throwing out of bounds error. Used by `indexWith` -errorIx :: (Show ix, Show ix') => String -> ix -> ix' -> a -errorIx fName sz ix = +-- | Helper function for throwing out of bounds error. Used by `indexAssert` +_errorIx :: (HasCallStack, Show ix, Show ix') => String -> ix -> ix' -> a +_errorIx fName sz ix = error $ - fName ++ - ": Index out of bounds: (" ++ show ix ++ ") for Array of size: (" ++ show sz ++ ")" -{-# NOINLINE errorIx #-} - + fName + ++ ": Index out of bounds: (" + ++ show ix + ++ ") for Array of size: (" + ++ show sz + ++ ")" +{-# NOINLINE _errorIx #-} diff --git a/src/Data/Massiv/Core/Index/Internal.hs b/src/Data/Massiv/Core/Index/Internal.hs --- a/src/Data/Massiv/Core/Index/Internal.hs +++ b/src/Data/Massiv/Core/Index/Internal.hs @@ -7,81 +7,112 @@ {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE LambdaCase #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE RankNTypes #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE TypeOperators #-} -{-# LANGUAGE TypeSynonymInstances #-} {-# LANGUAGE UndecidableInstances #-} {-# OPTIONS_GHC -Wno-unticked-promoted-constructors #-} -#if __GLASGOW_HASKELL__ < 820 -{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-} -#endif + -- | -- Module : Data.Massiv.Core.Index.Internal --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Core.Index.Internal - ( Sz(SafeSz) - , pattern Sz - , pattern Sz1 - , type Sz1 - , unSz - , zeroSz - , oneSz - , liftSz - , liftSz2 - , consSz - , unconsSz - , snocSz - , unsnocSz - , setSzM - , insertSzM - , pullOutSzM - , Dim(..) - , Dimension(DimN) - , pattern Dim1 - , pattern Dim2 - , pattern Dim3 - , pattern Dim4 - , pattern Dim5 - , IsIndexDimension - , IsDimValid - , ReportInvalidDim - , Lower - , Index(..) - , Ix0(..) - , type Ix1 - , pattern Ix1 - , IndexException(..) - , SizeException(..) - , ShapeException(..) - , showsPrecWrapped - ) where +module Data.Massiv.Core.Index.Internal ( + Sz (SafeSz), + pattern Sz, + pattern Sz1, + unSz, + zeroSz, + oneSz, + liftSz, + liftSz2, + consSz, + unconsSz, + snocSz, + unsnocSz, + setSzM, + insertSzM, + pullOutSzM, + mkSzM, + Dim (..), + Dimension (DimN), + pattern Dim1, + pattern Dim2, + pattern Dim3, + pattern Dim4, + pattern Dim5, + IsIndexDimension, + IsDimValid, + ReportInvalidDim, + Lower, + Index (..), + iterA_, + iterM_, + Ix0 (..), + type Ix1, + pattern Ix1, + IndexException (..), + SizeException (..), + ShapeException (..), + showsPrecWrapped, +) where import Control.DeepSeq -import Control.Exception (Exception(..), throw) -import Control.Monad.Catch (MonadThrow(..)) +import Control.Exception (Exception (..), throw) +import Control.Monad (void, when) +import Control.Monad.Catch (MonadThrow (..)) +import Control.Monad.ST +import Control.Scheduler import Data.Coerce -import Data.Massiv.Core.Iterator +import Data.Kind +import Data.Massiv.Core.Loop import Data.Typeable import GHC.TypeLits +import System.Random.Stateful --- | `Sz` provides type safety guarantees preventing mixup with index, which is used for looking into --- array cells, from the size, that describes total number of elements along each dimension in the --- array. Moreover the @Sz@ constructor will prevent creation of invalid sizes with negative numbers. +-- | `Sz` is the size of the array. It describes total number of elements along +-- each dimension in the array. It is a wrapper around an index of the same +-- dimension, however it provides type safety preventing mixup with +-- index. Moreover the @Sz@ constructor and others such as +-- `Data.Massiv.Core.Index.Sz1`, `Data.Massiv.Core.Index.Sz2`, ... that +-- are specialized to specific dimensions, prevent creation of invalid sizes with +-- negative values by clamping them to zero. -- +-- ====__Examples__ +-- +-- >>> import Data.Massiv.Array +-- >>> Sz (1 :> 2 :. 3) +-- Sz (1 :> 2 :. 3) +-- +-- `Sz` has a `Num` instance, which is very convenient: +-- +-- >>> Sz (1 :> 2 :. 3) + 5 +-- Sz (6 :> 7 :. 8) +-- +-- However subtraction can sometimes lead to surprising behavior, because size is not +-- allowed to take negative values it will be clamped at 0. +-- +-- >>> Sz (1 :> 2 :. 3) - 2 +-- Sz (0 :> 0 :. 1) +-- +-- __Warning__: It is always wrong to `negate` a size, thus it will result in an +-- error. For that reason also watch out for partially applied @(`Prelude.-` sz)@, which is +-- deugared into @`negate` sz@. See more info about it in +-- [#114](https://github.com/lehins/massiv/issues/114). +-- -- @since 0.3.0 -newtype Sz ix = - SafeSz ix - -- ^ Safe size constructor. It is unsafe to use it without making sure that it does not contain - -- negative components. Use `Data.Massiv.Core.Index.Sz` pattern instead. - -- - -- @since 0.3.0 +newtype Sz ix + = -- | Safe size constructor. It is unsafe to use it without making sure that it does not contain + -- negative components. Use `Data.Massiv.Core.Index.Sz` pattern instead. + -- + -- @since 0.3.0 + SafeSz ix deriving (Eq, Ord, NFData) -- | A safe bidirectional pattern synonym for `Sz` construction that will make sure that none of @@ -89,49 +120,80 @@ -- -- @since 0.3.0 pattern Sz :: Index ix => ix -> Sz ix -pattern Sz ix <- SafeSz ix where - Sz ix = SafeSz (liftIndex (max 0) ix) -{-# COMPLETE Sz #-} +pattern Sz ix <- SafeSz ix + where + Sz ix = SafeSz (liftIndex (max 0) ix) --- | 1-dimensional type synonym for size. --- --- @since 0.3.0 -type Sz1 = Sz Ix1 +{-# COMPLETE Sz #-} -- | 1-dimensional size constructor. Especially useful with literals: @(Sz1 5) == Sz (5 :: Int)@. -- -- @since 0.3.0 -pattern Sz1 :: Ix1 -> Sz1 -pattern Sz1 ix <- SafeSz ix where - Sz1 ix = SafeSz (max 0 ix) +pattern Sz1 :: Ix1 -> Sz Ix1 +pattern Sz1 ix <- SafeSz ix + where + Sz1 ix = SafeSz (max 0 ix) + {-# COMPLETE Sz1 #-} +instance (UniformRange ix, Index ix) => Uniform (Sz ix) where + uniformM g = SafeSz <$> uniformRM (pureIndex 0, pureIndex maxBound) g + {-# INLINE uniformM #-} +instance UniformRange ix => UniformRange (Sz ix) where + uniformRM (SafeSz l, SafeSz u) g = SafeSz <$> uniformRM (l, u) g + {-# INLINE uniformRM #-} +#if MIN_VERSION_random(1,3,0) + isInRange (SafeSz l, SafeSz u) (SafeSz k) = isInRange (l, u) k +#endif + +instance (UniformRange ix, Index ix) => Random (Sz ix) + instance Index ix => Show (Sz ix) where showsPrec n sz@(SafeSz usz) = showsPrecWrapped n (str ++) where str = - "Sz" ++ - case unDim (dimensions sz) of - 1 -> "1 " ++ show usz - _ -> " (" ++ shows usz ")" + "Sz" + ++ case unDim (dimensions sz) of + 1 -> "1 " ++ show usz + _ -> " (" ++ shows usz ")" +-- | Calling `negate` is an error. instance (Num ix, Index ix) => Num (Sz ix) where (+) x y = Sz (coerce x + coerce y) {-# INLINE (+) #-} (-) x y = Sz (coerce x - coerce y) {-# INLINE (-) #-} - (*) x y = SafeSz (coerce x * coerce y) + (*) x y = Sz (coerce x * coerce y) {-# INLINE (*) #-} abs !x = x {-# INLINE abs #-} - negate !_x = 0 + negate x + | x == zeroSz = x + | otherwise = + error $ + "Attempted to negate: " + ++ show x + ++ ", this can lead to unexpected behavior. See https://github.com/lehins/massiv/issues/114" {-# INLINE negate #-} signum x = SafeSz (signum (coerce x)) {-# INLINE signum #-} fromInteger = Sz . fromInteger {-# INLINE fromInteger #-} +-- | Construct size from index while checking its correctness. Throws +-- `SizeNegativeException` and `SizeOverflowException`. +-- +-- @since 0.6.0 +mkSzM :: (Index ix, MonadThrow m) => ix -> m (Sz ix) +mkSzM ix = do + let guardNegativeOverflow i !acc = do + when (i < 0) $ throwM $ SizeNegativeException (SafeSz ix) + let acc' = i * acc + when (acc' /= 0 && acc' < acc) $ throwM $ SizeOverflowException (SafeSz ix) + pure acc' + Sz ix <$ foldlIndex (\acc i -> acc >>= guardNegativeOverflow i) (pure 1) ix +{-# INLINE mkSzM #-} -- | Function for unwrapping `Sz`. -- @@ -172,7 +234,6 @@ oneSz = SafeSz (pureIndex 1) {-# INLINE oneSz #-} - -- | Same as `liftIndex`, but for `Sz` -- -- ==== __Example__ @@ -199,7 +260,6 @@ liftSz2 f sz1 sz2 = Sz (liftIndex2 f (coerce sz1) (coerce sz2)) {-# INLINE liftSz2 #-} - -- | Same as `consDim`, but for `Sz` -- -- ==== __Example__ @@ -209,11 +269,10 @@ -- Sz (1 :> 2 :. 3) -- -- @since 0.3.0 -consSz :: Index ix => Sz1 -> Sz (Lower ix) -> Sz ix +consSz :: Index ix => Sz Ix1 -> Sz (Lower ix) -> Sz ix consSz (SafeSz i) (SafeSz ix) = SafeSz (consDim i ix) {-# INLINE consSz #-} - -- | Same as `snocDim`, but for `Sz` -- -- ==== __Example__ @@ -223,7 +282,7 @@ -- Sz (2 :> 3 :. 1) -- -- @since 0.3.0 -snocSz :: Index ix => Sz (Lower ix) -> Sz1 -> Sz ix +snocSz :: Index ix => Sz (Lower ix) -> Sz Ix1 -> Sz ix snocSz (SafeSz i) (SafeSz ix) = SafeSz (snocDim i ix) {-# INLINE snocSz #-} @@ -266,7 +325,7 @@ -- (Sz1 1,Sz (2 :. 3)) -- -- @since 0.3.0 -unconsSz :: Index ix => Sz ix -> (Sz1, Sz (Lower ix)) +unconsSz :: Index ix => Sz ix -> (Sz Ix1, Sz (Lower ix)) unconsSz (SafeSz sz) = coerce (unconsDim sz) {-# INLINE unconsSz #-} @@ -279,7 +338,7 @@ -- (Sz (1 :. 2),Sz1 3) -- -- @since 0.3.0 -unsnocSz :: Index ix => Sz ix -> (Sz (Lower ix), Sz1) +unsnocSz :: Index ix => Sz ix -> (Sz (Lower ix), Sz Ix1) unsnocSz (SafeSz sz) = coerce (unsnocDim sz) {-# INLINE unsnocSz #-} @@ -296,15 +355,25 @@ pullOutSzM (SafeSz sz) = fmap coerce . pullOutDimM sz {-# INLINE pullOutSzM #-} - -- | A way to select Array dimension at a value level. -- -- @since 0.1.0 -newtype Dim = Dim { unDim :: Int } deriving (Eq, Ord, Num, Real, Integral, Enum) +newtype Dim = Dim {unDim :: Int} deriving (Eq, Ord, Num, Real, Integral, Enum, NFData) instance Show Dim where show (Dim d) = "(Dim " ++ show d ++ ")" +instance Uniform Dim where + uniformM g = Dim <$> uniformRM (1, maxBound) g + +instance UniformRange Dim where + uniformRM r g = Dim <$> uniformRM (coerce r) g +#if MIN_VERSION_random(1,3,0) + isInRange = isInRangeOrd +#endif + +instance Random Dim + -- | A way to select Array dimension at a type level. -- -- @since 0.2.4 @@ -341,26 +410,28 @@ pattern Dim5 :: Dimension 5 pattern Dim5 = DimN - -- | A type level constraint that ensures index is indeed valid and that supplied dimension can be -- safely used with it. -- -- @since 0.2.4 type IsIndexDimension ix n = (1 <= n, n <= Dimensions ix, Index ix, KnownNat n) - -- | This type family will always point to a type for a dimension that is one lower than the type -- argument. -- -- @since 0.1.0 -type family Lower ix :: * - +type family Lower ix :: Type type family ReportInvalidDim (dims :: Nat) (n :: Nat) isNotZero isLess :: Bool where ReportInvalidDim dims n True True = True ReportInvalidDim dims n True False = - TypeError (Text "Dimension " :<>: ShowType n :<>: Text " is higher than " :<>: - Text "the maximum expected " :<>: ShowType dims) + TypeError + ( Text "Dimension " + :<>: ShowType n + :<>: Text " is higher than " + :<>: Text "the maximum expected " + :<>: ShowType dims + ) ReportInvalidDim dims n False isLess = TypeError (Text "Zero dimensional indices are not supported") @@ -370,17 +441,19 @@ -- | This is bread and butter of multi-dimensional array indexing. It is unlikely that any of the -- functions in this class will be useful to a regular user, unless general algorithms are being -- implemented that do span multiple dimensions. -class ( Eq ix - , Ord ix - , Show ix - , NFData ix - , Eq (Lower ix) - , Ord (Lower ix) - , Show (Lower ix) - , NFData (Lower ix) - , KnownNat (Dimensions ix) - ) => - Index ix +class + ( Eq ix + , Ord ix + , Show ix + , NFData ix + , Typeable ix + , Eq (Lower ix) + , Ord (Lower ix) + , Show (Lower ix) + , NFData (Lower ix) + , KnownNat (Dimensions ix) + ) => + Index ix where -- | Type level information on how many dimensions this index has. -- @@ -469,23 +542,33 @@ -- | Perform a left fold over the index foldlIndex :: (a -> Int -> a) -> a -> ix -> a - default foldlIndex :: Index (Lower ix) => - (a -> Int -> a) -> a -> ix -> a + default foldlIndex + :: Index (Lower ix) + => (a -> Int -> a) + -> a + -> ix + -> a foldlIndex f !acc !ix = foldlIndex f (f acc i0) ixL where !(i0, ixL) = unconsDim ix {-# INLINE [1] foldlIndex #-} -- TODO: implement in terms of foldlIndex and pull out of the class + -- | Check whether index is positive and is within the size. -- -- @since 0.1.0 - isSafeIndex :: - Sz ix -- ^ Size - -> ix -- ^ Index + isSafeIndex + :: Sz ix + -- ^ Size + -> ix + -- ^ Index -> Bool - default isSafeIndex :: Index (Lower ix) => - Sz ix -> ix -> Bool + default isSafeIndex + :: Index (Lower ix) + => Sz ix + -> ix + -> Bool isSafeIndex sz !ix = isSafeIndex n0 i0 && isSafeIndex szL ixL where !(n0, szL) = unconsSz sz @@ -495,12 +578,13 @@ -- | Convert linear index from size and index -- -- @since 0.1.0 - toLinearIndex :: - Sz ix -- ^ Size - -> ix -- ^ Index - -> Int - default toLinearIndex :: Index (Lower ix) => - Sz ix -> ix -> Int + toLinearIndex + :: Sz ix + -- ^ Size + -> ix + -- ^ Index + -> Ix1 + default toLinearIndex :: Index (Lower ix) => Sz ix -> ix -> Ix1 toLinearIndex (SafeSz sz) !ix = toLinearIndex (SafeSz szL) ixL * n + i where !(szL, n) = unsnocDim sz @@ -508,12 +592,11 @@ {-# INLINE [1] toLinearIndex #-} -- | Convert linear index from size and index with an accumulator. Currently is useless and will - -- likley be removed in future versions. + -- likely be removed in future versions. -- -- @since 0.1.0 - toLinearIndexAcc :: Int -> ix -> ix -> Int - default toLinearIndexAcc :: Index (Lower ix) => - Int -> ix -> ix -> Int + toLinearIndexAcc :: Ix1 -> ix -> ix -> Ix1 + default toLinearIndexAcc :: Index (Lower ix) => Ix1 -> ix -> ix -> Ix1 toLinearIndexAcc !acc !sz !ix = toLinearIndexAcc (acc * n + i) szL ixL where !(n, szL) = unconsDim sz @@ -523,40 +606,47 @@ -- | Compute an index from size and linear index -- -- @since 0.1.0 - fromLinearIndex :: Sz ix -> Int -> ix - default fromLinearIndex :: Index (Lower ix) => - Sz ix -> Int -> ix - fromLinearIndex (SafeSz sz) k = consDim q ixL + fromLinearIndex :: Sz ix -> Ix1 -> ix + default fromLinearIndex :: Index (Lower ix) => Sz ix -> Ix1 -> ix + fromLinearIndex (SafeSz sz) !k = consDim q ixL where - !(q, ixL) = fromLinearIndexAcc (snd (unconsDim sz)) k + !(!q, !ixL) = fromLinearIndexAcc (snd (unconsDim sz)) k {-# INLINE [1] fromLinearIndex #-} -- | Compute an index from size and linear index using an accumulator, thus trying to optimize for -- tail recursion while getting the index computed. -- -- @since 0.1.0 - fromLinearIndexAcc :: ix -> Int -> (Int, ix) - default fromLinearIndexAcc :: Index (Lower ix) => - ix -> Int -> (Int, ix) - fromLinearIndexAcc ix' !k = (q, consDim r ixL) + fromLinearIndexAcc :: ix -> Ix1 -> (Int, ix) + default fromLinearIndexAcc :: Index (Lower ix) => ix -> Ix1 -> (Ix1, ix) + fromLinearIndexAcc !ix' !k = (q, consDim r ixL) where - !(m, ix) = unconsDim ix' - !(kL, ixL) = fromLinearIndexAcc ix k - !(q, r) = quotRem kL m + !(!m, !ix) = unconsDim ix' + !(!kL, !ixL) = fromLinearIndexAcc ix k + !(!q, !r) = quotRem kL m {-# INLINE [1] fromLinearIndexAcc #-} -- | A way to make sure index is withing the bounds for the supplied size. Takes two functions -- that will be invoked whenever index (2nd arg) is outsize the supplied size (1st arg) -- -- @since 0.1.0 - repairIndex :: - Sz ix -- ^ Size - -> ix -- ^ Index - -> (Sz Int -> Int -> Int) -- ^ Repair when below zero - -> (Sz Int -> Int -> Int) -- ^ Repair when higher than size + repairIndex + :: Sz ix + -- ^ Size -> ix - default repairIndex :: Index (Lower ix) => - Sz ix -> ix -> (Sz Int -> Int -> Int) -> (Sz Int -> Int -> Int) -> ix + -- ^ Index + -> (Sz Int -> Int -> Int) + -- ^ Repair when below zero + -> (Sz Int -> Int -> Int) + -- ^ Repair when higher than size + -> ix + default repairIndex + :: Index (Lower ix) + => Sz ix + -> ix + -> (Sz Int -> Int -> Int) + -> (Sz Int -> Int -> Int) + -> ix repairIndex sz !ix rBelow rOver = consDim (repairIndex n i rBelow rOver) (repairIndex szL ixL rBelow rOver) where @@ -567,17 +657,30 @@ -- | This function is what makes it possible to iterate over an array of any dimension. -- -- @since 0.1.0 - iterM :: - Monad m - => ix -- ^ Start index - -> ix -- ^ End index - -> ix -- ^ Increment - -> (Int -> Int -> Bool) -- ^ Continue iterating while predicate is True (eg. until end of row) - -> a -- ^ Initial value for an accumulator - -> (ix -> a -> m a) -- ^ Accumulator function + iterM + :: Monad m + => ix + -- ^ Start index + -> ix + -- ^ End index + -> ix + -- ^ Increment + -> (Int -> Int -> Bool) + -- ^ Continue iterating while predicate is True (eg. until end of row) + -> a + -- ^ Initial value for an accumulator + -> (ix -> a -> m a) + -- ^ Accumulator function -> m a - default iterM :: (Index (Lower ix), Monad m) => - ix -> ix -> ix -> (Int -> Int -> Bool) -> a -> (ix -> a -> m a) -> m a + default iterM + :: (Index (Lower ix), Monad m) + => ix + -> ix + -> ix + -> (Int -> Int -> Bool) + -> a + -> (ix -> a -> m a) + -> m a iterM !sIx eIx !incIx cond !acc f = loopM s (`cond` e) (+ inc) acc $ \ !i !acc0 -> iterM sIxL eIxL incIxL cond acc0 $ \ !ix -> f (consDim i ix) @@ -587,21 +690,288 @@ !(inc, incIxL) = unconsDim incIx {-# INLINE iterM #-} - -- TODO: Implement in terms of iterM, benchmark it and remove from `Index` - -- | Same as `iterM`, but don't bother with accumulator and return value. + iterRowMajorST + :: Int + -- ^ Scheduler multiplying factor. Must be positive + -> Scheduler s a + -- ^ Scheduler to use + -> ix + -- ^ Start index + -> ix + -- ^ Stride + -> Sz ix + -- ^ Size + -> a + -- ^ Initial accumulator + -> (a -> ST s (a, a)) + -- ^ Function that splits accumulator for each scheduled job. + -> (ix -> a -> ST s a) + -- ^ Action + -> ST s a + default iterRowMajorST + :: Index (Lower ix) + => Int + -> Scheduler s a + -> ix + -> ix + -> Sz ix + -> a + -> (a -> ST s (a, a)) + -> (ix -> a -> ST s a) + -> ST s a + iterRowMajorST !fact scheduler ixStart ixStride sz initAcc splitAcc f = do + let !(SafeSz n, szL@(SafeSz nL)) = unconsSz sz + if n > 0 + then do + let !(!start, !ixL) = unconsDim ixStart + !(!stride, !sL) = unconsDim ixStride + if numWorkers scheduler > 1 && fact > 1 && n < numWorkers scheduler * fact + then do + let !newFact = 1 + (fact `quot` n) + loopM start (< start + n * stride) (+ stride) initAcc $ \j acc -> + iterRowMajorST newFact scheduler ixL sL szL acc splitAcc (f . consDim j) + else splitWorkWithFactorST fact scheduler start stride n initAcc splitAcc $ + \_ _ chunkStartAdj chunkStopAdj acc -> + loopM chunkStartAdj (< chunkStopAdj) (+ stride) acc $ \j a -> + iterM ixL nL sL (<) a (f . consDim j) + else pure initAcc + {-# INLINE iterRowMajorST #-} + + -- | Similar to `iterM`, but no restriction on a Monad. -- + -- iterF (-10) 20 4 (<) [] (:) :: [Int] + -- [-10,-6,-2,2,6,10,14,18] + -- + -- @since 1.0.2 + iterF :: ix -> ix -> ix -> (Int -> Int -> Bool) -> f a -> (ix -> f a -> f a) -> f a + default iterF + :: Index (Lower ix) + => ix + -> ix + -> ix + -> (Int -> Int -> Bool) + -> f a + -> (ix -> f a -> f a) + -> f a + iterF !sIx !eIx !incIx cond initAct f = + loopF s (`cond` e) (+ inc) initAct $ \ !i g -> + iterF sIxL eIxL incIxL cond g (\ !ix -> f (consDim i ix)) + where + !(s, sIxL) = unconsDim sIx + !(e, eIxL) = unconsDim eIx + !(inc, incIxL) = unconsDim incIx + {-# INLINE iterF #-} + + -- | A single step in iteration + -- -- @since 0.1.0 - iterM_ :: Monad m => ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> m a) -> m () - default iterM_ :: (Index (Lower ix), Monad m) => - ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> m a) -> m () - iterM_ !sIx eIx !incIx cond f = - loopM_ s (`cond` e) (+ inc) $ \ !i -> iterM_ sIxL eIxL incIxL cond $ \ !ix -> f (consDim i ix) + stepNextMF :: ix -> ix -> ix -> (Int -> Int -> Bool) -> (Maybe ix -> f a) -> f a + default stepNextMF + :: Index (Lower ix) + => ix + -> ix + -> ix + -> (Int -> Int -> Bool) + -> (Maybe ix -> f a) + -> f a + stepNextMF !sIx !eIx !incIx cond f = + nextMaybeF s (`cond` e) (+ inc) $ \ !mni -> + stepNextMF sIxL eIxL incIxL cond $ \mIxN -> + f $! + case mIxN of + Just ixN -> Just $! consDim s ixN + Nothing -> + case mni of + Just ni -> Just $! consDim ni (pureIndex 0) + Nothing -> Nothing where !(s, sIxL) = unconsDim sIx !(e, eIxL) = unconsDim eIx !(inc, incIxL) = unconsDim incIx - {-# INLINE iterM_ #-} + {-# INLINE stepNextMF #-} + iterTargetRowMajorA_ + :: Applicative f + => Int + -- ^ Target linear index accumulator + -> Int + -- ^ Target linear index start + -> Sz ix + -- ^ Target size + -> ix + -- ^ Source start index + -> ix + -- ^ Source stride + -> (Ix1 -> ix -> f a) + -- ^ Action that accepts a linear index of the target, + -- multi-dimensional index of the source and accumulator + -> f () + default iterTargetRowMajorA_ + :: (Applicative f, Index (Lower ix)) + => Int + -> Int + -> Sz ix + -> ix + -> ix + -> (Ix1 -> ix -> f a) + -> f () + iterTargetRowMajorA_ !iAcc !iStart szRes ixStart ixStride f = do + let !(SafeSz nRes, !szL) = unconsSz szRes + !(!start, !ixL) = unconsDim ixStart + !(!stride, !sL) = unconsDim ixStride + iloopA_ (iAcc * nRes) start (< start + nRes * stride) (+ stride) $ \k j -> + iterTargetRowMajorA_ k iStart szL ixL sL $ \i jl -> f i (consDim j jl) + {-# INLINE iterTargetRowMajorA_ #-} + + iterTargetRowMajorAccM + :: Monad m + => Int + -- ^ Target linear index accumulator + -> Int + -- ^ Target linear index start + -> Sz ix + -- ^ Target size + -> ix + -- ^ Source start index + -> ix + -- ^ Source stride + -> a + -- ^ Accumulator + -> (Ix1 -> ix -> a -> m a) + -- ^ Action that accepts a linear index of the target, + -- multi-dimensional index of the source and accumulator + -> m a + default iterTargetRowMajorAccM + :: (Monad m, Index (Lower ix)) + => Int + -> Int + -> Sz ix + -> ix + -> ix + -> a + -> (Ix1 -> ix -> a -> m a) + -> m a + iterTargetRowMajorAccM !iAcc !iStart szRes ixStart ixStride initAcc f = do + let !(SafeSz nRes, !szL) = unconsSz szRes + !(!start, !ixL) = unconsDim ixStart + !(!stride, !sL) = unconsDim ixStride + iloopM (iAcc * nRes) start (< start + nRes * stride) (+ stride) initAcc $ \k j acc -> + iterTargetRowMajorAccM k iStart szL ixL sL acc $ \i jl -> f i (consDim j jl) + {-# INLINE iterTargetRowMajorAccM #-} + + iterTargetRowMajorAccST + :: Int + -- ^ Linear index accumulator + -> Int + -- ^ Scheduler multiplying factor. Must be positive + -> Scheduler s a + -- ^ Scheduler to use + -> Int + -- ^ Target linear index start + -> Sz ix + -- ^ Target size + -> ix + -- ^ Source start index + -> ix + -- ^ Source stride + -> a + -- ^ Initial accumulator + -> (a -> ST s (a, a)) + -- ^ Function that splits accumulator for each scheduled job. + -> (Ix1 -> ix -> a -> ST s a) + -- ^ Action + -> ST s a + default iterTargetRowMajorAccST + :: Index (Lower ix) + => Int + -> Int + -> Scheduler s a + -> Int + -> Sz ix + -> ix + -> ix + -> a + -> (a -> ST s (a, a)) + -> (Ix1 -> ix -> a -> ST s a) + -> ST s a + iterTargetRowMajorAccST !iAcc !fact scheduler iStart sz ixStart ixStride initAcc splitAcc f = do + let !(SafeSz n, nL) = unconsSz sz + if n > 0 + then do + let !(!start, !ixL) = unconsDim ixStart + !(!stride, !sL) = unconsDim ixStride + !iAccL = iAcc * n + if numWorkers scheduler > 1 && fact > 1 && n < numWorkers scheduler * fact + then do + let newFact = 1 + (fact `quot` n) + iloopM iAccL start (< start + n * stride) (+ stride) initAcc $ \k j acc -> do + iterTargetRowMajorAccST k newFact scheduler iStart nL ixL sL acc splitAcc $ \i -> + f i . consDim j + else splitWorkWithFactorST fact scheduler start stride n initAcc splitAcc $ + \chunkStart _ chunkStartAdj chunkStopAdj acc -> + iloopM (iAccL + chunkStart) chunkStartAdj (< chunkStopAdj) (+ stride) acc $ \k j a -> + iterTargetRowMajorAccM k iStart nL ixL sL a $ \i -> f i . consDim j + else pure initAcc + {-# INLINE iterTargetRowMajorAccST #-} + + iterTargetRowMajorAccST_ + :: Int + -- ^ Index accumulator + -> Int + -- ^ Scheduler multiplying factor. Must be positive + -> Scheduler s () + -- ^ Scheduler to use + -> Int + -- ^ Target linear start index + -> Sz ix + -- ^ Target size + -> ix + -- ^ Source start index + -> ix + -- ^ Source stride + -> a + -- ^ Initial accumulator + -> (a -> ST s (a, a)) + -- ^ Function that splits accumulator for each scheduled job. + -> (Ix1 -> ix -> a -> ST s a) + -- ^ Action + -> ST s () + default iterTargetRowMajorAccST_ + :: Index (Lower ix) + => Int + -> Int + -> Scheduler s () + -> Int + -> Sz ix + -> ix + -> ix + -> a + -> (a -> ST s (a, a)) + -> (Ix1 -> ix -> a -> ST s a) + -> ST s () + iterTargetRowMajorAccST_ !iAcc !fact scheduler iStart sz ixStart ixStride initAcc splitAcc f = do + let !(SafeSz n, szL) = unconsSz sz + when (n > 0) $ do + let !(!start, !ixL) = unconsDim ixStart + !(!stride, !sL) = unconsDim ixStride + !iAccL = iAcc * n + if numWorkers scheduler > 1 && fact > 1 && n < numWorkers scheduler * fact + then do + let !newFact = 1 + (fact `quot` n) + void $ iloopM iAccL start (< n * stride) (+ stride) initAcc $ \k j acc -> do + (accCur, accNext) <- splitAcc acc + scheduleWork_ scheduler $ + iterTargetRowMajorAccST_ k newFact scheduler iStart szL ixL sL accCur splitAcc $ \i -> + f i . consDim j + pure accNext + else void $ + splitWorkWithFactorST fact scheduler start stride n initAcc splitAcc $ + \chunkStart _ chunkStartAdj chunkStopAdj acc -> + void $ + iloopM (iAccL + chunkStart) chunkStartAdj (< chunkStopAdj) (+ stride) acc $ \k j a -> + iterTargetRowMajorAccM k iStart szL ixL sL a $ \i -> f i . consDim j + {-# INLINE iterTargetRowMajorAccST_ #-} + -- | Zero-dimension, i.e. a scalar. Can't really be used directly as there is no instance of -- `Index` for it, and is included for completeness. data Ix0 = Ix0 deriving (Eq, Ord, Show) @@ -628,10 +998,15 @@ -- @since 0.1.0 pattern Ix1 :: Int -> Ix1 pattern Ix1 i = i + {-# COMPLETE Ix1 #-} type instance Lower Int = Ix0 +-- This is needed to avoid GHC from doing redundant allocations +throwIndexZeroException :: Int -> a +throwIndexZeroException = throw . IndexZeroException +{-# NOINLINE throwIndexZeroException #-} instance Index Ix1 where type Dimensions Ix1 = 1 @@ -643,14 +1018,14 @@ {-# INLINE [1] isSafeIndex #-} toLinearIndex _ = id {-# INLINE [1] toLinearIndex #-} - toLinearIndexAcc !acc m i = acc * m + i + toLinearIndexAcc !acc m i = acc * m + i {-# INLINE [1] toLinearIndexAcc #-} fromLinearIndex _ = id {-# INLINE [1] fromLinearIndex #-} fromLinearIndexAcc n k = k `quotRem` n {-# INLINE [1] fromLinearIndexAcc #-} repairIndex k@(SafeSz ksz) !i rBelow rOver - | ksz <= 0 = throw $ IndexZeroException ksz + | ksz <= 0 = throwIndexZeroException ksz | i < 0 = rBelow k i | i >= ksz = rOver k i | otherwise = i @@ -666,8 +1041,8 @@ getDimM ix 1 = pure ix getDimM ix d = throwM $ IndexDimensionException ix d {-# INLINE [1] getDimM #-} - setDimM _ 1 ix = pure ix - setDimM ix d _ = throwM $ IndexDimensionException ix d + setDimM _ 1 ix = pure ix + setDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] setDimM #-} modifyDimM ix 1 f = pure (ix, f ix) modifyDimM ix d _ = throwM $ IndexDimensionException ix d @@ -676,7 +1051,7 @@ pullOutDimM ix d = throwM $ IndexDimensionException ix d {-# INLINE [1] pullOutDimM #-} insertDimM Ix0 1 i = pure i - insertDimM ix d _ = throwM $ IndexDimensionException ix d + insertDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] insertDimM #-} pureIndex i = i {-# INLINE [1] pureIndex #-} @@ -686,12 +1061,72 @@ {-# INLINE [1] liftIndex2 #-} foldlIndex f = f {-# INLINE [1] foldlIndex #-} - iterM k0 k1 inc cond = loopM k0 (`cond` k1) (+inc) + iterM k0 k1 inc cond = loopM k0 (`cond` k1) (+ inc) {-# INLINE iterM #-} - iterM_ k0 k1 inc cond = loopM_ k0 (`cond` k1) (+inc) - {-# INLINE iterM_ #-} + iterF k0 k1 inc cond = loopF k0 (`cond` k1) (+ inc) + {-# INLINE iterF #-} + stepNextMF k0 k1 inc cond = nextMaybeF k0 (`cond` k1) (+ inc) + {-# INLINE stepNextMF #-} + iterRowMajorST fact scheduler start step n = + iterLinearAccST fact scheduler start step (unSz n) + {-# INLINE iterRowMajorST #-} + iterTargetRowMajorA_ iAcc iStart (SafeSz nRes) start stride = + iloopA_ (iAcc * nRes + iStart) start (< start + nRes * stride) (+ stride) + {-# INLINE iterTargetRowMajorA_ #-} + + iterTargetRowMajorAccM iAcc iStart (SafeSz nRes) start stride = + iloopM (iAcc * nRes + iStart) start (< start + nRes * stride) (+ stride) + {-# INLINE iterTargetRowMajorAccM #-} + + iterTargetRowMajorAccST iAcc fact scheduler iStart sz start stride initAcc splitAcc action = do + let !n = unSz sz + !iAccL = iStart + iAcc * n + splitWorkWithFactorST fact scheduler start stride n initAcc splitAcc $ + \chunkStart _ chunkStartAdj chunkStopAdj acc -> + iloopM (iAccL + chunkStart) chunkStartAdj (< chunkStopAdj) (+ stride) acc action + {-# INLINE iterTargetRowMajorAccST #-} + + iterTargetRowMajorAccST_ iAcc fact scheduler iStart sz start stride initAcc splitAcc action = do + let !n = unSz sz + !iAccL = iStart + iAcc * n + void $ + splitWorkWithFactorST fact scheduler start stride n initAcc splitAcc $ + \chunkStart _ chunkStartAdj chunkStopAdj acc -> + void $ iloopM (iAccL + chunkStart) chunkStartAdj (< chunkStopAdj) (+ stride) acc action + {-# INLINE iterTargetRowMajorAccST_ #-} + +-- | Same as `iterM`, but don't bother with accumulator and return value. +-- +-- @since 0.1.0 +iterM_ :: (Index ix, Monad m) => ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> m a) -> m () +iterM_ sIx eIx incIx cond f = iterM sIx eIx incIx cond () $ \ !ix !a -> f ix >> pure a +{-# INLINE iterM_ #-} +{-# DEPRECATED iterM_ "In favor of more lax `iterA_`" #-} + +-- | Same as `iterM`, Iterate over a region with specific step, but using +-- `Applicative` instead of a `Monad` and don't bother with accumulator or return value. +-- +-- @since 1.0.2 +iterA_ + :: forall ix f a + . (Index ix, Applicative f) + => ix + -- ^ Starting index + -> ix + -- ^ Ending index (not included) + -> ix + -- ^ Stepping index + -> (Int -> Int -> Bool) + -- ^ Continuation function. Loop will stop on `False` + -> (ix -> f a) + -- ^ Action applied to an index. Result is ignored. + -> f () +iterA_ sIx eIx incIx cond f = + iterF sIx eIx incIx cond (pure ()) $ \ix go -> f ix *> go +{-# INLINE iterA_ #-} + -- | Exceptions that get thrown when there is a problem with an index, size or dimension. -- -- @since 0.3.0 @@ -699,7 +1134,7 @@ -- | Index contains a zero value along one of the dimensions. IndexZeroException :: Index ix => !ix -> IndexException -- | Dimension is out of reach. - IndexDimensionException :: (Show ix, Typeable ix) => !ix -> !Dim -> IndexException + IndexDimensionException :: (NFData ix, Eq ix, Show ix, Typeable ix) => !ix -> !Dim -> IndexException -- | Index is out of bounds. IndexOutOfBoundsException :: Index ix => !(Sz ix) -> !ix -> IndexException @@ -714,13 +1149,23 @@ instance Eq IndexException where e1 == e2 = case (e1, e2) of - (IndexZeroException i1, IndexZeroException i2) -> show i1 == show i2 - (IndexDimensionException i1 d1, IndexDimensionException i2 d2) -> - show i1 == show i2 && d1 == d2 - (IndexOutOfBoundsException sz1 i1, IndexOutOfBoundsException sz2 i2) -> - show sz1 == show sz2 && show i1 == show i2 + (IndexZeroException i1, IndexZeroException i2t) + | Just i2 <- cast i2t -> i1 == i2 + (IndexDimensionException i1 d1, IndexDimensionException i2t d2) + | Just i2 <- cast i2t -> i1 == i2 && d1 == d2 + (IndexOutOfBoundsException sz1 i1, IndexOutOfBoundsException sz2t i2t) + | Just i2 <- cast i2t + , Just sz2 <- cast sz2t -> + sz1 == sz2 && i1 == i2 _ -> False +instance NFData IndexException where + rnf = + \case + IndexZeroException i -> rnf i + IndexDimensionException i d -> i `deepseq` rnf d + IndexOutOfBoundsException sz i -> sz `deepseq` rnf i + instance Exception IndexException -- | Exception that indicates an issue with an array size. @@ -735,21 +1180,50 @@ SizeSubregionException :: Index ix => !(Sz ix) -> !ix -> !(Sz ix) -> SizeException -- | An array with the size cannot contain any elements. SizeEmptyException :: Index ix => !(Sz ix) -> SizeException + -- | Total number of elements is too large resulting in overflow. + -- + -- @since 0.6.0 + SizeOverflowException :: Index ix => !(Sz ix) -> SizeException + -- | At least one dimensions contain a negative value. + -- + -- @since 0.6.0 + SizeNegativeException :: Index ix => !(Sz ix) -> SizeException instance Eq SizeException where e1 == e2 = case (e1, e2) of - (SizeMismatchException sz1 sz1', SizeMismatchException sz2 sz2') -> - show sz1 == show sz2 && show sz1' == show sz2' - (SizeElementsMismatchException sz1 sz1', SizeElementsMismatchException sz2 sz2') -> - show sz1 == show sz2 && show sz1' == show sz2' - (SizeSubregionException sz1 i1 sz1', SizeSubregionException sz2 i2 sz2') -> - show sz1 == show sz2 && show i1 == show i2 && show sz1' == show sz2' - (SizeEmptyException sz1, SizeEmptyException sz2) -> show sz1 == show sz2 + (SizeMismatchException sz1 sz1', SizeMismatchException sz2t sz2t') + | Just sz2 <- cast sz2t + , Just sz2' <- cast sz2t' -> + sz1 == sz2 && sz1' == sz2' + (SizeElementsMismatchException sz1 sz1', SizeElementsMismatchException sz2t sz2t') + | Just sz2 <- cast sz2t + , Just sz2' <- cast sz2t' -> + sz1 == sz2 && sz1' == sz2' + (SizeSubregionException sz1 i1 sz1', SizeSubregionException sz2t i2t sz2t') + | Just sz2 <- cast sz2t + , Just i2 <- cast i2t + , Just sz2' <- cast sz2t' -> + sz1 == sz2 && i1 == i2 && sz1' == sz2' + (SizeEmptyException sz1, SizeEmptyException sz2t) + | Just sz2 <- cast sz2t -> sz1 == sz2 + (SizeOverflowException sz1, SizeOverflowException sz2t) + | Just sz2 <- cast sz2t -> sz1 == sz2 + (SizeNegativeException sz1, SizeNegativeException sz2t) + | Just sz2 <- cast sz2t -> sz1 == sz2 _ -> False -instance Exception SizeException +instance NFData SizeException where + rnf = + \case + SizeMismatchException sz sz' -> sz `deepseq` rnf sz' + SizeElementsMismatchException sz sz' -> sz `deepseq` rnf sz' + SizeSubregionException sz i sz' -> sz `deepseq` i `deepseq` rnf sz' + SizeEmptyException sz -> rnf sz + SizeOverflowException sz -> rnf sz + SizeNegativeException sz -> rnf sz +instance Exception SizeException instance Show SizeException where show (SizeMismatchException sz sz') = @@ -757,10 +1231,19 @@ show (SizeElementsMismatchException sz sz') = "SizeElementsMismatchException: (" ++ show sz ++ ") vs (" ++ show sz' ++ ")" show (SizeSubregionException sz' ix sz) = - "SizeSubregionException: (" ++ - show sz' ++ ") is to small for " ++ show ix ++ " (" ++ show sz ++ ")" + "SizeSubregionException: (" + ++ show sz' + ++ ") is to small for " + ++ show ix + ++ " (" + ++ show sz + ++ ")" show (SizeEmptyException sz) = "SizeEmptyException: (" ++ show sz ++ ") corresponds to an empty array" + show (SizeOverflowException sz) = + "SizeOverflowException: (" ++ show sz ++ ") is too big" + show (SizeNegativeException sz) = + "SizeNegativeException: (" ++ show sz ++ ") contains negative value" showsPrec n exc = showsPrecWrapped n (show exc ++) -- | Exception that can happen upon conversion of a ragged type array into the rectangular kind. Which @@ -768,21 +1251,37 @@ -- -- @since 0.3.0 data ShapeException - = DimTooShortException !Sz1 !Sz1 - | DimTooLongException - deriving Eq + = -- | Across a specific dimension there was not enough elements for the supplied size + DimTooShortException !Dim !(Sz Ix1) !(Sz Ix1) + | -- | Across a specific dimension there was too many elements for the supplied size + DimTooLongException !Dim !(Sz Ix1) !(Sz Ix1) + | -- | Expected an empty size, but the shape was not empty. + ShapeNonEmpty + deriving (Eq) instance Show ShapeException where - showsPrec _ DimTooLongException = ("DimTooLongException" ++) - showsPrec n (DimTooShortException sz sz') = - showsPrecWrapped - n - (("DimTooShortException: expected (" ++) . shows sz . ("), got (" ++) . shows sz' . (")" ++)) + showsPrec n = + \case + DimTooShortException d sz sz' -> showsShapeExc "DimTooShortException" d sz sz' + DimTooLongException d sz sz' -> showsShapeExc "DimTooLongException" d sz sz' + ShapeNonEmpty -> ("ShapeNonEmpty" ++) + where + showsShapeExc tyName d sz sz' = + showsPrecWrapped + n + ( (tyName ++) + . (" for " ++) + . shows d + . (": expected (" ++) + . shows sz + . ("), got (" ++) + . shows sz' + . (")" ++) + ) instance Exception ShapeException - showsPrecWrapped :: Int -> ShowS -> ShowS showsPrecWrapped n inner | n < 1 = inner - | otherwise = ('(':) . inner . (")" ++) + | otherwise = ('(' :) . inner . (")" ++) diff --git a/src/Data/Massiv/Core/Index/Iterator.hs b/src/Data/Massiv/Core/Index/Iterator.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Core/Index/Iterator.hs @@ -0,0 +1,539 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE ExplicitForAll #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MonoLocalBinds #-} +{-# LANGUAGE PatternSynonyms #-} + +-- | +-- Module : Data.Massiv.Core.Index.Iterator +-- Copyright : (c) Alexey Kuleshevich 2021-2022 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +module Data.Massiv.Core.Index.Iterator ( + Iterator (..), + + -- * Extra iterator functions + iterTargetAccST, + iterTargetAccST_, + iterTargetFullWithStrideAccST, + iterTargetFullWithStrideAccST_, + iterTargetST_, + iterTargetFullWithStrideST_, + + -- * Iterator implementations + RowMajor (RowMajor), + defRowMajor, + RowMajorLinear (RowMajorLinear), + defRowMajorLinear, + RowMajorUnbalanced (RowMajorUnbalanced), + defRowMajorUnbalanced, +) where + +import Control.Monad +import Control.Monad.ST +import Control.Scheduler +import Data.Massiv.Core.Index.Internal +import Data.Massiv.Core.Index.Stride +import Data.Massiv.Core.Loop + +class Iterator it where + {-# MINIMAL (iterTargetM, iterTargetA_, iterTargetWithStrideAccST, iterTargetWithStrideAccST_) #-} + + -- | Iterate over a target region using linear index with access to the source + -- index, which adjusted according to the stride. Use `iterTargetM` if you + -- need an accumulator. + -- + -- @since 1.0.2 + iterTargetA_ + :: (Index ix, Applicative f) + => it + -> Int + -- ^ Target linear index start + -> Sz ix + -- ^ Target size + -> ix + -- ^ Source start index + -> Stride ix + -- ^ Source stride + -> (Ix1 -> ix -> f a) + -- ^ Action that accepts a linear index of the target and multi-dimensional + -- index of the source. + -> f () + + -- | Iterate over a target region using linear index with access to the source + -- index, which adjusted according to the stride. + -- + -- @since 1.0.2 + iterTargetM + :: (Index ix, Monad m) + => it + -> Ix1 + -- ^ Target linear index start + -> Sz ix + -- ^ Target size + -> ix + -- ^ Source start index + -> Stride ix + -- ^ Source stride + -> a + -- ^ Accumulator + -> (Ix1 -> ix -> a -> m a) + -- ^ Action that accepts a linear index of the target, + -- multi-dimensional index of the source and accumulator + -> m a + + iterTargetWithStrideAccST + :: Index ix + => it + -> Scheduler s a + -- ^ Scheduler to use + -> Ix1 + -- ^ Target linear start index + -> Sz ix + -- ^ Target size + -> ix + -- ^ Source start index + -> Stride ix + -- ^ Source stride + -> a + -- ^ Initial accumulator + -> (a -> ST s (a, a)) + -- ^ Splitting action that produces new accumulators for separate worker threads. + -> (Ix1 -> ix -> a -> ST s a) + -- ^ Action + -> ST s a + + iterTargetWithStrideAccST_ + :: Index ix + => it + -> Scheduler s () + -- ^ Scheduler to use + -> Ix1 + -- ^ Target linear start index + -> Sz ix + -- ^ Target size + -> ix + -- ^ Start + -> Stride ix + -- ^ Stride + -> a + -- ^ Initial accumulator + -> (a -> ST s (a, a)) + -- ^ Splitting action that produces new accumulators for separate worker threads. + -> (Ix1 -> ix -> a -> ST s a) + -- ^ Action + -> ST s () + + -- | Iterate over a region with a monadic action and accumulator. + -- + -- @since 1.0.2 + iterFullM + :: (Index ix, Monad m) + => it + -> ix + -- ^ Source start index + -> Sz ix + -- ^ Source size + -> a + -- ^ Accumulator + -> (ix -> a -> m a) + -- ^ Action that accepts a linear index of the target, + -- multi-dimensional index of the source and accumulator + -> m a + iterFullM it start sz acc f = + iterTargetM it 0 sz start oneStride acc (const f) + {-# INLINE iterFullM #-} + + -- | Iterate over a region with an applicative action ignoring the result. + -- + -- @since 1.0.2 + iterFullA_ + :: (Index ix, Applicative f) + => it + -> ix + -- ^ Source start index + -> Sz ix + -- ^ Source size + -> (ix -> f a) + -- ^ Action that accepts a linear index of the target, + -- multi-dimensional index of the source and accumulator + -> f () + iterFullA_ it start sz f = + iterTargetA_ it 0 sz start oneStride (const f) + {-# INLINE iterFullA_ #-} + + -- | Iterate over a region in a ST monad with access to `Scheduler`. + iterFullAccST + :: Index ix + => it + -- ^ Scheduler multiplying factor. Must be positive + -> Scheduler s a + -- ^ Scheduler to use + -> ix + -- ^ Start index + -> Sz ix + -- ^ Size + -> a + -- ^ Initial accumulator + -> (a -> ST s (a, a)) + -- ^ Function that splits accumulator for each scheduled job. + -> (ix -> a -> ST s a) + -- ^ Action + -> ST s a + iterFullAccST it scheduler start sz acc splitAcc f = + iterTargetAccST it scheduler 0 sz start acc splitAcc (const f) + {-# INLINE iterFullAccST #-} + + iterTargetFullAccST + :: Index ix + => it + -> Scheduler s a + -- ^ Scheduler to use + -> Ix1 + -- ^ Target linear start index + -> Sz ix + -- ^ Target size + -> a + -- ^ Initial accumulator + -> (a -> ST s (a, a)) + -- ^ Function that splits accumulator for each scheduled job. + -> (Ix1 -> ix -> a -> ST s a) + -- ^ Action + -> ST s a + iterTargetFullAccST it scheduler iStart sz = + iterTargetFullWithStrideAccST it scheduler iStart sz oneStride + {-# INLINE iterTargetFullAccST #-} + + iterTargetFullAccST_ + :: Index ix + => it + -> Scheduler s () + -- ^ Scheduler to use + -> Ix1 + -- ^ Target linear start index + -> Sz ix + -- ^ Target size + -> a + -- ^ Initial accumulator + -> (a -> ST s (a, a)) + -- ^ Function that splits accumulator for each scheduled job. + -> (Ix1 -> ix -> a -> ST s a) + -- ^ Action + -> ST s () + iterTargetFullAccST_ it scheduler iStart sz = + iterTargetFullWithStrideAccST_ it scheduler iStart sz oneStride + {-# INLINE iterTargetFullAccST_ #-} + + iterTargetFullST_ + :: Index ix + => it + -> Scheduler s () + -- ^ Scheduler to use + -> Ix1 + -- ^ Target linear start index + -> Sz ix + -- ^ Target size + -> (Ix1 -> ix -> ST s ()) + -- ^ Action + -> ST s () + iterTargetFullST_ it scheduler iStart sz = + iterTargetST_ it scheduler iStart sz (pureIndex 0) + {-# INLINE iterTargetFullST_ #-} + + -- NOTE: this function does not have to be part of the class, but for some + -- reason it creates a severe regression when moved outside. + + -- | Iterate over a target array with a stride without an accumulator + iterTargetWithStrideST_ + :: Index ix + => it + -> Scheduler s () + -- ^ Scheduler to use + -> Ix1 + -- ^ Target linear start index + -> Sz ix + -- ^ Target size + -> ix + -- ^ Start + -> Stride ix + -- ^ Stride + -> (Ix1 -> ix -> ST s a) + -- ^ Action + -> ST s () + iterTargetWithStrideST_ it scheduler i sz ix stride action = + iterTargetWithStrideAccST_ it scheduler i sz ix stride () noSplit $ \j jx _ -> + void $ action j jx + {-# INLINE iterTargetWithStrideST_ #-} + +-- | Default iterator that parallelizes work in linear chunks. Supplied factor +-- will be used to schedule that many jobs per capability. +-- +-- @since 1.0.2 +newtype RowMajor = RowMajorInternal Int + +-- | Default row major iterator with multiplying factor set to @8@. +defRowMajor :: RowMajor +defRowMajor = RowMajorInternal 8 + +pattern RowMajor + :: Int + -- ^ Multiplier that will be used to scale number of jobs. + -> RowMajor +pattern RowMajor f <- RowMajorInternal f + where + RowMajor = RowMajorInternal . max 1 + +{-# COMPLETE RowMajor #-} + +instance Iterator RowMajor where + iterFullM _ start (Sz sz) = iterM start sz (pureIndex 1) (<) + {-# INLINE iterFullM #-} + iterFullA_ _ start (Sz sz) = iterA_ start sz (pureIndex 1) (<) + {-# INLINE iterFullA_ #-} + iterFullAccST (RowMajorInternal fact) scheduler startIx = + iterRowMajorST fact scheduler startIx (pureIndex 1) + {-# INLINE iterFullAccST #-} + iterTargetA_ _ i sz start (Stride stride) = + iterTargetRowMajorA_ 0 i sz start stride + {-# INLINE iterTargetA_ #-} + iterTargetM _ i sz start (Stride stride) = + iterTargetRowMajorAccM 0 i sz start stride + {-# INLINE iterTargetM #-} + iterTargetWithStrideAccST (RowMajor fact) scheduler i sz ix (Stride stride) = + iterTargetRowMajorAccST 0 fact scheduler i sz ix stride + {-# INLINE iterTargetWithStrideAccST #-} + iterTargetWithStrideAccST_ (RowMajor fact) scheduler i sz ix (Stride stride) = + iterTargetRowMajorAccST_ 0 fact scheduler i sz ix stride + {-# INLINE iterTargetWithStrideAccST_ #-} + +newtype RowMajorLinear = RowMajorLinear Int + +defRowMajorLinear :: RowMajorLinear +defRowMajorLinear = RowMajorLinear 8 + +instance Iterator RowMajorLinear where + iterTargetM _ iStart sz start (Stride stride) acc action = + loopM 0 (< totalElem sz) (+ 1) acc $ \i -> + action (iStart + i) (liftIndex2 (+) start (liftIndex2 (*) stride (fromLinearIndex sz i))) + {-# INLINE iterTargetM #-} + iterTargetA_ _ iStart sz start (Stride stride) action = + loopA_ 0 (< totalElem sz) (+ 1) $ \i -> + action (iStart + i) (liftIndex2 (+) start (liftIndex2 (*) stride (fromLinearIndex sz i))) + {-# INLINE iterTargetA_ #-} + iterTargetFullAccST it scheduler iStart sz acc splitAcc action = + let !(RowMajorLinear fact) = it + in iterLinearAccST fact scheduler iStart 1 (totalElem sz) acc splitAcc $ \ !i -> + action i (fromLinearIndex sz i) + {-# INLINE iterTargetFullAccST #-} + iterTargetFullAccST_ it scheduler iStart sz acc splitAcc action = + let !(RowMajorLinear fact) = it + in iterLinearAccST_ fact scheduler iStart 1 (totalElem sz) acc splitAcc $ \ !i -> + action i (fromLinearIndex sz i) + {-# INLINE iterTargetFullAccST_ #-} + iterTargetFullST_ it scheduler iStart sz action = + let !(RowMajorLinear fact) = it + in iterLinearST_ fact scheduler iStart 1 (totalElem sz) $ \ !i -> + action i (fromLinearIndex sz i) + {-# INLINE iterTargetFullST_ #-} + iterTargetWithStrideAccST it scheduler iStart sz start (Stride stride) acc spliAcc action = + let RowMajorLinear fact = it + in iterLinearAccST fact scheduler 0 1 (totalElem sz) acc spliAcc $ \i -> + action (iStart + i) $ + liftIndex2 (+) start (liftIndex2 (*) stride (fromLinearIndex sz i)) + {-# INLINE iterTargetWithStrideAccST #-} + iterTargetWithStrideAccST_ it scheduler iStart sz start (Stride stride) acc spliAcc action = + let RowMajorLinear fact = it + in iterLinearAccST_ fact scheduler 0 1 (totalElem sz) acc spliAcc $ \i -> + action (iStart + i) $ + liftIndex2 (+) start (liftIndex2 (*) stride (fromLinearIndex sz i)) + {-# INLINE iterTargetWithStrideAccST_ #-} + +-- | Parallelizing unbalanced computation (i.e. computing some elements of the +-- array is much more expensive then the others) it can be benefitial to +-- interleave iteration. Perfect example of this would be a ray tracer or the +-- Mandelbrot set. +-- +-- iteration without parallelization is equivalent to `RowMajor` +-- +-- @since 1.0.2 +newtype RowMajorUnbalanced = RowMajorUnbalancedInternal Int + +defRowMajorUnbalanced :: RowMajorUnbalanced +defRowMajorUnbalanced = RowMajorUnbalancedInternal 8 + +pattern RowMajorUnbalanced + :: Int + -- ^ Multiplier that will be used to scale number of jobs. + -> RowMajorUnbalanced +pattern RowMajorUnbalanced f <- RowMajorUnbalancedInternal f + where + RowMajorUnbalanced = RowMajorUnbalancedInternal . max 1 + +{-# COMPLETE RowMajorUnbalanced #-} + +instance Iterator RowMajorUnbalanced where + iterFullM (RowMajorUnbalanced fact) = iterFullM (RowMajor fact) + {-# INLINE iterFullM #-} + iterFullA_ (RowMajorUnbalanced fact) = iterFullA_ (RowMajor fact) + {-# INLINE iterFullA_ #-} + iterTargetM (RowMajorUnbalanced fact) = iterTargetM (RowMajor fact) + {-# INLINE iterTargetM #-} + iterTargetA_ (RowMajorUnbalanced fact) = iterTargetA_ (RowMajor fact) + {-# INLINE iterTargetA_ #-} + iterTargetWithStrideAccST = iterUnbalancedTargetWithStride loopM + {-# INLINE iterTargetWithStrideAccST #-} + iterTargetWithStrideAccST_ it scheduler iStart sz start stride acc splitAcc' action = + void $ + iterUnbalancedTargetWithStride innerLoop it scheduler iStart sz start stride acc splitAcc' action + where + innerLoop initial condition increment initAcc f = + void $ loopM initial condition increment initAcc f + {-# INLINE innerLoop #-} + {-# INLINE iterTargetWithStrideAccST_ #-} + +iterUnbalancedTargetWithStride + :: Index ix + => (Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> t) -> ST s b) + -> RowMajorUnbalanced + -> Scheduler s b + -> Int + -> Sz ix + -> ix + -> Stride ix + -> a + -> (a -> ST s (a, a)) + -> (Int -> ix -> t) + -> ST s a +iterUnbalancedTargetWithStride innerLoop it scheduler iStart sz start stride acc splitAcc action = + let RowMajorUnbalanced fact = it + !n = totalElem sz + !step = min (fact * numWorkers scheduler) n + in loopM 0 (< step) (+ 1) acc $ \ !istep !a -> do + (curAcc, nextAcc) <- splitAcc a + scheduleMassivWork scheduler $ + innerLoop istep (< n) (+ step) curAcc $ \i -> + action (iStart + i) $ + liftIndex2 (+) start (liftIndex2 (*) (unStride stride) (fromLinearIndex sz i)) + pure nextAcc +{-# INLINE iterUnbalancedTargetWithStride #-} + +noSplit :: Applicative m => () -> m ((), ()) +noSplit _ = pure ((), ()) + +iterTargetAccST + :: (Iterator it, Index ix) + => it + -> Scheduler s a + -- ^ Scheduler to use + -> Ix1 + -- ^ Target linear start index + -> Sz ix + -- ^ Target size + -> ix + -- ^ Source start + -> a + -> (a -> ST s (a, a)) + -> (Ix1 -> ix -> a -> ST s a) + -- ^ Action + -> ST s a +iterTargetAccST it scheduler iStart sz ix = + iterTargetWithStrideAccST it scheduler iStart sz ix oneStride +{-# INLINE iterTargetAccST #-} + +iterTargetAccST_ + :: (Iterator it, Index ix) + => it + -> Scheduler s () + -- ^ Scheduler to use + -> Ix1 + -- ^ Target linear start index + -> Sz ix + -- ^ Target size + -> ix + -- ^ Source start + -> a + -> (a -> ST s (a, a)) + -> (Ix1 -> ix -> a -> ST s a) + -- ^ Action + -> ST s () +iterTargetAccST_ it scheduler iStart sz ix = + iterTargetWithStrideAccST_ it scheduler iStart sz ix oneStride +{-# INLINE iterTargetAccST_ #-} + +iterTargetFullWithStrideST_ + :: (Iterator it, Index ix) + => it + -> Scheduler s () + -- ^ Scheduler to use + -> Ix1 + -- ^ Target linear start index + -> Sz ix + -- ^ Target size + -> Stride ix + -- ^ Stride + -> (Ix1 -> ix -> ST s ()) + -- ^ Action + -> ST s () +iterTargetFullWithStrideST_ it scheduler iStart sz = + iterTargetWithStrideST_ it scheduler iStart sz (pureIndex 0) +{-# INLINE iterTargetFullWithStrideST_ #-} + +iterTargetST_ + :: (Iterator it, Index ix) + => it + -> Scheduler s () + -- ^ Scheduler to use + -> Ix1 + -- ^ Target linear start index + -> Sz ix + -- ^ Target size + -> ix + -- ^ Start + -> (Ix1 -> ix -> ST s ()) + -- ^ Action + -> ST s () +iterTargetST_ it scheduler iStart sz ix = + iterTargetWithStrideST_ it scheduler iStart sz ix oneStride +{-# INLINE iterTargetST_ #-} + +iterTargetFullWithStrideAccST + :: (Iterator it, Index ix) + => it + -> Scheduler s a + -- ^ Scheduler to use + -> Ix1 + -- ^ Target linear start index + -> Sz ix + -- ^ Target size + -> Stride ix + -- ^ Stride + -> a + -> (a -> ST s (a, a)) + -> (Ix1 -> ix -> a -> ST s a) + -- ^ Action + -> ST s a +iterTargetFullWithStrideAccST it scheduler iStart sz = + iterTargetWithStrideAccST it scheduler iStart sz (pureIndex 0) +{-# INLINE iterTargetFullWithStrideAccST #-} + +iterTargetFullWithStrideAccST_ + :: (Iterator it, Index ix) + => it + -> Scheduler s () + -- ^ Scheduler to use + -> Ix1 + -- ^ Target linear start index + -> Sz ix + -- ^ Target size + -> Stride ix + -- ^ Stride + -> a + -> (a -> ST s (a, a)) + -> (Ix1 -> ix -> a -> ST s a) + -- ^ Action + -> ST s () +iterTargetFullWithStrideAccST_ it scheduler iStart sz = + iterTargetWithStrideAccST_ it scheduler iStart sz (pureIndex 0) +{-# INLINE iterTargetFullWithStrideAccST_ #-} diff --git a/src/Data/Massiv/Core/Index/Ix.hs b/src/Data/Massiv/Core/Index/Ix.hs --- a/src/Data/Massiv/Core/Index/Ix.hs +++ b/src/Data/Massiv/Core/Index/Ix.hs @@ -1,64 +1,62 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE CPP #-} -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilyDependencies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} + -- | -- Module : Data.Massiv.Core.Index.Ix --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Core.Index.Ix - ( Ix - , IxN((:>)) - , type Sz - , pattern Sz - , type Ix1 - , pattern Ix1 - , type Sz1 - , pattern Sz1 - , type Ix2(Ix2, (:.)) - , type Sz2 - , pattern Sz2 - , type Ix3 - , pattern Ix3 - , type Sz3 - , pattern Sz3 - , type Ix4 - , pattern Ix4 - , type Sz4 - , pattern Sz4 - , type Ix5 - , pattern Ix5 - , type Sz5 - , pattern Sz5 - ) where - -import Control.Monad.Catch (MonadThrow(..)) -import Control.DeepSeq -import Control.Monad (liftM) -import Data.Massiv.Core.Index.Internal -import Data.Monoid ((<>)) -import Data.Proxy -import qualified Data.Vector.Generic as V -import qualified Data.Vector.Generic.Mutable as VM -import qualified Data.Vector.Unboxed as VU -import GHC.TypeLits +module Data.Massiv.Core.Index.Ix ( + Ix, + IxN ((:>)), + type Sz, + pattern Sz, + type Ix1, + pattern Ix1, + pattern Sz1, + type Ix2 (Ix2, (:.)), + pattern Sz2, + type Ix3, + pattern Ix3, + pattern Sz3, + type Ix4, + pattern Ix4, + pattern Sz4, + type Ix5, + pattern Ix5, + pattern Sz5, + HighIxN, +) where +import Control.DeepSeq +import Control.Monad.Catch (MonadThrow (..)) +import Data.Massiv.Core.Index.Internal +import Data.Massiv.Core.Loop +import Data.Proxy +import qualified Data.Vector.Generic as V +import qualified Data.Vector.Generic.Mutable as VM +import qualified Data.Vector.Unboxed as VU +import qualified GHC.Arr as I +import GHC.TypeLits +import System.Random.Stateful +#if !MIN_VERSION_base(4,11,0) +import Data.Semigroup +#endif infixr 5 :>, :. - -- | 2-dimensional index. This is also a base index for higher dimensions. -- -- @since 0.1.0 @@ -69,18 +67,15 @@ -- @since 0.1.0 pattern Ix2 :: Int -> Int -> Ix2 pattern Ix2 i2 i1 = i2 :. i1 -{-# COMPLETE Ix2 #-} --- | 2-dimensional size type synonym. --- --- @since 0.3.0 -type Sz2 = Sz Ix2 +{-# COMPLETE Ix2 #-} -- | 2-dimensional size constructor. @(Sz2 i j) == Sz (i :. j)@ -- -- @since 0.3.0 -pattern Sz2 :: Int -> Int -> Sz2 +pattern Sz2 :: Int -> Int -> Sz Ix2 pattern Sz2 i2 i1 = Sz (i2 :. i1) + {-# COMPLETE Sz2 #-} -- | 3-dimensional type synonym. Useful as a alternative to enabling @DataKinds@ and using type @@ -94,18 +89,15 @@ -- @since 0.1.0 pattern Ix3 :: Int -> Int -> Int -> Ix3 pattern Ix3 i3 i2 i1 = i3 :> i2 :. i1 -{-# COMPLETE Ix3 #-} --- | 3-dimensional size type synonym. --- --- @since 0.3.0 -type Sz3 = Sz Ix3 +{-# COMPLETE Ix3 #-} -- | 3-dimensional size constructor. @(Sz3 i j k) == Sz (i :> j :. k)@ -- -- @since 0.3.0 -pattern Sz3 :: Int -> Int -> Int -> Sz3 +pattern Sz3 :: Int -> Int -> Int -> Sz Ix3 pattern Sz3 i3 i2 i1 = Sz (i3 :> i2 :. i1) + {-# COMPLETE Sz3 #-} -- | 4-dimensional type synonym. @@ -118,18 +110,15 @@ -- @since 0.1.0 pattern Ix4 :: Int -> Int -> Int -> Int -> Ix4 pattern Ix4 i4 i3 i2 i1 = i4 :> i3 :> i2 :. i1 -{-# COMPLETE Ix4 #-} --- | 4-dimensional size type synonym. --- --- @since 0.3.0 -type Sz4 = Sz Ix4 +{-# COMPLETE Ix4 #-} -- | 4-dimensional size constructor. @(Sz4 i j k l) == Sz (i :> j :> k :. l)@ -- -- @since 0.3.0 -pattern Sz4 :: Int -> Int -> Int -> Int -> Sz4 +pattern Sz4 :: Int -> Int -> Int -> Int -> Sz Ix4 pattern Sz4 i4 i3 i2 i1 = Sz (i4 :> i3 :> i2 :. i1) + {-# COMPLETE Sz4 #-} -- | 5-dimensional type synonym. @@ -142,18 +131,15 @@ -- @since 0.1.0 pattern Ix5 :: Int -> Int -> Int -> Int -> Int -> Ix5 pattern Ix5 i5 i4 i3 i2 i1 = i5 :> i4 :> i3 :> i2 :. i1 -{-# COMPLETE Ix5 #-} --- | 5-dimensional size type synonym. --- --- @since 0.3.0 -type Sz5 = Sz Ix5 +{-# COMPLETE Ix5 #-} -- | 5-dimensional size constructor. @(Sz5 i j k l m) == Sz (i :> j :> k :> l :. m)@ -- -- @since 0.3.0 -pattern Sz5 :: Int -> Int -> Int -> Int -> Int -> Sz5 +pattern Sz5 :: Int -> Int -> Int -> Int -> Int -> Sz Ix5 pattern Sz5 i5 i4 i3 i2 i1 = Sz (i5 :> i4 :> i3 :> i2 :. i1) + {-# COMPLETE Sz5 #-} -- | n-dimensional index. Needs a base case, which is the `Ix2`. @@ -170,10 +156,9 @@ Ix 2 = Ix2 Ix n = IxN n - type instance Lower Ix2 = Ix1 -type instance Lower (IxN n) = Ix (n - 1) +type instance Lower (IxN n) = Ix (n - 1) instance Show Ix2 where showsPrec n (i :. j) = showsPrecWrapped n (shows i . (" :. " ++) . shows j) @@ -181,7 +166,64 @@ instance Show (Ix (n - 1)) => Show (IxN n) where showsPrec n (i :> ix) = showsPrecWrapped n (shows i . (" :> " ++) . shows ix) +instance Uniform Ix2 where + uniformM g = (:.) <$> uniformM g <*> uniformM g + {-# INLINE uniformM #-} +instance UniformRange Ix2 where + uniformRM (l1 :. l2, u1 :. u2) g = (:.) <$> uniformRM (l1, u1) g <*> uniformRM (l2, u2) g + {-# INLINE uniformRM #-} +#if MIN_VERSION_random(1,3,0) + isInRange (l1 :. l2, u1 :. u2) (i1 :. i2) = + isInRangeOrd (l1, u1) i1 && isInRangeOrd (l2, u2) i2 +#endif + +instance Random Ix2 + +instance Uniform (Ix (n - 1)) => Uniform (IxN n) where + uniformM g = (:>) <$> uniformM g <*> uniformM g + {-# INLINE uniformM #-} + +instance UniformRange (Ix (n - 1)) => UniformRange (IxN n) where + uniformRM (l1 :> l2, u1 :> u2) g = (:>) <$> uniformRM (l1, u1) g <*> uniformRM (l2, u2) g + {-# INLINE uniformRM #-} +#if MIN_VERSION_random(1,3,0) + isInRange (l1 :> l2, u1 :> u2) (i1 :> i2) = + isInRangeOrd (l1, u1) i1 && isInRange (l2, u2) i2 +#endif + +instance Random (Ix (n - 1)) => Random (IxN n) where + random g = + case random g of + (i, g') -> + case random g' of + (n, g'') -> (i :> n, g'') + {-# INLINE random #-} + randomR (l1 :> l2, u1 :> u2) g = + case randomR (l1, u1) g of + (i, g') -> + case randomR (l2, u2) g' of + (n, g'') -> (i :> n, g'') + {-# INLINE randomR #-} + +instance I.Ix Ix2 where + range (i1 :. j1, i2 :. j2) = [i :. j | i <- [i1 .. i2], j <- [j1 .. j2]] + {-# INLINE range #-} + unsafeIndex (l1 :. l2, u1 :. u2) (i1 :. i2) = + I.unsafeIndex (l1, u1) i1 * I.unsafeRangeSize (l2, u2) + I.unsafeIndex (l2, u2) i2 + {-# INLINE unsafeIndex #-} + inRange (l1 :. l2, u1 :. u2) (i1 :. i2) = I.inRange (l1, u1) i1 && I.inRange (l2, u2) i2 + {-# INLINE inRange #-} + +instance I.Ix (Ix (n - 1)) => I.Ix (IxN n) where + range (i1 :> j1, i2 :> j2) = [i :> j | i <- [i1 .. i2], j <- I.range (j1, j2)] + {-# INLINE range #-} + unsafeIndex (l1 :> l2, u1 :> u2) (i1 :> i2) = + I.unsafeIndex (l1, u1) i1 * I.unsafeRangeSize (l2, u2) + I.unsafeIndex (l2, u2) i2 + {-# INLINE unsafeIndex #-} + inRange (l1 :> l2, u1 :> u2) (i1 :> i2) = I.inRange (l1, u1) i1 && I.inRange (l2, u2) i2 + {-# INLINE inRange #-} + instance Num Ix2 where (+) = liftIndex2 (+) {-# INLINE [1] (+) #-} @@ -214,15 +256,7 @@ fromInteger = pureIndex . fromInteger {-# INLINE [1] fromInteger #-} - -instance {-# OVERLAPPABLE #-} ( 1 <= n - , 4 <= n - , KnownNat n - , KnownNat (n - 1) - , Index (Ix (n - 1)) - , IxN (n - 1) ~ Ix (n - 1) - ) => - Num (IxN n) where +instance {-# OVERLAPPABLE #-} HighIxN n => Num (IxN n) where (+) = liftIndex2 (+) {-# INLINE [1] (+) #-} (-) = liftIndex2 (-) @@ -238,8 +272,6 @@ fromInteger = pureIndex . fromInteger {-# INLINE [1] fromInteger #-} - - instance Bounded Ix2 where minBound = pureIndex minBound {-# INLINE minBound #-} @@ -252,14 +284,7 @@ maxBound = pureIndex maxBound {-# INLINE maxBound #-} -instance {-# OVERLAPPABLE #-} ( 1 <= n - , 4 <= n - , KnownNat n - , KnownNat (n - 1) - , Index (Ix (n - 1)) - , IxN (n - 1) ~ Ix (n - 1) - ) => - Bounded (IxN n) where +instance {-# OVERLAPPABLE #-} HighIxN n => Bounded (IxN n) where minBound = pureIndex minBound {-# INLINE minBound #-} maxBound = pureIndex maxBound @@ -271,21 +296,18 @@ instance NFData (IxN n) where rnf ix = ix `seq` () - instance Eq Ix2 where - (i1 :. j1) == (i2 :. j2) = i1 == i2 && j1 == j2 + (i1 :. j1) == (i2 :. j2) = i1 == i2 && j1 == j2 instance Eq (Ix (n - 1)) => Eq (IxN n) where (i1 :> ix1) == (i2 :> ix2) = i1 == i2 && ix1 == ix2 - instance Ord Ix2 where compare (i1 :. j1) (i2 :. j2) = compare i1 i2 <> compare j1 j2 instance Ord (Ix (n - 1)) => Ord (IxN n) where compare (i1 :> ix1) (i2 :> ix2) = compare i1 i2 <> compare ix1 ix2 - instance Index Ix2 where type Dimensions Ix2 = 2 dimensions _ = 2 @@ -296,8 +318,9 @@ {-# INLINE [1] isSafeIndex #-} toLinearIndex (SafeSz (_ :. k1)) (i2 :. i1) = k1 * i2 + i1 {-# INLINE [1] toLinearIndex #-} - fromLinearIndex (SafeSz (_ :. k1)) i = case i `quotRem` k1 of - (i2, i1) -> i2 :. i1 + fromLinearIndex (SafeSz (_ :. k1)) i = + case i `quotRem` k1 of + (i2, i1) -> i2 :. i1 {-# INLINE [1] fromLinearIndex #-} consDim = (:.) {-# INLINE [1] consDim #-} @@ -307,21 +330,21 @@ {-# INLINE [1] snocDim #-} unsnocDim (i2 :. i1) = (i2, i1) {-# INLINE [1] unsnocDim #-} - getDimM (i2 :. _) 2 = pure i2 - getDimM ( _ :. i1) 1 = pure i1 - getDimM ix d = throwM $ IndexDimensionException ix d + getDimM (i2 :. _) 2 = pure i2 + getDimM (_ :. i1) 1 = pure i1 + getDimM ix d = throwM $ IndexDimensionException ix d {-# INLINE [1] getDimM #-} - setDimM ( _ :. i1) 2 i2 = pure (i2 :. i1) - setDimM (i2 :. _) 1 i1 = pure (i2 :. i1) - setDimM ix d _ = throwM $ IndexDimensionException ix d + setDimM (_ :. i1) 2 i2 = pure (i2 :. i1) + setDimM (i2 :. _) 1 i1 = pure (i2 :. i1) + setDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] setDimM #-} pullOutDimM (i2 :. i1) 2 = pure (i2, i1) pullOutDimM (i2 :. i1) 1 = pure (i1, i2) - pullOutDimM ix d = throwM $ IndexDimensionException ix d + pullOutDimM ix d = throwM $ IndexDimensionException ix d {-# INLINE [1] pullOutDimM #-} insertDimM i1 2 i2 = pure (i2 :. i1) insertDimM i2 1 i1 = pure (i2 :. i1) - insertDimM ix d _ = throwM $ IndexDimensionException ix d + insertDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] insertDimM #-} pureIndex i = i :. i {-# INLINE [1] pureIndex #-} @@ -332,7 +355,10 @@ repairIndex (SafeSz (k :. szL)) (i :. ixL) rBelow rOver = repairIndex (SafeSz k) i rBelow rOver :. repairIndex (SafeSz szL) ixL rBelow rOver {-# INLINE [1] repairIndex #-} - + iterF (s :. sIxL) (e :. eIxL) (inc :. incIxL) cond initAct f = + loopF s (`cond` e) (+ inc) initAct $ \ !i g -> + loopF sIxL (`cond` eIxL) (+ incIxL) g $ \ !j -> f (i :. j) + {-# INLINE iterF #-} instance {-# OVERLAPPING #-} Index (IxN 3) where type Dimensions Ix3 = 3 @@ -343,10 +369,15 @@ isSafeIndex (SafeSz (k3 :> k2 :. k1)) (i3 :> i2 :. i1) = 0 <= i3 && 0 <= i2 && 0 <= i1 && i3 < k3 && i2 < k2 && i1 < k1 {-# INLINE [1] isSafeIndex #-} - toLinearIndex (SafeSz (_ :> k2 :. k1)) (i3 :> i2 :. i1) = (k2 * i3 + i2) * k1 + i1 + toLinearIndex (SafeSz (_ :> k2 :. k1)) (i3 :> i2 :. i1) = (i3 * k2 + i2) * k1 + i1 {-# INLINE [1] toLinearIndex #-} fromLinearIndex (SafeSz (_ :> ix)) i = let !(q, ixL) = fromLinearIndexAcc ix i in q :> ixL - {-# INLINE [1] fromLinearIndex #-} + {-# INLINE fromLinearIndex #-} + fromLinearIndexAcc (m :> ix) !k = (q, r :> ixL) + where + !(!kL, !ixL) = fromLinearIndexAcc ix k + !(!q, !r) = quotRem kL m + {-# INLINE fromLinearIndexAcc #-} consDim = (:>) {-# INLINE [1] consDim #-} unconsDim (i3 :> ix) = (i3, ix) @@ -355,25 +386,25 @@ {-# INLINE [1] snocDim #-} unsnocDim (i3 :> i2 :. i1) = (i3 :. i2, i1) {-# INLINE [1] unsnocDim #-} - getDimM (i3 :> _ :. _) 3 = pure i3 - getDimM ( _ :> i2 :. _) 2 = pure i2 - getDimM ( _ :> _ :. i1) 1 = pure i1 - getDimM ix d = throwM $ IndexDimensionException ix d + getDimM (i3 :> _ :. _) 3 = pure i3 + getDimM (_ :> i2 :. _) 2 = pure i2 + getDimM (_ :> _ :. i1) 1 = pure i1 + getDimM ix d = throwM $ IndexDimensionException ix d {-# INLINE [1] getDimM #-} - setDimM ( _ :> i2 :. i1) 3 i3 = pure (i3 :> i2 :. i1) - setDimM (i3 :> _ :. i1) 2 i2 = pure (i3 :> i2 :. i1) - setDimM (i3 :> i2 :. _) 1 i1 = pure (i3 :> i2 :. i1) - setDimM ix d _ = throwM $ IndexDimensionException ix d + setDimM (_ :> i2 :. i1) 3 i3 = pure (i3 :> i2 :. i1) + setDimM (i3 :> _ :. i1) 2 i2 = pure (i3 :> i2 :. i1) + setDimM (i3 :> i2 :. _) 1 i1 = pure (i3 :> i2 :. i1) + setDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] setDimM #-} pullOutDimM (i3 :> i2 :. i1) 3 = pure (i3, i2 :. i1) pullOutDimM (i3 :> i2 :. i1) 2 = pure (i2, i3 :. i1) pullOutDimM (i3 :> i2 :. i1) 1 = pure (i1, i3 :. i2) - pullOutDimM ix d = throwM $ IndexDimensionException ix d + pullOutDimM ix d = throwM $ IndexDimensionException ix d {-# INLINE [1] pullOutDimM #-} insertDimM (i2 :. i1) 3 i3 = pure (i3 :> i2 :. i1) insertDimM (i3 :. i1) 2 i2 = pure (i3 :> i2 :. i1) insertDimM (i3 :. i2) 1 i1 = pure (i3 :> i2 :. i1) - insertDimM ix d _ = throwM $ IndexDimensionException ix d + insertDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] insertDimM #-} pureIndex i = i :> i :. i {-# INLINE [1] pureIndex #-} @@ -384,15 +415,35 @@ repairIndex (SafeSz (n :> szL)) (i :> ixL) rBelow rOver = repairIndex (SafeSz n) i rBelow rOver :> repairIndex (SafeSz szL) ixL rBelow rOver {-# INLINE [1] repairIndex #-} + iterTargetRowMajorAccM iAcc iStart sz (b3 :> b2 :. b1) (s3 :> s2 :. s1) initAcc action = + let n = totalElem sz + iShift = iStart + iAcc * n + in loopM 0 (< n) (+ 1) initAcc $ \ !i !acc -> + let (i3 :> i2 :. i1) = fromLinearIndex sz i + in action (iShift + i) ((b3 + s3 * i3) :> (b2 + s2 * i2) :. (b1 + s1 * i1)) acc + {-# INLINE iterTargetRowMajorAccM #-} + iterTargetRowMajorAccST_ iAcc fact scheduler iStart sz (b3 :> b2 :. b1) (s3 :> s2 :. s1) acc splitAcc action = + let n = totalElem sz + iShift = iStart + iAcc * n + in iterLinearAccST_ fact scheduler 0 1 n acc splitAcc $ \ !i -> + let (i3 :> i2 :. i1) = fromLinearIndex sz i + in action (iShift + i) ((b3 + s3 * i3) :> (b2 + s2 * i2) :. (b1 + s1 * i1)) + {-# INLINE iterTargetRowMajorAccST_ #-} + iterTargetRowMajorAccST iAcc fact scheduler iStart sz (b3 :> b2 :. b1) (s3 :> s2 :. s1) acc splitAcc action = + let n = totalElem sz + iShift = iStart + iAcc * n + in iterLinearAccST fact scheduler 0 1 n acc splitAcc $ \ !i -> + let (i3 :> i2 :. i1) = fromLinearIndex sz i + in action (iShift + i) ((b3 + s3 * i3) :> (b2 + s2 * i2) :. (b1 + s1 * i1)) + {-# INLINE iterTargetRowMajorAccST #-} -instance {-# OVERLAPPABLE #-} ( 1 <= n - , 4 <= n - , KnownNat n - , KnownNat (n - 1) - , Index (Ix (n - 1)) - , IxN (n - 1) ~ Ix (n - 1) - ) => - Index (IxN n) where +-- | Constraint synonym that encapsulates all constraints needed for dimension 4 and higher. +-- +-- @since 1.0.0 +type HighIxN n = + (4 <= n, KnownNat n, KnownNat (n - 1), Index (IxN (n - 1)), IxN (n - 1) ~ Ix (n - 1)) + +instance {-# OVERLAPPABLE #-} HighIxN n => Index (IxN n) where type Dimensions (IxN n) = n dimensions _ = fromInteger $ natVal (Proxy :: Proxy n) {-# INLINE [1] dimensions #-} @@ -419,12 +470,12 @@ pullOutDimM ix@(i :> ixl) d | d == dimensions (Proxy :: Proxy (IxN n)) = pure (i, ixl) | otherwise = - maybe (throwM $ IndexDimensionException ix d) (pure . fmap (i :>)) (pullOutDimM ixl d) + maybe (throwM $ IndexDimensionException ix d) (pure . fmap (i :>)) (pullOutDimM ixl d) {-# INLINE [1] pullOutDimM #-} insertDimM ix@(i :> ixl) d di | d == dimensions (Proxy :: Proxy (IxN n)) = pure (di :> ix) | otherwise = - maybe (throwM $ IndexDimensionException ix d) (pure . (i :>)) (insertDimM ixl d di) + maybe (throwM $ IndexDimensionException ix d) (pure . (i :>)) (insertDimM ixl d di) {-# INLINE [1] insertDimM #-} pureIndex i = i :> (pureIndex i :: Ix (n - 1)) {-# INLINE [1] pureIndex #-} @@ -436,8 +487,6 @@ repairIndex (SafeSz n) i rBelow rOver :> repairIndex (SafeSz szL) ixL rBelow rOver {-# INLINE [1] repairIndex #-} - - ---- Unbox Ix -- | Unboxing of a `Ix2`. @@ -452,11 +501,11 @@ {-# INLINE basicUnsafeSlice #-} basicOverlaps (MV_Ix2 mvec) (MV_Ix2 mvec') = VM.basicOverlaps mvec mvec' {-# INLINE basicOverlaps #-} - basicUnsafeNew len = MV_Ix2 `liftM` VM.basicUnsafeNew len + basicUnsafeNew len = MV_Ix2 <$> VM.basicUnsafeNew len {-# INLINE basicUnsafeNew #-} - basicUnsafeReplicate len (i :. j) = MV_Ix2 `liftM` VM.basicUnsafeReplicate len (i, j) + basicUnsafeReplicate len (i :. j) = MV_Ix2 <$> VM.basicUnsafeReplicate len (i, j) {-# INLINE basicUnsafeReplicate #-} - basicUnsafeRead (MV_Ix2 mvec) idx = uncurry (:.) `liftM` VM.basicUnsafeRead mvec idx + basicUnsafeRead (MV_Ix2 mvec) idx = uncurry (:.) <$> VM.basicUnsafeRead mvec idx {-# INLINE basicUnsafeRead #-} basicUnsafeWrite (MV_Ix2 mvec) idx (i :. j) = VM.basicUnsafeWrite mvec idx (i, j) {-# INLINE basicUnsafeWrite #-} @@ -468,42 +517,37 @@ {-# INLINE basicUnsafeCopy #-} basicUnsafeMove (MV_Ix2 mvec) (MV_Ix2 mvec') = VM.basicUnsafeMove mvec mvec' {-# INLINE basicUnsafeMove #-} - basicUnsafeGrow (MV_Ix2 mvec) len = MV_Ix2 `liftM` VM.basicUnsafeGrow mvec len + basicUnsafeGrow (MV_Ix2 mvec) len = MV_Ix2 <$> VM.basicUnsafeGrow mvec len {-# INLINE basicUnsafeGrow #-} #if MIN_VERSION_vector(0,11,0) basicInitialize (MV_Ix2 mvec) = VM.basicInitialize mvec {-# INLINE basicInitialize #-} #endif - newtype instance VU.Vector Ix2 = V_Ix2 (VU.Vector (Int, Int)) instance V.Vector VU.Vector Ix2 where - basicUnsafeFreeze (MV_Ix2 mvec) = V_Ix2 `liftM` V.basicUnsafeFreeze mvec + basicUnsafeFreeze (MV_Ix2 mvec) = V_Ix2 <$> V.basicUnsafeFreeze mvec {-# INLINE basicUnsafeFreeze #-} - basicUnsafeThaw (V_Ix2 vec) = MV_Ix2 `liftM` V.basicUnsafeThaw vec + basicUnsafeThaw (V_Ix2 vec) = MV_Ix2 <$> V.basicUnsafeThaw vec {-# INLINE basicUnsafeThaw #-} basicLength (V_Ix2 vec) = V.basicLength vec {-# INLINE basicLength #-} basicUnsafeSlice idx len (V_Ix2 vec) = V_Ix2 (V.basicUnsafeSlice idx len vec) {-# INLINE basicUnsafeSlice #-} - basicUnsafeIndexM (V_Ix2 vec) idx = uncurry (:.) `liftM` V.basicUnsafeIndexM vec idx + basicUnsafeIndexM (V_Ix2 vec) idx = uncurry (:.) <$> V.basicUnsafeIndexM vec idx {-# INLINE basicUnsafeIndexM #-} basicUnsafeCopy (MV_Ix2 mvec) (V_Ix2 vec) = V.basicUnsafeCopy mvec vec {-# INLINE basicUnsafeCopy #-} elemseq _ = seq {-# INLINE elemseq #-} - - ---- Unbox Ix - - -- | Unboxing of a `IxN`. instance (3 <= n, VU.Unbox (Ix (n - 1))) => VU.Unbox (IxN n) -newtype instance VU.MVector s (IxN n) = MV_IxN (VU.MVector s Int, VU.MVector s (Ix (n-1))) +newtype instance VU.MVector s (IxN n) = MV_IxN (VU.MVector s Int, VU.MVector s (Ix (n - 1))) instance (3 <= n, VU.Unbox (Ix (n - 1))) => VM.MVector VU.MVector (IxN n) where basicLength (MV_IxN (_, mvec)) = VM.basicLength mvec @@ -554,8 +598,7 @@ {-# INLINE basicInitialize #-} #endif - -newtype instance VU.Vector (IxN n) = V_IxN (VU.Vector Int, VU.Vector (Ix (n-1))) +newtype instance VU.Vector (IxN n) = V_IxN (VU.Vector Int, VU.Vector (Ix (n - 1))) instance (3 <= n, VU.Unbox (Ix (n - 1))) => V.Vector VU.Vector (IxN n) where basicUnsafeFreeze (MV_IxN (mvec1, mvec)) = do diff --git a/src/Data/Massiv/Core/Index/Stride.hs b/src/Data/Massiv/Core/Index/Stride.hs --- a/src/Data/Massiv/Core/Index/Stride.hs +++ b/src/Data/Massiv/Core/Index/Stride.hs @@ -2,29 +2,26 @@ {-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE PatternSynonyms #-} -#if __GLASGOW_HASKELL__ < 820 -{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-} -#endif -- | -- Module : Data.Massiv.Core.Index.Stride --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Core.Index.Stride - ( Stride(SafeStride) - , pattern Stride - , unStride - , oneStride - , toLinearIndexStride - , strideStart - , strideSize - ) where +module Data.Massiv.Core.Index.Stride ( + Stride (SafeStride), + pattern Stride, + unStride, + oneStride, + toLinearIndexStride, + strideStart, + strideSize, +) where -import Control.DeepSeq +import Control.DeepSeq (NFData) import Data.Massiv.Core.Index.Internal +import System.Random.Stateful (Random, Uniform (..), UniformRange (..)) -- | Stride provides a way to ignore elements of an array if an index is divisible by a -- corresponding value in a stride. So, for a @Stride (i :. j)@ only elements with indices will be @@ -54,21 +51,33 @@ -- @since 0.2.1 newtype Stride ix = SafeStride ix deriving (Eq, Ord, NFData) - -- | A safe bidirectional pattern synonym for `Stride` construction that will make sure stride -- elements are always positive. -- -- @since 0.2.1 pattern Stride :: Index ix => ix -> Stride ix -pattern Stride ix <- SafeStride ix where - Stride ix = SafeStride (liftIndex (max 1) ix) -{-# COMPLETE Stride #-} +pattern Stride ix <- SafeStride ix + where + Stride ix = SafeStride (liftIndex (max 1) ix) +{-# COMPLETE Stride #-} instance Index ix => Show (Stride ix) where showsPrec n (SafeStride ix) = showsPrecWrapped n (("Stride " ++) . showsPrec 1 ix) +instance (UniformRange ix, Index ix) => Uniform (Stride ix) where + uniformM g = SafeStride <$> uniformRM (pureIndex 1, pureIndex maxBound) g + {-# INLINE uniformM #-} +instance UniformRange ix => UniformRange (Stride ix) where + uniformRM (SafeStride l, SafeStride u) g = SafeStride <$> uniformRM (l, u) g + {-# INLINE uniformRM #-} +#if MIN_VERSION_random(1,3,0) + isInRange (SafeStride l, SafeStride u) (SafeStride k) = isInRange (l, u) k +#endif + +instance (UniformRange ix, Index ix) => Random (Stride ix) + -- | Just a helper function for unwrapping `Stride`. -- -- @since 0.2.1 @@ -98,21 +107,21 @@ -- | Compute linear index with stride using the original size and index -- -- @since 0.2.1 -toLinearIndexStride :: - Index ix - => Stride ix -- ^ Stride - -> Sz ix -- ^ Size - -> ix -- ^ Index +toLinearIndexStride + :: Index ix + => Stride ix + -- ^ Stride + -> Sz ix + -- ^ Size + -> ix + -- ^ Index -> Int toLinearIndexStride (SafeStride stride) sz ix = toLinearIndex sz (liftIndex2 div ix stride) {-# INLINE toLinearIndexStride #-} - -- | A default stride of @1@, where all elements are kept -- -- @since 0.2.1 oneStride :: Index ix => Stride ix oneStride = SafeStride (pureIndex 1) {-# INLINE oneStride #-} - - diff --git a/src/Data/Massiv/Core/Index/Tuple.hs b/src/Data/Massiv/Core/Index/Tuple.hs --- a/src/Data/Massiv/Core/Index/Tuple.hs +++ b/src/Data/Massiv/Core/Index/Tuple.hs @@ -1,42 +1,51 @@ -{-# OPTIONS_GHC -fno-warn-orphans #-} {-# LANGUAGE BangPatterns #-} {-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE TypeFamilies #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} + -- | -- Module : Data.Massiv.Core.Index.Tuple --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Core.Index.Tuple - ( -- * Tuple based indices - -- ** 1-dimensional - Ix1T - -- ** 2-dimensional - , Ix2T - , toIx2 - , fromIx2 - -- ** 3-dimensional - , Ix3T - , toIx3 - , fromIx3 - -- ** 4-dimensional - , Ix4T - , toIx4 - , fromIx4 - -- ** 5-dimensional - , Ix5T - , toIx5 - , fromIx5 - ) where +module Data.Massiv.Core.Index.Tuple ( + -- * Tuple based indices -import Control.Monad.Catch (MonadThrow(..)) -import Data.Massiv.Core.Index.Internal (Index(..), IndexException(..), Lower, - Sz(..)) + -- ** 1-dimensional + Ix1T, + + -- ** 2-dimensional + Ix2T, + toIx2, + fromIx2, + + -- ** 3-dimensional + Ix3T, + toIx3, + fromIx3, + + -- ** 4-dimensional + Ix4T, + toIx4, + fromIx4, + + -- ** 5-dimensional + Ix5T, + toIx5, + fromIx5, +) where + +import Control.Monad.Catch (MonadThrow (..)) +import Data.Massiv.Core.Index.Internal ( + Index (..), + IndexException (..), + Lower, + Sz (..), + ) import Data.Massiv.Core.Index.Ix -- | Another 1-dimensional index type synonym for `Int`, same as `Ix1` and is here just for @@ -56,11 +65,12 @@ type Ix5T = (Int, Int, Int, Int, Int) type instance Lower Ix2T = Ix1T + type instance Lower Ix3T = Ix2T -type instance Lower Ix4T = Ix3T -type instance Lower Ix5T = Ix4T +type instance Lower Ix4T = Ix3T +type instance Lower Ix5T = Ix4T -- | Convert an `Int` tuple to `Ix2` -- @@ -178,39 +188,38 @@ {-# INLINE [1] snocDim #-} unsnocDim = id {-# INLINE [1] unsnocDim #-} - getDimM (i2, _) 2 = pure i2 - getDimM ( _, i1) 1 = pure i1 - getDimM ix d = throwM $ IndexDimensionException ix d + getDimM (i2, _) 2 = pure i2 + getDimM (_, i1) 1 = pure i1 + getDimM ix d = throwM $ IndexDimensionException ix d {-# INLINE [1] getDimM #-} setDimM (_, i1) 2 i2 = pure (i2, i1) setDimM (i2, _) 1 i1 = pure (i2, i1) - setDimM ix d _ = throwM $ IndexDimensionException ix d + setDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] setDimM #-} - modifyDimM (i2, i1) 2 f = pure (i2, (f i2, i1)) - modifyDimM (i2, i1) 1 f = pure (i1, ( i2, f i1)) - modifyDimM ix d _ = throwM $ IndexDimensionException ix d + modifyDimM (i2, i1) 2 f = pure (i2, (f i2, i1)) + modifyDimM (i2, i1) 1 f = pure (i1, (i2, f i1)) + modifyDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] modifyDimM #-} pullOutDimM (i2, i1) 2 = pure (i2, i1) pullOutDimM (i2, i1) 1 = pure (i1, i2) - pullOutDimM ix d = throwM $ IndexDimensionException ix d + pullOutDimM ix d = throwM $ IndexDimensionException ix d {-# INLINE [1] pullOutDimM #-} insertDimM i1 2 i2 = pure (i2, i1) insertDimM i2 1 i1 = pure (i2, i1) - insertDimM ix d _ = throwM $ IndexDimensionException ix d + insertDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] insertDimM #-} pureIndex i = (i, i) {-# INLINE [1] pureIndex #-} liftIndex2 f (i2, i1) (i2', i1') = (f i2 i2', f i1 i1') {-# INLINE [1] liftIndex2 #-} - -- | -- @since 0.1.0 instance Index Ix3T where type Dimensions Ix3T = 3 dimensions _ = 3 {-# INLINE [1] dimensions #-} - totalElem (SafeSz (k3, k2, k1)) = k3 * k2 * k1 + totalElem (SafeSz (k3, k2, k1)) = k3 * k2 * k1 {-# INLINE [1] totalElem #-} consDim i3 (i2, i1) = (i3, i2, i1) {-# INLINE [1] consDim #-} @@ -220,36 +229,35 @@ {-# INLINE [1] snocDim #-} unsnocDim (i3, i2, i1) = ((i3, i2), i1) {-# INLINE [1] unsnocDim #-} - getDimM (i3, _, _) 3 = pure i3 - getDimM ( _, i2, _) 2 = pure i2 - getDimM ( _, _, i1) 1 = pure i1 - getDimM ix d = throwM $ IndexDimensionException ix d + getDimM (i3, _, _) 3 = pure i3 + getDimM (_, i2, _) 2 = pure i2 + getDimM (_, _, i1) 1 = pure i1 + getDimM ix d = throwM $ IndexDimensionException ix d {-# INLINE [1] getDimM #-} - setDimM ( _, i2, i1) 3 i3 = pure (i3, i2, i1) - setDimM (i3, _, i1) 2 i2 = pure (i3, i2, i1) - setDimM (i3, i2, _) 1 i1 = pure (i3, i2, i1) - setDimM ix d _ = throwM $ IndexDimensionException ix d + setDimM (_, i2, i1) 3 i3 = pure (i3, i2, i1) + setDimM (i3, _, i1) 2 i2 = pure (i3, i2, i1) + setDimM (i3, i2, _) 1 i1 = pure (i3, i2, i1) + setDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] setDimM #-} - modifyDimM (i3, i2, i1) 3 f = pure (i3, (f i3, i2, i1)) - modifyDimM (i3, i2, i1) 2 f = pure (i2, ( i3, f i2, i1)) - modifyDimM (i3, i2, i1) 1 f = pure (i1, ( i3, i2, f i1)) - modifyDimM ix d _ = throwM $ IndexDimensionException ix d + modifyDimM (i3, i2, i1) 3 f = pure (i3, (f i3, i2, i1)) + modifyDimM (i3, i2, i1) 2 f = pure (i2, (i3, f i2, i1)) + modifyDimM (i3, i2, i1) 1 f = pure (i1, (i3, i2, f i1)) + modifyDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] modifyDimM #-} pullOutDimM (i3, i2, i1) 3 = pure (i3, (i2, i1)) pullOutDimM (i3, i2, i1) 2 = pure (i2, (i3, i1)) pullOutDimM (i3, i2, i1) 1 = pure (i1, (i3, i2)) - pullOutDimM ix d = throwM $ IndexDimensionException ix d + pullOutDimM ix d = throwM $ IndexDimensionException ix d {-# INLINE [1] pullOutDimM #-} insertDimM (i2, i1) 3 i3 = pure (i3, i2, i1) insertDimM (i3, i1) 2 i2 = pure (i3, i2, i1) insertDimM (i3, i2) 1 i1 = pure (i3, i2, i1) - insertDimM ix d _ = throwM $ IndexDimensionException ix d + insertDimM ix d _ = throwM $ IndexDimensionException ix d pureIndex i = (i, i, i) {-# INLINE [1] pureIndex #-} liftIndex2 f (i3, i2, i1) (i3', i2', i1') = (f i3 i3', f i2 i2', f i1 i1') {-# INLINE [1] liftIndex2 #-} - instance Index Ix4T where type Dimensions Ix4T = 4 dimensions _ = 4 @@ -264,42 +272,41 @@ {-# INLINE [1] snocDim #-} unsnocDim (i4, i3, i2, i1) = ((i4, i3, i2), i1) {-# INLINE [1] unsnocDim #-} - getDimM (i4, _, _, _) 4 = pure i4 - getDimM ( _, i3, _, _) 3 = pure i3 - getDimM ( _, _, i2, _) 2 = pure i2 - getDimM ( _, _, _, i1) 1 = pure i1 - getDimM ix d = throwM $ IndexDimensionException ix d + getDimM (i4, _, _, _) 4 = pure i4 + getDimM (_, i3, _, _) 3 = pure i3 + getDimM (_, _, i2, _) 2 = pure i2 + getDimM (_, _, _, i1) 1 = pure i1 + getDimM ix d = throwM $ IndexDimensionException ix d {-# INLINE [1] getDimM #-} - setDimM ( _, i3, i2, i1) 4 i4 = pure (i4, i3, i2, i1) - setDimM (i4, _, i2, i1) 3 i3 = pure (i4, i3, i2, i1) - setDimM (i4, i3, _, i1) 2 i2 = pure (i4, i3, i2, i1) - setDimM (i4, i3, i2, _) 1 i1 = pure (i4, i3, i2, i1) - setDimM ix d _ = throwM $ IndexDimensionException ix d + setDimM (_, i3, i2, i1) 4 i4 = pure (i4, i3, i2, i1) + setDimM (i4, _, i2, i1) 3 i3 = pure (i4, i3, i2, i1) + setDimM (i4, i3, _, i1) 2 i2 = pure (i4, i3, i2, i1) + setDimM (i4, i3, i2, _) 1 i1 = pure (i4, i3, i2, i1) + setDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] setDimM #-} - modifyDimM (i4, i3, i2, i1) 4 f = pure (i4, (f i4, i3, i2, i1)) - modifyDimM (i4, i3, i2, i1) 3 f = pure (i3, ( i4, f i3, i2, i1)) - modifyDimM (i4, i3, i2, i1) 2 f = pure (i2, ( i4, i3, f i2, i1)) - modifyDimM (i4, i3, i2, i1) 1 f = pure (i1, ( i4, i3, i2, f i1)) - modifyDimM ix d _ = throwM $ IndexDimensionException ix d + modifyDimM (i4, i3, i2, i1) 4 f = pure (i4, (f i4, i3, i2, i1)) + modifyDimM (i4, i3, i2, i1) 3 f = pure (i3, (i4, f i3, i2, i1)) + modifyDimM (i4, i3, i2, i1) 2 f = pure (i2, (i4, i3, f i2, i1)) + modifyDimM (i4, i3, i2, i1) 1 f = pure (i1, (i4, i3, i2, f i1)) + modifyDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] modifyDimM #-} pullOutDimM (i4, i3, i2, i1) 4 = pure (i4, (i3, i2, i1)) pullOutDimM (i4, i3, i2, i1) 3 = pure (i3, (i4, i2, i1)) pullOutDimM (i4, i3, i2, i1) 2 = pure (i2, (i4, i3, i1)) pullOutDimM (i4, i3, i2, i1) 1 = pure (i1, (i4, i3, i2)) - pullOutDimM ix d = throwM $ IndexDimensionException ix d + pullOutDimM ix d = throwM $ IndexDimensionException ix d {-# INLINE [1] pullOutDimM #-} insertDimM (i3, i2, i1) 4 i4 = pure (i4, i3, i2, i1) insertDimM (i4, i2, i1) 3 i3 = pure (i4, i3, i2, i1) insertDimM (i4, i3, i1) 2 i2 = pure (i4, i3, i2, i1) insertDimM (i4, i3, i2) 1 i1 = pure (i4, i3, i2, i1) - insertDimM ix d _ = throwM $ IndexDimensionException ix d + insertDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] insertDimM #-} pureIndex i = (i, i, i, i) {-# INLINE [1] pureIndex #-} liftIndex2 f (i4, i3, i2, i1) (i4', i3', i2', i1') = (f i4 i4', f i3 i3', f i2 i2', f i1 i1') {-# INLINE [1] liftIndex2 #-} - instance Index Ix5T where type Dimensions Ix5T = 5 dimensions _ = 5 @@ -314,40 +321,40 @@ {-# INLINE [1] snocDim #-} unsnocDim (i5, i4, i3, i2, i1) = ((i5, i4, i3, i2), i1) {-# INLINE [1] unsnocDim #-} - getDimM (i5, _, _, _, _) 5 = pure i5 - getDimM ( _, i4, _, _, _) 4 = pure i4 - getDimM ( _, _, i3, _, _) 3 = pure i3 - getDimM ( _, _, _, i2, _) 2 = pure i2 - getDimM ( _, _, _, _, i1) 1 = pure i1 - getDimM ix d = throwM $ IndexDimensionException ix d + getDimM (i5, _, _, _, _) 5 = pure i5 + getDimM (_, i4, _, _, _) 4 = pure i4 + getDimM (_, _, i3, _, _) 3 = pure i3 + getDimM (_, _, _, i2, _) 2 = pure i2 + getDimM (_, _, _, _, i1) 1 = pure i1 + getDimM ix d = throwM $ IndexDimensionException ix d {-# INLINE [1] getDimM #-} - setDimM ( _, i4, i3, i2, i1) 5 i5 = pure (i5, i4, i3, i2, i1) - setDimM (i5, _, i3, i2, i1) 4 i4 = pure (i5, i4, i3, i2, i1) - setDimM (i5, i4, _, i2, i1) 3 i3 = pure (i5, i4, i3, i2, i1) - setDimM (i5, i4, i3, _, i1) 2 i2 = pure (i5, i4, i3, i2, i1) - setDimM (i5, i4, i3, i2, _) 1 i1 = pure (i5, i4, i3, i2, i1) - setDimM ix d _ = throwM $ IndexDimensionException ix d + setDimM (_, i4, i3, i2, i1) 5 i5 = pure (i5, i4, i3, i2, i1) + setDimM (i5, _, i3, i2, i1) 4 i4 = pure (i5, i4, i3, i2, i1) + setDimM (i5, i4, _, i2, i1) 3 i3 = pure (i5, i4, i3, i2, i1) + setDimM (i5, i4, i3, _, i1) 2 i2 = pure (i5, i4, i3, i2, i1) + setDimM (i5, i4, i3, i2, _) 1 i1 = pure (i5, i4, i3, i2, i1) + setDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] setDimM #-} - modifyDimM (i5, i4, i3, i2, i1) 5 f = pure (i5, (f i5, i4, i3, i2, i1)) - modifyDimM (i5, i4, i3, i2, i1) 4 f = pure (i4, ( i5, f i4, i3, i2, i1)) - modifyDimM (i5, i4, i3, i2, i1) 3 f = pure (i3, ( i5, i4, f i3, i2, i1)) - modifyDimM (i5, i4, i3, i2, i1) 2 f = pure (i2, ( i5, i4, i3, f i2, i1)) - modifyDimM (i5, i4, i3, i2, i1) 1 f = pure (i1, ( i5, i4, i3, i2, f i1)) - modifyDimM ix d _ = throwM $ IndexDimensionException ix d + modifyDimM (i5, i4, i3, i2, i1) 5 f = pure (i5, (f i5, i4, i3, i2, i1)) + modifyDimM (i5, i4, i3, i2, i1) 4 f = pure (i4, (i5, f i4, i3, i2, i1)) + modifyDimM (i5, i4, i3, i2, i1) 3 f = pure (i3, (i5, i4, f i3, i2, i1)) + modifyDimM (i5, i4, i3, i2, i1) 2 f = pure (i2, (i5, i4, i3, f i2, i1)) + modifyDimM (i5, i4, i3, i2, i1) 1 f = pure (i1, (i5, i4, i3, i2, f i1)) + modifyDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] modifyDimM #-} pullOutDimM (i5, i4, i3, i2, i1) 5 = pure (i5, (i4, i3, i2, i1)) pullOutDimM (i5, i4, i3, i2, i1) 4 = pure (i4, (i5, i3, i2, i1)) pullOutDimM (i5, i4, i3, i2, i1) 3 = pure (i3, (i5, i4, i2, i1)) pullOutDimM (i5, i4, i3, i2, i1) 2 = pure (i2, (i5, i4, i3, i1)) pullOutDimM (i5, i4, i3, i2, i1) 1 = pure (i1, (i5, i4, i3, i2)) - pullOutDimM ix d = throwM $ IndexDimensionException ix d + pullOutDimM ix d = throwM $ IndexDimensionException ix d {-# INLINE [1] pullOutDimM #-} insertDimM (i4, i3, i2, i1) 5 i5 = pure (i5, i4, i3, i2, i1) insertDimM (i5, i3, i2, i1) 4 i4 = pure (i5, i4, i3, i2, i1) insertDimM (i5, i4, i2, i1) 3 i3 = pure (i5, i4, i3, i2, i1) insertDimM (i5, i4, i3, i1) 2 i2 = pure (i5, i4, i3, i2, i1) insertDimM (i5, i4, i3, i2) 1 i1 = pure (i5, i4, i3, i2, i1) - insertDimM ix d _ = throwM $ IndexDimensionException ix d + insertDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] insertDimM #-} pureIndex i = (i, i, i, i, i) {-# INLINE [1] pureIndex #-} diff --git a/src/Data/Massiv/Core/Iterator.hs b/src/Data/Massiv/Core/Iterator.hs deleted file mode 100644 --- a/src/Data/Massiv/Core/Iterator.hs +++ /dev/null @@ -1,173 +0,0 @@ -{-# LANGUAGE BangPatterns #-} --- | --- Module : Data.Massiv.Core.Iterator --- Copyright : (c) Alexey Kuleshevich 2018-2019 --- License : BSD3 --- Maintainer : Alexey Kuleshevich --- Stability : experimental --- Portability : non-portable --- -module Data.Massiv.Core.Iterator - ( loop - , loopA_ - , loopM - , loopM_ - , loopDeepM - , splitLinearly - , splitLinearlyWith_ - , splitLinearlyWithM_ - , splitLinearlyWithStartAtM_ - , splitLinearlyWithStatefulM_ - ) where - -import Control.Scheduler - --- | Efficient loop with an accumulator --- --- @since 0.1.0 -loop :: Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> a) -> a -loop !init' condition increment !initAcc f = go init' initAcc - where - go !step !acc - | condition step = go (increment step) (f step acc) - | otherwise = acc -{-# INLINE loop #-} - - --- | Efficient monadic loop with an accumulator --- --- >>> loopM 1 (< 20) (+ 2) [] (\i a -> Just (i:a)) --- Just [19,17,15,13,11,9,7,5,3,1] --- --- @since 0.1.0 -loopM :: Monad m => Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a -loopM !init' condition increment !initAcc f = go init' initAcc - where - go !step !acc - | condition step = f step acc >>= go (increment step) - | otherwise = return acc -{-# INLINE loopM #-} - - --- | Efficient monadic loop. Result of each iteration is discarded. --- --- @since 0.1.0 -loopM_ :: Monad m => Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m () -loopM_ !init' condition increment f = go init' - where - go !step - | condition step = f step >> go (increment step) - | otherwise = pure () - -{-# INLINE loopM_ #-} - - --- | Efficient Applicative loop. Result of each iteration is discarded. --- --- @since 0.3.0 -loopA_ :: Applicative f => Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f a) -> f () -loopA_ !init' condition increment f = go init' - where - go !step - | condition step = f step *> go (increment step) - | otherwise = pure () -{-# INLINE loopA_ #-} - - --- | Similar to `loopM`, but slightly less efficient monadic loop with an accumulator that reverses --- the direction of action application. eg: --- --- >>> loopDeepM 1 (< 20) (+ 2) [] (\i a -> Just (i:a)) --- Just [1,3,5,7,9,11,13,15,17,19] --- --- Equivalent to: --- --- >>> loopM 19 (>= 1) (subtract 2) [] (\i a -> Just (i:a)) --- Just [1,3,5,7,9,11,13,15,17,19] --- --- @since 0.1.0 -loopDeepM :: Monad m => Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a -loopDeepM !init' condition increment !initAcc f = go init' initAcc - where - go !step !acc - | condition step = go (increment step) acc >>= f step - | otherwise = return acc -{-# INLINE loopDeepM #-} - - --- | Divide length in chunks and apply a function to the computed results --- --- @since 0.2.1 -splitLinearly :: Int -- ^ Number of chunks - -> Int -- ^ Total length - -> (Int -> Int -> a) -- ^ Function that accepts a chunk length and slack start index - -> a -splitLinearly numChunks totalLength action = action chunkLength slackStart - where - !chunkLength = totalLength `quot` numChunks - !slackStart = chunkLength * numChunks -{-# INLINE splitLinearly #-} - - --- | Interator that can be used to split computation amongst different workers. For monadic --- generator see `splitLinearlyWithM_`. --- --- @since 0.2.1 -splitLinearlyWith_ :: - Monad m => Scheduler m () -> Int -> (Int -> b) -> (Int -> b -> m ()) -> m () -splitLinearlyWith_ scheduler totalLength index = - splitLinearlyWithM_ scheduler totalLength (pure . index) -{-# INLINE splitLinearlyWith_ #-} - - --- | Interator that can be used to split computation jobs --- --- @since 0.2.6 -splitLinearlyWithM_ :: - Monad m => Scheduler m () -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m () -splitLinearlyWithM_ scheduler totalLength make write = - splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do - loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start -> - scheduleWork_ scheduler $ - loopM_ start (< (start + chunkLength)) (+ 1) $ \ !k -> make k >>= write k - scheduleWork_ scheduler $ loopM_ slackStart (< totalLength) (+ 1) $ \ !k -> make k >>= write k -{-# INLINE splitLinearlyWithM_ #-} - - --- | Interator that can be used to split computation jobs --- --- @since 0.3.0 -splitLinearlyWithStartAtM_ :: - Monad m => Scheduler m () -> Int -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m () -splitLinearlyWithStartAtM_ scheduler startAt totalLength make write = - splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do - loopM_ startAt (< (slackStart + startAt)) (+ chunkLength) $ \ !start -> - scheduleWork_ scheduler $ - loopM_ start (< (start + chunkLength)) (+ 1) $ \ !k -> make k >>= write k - scheduleWork_ scheduler $ - loopM_ (slackStart + startAt) (< (totalLength + startAt)) (+ 1) $ \ !k -> make k >>= write k -{-# INLINE splitLinearlyWithStartAtM_ #-} - - - --- | Interator that can be used to split computation jobs, while using a stateful scheduler. --- --- @since 0.3.4 -splitLinearlyWithStatefulM_ :: - Monad m - => SchedulerWS s m () - -> Int -- ^ Total linear length - -> (Int -> s -> m b) -- ^ Element producing action - -> (Int -> b -> m c) -- ^ Element storing action - -> m () -splitLinearlyWithStatefulM_ schedulerWS totalLength make store = - let nWorkers = numWorkers (unwrapSchedulerWS schedulerWS) - in splitLinearly nWorkers totalLength $ \chunkLength slackStart -> do - loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start -> - scheduleWorkState_ schedulerWS $ \s -> - loopM_ start (< (start + chunkLength)) (+ 1) $ \ !k -> - make k s >>= store k - scheduleWorkState_ schedulerWS $ \s -> - loopM_ slackStart (< totalLength) (+ 1) $ \ !k -> - make k s >>= store k -{-# INLINE splitLinearlyWithStatefulM_ #-} 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 @@ -1,233 +1,218 @@ {-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE LambdaCase #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE RecordWildCards #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} {-# LANGUAGE UndecidableInstances #-} {-# OPTIONS_GHC -fno-warn-orphans #-} + -- | -- Module : Data.Massiv.Core.List --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Core.List - ( LN - , L(..) - , Array(..) - , toListArray - , showsArrayPrec - , showArrayList - , ListItem - ) where +module Data.Massiv.Core.List ( + L (..), + Array (..), + List (..), + toListArray, + showsArrayPrec, + showArrayList, + ListItem, +) where -import Control.Exception import Control.Monad (unless, when) import Control.Scheduler import Data.Coerce -import Data.Foldable (foldr') +import Data.Functor.Identity +import Data.Kind import qualified Data.List as L -import qualified Data.Massiv.Array.Manifest.Vector.Stream as S import Data.Massiv.Core.Common +import qualified Data.Massiv.Vector.Stream as S +import Data.Monoid import Data.Typeable -import GHC.Exts +import GHC.Exts (IsList (..)) +import GHC.TypeLits import System.IO.Unsafe (unsafePerformIO) -data LN - -type family ListItem ix e :: * where +type family ListItem ix e :: Type where ListItem Ix1 e = e - ListItem ix e = [ListItem (Lower ix) e] - -type instance NestedStruct LN ix e = [ListItem ix e] - -newtype instance Array LN ix e = List { unList :: [Elt LN ix e] } - - -instance Construct LN Ix1 e where - setComp _ = id - {-# INLINE setComp #-} - makeArray _ (Sz n) f = coerce (fmap f [0 .. n - 1]) - {-# INLINE makeArray #-} - makeArrayLinear _ (Sz n) f = coerce (fmap f [0 .. n - 1]) - {-# INLINE makeArrayLinear #-} - -instance {-# OVERLAPPING #-} Nested LN Ix1 e where - fromNested = coerce - {-# INLINE fromNested #-} - toNested = coerce - {-# INLINE toNested #-} + ListItem ix e = [ListItem (Lower ix) e] -instance ( Elt LN ix e ~ Array LN (Lower ix) e - , ListItem ix e ~ [ListItem (Lower ix) e] - , Coercible (Elt LN ix e) (ListItem ix e) - ) => - Nested LN ix e where - fromNested = coerce - {-# INLINE fromNested #-} - toNested = coerce - {-# INLINE toNested #-} +type family Elt ix e :: Type where + Elt Ix1 e = e + Elt ix e = List (Lower ix) e +newtype List ix e = List {unList :: [Elt ix e]} -instance Nested LN ix e => IsList (Array LN ix e) where - type Item (Array LN ix e) = ListItem ix e - fromList = fromNested +instance Coercible (Elt ix e) (ListItem ix e) => IsList (List ix e) where + type Item (List ix e) = ListItem ix e + fromList = coerce {-# INLINE fromList #-} - toList = toNested + toList = coerce {-# INLINE toList #-} - data L = L -type instance NestedStruct L ix e = Array LN ix e - -data instance Array L ix e = LArray { lComp :: Comp - , lData :: !(Array LN ix e) } - - -instance Nested L ix e where - fromNested = LArray Seq - {-# INLINE fromNested #-} - toNested = lData - {-# INLINE toNested #-} - +data instance Array L ix e = LArray + { lComp :: Comp + , lData :: !(List ix e) + } -instance Nested LN ix e => IsList (Array L ix e) where +instance Coercible (Elt ix e) (ListItem ix e) => IsList (Array L ix e) where type Item (Array L ix e) = ListItem ix e - fromList = LArray Seq . fromNested + fromList = LArray Seq . coerce {-# INLINE fromList #-} - toList = toNested . lData + toList = coerce . lData {-# INLINE toList #-} -instance {-# OVERLAPPING #-} Ragged L Ix1 e where +lengthHintList :: [a] -> LengthHint +lengthHintList = + \case + [] -> LengthExact zeroSz + _ -> LengthUnknown +{-# INLINE lengthHintList #-} + +instance Shape L Ix1 where + linearSize = outerLength + {-# INLINE linearSize #-} + linearSizeHint = lengthHintList . unList . lData + {-# INLINE linearSizeHint #-} isNull = null . unList . lData {-# INLINE isNull #-} - emptyR comp = LArray comp (List []) - {-# INLINE emptyR #-} - edgeSize = SafeSz . length . unList . lData - {-# INLINE edgeSize #-} - consR x arr = arr { lData = coerce (x : coerce (lData arr)) } - {-# INLINE consR #-} - unconsR LArray {..} = - case L.uncons $ coerce lData of - Nothing -> Nothing - Just (x, xs) -> Just (x, LArray lComp (coerce xs)) - {-# INLINE unconsR #-} + outerSize = linearSize + {-# INLINE outerSize #-} + +instance Shape L Ix2 where + linearSize = SafeSz . getSum . foldMap (Sum . length . unList) . unList . lData + {-# INLINE linearSize #-} + linearSizeHint = lengthHintList . unList . lData + {-# INLINE linearSizeHint #-} + isNull = getAll . foldMap (All . null . unList) . unList . lData + {-# INLINE isNull #-} + outerSize arr = + case unList (lData arr) of + [] -> zeroSz + (x : xs) -> SafeSz ((1 + length xs) :. length (unList x)) + {-# INLINE outerSize #-} + +instance (Shape L (Ix (n - 1)), Index (IxN n)) => Shape L (IxN n) where + linearSize = SafeSz . getSum . foldMap (Sum . unSz . linearSize . LArray Seq) . unList . lData + {-# INLINE linearSize #-} + linearSizeHint = lengthHintList . unList . lData + {-# INLINE linearSizeHint #-} + isNull = getAll . foldMap (All . isNull . LArray Seq) . unList . lData + {-# INLINE isNull #-} + outerSize arr = + case unList (lData arr) of + [] -> zeroSz + (x : xs) -> SafeSz ((1 + length xs) :> unSz (outerSize (LArray Seq x))) + {-# INLINE outerSize #-} + +outerLength :: Array L ix e -> Sz Int +outerLength = SafeSz . length . unList . lData +{-# INLINE outerLength #-} + +instance Ragged L Ix1 e where flattenRagged = id {-# INLINE flattenRagged #-} generateRaggedM !comp !k f = do - xs <- loopDeepM 0 (< coerce k) (+ 1) [] $ \i acc -> do - e <- f i - return (e:acc) + xs <- + loopDeepM 0 (< coerce k) (+ 1) [] $ \i acc -> do + e <- f i + return (e : acc) return $ LArray comp $ coerce xs {-# INLINE generateRaggedM #-} - loadRagged using uWrite start end sz xs = - using $ do - leftOver <- - loopM start (< end) (+ 1) xs $ \i xs' -> - case unconsR xs' of - Nothing -> return $! throw (DimTooShortException sz (outerLength xs)) - Just (y, ys) -> uWrite i y >> return ys - unless (isNull leftOver) (return $! throw DimTooLongException) - {-# INLINE loadRagged #-} + loadRaggedST _scheduler xs uWrite start end sz = go (unList (lData xs)) start + where + go (y : ys) i + | i < end = uWrite i y >> go ys (i + 1) + | otherwise = throwM $ DimTooLongException 1 sz (outerLength xs) + go [] i = when (i /= end) $ throwM $ DimTooShortException 1 sz (outerLength xs) + {-# INLINE loadRaggedST #-} raggedFormat f _ arr = L.concat $ "[ " : L.intersperse ", " (map f (coerce (lData arr))) ++ [" ]"] - -instance (Index ix, Ragged L ix e) => Load L ix e where - size = coerce . edgeSize - {-# INLINE size #-} - getComp = lComp - {-# INLINE getComp #-} - loadArrayM scheduler arr uWrite = - loadRagged (scheduleWork scheduler) uWrite 0 (totalElem sz) sz arr - where !sz = edgeSize arr - {-# INLINE loadArrayM #-} - - -instance (Index ix, Load L ix e, Ragged L ix e) => Load LN ix e where - size = edgeSize . LArray Seq - {-# INLINE size #-} - getComp _ = Seq - {-# INLINE getComp #-} - loadArrayM scheduler arr uWrite = - loadRagged (scheduleWork scheduler) uWrite 0 (totalElem sz) sz arrL +instance (Shape L ix, Ragged L ix e) => Load L ix e where + makeArray comp sz f = runIdentity $ generateRaggedM comp sz (pure . f) + {-# INLINE makeArray #-} + iterArrayLinearST_ scheduler arr uWrite = + loadRaggedST scheduler arr uWrite 0 (totalElem sz) sz where - !arrL = LArray Seq arr - !sz = size arrL - {-# INLINE loadArrayM #-} - - - -outerLength :: Array L ix e -> Sz Int -outerLength = SafeSz . length . unList . lData + !sz = outerSize arr + {-# INLINE iterArrayLinearST_ #-} -instance ( Index ix - , Index (Lower ix) - , Ragged L (Lower ix) e - , Elt L ix e ~ Array L (Lower ix) e - , Elt LN ix e ~ Array LN (Lower ix) e - , Coercible (Elt LN ix e) [Elt LN (Lower ix) e] - ) => - Ragged L ix e where - isNull = null . unList . lData - {-# INLINE isNull #-} - emptyR comp = LArray comp (List []) - {-# INLINE emptyR #-} - edgeSize arr = - SafeSz - (consDim (length (unList (lData arr))) $ - case unconsR arr of - Nothing -> zeroIndex - Just (x, _) -> coerce (edgeSize x)) - {-# INLINE edgeSize #-} - consR (LArray _ x) arr = newArr +instance Ragged L Ix2 e where + generateRaggedM = unsafeGenerateParM + {-# INLINE generateRaggedM #-} + flattenRagged arr = LArray{lComp = lComp arr, lData = coerce xs} where - newArr = arr {lData = coerce (x : coerce (lData arr))} - {-# INLINE consR #-} - unconsR LArray {..} = - case L.uncons (coerce lData) of - Nothing -> Nothing - Just (x, xs) -> - let newArr = LArray lComp (coerce xs) - newX = LArray lComp x - in Just (newX, newArr) - {-# INLINE unconsR #-} - -- generateRaggedM Seq !sz f = do - -- let !(k, szL) = unconsSz sz - -- loopDeepM 0 (< coerce k) (+ 1) (emptyR Seq) $ \i acc -> do - -- e <- generateRaggedM Seq szL (\ !ixL -> f (consDim i ixL)) - -- return (cons e acc) + xs = concatMap (unList . lData . flattenRagged . LArray (lComp arr)) (unList (lData arr)) + {-# INLINE flattenRagged #-} + loadRaggedST scheduler xs uWrite start end sz + | isZeroSz sz = when (isNotNull (flattenRagged xs)) (throwM ShapeNonEmpty) + | otherwise = do + let (k, szL) = unconsSz sz + step = totalElem szL + leftOver <- + loopM start (< end) (+ step) (coerce (lData xs)) $ \i zs -> + case zs of + [] -> throwM (DimTooShortException 2 k (outerLength xs)) + (y : ys) -> do + scheduleWork_ scheduler $ + let end' = i + step + go (a : as) j + | j < end' = uWrite j a >> go as (j + 1) + | otherwise = throwM $ DimTooLongException 1 szL (Sz (length y)) + go [] j = when (j /= end') $ throwM (DimTooShortException 1 szL (Sz (length y))) + in go y i + pure ys + unless (null leftOver) $ throwM $ DimTooLongException 2 k (outerLength xs) + {-# INLINE loadRaggedST #-} + raggedFormat f sep (LArray comp xs) = + showN (\s y -> raggedFormat f s (LArray comp y :: Array L Ix1 e)) sep (coerce xs) + +instance + ( Shape L (IxN n) + , Ragged L (Ix (n - 1)) e + , Coercible (Elt (Ix (n - 1)) e) (ListItem (Ix (n - 1)) e) + ) + => Ragged L (IxN n) e + where generateRaggedM = unsafeGenerateParM {-# INLINE generateRaggedM #-} - flattenRagged arr = LArray {lComp = lComp arr, lData = coerce xs} + flattenRagged arr = LArray{lComp = lComp arr, lData = coerce xs} where xs = concatMap (unList . lData . flattenRagged . LArray (lComp arr)) (unList (lData arr)) {-# INLINE flattenRagged #-} - loadRagged using uWrite start end sz xs = do - let (k, szL) = unconsSz sz - step = totalElem szL - isZero = totalElem sz == 0 - when (isZero && not (isNull (flattenRagged xs))) (return $! throw DimTooLongException) - unless isZero $ do - leftOver <- - loopM start (< end) (+ step) xs $ \i zs -> - case unconsR zs of - Nothing -> return $! throw (DimTooShortException k (outerLength xs)) - Just (y, ys) -> do - _ <- loadRagged using uWrite i (i + step) szL y - return ys - unless (isNull leftOver) (return $! throw DimTooLongException) - {-# INLINE loadRagged #-} + loadRaggedST scheduler xs uWrite start end sz + | isZeroSz sz = when (isNotNull (flattenRagged xs)) (throwM ShapeNonEmpty) + | otherwise = do + let (k, szL) = unconsSz sz + step = totalElem szL + subScheduler + | end - start < numWorkers scheduler * step = scheduler + | otherwise = trivialScheduler_ + leftOver <- + loopM start (< end) (+ step) (unList (lData xs)) $ \i zs -> + case zs of + [] -> throwM (DimTooShortException (dimensions sz) k (outerLength xs)) + (y : ys) -> do + scheduleWork_ scheduler $ + loadRaggedST subScheduler (LArray Seq y) uWrite i (i + step) szL + pure ys + unless (null leftOver) $ throwM $ DimTooLongException (dimensions sz) k (outerLength xs) + {-# INLINE loadRaggedST #-} raggedFormat f sep (LArray comp xs) = - showN (\s y -> raggedFormat f s (LArray comp y :: Array L (Lower ix) e)) sep (coerce xs) + showN (\s y -> raggedFormat f s (LArray comp y :: Array L (Ix (n - 1)) e)) sep (coerce xs) -unsafeGenerateParM :: - (Elt LN ix e ~ Array LN (Lower ix) e, Index ix, Monad m, Ragged L (Lower ix) e) +unsafeGenerateParM + :: (Elt ix e ~ List (Lower ix) e, Index ix, Monad m, Ragged L (Lower ix) e) => Comp -> Sz ix -> (ix -> m e) @@ -236,153 +221,125 @@ res <- sequence $ unsafePerformIO $ do let !(ksz, szL) = unconsSz sz !k = unSz ksz - withScheduler comp $ \ scheduler -> - splitLinearly (numWorkers scheduler) k $ \ chunkLength slackStart -> do - loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start -> + withScheduler comp $ \scheduler -> + splitLinearly (numWorkers scheduler) k $ \chunkLength slackStart -> do + loopA_ 0 (< slackStart) (+ chunkLength) $ \ !start -> scheduleWork scheduler $ do res <- loopDeepM start (< (start + chunkLength)) (+ 1) [] $ \i acc -> - return (fmap lData (generateRaggedM Seq szL (\ !ixL -> f (consDim i ixL))):acc) + return (fmap lData (generateRaggedM Seq szL (\ !ixL -> f (consDim i ixL))) : acc) return $! sequence res when (slackStart < k) $ scheduleWork scheduler $ do res <- loopDeepM slackStart (< k) (+ 1) [] $ \i acc -> - return (fmap lData (generateRaggedM Seq szL (\ !ixL -> f (consDim i ixL))):acc) + return (fmap lData (generateRaggedM Seq szL (\ !ixL -> f (consDim i ixL))) : acc) return $! sequence res return $ LArray comp $ List $ concat res {-# INLINE unsafeGenerateParM #-} - -instance {-# OVERLAPPING #-} Construct L Ix1 e where - setComp c arr = arr { lComp = c } - {-# INLINE setComp #-} - makeArray comp sz f = LArray comp $ List $ unsafePerformIO $ - withScheduler comp $ \scheduler -> - loopM_ 0 (< coerce sz) (+ 1) (scheduleWork scheduler . return . f) - {-# INLINE makeArray #-} - - -instance ( Index ix - , Ragged L ix e - , Ragged L (Lower ix) e - , Elt L ix e ~ Array L (Lower ix) e - ) => - Construct L ix e where - setComp c arr = arr {lComp = c} +instance Strategy L where + setComp c arr = arr{lComp = c} {-# INLINE setComp #-} - makeArray = unsafeGenerateN - {-# INLINE makeArray #-} - - -- TODO: benchmark against using unsafeGenerateM directly -unsafeGenerateN :: - ( Ragged r ix e - , Ragged r (Lower ix) e - , Elt r ix e ~ Array r (Lower ix) e ) - => Comp - -> Sz ix - -> (ix -> e) - -> Array r ix e -unsafeGenerateN comp sz f = unsafePerformIO $ do - let !(m, szL) = unconsSz sz - xs <- withScheduler comp $ \scheduler -> - loopM_ 0 (< coerce m) (+ 1) $ \i -> scheduleWork scheduler $ - generateRaggedM comp szL $ \ix -> return $ f (consDim i ix) - return $! foldr' consR (emptyR comp) xs -{-# INLINE unsafeGenerateN #-} + getComp = lComp + {-# INLINE getComp #-} + repr = L +-- -- TODO: benchmark against using unsafeGenerateM directly +-- unsafeGenerateN :: +-- ( Ragged r ix e +-- , Ragged r (Lower ix) e +-- , Elt r ix e ~ Array r (Lower ix) e ) +-- => Comp +-- -> Sz ix +-- -> (ix -> e) +-- -> Array r ix e +-- unsafeGenerateN comp sz f = unsafePerformIO $ do +-- let !(m, szL) = unconsSz sz +-- xs <- withScheduler comp $ \scheduler -> +-- loopM_ 0 (< coerce m) (+ 1) $ \i -> scheduleWork scheduler $ +-- generateRaggedM comp szL $ \ix -> return $ f (consDim i ix) +-- return $! foldr' consR (emptyR comp) xs +-- {-# INLINE unsafeGenerateN #-} -- | Construct an array backed by linked lists from any source array -- -- @since 0.4.0 -toListArray :: (Construct L ix e, Source r ix e) - => Array r ix e - -> Array L ix e -toListArray !arr = makeArray (getComp arr) (size arr) (unsafeIndex arr) +toListArray :: (Ragged L ix e, Shape r ix, Source r e) => Array r ix e -> Array L ix e +toListArray !arr = makeArray (getComp arr) (outerSize arr) (unsafeIndex arr) {-# INLINE toListArray #-} - - instance (Ragged L ix e, Show e) => Show (Array L ix e) where - showsPrec = showsArrayLAsPrec (Proxy :: Proxy L) - -instance (Ragged L ix e, Show e) => Show (Array LN ix e) where - show arr = " " ++ raggedFormat show "\n " arrL - where arrL = fromNested arr :: Array L ix e + showsPrec n arr = showsArrayLAsPrec (Proxy :: Proxy L) (outerSize arr) n arr +instance (Ragged L ix e, Show e) => Show (List ix e) where + show xs = " " ++ raggedFormat show "\n " arrL + where + arrL = LArray Seq xs :: Array L ix e showN :: (String -> a -> String) -> String -> [a] -> String -showN _ _ [] = "[ ]" +showN _ _ [] = "[ ]" showN fShow lnPrefix ls = L.concat - (["[ "] ++ - L.intersperse (lnPrefix ++ ", ") (map (fShow (lnPrefix ++ " ")) ls) ++ [lnPrefix, "]"]) - + ( ["[ "] + ++ L.intersperse (lnPrefix ++ ", ") (map (fShow (lnPrefix ++ " ")) ls) + ++ [lnPrefix, "]"] + ) -showsArrayLAsPrec :: - forall r ix e. (Ragged L ix e, Typeable r, Show e) +showsArrayLAsPrec + :: forall r ix e + . (Ragged L ix e, Typeable r, Show e) => Proxy r + -> Sz ix -> Int -> Array L ix e -- Array to show -> ShowS -showsArrayLAsPrec pr n arr = - opp . - ("Array " ++) . - showsTypeRep (typeRep pr) . - (' ':) . - showsPrec 1 (getComp arr) . (" (" ++) . shows (size arr) . (")\n" ++) . shows lnarr . clp +showsArrayLAsPrec pr sz n arr = + opp + . ("Array " ++) + . showsTypeRep (typeRep pr) + . (' ' :) + . showsPrec 1 (getComp arr) + . (" (" ++) + . shows sz + . (")\n" ++) + . shows lnarr + . clp where (opp, clp) = if n == 0 then (id, id) - else (('(':), ("\n)" ++)) - lnarr = toNested arr + else (('(' :), ("\n)" ++)) + lnarr = lData arr -- | Helper function for declaring `Show` instances for arrays -- -- @since 0.4.0 -showsArrayPrec :: - forall r r' ix ix' e. (Ragged L ix' e, Load r ix e, Source r' ix' e, Show e) - => (Array r ix e -> Array r' ix' e) -- ^ Modifier +showsArrayPrec + :: forall r r' ix e + . (Ragged L ix e, Load r ix e, Load r' ix e, Source r' e, Show e) + => (Array r ix e -> Array r' ix e) + -- ^ Modifier -> Int -> Array r ix e -- Array to show -> ShowS -showsArrayPrec f n arr = showsArrayLAsPrec (Proxy :: Proxy r) n larr +showsArrayPrec f n arr = showsArrayLAsPrec (Proxy :: Proxy r) sz n larr where + sz = size arr' arr' = f arr - larr = makeArray (getComp arr') (size arr') (evaluate' arr') :: Array L ix' e - + larr = makeArray (getComp arr') sz (evaluate' arr') :: Array L ix e -- | Helper function for declaring `Show` instances for arrays -- -- @since 0.4.0 showArrayList :: Show arr => [arr] -> String -> String -showArrayList arrs = ('[':) . go arrs . (']':) +showArrayList arrs = ('[' :) . go arrs . (']' :) where - go [] = id - go [x] = (' ':) . shows x . ('\n':) - go (x:xs) = (' ':) . shows x . ("\n," ++) . go xs - - -instance {-# OVERLAPPING #-} OuterSlice L Ix1 e where - unsafeOuterSlice (LArray _ xs) = (coerce xs !!) - {-# INLINE unsafeOuterSlice #-} - - -instance Ragged L ix e => OuterSlice L ix e where - unsafeOuterSlice arr' i = go 0 arr' - where - go n arr = - case unconsR arr of - Nothing -> throw $ IndexOutOfBoundsException (Sz (headDim (unSz (size arr')))) i - Just (x, _) | n == i -> x - Just (_, xs) -> go (n + 1) xs - {-# INLINE unsafeOuterSlice #-} - - -instance Stream LN Ix1 e where - toStream = S.fromList . coerce - {-# INLINE toStream #-} + go [] = id + go [x] = (' ' :) . shows x . ('\n' :) + go (x : xs) = (' ' :) . shows x . ("\n," ++) . go xs -instance Ragged L ix e => Stream L ix e where - toStream = S.fromList . coerce . lData . flattenRagged +instance Stream L Ix1 e where + toStream = S.fromList . unList . lData {-# INLINE toStream #-} + toStreamIx = S.indexed . S.fromList . unList . lData + {-# INLINE toStreamIx #-} diff --git a/src/Data/Massiv/Core/Loop.hs b/src/Data/Massiv/Core/Loop.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Core/Loop.hs @@ -0,0 +1,474 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MonoLocalBinds #-} +{-# LANGUAGE ScopedTypeVariables #-} + +-- | +-- Module : Data.Massiv.Core.Loop +-- Copyright : (c) Alexey Kuleshevich 2018-2022 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +module Data.Massiv.Core.Loop ( + loop, + loopF, + nextMaybeF, + loopA, + loopA_, + loopM, + loopM_, + iloopM, + iloopA_, + loopNextM, + loopNextA_, + loopDeepM, + splitLinearly, + splitLinearlyM, + splitLinearlyM_, + splitLinearlyWith_, + splitLinearlyWithM_, + splitLinearlyWithStartAtM_, + splitLinearlyWithStatefulM_, + iterLinearST_, + iterLinearAccST_, + iterLinearAccST, + splitNumChunks, + stepStartAdjust, + + -- * Experimental + splitWorkWithFactorST, + scheduleMassivWork, + withMassivScheduler_, +) where + +import Control.Monad (void, when) +import Control.Monad.IO.Unlift (MonadUnliftIO (..)) +import Control.Monad.Primitive +import Control.Monad.ST (ST) +import Control.Scheduler ( + Comp (..), + Scheduler, + SchedulerWS, + numWorkers, + scheduleWork, + scheduleWorkState_, + scheduleWork_, + trivialScheduler_, + unwrapSchedulerWS, + withScheduler_, + ) +import Control.Scheduler.Global (globalScheduler, withGlobalScheduler_) +import Data.Coerce +import Data.Functor.Identity + +-- | Efficient loop with an accumulator +-- +-- @since 0.1.0 +loop :: Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> a) -> a +loop initial condition increment initAcc f = + runIdentity (loopM initial condition increment initAcc (coerce f)) +{-# INLINE loop #-} + +-- | Efficient monadic loop with an accumulator +-- +-- >>> loopM 1 (< 20) (+ 2) [] (\i a -> Just (i:a)) +-- Just [19,17,15,13,11,9,7,5,3,1] +-- +-- @since 0.1.0 +loopM :: Monad m => Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a +loopM !initial condition increment !initAcc f = + go initial initAcc + where + go !step !acc + | condition step = f step acc >>= go (increment step) + | otherwise = pure acc +{-# INLINE loopM #-} + +-- | Efficient monadic loop with an accumulator and extra linear index incremented by 1. +-- +-- >>> iloopM 100 1 (< 20) (+ 2) [] (\i ix a -> Just ((i, ix) : a)) +-- Just [(109,19),(108,17),(107,15),(106,13),(105,11),(104,9),(103,7),(102,5),(101,3),(100,1)] +-- +-- @since 1.0.2 +iloopM + :: Monad m => Int -> Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> Int -> a -> m a) -> m a +iloopM !istart !initIx condition increment !initAcc f = go istart initIx initAcc + where + go !i !step !acc + | condition step = f i step acc >>= go (i + 1) (increment step) + | otherwise = pure acc +{-# INLINE iloopM #-} + +-- | Efficient monadic loop. Result of each iteration is discarded. +-- +-- @since 0.1.0 +loopM_ :: Monad m => Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m () +loopM_ !initial condition increment f = go initial + where + go !step + | condition step = f step >> go (increment step) + | otherwise = pure () +-- loopF initial condition increment (pure ()) (\i ma -> f i >> ma) +{-# INLINE loopM_ #-} +{-# DEPRECATED loopM_ "In favor of `loopA_`" #-} + +-- | Efficient monadic loop with extra linear index incremented by 1. +-- +-- >>> iloopA_ 100 1 (< 10) (+ 2) (\i ix -> print (i, ix)) +-- (100,1) +-- (101,3) +-- (102,5) +-- (103,7) +-- (104,9) +-- +-- @since 1.0.2 +iloopA_ + :: Applicative f => Int -> Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> Int -> f a) -> f () +iloopA_ !istart !initIx condition increment f = go istart initIx + where + go !i !step + | condition step = f i step *> go (i + 1) (increment step) + | otherwise = pure () +{-# INLINE iloopA_ #-} + +-- | Similar to `loopM_` except the action accepts not only the value for current step, +-- but also for the next one as well. +-- +-- @since 1.0.2 +loopNextA_ :: Applicative f => Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> Int -> f a) -> f () +loopNextA_ !initial condition increment f = go initial + where + go !step + | condition step = + let !next = increment step + in f step next *> go next + | otherwise = pure () +{-# INLINE loopNextA_ #-} + +-- | Similar to `loopM_` except the action accepts not only the value for current step, +-- but also for the next one as well. +-- +-- @since 1.0.2 +loopNextM :: Monad m => Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> Int -> a -> m a) -> m a +loopNextM !initial condition increment !initAcc f = go initial initAcc + where + go !step !acc + | condition step = + let !next = increment step + in f step next acc >>= go next + | otherwise = pure acc +{-# INLINE loopNextM #-} + +-- | Efficient Applicative loop. Result of each iteration is discarded. +-- +-- > loopA_ initial cond incr f === loopA initial cond incr (pure ()) (\i -> id <$ f i) +-- +-- @since 1.0.2 +loopA_ :: Applicative f => Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f a) -> f () +loopA_ !initial condition increment f = + loopF initial condition increment (pure ()) (\i ma -> f i *> ma) +{-# INLINE loopA_ #-} + +-- | Applicative loop. Use monadic `loopM` when possible, since it will be more efficient. +-- +-- @since 0.3.0 +loopA :: Applicative f => Int -> (Int -> Bool) -> (Int -> Int) -> f b -> (Int -> f (b -> b)) -> f b +loopA !initial condition increment lastAction f = + loopF initial condition increment lastAction (\i ma -> f i <*> ma) +{-# INLINE loopA #-} + +loopF :: Int -> (Int -> Bool) -> (Int -> Int) -> f a -> (Int -> f a -> f a) -> f a +loopF !initial condition increment lastAction f = go initial + where + go !step + | condition step = f step (go (increment step)) + | otherwise = lastAction +{-# INLINE loopF #-} + +nextMaybeF :: Int -> (Int -> Bool) -> (Int -> Int) -> (Maybe Int -> f a) -> f a +nextMaybeF !cur condition increment f = + let !i = increment cur + in f $! if condition i then Just i else Nothing +{-# INLINE nextMaybeF #-} + +-- | Similar to `loopM`, but way less efficient monadic loop with an accumulator that reverses +-- the direction of action application. eg: +-- +-- >>> loopDeepM 1 (< 20) (+ 2) [] (\i a -> Just (i:a)) +-- Just [1,3,5,7,9,11,13,15,17,19] +-- +-- Equivalent to: +-- +-- >>> loopM 19 (>= 1) (subtract 2) [] (\i a -> Just (i:a)) +-- Just [1,3,5,7,9,11,13,15,17,19] +-- +-- @since 0.1.0 +loopDeepM :: Monad m => Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a +loopDeepM !initial condition increment !initAcc f = + loopF initial condition increment (pure initAcc) (\i ma -> ma >>= f i) +{-# INLINE loopDeepM #-} + +-- | Divide length in chunks and apply a function to the computed results +-- +-- @since 0.2.1 +splitLinearly + :: Int + -- ^ Number of chunks + -> Int + -- ^ Total length + -> (Int -> Int -> a) + -- ^ Function that accepts a chunk length and slack start index + -> a +splitLinearly numChunks totalLength action = action chunkLength slackStart + where + !chunkLength = totalLength `quot` numChunks + !slackStart = chunkLength * numChunks +{-# INLINE splitLinearly #-} + +-- | Iterator that expects an action that accepts starting linear index as well as the ending +-- +-- @since 0.5.7 +splitLinearlyM_ + :: MonadPrimBase s m => Scheduler s () -> Int -> (Int -> Int -> m ()) -> m () +splitLinearlyM_ scheduler totalLength action = + splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do + loopNextA_ 0 (< slackStart) (+ chunkLength) $ \start next -> + scheduleWork_ scheduler $ action start next + when (slackStart < totalLength) $ + scheduleWork_ scheduler $ + action slackStart totalLength +{-# INLINE splitLinearlyM_ #-} + +-- | Iterator that expects an action that accepts starting linear index as well as the ending +-- +-- @since 1.0.2 +splitLinearlyM + :: MonadPrimBase s m => Scheduler s a -> Int -> (Int -> Int -> m a) -> m () +splitLinearlyM scheduler totalLength action = + splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do + loopNextA_ 0 (< slackStart) (+ chunkLength) $ \start next -> + scheduleWork scheduler (action start next) + when (slackStart < totalLength) $ + scheduleWork scheduler (action slackStart totalLength) +{-# INLINE splitLinearlyM #-} + +-- | Iterator that can be used to split computation amongst different workers. For monadic +-- generator see `splitLinearlyWithM_`. +-- +-- @since 0.2.1 +splitLinearlyWith_ + :: MonadPrimBase s m => Scheduler s () -> Int -> (Int -> b) -> (Int -> b -> m ()) -> m () +splitLinearlyWith_ scheduler totalLength index = + splitLinearlyWithM_ scheduler totalLength (pure . index) +{-# INLINE splitLinearlyWith_ #-} + +-- | Iterator that can be used to split computation jobs +-- +-- @since 0.2.6 +splitLinearlyWithM_ + :: MonadPrimBase s m => Scheduler s () -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m () +splitLinearlyWithM_ scheduler totalLength make write = + splitLinearlyM_ scheduler totalLength go + where + go start end = loopM_ start (< end) (+ 1) $ \k -> make k >>= write k + {-# INLINE go #-} +{-# INLINE splitLinearlyWithM_ #-} + +-- | Iterator that can be used to split computation jobs +-- +-- @since 0.3.0 +splitLinearlyWithStartAtM_ + :: MonadPrimBase s m => Scheduler s () -> Int -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m () +splitLinearlyWithStartAtM_ scheduler startAt totalLength make write = + splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do + loopM_ startAt (< (slackStart + startAt)) (+ chunkLength) $ \ !start -> + scheduleWork_ scheduler $ + loopM_ start (< (start + chunkLength)) (+ 1) $ + \ !k -> make k >>= write k + when (slackStart < totalLength) $ + scheduleWork_ scheduler $ + loopM_ (slackStart + startAt) (< (totalLength + startAt)) (+ 1) $ + \ !k -> make k >>= write k +{-# INLINE splitLinearlyWithStartAtM_ #-} + +-- | Iterator that can be used to split computation jobs, while using a stateful scheduler. +-- +-- @since 0.3.4 +splitLinearlyWithStatefulM_ + :: MonadUnliftIO m + => SchedulerWS ws () + -> Int + -- ^ Total linear length + -> (Int -> ws -> m b) + -- ^ Element producing action + -> (Int -> b -> m c) + -- ^ Element storing action + -> m () +splitLinearlyWithStatefulM_ schedulerWS totalLength make store = + let nWorkers = numWorkers (unwrapSchedulerWS schedulerWS) + in withRunInIO $ \run -> + splitLinearly nWorkers totalLength $ \chunkLength slackStart -> do + loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start -> + scheduleWorkState_ schedulerWS $ \s -> + loopM_ start (< (start + chunkLength)) (+ 1) $ \ !k -> + run (make k s >>= store k) + scheduleWorkState_ schedulerWS $ \s -> + loopM_ slackStart (< totalLength) (+ 1) $ \ !k -> + run (make k s >>= store k) +{-# INLINE splitLinearlyWithStatefulM_ #-} + +-- | This is a major helper function for fair splitting and parallelization of +-- work with ability to use some arbitrary accumulator and splittable seed +-- +-- @since 1.0.2 +splitWorkWithFactorST + :: Int + -- ^ Multiplying factor to be applied to number of workers for number + -- of jobs to schedule. Higher the factor, more jobs will be + -- scheduled. Only positive values are valid. + -> Scheduler s a + -> Int + -- ^ Starting index + -> Int + -- ^ Stepping value. Can be negative, but must not be zero. + -> Int + -- ^ Total number of steps to be taken + -> b + -- ^ Initial value for an accumulator + -> (b -> ST s (b, b)) + -- ^ An action to split accumulator for multiple threads + -> (Int -> Int -> Int -> Int -> b -> ST s a) + -- ^ A job to be scheduled. Accepts: + -- + -- * Chunk index start + -- * Chunk length + -- * Chunk start index adjusted for supplied start and stepping value + -- * Chunk stop index adjusted for supplied start and stepping value + -> ST s b +splitWorkWithFactorST fact scheduler start step totalLength initAcc splitAcc action = do + let !(chunkLength, slackStart) = splitNumChunks fact (numWorkers scheduler) totalLength + slackAcc <- + loopM 0 (< slackStart) (+ chunkLength) initAcc $ \ !chunkStart !acc -> do + (accCur, accNext) <- splitAcc acc + scheduleMassivWork scheduler $ do + let !chunkStartAdj = start + chunkStart * step + !chunkStopAdj = chunkStartAdj + chunkLength * step + action chunkStart chunkLength chunkStartAdj chunkStopAdj accCur + pure accNext + let !slackLength = totalLength - slackStart + if slackLength > 0 + then do + (curAcc, nextAcc) <- splitAcc slackAcc + scheduleMassivWork scheduler $ do + let !slackStartAdj = start + slackStart * step + !slackStopAdj = slackStartAdj + slackLength * step + action slackStart slackLength slackStartAdj slackStopAdj curAcc + pure nextAcc + else pure slackAcc +{-# INLINE splitWorkWithFactorST #-} + +-- | Linear iterator that supports multiplying factor +-- +-- @since 1.0.2 +iterLinearST_ + :: Int + -> Scheduler s () + -> Int + -> Int + -> Int + -> (Int -> ST s a) + -> ST s () +iterLinearST_ fact scheduler start step n action = do + let totalLength = (n - start) `quot` step + splitWorkWithFactorST fact scheduler start step totalLength () (\_ -> pure ((), ())) $ + \_ _ chunkStartAdj chunkStopAdj _ -> + loopA_ chunkStartAdj (< chunkStopAdj) (+ step) action +{-# INLINE iterLinearST_ #-} + +-- | Linear iterator that supports multiplying factor and accumulator, but the results are discarded. +-- +-- @since 1.0.2 +iterLinearAccST_ + :: Int + -> Scheduler s () + -> Int + -> Int + -> Int + -> a + -> (a -> ST s (a, a)) + -> (Int -> a -> ST s a) + -> ST s () +iterLinearAccST_ fact scheduler start step n initAcc splitAcc action = do + let totalLength = (n - start) `quot` step + void $ + splitWorkWithFactorST fact scheduler start step totalLength initAcc splitAcc $ + \_ _ chunkStartAdj chunkStopAdj accCur -> + void $ loopM chunkStartAdj (< chunkStopAdj) (+ step) accCur action +{-# INLINE iterLinearAccST_ #-} + +-- | Linear iterator that supports multiplying factor and accumulator. Results +-- of actions are stored in the scheduler. +-- +-- @since 1.0.2 +iterLinearAccST + :: Int + -> Scheduler s a + -> Int + -> Int + -- ^ Step. Must be non-zero + -> Int + -> a + -> (a -> ST s (a, a)) + -> (Int -> a -> ST s a) + -> ST s a +iterLinearAccST fact scheduler start step n initAcc splitAcc action = do + let totalLength = (n - start) `quot` step + splitWorkWithFactorST fact scheduler start step totalLength initAcc splitAcc $ + \_ _ chunkStartAdj chunkStopAdj accCur -> + loopM chunkStartAdj (< chunkStopAdj) (+ step) accCur action +{-# INLINE iterLinearAccST #-} + +-- | Helper for figuring out the chunk length and slack start +splitNumChunks :: Int -> Int -> Int -> (Int, Int) +splitNumChunks fact nw totalLength = + let maxNumChunks = nw * max 1 fact + !numChunks + | nw == 1 || totalLength <= 0 = 1 -- Optimize for Seq and avoid `quot` by 0. + | totalLength <= nw = totalLength + | totalLength >= maxNumChunks = maxNumChunks + | otherwise = nw + !chunkLength = totalLength `quot` numChunks + !slackStart = chunkLength * numChunks + in (chunkLength, slackStart) + +-- | Helper for adjusting stride of a chunk +stepStartAdjust :: Int -> Int -> Int +stepStartAdjust step ix = ix + ((step - (ix `mod` step)) `mod` step) +{-# INLINE stepStartAdjust #-} + +-- | Internal version of a `scheduleWork` that will be replaced by +-- `scheduleWork_` by the compiler whenever action produces `()` +scheduleMassivWork :: PrimBase m => Scheduler (PrimState m) a -> m a -> m () +scheduleMassivWork = scheduleWork +{-# INLINE [0] scheduleMassivWork #-} + +{-# RULES +"scheduleWork/scheduleWork_/ST" forall (scheduler :: Scheduler s ()) (action :: ST s ()). + scheduleMassivWork scheduler action = + scheduleWork_ scheduler action +"scheduleWork/scheduleWork_/IO" forall (scheduler :: Scheduler RealWorld ()) (action :: IO ()). + scheduleMassivWork scheduler action = + scheduleWork_ scheduler action + #-} + +-- | Selects an optimal scheduler for the supplied strategy, but it works only in `IO` +-- +-- @since 1.0.0 +withMassivScheduler_ :: Comp -> (Scheduler RealWorld () -> IO ()) -> IO () +withMassivScheduler_ comp f = + case comp of + Par -> withGlobalScheduler_ globalScheduler f + Seq -> f trivialScheduler_ + _ -> withScheduler_ comp f +{-# INLINE withMassivScheduler_ #-} diff --git a/src/Data/Massiv/Core/Operations.hs b/src/Data/Massiv/Core/Operations.hs --- a/src/Data/Massiv/Core/Operations.hs +++ b/src/Data/Massiv/Core/Operations.hs @@ -1,43 +1,92 @@ -{-# LANGUAGE DefaultSignatures #-} +{-# LANGUAGE BangPatterns #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} + -- | -- Module : Data.Massiv.Core.Operations --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2019-2022 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable -module Data.Massiv.Core.Operations - ( Numeric(..) - , NumericFloat(..) - ) where +module Data.Massiv.Core.Operations ( + FoldNumeric (..), + defaultPowerSumArray, + defaultUnsafeDotProduct, + defaultFoldArray, + Numeric (..), + defaultUnsafeLiftArray, + defaultUnsafeLiftArray2, + NumericFloat (..), +) where import Data.Massiv.Core.Common --- import Data.Massiv.Array.Ops.Fold.Internal +class (Size r, Num e) => FoldNumeric r e where + {-# MINIMAL foldArray, powerSumArray, unsafeDotProduct #-} -class Num e => Numeric r e where + -- | Compute sum of all elements in the array + -- + -- @since 0.5.6 + sumArray :: Index ix => Array r ix e -> e + sumArray = foldArray (+) 0 + {-# INLINE sumArray #-} - {-# MINIMAL unsafeLiftArray, unsafeLiftArray2 #-} + -- | Compute product of all elements in the array + -- + -- @since 0.5.6 + productArray :: Index ix => Array r ix e -> e + productArray = foldArray (*) 1 + {-# INLINE productArray #-} - -- sumArray :: Array r Ix1 e -> e - -- default sumArray :: Source r Ix1 e => Array r Ix1 e -> e - -- sumArray = foldlS (+) 0 - -- {-# INLINE sumArray #-} + -- | Raise each element in the array to some non-negative power and sum the results + -- + -- @since 0.5.7 + powerSumArray :: Index ix => Array r ix e -> Int -> e - -- productArray :: Array r Ix1 e -> e - -- default productArray :: Source r Ix1 e => Array r Ix1 e -> e - -- productArray = foldlS (*) 1 - -- {-# INLINE productArray #-} + -- | Compute dot product without any extraneous checks + -- + -- @since 0.5.6 + unsafeDotProduct :: Index ix => Array r ix e -> Array r ix e -> e - -- -- | Raise each element in the array to some non-negative power and sum the results - -- powerSumArray :: Array r Ix1 e -> Int -> e + -- | Fold over an array + -- + -- @since 0.5.6 + foldArray :: Index ix => (e -> e -> e) -> e -> Array r ix e -> e - -- unsafeDotProduct :: Array r Ix1 e -> Array r Ix1 e -> e +defaultUnsafeDotProduct + :: (Num e, Index ix, Source r e) => Array r ix e -> Array r ix e -> e +defaultUnsafeDotProduct a1 a2 = go 0 0 + where + !len = totalElem (size a1) + go !acc i + | i < len = go (acc + unsafeLinearIndex a1 i * unsafeLinearIndex a2 i) (i + 1) + | otherwise = acc +{-# INLINE defaultUnsafeDotProduct #-} +defaultPowerSumArray :: (Index ix, Source r e, Num e) => Array r ix e -> Int -> e +defaultPowerSumArray arr p = go 0 0 + where + !len = totalElem (size arr) + go !acc i + | i < len = go (acc + unsafeLinearIndex arr i ^ p) (i + 1) + | otherwise = acc +{-# INLINE defaultPowerSumArray #-} + +defaultFoldArray :: (Index ix, Source r e) => (e -> e -> e) -> e -> Array r ix e -> e +defaultFoldArray f !initAcc arr = go initAcc 0 + where + !len = totalElem (size arr) + go !acc i + | i < len = go (f acc (unsafeLinearIndex arr i)) (i + 1) + | otherwise = acc +{-# INLINE defaultFoldArray #-} + +class FoldNumeric r e => Numeric r e where + {-# MINIMAL unsafeLiftArray, unsafeLiftArray2 #-} + plusScalar :: Index ix => Array r ix e -> e -> Array r ix e plusScalar arr e = unsafeLiftArray (+ e) arr {-# INLINE plusScalar #-} @@ -46,6 +95,10 @@ minusScalar arr e = unsafeLiftArray (subtract e) arr {-# INLINE minusScalar #-} + scalarMinus :: Index ix => e -> Array r ix e -> Array r ix e + scalarMinus e arr = unsafeLiftArray (e -) arr + {-# INLINE scalarMinus #-} + multiplyScalar :: Index ix => Array r ix e -> e -> Array r ix e multiplyScalar arr e = unsafeLiftArray (* e) arr {-# INLINE multiplyScalar #-} @@ -66,24 +119,44 @@ multiplicationPointwise = unsafeLiftArray2 (*) {-# INLINE multiplicationPointwise #-} + -- TODO: + -- - rename to powerScalar + -- - add? powerPointwise :: Array r ix e -> Array r ix Int -> Array r ix e + -- | Raise each element of the array to the power powerPointwise :: Index ix => Array r ix e -> Int -> Array r ix e powerPointwise arr pow = unsafeLiftArray (^ pow) arr {-# INLINE powerPointwise #-} - - unsafeLiftArray :: Index ix => (a -> e) -> Array r ix a -> Array r ix e + unsafeLiftArray :: Index ix => (e -> e) -> Array r ix e -> Array r ix e - unsafeLiftArray2 :: Index ix => (a -> b -> e) -> Array r ix a -> Array r ix b -> Array r ix e + unsafeLiftArray2 :: Index ix => (e -> e -> e) -> Array r ix e -> Array r ix e -> Array r ix e +defaultUnsafeLiftArray + :: (Load r ix e, Source r e) => (e -> e) -> Array r ix e -> Array r ix e +defaultUnsafeLiftArray f arr = makeArrayLinear (getComp arr) (size arr) (f . unsafeLinearIndex arr) +{-# INLINE defaultUnsafeLiftArray #-} +defaultUnsafeLiftArray2 + :: (Load r ix e, Source r e) + => (e -> e -> e) + -> Array r ix e + -> Array r ix e + -> Array r ix e +defaultUnsafeLiftArray2 f a1 a2 = + makeArrayLinear (getComp a1 <> getComp a2) (size a1) $ \ !i -> + f (unsafeLinearIndex a1 i) (unsafeLinearIndex a2 i) +{-# INLINE defaultUnsafeLiftArray2 #-} class (Numeric r e, Floating e) => NumericFloat r e where - divideScalar :: Index ix => Array r ix e -> e -> Array r ix e divideScalar arr e = unsafeLiftArray (/ e) arr {-# INLINE divideScalar #-} + scalarDivide :: Index ix => e -> Array r ix e -> Array r ix e + scalarDivide e arr = unsafeLiftArray (e /) arr + {-# INLINE scalarDivide #-} + divisionPointwise :: Index ix => Array r ix e -> Array r ix e -> Array r ix e divisionPointwise = unsafeLiftArray2 (/) {-# INLINE divisionPointwise #-} @@ -96,14 +169,13 @@ sqrtPointwise = unsafeLiftArray sqrt {-# INLINE sqrtPointwise #-} - -- floorPointwise :: (Index ix, Integral a) => Array r ix e -> Array r ix a - -- floorPointwise = unsafeLiftArray floor - -- {-# INLINE floorPointwise #-} - - -- ceilingPointwise :: (Index ix, Integral a) => Array r ix e -> Array r ix a - -- ceilingPointwise = unsafeLiftArray ceiling - -- {-# INLINE ceilingPointwise #-} +-- floorPointwise :: (Index ix, Integral a) => Array r ix e -> Array r ix a +-- floorPointwise = unsafeLiftArray floor +-- {-# INLINE floorPointwise #-} +-- ceilingPointwise :: (Index ix, Integral a) => Array r ix e -> Array r ix a +-- ceilingPointwise = unsafeLiftArray ceiling +-- {-# INLINE ceilingPointwise #-} -- class Equality r e where @@ -111,7 +183,6 @@ -- unsafeEqPointwise :: Index ix => Array r ix e -> Array r ix e -> Array r ix Bool - -- class Relation r e where -- unsafePointwiseLT :: Array r ix e -> Array r ix e -> Array r ix Bool @@ -126,5 +197,3 @@ -- unsafeMinimum :: Array r ix e -> e -- unsafeMaximum :: Array r ix e -> e - - diff --git a/src/Data/Massiv/Vector.hs b/src/Data/Massiv/Vector.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Vector.hs @@ -0,0 +1,2840 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE ExplicitForAll #-} +{-# LANGUAGE FlexibleContexts #-} +{-# OPTIONS_GHC -fno-warn-duplicate-exports #-} + +-- | +-- Module : Data.Massiv.Vector +-- Copyright : (c) Alexey Kuleshevich 2020-2022 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +module Data.Massiv.Vector ( + Vector, + MVector, + + -- * Accessors + + -- *** Size + slength, + maxLinearSize, + size, + isNull, + isNotNull, + + -- *** Indexing + (!?), + (!), + index, + index', + head', + shead', + last', + + -- *** Monadic Indexing + indexM, + headM, + sheadM, + lastM, + unconsM, + unsnocM, + + -- ** Slicing + slice, + slice', + sliceM, + sslice, + sliceAt, + sliceAt', + sliceAtM, + + -- *** Init + init, + init', + initM, + + -- *** Tail + tail, + tail', + tailM, + + -- *** Take + take, + take', + takeM, + takeWhile, + stake, + + -- *** Drop + drop, + dropWhile, + drop', + dropM, + sdrop, + + -- * Construction + + -- ** Initialization + empty, + sempty, + singleton, + ssingleton, + cons, + snoc, + A.replicate, + sreplicate, + generate, + sgenerate, + -- , iterateN + -- , iiterateN + siterate, + siterateN, + + -- ** Monadic initialization + sreplicateM, + sgenerateM, + siterateNM, + -- , create + -- , createT + + -- ** Unfolding + sunfoldr, + sunfoldrM, + sunfoldrN, + sunfoldrNM, + sunfoldrExactN, + sunfoldrExactNM, + -- , constructN + -- , constructrN + + -- ** Enumeration + (...), + (..:), + enumFromN, + senumFromN, + enumFromStepN, + senumFromStepN, + + -- ** Concatenation + + -- , consS -- cons + -- , snocS -- snoc + sappend, -- (++) + sconcat, -- concat + -- -- ** Restricitng memory usage + -- , force + -- -- * Modifying + -- -- ** Bulk updates + -- , (//) + -- , update_ + -- -- ** Accumulations + -- , accum + -- , accumulate_ + -- -- ** Permutations + -- , reverse + -- , backpermute + -- -- ** Manifest updates + -- , modify + -- -- * Elementwise + -- -- ** Mapping + smap, + simap, + -- , sconcatMap + + -- ** Monadic mapping + straverse, + sitraverse, + smapM, + smapM_, + simapM, + simapM_, + sforM, + sforM_, + siforM, + siforM_, + + -- ** Zipping + szip, + szip3, + szip4, + szip5, + szip6, + szipWith, + szipWith3, + szipWith4, + szipWith5, + szipWith6, + sizipWith, + sizipWith3, + sizipWith4, + sizipWith5, + sizipWith6, + + -- ** Monadic zipping + szipWithM, + szipWith3M, + szipWith4M, + szipWith5M, + szipWith6M, + sizipWithM, + sizipWith3M, + sizipWith4M, + sizipWith5M, + sizipWith6M, + szipWithM_, + szipWith3M_, + szipWith4M_, + szipWith5M_, + szipWith6M_, + sizipWithM_, + sizipWith3M_, + sizipWith4M_, + sizipWith5M_, + sizipWith6M_, + + -- * Predicates + + -- ** Filtering + sfilter, + sifilter, + sfilterM, + sifilterM, + -- , uniq -- sunique? + smapMaybe, + smapMaybeM, + scatMaybes, + simapMaybe, + simapMaybeM, + -- , stakeWhile + -- , sdropWhile + -- -- ** Partitioning + -- , partition + -- , unstablePartition + -- , partitionWith + -- , span + -- , break + -- -- ** Searching + -- , elem + -- , notElem + -- , find + findIndex, + -- , findIndices + -- , elemIndex + -- , elemIndices + + -- * Folding + sfoldl, + sfoldlM, + sfoldlM_, + sifoldl, + sifoldlM, + sifoldlM_, + sfoldl1', + sfoldl1M, + sfoldl1M_, + + -- ** Specialized folds + sor, + sand, + sall, + sany, + ssum, + sproduct, + smaximum', + smaximumM, + -- , maximumBy + sminimum', + sminimumM, + -- , minimumBy + -- , minIndex + -- , minIndexBy + -- , maxIndex + -- , maxIndexBy + + -- ** Scanning + sprescanl, + spostscanl, + spostscanlAcc, + sscanl, + sscanl1, + -- sprescanr, + -- spostscanr, + -- sscanr, + -- sscanr1, + + -- * Conversions + + -- ** Lists + stoList, + fromList, + sfromList, + sfromListN, + + -- * Computation + compute, + computeS, + computeIO, + computePrimM, + computeAs, + computeProxy, + computeSource, + computeWithStride, + computeWithStrideAs, + clone, + convert, + convertAs, + convertProxy, + + -- ** Re-exports + module Data.Massiv.Core, + module Data.Massiv.Array.Delayed, + module Data.Massiv.Array.Manifest, + module Data.Massiv.Array.Mutable, +) where + +import Control.Monad hiding (filterM, replicateM) +import Data.Coerce +import Data.Massiv.Array.Delayed +import Data.Massiv.Array.Delayed.Pull +import Data.Massiv.Array.Delayed.Push +import Data.Massiv.Array.Delayed.Stream +import Data.Massiv.Array.Manifest +import Data.Massiv.Array.Manifest.Internal +import Data.Massiv.Array.Manifest.List (fromList) +import Data.Massiv.Array.Mutable +import Data.Massiv.Array.Ops.Construct +import qualified Data.Massiv.Array.Ops.Construct as A (replicate) +import Data.Massiv.Core +import Data.Massiv.Core.Common +import qualified Data.Massiv.Vector.Stream as S +import Data.Massiv.Vector.Unsafe +import Data.Maybe +import Prelude hiding ( + drop, + dropWhile, + init, + length, + null, + replicate, + splitAt, + tail, + take, + takeWhile, + ) + +-- ========= -- +-- Accessors -- +-- ========= -- + +------------------------ +-- Length information -- +------------------------ + +-- | /O(1)/ - Get the length of a `Stream` array, but only if it is known exactly in +-- constant time without looking at any of the elements in the array. +-- +-- /Related/: `maxLinearSize`, `size`, `elemsCount` and `totalElem` +-- +-- ==== __Examples__ +-- +-- >>> slength $ sfromList [] +-- Nothing +-- >>> slength $ sreplicate 5 () +-- Just (Sz1 5) +-- >>> slength $ makeArrayLinearR D Seq (Sz1 5) id +-- Just (Sz1 5) +-- >>> slength $ sunfoldr (\x -> Just (x, x)) (0 :: Int) +-- Nothing +-- >>> slength $ sunfoldrN 10 (\x -> Just (x, x)) (0 :: Int) +-- Nothing +-- >>> slength $ sunfoldrExactN 10 (\x -> (x, x)) (0 :: Int) +-- Just (Sz1 10) +-- +-- /__Similar__/: +-- +-- [@Data.Foldable.`Data.Foldable.length`@] For some data structures, like a list for +-- example, it is an /O(n)/ operation, because there is a need to evaluate the full spine +-- and possibly even the elements in order to get the full length. With `Stream` vectors +-- that is not always the case. +-- +-- [@Data.Vector.Generic.`Data.Vector.Generic.length`@] In the vector package this +-- function will always break fusion, unless it is the only operation that is applied to +-- the vector. +-- +-- @since 0.5.0 +slength + :: forall r ix e + . Stream r ix e + => Array r ix e + -> Maybe Sz1 +slength v = + case stepsSize (toStream v) of + LengthExact sz -> Just sz + _ -> Nothing +{-# INLINE slength #-} + +-------------- +-- Indexing -- +-------------- + +-- | /O(1)/ - Get the first element of a `Source` vector. Throws an error on empty. +-- +-- /Related/: 'shead'', `headM`, `sheadM`, `unconsM`. +-- +-- ==== __Examples__ +-- +-- >>> head' (Ix1 10 ..: 10000000000000) +-- 10 +-- +-- /__Similar__/: +-- +-- [@Data.List.`Data.List.head`@] Also constant time and partial. Fusion is broken if +-- there other consumers of the list. +-- +-- [@Data.Vector.Generic.`Data.Vector.Generic.head`@] Also constant time and partial. Will +-- cause materialization of the full vector if any other function is applied to the vector. +-- +-- @since 0.5.0 +head' + :: forall r e + . (HasCallStack, Source r e) + => Vector r e + -> e +head' = throwEither . headM +{-# INLINE head' #-} + +-- | /O(1)/ - Get the first element of a `Source` vector. +-- +-- /Related/: 'head'', 'shead'', `sheadM`, `unconsM`. +-- +-- /__Throws Exceptions__/: `SizeEmptyException` when array is empty +-- +-- ==== __Examples__ +-- +-- >>> headM (Ix1 10 ..: 10000000000000) +-- 10 +-- >>> headM (Ix1 10 ..: 10000000000000) :: Maybe Int +-- Just 10 +-- >>> headM (empty :: Array D Ix1 Int) :: Maybe Int +-- Nothing +-- >>> either show (const "") $ headM (Ix1 10 ..: 10) +-- "SizeEmptyException: (Sz1 0) corresponds to an empty array" +-- +-- /__Similar__/: +-- +-- [@Data.Maybe.`Data.Maybe.listToMaybe`@] It also a safe way to get the head of the list, +-- except it is restricted to `Maybe` +-- +-- @since 0.5.0 +headM + :: forall r e m + . (Source r e, MonadThrow m) + => Vector r e + -> m e +headM v + | elemsCount v == 0 = throwM $ SizeEmptyException (size v) + | otherwise = pure $ unsafeLinearIndex v 0 +{-# INLINE headM #-} + +-- | /O(1)/ - Get the first element of a `Stream` vector. Throws an error on empty. +-- +-- /Related/: 'head'', `headM`, `sheadM`, `unconsM`. +-- +-- ==== __Examples__ +-- +-- >>> shead' $ sunfoldr (\x -> Just (x, x)) (0 :: Int) +-- 0 +-- >>> shead' (Ix1 3 ... 5) +-- 3 +-- +-- @since 0.5.0 +shead' + :: forall r e + . (HasCallStack, Stream r Ix1 e) + => Vector r e + -> e +shead' = throwEither . sheadM +{-# INLINE shead' #-} + +-- | /O(1)/ - Get the first element of a `Stream` vector. +-- +-- /Related/: 'head'', 'shead'', `headM`, `unconsM`. +-- +-- /__Throws Exceptions__/: `SizeEmptyException` +-- +-- ==== __Examples__ +-- +-- >>> maybe 101 id $ sheadM (empty :: Vector D Int) +-- 101 +-- >>> maybe 101 id $ sheadM (singleton 202 :: Vector D Int) +-- 202 +-- >>> sheadM $ sunfoldr (\x -> Just (x, x)) (0 :: Int) +-- 0 +-- >>> x <- sheadM $ sunfoldr (\_ -> Nothing) (0 :: Int) +-- *** Exception: SizeEmptyException: (Sz1 0) corresponds to an empty array +-- +-- @since 0.5.0 +sheadM + :: forall r e m + . (Stream r Ix1 e, MonadThrow m) + => Vector r e + -> m e +sheadM v = + case S.unId (S.headMaybe (toStream v)) of + Nothing -> throwM $ SizeEmptyException (zeroSz :: Sz1) + Just e -> pure e +{-# INLINE sheadM #-} + +-- | /O(1)/ - Take one element off of the `Source` vector from the left side, as well as +-- the remaining part of the vector in delayed `D` representation. +-- +-- /Related/: 'head'', 'shead'', `headM`, `sheadM`, `cons` +-- +-- /__Throws Exceptions__/: `SizeEmptyException` +-- +-- ==== __Examples__ +-- +-- >>> unconsM (fromList Seq [1,2,3] :: Array P Ix1 Int) +-- (1,Array P Seq (Sz1 2) +-- [ 2, 3 ]) +-- +-- /__Similar__/: +-- +-- [@Data.List.`Data.List.uncons`@] Same concept, except it is restricted to `Maybe` instead of +-- the more general `MonadThrow` +-- +-- @since 0.3.0 +unconsM + :: forall r e m + . (MonadThrow m, Source r e) + => Vector r e + -> m (e, Vector r e) +unconsM arr + | 0 == totalElem sz = throwM $ SizeEmptyException sz + | otherwise = pure (unsafeLinearIndex arr 0, unsafeLinearSlice 1 (SafeSz (unSz sz - 1)) arr) + where + !sz = size arr +{-# INLINE unconsM #-} + +-- | /O(1)/ - Take one element off of the vector from the right side, as well as the +-- remaining part of the vector. +-- +-- /Related/: 'last'', `lastM`, `snoc` +-- +-- /__Throws Exceptions__/: `SizeEmptyException` +-- +-- ==== __Examples__ +-- +-- >>> unsnocM (fromList Seq [1,2,3] :: Array P Ix1 Int) +-- (Array P Seq (Sz1 2) +-- [ 1, 2 ],3) +-- +-- @since 0.3.0 +unsnocM + :: forall r e m + . (MonadThrow m, Source r e) + => Vector r e + -> m (Vector r e, e) +unsnocM arr + | 0 == totalElem sz = throwM $ SizeEmptyException sz + | otherwise = pure (unsafeLinearSlice 0 (SafeSz k) arr, unsafeLinearIndex arr k) + where + !sz = size arr + !k = unSz sz - 1 +{-# INLINE unsnocM #-} + +-- | /O(1)/ - Get the last element of a `Source` vector. Throws an error on empty. +-- +-- /Related/: `lastM`, `unsnocM` +-- +-- ==== __Examples__ +-- +-- >>> last' (Ix1 10 ... 10000000000000) +-- 10000000000000 +-- +-- /__Similar__/: +-- +-- [@Data.List.`Data.List.last`@] Also partial, but it has /O(n)/ complexity. Fusion is +-- broken if there other consumers of the list. +-- +-- [@Data.Vector.Generic.`Data.Vector.Generic.last`@] Also constant time and partial. Will +-- cause materialization of the full vector if any other function is applied to the vector. +-- +-- @since 0.5.0 +last' :: forall r e. (HasCallStack, Source r e) => Vector r e -> e +last' = throwEither . lastM +{-# INLINE last' #-} + +-- | /O(1)/ - Get the last element of a `Source` vector. +-- +-- /Related/: 'last'', `unsnocM` +-- +-- /__Throws Exceptions__/: `SizeEmptyException` +-- +-- ==== __Examples__ +-- +-- >>> lastM (Ix1 10 ... 10000000000000) +-- 10000000000000 +-- >>> lastM (Ix1 10 ... 10000000000000) :: Maybe Int +-- Just 10000000000000 +-- >>> either show (const "") $ lastM (fromList Seq [] :: Array P Ix1 Int) +-- "SizeEmptyException: (Sz1 0) corresponds to an empty array" +-- +-- @since 0.5.0 +lastM :: forall r e m. (Source r e, MonadThrow m) => Vector r e -> m e +lastM v + | k == 0 = throwM $ SizeEmptyException (size v) + | otherwise = pure $ unsafeLinearIndex v (k - 1) + where + k = unSz (size v) +{-# INLINE lastM #-} + +-- | /O(1)/ - Take a slice of a `Source` vector. Never fails, instead adjusts the indices. +-- +-- ==== __Examples__ +-- +-- >>> slice 10 5 (Ix1 0 ... 10000000000000) +-- Array D Seq (Sz1 5) +-- [ 10, 11, 12, 13, 14 ] +-- >>> slice (-10) 5 (Ix1 0 ... 10000000000000) +-- Array D Seq (Sz1 5) +-- [ 0, 1, 2, 3, 4 ] +-- >>> slice 9999999999998 50 (Ix1 0 ... 10000000000000) +-- Array D Seq (Sz1 3) +-- [ 9999999999998, 9999999999999, 10000000000000 ] +-- +-- @since 0.5.0 +slice :: forall r e. Source r e => Ix1 -> Sz1 -> Vector r e -> Vector r e +slice !i (Sz k) v = unsafeLinearSlice i' newSz v + where + !i' = min n (max 0 i) + !newSz = SafeSz (min (n - i') k) + Sz n = size v +{-# INLINE slice #-} + +-- | /O(1)/ - Take a slice of a `Source` vector. Throws an error on incorrect indices. +-- +-- ==== __Examples__ +-- +-- >>> slice' 10 5 (Ix1 0 ... 100) +-- Array D Seq (Sz1 5) +-- [ 10, 11, 12, 13, 14 ] +-- >>> slice' 9999999999998 3 (Ix1 0 ... 10000000000000) +-- Array D Seq (Sz1 3) +-- [ 9999999999998, 9999999999999, 10000000000000 ] +-- +-- @since 0.5.0 +slice' :: forall r e. (HasCallStack, Source r e) => Ix1 -> Sz1 -> Vector r e -> Vector r e +slice' i k = throwEither . sliceM i k +{-# INLINE slice' #-} + +-- | /O(1)/ - Take a slice of a `Source` vector. Throws an error on incorrect indices. +-- +-- /__Throws Exceptions__/: `SizeSubregionException` +-- +-- ==== __Examples__ +-- +-- >>> sliceM 10 5 (Ix1 0 ... 100) +-- Array D Seq (Sz1 5) +-- [ 10, 11, 12, 13, 14 ] +-- >>> sliceM (-10) 5 (Ix1 0 ... 100) +-- *** Exception: SizeSubregionException: (Sz1 101) is to small for -10 (Sz1 5) +-- >>> sliceM 98 50 (Ix1 0 ... 100) +-- *** Exception: SizeSubregionException: (Sz1 101) is to small for 98 (Sz1 50) +-- >>> sliceM 9999999999998 3 (Ix1 0 ... 10000000000000) +-- Array D Seq (Sz1 3) +-- [ 9999999999998, 9999999999999, 10000000000000 ] +-- +-- @since 0.5.0 +sliceM + :: forall r e m + . (Source r e, MonadThrow m) + => Ix1 + -- ^ Starting index + -> Sz1 + -- ^ Number of elements to take from the Source vector + -> Vector r e + -- ^ Source vector to take a slice from + -> m (Vector r e) +sliceM i newSz@(Sz k) v + | i >= 0 && k <= n - i = pure $ unsafeLinearSlice i newSz v + | otherwise = throwM $ SizeSubregionException sz i newSz + where + sz@(Sz n) = size v +{-# INLINE sliceM #-} + +-- | Take a slice of a `Stream` vector. Never fails, instead adjusts the indices. +-- +-- ==== __Examples__ +-- +-- >>> sslice 10 5 (Ix1 0 ... 10000000000000) +-- Array DS Seq (Sz1 5) +-- [ 10, 11, 12, 13, 14 ] +-- >>> sslice 10 5 (sfromList [0 :: Int .. ]) +-- Array DS Seq (Sz1 5) +-- [ 10, 11, 12, 13, 14 ] +-- >>> sslice (-10) 5 (Ix1 0 ... 10000000000000) +-- Array DS Seq (Sz1 5) +-- [ 0, 1, 2, 3, 4 ] +-- +-- Unlike `slice` it has to iterate through each element until the staring index is reached, +-- therefore something like @sslice 9999999999998 50 (Ix1 0 ... 10000000000000)@ will not +-- be feasable. +-- +-- >>> import System.Timeout (timeout) +-- >>> let smallArr = sslice 9999999999998 50 (Ix1 0 ... 10000000000000) +-- >>> timeout 500000 (computeIO smallArr :: IO (Array P Ix1 Int)) +-- Nothing +-- +-- @since 0.5.0 +sslice + :: forall r e + . Stream r Ix1 e + => Ix1 + -- ^ Starting index + -> Sz1 + -- ^ Number of elements to take from the stream vector + -> Vector r e + -- ^ Stream vector to take a slice from + -> Vector DS e +sslice !i !k = fromSteps . S.slice i k . S.toStream +{-# INLINE sslice #-} + +-- | /O(1)/ - Get a vector without the last element. Never fails. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> A.init (0 ..: 10) +-- Array D Seq (Sz1 9) +-- [ 0, 1, 2, 3, 4, 5, 6, 7, 8 ] +-- >>> A.init (empty :: Array D Ix1 Int) +-- Array D Seq (Sz1 0) +-- [ ] +-- +-- @since 0.5.0 +init :: forall r e. Source r e => Vector r e -> Vector r e +init v = unsafeLinearSlice 0 (Sz (coerce (size v) - 1)) v +{-# INLINE init #-} + +-- | /O(1)/ - Get a vector without the last element. Throws an error on empty +-- +-- ==== __Examples__ +-- +-- >>> init' (0 ..: 10) +-- Array D Seq (Sz1 9) +-- [ 0, 1, 2, 3, 4, 5, 6, 7, 8 ] +-- +-- @since 0.5.0 +init' :: forall r e. (HasCallStack, Source r e) => Vector r e -> Vector r e +init' = throwEither . initM +{-# INLINE init' #-} + +-- | /O(1)/ - Get a vector without the last element. Throws an error on empty +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> initM (0 ..: 10) +-- Array D Seq (Sz1 9) +-- [ 0, 1, 2, 3, 4, 5, 6, 7, 8 ] +-- >>> maybe 0 A.sum $ initM (0 ..: 10) +-- 36 +-- >>> maybe 0 A.sum $ initM (empty :: Array D Ix1 Int) +-- 0 +-- +-- @since 0.5.0 +initM :: forall r e m. (Source r e, MonadThrow m) => Vector r e -> m (Vector r e) +initM v = do + when (elemsCount v == 0) $ throwM $ SizeEmptyException $ size v + pure $ unsafeInit v +{-# INLINE initM #-} + +-- | /O(1)/ - Get a vector without the first element. Never fails +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> A.tail (0 ..: 10) +-- Array D Seq (Sz1 9) +-- [ 1, 2, 3, 4, 5, 6, 7, 8, 9 ] +-- >>> A.tail (empty :: Array D Ix1 Int) +-- Array D Seq (Sz1 0) +-- [ ] +-- +-- @since 0.5.0 +tail :: forall r e. Source r e => Vector r e -> Vector r e +tail = drop oneSz +{-# INLINE tail #-} + +-- | /O(1)/ - Get a vector without the first element. Throws an error on empty +-- +-- ==== __Examples__ +-- +-- λ> tail' (0 ..: 10) +-- Array D Seq (Sz1 9) +-- [ 1, 2, 3, 4, 5, 6, 7, 8, 9 ] +-- λ> tail' (empty :: Array D Ix1 Int) +-- Array D *** Exception: SizeEmptyException: (Sz1 0) corresponds to an empty array +-- +-- @since 0.5.0 +tail' :: forall r e. (HasCallStack, Source r e) => Vector r e -> Vector r e +tail' = throwEither . tailM +{-# INLINE tail' #-} + +-- | /O(1)/ - Get the vector without the first element. Throws an error on empty +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> tailM (0 ..: 10) +-- Array D Seq (Sz1 9) +-- [ 1, 2, 3, 4, 5, 6, 7, 8, 9 ] +-- >>> maybe 0 A.sum $ tailM (0 ..: 10) +-- 45 +-- >>> maybe 0 A.sum $ tailM (empty :: Array D Ix1 Int) +-- 0 +-- +-- @since 0.5.0 +tailM :: forall r e m. (Source r e, MonadThrow m) => Vector r e -> m (Vector r e) +tailM v = do + when (elemsCount v == 0) $ throwM $ SizeEmptyException $ size v + pure $ unsafeTail v +{-# INLINE tailM #-} + +-- | /O(1)/ - Take first @n@ elements from a vector. This function never fails and has +-- similar semantics as the `Data.List.take` for lists. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> A.take 5 (0 ..: 10) +-- Array D Seq (Sz1 5) +-- [ 0, 1, 2, 3, 4 ] +-- >>> A.take 0 (0 ..: 10) +-- Array D Seq (Sz1 0) +-- [ ] +-- >>> A.take 100 (0 ..: 10) +-- Array D Seq (Sz1 10) +-- [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9 ] +-- +-- @since 0.5.0 +take :: Source r e => Sz1 -> Vector r e -> Vector r e +take k = fst . sliceAt k +{-# INLINE take #-} + +-- | Slice a manifest vector in such a way that it will contain all initial elements that +-- satisfy the supplied predicate. +-- +-- @since 0.5.5 +takeWhile :: Manifest r e => (e -> Bool) -> Vector r e -> Vector r e +takeWhile f v = take (go 0) v + where + !k = elemsCount v + go !i + | i < k && f (unsafeLinearIndex v i) = go (i + 1) + | otherwise = SafeSz i +{-# INLINE takeWhile #-} + +-- | /O(1)/ - Get the vector with the first @n@ elements. Throws an error size is less +-- than @n@. +-- +-- ==== __Examples__ +-- +-- >>> take' 0 (0 ..: 0) +-- Array D Seq (Sz1 0) +-- [ ] +-- >>> take' 5 (0 ..: 10) +-- Array D Seq (Sz1 5) +-- [ 0, 1, 2, 3, 4 ] +-- +-- @since 0.5.0 +take' :: forall r e. (HasCallStack, Source r e) => Sz1 -> Vector r e -> Vector r e +take' k = throwEither . takeM k +{-# INLINE take' #-} + +-- | /O(1)/ - Get the vector with the first @n@ elements. Throws an error size is less than @n@ +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> takeM 5 (0 ..: 10) +-- Array D Seq (Sz1 5) +-- [ 0, 1, 2, 3, 4 ] +-- >>> maybe 0 A.sum $ takeM 5 (0 ..: 10) +-- 10 +-- >>> maybe (-1) A.sum $ takeM 15 (0 ..: 10) +-- -1 +-- >>> takeM 15 (0 ..: 10) +-- *** Exception: SizeSubregionException: (Sz1 10) is to small for 0 (Sz1 15) +-- +-- @since 0.5.0 +takeM :: forall r e m. (Source r e, MonadThrow m) => Sz1 -> Vector r e -> m (Vector r e) +takeM k v = do + let sz = size v + when (k > sz) $ throwM $ SizeSubregionException sz 0 k + pure $ unsafeTake k v +{-# INLINE takeM #-} + +-- | /O(1)/ - Create a `Stream` vector with the first @n@ elements. Never fails +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +stake :: forall r e. Stream r Ix1 e => Sz1 -> Vector r e -> Vector DS e +stake n = fromSteps . S.take n . S.toStream +{-# INLINE stake #-} + +-- | /O(1)/ - Drop @n@ elements from a vector. This function never fails and has +-- similar semantics as the `Data.List.drop` for lists. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> v = makeVectorR D Seq 10 id +-- >>> v +-- Array D Seq (Sz1 10) +-- [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9 ] +-- >>> A.drop 5 v +-- Array D Seq (Sz1 5) +-- [ 5, 6, 7, 8, 9 ] +-- >>> A.drop 25 v +-- Array D Seq (Sz1 0) +-- [ ] +-- +-- @since 0.5.0 +drop :: forall r e. Source r e => Sz1 -> Vector r e -> Vector r e +drop k = snd . sliceAt k +{-# INLINE drop #-} + +-- | Slice a manifest vector in such a way that it will not contain all initial elements +-- that satisfy the supplied predicate. +-- +-- @since 0.5.5 +dropWhile :: forall r e. Manifest r e => (e -> Bool) -> Vector r e -> Vector r e +dropWhile f v = drop (go 0) v + where + !k = elemsCount v + go !i + | i < k && f (unsafeLinearIndex v i) = go (i + 1) + | otherwise = SafeSz i +{-# INLINE dropWhile #-} + +-- | Keep all but the first @n@ elements from the delayed stream vector. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sdrop :: forall r e. Stream r Ix1 e => Sz1 -> Vector r e -> Vector DS e +sdrop n = fromSteps . S.drop n . S.toStream +{-# INLINE sdrop #-} + +-- | /O(1)/ - Drop @n@ elements from a vector. Unlike `drop`, this function will +-- produce an error when supplied number of elements to drop is larger than size +-- of the supplied vector +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +drop' :: forall r e. (HasCallStack, Source r e) => Sz1 -> Vector r e -> Vector r e +drop' k = throwEither . dropM k +{-# INLINE drop' #-} + +-- | +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +dropM :: forall r e m. (Source r e, MonadThrow m) => Sz1 -> Vector r e -> m (Vector r e) +dropM k@(Sz d) v = do + let sz@(Sz n) = size v + when (k > sz) $ throwM $ SizeSubregionException sz d (SafeSz (n - d)) + pure $ unsafeLinearSlice d (SafeSz (n - d)) v +{-# INLINE dropM #-} + +-- | Same as 'sliceAt'', except it never fails. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sliceAt :: forall r e. Source r e => Sz1 -> Vector r e -> (Vector r e, Vector r e) +sliceAt (Sz k) v = (unsafeTake d v, unsafeDrop d v) + where + !n = coerce (size v) + !d = SafeSz (min k n) +{-# INLINE sliceAt #-} + +-- | Same as 'Data.Massiv.Array.splitAt'', except for a flat vector. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sliceAt' :: (HasCallStack, Source r e) => Sz1 -> Vector r e -> (Vector r e, Vector r e) +sliceAt' k = throwEither . sliceAtM k +{-# INLINE sliceAt' #-} + +-- | Same as `Data.Massiv.Array.splitAtM`, except for a flat vector. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sliceAtM + :: forall r e m. (Source r e, MonadThrow m) => Sz1 -> Vector r e -> m (Vector r e, Vector r e) +sliceAtM k v = do + l <- takeM k v + pure (l, unsafeDrop k v) +{-# INLINE sliceAtM #-} + +-- | Create an empty delayed stream vector +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sempty :: Vector DS e +sempty = DSArray S.empty +{-# INLINE sempty #-} + +-- | Create a delayed stream vector with a single element +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +ssingleton :: e -> Vector DS e +ssingleton = DSArray . S.singleton +{-# INLINE ssingleton #-} + +-- | /O(1)/ - Add an element to the vector from the left side +-- +-- @since 0.3.0 +cons :: forall r e. (Size r, Load r Ix1 e) => e -> Vector r e -> Vector DL e +cons e v = + let dv = toLoadArray v + load scheduler startAt uWrite uSet = + uWrite startAt e >> dlLoad dv scheduler (startAt + 1) uWrite uSet + {-# INLINE load #-} + in dv{dlSize = SafeSz (1 + unSz (dlSize dv)), dlLoad = load} +{-# INLINE cons #-} + +-- | /O(1)/ - Add an element to the vector from the right side +-- +-- @since 0.3.0 +snoc :: forall r e. (Size r, Load r Ix1 e) => Vector r e -> e -> Vector DL e +snoc v e = + let dv = toLoadArray v + !k = unSz (size dv) + load scheduler startAt uWrite uSet = + dlLoad dv scheduler startAt uWrite uSet >> uWrite (k + startAt) e + {-# INLINE load #-} + in dv{dlSize = SafeSz (1 + k), dlLoad = load} +{-# INLINE snoc #-} + +-- | Replicate the same element @n@ times +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sreplicate :: Sz1 -> e -> Vector DS e +sreplicate n = DSArray . S.replicate n +{-# INLINE sreplicate #-} + +-- | Create a delayed vector of length @n@ with a function that maps an index to an +-- element. Same as `makeLinearArray` +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +generate :: Comp -> Sz1 -> (Ix1 -> e) -> Vector D e +generate = makeArrayLinear +{-# INLINE generate #-} + +-- | Create a delayed stream vector of length @n@ with a function that maps an index to an +-- element. Same as `makeLinearArray` +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sgenerate :: Sz1 -> (Ix1 -> e) -> Vector DS e +sgenerate n = DSArray . S.generate n +{-# INLINE sgenerate #-} + +-- | Create a delayed stream vector of infinite length by repeatedly applying a function to the +-- initial value. +-- +-- ==== __Examples__ +-- +-- >>> stake 10 $ siterate succ 'a' +-- Array DS Seq (Sz1 10) +-- [ 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j' ] +-- +-- @since 0.5.2 +siterate :: (e -> e) -> e -> Vector DS e +siterate f = fromSteps . S.unfoldr (\a -> Just (a, f a)) +{-# INLINE siterate #-} + +-- | Create a delayed stream vector of length @n@ by repeatedly applying a function to the +-- initial value. +-- +-- ==== __Examples__ +-- +-- >>> siterateN 10 succ 'a' +-- Array DS Seq (Sz1 10) +-- [ 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j' ] +-- +-- @since 0.5.0 +siterateN :: Sz1 -> (e -> e) -> e -> Vector DS e +siterateN n f a = fromSteps $ S.iterateN n f a +{-# INLINE siterateN #-} + +-- | Create a vector by using the same monadic action @n@ times +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sreplicateM :: forall e m. Monad m => Sz1 -> m e -> m (Vector DS e) +sreplicateM n f = fromStepsM $ S.replicateM n f +{-# INLINE sreplicateM #-} + +-- | Create a delayed stream vector of length @n@ with a monadic action that from an index +-- generates an element. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sgenerateM :: forall e m. Monad m => Sz1 -> (Ix1 -> m e) -> m (Vector DS e) +sgenerateM n f = fromStepsM $ S.generateM n f +{-# INLINE sgenerateM #-} + +-- | Create a delayed stream vector of length @n@ by repeatedly apply a monadic action to +-- the initial value. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +siterateNM :: forall e m. Monad m => Sz1 -> (e -> m e) -> e -> m (Vector DS e) +siterateNM n f a = fromStepsM $ S.iterateNM n f a +{-# INLINE siterateNM #-} + +-- | Right unfolding function. Useful when it is unknown ahead of time how many +-- elements a vector will have. +-- +-- ====__Example__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> sunfoldr (\i -> if i < 9 then Just (i * i, i + 1) else Nothing) (0 :: Int) +-- Array DS Seq (Sz1 9) +-- [ 0, 1, 4, 9, 16, 25, 36, 49, 64 ] +-- +-- @since 0.5.0 +sunfoldr :: forall e s. (s -> Maybe (e, s)) -> s -> Vector DS e +sunfoldr f = DSArray . S.unfoldr f +{-# INLINE sunfoldr #-} + +-- | /O(n)/ - Right unfolding function with at most @n@ number of elements. +-- +-- ==== __Example__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> sunfoldrN 9 (\i -> Just (i*i, i + 1)) (0 :: Int) +-- Array DS Seq (Sz1 9) +-- [ 0, 1, 4, 9, 16, 25, 36, 49, 64 ] +-- +-- @since 0.5.0 +sunfoldrN + :: forall e s + . Sz1 + -- ^ @n@ - maximum number of elements that the vector will have + -> (s -> Maybe (e, s)) + -- ^ Unfolding function. Stops when `Nothing` is returned or maximum number of elements + -- is reached. + -> s + -- ^ Inititial element. + -> Vector DS e +sunfoldrN n f = DSArray . S.unfoldrN n f +{-# INLINE sunfoldrN #-} + +-- | /O(n)/ - Same as `sunfoldr`, but with monadic generating function. +-- +-- ==== __Examples__ +-- +-- >>> import Control.Monad (when, guard) +-- >>> sunfoldrM (\i -> when (i == 0) (Left "Zero denominator") >> Right (guard (i < 5) >> Just (100 `div` i, i + 1))) (-10 :: Int) +-- Left "Zero denominator" +-- >>> sunfoldrM (\i -> when (i == 0) (Left "Zero denominator") >> Right (guard (i < -5) >> Just (100 `div` i, i + 1))) (-10 :: Int) +-- Right (Array DS Seq (Sz1 5) +-- [ -10, -12, -13, -15, -17 ] +-- ) +-- +-- @since 0.5.0 +sunfoldrM :: forall e s m. Monad m => (s -> m (Maybe (e, s))) -> s -> m (Vector DS e) +sunfoldrM f = fromStepsM . S.unfoldrM f +{-# INLINE sunfoldrM #-} + +-- | /O(n)/ - Same as `sunfoldrN`, but with monadic generating function. +-- +-- ==== __Examples__ +-- +-- >>> import Control.Monad (guard) +-- >>> sunfoldrNM 6 (\i -> print i >> pure (guard (i < 5) >> Just (i * i, i + 1))) (10 :: Int) +-- 10 +-- Array DS Seq (Sz1 0) +-- [ ] +-- >>> sunfoldrNM 6 (\i -> print i >> pure (guard (i < 15) >> Just (i * i, i + 1))) (10 :: Int) +-- 10 +-- 11 +-- 12 +-- 13 +-- 14 +-- 15 +-- Array DS Seq (Sz1 5) +-- [ 100, 121, 144, 169, 196 ] +-- +-- +-- @since 0.5.0 +sunfoldrNM :: forall e s m. Monad m => Sz1 -> (s -> m (Maybe (e, s))) -> s -> m (Vector DS e) +sunfoldrNM (Sz n) f = fromStepsM . S.unfoldrNM n f +{-# INLINE sunfoldrNM #-} + +-- | /O(n)/ - Similar to `sunfoldrN`, except the length of the resulting vector will be exactly @n@ +-- +-- ==== __Examples__ +-- +-- >>> sunfoldrExactN 10 (\i -> (i * i, i + 1)) (10 :: Int) +-- Array DS Seq (Sz1 10) +-- [ 100, 121, 144, 169, 196, 225, 256, 289, 324, 361 ] +-- +-- @since 0.5.0 +sunfoldrExactN :: forall e s. Sz1 -> (s -> (e, s)) -> s -> Vector DS e +sunfoldrExactN n f = fromSteps . S.unfoldrExactN n f +{-# INLINE sunfoldrExactN #-} + +-- | /O(n)/ - Similar to `sunfoldrNM`, except the length of the resulting vector will be exactly @n@ +-- +-- ==== __Examples__ +-- +-- λ> sunfoldrExactNM 11 (\i -> pure (100 `div` i, i + 1)) (-10 :: Int) +-- Array DS *** Exception: divide by zero +-- λ> sunfoldrExactNM 11 (\i -> guard (i /= 0) >> Just (100 `div` i, i + 1)) (-10 :: Int) +-- Nothing +-- λ> sunfoldrExactNM 9 (\i -> guard (i /= 0) >> Just (100 `div` i, i + 1)) (-10 :: Int) +-- Just (Array DS Seq (Sz1 9) +-- [ -10, -12, -13, -15, -17, -20, -25, -34, -50 ] +-- ) +-- +-- @since 0.5.0 +sunfoldrExactNM :: forall e s m. Monad m => Sz1 -> (s -> m (e, s)) -> s -> m (Vector DS e) +sunfoldrExactNM n f = fromStepsM . S.unfoldrExactNM n f +{-# INLINE sunfoldrExactNM #-} + +-- | /O(n)/ - Enumerate from a starting number @x@ exactly @n@ times with a step @1@. +-- +-- /Related/: `senumFromStepN`, `enumFromN`, `enumFromStepN`, `rangeSize`, +-- `rangeStepSize`, `range`, 'rangeStep'' +-- +-- ==== __Examples__ +-- +-- >>> senumFromN (10 :: Int) 9 +-- Array DS Seq (Sz1 9) +-- [ 10, 11, 12, 13, 14, 15, 16, 17, 18 ] +-- +-- /__Similar__/: +-- +-- [@Prelude.`Prelude.enumFromTo`@] Very similar to @[x .. x + n - 1]@, except that +-- `senumFromN` is faster and it only works for `Num` and not for `Enum` elements +-- +-- [@Data.Vector.Generic.`Data.Vector.Generic.enumFromN`@] Uses exactly the same +-- implementation underneath. +-- +-- @since 0.5.0 +senumFromN + :: Num e + => e + -- ^ @x@ - starting number + -> Sz1 + -- ^ @n@ - length of resulting vector + -> Vector DS e +senumFromN x n = DSArray $ S.enumFromStepN x 1 n +{-# INLINE senumFromN #-} + +-- | /O(n)/ - Enumerate from a starting number @x@ exactly @n@ times with a custom step value @dx@ +-- +-- ==== __Examples__ +-- +-- >>> senumFromStepN (5 :: Int) 2 10 +-- Array DS Seq (Sz1 10) +-- [ 5, 7, 9, 11, 13, 15, 17, 19, 21, 23 ] +-- +-- __/Similar/__: +-- +-- [@Prelude.`Prelude.enumFrom`@] Just like @take n [x, x + dx ..]@, except that +-- `senumFromN` is faster and it only works for `Num` and not for `Enum` elements +-- +-- [@Data.Vector.Generic.`Data.Vector.Generic.enumFromStepN`@] Uses exactly the same +-- implementation underneath. +-- +-- @since 0.5.0 +senumFromStepN + :: Num e + => e + -- ^ @x@ - starting number + -> e + -- ^ @dx@ - Step + -> Sz1 + -- ^ @n@ - length of resulting vector + -> Vector DS e +senumFromStepN x step n = DSArray $ S.enumFromStepN x step n +{-# INLINE senumFromStepN #-} + +-- | Append two vectors together +-- +-- /Related/: `appendM`, `appendOuterM`, +-- +-- ==== __Examples__ +-- +-- λ> sappend (1 ..: 6) (senumFromStepN 6 (-1) 6) +-- Array DS Seq (Sz1 11) +-- [ 1, 2, 3, 4, 5, 6, 5, 4, 3, 2, 1 ] +-- +-- __/Similar/__: +-- +-- [@Data.Semigroup.`Data.Semigroup.<>`@] `DS` and `DL` arrays have instances for +-- `Semigroup`, so they will work in a similar fashion. `sappend` differs in that it accepts +-- `Stream` arrays with possibly different representations. +-- +-- [@Data.List.`Data.List.++`@] Same operation, but for lists. +-- +-- [@Data.Vector.Generic.`Data.Vector.Generic.++`@] Uses exactly the same implementation +-- underneath as `sappend`, except that it cannot append two vectors with different +-- memory representations. +-- +-- @since 0.5.0 +sappend + :: forall r1 r2 e + . (Stream r1 Ix1 e, Stream r2 Ix1 e) + => Vector r1 e + -> Vector r2 e + -> Vector DS e +sappend a1 a2 = fromSteps (toStream a1 `S.append` toStream a2) +{-# INLINE sappend #-} + +-- | Concat vectors together +-- +-- /Related/: `concatM`, `concatOuterM`, +-- +-- ==== __Examples__ +-- +-- >>> sconcat [2 ... 6, empty, singleton 1, generate Seq 5 id] +-- Array DS Seq (Sz1 11) +-- [ 2, 3, 4, 5, 6, 1, 0, 1, 2, 3, 4 ] +-- >>> sconcat [senumFromN 2 5, sempty, ssingleton 1, sgenerate 5 id] +-- Array DS Seq (Sz1 11) +-- [ 2, 3, 4, 5, 6, 1, 0, 1, 2, 3, 4 ] +-- +-- __/Similar/__: +-- +-- [@Data.Monoid.`Data.Monoid.mconcat`@] `DS` and `DL` arrays have instances for `Monoid`, so +-- they will work in a similar fashion. `sconcat` differs in that it accepts `Stream` +-- arrays of other representations. +-- +-- [@Data.List.`Data.List.concat`@] Same operation, but for lists. +-- +-- [@Data.Vector.Generic.`Data.Vector.Generic.concat`@] Uses exactly the same +-- implementation underneath as `sconcat`. +-- +-- @since 0.5.0 +sconcat :: forall r e. Stream r Ix1 e => [Vector r e] -> Vector DS e +sconcat = DSArray . foldMap toStream +{-# INLINE sconcat #-} + +-- | Convert a list to a delayed stream vector +-- +-- /Related/: `fromList`, `fromListN`, `sfromListN` +-- +-- ==== __Examples__ +-- +-- >>> sfromList ([] :: [Int]) +-- Array DS Seq (Sz1 0) +-- [ ] +-- >>> sfromList ([1,2,3] :: [Int]) +-- Array DS Seq (Sz1 3) +-- [ 1, 2, 3 ] +-- +-- @since 0.5.0 +sfromList :: [e] -> Vector DS e +sfromList = fromSteps . S.fromList +{-# INLINE sfromList #-} + +-- | Convert a list to a delayed stream vector. Length of the resulting vector will be at +-- most @n@. This version isn't really more efficient then `sfromList`, but there is +-- `Data.Massiv.Array.Unsafe.unsafeFromListN` +-- +-- /Related/: `fromList`, `fromListN`, `sfromList` +-- +-- ==== __Examples__ +-- +-- >>> sfromListN 10 [1 :: Int ..] +-- Array DS Seq (Sz1 10) +-- [ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 ] +-- >>> sfromListN 10 [1 :: Int .. 5] +-- Array DS Seq (Sz1 5) +-- [ 1, 2, 3, 4, 5 ] +-- +-- @since 0.5.1 +sfromListN :: Sz1 -> [e] -> Vector DS e +sfromListN (Sz n) = fromSteps . S.fromListN n +{-# INLINE sfromListN #-} + +-- | Convert an array to a list by the means of a delayed stream vector. +-- +-- /Related/: `toList` +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +stoList :: forall r ix e. Stream r ix e => Array r ix e -> [e] +stoList = S.toList . toStream +{-# INLINE stoList #-} + +-- | Sequentially filter out elements from the array according to the supplied predicate. +-- +-- ==== __Example__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> arr = makeArrayR D Seq (Sz2 3 4) fromIx2 +-- >>> arr +-- Array D Seq (Sz (3 :. 4)) +-- [ [ (0,0), (0,1), (0,2), (0,3) ] +-- , [ (1,0), (1,1), (1,2), (1,3) ] +-- , [ (2,0), (2,1), (2,2), (2,3) ] +-- ] +-- >>> sfilter (even . fst) arr +-- Array DS Seq (Sz1 8) +-- [ (0,0), (0,1), (0,2), (0,3), (2,0), (2,1), (2,2), (2,3) ] +-- +-- @since 0.5.0 +sfilter :: forall r ix e. S.Stream r ix e => (e -> Bool) -> Array r ix e -> Vector DS e +sfilter f = DSArray . S.filter f . S.toStream +{-# INLINE sfilter #-} + +-- | Similar to `sfilter`, but filter with an index aware function. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sifilter :: forall r ix e. Stream r ix e => (ix -> e -> Bool) -> Array r ix e -> Vector DS e +sifilter f = + simapMaybe $ \ix e -> + if f ix e + then Just e + else Nothing +{-# INLINE sifilter #-} + +-- | Sequentially filter out elements from the array according to the supplied applicative predicate. +-- +-- ==== __Example__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> arr = makeArrayR D Seq (Sz2 3 4) fromIx2 +-- >>> arr +-- Array D Seq (Sz (3 :. 4)) +-- [ [ (0,0), (0,1), (0,2), (0,3) ] +-- , [ (1,0), (1,1), (1,2), (1,3) ] +-- , [ (2,0), (2,1), (2,2), (2,3) ] +-- ] +-- >>> sfilterM (Just . odd . fst) arr +-- Just (Array DS Seq (Sz1 4) +-- [ (1,0), (1,1), (1,2), (1,3) ] +-- ) +-- >>> sfilterM (\ix@(_, j) -> print ix >> return (even j)) arr +-- (0,0) +-- (0,1) +-- (0,2) +-- (0,3) +-- (1,0) +-- (1,1) +-- (1,2) +-- (1,3) +-- (2,0) +-- (2,1) +-- (2,2) +-- (2,3) +-- Array DS Seq (Sz1 6) +-- [ (0,0), (0,2), (1,0), (1,2), (2,0), (2,2) ] +-- +-- @since 0.5.0 +sfilterM + :: forall r ix e f + . (S.Stream r ix e, Applicative f) + => (e -> f Bool) + -> Array r ix e + -> f (Vector DS e) +sfilterM f arr = DSArray <$> S.filterA f (S.toStream arr) +{-# INLINE sfilterM #-} + +-- | Similar to `filterM`, but filter with an index aware function. +-- +-- Corresponds to: @`filterM` (uncurry f) . `simap` (,)@ +-- +-- @since 0.5.0 +sifilterM + :: forall r ix e f + . (Stream r ix e, Applicative f) + => (ix -> e -> f Bool) + -> Array r ix e + -> f (Vector DS e) +sifilterM f = + simapMaybeM $ \ix e -> + ( \p -> + if p + then Just e + else Nothing + ) + <$> f ix e +{-# INLINE sifilterM #-} + +-- | Apply a function to each element of the array, while discarding `Nothing` and +-- keeping the `Maybe` result. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +smapMaybe :: forall r ix a b. S.Stream r ix a => (a -> Maybe b) -> Array r ix a -> Vector DS b +smapMaybe f = DSArray . S.mapMaybe f . S.toStream +{-# INLINE smapMaybe #-} + +-- | Similar to `smapMaybe`, but map with an index aware function. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +simapMaybe + :: forall r ix a b + . Stream r ix a + => (ix -> a -> Maybe b) + -> Array r ix a + -> Vector DS b +simapMaybe f = DSArray . S.mapMaybe (uncurry f) . toStreamIx +{-# INLINE simapMaybe #-} + +-- | Similar to `smapMaybeM`, but map with an index aware function. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +simapMaybeM + :: forall r ix a b f + . (Stream r ix a, Applicative f) + => (ix -> a -> f (Maybe b)) + -> Array r ix a + -> f (Vector DS b) +simapMaybeM f = fmap DSArray . S.mapMaybeA (uncurry f) . toStreamIx +{-# INLINE simapMaybeM #-} + +-- | Keep all `Maybe`s and discard the `Nothing`s. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +scatMaybes :: forall r ix a. S.Stream r ix (Maybe a) => Array r ix (Maybe a) -> Vector DS a +scatMaybes = smapMaybe id +{-# INLINE scatMaybes #-} + +-- | Similar to `smapMaybe`, but with the `Applicative` function. +-- +-- Similar to @mapMaybe id <$> mapM f arr@ +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +smapMaybeM + :: forall r ix a b f + . (S.Stream r ix a, Applicative f) + => (a -> f (Maybe b)) + -> Array r ix a + -> f (Vector DS b) +smapMaybeM f = fmap DSArray . S.mapMaybeA f . S.toStream +{-# INLINE smapMaybeM #-} + +-- | Map a function over a stream vector +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +smap + :: forall r ix a b + . S.Stream r ix a + => (a -> b) + -> Array r ix a + -> Vector DS b +smap f = fromSteps . S.map f . S.toStream +{-# INLINE smap #-} + +-- | Map an index aware function over a stream vector +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +simap + :: forall r ix a b + . S.Stream r ix a + => (ix -> a -> b) + -> Array r ix a + -> Vector DS b +simap f = fromSteps . S.map (uncurry f) . S.toStreamIx +{-# INLINE simap #-} + +-- | Traverse a stream vector with an applicative function. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +straverse + :: forall r ix a b f + . (S.Stream r ix a, Applicative f) + => (a -> f b) + -> Array r ix a + -> f (Vector DS b) +straverse f = fmap fromSteps . S.traverse f . S.toStream +{-# INLINE straverse #-} + +-- | Traverse a stream vector with an index aware applicative function. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sitraverse + :: forall r ix a b f + . (S.Stream r ix a, Applicative f) + => (ix -> a -> f b) + -> Array r ix a + -> f (Vector DS b) +sitraverse f = fmap fromSteps . S.traverse (uncurry f) . S.toStreamIx +{-# INLINE sitraverse #-} + +-- | Traverse a stream vector with a monadic function. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +smapM + :: forall r ix a b m + . (S.Stream r ix a, Monad m) + => (a -> m b) + -> Array r ix a + -> m (Vector DS b) +smapM f = fromStepsM . S.mapM f . S.transStepsId . S.toStream +{-# INLINE smapM #-} + +-- | Traverse a stream vector with a monadic index aware function. +-- +-- Corresponds to: @mapM (uncurry f) . imap (,) v@ +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +simapM + :: forall r ix a b m + . (S.Stream r ix a, Monad m) + => (ix -> a -> m b) + -> Array r ix a + -> m (Vector DS b) +simapM f = fromStepsM . S.mapM (uncurry f) . S.transStepsId . S.toStreamIx +{-# INLINE simapM #-} + +-- | Traverse a stream vector with a monadic function, while discarding the result +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +smapM_ + :: forall r ix a b m + . (S.Stream r ix a, Monad m) + => (a -> m b) + -> Array r ix a + -> m () +smapM_ f = S.mapM_ f . S.transStepsId . S.toStream +{-# INLINE smapM_ #-} + +-- | Traverse a stream vector with a monadic index aware function, while discarding the result +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +simapM_ + :: forall r ix a b m + . (S.Stream r ix a, Monad m) + => (ix -> a -> m b) + -> Array r ix a + -> m () +simapM_ f = S.mapM_ (uncurry f) . S.transStepsId . S.toStreamIx +{-# INLINE simapM_ #-} + +-- | Same as `smapM`, but with arguments flipped. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sforM + :: forall r ix a b m + . (S.Stream r ix a, Monad m) + => Array r ix a + -> (a -> m b) + -> m (Vector DS b) +sforM = flip smapM +{-# INLINE sforM #-} + +-- | Same as `simapM`, but with arguments flipped. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +siforM + :: forall r ix a b m + . (S.Stream r ix a, Monad m) + => Array r ix a + -> (ix -> a -> m b) + -> m (Vector DS b) +siforM = flip simapM +{-# INLINE siforM #-} + +-- | Same as `smapM_`, but with arguments flipped. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sforM_ :: (S.Stream r ix a, Monad m) => Array r ix a -> (a -> m b) -> m () +sforM_ = flip smapM_ +{-# INLINE sforM_ #-} + +-- | Same as `simapM_`, but with arguments flipped. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +siforM_ + :: forall r ix a b m + . (S.Stream r ix a, Monad m) + => Array r ix a + -> (ix -> a -> m b) + -> m () +siforM_ = flip simapM_ +{-# INLINE siforM_ #-} + +-- | Zip two vectors together into a vector. The length of a resulting vector will +-- be the smallest length of the supplied vectors. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +szip + :: forall ra rb a b + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b) + => Vector ra a + -> Vector rb b + -> Vector DS (a, b) +szip = szipWith (,) +{-# INLINE szip #-} + +-- | Zip three vectors together into a vector. The length of a resulting vector will +-- be the smallest length of the supplied vectors. +-- +-- @since 0.5.0 +szip3 + :: forall ra rb rc a b c + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c) + => Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector DS (a, b, c) +szip3 = szipWith3 (,,) +{-# INLINE szip3 #-} + +-- | Zip four vectors together into a vector. The length of a resulting vector will +-- be the smallest length of the supplied vectors. +-- +-- @since 0.5.0 +szip4 + :: forall ra rb rc rd a b c d + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, S.Stream rd Ix1 d) + => Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector DS (a, b, c, d) +szip4 = szipWith4 (,,,) +{-# INLINE szip4 #-} + +-- | Zip five vectors together into a vector. The length of a resulting vector will +-- be the smallest length of the supplied vectors. +-- +-- @since 0.5.0 +szip5 + :: forall ra rb rc rd re a b c d e + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, S.Stream rd Ix1 d, S.Stream re Ix1 e) + => Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector re e + -> Vector DS (a, b, c, d, e) +szip5 = szipWith5 (,,,,) +{-# INLINE szip5 #-} + +-- | Zip six vectors together into a vector. The length of a resulting vector will +-- be the smallest length of the supplied vectors. +-- +-- @since 0.5.0 +szip6 + :: forall ra rb rc rd re rf a b c d e f + . ( S.Stream ra Ix1 a + , S.Stream rb Ix1 b + , S.Stream rc Ix1 c + , S.Stream rd Ix1 d + , S.Stream re Ix1 e + , S.Stream rf Ix1 f + ) + => Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector re e + -> Vector rf f + -> Vector DS (a, b, c, d, e, f) +szip6 = szipWith6 (,,,,,) +{-# INLINE szip6 #-} + +-- | Zip two vectors together with a binary function into a vector. The length +-- of a resulting vector will be the smallest length of the supplied vectors. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +szipWith + :: forall ra rb a b c + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b) + => (a -> b -> c) + -> Vector ra a + -> Vector rb b + -> Vector DS c +szipWith f v1 v2 = fromSteps $ S.zipWith f (S.toStream v1) (S.toStream v2) +{-# INLINE szipWith #-} + +-- | Zip three vectors together with a ternary function into a vector. The length +-- of a resulting vector will be the smallest length of the supplied vectors. +-- +-- @since 0.5.0 +szipWith3 + :: forall ra rb rc a b c d + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c) + => (a -> b -> c -> d) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector DS d +szipWith3 f v1 v2 v3 = fromSteps $ S.zipWith3 f (S.toStream v1) (S.toStream v2) (S.toStream v3) +{-# INLINE szipWith3 #-} + +-- | Zip four vectors together with a quaternary function into a vector. The length +-- of a resulting vector will be the smallest length of the supplied vectors. +-- +-- @since 0.5.0 +szipWith4 + :: forall ra rb rc rd a b c d e + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, S.Stream rd Ix1 d) + => (a -> b -> c -> d -> e) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector DS e +szipWith4 f v1 v2 v3 v4 = + fromSteps $ S.zipWith4 f (S.toStream v1) (S.toStream v2) (S.toStream v3) (S.toStream v4) +{-# INLINE szipWith4 #-} + +-- | Zip five vectors together with a quinary function into a vector. The length +-- of a resulting vector will be the smallest length of the supplied vectors. +-- +-- @since 0.5.0 +szipWith5 + :: forall ra rb rc rd re a b c d e f + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, S.Stream rd Ix1 d, S.Stream re Ix1 e) + => (a -> b -> c -> d -> e -> f) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector re e + -> Vector DS f +szipWith5 f v1 v2 v3 v4 v5 = + fromSteps $ + S.zipWith5 f (S.toStream v1) (S.toStream v2) (S.toStream v3) (S.toStream v4) (S.toStream v5) +{-# INLINE szipWith5 #-} + +-- | Zip six vectors together with a senary function into a vector. The length +-- of a resulting vector will be the smallest length of the supplied vectors. +-- +-- @since 0.5.0 +szipWith6 + :: forall ra rb rc rd re rf a b c d e f g + . ( S.Stream ra Ix1 a + , S.Stream rb Ix1 b + , S.Stream rc Ix1 c + , S.Stream rd Ix1 d + , S.Stream re Ix1 e + , S.Stream rf Ix1 f + ) + => (a -> b -> c -> d -> e -> f -> g) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector re e + -> Vector rf f + -> Vector DS g +szipWith6 f v1 v2 v3 v4 v5 v6 = + fromSteps $ + S.zipWith6 + f + (S.toStream v1) + (S.toStream v2) + (S.toStream v3) + (S.toStream v4) + (S.toStream v5) + (S.toStream v6) +{-# INLINE szipWith6 #-} + +-- | Just like `szipWith`, zip two vectors together, but with an index aware +-- function. The length of a resulting vector will be the smallest length of the +-- supplied vectors. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sizipWith + :: forall ra rb a b c + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b) + => (Ix1 -> a -> b -> c) + -> Vector ra a + -> Vector rb b + -> Vector DS c +sizipWith f v1 v2 = fromSteps $ S.zipWith (uncurry f) (S.toStreamIx v1) (S.toStream v2) +{-# INLINE sizipWith #-} + +-- | Just like `szipWith3`, zip three vectors together, but with an index aware +-- function. The length of a resulting vector will be the smallest length of the +-- supplied vectors. +-- +-- @since 0.5.0 +sizipWith3 + :: forall ra rb rc a b c d + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c) + => (Ix1 -> a -> b -> c -> d) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector DS d +sizipWith3 f v1 v2 v3 = + fromSteps $ S.zipWith3 (uncurry f) (S.toStreamIx v1) (S.toStream v2) (S.toStream v3) +{-# INLINE sizipWith3 #-} + +-- | Just like `szipWith4`, zip four vectors together, but with an index aware +-- function. The length of a resulting vector will be the smallest +-- length of the supplied vectors. +-- +-- @since 0.5.0 +sizipWith4 + :: forall ra rb rc rd a b c d e + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, S.Stream rd Ix1 d) + => (Ix1 -> a -> b -> c -> d -> e) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector DS e +sizipWith4 f v1 v2 v3 v4 = + fromSteps $ + S.zipWith4 (uncurry f) (S.toStreamIx v1) (S.toStream v2) (S.toStream v3) (S.toStream v4) +{-# INLINE sizipWith4 #-} + +-- | Just like `szipWith5`, zip five vectors together, but with an index aware +-- function. The length of a resulting vector will be the smallest length of the +-- supplied vectors. +-- +-- @since 0.5.0 +sizipWith5 + :: forall ra rb rc rd re a b c d e f + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, S.Stream rd Ix1 d, S.Stream re Ix1 e) + => (Ix1 -> a -> b -> c -> d -> e -> f) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector re e + -> Vector DS f +sizipWith5 f v1 v2 v3 v4 v5 = + fromSteps $ + S.zipWith5 + (uncurry f) + (S.toStreamIx v1) + (S.toStream v2) + (S.toStream v3) + (S.toStream v4) + (S.toStream v5) +{-# INLINE sizipWith5 #-} + +-- | Just like `szipWith6`, zip six vectors together, but with an index aware +-- function. The length of a resulting vector will be the smallest length of the +-- supplied vectors. +-- +-- @since 0.5.0 +sizipWith6 + :: forall ra rb rc rd re rf a b c d e f g + . ( S.Stream ra Ix1 a + , S.Stream rb Ix1 b + , S.Stream rc Ix1 c + , S.Stream rd Ix1 d + , S.Stream re Ix1 e + , S.Stream rf Ix1 f + ) + => (Ix1 -> a -> b -> c -> d -> e -> f -> g) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector re e + -> Vector rf f + -> Vector DS g +sizipWith6 f v1 v2 v3 v4 v5 v6 = + fromSteps $ + S.zipWith6 + (uncurry f) + (S.toStreamIx v1) + (S.toStream v2) + (S.toStream v3) + (S.toStream v4) + (S.toStream v5) + (S.toStream v6) +{-# INLINE sizipWith6 #-} + +-- | Zip two vectors together with a binary monadic action into a vector. The +-- length of a resulting vector will be the smallest length of the supplied +-- vectors. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +szipWithM + :: forall ra rb a b c m + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, Monad m) + => (a -> b -> m c) + -> Vector ra a + -> Vector rb b + -> m (Vector DS c) +szipWithM f v1 v2 = fromStepsM $ S.zipWithM f (toStreamM v1) (toStreamM v2) +{-# INLINE szipWithM #-} + +-- | Zip three vectors together with a ternary monadic action into a vector. The +-- length of a resulting vector will be the smallest length of the supplied +-- vectors. +-- +-- @since 0.5.0 +szipWith3M + :: forall ra rb rc a b c d m + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, Monad m) + => (a -> b -> c -> m d) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> m (Vector DS d) +szipWith3M f v1 v2 v3 = fromStepsM $ S.zipWith3M f (toStreamM v1) (toStreamM v2) (toStreamM v3) +{-# INLINE szipWith3M #-} + +-- | Zip four vectors together with a quaternary monadic action into a vector. The +-- length of a resulting vector will be the smallest length of the supplied +-- vectors. +-- +-- @since 0.5.0 +szipWith4M + :: forall ra rb rc rd a b c d e m + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, S.Stream rd Ix1 d, Monad m) + => (a -> b -> c -> d -> m e) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> m (Vector DS e) +szipWith4M f v1 v2 v3 v4 = + fromStepsM $ S.zipWith4M f (toStreamM v1) (toStreamM v2) (toStreamM v3) (toStreamM v4) +{-# INLINE szipWith4M #-} + +-- | Zip five vectors together with a quinary monadic action into a vector. The +-- length of a resulting vector will be the smallest length of the supplied +-- vectors. +-- +-- @since 0.5.0 +szipWith5M + :: forall ra rb rc rd re a b c d e f m + . ( S.Stream ra Ix1 a + , S.Stream rb Ix1 b + , S.Stream rc Ix1 c + , S.Stream rd Ix1 d + , S.Stream re Ix1 e + , Monad m + ) + => (a -> b -> c -> d -> e -> m f) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector re e + -> m (Vector DS f) +szipWith5M f v1 v2 v3 v4 v5 = + fromStepsM $ + S.zipWith5M f (toStreamM v1) (toStreamM v2) (toStreamM v3) (toStreamM v4) (toStreamM v5) +{-# INLINE szipWith5M #-} + +-- | Zip six vectors together with a senary monadic action into a vector. The +-- length of a resulting vector will be the smallest length of the supplied +-- vectors. +-- +-- @since 0.5.0 +szipWith6M + :: forall ra rb rc rd re rf a b c d e f g m + . ( S.Stream ra Ix1 a + , S.Stream rb Ix1 b + , S.Stream rc Ix1 c + , S.Stream rd Ix1 d + , S.Stream re Ix1 e + , S.Stream rf Ix1 f + , Monad m + ) + => (a -> b -> c -> d -> e -> f -> m g) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector re e + -> Vector rf f + -> m (Vector DS g) +szipWith6M f v1 v2 v3 v4 v5 v6 = + fromStepsM $ + S.zipWith6M + f + (toStreamM v1) + (toStreamM v2) + (toStreamM v3) + (toStreamM v4) + (toStreamM v5) + (toStreamM v6) +{-# INLINE szipWith6M #-} + +-- | Just like `szipWithM`, zip two vectors together, but with an index aware +-- monadic action. The length of a resulting vector will be the smallest length of the +-- supplied vectors. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sizipWithM + :: forall ra rb a b c m + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, Monad m) + => (Ix1 -> a -> b -> m c) + -> Vector ra a + -> Vector rb b + -> m (Vector DS c) +sizipWithM f v1 v2 = fromStepsM $ S.zipWithM (uncurry f) (toStreamIxM v1) (toStreamM v2) +{-# INLINE sizipWithM #-} + +-- | Just like `szipWith3M`, zip three vectors together, but with an index aware +-- monadic action. The length of a resulting vector will be the smallest length of the +-- supplied vectors. +-- +-- @since 0.5.0 +sizipWith3M + :: forall ra rb rc a b c d m + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, Monad m) + => (Ix1 -> a -> b -> c -> m d) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> m (Vector DS d) +sizipWith3M f v1 v2 v3 = + fromStepsM $ S.zipWith3M (uncurry f) (toStreamIxM v1) (toStreamM v2) (toStreamM v3) +{-# INLINE sizipWith3M #-} + +-- | Just like `szipWith4M`, zip four vectors together, but with an index aware +-- monadic action. The length of a resulting vector will be the smallest length of the +-- supplied vectors. +-- +-- @since 0.5.0 +sizipWith4M + :: forall ra rb rc rd a b c d e m + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, S.Stream rd Ix1 d, Monad m) + => (Ix1 -> a -> b -> c -> d -> m e) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> m (Vector DS e) +sizipWith4M f v1 v2 v3 v4 = + fromStepsM $ + S.zipWith4M (uncurry f) (toStreamIxM v1) (toStreamM v2) (toStreamM v3) (toStreamM v4) +{-# INLINE sizipWith4M #-} + +-- | Just like `szipWith6M`, zip five vectors together, but with an index aware +-- monadic action. The length of a resulting vector will be the smallest length of the +-- supplied vectors. +-- +-- @since 0.5.0 +sizipWith5M + :: forall ra rb rc rd re a b c d e f m + . ( S.Stream ra Ix1 a + , S.Stream rb Ix1 b + , S.Stream rc Ix1 c + , S.Stream rd Ix1 d + , S.Stream re Ix1 e + , Monad m + ) + => (Ix1 -> a -> b -> c -> d -> e -> m f) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector re e + -> m (Vector DS f) +sizipWith5M f v1 v2 v3 v4 v5 = + fromStepsM $ + S.zipWith5M + (uncurry f) + (toStreamIxM v1) + (toStreamM v2) + (toStreamM v3) + (toStreamM v4) + (toStreamM v5) +{-# INLINE sizipWith5M #-} + +-- | Just like `szipWith6M`, zip six vectors together, but with an index aware +-- monadic action. The length of a resulting vector will be the smallest length of the +-- supplied vectors. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sizipWith6M + :: forall ra rb rc rd re rf a b c d e f g m + . ( S.Stream ra Ix1 a + , S.Stream rb Ix1 b + , S.Stream rc Ix1 c + , S.Stream rd Ix1 d + , S.Stream re Ix1 e + , S.Stream rf Ix1 f + , Monad m + ) + => (Ix1 -> a -> b -> c -> d -> e -> f -> m g) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector re e + -> Vector rf f + -> m (Vector DS g) +sizipWith6M f v1 v2 v3 v4 v5 v6 = + fromStepsM $ + S.zipWith6M + (uncurry f) + (toStreamIxM v1) + (toStreamM v2) + (toStreamM v3) + (toStreamM v4) + (toStreamM v5) + (toStreamM v6) +{-# INLINE sizipWith6M #-} + +-- | Similar to `szipWithM`, zip two vectors together with a binary monadic +-- action, while discarding its result. The action will be invoked as many times as +-- the length of the smallest vector. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +szipWithM_ + :: forall ra rb a b c m + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, Monad m) + => (a -> b -> m c) + -> Vector ra a + -> Vector rb b + -> m () +szipWithM_ f v1 v2 = S.zipWithM_ f (toStreamM v1) (toStreamM v2) +{-# INLINE szipWithM_ #-} + +-- | Similar to `szipWith3M`, zip three vectors together with a ternary monadic +-- action, while discarding its result. The action will be invoked as many times as +-- the length of the smallest vector. +-- +-- @since 0.5.0 +szipWith3M_ + :: forall ra rb rc a b c d m + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, Monad m) + => (a -> b -> c -> m d) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> m () +szipWith3M_ f v1 v2 v3 = S.zipWith3M_ f (toStreamM v1) (toStreamM v2) (toStreamM v3) +{-# INLINE szipWith3M_ #-} + +-- | Similar to `szipWith4M`, zip four vectors together with a quaternary monadic +-- action, while discarding its result. The action will be invoked as many times as +-- the length of the smallest vector. +-- +-- @since 0.5.0 +szipWith4M_ + :: forall ra rb rc rd a b c d e m + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, S.Stream rd Ix1 d, Monad m) + => (a -> b -> c -> d -> m e) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> m () +szipWith4M_ f v1 v2 v3 v4 = + S.zipWith4M_ f (toStreamM v1) (toStreamM v2) (toStreamM v3) (toStreamM v4) +{-# INLINE szipWith4M_ #-} + +-- | Similar to `szipWith5M`, zip five vectors together with a quinary monadic +-- action, while discarding its result. The action will be invoked as many times as +-- the length of the smallest vector. +-- +-- @since 0.5.0 +szipWith5M_ + :: forall ra rb rc rd re a b c d e f m + . ( S.Stream ra Ix1 a + , S.Stream rb Ix1 b + , S.Stream rc Ix1 c + , S.Stream rd Ix1 d + , S.Stream re Ix1 e + , Monad m + ) + => (a -> b -> c -> d -> e -> m f) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector re e + -> m () +szipWith5M_ f v1 v2 v3 v4 v5 = + S.zipWith5M_ f (toStreamM v1) (toStreamM v2) (toStreamM v3) (toStreamM v4) (toStreamM v5) +{-# INLINE szipWith5M_ #-} + +-- | Similar to `szipWith6M`, zip six vectors together with a senary monadic +-- action, while discarding its result. The action will be invoked as many times as +-- the length of the smallest vector. +-- +-- @since 0.5.0 +szipWith6M_ + :: forall ra rb rc rd re rf a b c d e f g m + . ( S.Stream ra Ix1 a + , S.Stream rb Ix1 b + , S.Stream rc Ix1 c + , S.Stream rd Ix1 d + , S.Stream re Ix1 e + , S.Stream rf Ix1 f + , Monad m + ) + => (a -> b -> c -> d -> e -> f -> m g) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector re e + -> Vector rf f + -> m () +szipWith6M_ f v1 v2 v3 v4 v5 v6 = + S.zipWith6M_ + f + (toStreamM v1) + (toStreamM v2) + (toStreamM v3) + (toStreamM v4) + (toStreamM v5) + (toStreamM v6) +{-# INLINE szipWith6M_ #-} + +-- | Same as `szipWithM_`, zip two vectors together, but with an index aware +-- monadic action. The action will be invoked as many times as the length of the +-- smallest vector. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sizipWithM_ + :: forall ra rb a b c m + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, Monad m) + => (Ix1 -> a -> b -> m c) + -> Vector ra a + -> Vector rb b + -> m () +sizipWithM_ f v1 v2 = S.zipWithM_ (uncurry f) (toStreamIxM v1) (toStreamM v2) +{-# INLINE sizipWithM_ #-} + +-- | Same as `szipWith3M_`, zip three vectors together, but with an index aware +-- monadic action. The action will be invoked as many times as the length of the +-- smallest vector. +-- +-- @since 0.5.0 +sizipWith3M_ + :: forall ra rb rc a b c d m + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, Monad m) + => (Ix1 -> a -> b -> c -> m d) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> m () +sizipWith3M_ f v1 v2 v3 = S.zipWith3M_ (uncurry f) (toStreamIxM v1) (toStreamM v2) (toStreamM v3) +{-# INLINE sizipWith3M_ #-} + +-- | Same as `szipWith4M_`, zip four vectors together, but with an index aware +-- monadic action. The action will be invoked as many times as the length of the +-- smallest vector. +-- +-- @since 0.5.0 +sizipWith4M_ + :: forall ra rb rc rd a b c d e m + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, S.Stream rd Ix1 d, Monad m) + => (Ix1 -> a -> b -> c -> d -> m e) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> m () +sizipWith4M_ f v1 v2 v3 v4 = + S.zipWith4M_ (uncurry f) (toStreamIxM v1) (toStreamM v2) (toStreamM v3) (toStreamM v4) +{-# INLINE sizipWith4M_ #-} + +-- | Same as `szipWith5M_`, zip five vectors together, but with an index aware +-- monadic action. The action will be invoked as many times as the length of the +-- smallest vector. +-- +-- @since 0.5.0 +sizipWith5M_ + :: forall ra rb rc rd re a b c d e f m + . ( S.Stream ra Ix1 a + , S.Stream rb Ix1 b + , S.Stream rc Ix1 c + , S.Stream rd Ix1 d + , S.Stream re Ix1 e + , Monad m + ) + => (Ix1 -> a -> b -> c -> d -> e -> m f) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector re e + -> m () +sizipWith5M_ f v1 v2 v3 v4 v5 = + S.zipWith5M_ + (uncurry f) + (toStreamIxM v1) + (toStreamM v2) + (toStreamM v3) + (toStreamM v4) + (toStreamM v5) +{-# INLINE sizipWith5M_ #-} + +-- | Same as `szipWith6M_`, zip six vectors together, but with an index aware +-- monadic action. The action will be invoked as many times as the length of the +-- smallest vector. +-- +-- @since 0.5.0 +sizipWith6M_ + :: forall ra rb rc rd re rf a b c d e f g m + . ( S.Stream ra Ix1 a + , S.Stream rb Ix1 b + , S.Stream rc Ix1 c + , S.Stream rd Ix1 d + , S.Stream re Ix1 e + , S.Stream rf Ix1 f + , Monad m + ) + => (Ix1 -> a -> b -> c -> d -> e -> f -> m g) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector re e + -> Vector rf f + -> m () +sizipWith6M_ f v1 v2 v3 v4 v5 v6 = + S.zipWith6M_ + (uncurry f) + (toStreamIxM v1) + (toStreamM v2) + (toStreamM v3) + (toStreamM v4) + (toStreamM v5) + (toStreamM v6) +{-# INLINE sizipWith6M_ #-} + +-- | Streaming fold over an array in a row-major fashion with a left biased +-- function and a strict accumulator. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sfoldl + :: forall r ix e a + . Stream r ix e + => (a -> e -> a) + -> a + -> Array r ix e + -> a +sfoldl f acc = S.unId . S.foldl f acc . toStream +{-# INLINE sfoldl #-} + +-- | +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sfoldlM + :: forall r ix e a m + . (Stream r ix e, Monad m) + => (a -> e -> m a) + -> a + -> Array r ix e + -> m a +sfoldlM f acc = S.foldlM f acc . S.transStepsId . toStream +{-# INLINE sfoldlM #-} + +-- | +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sfoldlM_ + :: forall r ix e a m + . (Stream r ix e, Monad m) + => (a -> e -> m a) + -> a + -> Array r ix e + -> m () +sfoldlM_ f acc = void . sfoldlM f acc +{-# INLINE sfoldlM_ #-} + +-- | +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sfoldl1' + :: forall r ix e + . (HasCallStack, Stream r ix e) + => (e -> e -> e) + -> Array r ix e + -> e +sfoldl1' f = throwEither . sfoldl1M (\e -> pure . f e) +{-# INLINE sfoldl1' #-} + +-- | +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sfoldl1M + :: forall r ix e m + . (Stream r ix e, MonadThrow m) + => (e -> e -> m e) + -> Array r ix e + -> m e +sfoldl1M f arr = do + let str = S.transStepsId $ toStream arr + isNullStream <- S.null str + when isNullStream $ throwM $ SizeEmptyException (outerSize arr) + S.foldl1M f str +{-# INLINE sfoldl1M #-} + +-- | +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sfoldl1M_ + :: forall r ix e m + . (Stream r ix e, MonadThrow m) + => (e -> e -> m e) + -> Array r ix e + -> m () +sfoldl1M_ f = void . sfoldl1M f +{-# INLINE sfoldl1M_ #-} + +-- | +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sifoldl + :: forall r ix e a + . Stream r ix e + => (a -> ix -> e -> a) + -> a + -> Array r ix e + -> a +sifoldl f acc = S.unId . S.foldl (\a (ix, e) -> f a ix e) acc . toStreamIx +{-# INLINE sifoldl #-} + +-- | +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sifoldlM + :: forall r ix e a m + . (Stream r ix e, Monad m) + => (a -> ix -> e -> m a) + -> a + -> Array r ix e + -> m a +sifoldlM f acc = S.foldlM (\a (ix, e) -> f a ix e) acc . S.transStepsId . toStreamIx +{-# INLINE sifoldlM #-} + +-- | +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sifoldlM_ + :: forall r ix e a m + . (Stream r ix e, Monad m) + => (a -> ix -> e -> m a) + -> a + -> Array r ix e + -> m () +sifoldlM_ f acc = void . sifoldlM f acc +{-# INLINE sifoldlM_ #-} + +-- | +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sor + :: forall r ix + . Stream r ix Bool + => Array r ix Bool + -> Bool +sor = S.unId . S.or . toStream +{-# INLINE sor #-} + +-- | +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sand :: forall r ix. Stream r ix Bool => Array r ix Bool -> Bool +sand = S.unId . S.and . toStream +{-# INLINE sand #-} + +-- | +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sany :: forall r ix e. Stream r ix e => (e -> Bool) -> Array r ix e -> Bool +sany f = S.unId . S.or . S.map f . toStream +{-# INLINE sany #-} + +-- | +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sall :: forall r ix e. Stream r ix e => (e -> Bool) -> Array r ix e -> Bool +sall f = S.unId . S.and . S.map f . toStream +{-# INLINE sall #-} + +-- | Add all elements of the array together +-- +-- /Related/: `sum`. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Vector as V +-- >>> V.ssum $ V.sfromList [10, 3, 70, 5 :: Int] +-- 88 +-- +-- @since 0.5.0 +ssum :: forall r ix e. (Num e, Stream r ix e) => Array r ix e -> e +ssum = sfoldl (+) 0 +{-# INLINE ssum #-} + +-- | Multiply all elements of the array together +-- +-- /Related/: `product`. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Vector as V +-- >>> V.sproduct $ V.sfromList [10, 3, 70, 5 :: Int] +-- 10500 +-- +-- @since 0.5.0 +sproduct :: forall r ix e. (Num e, Stream r ix e) => Array r ix e -> e +sproduct = sfoldl (*) 1 +{-# INLINE sproduct #-} + +-- | /O(n)/ - Find the largest value in the array. Throws an error on empty. +-- +-- /Related/: `smaximumM`, `maximum`, `maximumM`. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Vector as V +-- >>> V.smaximum' $ V.sfromList [10, 3, 70, 5 :: Int] +-- 70 +-- +-- @since 0.5.0 +smaximum' :: forall r ix e. (HasCallStack, Ord e, Stream r ix e) => Array r ix e -> e +smaximum' = sfoldl1' max +{-# INLINE smaximum' #-} + +-- | /O(n)/ - Find the largest value in the array. +-- +-- /Related/: `smaximum`, `maximum`, `maximumM`. +-- +-- /__Throws Exceptions__/: `SizeEmptyException` when array is empty +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Vector as V +-- >>> V.smaximumM $ V.sfromList [10, 3, 70, 5 :: Int] +-- 70 +-- >>> V.smaximumM (V.empty :: Vector D Int) :: Maybe Int +-- Nothing +-- +-- @since 0.5.0 +smaximumM :: forall r ix e m. (Ord e, Stream r ix e, MonadThrow m) => Array r ix e -> m e +smaximumM = sfoldl1M (\e acc -> pure (max e acc)) +{-# INLINE smaximumM #-} + +-- | /O(n)/ - Find the smallest value in the array. Throws an error on empty. +-- +-- /Related/: `sminimumM`, `minimum`, `minimumM`. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Vector as V +-- >>> V.sminimum' $ V.sfromList [10, 3, 70, 5 :: Int] +-- 3 +-- +-- @since 0.5.0 +sminimum' :: forall r ix e. (HasCallStack, Ord e, Stream r ix e) => Array r ix e -> e +sminimum' = sfoldl1' min +{-# INLINE sminimum' #-} + +-- | /O(n)/ - Find the smallest value in the array. +-- +-- /Related/: 'sminimum'', `minimum`, `minimumM`. +-- +-- /__Throws Exceptions__/: `SizeEmptyException` when array is empty +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Vector as V +-- >>> V.sminimumM $ V.sfromList [10, 3, 70, 5 :: Int] +-- 3 +-- >>> V.sminimumM (V.empty :: Array D Ix2 Int) :: Maybe Int +-- Nothing +-- +-- @since 0.5.0 +sminimumM :: forall r ix e m. (Ord e, Stream r ix e, MonadThrow m) => Array r ix e -> m e +sminimumM = sfoldl1M (\e acc -> pure (min e acc)) +{-# INLINE sminimumM #-} + +-- | /O(n)/ - left scan with strict accumulator. First element is the value of the +-- accumulator. Last element is not included. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Vector +-- >>> sprescanl min 6 $ sfromList [10, 5, 70, 3 :: Int] +-- Array DS Seq (Sz1 4) +-- [ 6, 6, 5, 5 ] +-- >>> sprescanl (+) 0 $ sfromList [10, 5, 70, 3 :: Int] +-- Array DS Seq (Sz1 4) +-- [ 0, 10, 15, 85 ] +-- +-- @since 1.0.3 +sprescanl :: Stream r ix e => (a -> e -> a) -> a -> Array r ix e -> Vector DS a +sprescanl f acc = DSArray . S.prescanlM (\a b -> pure (f a b)) acc . toStream +{-# INLINE sprescanl #-} + +-- | /O(n)/ - left scan with strict accumulator. First element is the result of applying +-- the supplied function. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Vector +-- >>> spostscanl min 6 $ sfromList [10, 5, 70, 3 :: Int] +-- Array DS Seq (Sz1 4) +-- [ 6, 5, 5, 3 ] +-- >>> spostscanl (+) 0 $ sfromList [10, 5, 70, 3 :: Int] +-- Array DS Seq (Sz1 4) +-- [ 10, 15, 85, 88 ] +-- +-- @since 1.0.3 +spostscanl :: Stream r ix e => (a -> e -> a) -> a -> Array r ix e -> Vector DS a +spostscanl f acc = DSArray . S.postscanlM (\a b -> pure (f a b)) acc . toStream +{-# INLINE spostscanl #-} + +-- | /O(n)/ - Just like `spostscanl` except it is possible to produce a vector with an +-- element type that differes from accumulator type. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Vector +-- >>> spostscanlAcc (\x y -> if x < y then (True, x) else (False, y)) 6 $ sfromList [10, 5, 70, 3 :: Int] +-- Array DS Seq (Sz1 4) +-- [ True, False, True, False ] +-- +-- @since 1.0.3 +spostscanlAcc :: Stream r ix e => (c -> e -> (a, c)) -> c -> Array r ix e -> Vector DS a +spostscanlAcc f acc = DSArray . S.postscanlAccM (\a b -> pure (f a b)) acc . toStream +{-# INLINE spostscanlAcc #-} + +-- | /O(n)/ - left scan with strict accumulator. First element is the value of the accumulator. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Vector +-- >>> sscanl min 6 $ sfromList [10, 5, 70, 3 :: Int] +-- Array DS Seq (Sz1 5) +-- [ 6, 6, 5, 5, 3 ] +-- >>> sscanl (+) 0 $ sfromList [10, 5, 70, 3 :: Int] +-- Array DS Seq (Sz1 5) +-- [ 0, 10, 15, 85, 88 ] +-- +-- @since 1.0.3 +sscanl :: Stream r ix e => (a -> e -> a) -> a -> Array r ix e -> Vector DS a +sscanl f acc = DSArray . S.scanlM (\a b -> pure (f a b)) acc . toStream +{-# INLINE sscanl #-} + +-- | /O(n)/ - left scan with strict accumulator and no initial value for the accumulator. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Vector +-- >>> sscanl1 min $ sfromList [10, 5, 70, 3 :: Int] +-- Array DS Seq (Sz1 4) +-- [ 10, 5, 5, 3 ] +-- >>> sscanl1 (+) $ sfromList [10, 5, 70, 3 :: Int] +-- Array DS Seq (Sz1 4) +-- [ 10, 15, 85, 88 ] +-- >>> sscanl1 (+) $ sfromList ([] :: [Int]) +-- Array DS Seq (Sz1 0) +-- [ ] +-- +-- @since 1.0.3 +sscanl1 :: Stream r ix e => (e -> e -> e) -> Array r ix e -> Vector DS e +sscanl1 f = DSArray . S.scanl1M (\a b -> pure (f a b)) . toStream +{-# INLINE sscanl1 #-} diff --git a/src/Data/Massiv/Vector/Stream.hs b/src/Data/Massiv/Vector/Stream.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Vector/Stream.hs @@ -0,0 +1,934 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE TypeFamilies #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# OPTIONS_HADDOCK hide, not-home #-} + +-- | +-- Module : Data.Massiv.Vector.Stream +-- Copyright : (c) Alexey Kuleshevich 2019-2022 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +module Data.Massiv.Vector.Stream ( + -- | This module has a similar purpose as the 'Data.Vector.Fusion.Bundle.Monadic', but + -- quite a bit simpler. + -- + -- __Important__ - This module is still experimental, as such it is considered + -- internal and exported for the curious users only. + Steps (..), + Stream (..), + + -- * Conversion + steps, + isteps, + consume, + fromStream, + fromStreamM, + fromStreamExactM, + unstreamExact, + unstreamMax, + unstreamMaxM, + unstreamUnknown, + unstreamUnknownM, + unstreamIntoM, + + -- * Operations on Steps + length, + null, + empty, + singleton, + generate, + headMaybe, + last, + cons, + uncons, + snoc, + drop, + take, + slice, + iterateN, + iterateNM, + replicate, + replicateM, + generateM, + traverse, + map, + mapM, + mapM_, + indexed, + concatMap, + append, + zipWith, + zipWith3, + zipWith4, + zipWith5, + zipWith6, + zipWithM, + zipWith3M, + zipWith4M, + zipWith5M, + zipWith6M, + zipWithM_, + zipWith3M_, + zipWith4M_, + zipWith5M_, + zipWith6M_, + + -- ** Folding + foldl, + foldl1, + foldlM, + foldl1M, + foldlLazy, + foldl1Lazy, + foldlLazyM, + foldl1LazyM, + foldrLazy, + foldr1Lazy, + foldrLazyM, + foldr1LazyM, + or, + and, + + -- ** Unfolding + unfoldr, + unfoldrN, + unsafeUnfoldrN, + unfoldrM, + unfoldrNM, + unsafeUnfoldrNM, + unfoldrExactN, + unfoldrExactNM, + + -- ** Scanning + prescanlM, + postscanlM, + postscanlAccM, + scanlM, + scanl1M, + + -- ** Enumeration + enumFromStepN, + + -- * Lists + toList, + fromList, + fromListN, + unsafeFromListN, + + -- ** Filter + mapMaybe, + mapMaybeA, + mapMaybeM, + filter, + filterA, + filterM, + + -- * Transformations + transSteps, + transStepsId, + + -- * Useful re-exports + module Data.Vector.Fusion.Util, + Id (..), +) where + +import qualified Control.Monad as M +import Control.Monad.ST +import Data.Coerce +import qualified Data.Foldable as F +import Data.Massiv.Core.Common hiding (empty, replicate, singleton) +import Data.Maybe (catMaybes) +import qualified Data.Stream.Monadic as S +import qualified Data.Traversable as Traversable (traverse) +import qualified Data.Vector.Fusion.Bundle.Size as B +import Data.Vector.Fusion.Util +import qualified GHC.Exts (IsList (..)) +import Prelude hiding ( + and, + concatMap, + drop, + filter, + foldl, + foldl1, + foldr, + foldr1, + length, + map, + mapM, + mapM_, + null, + or, + replicate, + take, + traverse, + zipWith, + zipWith3, + ) + +instance Monad m => Functor (Steps m) where + fmap f str = str{stepsStream = S.map f (stepsStream str)} + {-# INLINE fmap #-} + (<$) e str = + case stepsSize str of + LengthExact n -> str{stepsStream = S.replicate (coerce n) e} + _ -> fmap (const e) str + {-# INLINE (<$) #-} + +instance Monad m => Semigroup (Steps m e) where + (<>) = append + {-# INLINE (<>) #-} + +instance Monad m => Monoid (Steps m e) where + mempty = empty + {-# INLINE mempty #-} +#if !MIN_VERSION_base(4,11,0) + mappend = append + {-# INLINE mappend #-} +#endif + +instance GHC.Exts.IsList (Steps Id e) where + type Item (Steps Id e) = e + toList = toList + {-# INLINE toList #-} + fromList = fromList + {-# INLINE fromList #-} + fromListN n = (`Steps` LengthMax (Sz n)) . S.fromListN n + {-# INLINE fromListN #-} + +instance Foldable (Steps Id) where + foldr f acc = unId . foldrLazy f acc + {-# INLINE foldr #-} + foldl f acc = unId . foldlLazy f acc + {-# INLINE foldl #-} + foldl' f acc = unId . foldl f acc + {-# INLINE foldl' #-} + foldr1 f = unId . foldr1Lazy f + {-# INLINE foldr1 #-} + foldl1 f = unId . foldl1Lazy f + {-# INLINE foldl1 #-} + toList = toList + {-# INLINE toList #-} + length = unId . length + {-# INLINE length #-} + null = unId . null + {-# INLINE null #-} + sum = unId . foldl (+) 0 + {-# INLINE sum #-} + product = unId . foldl (*) 1 + {-# INLINE product #-} + maximum = unId . foldl1 max + {-# INLINE maximum #-} + minimum = unId . foldl1 min + {-# INLINE minimum #-} + +steps :: forall r ix e m. (Monad m, Index ix, Source r e) => Array r ix e -> Steps m e +steps !arr = + case unsafePrefIndex arr of + PrefIndex gix -> gix <$> ixRangeSteps (size arr) + PrefIndexLinear gi -> + Steps (S.Stream step 0) (LengthExact (coerce k)) + where + !k = totalElem $ size arr + step !i + | i < k = pure $ S.Yield (gi i) (i + 1) + | otherwise = pure S.Done + {-# INLINE [0] step #-} +{-# INLINE [1] steps #-} + +ixRangeSteps :: forall ix m. (Monad m, Index ix) => Sz ix -> Steps m ix +ixRangeSteps sz = Steps (S.Stream step initStep) (LengthExact k) + where + !k = toLinearSz sz + !initStep = if k == zeroSz then Nothing else Just zeroIndex + step (Just ix) = stepNextMF ix (unSz sz) oneIndex (<) $ \mIx -> pure $ S.Yield ix mIx + step Nothing = pure S.Done + {-# INLINE [0] step #-} +{-# INLINE [1] ixRangeSteps #-} + +isteps :: forall r ix e m. (Monad m, Index ix, Source r e) => Array r ix e -> Steps m (ix, e) +isteps !arr = + case unsafePrefIndex arr of + PrefIndex gix -> (\ !ix -> let e = gix ix in e `seq` (ix, e)) <$> ixRangeSteps sz + PrefIndexLinear gi -> + let k = totalElem sz + step i + | i < k = + let e = gi i + in e `seq` pure $ S.Yield (fromLinearIndex sz i, e) (i + 1) + | otherwise = pure S.Done + {-# INLINE [0] step #-} + in Steps (S.Stream step 0) (LengthExact (coerce k)) + where + !sz = size arr +{-# INLINE isteps #-} + +fromStream :: forall r e. Manifest r e => B.Size -> S.Stream Id e -> Vector r e +fromStream sz str = + case B.upperBound sz of + Nothing -> unstreamUnknown str + Just k -> unstreamMax k str +{-# INLINE fromStream #-} + +fromStreamM :: forall r e m. (Monad m, Manifest r e) => B.Size -> S.Stream m e -> m (Vector r e) +fromStreamM sz str = do + xs <- S.toList str + case B.upperBound sz of + Nothing -> pure $! unstreamUnknown (S.fromList xs) + Just k -> pure $! unstreamMax k (S.fromList xs) +{-# INLINE fromStreamM #-} + +fromStreamExactM + :: forall r ix e m + . (Monad m, Manifest r e, Index ix) + => Sz ix + -> S.Stream m e + -> m (Array r ix e) +fromStreamExactM sz str = do + xs <- S.toList str + pure $! unstreamExact sz (S.fromList xs) +{-# INLINE fromStreamExactM #-} + +unstreamIntoM + :: (Manifest r a, PrimMonad m) + => MVector (PrimState m) r a + -> LengthHint + -> S.Stream Id a + -> m (MVector (PrimState m) r a) +unstreamIntoM marr sz str = + case sz of + LengthExact _ -> marr <$ unstreamMaxM marr str + LengthMax _ -> unsafeLinearShrink marr . SafeSz =<< unstreamMaxM marr str + LengthUnknown -> unstreamUnknownM marr str +{-# INLINE unstreamIntoM #-} + +unstreamMax + :: forall r e + . Manifest r e + => Int + -> S.Stream Id e + -> Vector r e +unstreamMax kMax str = + runST $ do + marr <- unsafeNew (SafeSz kMax) + k <- unstreamMaxM marr str + marrShrunk <- + if k == kMax + then pure marr + else unsafeLinearShrink marr (SafeSz k) + unsafeFreeze Seq marrShrunk +{-# INLINE unstreamMax #-} + +unstreamMaxM + :: (Manifest r a, Index ix, PrimMonad m) => MArray (PrimState m) r ix a -> S.Stream Id a -> m Int +unstreamMaxM marr = S.foldlM' fillAtIndex 0 . S.trans (pure . unId) + where + fillAtIndex i x = (i + 1) <$ unsafeLinearWrite marr i x + {-# INLINE fillAtIndex #-} +{-# INLINE unstreamMaxM #-} + +unstreamUnknown :: Manifest r a => S.Stream Id a -> Vector r a +unstreamUnknown str = + runST $ do + marr <- unsafeNew zeroSz + unstreamUnknownM marr str >>= unsafeFreeze Seq +{-# INLINE unstreamUnknown #-} + +unstreamUnknownM + :: (Manifest r a, PrimMonad m) + => MVector (PrimState m) r a + -> S.Stream Id a + -> m (MVector (PrimState m) r a) +unstreamUnknownM marr str = do + (marr', k) <- S.foldlM' fillAtIndex (marr, 0) $ S.trans (pure . unId) str + if k < unSz (sizeOfMArray marr') + then unsafeLinearShrink marr' (SafeSz k) + else pure marr' + where + fillAtIndex (!ma, !i) x = do + let k = unSz (sizeOfMArray ma) + ma' <- + if i < k + then pure ma + else unsafeLinearGrow ma (SafeSz (max 1 k * 2)) + (ma', i + 1) <$ unsafeLinearWrite ma' i x + {-# INLINE fillAtIndex #-} +{-# INLINE unstreamUnknownM #-} + +unstreamExact + :: forall r ix e + . (Manifest r e, Index ix) + => Sz ix + -> S.Stream Id e + -> Array r ix e +unstreamExact sz str = + runST $ do + marr <- unsafeNew sz + _ <- unstreamMaxM marr str + unsafeFreeze Seq marr +{-# INLINE unstreamExact #-} + +length :: Monad m => Steps m a -> m Int +length (Steps str sz) = + case sz of + LengthExact k -> pure $ coerce k + _ -> S.length str +{-# INLINE length #-} + +null :: Monad m => Steps m a -> m Bool +null (Steps str sz) = + case sz of + LengthExact k -> pure (k == zeroSz) + _ -> S.null str +{-# INLINE null #-} + +empty :: Monad m => Steps m e +empty = Steps S.empty (LengthExact zeroSz) +{-# INLINE empty #-} + +singleton :: Monad m => e -> Steps m e +singleton e = Steps (S.singleton e) (LengthExact oneSz) +{-# INLINE singleton #-} + +generate :: Monad m => Sz1 -> (Int -> e) -> Steps m e +generate k f = Steps (S.generate (coerce k) f) (LengthExact k) +{-# INLINE generate #-} + +-- | First element of the 'Stream' or error if empty +headMaybe :: Monad m => Steps m a -> m (Maybe a) +headMaybe (Steps (S.Stream step t) _) = headMaybeLoop S.SPEC t + where + headMaybeLoop !_ s = do + r <- step s + case r of + S.Yield x _ -> pure $ Just x + S.Skip s' -> headMaybeLoop S.SPEC s' + S.Done -> pure Nothing + {-# INLINE [0] headMaybeLoop #-} +{-# INLINE headMaybe #-} + +cons :: Monad m => e -> Steps m e -> Steps m e +cons e (Steps str k) = Steps (S.cons e str) (k `addInt` 1) +{-# INLINE cons #-} + +-- | First element of the `Steps` or `Nothing` if empty +uncons :: Monad m => Steps m e -> m (Maybe (e, Steps m e)) +uncons sts = (\mx -> (,drop oneSz sts) <$> mx) <$> headMaybe sts +{-# INLINE uncons #-} + +snoc :: Monad m => Steps m e -> e -> Steps m e +snoc (Steps str k) e = Steps (S.snoc str e) (k `addInt` 1) +{-# INLINE snoc #-} + +traverse :: (Monad m, Applicative f) => (e -> f a) -> Steps Id e -> f (Steps m a) +traverse f (Steps str k) = (`Steps` k) <$> liftListA (Traversable.traverse f) str +{-# INLINE traverse #-} + +append :: Monad m => Steps m e -> Steps m e -> Steps m e +append (Steps str1 k1) (Steps str2 k2) = Steps (str1 S.++ str2) (k1 `addLengthHint` k2) +{-# INLINE append #-} + +map :: Monad m => (e -> a) -> Steps m e -> Steps m a +map f (Steps str k) = Steps (S.map f str) k +{-# INLINE map #-} + +indexed :: Monad m => Steps m e -> Steps m (Int, e) +indexed (Steps str k) = Steps (S.indexed str) k +{-# INLINE indexed #-} + +mapM :: Monad m => (e -> m a) -> Steps m e -> Steps m a +mapM f (Steps str k) = Steps (S.mapM f str) k +{-# INLINE mapM #-} + +mapM_ :: Monad m => (e -> m a) -> Steps m e -> m () +mapM_ f (Steps str _) = S.mapM_ f str +{-# INLINE mapM_ #-} + +zipWith :: Monad m => (a -> b -> e) -> Steps m a -> Steps m b -> Steps m e +zipWith f (Steps sa ka) (Steps sb kb) = Steps (S.zipWith f sa sb) (minLengthHint ka kb) +{-# INLINE zipWith #-} + +zipWith3 :: Monad m => (a -> b -> c -> d) -> Steps m a -> Steps m b -> Steps m c -> Steps m d +zipWith3 f (Steps sa ka) (Steps sb kb) (Steps sc kc) = + Steps (S.zipWith3 f sa sb sc) (minLengthHint ka (minLengthHint kb kc)) +{-# INLINE zipWith3 #-} + +zipWith4 + :: Monad m => (a -> b -> c -> d -> e) -> Steps m a -> Steps m b -> Steps m c -> Steps m d -> Steps m e +zipWith4 f (Steps sa ka) (Steps sb kb) (Steps sc kc) (Steps sd kd) = + Steps (S.zipWith4 f sa sb sc sd) (minLengthHint ka (minLengthHint kb (minLengthHint kc kd))) +{-# INLINE zipWith4 #-} + +zipWith5 + :: Monad m + => (a -> b -> c -> d -> e -> f) + -> Steps m a + -> Steps m b + -> Steps m c + -> Steps m d + -> Steps m e + -> Steps m f +zipWith5 f (Steps sa ka) (Steps sb kb) (Steps sc kc) (Steps sd kd) (Steps se ke) = + Steps + (S.zipWith5 f sa sb sc sd se) + (minLengthHint ka (minLengthHint kb (minLengthHint kc (minLengthHint kd ke)))) +{-# INLINE zipWith5 #-} + +zipWith6 + :: Monad m + => (a -> b -> c -> d -> e -> f -> g) + -> Steps m a + -> Steps m b + -> Steps m c + -> Steps m d + -> Steps m e + -> Steps m f + -> Steps m g +zipWith6 f (Steps sa ka) (Steps sb kb) (Steps sc kc) (Steps sd kd) (Steps se ke) (Steps sf kf) = + Steps + (S.zipWith6 f sa sb sc sd se sf) + (minLengthHint ka (minLengthHint kb (minLengthHint kc (minLengthHint kd (minLengthHint ke kf))))) +{-# INLINE zipWith6 #-} + +zipWithM :: Monad m => (a -> b -> m c) -> Steps m a -> Steps m b -> Steps m c +zipWithM f (Steps sa ka) (Steps sb kb) = Steps (S.zipWithM f sa sb) (minLengthHint ka kb) +{-# INLINE zipWithM #-} + +zipWith3M :: Monad m => (a -> b -> c -> m d) -> Steps m a -> Steps m b -> Steps m c -> Steps m d +zipWith3M f (Steps sa ka) (Steps sb kb) (Steps sc kc) = + Steps (S.zipWith3M f sa sb sc) (minLengthHint ka (minLengthHint kb kc)) +{-# INLINE zipWith3M #-} + +zipWith4M + :: Monad m + => (a -> b -> c -> d -> m e) + -> Steps m a + -> Steps m b + -> Steps m c + -> Steps m d + -> Steps m e +zipWith4M f (Steps sa ka) (Steps sb kb) (Steps sc kc) (Steps sd kd) = + Steps (S.zipWith4M f sa sb sc sd) (minLengthHint ka (minLengthHint kb (minLengthHint kc kd))) +{-# INLINE zipWith4M #-} + +zipWith5M + :: Monad m + => (a -> b -> c -> d -> e -> m f) + -> Steps m a + -> Steps m b + -> Steps m c + -> Steps m d + -> Steps m e + -> Steps m f +zipWith5M f (Steps sa ka) (Steps sb kb) (Steps sc kc) (Steps sd kd) (Steps se ke) = + Steps + (S.zipWith5M f sa sb sc sd se) + (minLengthHint ka (minLengthHint kb (minLengthHint kc (minLengthHint kd ke)))) +{-# INLINE zipWith5M #-} + +zipWith6M + :: Monad m + => (a -> b -> c -> d -> e -> f -> m g) + -> Steps m a + -> Steps m b + -> Steps m c + -> Steps m d + -> Steps m e + -> Steps m f + -> Steps m g +zipWith6M f (Steps sa ka) (Steps sb kb) (Steps sc kc) (Steps sd kd) (Steps se ke) (Steps sf kf) = + Steps + (S.zipWith6M f sa sb sc sd se sf) + (minLengthHint ka (minLengthHint kb (minLengthHint kc (minLengthHint kd (minLengthHint ke kf))))) +{-# INLINE zipWith6M #-} + +zipWithM_ :: Monad m => (a -> b -> m c) -> Steps m a -> Steps m b -> m () +zipWithM_ f (Steps str1 _) (Steps str2 _) = S.zipWithM_ f str1 str2 +{-# INLINE zipWithM_ #-} + +zipWith3M_ :: Monad m => (a -> b -> c -> m d) -> Steps m a -> Steps m b -> Steps m c -> m () +zipWith3M_ f sa sb sc = consume $ zipWith3M f sa sb sc +{-# INLINE zipWith3M_ #-} + +zipWith4M_ + :: Monad m + => (a -> b -> c -> d -> m e) + -> Steps m a + -> Steps m b + -> Steps m c + -> Steps m d + -> m () +zipWith4M_ f sa sb sc sd = consume $ zipWith4M f sa sb sc sd +{-# INLINE zipWith4M_ #-} + +zipWith5M_ + :: Monad m + => (a -> b -> c -> d -> e -> m f) + -> Steps m a + -> Steps m b + -> Steps m c + -> Steps m d + -> Steps m e + -> m () +zipWith5M_ f sa sb sc sd se = consume $ zipWith5M f sa sb sc sd se +{-# INLINE zipWith5M_ #-} + +zipWith6M_ + :: Monad m + => (a -> b -> c -> d -> e -> f -> m g) + -> Steps m a + -> Steps m b + -> Steps m c + -> Steps m d + -> Steps m e + -> Steps m f + -> m () +zipWith6M_ f sa sb sc sd se sf = consume $ zipWith6M f sa sb sc sd se sf +{-# INLINE zipWith6M_ #-} + +consume :: Monad m => Steps m a -> m () +consume (Steps (S.Stream step t) _) = consumeLoop S.SPEC t + where + consumeLoop !_ s = do + r <- step s + case r of + S.Yield _ s' -> consumeLoop S.SPEC s' + S.Skip s' -> consumeLoop S.SPEC s' + S.Done -> pure () +{-# INLINE consume #-} + +transStepsId :: Monad m => Steps Id e -> Steps m e +transStepsId (Steps sts k) = Steps (S.trans (pure . unId) sts) k +{-# INLINE transStepsId #-} + +transSteps :: (Monad m, Monad n) => Steps m e -> m (Steps n e) +transSteps (Steps strM sz@(LengthExact _)) = (`Steps` sz) <$> transListM strM +transSteps (Steps strM _) = do + (n, strN) <- transListNM strM + pure (Steps strN (LengthExact n)) +{-# INLINE transSteps #-} + +foldl :: Monad m => (b -> a -> b) -> b -> Steps m a -> m b +foldl f acc = S.foldl' f acc . stepsStream +{-# INLINE foldl #-} + +foldl1 :: Monad m => (a -> a -> a) -> Steps m a -> m a +foldl1 f = S.foldl1' f . stepsStream +{-# INLINE foldl1 #-} + +foldlM :: Monad m => (a -> b -> m a) -> a -> Steps m b -> m a +foldlM f acc = S.foldlM' f acc . stepsStream +{-# INLINE foldlM #-} + +foldl1M :: Monad m => (a -> a -> m a) -> Steps m a -> m a +foldl1M f (Steps sts _) = S.foldl1M' f sts +{-# INLINE foldl1M #-} + +foldrLazy :: Monad m => (a -> b -> b) -> b -> Steps m a -> m b +foldrLazy f acc = S.foldr f acc . stepsStream +{-# INLINE foldrLazy #-} + +foldr1Lazy :: Monad m => (a -> a -> a) -> Steps m a -> m a +foldr1Lazy f = S.foldr1 f . stepsStream +{-# INLINE foldr1Lazy #-} + +foldlLazy :: Monad m => (b -> a -> b) -> b -> Steps m a -> m b +foldlLazy f acc = S.foldl f acc . stepsStream +{-# INLINE foldlLazy #-} + +foldl1Lazy :: Monad m => (a -> a -> a) -> Steps m a -> m a +foldl1Lazy f = S.foldl1 f . stepsStream +{-# INLINE foldl1Lazy #-} + +foldlLazyM :: Monad m => (a -> b -> m a) -> a -> Steps m b -> m a +foldlLazyM f acc = S.foldlM f acc . stepsStream +{-# INLINE foldlLazyM #-} + +foldl1LazyM :: Monad m => (a -> a -> m a) -> Steps m a -> m a +foldl1LazyM f (Steps sts _) = S.foldl1M f sts +{-# INLINE foldl1LazyM #-} + +foldrLazyM :: Monad m => (b -> a -> m a) -> a -> Steps m b -> m a +foldrLazyM f acc (Steps sts _) = S.foldrM f acc sts +{-# INLINE foldrLazyM #-} + +foldr1LazyM :: Monad m => (a -> a -> m a) -> Steps m a -> m a +foldr1LazyM f = S.foldr1M f . stepsStream +{-# INLINE foldr1LazyM #-} + +or :: Monad m => Steps m Bool -> m Bool +or = S.or . stepsStream +{-# INLINE or #-} + +and :: Monad m => Steps m Bool -> m Bool +and = S.and . stepsStream +{-# INLINE and #-} + +mapMaybe :: Monad m => (a -> Maybe e) -> Steps m a -> Steps m e +mapMaybe f (Steps str k) = Steps (S.mapMaybe f str) (toLengthMax k) +{-# INLINE mapMaybe #-} + +concatMap :: Monad m => (a -> Steps m e) -> Steps m a -> Steps m e +concatMap f (Steps str _) = Steps (S.concatMap (stepsStream . f) str) LengthUnknown +{-# INLINE concatMap #-} + +mapMaybeA :: (Monad m, Applicative f) => (a -> f (Maybe e)) -> Steps Id a -> f (Steps m e) +mapMaybeA f (Steps str k) = (`Steps` toLengthMax k) <$> liftListA (mapMaybeListA f) str +{-# INLINE mapMaybeA #-} + +mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Steps m a -> Steps m b +mapMaybeM f (Steps str k) = Steps (mapMaybeStreamM f str) (toLengthMax k) +{-# INLINE mapMaybeM #-} + +mapMaybeListA :: Applicative f => (a -> f (Maybe b)) -> [a] -> f [b] +mapMaybeListA f = fmap catMaybes . Traversable.traverse f +{-# INLINE mapMaybeListA #-} + +mapMaybeStreamM :: Monad m => (a -> m (Maybe b)) -> S.Stream m a -> S.Stream m b +mapMaybeStreamM f (S.Stream step t) = S.Stream step' t + where + step' s = do + r <- step s + case r of + S.Yield x s' -> do + b <- f x + pure $ + case b of + Nothing -> S.Skip s' + Just b' -> S.Yield b' s' + S.Skip s' -> pure $ S.Skip s' + S.Done -> pure S.Done + {-# INLINE [0] step' #-} +{-# INLINE mapMaybeStreamM #-} + +filter :: Monad m => (a -> Bool) -> Steps m a -> Steps m a +filter f (Steps str k) = Steps (S.filter f str) (toLengthMax k) +{-# INLINE filter #-} + +filterA :: (Monad m, Applicative f) => (e -> f Bool) -> Steps Id e -> f (Steps m e) +filterA f (Steps str k) = (`Steps` toLengthMax k) <$> liftListA (M.filterM f) str +{-# INLINE filterA #-} + +filterM :: Monad m => (e -> m Bool) -> Steps m e -> Steps m e +filterM f (Steps str k) = Steps (S.filterM f str) (toLengthMax k) +{-# INLINE filterM #-} + +take :: Monad m => Sz1 -> Steps m a -> Steps m a +take n (Steps str sz) = + Steps (S.take (coerce n) str) $! + case sz of + LengthExact k -> LengthExact (inline0 min n k) + LengthMax k -> LengthMax (inline0 min n k) + LengthUnknown -> LengthUnknown +{-# INLINE take #-} + +drop :: Monad m => Sz1 -> Steps m a -> Steps m a +drop n (Steps str k) = Steps (S.drop (coerce n) str) (k `subtractLengthHint` LengthExact n) +{-# INLINE drop #-} + +slice :: Monad m => Int -> Sz1 -> Steps m a -> Steps m a +slice i k (Steps str _) = Steps (S.slice i (coerce k) str) (LengthMax k) +{-# INLINE slice #-} + +iterateN :: Monad m => Sz1 -> (a -> a) -> a -> Steps m a +iterateN n f a = Steps (S.iterateN (coerce n) f a) (LengthExact n) +{-# INLINE iterateN #-} + +iterateNM :: Monad m => Sz1 -> (a -> m a) -> a -> Steps m a +iterateNM n f a = Steps (S.iterateNM (coerce n) f a) (LengthExact n) +{-# INLINE iterateNM #-} + +replicate :: Monad m => Sz1 -> a -> Steps m a +replicate n a = Steps (S.replicate (coerce n) a) (LengthExact n) +{-# INLINE replicate #-} + +replicateM :: Monad m => Sz1 -> m a -> Steps m a +replicateM n f = Steps (S.replicateM (coerce n) f) (LengthExact n) +{-# INLINE replicateM #-} + +generateM :: Monad m => Sz1 -> (Int -> m a) -> Steps m a +generateM n f = Steps (S.generateM (coerce n) f) (LengthExact n) +{-# INLINE generateM #-} + +unfoldr :: Monad m => (s -> Maybe (e, s)) -> s -> Steps m e +unfoldr f e0 = Steps (S.unfoldr f e0) LengthUnknown +{-# INLINE unfoldr #-} + +unfoldrN :: Monad m => Sz1 -> (s -> Maybe (e, s)) -> s -> Steps m e +unfoldrN n f e0 = Steps (S.unfoldrN (coerce n) f e0) LengthUnknown +{-# INLINE unfoldrN #-} + +unsafeUnfoldrN :: Monad m => Sz1 -> (s -> Maybe (e, s)) -> s -> Steps m e +unsafeUnfoldrN n f e0 = Steps (S.unfoldrN (coerce n) f e0) (LengthMax n) +{-# INLINE unsafeUnfoldrN #-} + +unfoldrM :: Monad m => (s -> m (Maybe (e, s))) -> s -> Steps m e +unfoldrM f e0 = Steps (S.unfoldrM f e0) LengthUnknown +{-# INLINE unfoldrM #-} + +unfoldrNM :: Monad m => Int -> (s -> m (Maybe (e, s))) -> s -> Steps m e +unfoldrNM n f e0 = Steps (S.unfoldrNM n f e0) LengthUnknown +{-# INLINE unfoldrNM #-} + +unsafeUnfoldrNM :: Monad m => Sz1 -> (s -> m (Maybe (e, s))) -> s -> Steps m e +unsafeUnfoldrNM n f e0 = Steps (S.unfoldrNM (coerce n) f e0) (LengthMax n) +{-# INLINE unsafeUnfoldrNM #-} + +unfoldrExactN :: Monad m => Sz1 -> (s -> (a, s)) -> s -> Steps m a +unfoldrExactN n f = unfoldrExactNM n (pure . f) +{-# INLINE unfoldrExactN #-} + +unfoldrExactNM :: Monad m => Sz1 -> (s -> m (a, s)) -> s -> Steps m a +unfoldrExactNM n f t = Steps (S.Stream step (t, unSz n)) (LengthExact n) + where + step (s, i) + | i <= 0 = pure S.Done + | otherwise = fmap (\(x, s') -> S.Yield x (s', i - 1)) (f s) + {-# INLINE [0] step #-} +{-# INLINE unfoldrExactNM #-} + +enumFromStepN :: (Num a, Monad m) => a -> a -> Sz1 -> Steps m a +enumFromStepN x step k = Steps (S.enumFromStepN x step (coerce k)) (LengthExact k) +{-# INLINE enumFromStepN #-} + +toList :: Steps Id e -> [e] +toList (Steps str _) = unId (S.toList str) +{-# INLINE toList #-} + +fromList :: Monad m => [e] -> Steps m e +fromList = (`Steps` LengthUnknown) . S.fromList +{-# INLINE fromList #-} + +fromListN :: Monad m => Int -> [e] -> Steps m e +fromListN n = (`Steps` LengthUnknown) . S.fromListN n +{-# INLINE fromListN #-} + +unsafeFromListN :: Monad m => Sz1 -> [e] -> Steps m e +unsafeFromListN n = (`Steps` LengthMax n) . S.fromListN (coerce n) +{-# INLINE unsafeFromListN #-} + +liftListA :: (Monad m, Functor f) => ([a] -> f [b]) -> S.Stream Id a -> f (S.Stream m b) +liftListA f str = S.fromList <$> f (unId (S.toList str)) +{-# INLINE liftListA #-} + +transListM :: (Monad m, Monad n) => S.Stream m a -> m (S.Stream n a) +transListM str = do + xs <- S.toList str + pure $ S.fromList xs +{-# INLINE transListM #-} + +transListNM :: (Monad m, Monad n) => S.Stream m a -> m (Sz1, S.Stream n a) +transListNM str = do + (n, xs) <- toListN str + pure (coerce n, S.fromList xs) +{-# INLINE transListNM #-} + +toListN :: Monad m => S.Stream m a -> m (Int, [a]) +toListN = S.foldr (\x (i, xs) -> (i + 1, x : xs)) (0, []) +{-# INLINE toListN #-} + +addHint :: (Sz1 -> LengthHint) -> Int -> Int -> LengthHint +addHint hint m n + | k == coerce sz = hint sz + | otherwise = LengthUnknown -- overflow + where + k = m + n + sz = Sz k +{-# INLINE addHint #-} + +addInt :: LengthHint -> Int -> LengthHint +addInt (LengthExact m) n = addHint LengthExact (coerce m) (coerce n) +addInt (LengthMax m) n = addHint LengthExact (coerce m) n +addInt _ _ = LengthUnknown +{-# INLINE addInt #-} + +addLengthHint :: LengthHint -> LengthHint -> LengthHint +addLengthHint (LengthExact m) (LengthExact n) = addHint LengthExact (coerce m) (coerce n) +addLengthHint (LengthMax m) (LengthExact n) = addHint LengthMax (coerce m) (coerce n) +addLengthHint (LengthExact m) (LengthMax n) = addHint LengthMax (coerce m) (coerce n) +addLengthHint (LengthMax m) (LengthMax n) = addHint LengthMax (coerce m) (coerce n) +addLengthHint _ _ = LengthUnknown +{-# INLINE addLengthHint #-} + +subtractLengthHint :: LengthHint -> LengthHint -> LengthHint +subtractLengthHint (LengthExact m) (LengthExact n) = LengthExact (m - n) +subtractLengthHint (LengthMax m) (LengthExact n) = LengthMax (m - n) +subtractLengthHint (LengthExact m) (LengthMax _) = LengthMax m +subtractLengthHint (LengthMax m) (LengthMax _) = LengthMax m +subtractLengthHint _ _ = LengthUnknown +{-# INLINE subtractLengthHint #-} + +minLengthHint :: LengthHint -> LengthHint -> LengthHint +minLengthHint (LengthExact m) (LengthExact n) = LengthExact (inline0 min m n) +minLengthHint (LengthExact m) (LengthMax n) = LengthMax (inline0 min m n) +minLengthHint (LengthExact m) LengthUnknown = LengthMax m +minLengthHint (LengthMax m) (LengthExact n) = LengthMax (inline0 min m n) +minLengthHint (LengthMax m) (LengthMax n) = LengthMax (inline0 min m n) +minLengthHint (LengthMax m) LengthUnknown = LengthMax m +minLengthHint LengthUnknown (LengthExact n) = LengthMax n +minLengthHint LengthUnknown (LengthMax n) = LengthMax n +minLengthHint LengthUnknown LengthUnknown = LengthUnknown +{-# INLINE minLengthHint #-} + +toLengthMax :: LengthHint -> LengthHint +toLengthMax (LengthExact n) = LengthMax n +toLengthMax (LengthMax n) = LengthMax n +toLengthMax LengthUnknown = LengthUnknown +{-# INLINE toLengthMax #-} + +-- | Prefix scan with strict accumulator and a monadic operator +prescanlM :: Monad m => (a -> b -> m a) -> a -> Steps m b -> Steps m a +prescanlM f acc ss = ss{stepsStream = S.prescanlM' f acc (stepsStream ss)} +{-# INLINE prescanlM #-} + +-- | Suffix scan with a monadic operator +postscanlM :: Monad m => (a -> b -> m a) -> a -> Steps m b -> Steps m a +postscanlM f acc ss = ss{stepsStream = S.postscanlM' f acc (stepsStream ss)} +{-# INLINE postscanlM #-} + +-- | Suffix scan with a monadic operator +postscanlAccM :: Monad m => (c -> b -> m (a, c)) -> c -> Steps m b -> Steps m a +postscanlAccM f acc ss = ss{stepsStream = postscanlAccStreamM f acc (stepsStream ss)} +{-# INLINE postscanlAccM #-} + +-- | Suffix scan with strict acccumulator and a monadic operator +postscanlAccStreamM :: Monad m => (c -> b -> m (a, c)) -> c -> S.Stream m b -> S.Stream m a +postscanlAccStreamM f w (S.Stream step t) = w `seq` S.Stream step' (t, w) + where + step' (s, x) = + x `seq` + do + r <- step s + case r of + S.Yield y s' -> do + (a, z) <- f x y + z `seq` return (S.Yield a (s', z)) + S.Skip s' -> return $ S.Skip (s', x) + S.Done -> return S.Done + {-# INLINE [0] step' #-} +{-# INLINE postscanlAccStreamM #-} + +-- | Haskell-style scan with a monadic operator +scanlM :: Monad m => (a -> b -> m a) -> a -> Steps m b -> Steps m a +scanlM f acc Steps{stepsStream, stepsSize} = + Steps + { stepsStream = S.scanlM' f acc stepsStream + , stepsSize = addLengthHint (LengthExact 1) stepsSize + } +{-# INLINE scanlM #-} + +-- | Initial-value free scan over a 'Stream' with a strict accumulator +-- and a monadic operator +scanl1M :: Monad m => (a -> a -> m a) -> Steps m a -> Steps m a +scanl1M f ss = ss{stepsStream = S.scanl1M' f (stepsStream ss)} +{-# INLINE scanl1M #-} diff --git a/src/Data/Massiv/Vector/Unsafe.hs b/src/Data/Massiv/Vector/Unsafe.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Vector/Unsafe.hs @@ -0,0 +1,172 @@ +{-# LANGUAGE FlexibleContexts #-} + +-- | +-- Module : Data.Massiv.Vector.Unsafe +-- Copyright : (c) Alexey Kuleshevich 2020-2022 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +module Data.Massiv.Vector.Unsafe ( + -- * Vector + + -- ** Accessors + + -- *** Indexing + unsafeHead, + unsafeLast, + + -- *** Monadic Indexing + unsafeIndexM, + unsafeHeadM, + unsafeLastM, + + -- *** Slicing + unsafeInit, + unsafeTail, + unsafeTake, + unsafeDrop, + -- -- ** Modifying + -- -- *** Bulk updates + -- , unsafeUpdate + -- , unsafeUpdate_ + -- -- *** Accumulation + -- , unsafeAccum + -- , unsafeAccumulate_ + -- , unsafeBackpermute + -- -- ** Predicates + -- , unsafePartition + + -- ** Unbounded streams + unsafeUnfoldrN, + unsafeUnfoldrNM, + unsafeFromListN, +) where + +import Data.Coerce +import Data.Massiv.Array.Delayed.Stream +import Data.Massiv.Core.Common +import qualified Data.Massiv.Vector.Stream as S + +-- ========= -- +-- Accessors -- +-- ========= -- + +-------------- +-- Indexing -- +-------------- + +-- | +-- +-- @since 0.5.0 +unsafeHead :: Source r e => Vector r e -> e +unsafeHead = (`unsafeLinearIndex` 0) +{-# INLINE unsafeHead #-} + +-- | +-- +-- @since 0.5.0 +unsafeLast :: Source r e => Vector r e -> e +unsafeLast v = unsafeLinearIndex v (unSz (size v) - 1) +{-# INLINE unsafeLast #-} + +---------------------- +-- Monadic indexing -- +---------------------- + +-- | +-- +-- @since 0.5.0 +unsafeIndexM :: (Source r e, Monad m) => Vector r e -> Ix1 -> m e +unsafeIndexM v i = pure $! unsafeLinearIndex v i +{-# INLINE unsafeIndexM #-} + +-- | +-- +-- @since 0.5.0 +unsafeHeadM :: (Monad m, Source r e) => Vector r e -> m e +unsafeHeadM v = pure $! unsafeHead v +{-# INLINE unsafeHeadM #-} + +-- | +-- +-- @since 0.5.0 +unsafeLastM :: (Monad m, Source r e) => Vector r e -> m e +unsafeLastM v = pure $! unsafeLast v +{-# INLINE unsafeLastM #-} + +------------- +-- Slicing -- +------------- + +-- | +-- +-- @since 0.5.0 +unsafeInit :: Source r e => Vector r e -> Vector r e +unsafeInit v = unsafeLinearSlice 0 (SafeSz (coerce (size v) - 1)) v +{-# INLINE unsafeInit #-} + +-- | +-- +-- @since 0.5.0 +unsafeTail :: Source r e => Vector r e -> Vector r e +unsafeTail = unsafeDrop oneSz +{-# INLINE unsafeTail #-} + +-- | +-- +-- @since 0.5.0 +unsafeTake :: Source r e => Sz1 -> Vector r e -> Vector r e +unsafeTake = unsafeLinearSlice 0 +{-# INLINE unsafeTake #-} + +-- | +-- +-- @since 0.5.0 +unsafeDrop :: Source r e => Sz1 -> Vector r e -> Vector r e +unsafeDrop (Sz d) v = unsafeLinearSlice d (SafeSz (coerce (size v) - d)) v +{-# INLINE unsafeDrop #-} + +-- | /O(n)/ - Convert a list of a known length to a delayed stream vector. +-- +-- /Unsafe/ - This function is unsafe because it will allocate enough space in memory for +-- @n@ elements ahead of time, regardless of the actual size of the list. Supplying @n@ +-- that is too big will result in an asynchronous `Control.Exception.Base.HeapOverflow` +-- exception. +-- +-- @since 0.5.1 +unsafeFromListN :: Sz1 -> [e] -> Vector DS e +unsafeFromListN n = fromSteps . S.unsafeFromListN n +{-# INLINE unsafeFromListN #-} + +-- | /O(n)/ - Right unfolding function with at most @n@ number of elements. +-- +-- /Unsafe/ - This function is unsafe because it will allocate enough space in memory for +-- @n@ elements ahead of time, regardless of when unfolding function returns a +-- `Nothing`. Supplying @n@ that is too big will result in an asynchronous +-- `Control.Exception.Base.HeapOverflow` exception. +-- +-- @since 0.5.1 +unsafeUnfoldrN + :: Sz1 + -- ^ @n@ - maximum number of elements that the vector will have + -> (s -> Maybe (e, s)) + -- ^ Unfolding function. Stops when `Nothing` is returned or maximum number of elements + -- is reached. + -> s + -- ^ Inititial element. + -> Vector DS e +unsafeUnfoldrN n f = DSArray . S.unsafeUnfoldrN n f +{-# INLINE unsafeUnfoldrN #-} + +-- | /O(n)/ - Same as `unsafeUnfoldrN`, but with monadic generating function. +-- +-- /Unsafe/ - This function is unsafe because it will allocate enough space in memory for +-- @n@ elements ahead of time, regardless of when unfolding function returns a +-- `Nothing`. Supplying @n@ that is too big will result in an asynchronous +-- `Control.Exception.Base.HeapOverflow` exception. +-- +-- @since 0.5.1 +unsafeUnfoldrNM :: Monad m => Sz1 -> (s -> m (Maybe (e, s))) -> s -> m (Vector DS e) +unsafeUnfoldrNM n f = fromStepsM . S.unsafeUnfoldrNM n f +{-# INLINE unsafeUnfoldrNM #-} diff --git a/tests/doctests.hs b/tests/doctests.hs --- a/tests/doctests.hs +++ b/tests/doctests.hs @@ -1,12 +1,19 @@ +{-# LANGUAGE CPP #-} + module Main where -import Build_doctests (flags, pkgs, module_sources) -import Data.Foldable (traverse_) +#if __GLASGOW_HASKELL__ >= 802 && __GLASGOW_HASKELL__ < 810 + import Test.DocTest (doctest) main :: IO () -main = do - traverse_ putStrLn args - doctest args - where - args = flags ++ pkgs ++ module_sources +main = doctest ["-Iinclude","src"] + +#else + +-- TODO: fix doctest support +main :: IO () +main = + putStrLn "\nDoctests are not supported for ghc version 8.2 and prior as well as 8.10 and newer\n" + +#endif