diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,58 @@ +# 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_`, @@ -6,6 +61,7 @@ * 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: @@ -54,6 +110,14 @@ * 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 `> + 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 +λ> 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`: + +* [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/massiv.cabal b/massiv.cabal --- a/massiv.cabal +++ b/massiv.cabal @@ -1,5 +1,5 @@ name: massiv -version: 0.6.1.0 +version: 1.0.0.0 synopsis: Massiv (Массив) is an Array Library. description: Multi-dimensional Arrays with fusion, stencils and parallel computation. homepage: https://github.com/lehins/massiv @@ -20,6 +20,7 @@ , GHC == 8.6.5 , GHC == 8.8.4 , GHC == 8.10.2 + , GHC == 9.0.1 flag unsafe-checks description: Enable all the bounds checks for unsafe functions at the cost of @@ -81,8 +82,9 @@ , bytestring , deepseq , exceptions - , scheduler >= 1.5.0 - , primitive + , scheduler >= 2.0.0 && < 3.0.0 + , primitive >= 0.7.1.0 + , random >= 1.2.0 , unliftio-core , vector @@ -97,6 +99,7 @@ -Wincomplete-record-updates -Wincomplete-uni-patterns -Wredundant-constraints + -Wno-simplifiable-class-constraints test-suite doctests type: exitcode-stdio-1.0 @@ -107,7 +110,8 @@ , QuickCheck , massiv , mersenne-random-pure64 - , random + , random >= 1.2.0 + , mwc-random >= 0.15.0.1 , splitmix >= 0.0.1 , template-haskell default-language: Haskell2010 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 @@ -24,7 +24,8 @@ -- 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`, is also a boxed type, but lazy. It's elements are not evaluated. +-- * `BL` - Similar to `B`, is also a boxed type, but lazy. It's 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 @@ -36,9 +37,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 @@ -64,15 +62,15 @@ -- -- 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 @@ -85,6 +83,7 @@ -- * Compute , getComp , setComp + , appComp , compute , computeS , computeP @@ -108,6 +107,8 @@ , elemsCount , isEmpty , isNotEmpty + , isNull + , isNotNull -- * Indexing , (!?) , (!) 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 @@ -11,6 +11,8 @@ -- ** Delayed Pull Array D(..) , delay + , liftArray2' + , liftArray2M -- ** Delayed Push Array , DL(..) , toLoadArray 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 @@ -15,6 +15,7 @@ -- module Data.Massiv.Array.Delayed.Interleaved ( DI(..) + , Array(..) , toInterleaved , fromInterleaved ) where @@ -30,52 +31,52 @@ 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 +instance Strategy DI where setComp c arr = arr { diArray = (diArray arr) { dComp = c } } {-# INLINE setComp #-} + getComp = dComp . diArray + {-# INLINE getComp #-} - 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 = + makeArray c sz = DIArray . makeArray c sz + {-# INLINE makeArray #-} + iterArrayLinearST_ 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 #-} + {-# INLINE iterArrayLinearST_ #-} instance Index ix => StrideLoad DI ix 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 #-} + {-# 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 @@ -21,24 +21,30 @@ , eqArrays , compareArrays , imap - , liftArray2Matching + , liftArray2' + , liftArray2M + , unsafeExtract + , unsafeSlice + , unsafeInnerSlice ) where -import Control.Applicative +import Control.Applicative import qualified Data.Foldable as F -import Data.Massiv.Array.Ops.Fold.Internal as A -import Data.Massiv.Vector.Stream as S (steps) -import Data.Massiv.Core.Common -import Data.Massiv.Core.Operations -import Data.Massiv.Core.List (L, showArrayList, showsArrayPrec) -import GHC.Base (build) -import Prelude hiding (zipWith) +import Data.Massiv.Array.Ops.Fold.Internal as A +import Data.Massiv.Core.Common as A +import Data.Massiv.Core.List (L, showArrayList, showsArrayPrec) +import Data.Massiv.Core.Operations +import Data.Massiv.Vector.Stream as S (steps) +import GHC.Base (build) +import Prelude hiding (zipWith) #include "massiv.h" + -- | Delayed representation. data D = D deriving Show + data instance Array D ix e = DArray { dComp :: !Comp , dSize :: !(Sz ix) , dIndex :: ix -> e } @@ -47,59 +53,59 @@ 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)) {-# 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 +instance Strategy D where setComp c arr = arr { dComp = c } {-# INLINE setComp #-} - - makeArray = DArray - {-# INLINE makeArray #-} - + getComp = dComp + {-# INLINE getComp #-} -instance Index ix => Source D ix e where +instance Source D e where unsafeIndex = INDEX_CHECK("(Source D ix e).unsafeIndex", size, dIndex) {-# INLINE unsafeIndex #-} + + unsafeOuterSlice !arr !szL !i = DArray (dComp arr) szL (unsafeIndex arr . consDim i) + {-# INLINE unsafeOuterSlice #-} + unsafeLinearSlice !o !sz arr = DArray (dComp arr) sz $ \ !i -> unsafeIndex arr (fromLinearIndex (size arr) (i + o)) {-# 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 #-} +-- | /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 = + DArray (getComp arr) newSz (unsafeIndex arr . liftIndex2 (+) sIx) +{-# INLINE unsafeExtract #-} +-- | /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 #-} -instance (Elt D ix e ~ Array D (Lower ix) e, Index ix) => OuterSlice D ix e where +-- | /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 (unsafeIndex arr . (`snocDim` i)) +{-# INLINE unsafeInnerSlice #-} - unsafeOuterSlice !arr !i = - DArray (dComp arr) (snd (unconsSz (size arr))) (\ !ix -> unsafeIndex arr (consDim i ix)) - {-# INLINE unsafeOuterSlice #-} -instance (Elt D ix e ~ Array D (Lower ix) e, Index ix) => InnerSlice D ix e where - unsafeInnerSlice !arr (szL, _) !i = - DArray (dComp arr) szL (\ !ix -> unsafeIndex arr (snocDim ix i)) - {-# INLINE unsafeInnerSlice #-} - - instance (Eq e, Index ix) => Eq (Array D ix e) where (==) = eqArrays (==) {-# INLINE (==) #-} @@ -118,9 +124,12 @@ instance Index ix => Applicative (Array D ix) where pure = singleton {-# INLINE pure #-} - liftA2 = liftArray2Matching + (<*>) = 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 @@ -147,12 +156,11 @@ 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 = DArray + {-# INLINE makeArray #-} + iterArrayLinearST_ !scheduler !arr = + splitLinearlyWith_ scheduler (elemsCount arr) (unsafeLinearIndex arr) + {-# INLINE iterArrayLinearST_ #-} instance Index ix => StrideLoad D ix e @@ -163,60 +171,13 @@ {-# INLINE toStreamIx #-} -- | Map an index aware function over an array -imap :: Source r ix e' => (ix -> e' -> e) -> Array r ix e' -> Array D ix e +-- +-- @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 = DArray (getComp arr) (size arr) (\ !ix -> f ix (unsafeIndex arr ix)) {-# INLINE imap #-} -instance (Index ix, Num e) => Num (Array D ix e) where - (+) = liftArray2Matching (+) - {-# INLINE (+) #-} - (-) = liftArray2Matching (-) - {-# INLINE (-) #-} - (*) = liftArray2Matching (*) - {-# 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 - (/) = liftArray2Matching (/) - {-# INLINE (/) #-} - fromRational = singleton . fromRational - {-# INLINE fromRational #-} - - -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 #-} - - instance Num e => FoldNumeric D e where unsafeDotProduct = defaultUnsafeDotProduct {-# INLINE unsafeDotProduct #-} @@ -239,7 +200,7 @@ -- | /O(1)/ Conversion from a source array to `D` representation. -delay :: Source r ix e => Array r ix e -> Array D ix e +delay :: (Index ix, Source r e) => Array r ix e -> Array D ix e delay arr = DArray (getComp arr) (size arr) (unsafeIndex arr) {-# INLINE [1] delay #-} @@ -247,24 +208,30 @@ "delay" [~1] forall (arr :: Array D ix e) . delay arr = arr #-} --- | /O(min (n1, n2))/ - Compute array equality by applying a comparing function to each element. +-- | 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 :: (Source r1 ix e1, Source r2 ix e2) => - (e1 -> e2 -> Bool) -> Array r1 ix e1 -> Array r2 ix e2 -> Bool +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 = - (size arr1 == size arr2) && - not (A.any not - (DArray (getComp arr1 <> getComp arr2) (size arr1) $ \ix -> - f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix))) + let sz1 = size arr1 + sz2 = size arr2 + in (sz1 == sz2 && + not + (A.any + not + (DArray (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. -- -- @since 0.5.7 -compareArrays :: (Source r1 ix e1, Source r2 ix e2) => +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) <> @@ -273,45 +240,35 @@ f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)) {-# INLINE compareArrays #-} - -liftArray2Matching - :: (Source r1 ix a, Source r2 ix b) +-- | 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 -liftArray2Matching f !arr1 !arr2 +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 $ DArray (getComp arr1 <> getComp arr2) sz1 (\ !ix -> f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)) - | otherwise = throw $ SizeMismatchException (size arr1) (size arr2) + | isZeroSz sz1 && isZeroSz sz2 = pure A.empty + | otherwise = throwM $ SizeMismatchException (size arr1) (size arr2) where sz1 = size arr1 sz2 = size arr2 -{-# INLINE liftArray2Matching #-} - - --- -- | 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 #-} - --- -- | 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 #-} - +{-# INLINE liftArray2M #-} 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 @@ -20,6 +20,7 @@ module Data.Massiv.Array.Delayed.Push ( DL(..) , Array(..) + , Loader , toLoadArray , makeLoadArrayS , makeLoadArray @@ -41,34 +42,36 @@ -- | Delayed load representation. Also known as Push array. 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) - , dlLoad :: forall m . Monad m - => Scheduler m () - -> Ix1 -- start loading at this linear index - -> (Ix1 -> e -> m ()) -- linear element writing action - -> (Ix1 -> Sz1 -> e -> m ()) -- linear region setting action - -> m () + , dlLoad :: Loader e } -instance Index ix => Construct DL ix e where +instance Strategy DL where + getComp = dlComp + {-# INLINE getComp #-} setComp c arr = arr {dlComp = c} {-# INLINE setComp #-} - makeArrayLinear comp sz f = DLArray comp sz load - where - load :: Monad m => - Scheduler m () -> Ix1 -> (Ix1 -> e -> m ()) -> (Ix1 -> Sz1 -> e -> m ()) -> m () - 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 #-} -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 @@ -91,8 +94,8 @@ DLArray {dlComp = foldMap getComp arrs, dlSize = SafeSz k, dlLoad = load} where !k = F.foldl' (+) 0 (unSz . size <$> arrs) - load :: Monad m => - Scheduler m () -> Ix1 -> (Ix1 -> e -> m ()) -> (Ix1 -> Sz1 -> e -> m ()) -> m () + 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 @@ -110,8 +113,8 @@ where !k1 = unSz sz1 !k2 = unSz sz2 - load :: Monad n => - Scheduler n () -> Ix1 -> (Ix1 -> e -> n ()) -> (Ix1 -> Sz1 -> e -> n ()) -> n () + 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 @@ -132,9 +135,9 @@ (!i2, !szl2) = unconsSz sz2 unless (szl1 == szl2) $ throwM $ SizeMismatchException sz1 sz2 pure $ - DLArray {dlComp = c1 <> c2, dlSize = consSz (i1 + i2) szl1, dlLoad = load} + DLArray {dlComp = c1 <> c2, dlSize = consSz (liftSz2 (+) i1 i2) szl1, dlLoad = load} where - load :: Monad n => Scheduler n () -> Ix1 -> (Ix1 -> e -> n ()) -> (Ix1 -> Sz1 -> e -> n ()) -> n () + 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 @@ -171,8 +174,8 @@ -> Array DL ix e makeLoadArrayS sz defVal writer = DLArray Seq sz load where - load :: Monad m => - Scheduler m () -> Ix1 -> (Ix1 -> e -> m ()) -> (Ix1 -> Sz1 -> e -> m ()) -> m () + 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 @@ -197,15 +200,15 @@ -- ^ 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 load where - load :: Monad m => - Scheduler m () -> Ix1 -> (Ix1 -> e -> m ()) -> (Ix1 -> Sz1 -> e -> m ()) -> m () + 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 @@ -232,7 +235,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 () -> Ix1 -> (Ix1 -> 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 @@ -244,8 +247,7 @@ -> Array DL ix e unsafeMakeLoadArray comp sz mDefVal writer = DLArray comp sz load where - load :: Monad m => - Scheduler m () -> Ix1 -> (Ix1 -> e -> m ()) -> (Ix1 -> Sz1 -> e -> m ()) -> m () + load :: Loader e load scheduler startAt uWrite uSet = do S.traverse_ (uSet startAt (toLinearSz sz)) mDefVal writer scheduler startAt uWrite @@ -261,12 +263,12 @@ => Comp -> Sz ix -> Maybe e - -> (forall m. Monad m => Scheduler m () -> (Ix1 -> e -> m ()) -> m ()) + -> (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 :: Monad m => - Scheduler m () -> Ix1 -> (Ix1 -> e -> m ()) -> (Ix1 -> Sz1 -> e -> m ()) -> m () + 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)) @@ -277,16 +279,16 @@ -- -- @since 0.3.0 toLoadArray :: - forall r ix e. Load r ix e + 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 :: Monad m => - Scheduler m () -> Ix1 -> (Ix1 -> e -> m ()) -> (Ix1 -> Sz1 -> e -> m ()) -> m () + load :: forall s. + Scheduler s () -> Ix1 -> (Ix1 -> e -> ST s ()) -> (Ix1 -> Sz1 -> e -> ST s ()) -> ST s () load scheduler !startAt dlWrite dlSet = - loadArrayWithSetM scheduler arr (dlWrite . (+ startAt)) (\offset -> dlSet (offset + startAt)) + iterArrayLinearWithSetST_ scheduler arr (dlWrite . (+ startAt)) (\offset -> dlSet (offset + startAt)) {-# INLINE load #-} {-# INLINE[1] toLoadArray #-} {-# RULES "toLoadArray/id" toLoadArray = id #-} @@ -295,28 +297,32 @@ -- -- @since 0.3.0 fromStrideLoad :: - forall r ix e. StrideLoad r 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 load where - !newsz = strideSize stride (size arr) - load :: Monad m => - Scheduler m () -> Ix1 -> (Ix1 -> e -> m ()) -> (Ix1 -> Sz1 -> e -> m ()) -> m () + !newsz = strideSize stride (outerSize arr) + load :: Loader e load scheduler !startAt dlWrite _ = - loadArrayWithStrideM scheduler stride newsz arr (\ !i -> dlWrite (i + startAt)) + 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 #-} - loadArrayWithSetM scheduler DLArray {dlLoad} = dlLoad scheduler 0 - {-# INLINE loadArrayWithSetM #-} + 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 Index ix => Functor (Array DL ix) where fmap f arr = arr {dlLoad = loadFunctor arr f} @@ -327,21 +333,20 @@ overwriteFunctor :: forall ix a b. Index ix => a -> Array DL ix b -> Array DL ix a overwriteFunctor e arr = arr {dlLoad = load} where - load :: Scheduler m () -> Ix1 -> (Ix1 -> a -> m ()) -> (Ix1 -> Sz1 -> a -> m ()) -> m () + load :: Loader a load _ !startAt _ dlSet = dlSet startAt (linearSize arr) e {-# INLINE load #-} {-# INLINE overwriteFunctor #-} loadFunctor :: - Monad m - => Array DL ix a + Array DL ix a -> (a -> b) - -> Scheduler m () + -> Scheduler s () -> Ix1 - -> (Ix1 -> b -> m ()) - -> (Ix1 -> Sz1 -> b -> m ()) - -> m () + -> (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 @@ -23,7 +23,7 @@ ) where import Control.Applicative -import Control.Monad (void) +import Control.Monad.ST import Data.Coerce import Data.Foldable import Data.Massiv.Array.Delayed.Pull @@ -31,7 +31,6 @@ import Data.Massiv.Core.Common import GHC.Exts import Prelude hiding (take, drop) -import Data.Vector.Fusion.Bundle.Size (upperBound) -- | 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 @@ -45,25 +44,45 @@ -- | /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 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 (Array DS Ix1 e) +fromStepsM :: Monad m => Steps m e -> m (Vector DS e) fromStepsM = fmap DSArray . S.transSteps {-# INLINE fromStepsM #-} +instance Shape DS Ix1 where + linearSizeHint = stepsSize . dsArray + {-# INLINE linearSizeHint #-} + + linearSize = SafeSz . unId . S.length . dsArray + {-# INLINE linearSize #-} + + outerSize = linearSize + {-# INLINE outerSize #-} + + isNull = S.unId . S.null . coerce + {-# INLINE isNull #-} + + +--TODO remove +instance Strategy DS where + getComp _ = Seq + setComp _ = id + + instance Functor (Array DS Ix1) where fmap f = coerce . S.map f . dsArray {-# INLINE fmap #-} @@ -114,8 +133,6 @@ 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 (<>) #-} @@ -129,9 +146,9 @@ 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 #-} @@ -147,7 +164,7 @@ -- | 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[1] toStreamArray #-} {-# RULES "toStreamArray/id" toStreamArray = id #-} @@ -167,45 +184,24 @@ {-# INLINE toStreamIxM #-} -instance Construct DS Ix1 e where - setComp _ arr = arr - {-# INLINE setComp #-} - - makeArrayLinear _ (Sz k) = fromSteps . S.generate k - {-# INLINE makeArrayLinear #-} - - -instance Extract DS Ix1 e where - unsafeExtract sIx newSz = fromSteps . S.slice sIx (unSz newSz) . dsArray - {-# INLINE unsafeExtract #-} - -- | /O(n)/ - `size` implementation. instance Load DS Ix1 e where - size = coerce . S.unId . S.length . coerce - {-# INLINE size #-} - maxSize = coerce . upperBound . stepsSize . dsArray - {-# INLINE maxSize #-} - - isEmpty = S.unId . S.null . coerce - {-# INLINE isEmpty #-} - - 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 #-} + {-# INLINE unsafeLoadIntoST #-} - unsafeLoadIntoM marr arr = liftIO $ unsafeLoadIntoS marr arr - {-# INLINE unsafeLoadIntoM #-} + unsafeLoadIntoIO marr arr = stToIO $ unsafeLoadIntoST marr arr + {-# INLINE unsafeLoadIntoIO #-} -- cons :: e -> Array DS Ix1 e -> Array DS Ix1 e @@ -223,13 +219,13 @@ -- 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 #-} +-- {-# 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 @@ -62,15 +62,11 @@ showsPrec = showsArrayPrec (computeAs B) showList = showArrayList - -instance Index ix => Construct DW ix e where - +instance Strategy DW where setComp c arr = arr { dwArray = (dwArray arr) { dComp = c } } {-# INLINE setComp #-} - - makeArray c sz f = DWArray (makeArray c sz f) Nothing - {-# INLINE makeArray #-} - + getComp = dComp . dwArray + {-# INLINE getComp #-} instance Functor (Array DW ix) where @@ -122,7 +118,7 @@ -- -- @since 0.1.3 makeWindowedArray - :: Source r ix e + :: (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 @@ -138,7 +134,7 @@ -- -- @since 0.3.0 insertWindow - :: Source D ix e + :: 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 @@ -209,13 +205,18 @@ return (\from to -> with $ iterM_ from to 1 (<) $ \ !i -> uWrite i (indexW i), it, wEnd) {-# INLINE loadWithIx1 #-} +instance Index ix => Shape DW ix where + maxLinearSize = Just . linearSize + {-# INLINE maxLinearSize #-} + linearSize = SafeSz . totalElem . dSize . dwArray + {-# INLINE linearSize #-} + outerSize = dSize . dwArray + {-# 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 (makeArray 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 -> @@ -224,10 +225,10 @@ 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 + 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 loopM_ 0 (< numWorkers scheduler) (+ 1) $ \ !wid -> @@ -236,7 +237,7 @@ when (slackWidth > 0) $ let !itSlack = numWorkers scheduler * chunkWidth + wStart in loadWindow (itSlack, itSlack + slackWidth) - {-# INLINE loadArrayWithStrideM #-} + {-# INLINE iterArrayLinearWithStrideST_ #-} loadArrayWithIx1 :: (Monad m) @@ -276,7 +277,7 @@ let ib :. jb = (wm + it) :. (wn + jt) !blockHeight = maybe 1 (min 7 . max 1) mUnrollHeight stride = oneStride - !sz = strideSize stride $ size arr + !sz = strideSize stride $ outerSize arr writeB !ix = uWrite (toLinearIndex sz ix) (indexB ix) {-# INLINE writeB #-} writeW !ix = uWrite (toLinearIndex sz ix) (indexW ix) @@ -335,42 +336,38 @@ instance Load DW Ix2 e where - size = dSize . dwArray - {-# INLINE size #-} - getComp = dComp . dwArray - {-# INLINE getComp #-} - loadArrayM scheduler arr uWrite = + makeArray c sz f = DWArray (makeArray c sz f) Nothing + {-# INLINE makeArray #-} + iterArrayLinearST_ scheduler arr uWrite = loadWithIx2 (scheduleWork scheduler) arr uWrite >>= uncurry (loadWindowIx2 (numWorkers scheduler)) - {-# INLINE loadArrayM #-} + {-# INLINE iterArrayLinearST_ #-} instance StrideLoad DW Ix2 e where - loadArrayWithStrideM scheduler stride sz arr uWrite = + iterArrayLinearWithStrideST_ scheduler stride sz arr uWrite = loadArrayWithIx2 (scheduleWork scheduler) arr stride sz uWrite >>= uncurry (loadWindowIx2 (numWorkers scheduler)) - {-# INLINE loadArrayWithStrideM #-} + {-# 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 (makeArray 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 () + (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 @@ -393,7 +390,7 @@ DWArray {dwArray = DArray Seq lowerSourceSize (indexBorder . consDim i), dwWindow = ($ i) <$> mw} loadLower mw !i = - loadArrayWithStrideM + iterArrayLinearWithStrideST_ scheduler (Stride lowerStrideIx) lowerSize @@ -412,11 +409,11 @@ loadWithIxN :: - (Index ix, Monad m, Load DW (Lower ix) e) - => Scheduler m () + (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 @@ -436,7 +433,7 @@ 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)) + iterArrayLinearST_ 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)) @@ -479,79 +476,3 @@ -- 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 @@ -17,8 +17,6 @@ module Data.Massiv.Array.Manifest ( -- * Manifest Manifest - , toManifest - , M -- * Boxed , B(..) , BL(..) @@ -123,15 +121,16 @@ import Data.Massiv.Core.Common 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 + Load r Ix1 Word8 + => 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) + -> Vector r Word8 +fromByteString comp bs = makeArrayLinear comp (SafeSz (S.length bs)) (SU.unsafeIndex bs) {-# INLINE fromByteString #-} -- | /O(n)/ - Convert any source array into a strict `ByteString`. In case when the source array is @@ -154,28 +153,28 @@ -- | /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 :: Array S ix Word8 -> Builder +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 #-} @@ -193,7 +192,7 @@ -- after it was applyied to all elements of the array. -- -- @since 0.5.5 -findIndex :: Manifest r ix e => (e -> Bool) -> Array r ix e -> Maybe ix +findIndex :: (Index ix, Manifest r e) => (e -> Bool) -> Array r ix e -> Maybe ix findIndex f arr = go 0 where !sz = size arr @@ -211,7 +210,7 @@ -- programs. -- -- @since 0.5.9 -mallocCompute :: forall r ix e. (Source r ix e, Storable e) => Array r ix e -> IO (Array S ix e) +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 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 @@ -24,6 +24,7 @@ , N , pattern N , Array(..) + , MArray(..) , wrapLazyArray , unwrapLazyArray , unwrapNormalForm @@ -61,7 +62,7 @@ import qualified Data.Foldable as F (Foldable(..)) 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 (computeAs) import Data.Massiv.Array.Manifest.List as L import Data.Massiv.Array.Mutable import Data.Massiv.Array.Ops.Fold @@ -80,16 +81,6 @@ #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 -- ---------------- @@ -103,6 +94,8 @@ , blOffset :: {-# UNPACK #-} !Int , blData :: {-# 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 BL ix e) where showsPrec = showsArrayPrec id @@ -129,71 +122,40 @@ compare = compareArrays compare {-# INLINE compare #-} -instance Index ix => Construct BL ix e where +instance Strategy BL where setComp c arr = arr { blComp = c } {-# INLINE setComp #-} + getComp = blComp + {-# INLINE getComp #-} - makeArrayLinear !comp !sz f = unsafePerformIO $ generateArrayLinear comp sz (pure . f) - {-# INLINE makeArrayLinear #-} - replicate comp sz e = runST (newMArray sz e >>= unsafeFreeze comp) - {-# INLINE replicate #-} - -instance Index ix => Source BL ix e where +instance Source BL e where unsafeLinearIndex (BLArray _ _sz o a) i = INDEX_CHECK("(Source BL ix e).unsafeLinearIndex", - SafeSz . sizeofArray, A.indexArray) a (i + o) + SafeSz . A.sizeofArray, A.indexArray) a (i + o) {-# INLINE unsafeLinearIndex #-} - unsafeLinearSlice i k (BLArray c _ o a) = BLArray c k (o + i) a - {-# INLINE unsafeLinearSlice #-} - - -instance Index ix => Resize BL ix where - unsafeResize !sz !arr = arr { blSize = sz } - {-# INLINE unsafeResize #-} - -instance Index ix => Extract BL ix e where - unsafeExtract !sIx !newSz !arr = unsafeExtract sIx newSz (toManifest arr) - {-# INLINE unsafeExtract #-} - - -instance ( Index ix - , Index (Lower ix) - , Elt M ix e ~ Array M (Lower ix) e - , Elt BL ix e ~ Array M (Lower ix) e - ) => - OuterSlice BL ix e where - unsafeOuterSlice arr = unsafeOuterSlice (toManifest arr) + unsafeOuterSlice (BLArray c _ o a) szL i = BLArray c szL (i * totalElem szL + o) a {-# INLINE unsafeOuterSlice #-} -instance ( Index ix - , Index (Lower ix) - , Elt M ix e ~ Array M (Lower ix) e - , Elt BL ix e ~ Array M (Lower ix) e - ) => - InnerSlice BL ix e where - unsafeInnerSlice arr = unsafeInnerSlice (toManifest arr) - {-# INLINE unsafeInnerSlice #-} - -instance {-# OVERLAPPING #-} Slice BL Ix1 e where - unsafeSlice arr i _ _ = pure (unsafeLinearIndex arr i) - {-# INLINE unsafeSlice #-} - + unsafeLinearSlice i k (BLArray c _ o a) = BLArray c k (o + i) a + {-# INLINE unsafeLinearSlice #-} -instance Index ix => Manifest BL ix e where +instance Manifest BL e where unsafeLinearIndexM (BLArray _ _sz o a) i = INDEX_CHECK("(Manifest BL ix e).unsafeLinearIndexM", - SafeSz . sizeofArray, A.indexArray) a (i + o) + SafeSz . A.sizeofArray, A.indexArray) a (i + o) {-# INLINE unsafeLinearIndexM #-} + sizeOfMArray (MBLArray sz _ _) = sz + {-# INLINE sizeOfMArray #-} -instance Index ix => Mutable BL ix e where - data MArray s BL ix e = MBLArray !(Sz ix) {-# UNPACK #-} !Int {-# UNPACK #-} !(A.MutableArray s e) + unsafeResizeMArray sz (MBLArray _ off marr) = MBLArray sz off marr + {-# INLINE unsafeResizeMArray #-} - msize (MBLArray sz _ _) = sz - {-# INLINE msize #-} + 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 #-} @@ -211,24 +173,37 @@ {-# INLINE newMArray #-} unsafeLinearRead (MBLArray _ o ma) i = - INDEX_CHECK("(Mutable BL ix e).unsafeLinearRead", - SafeSz . sizeofMutableArray, A.readArray) ma (i + o) + INDEX_CHECK("(Manifest BL ix e).unsafeLinearRead", + SafeSz . A.sizeofMutableArray, A.readArray) ma (i + o) {-# INLINE unsafeLinearRead #-} unsafeLinearWrite (MBLArray _sz o ma) i e = e `seq` - INDEX_CHECK("(Mutable BL ix e).unsafeLinearWrite", - SafeSz . sizeofMutableArray, A.writeArray) ma (i + o) e + INDEX_CHECK("(Manifest BL ix e).unsafeLinearWrite", + SafeSz . A.sizeofMutableArray, A.writeArray) ma (i + o) e {-# INLINE unsafeLinearWrite #-} -instance Index ix => Load BL ix e where - type R BL = M +instance Size BL where size = blSize {-# INLINE size #-} - getComp = blComp - {-# INLINE getComp #-} - loadArrayM !scheduler !arr = splitLinearlyWith_ scheduler (elemsCount arr) (unsafeLinearIndex arr) - {-# INLINE loadArrayM #-} + unsafeResize !sz !arr = arr { blSize = sz } + {-# INLINE unsafeResize #-} + +instance Index ix => Shape BL ix where + maxLinearSize = Just . SafeSz . elemsCount + {-# INLINE maxLinearSize #-} + +instance Index ix => Load BL ix e where + 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 => StrideLoad BL ix e instance Index ix => Stream BL ix e where @@ -270,12 +245,7 @@ traverse = traverseA {-# INLINE traverse #-} -instance ( IsList (Array L ix e) - , Nested LN ix e - , Nested L ix e - , Ragged L ix e - ) => - IsList (Array BL ix e) where +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 #-} @@ -308,6 +278,8 @@ 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 @@ -324,67 +296,48 @@ compare = compareArrays compare {-# INLINE compare #-} -instance Index ix => Construct B ix e where - setComp c = coerce (\arr -> arr { blComp = c }) - {-# INLINE setComp #-} - makeArrayLinear !comp !sz f = unsafePerformIO $ generateArrayLinear comp sz (pure . f) - {-# INLINE makeArrayLinear #-} - - replicate comp sz e = runST (newMArray sz e >>= unsafeFreeze comp) - {-# INLINE replicate #-} - -instance Index ix => Source B ix e where - unsafeLinearIndex arr = unsafeLinearIndex (coerce arr :: Array BL ix e) +instance Source B e where + unsafeLinearIndex arr = unsafeLinearIndex (toLazyArray arr) {-# INLINE unsafeLinearIndex #-} - unsafeLinearSlice i k arr = coerce (unsafeLinearSlice i k (coerce arr :: Array BL ix e)) + unsafeLinearSlice i k arr = coerce (unsafeLinearSlice i k (toLazyArray arr)) {-# INLINE unsafeLinearSlice #-} - -instance Index ix => Resize B ix where - unsafeResize sz = coerce (\arr -> arr { blSize = sz }) - {-# INLINE unsafeResize #-} + unsafeOuterSlice arr i = coerce (unsafeOuterSlice (toLazyArray arr) i) + {-# INLINE unsafeOuterSlice #-} -instance Index ix => Extract B ix e where - unsafeExtract !sIx !newSz !arr = unsafeExtract sIx newSz (toManifest arr) - {-# INLINE unsafeExtract #-} +instance Strategy B where + getComp = blComp . coerce + {-# INLINE getComp #-} + setComp c arr = coerceBoxedArray (coerce arr) { blComp = c } + {-# INLINE setComp #-} -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 #-} - -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 #-} +instance Index ix => Shape B ix where + maxLinearSize = Just . SafeSz . elemsCount + {-# INLINE maxLinearSize #-} -instance {-# OVERLAPPING #-} Slice B Ix1 e where - unsafeSlice arr i _ _ = pure (unsafeLinearIndex arr i) - {-# INLINE unsafeSlice #-} +instance Size B where + size = blSize . coerce + {-# INLINE size #-} + unsafeResize sz = coerce (\arr -> arr { blSize = sz }) + {-# INLINE unsafeResize #-} -instance Index ix => Manifest B ix e where +instance Manifest B e where unsafeLinearIndexM = coerce unsafeLinearIndexM {-# INLINE unsafeLinearIndexM #-} + sizeOfMArray = sizeOfMArray . coerce + {-# INLINE sizeOfMArray #-} -instance Index ix => Mutable B ix e where - newtype MArray s B ix e = MBArray (MArray s BL ix e) + unsafeResizeMArray sz = MBArray . unsafeResizeMArray sz . coerce + {-# INLINE unsafeResizeMArray #-} - msize = msize . coerce - {-# INLINE msize #-} + unsafeLinearSliceMArray i k = MBArray . unsafeLinearSliceMArray i k . coerce + {-# INLINE unsafeLinearSliceMArray #-} unsafeThaw arr = MBArray <$> unsafeThaw (coerce arr) {-# INLINE unsafeThaw #-} @@ -408,14 +361,15 @@ {-# INLINE unsafeLinearWrite #-} instance Index ix => Load B ix e where - type R B = M - size = blSize . coerce - {-# INLINE size #-} - getComp = blComp . coerce - {-# INLINE getComp #-} - loadArrayM scheduler = coerce (loadArrayM scheduler) - {-# INLINE loadArrayM #-} + 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 @@ -457,12 +411,7 @@ 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 #-} @@ -493,16 +442,20 @@ data BN = BN deriving Show -- | Type and pattern `N` have been added for backwards compatibility and will be replaced --- in the future in favor of `BN` +-- in the future in favor of `BN`. +-- +-- /Deprecated/ - since 1.0.0 type N = BN pattern N :: N pattern N = BN {-# COMPLETE N #-} +{-# DEPRECATED N "In favor of more consistently named `BN`" #-} -newtype instance Array N ix e = BNArray { bArray :: Array BL ix e } +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 bArray + showsPrec = showsArrayPrec coerce showList = showArrayList -- | /O(1)/ - `BN` is already in normal form @@ -518,67 +471,44 @@ compare = compareArrays compare {-# INLINE compare #-} - -instance (Index ix, NFData e) => Construct BN ix e where - setComp c (BNArray arr) = BNArray (arr {blComp = c}) +instance Strategy N where + setComp c = coerce (setComp c) {-# INLINE setComp #-} - makeArrayLinear !comp !sz f = unsafePerformIO $ generateArrayLinear comp sz (pure . f) - {-# INLINE makeArrayLinear #-} - replicate comp sz e = runST (newMArray sz e >>= unsafeFreeze comp) - {-# INLINE replicate #-} + getComp = blComp . coerce + {-# INLINE getComp #-} -instance (Index ix, NFData e) => Source BN ix e where +instance NFData e => Source BN e where unsafeLinearIndex (BNArray arr) = unsafeLinearIndex arr {-# INLINE unsafeLinearIndex #-} - unsafeLinearSlice i k (BNArray a) = BNArray $ unsafeLinearSlice i k a + unsafeLinearSlice i k (BNArray a) = coerce (unsafeLinearSlice i k a) {-# INLINE unsafeLinearSlice #-} - - -instance Index ix => Resize BN ix where - unsafeResize !sz = BNArray . unsafeResize sz . bArray - {-# INLINE unsafeResize #-} - -instance (Index ix, NFData e) => Extract BN 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 BN ix e ~ Array M (Lower ix) e - ) => - OuterSlice BN ix e where - unsafeOuterSlice = unsafeOuterSlice . toManifest + 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 BN ix e ~ Array M (Lower ix) e - ) => - InnerSlice BN ix e where - unsafeInnerSlice = unsafeInnerSlice . toManifest - {-# INLINE unsafeInnerSlice #-} -instance {-# OVERLAPPING #-} NFData e => Slice BN 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 BN ix e where + unsafeResize !sz = coerce . unsafeResize sz . coerce + {-# INLINE unsafeResize #-} +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 BN ix e where - newtype MArray s BN ix e = MBNArray (MArray s BL ix e) + unsafeResizeMArray sz = coerce . unsafeResizeMArray sz . coerce + {-# INLINE unsafeResizeMArray #-} - msize = msize . coerce - {-# INLINE msize #-} + unsafeLinearSliceMArray i k = MBNArray . unsafeLinearSliceMArray i k . coerce + {-# INLINE unsafeLinearSliceMArray #-} unsafeThaw arr = MBNArray <$> unsafeThaw (coerce arr) {-# INLINE unsafeThaw #-} @@ -602,13 +532,13 @@ {-# INLINE unsafeLinearWrite #-} instance (Index ix, NFData e) => Load BN ix e where - type R BN = M - size = blSize . coerce - {-# INLINE size #-} - getComp = blComp . coerce - {-# INLINE getComp #-} - loadArrayM !scheduler !arr = splitLinearlyWith_ scheduler (elemsCount arr) (unsafeLinearIndex arr) - {-# INLINE loadArrayM #-} + 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 BN ix e @@ -619,13 +549,7 @@ {-# INLINE toStreamIx #-} -instance ( NFData e - , IsList (Array L ix e) - , Nested LN ix e - , Nested L ix e - , Ragged L ix e - ) => - IsList (Array BN ix e) where +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 #-} @@ -688,7 +612,7 @@ -- -- @since 0.6.0 wrapLazyArray :: A.Array e -> Vector BL e -wrapLazyArray a = BLArray Seq (SafeSz (sizeofArray a)) 0 a +wrapLazyArray a = BLArray Seq (SafeSz (A.sizeofArray a)) 0 a {-# INLINE wrapLazyArray #-} @@ -795,7 +719,7 @@ -> A.MutableArray (PrimState m) e -> m (MArray (PrimState m) BL Ix1 e) fromMutableArraySeq with ma = do - let !sz = sizeofMutableArray ma + let !sz = A.sizeofMutableArray ma loopM_ 0 (< sz) (+ 1) (A.readArray ma >=> (`with` return ())) return $! MBLArray (SafeSz sz) 0 ma {-# INLINE fromMutableArraySeq #-} 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,8 +1,8 @@ {-# LANGUAGE BangPatterns #-} -{-# LANGUAGE CPP #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MagicHash #-} +{-# LANGUAGE MonoLocalBinds #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} @@ -17,10 +17,9 @@ -- Portability : non-portable -- module Data.Massiv.Array.Manifest.Internal - ( M - , Manifest(..) + ( Manifest(..) , Array(..) - , toManifest + , flattenMArray , compute , computeS , computeP @@ -38,190 +37,33 @@ , gcastArr , fromRaggedArrayM , fromRaggedArray' - , sizeofArray - , sizeofMutableArray + , unsafeLoadIntoS + , unsafeLoadIntoM , iterateUntil , iterateUntilM ) where import Control.Exception (try) +import Control.DeepSeq import Control.Monad.ST +import Control.Monad.Primitive 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 as A import Data.Massiv.Array.Mutable.Internal (unsafeCreateArray_) -import Data.Massiv.Vector.Stream as S (steps, isteps) import Data.Massiv.Core.Common import Data.Massiv.Core.List -import Data.Massiv.Core.Operations 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 - (==) = eqArrays (==) - {-# INLINE (==) #-} - -instance (Ord e, Index ix) => Ord (Array M ix e) where - compare = compareArrays 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 #-} - elem e = A.any (e ==) - {-# INLINE elem #-} - toList arr = build (\ c n -> foldrFB c n arr) - {-# INLINE toList #-} - - -instance Index ix => Source M ix e where - unsafeLinearIndex = mLinearIndex - {-# INLINE unsafeLinearIndex #-} - unsafeLinearSlice ix sz arr = unsafeExtract ix sz (unsafeResize sz arr) - {-# INLINE unsafeLinearSlice #-} - - -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 #-} - toStreamIx = S.isteps - {-# INLINE toStreamIx #-} - - -instance Num e => FoldNumeric M e where - unsafeDotProduct = defaultUnsafeDotProduct - {-# INLINE unsafeDotProduct #-} - powerSumArray = defaultPowerSumArray - {-# INLINE powerSumArray #-} - foldArray = defaultFoldArray - {-# INLINE foldArray #-} - -- | 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 #-} @@ -229,7 +71,7 @@ -- 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 #-} @@ -240,7 +82,7 @@ -- -- @since 0.5.4 computeP :: - forall r ix e r'. (Mutable r ix e, Construct r' ix e, Load r' ix e) + 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) @@ -253,7 +95,7 @@ -- -- @since 0.4.5 computeIO :: - forall r ix e r' m. (Mutable r ix e, Load r' ix e, MonadIO m) + 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)) @@ -264,7 +106,7 @@ -- -- @since 0.4.5 computePrimM :: - forall r ix e r' m. (Mutable r ix e, Load r' ix e, PrimMonad m) + 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) @@ -280,7 +122,7 @@ -- 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 #-} @@ -300,7 +142,7 @@ -- [ 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 #-} @@ -309,16 +151,16 @@ -- 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) +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 arr) (\Refl -> arr) (eqT :: Maybe (r' :~: r)) +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 #-} @@ -333,7 +175,7 @@ -- 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) +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 #-} @@ -341,7 +183,7 @@ -- | 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) +convertAs :: (Manifest r e, Load r' ix e) => r -> Array r' ix e -> Array r ix e convertAs _ = convert {-# INLINE convertAs #-} @@ -351,7 +193,7 @@ -- proxy argument. -- -- @since 0.1.1 -convertProxy :: (Mutable r ix e, Load 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 #-} @@ -362,48 +204,48 @@ -- -- @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) + 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 (withMassivScheduler_ (getComp arr) $ \scheduler -> - loadRagged (scheduleWork scheduler) (unsafeLinearWrite marr) 0 (totalElem sz) sz arr) + 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) + 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 +-- /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) + 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) + let !sz = strideSize stride (outerSize arr) unsafeCreateArray_ (getComp arr) sz $ \scheduler marr -> - loadArrayWithStrideM scheduler stride sz arr (unsafeLinearWrite marr) + stToIO $ iterArrayLinearWithStrideST_ scheduler stride sz arr (unsafeLinearWrite marr) {-# INLINE computeWithStride #-} @@ -411,12 +253,36 @@ -- -- @since 0.3.0 computeWithStrideAs :: - (Mutable r ix e, StrideLoad r' ix e) => r -> Stride ix -> Array r' ix e -> Array r ix e + (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 -- allocated, regardless of the number of iterations. If the size does change from one @@ -426,8 +292,8 @@ -- ====__Example__ -- -- >>> import Data.Massiv.Array --- >>> a = computeAs P $ makeLoadArrayS (Sz2 8 8) (0 :: Int) $ \ w -> () <$ w (0 :. 0) 1 --- >>> 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 ] @@ -438,9 +304,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 ] @@ -454,7 +320,7 @@ -- -- @since 0.3.6 iterateUntil :: - (Load r' ix e, Mutable r ix e) + (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. @@ -463,85 +329,59 @@ -- 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 comp marr' -> convergence n a <$> unsafeFreeze comp marr') - 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 +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) + (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 -> 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 -> Comp -> MArray (PrimState m) r ix e -> m Bool) - -> (Int -> Array r ix e -> Array r' ix e) + (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) - !comp = getComp loadArr - marr' <- - if k == mk - then pure marr - else if k < mk - then unsafeLinearShrink marr sz - else unsafeLinearGrow marr sz - computeInto marr' loadArr - shouldStop <- convergence n arr comp marr' - arr' <- unsafeFreeze comp 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 @@ -13,8 +13,7 @@ -- Portability : non-portable -- module Data.Massiv.Array.Manifest.List - ( - -- ** List + ( -- ** List fromList , fromListsM , fromLists' @@ -30,16 +29,16 @@ 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 (build, IsList(..)) -- | Convert a flat list into a vector -- -- @since 0.1.0 fromList :: - forall r e. Mutable r Ix1 e + forall r e. Manifest r e => Comp -- ^ Computation startegy to use -> [e] -- ^ Flat list - -> Array r Ix1 e + -> Vector r e fromList = fromLists' {-# INLINE fromList #-} @@ -75,27 +74,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 #-}@ -- 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 @@ -116,26 +126,17 @@ -- , [ 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__ @@ -145,8 +146,8 @@ -- [(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 #-} @@ -171,10 +172,11 @@ -- [[[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 #-} @@ -191,7 +193,7 @@ -- [[(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 #-} @@ -200,7 +202,8 @@ -- 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,12 +212,17 @@ -- -- @since 0.1.0 toLists4 :: - ( Index (Lower (Lower (Lower ix))) - , Index (Lower (Lower ix)) + ( 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 @@ -20,6 +20,7 @@ module Data.Massiv.Array.Manifest.Primitive ( P(..) , Array(..) + , MArray(..) , Prim , toPrimitiveVector , toPrimitiveMVector @@ -83,6 +84,9 @@ , 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 @@ -103,92 +107,52 @@ compare = compareArrays compare {-# INLINE compare #-} -instance (Prim e, Index ix) => Construct P ix e where +instance Strategy P where + getComp = pComp + {-# INLINE getComp #-} setComp c arr = arr { pComp = c } {-# INLINE setComp #-} - makeArrayLinear !comp !sz f = unsafePerformIO $ generateArrayLinear comp sz (pure . f) - {-# INLINE makeArrayLinear #-} - replicate comp !sz !e = runST (newMArray sz e >>= unsafeFreeze comp) - {-# INLINE replicate #-} +instance Index ix => Shape P ix where + maxLinearSize = Just . SafeSz . elemsCount + {-# INLINE maxLinearSize #-} -instance (Prim e, Index ix) => Source P ix e where +instance Size P where + size = pSize + {-# INLINE size #-} + unsafeResize !sz !arr = arr { pSize = sz } + {-# INLINE unsafeResize #-} + +instance Prim e => Source P e where unsafeLinearIndex _arr@(PArray _ _ o a) i = INDEX_CHECK("(Source P ix e).unsafeLinearIndex", SafeSz . elemsBA _arr, indexByteArray) a (i + o) {-# INLINE unsafeLinearIndex #-} - unsafeLinearSlice i k (PArray c _ o a) = PArray c k (i + o) a - {-# INLINE unsafeLinearSlice #-} - - -instance Index ix => Resize P ix where - 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 + unsafeOuterSlice (PArray c _ o a) szL i = + PArray c szL (i * totalElem szL + o) a {-# INLINE unsafeOuterSlice #-} -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) - {-# INLINE unsafeOuterSlice #-} + unsafeLinearSlice i k (PArray c _ o a) = PArray c k (i + o) a + {-# INLINE unsafeLinearSlice #-} -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 +instance Prim e => Manifest P e where unsafeLinearIndexM _pa@(PArray _ _sz o a) i = INDEX_CHECK("(Manifest P ix e).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 #-} !Int {-# 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 o a) = MPArray sz o <$> unsafeThawByteArray a {-# INLINE unsafeThaw #-} @@ -209,12 +173,12 @@ {-# INLINE initialize #-} unsafeLinearRead _mpa@(MPArray _sz o ma) i = - INDEX_CHECK("(Mutable P ix e).unsafeLinearRead", + INDEX_CHECK("(Manifest P ix e).unsafeLinearRead", const (Sz (totalElem _sz)), readByteArray) ma (i + o) {-# INLINE unsafeLinearRead #-} unsafeLinearWrite _mpa@(MPArray _sz o ma) i = - INDEX_CHECK("(Mutable P ix e).unsafeLinearWrite", + INDEX_CHECK("(Manifest P ix e).unsafeLinearWrite", const (Sz (totalElem _sz)), writeByteArray) ma (i + o) {-# INLINE unsafeLinearWrite #-} @@ -237,19 +201,20 @@ {-# INLINE unsafeLinearShrink #-} unsafeLinearGrow (MPArray _ o ma) sz = - MPArray sz o <$> resizeMutableByteArrayCompat ma ((o + totalElem sz) * sizeOf (undefined :: e)) + 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 = + 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 @@ -278,13 +243,7 @@ instance (Prim e, Floating e) => NumericFloat P e -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, 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 #-} @@ -647,27 +606,3 @@ mba (o + toLinearIndex sz ix) {-# INLINE unsafeAtomicXorIntArray #-} - - -#if !MIN_VERSION_primitive(0,7,1) -shrinkMutableByteArray :: forall m. (PrimMonad m) - => MutableByteArray (PrimState m) - -> Int -- ^ new size - -> m () -shrinkMutableByteArray (MutableByteArray arr#) (I# n#) - = primitive_ (shrinkMutableByteArray# arr# n#) -{-# INLINE shrinkMutableByteArray #-} -#endif - -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 @@ -18,6 +18,7 @@ module Data.Massiv.Array.Manifest.Storable ( S (..) , Array(..) + , MArray(..) , Storable , toStorableVector , toStorableMVector @@ -36,51 +37,57 @@ import Control.DeepSeq (NFData(..), deepseq) import Control.Exception +import Control.Monad import Control.Monad.IO.Unlift -import Control.Monad.Primitive (unsafePrimToPrim) +import Control.Monad.Primitive import Data.Massiv.Array.Delayed.Pull (compareArrays, eqArrays) import Data.Massiv.Array.Manifest.Internal import Data.Massiv.Array.Manifest.List as A -import Data.Massiv.Array.Manifest.Primitive (shrinkMutableByteArray) 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 Data.Primitive.ByteArray (MutableByteArray(..)) -import qualified Data.Vector.Generic.Mutable as VGM -import qualified Data.Vector.Storable as VS -import qualified Data.Vector.Storable.Mutable as MVS -import Foreign.ForeignPtr (newForeignPtr, withForeignPtr) +import Data.Primitive.Ptr (setPtr) +import Data.Primitive.ByteArray +import Foreign.ForeignPtr import Foreign.Marshal.Alloc import Foreign.Marshal.Array (advancePtr, copyArray) import Foreign.Ptr import Foreign.Storable import GHC.Exts as GHC (IsList(..)) -import GHC.ForeignPtr (ForeignPtr(..), ForeignPtrContents(..)) +import GHC.ForeignPtr import Prelude hiding (mapM) import System.IO.Unsafe (unsafePerformIO) +import Data.Word +import Unsafe.Coerce +import qualified Data.Vector.Generic.Mutable as MVG +import qualified Data.Vector.Storable as VS +import qualified Data.Vector.Storable.Mutable as MVS + #include "massiv.h" -- | Representation for `Storable` elements 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) } +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 NFData ix => NFData (MArray s S ix e) where - rnf (MSArray sz mv) = sz `deepseq` mv `deepseq` () + rnf (MSArray sz _mv) = sz `deepseq` () {-# INLINE rnf #-} instance (Storable e, Eq e, Index ix) => Eq (Array S ix e) where @@ -91,98 +98,95 @@ compare = compareArrays compare {-# INLINE compare #-} -instance (Storable e, Index ix) => Construct S ix e where +instance Strategy S where + getComp = sComp + {-# INLINE getComp #-} setComp c arr = arr { sComp = c } {-# INLINE setComp #-} - makeArrayLinear !comp !sz f = unsafePerformIO $ generateArrayLinear comp sz (pure . f) - {-# INLINE makeArrayLinear #-} - - replicate comp !sz !e = runST (newMArray sz e >>= unsafeFreeze comp) - {-# INLINE replicate #-} +advanceForeignPtr :: forall e . Storable e => ForeignPtr e -> Int -> ForeignPtr e +advanceForeignPtr fp i = plusForeignPtr 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, Index ix) => Source S ix e where - unsafeLinearIndex (SArray _ _ v) = - INDEX_CHECK("(Source S ix e).unsafeLinearIndex", Sz . VS.length, VS.unsafeIndex) v +instance Storable e => Source S e where + unsafeLinearIndex (SArray _ _sz fp) = + INDEX_CHECK("(Source S ix e).unsafeLinearIndex", const (toLinearSz _sz), indexForeignPtr) fp {-# INLINE unsafeLinearIndex #-} - unsafeLinearSlice i k (SArray c _ v) = SArray c k $ VS.unsafeSlice i (unSz k) v - {-# INLINE unsafeLinearSlice #-} -instance Index ix => Resize S ix where - unsafeResize !sz !arr = arr { sSize = sz } - {-# INLINE unsafeResize #-} - -instance (Storable e, Index ix) => Extract S ix e where - unsafeExtract !sIx !newSz !arr = unsafeExtract sIx newSz (toManifest arr) - {-# INLINE unsafeExtract #-} - - - -instance ( 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 ( 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 #-} + unsafeLinearSlice i k (SArray c _ fp) = + SArray c k $ advanceForeignPtr fp i + {-# INLINE unsafeLinearSlice #-} -instance {-# OVERLAPPING #-} Storable e => Slice S Ix1 e where - unsafeSlice arr i _ _ = pure (unsafeLinearIndex arr i) - {-# INLINE unsafeSlice #-} +instance Index ix => Shape S ix where + maxLinearSize = Just . SafeSz . elemsCount + {-# INLINE maxLinearSize #-} +instance Size S where + size = sSize + {-# INLINE size #-} + unsafeResize !sz !arr = arr { sSize = sz } + {-# INLINE unsafeResize #-} -instance (Index ix, Storable e) => Manifest S ix e where - 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) = + INDEX_CHECK("(Source S ix e).unsafeLinearIndex", const (toLinearSz _sz), indexForeignPtr) fp {-# INLINE unsafeLinearIndexM #-} + sizeOfMArray (MSArray sz _) = sz + {-# INLINE sizeOfMArray #-} -instance (Index ix, 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 $ + INDEX_CHECK("(Manifest S ix e).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 $ + INDEX_CHECK("(Manifest S ix e).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 = + MVG.basicSet (MVS.unsafeFromForeignPtr0 (advanceForeignPtr fp i) (unSz k)) {-# 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 @@ -196,33 +200,30 @@ 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 + makeArrayLinear !comp !sz f = unsafePerformIO $ generateArrayLinear comp sz (pure . f) + {-# INLINE makeArrayLinear #-} + replicate comp !sz !e = runST (newMArray sz e >>= unsafeFreeze comp) + {-# INLINE replicate #-} -instance (Index ix, 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 #-} + iterArrayLinearST_ !scheduler !arr = + splitLinearlyWith_ scheduler (elemsCount arr) (unsafeLinearIndex arr) + {-# INLINE iterArrayLinearST_ #-} instance (Index ix, Storable e) => StrideLoad S ix e @@ -250,13 +251,7 @@ instance (Storable e, Floating e) => NumericFloat S e -instance ( 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, 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 #-} @@ -267,69 +262,75 @@ -- referential transparency. -- -- @since 0.1.3 -unsafeWithPtr :: (MonadUnliftIO m, 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, 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 = + unsafeCoerce $ -- this hack is needed to workaround the redundant Storable constraint + -- see haskell/vector#394 + 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 :: Storable e => Comp -> VS.Vector e -> Array S Ix1 e -fromStorableVector comp v = SArray {sComp = comp, sSize = SafeSz (VS.length v), sData = v} +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 -> MArray s S Ix1 e -fromStorableMVector mv@(MVS.MVector len _) = MSArray (SafeSz len) mv +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 :: 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 :: 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 :: 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. @@ -337,7 +338,7 @@ -- @since 0.3.0 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 #-} @@ -345,9 +346,8 @@ -- modify the pointer, unless the array gets frozen prior to modification. -- -- @since 0.3.0 -unsafeMArrayFromForeignPtr0 :: 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 #-} @@ -356,8 +356,7 @@ -- -- @since 0.3.0 unsafeMArrayFromForeignPtr :: Storable e => ForeignPtr e -> Int -> Sz1 -> MArray s S Ix1 e -unsafeMArrayFromForeignPtr fp offset sz = - MSArray sz (MVS.unsafeFromForeignPtr fp offset (unSz sz)) +unsafeMArrayFromForeignPtr fp offset sz = MSArray sz (advanceForeignPtr fp offset) {-# INLINE unsafeMArrayFromForeignPtr #-} @@ -366,13 +365,24 @@ -- -- @since 0.5.9 unsafeMallocMArray :: - forall ix e m. (Index ix, Storable e, MonadIO m) + forall ix e m. (Index ix, Storable e, PrimMonad m) => Sz ix - -> m (MArray RealWorld S ix e) -unsafeMallocMArray sz = liftIO $ do + -> 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 (MVS.unsafeFromForeignPtr0 foreignPtr n) + 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 @@ -16,8 +16,9 @@ -- module Data.Massiv.Array.Manifest.Unboxed ( U (..) - , VU.Unbox + , Unbox , Array(..) + , MArray(..) , toUnboxedVector , toUnboxedMVector , fromUnboxedVector @@ -26,13 +27,13 @@ import Control.DeepSeq (NFData(..), deepseq) import Data.Massiv.Array.Delayed.Pull (eqArrays, compareArrays) -import Data.Massiv.Array.Manifest.Internal (M, toManifest) import Data.Massiv.Array.Manifest.List as A import Data.Massiv.Vector.Stream as S (steps, isteps) import Data.Massiv.Array.Mutable import Data.Massiv.Core.Common import Data.Massiv.Core.List import Data.Massiv.Core.Operations +import Data.Vector.Unboxed (Unbox) import qualified Data.Vector.Generic.Mutable as VGM import qualified Data.Vector.Unboxed as VU import qualified Data.Vector.Unboxed.Mutable as MVU @@ -49,8 +50,9 @@ , uSize :: !(Sz ix) , uData :: !(VU.Vector e) } +data instance MArray s U ix e = MUArray !(Sz ix) !(VU.MVector s e) -instance (Ragged L ix e, Show e, VU.Unbox e) => Show (Array U ix e) where +instance (Ragged L ix e, Show e, Unbox e) => Show (Array U ix e) where showsPrec = showsArrayPrec id showList = showArrayList @@ -62,110 +64,76 @@ rnf (MUArray sz mv) = sz `deepseq` mv `deepseq` () {-# INLINE rnf #-} -instance (VU.Unbox e, Index ix) => Construct U ix e where +instance Strategy U where + getComp = uComp + {-# INLINE getComp #-} setComp c arr = arr { uComp = c } {-# INLINE setComp #-} - makeArrayLinear !comp !sz f = unsafePerformIO $ generateArrayLinear comp sz (pure . f) - {-# INLINE makeArrayLinear #-} - replicate comp !sz !e = runST (newMArray sz e >>= unsafeFreeze comp) - {-# INLINE replicate #-} - - -instance (VU.Unbox e, Eq e, Index ix) => Eq (Array U ix e) where +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 +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 {-# INLINE unsafeLinearIndex #-} + + unsafeOuterSlice (UArray c _ v) szL i = + let k = totalElem szL + in UArray c szL $ VU.unsafeSlice (i * k) k v + {-# INLINE unsafeOuterSlice #-} + unsafeLinearSlice i k (UArray c _ v) = UArray c k $ VU.unsafeSlice i (unSz k) v {-# INLINE unsafeLinearSlice #-} +instance Index ix => Shape U ix where + maxLinearSize = Just . SafeSz . elemsCount + {-# INLINE maxLinearSize #-} -instance Index ix => Resize U ix where +instance Size U where + size = uSize + {-# INLINE size #-} unsafeResize !sz !arr = arr { uSize = sz } {-# INLINE unsafeResize #-} -instance (VU.Unbox e, Index ix) => Extract U ix e where - unsafeExtract !sIx !newSz !arr = unsafeExtract sIx newSz (toManifest arr) - {-# INLINE unsafeExtract #-} +instance (Unbox e, Index ix) => Load U ix e where + makeArrayLinear !comp !sz f = unsafePerformIO $ generateArrayLinear comp sz (pure . f) + {-# INLINE makeArrayLinear #-} -instance (VU.Unbox e, Index ix) => Load U ix e where - type R U = M - size = uSize - {-# INLINE size #-} - getComp = uComp - {-# 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 (VU.Unbox e, Index ix) => StrideLoad U ix e + iterArrayLinearST_ !scheduler !arr = + splitLinearlyWith_ scheduler (elemsCount arr) (unsafeLinearIndex arr) + {-# INLINE iterArrayLinearST_ #-} +instance (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 #-} -instance {-# OVERLAPPING #-} VU.Unbox e => OuterSlice U Ix1 e where - unsafeOuterSlice = unsafeLinearIndex - {-# INLINE unsafeOuterSlice #-} - -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 #-} - -instance {-# OVERLAPPING #-} VU.Unbox e => InnerSlice U Ix1 e where - unsafeInnerSlice arr _ = unsafeLinearIndex arr - {-# INLINE unsafeInnerSlice #-} - -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 #-} - -instance (VU.Unbox e, Index ix) => Manifest U ix e where +instance Unbox e => Manifest U e where unsafeLinearIndexM (UArray _ _ v) = INDEX_CHECK("(Manifest U ix e).unsafeLinearIndexM", Sz . 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 #-} @@ -184,31 +152,25 @@ {-# INLINE unsafeLinearCopy #-} unsafeLinearRead (MUArray _ mv) = - INDEX_CHECK("(Mutable U ix e).unsafeLinearRead", Sz . MVU.length, MVU.unsafeRead) mv + INDEX_CHECK("(Manifest U ix e).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 + INDEX_CHECK("(Manifest U ix e).unsafeLinearWrite", Sz . MVU.length, MVU.unsafeWrite) mv {-# INLINE unsafeLinearWrite #-} unsafeLinearGrow (MUArray _ mv) sz = MUArray sz <$> MVU.unsafeGrow mv (totalElem sz) {-# INLINE unsafeLinearGrow #-} -instance (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 #-} @@ -251,7 +213,7 @@ -- | /O(1)/ - Wrap an unboxed vector and produce an unboxed flat array. -- -- @since 0.6.0 -fromUnboxedVector :: VU.Unbox e => Comp -> VU.Vector e -> Array U Ix1 e +fromUnboxedVector :: VU.Unbox e => Comp -> VU.Vector e -> Vector U e fromUnboxedVector comp v = UArray comp (SafeSz (VU.length v)) v {-# INLINE fromUnboxedVector #-} @@ -259,6 +221,6 @@ -- | /O(1)/ - Wrap an unboxed mutable vector and produce a mutable unboxed flat array. -- -- @since 0.5.0 -fromUnboxedMVector :: VU.Unbox e => VU.MVector s e -> MArray s U Ix1 e +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 @@ -23,12 +23,12 @@ ) 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,19 +39,19 @@ 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 = 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 @@ -59,7 +59,7 @@ -- 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, Mutable r ix e, ARepr v ~ r) +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 @@ -72,7 +72,7 @@ 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} + 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} @@ -90,12 +90,7 @@ -- -- @since 0.3.0 fromVectorM :: - ( MonadThrow m - , Typeable v - , VG.Vector v a - , Mutable (ARepr v) ix a - , Mutable r ix a - ) + (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 @@ -113,19 +108,19 @@ -- -- @since 0.3.0 fromVector' :: - (Typeable v, VG.Vector v a, Mutable (ARepr v) ix a, Mutable r ix a) + (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 -> 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) + forall v r ix e. (Manifest r e, Index ix, VRepr r ~ v) => Array r ix e -> Maybe (v e) castToVector arr = @@ -135,14 +130,14 @@ return $ uData uArr , do Refl <- eqT :: Maybe (r :~: S) sArr <- gcastArr arr - return $ sData sArr + 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 :~: N) + , do Refl <- eqT :: Maybe (r :~: BN) bArr <- gcastArr arr return $ toBoxedVector $ toLazyArray $ unwrapNormalForm bArr , do Refl <- eqT :: Maybe (r :~: BL) @@ -161,6 +156,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] @@ -175,8 +171,9 @@ -- toVector :: forall r ix e v. - ( Manifest r ix e - , Mutable (ARepr v) ix e + ( Manifest r e + , Load r ix e + , Manifest (ARepr v) e , VG.Vector v e , VRepr (ARepr v) ~ v ) 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,6 +1,7 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE LambdaCase #-} +{-# LANGUAGE MonoLocalBinds #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} @@ -14,25 +15,27 @@ -- module Data.Massiv.Array.Mutable ( -- ** Size - msize + sizeOfMArray + , msize + , resizeMArrayM + , flattenMArray + , outerSliceMArrayM + , outerSlicesMArray -- ** Element-wise mutation , read , readM - , read' , write , write_ , writeM - , write' , modify , modify_ , modifyM , modifyM_ - , modify' , swap , swap_ , swapM , swapM_ - , swap' + , zipSwapM_ -- ** Operations on @MArray@ -- *** Immutable conversion , thaw @@ -40,7 +43,6 @@ , freeze , freezeS -- *** Create mutable - , new , newMArray , newMArray' , makeMArray @@ -78,6 +80,8 @@ , iforPrimM_ , iforLinearPrimM , iforLinearPrimM_ + , for2PrimM_ + , ifor2PrimM_ -- *** Modify , withMArray , withMArray_ @@ -94,7 +98,7 @@ , initialize , initializeNew -- ** Computation - , Mutable + , Manifest , MArray , RealWorld , computeInto @@ -107,28 +111,123 @@ import Data.Maybe (fromMaybe) import Control.Monad (void, when, unless, (>=>)) import Control.Monad.ST +import Control.Monad.Primitive import Control.Scheduler import Data.Massiv.Core.Common import Data.Massiv.Array.Mutable.Internal +import Data.Massiv.Array.Delayed.Pull (D) 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. +-- | /O(1)/ - Change the size of a mutable array. Throws +-- `SizeElementsMismatchException` if total number of elements does not match +-- the supplied array. -- --- @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 #-} -{-# DEPRECATED new "In favor of a more robust and safer `newMArray` or a more consistently named `newMArray'`" #-} +-- @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 in will be a thunk with `Uninitialized` exception, while for others it will be --- simply zeros. This is a partial function. +-- boxed arrays it will be a thunk with `Uninitialized` exception, while for others it will be +-- simply zeros. -- -- ==== __Examples__ -- @@ -148,12 +247,12 @@ -- [ [ 0, 0, 0, 0, 0, 0 ] -- , [ 0, 0, 0, 0, 0, 0 ] -- ] --- >>> newMArray' @B @_ @Int (Sz2 2 6) >>= (`readM` 1) +-- >>> newMArray' @B @_ @Int (Sz2 2 6) >>= freezeS -- *** Exception: Uninitialized -- -- @since 0.6.0 newMArray' :: - forall r ix e m. (Mutable r ix e, PrimMonad m) + 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 @@ -178,7 +277,7 @@ -- ] -- -- @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 :: 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 @@ -210,7 +309,7 @@ -- -- @since 0.3.0 thawS :: - forall r ix e m. (Mutable r ix e, PrimMonad m) + 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 @@ -226,7 +325,7 @@ -- ==== __Example__ -- -- >>> import Data.Massiv.Array --- >>> marr <- newMArray @P @_ @Int (Sz2 2 6) 0 +-- >>> 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)) @@ -236,13 +335,13 @@ -- -- @since 0.1.0 freeze :: - forall r ix e m. (Mutable r ix e, MonadIO m) + 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 = msize smarr + let sz = sizeOfMArray smarr totalLength = totalElem sz tmarr <- unsafeNew sz withMassivScheduler_ comp $ \scheduler -> @@ -262,33 +361,31 @@ -- -- @since 0.3.0 freezeS :: - forall r ix e m. (Mutable r ix e, PrimMonad m) + 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 = msize smarr + let sz = sizeOfMArray smarr tmarr <- unsafeNew sz unsafeLinearCopy smarr 0 tmarr 0 (SafeSz (totalElem sz)) unsafeFreeze Seq tmarr {-# INLINE freezeS #-} - -unsafeNewUninitialized :: - (Load r' ix e, Mutable r ix e, PrimMonad m) => Array r' ix e -> m (MArray (PrimState m) r ix e) -unsafeNewUninitialized !arr = unsafeNew (fromMaybe zeroSz (maxSize arr)) -{-# INLINE unsafeNewUninitialized #-} - +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, Mutable r ix e, PrimMonad m) + 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 <- unsafeNewUninitialized arr - unsafeLoadIntoS marr arr + marr <- unsafeNewUpper arr + stToPrim $ unsafeLoadIntoST marr arr {-# INLINE loadArrayS #-} @@ -296,13 +393,13 @@ -- -- @since 0.3.0 loadArray :: - forall r ix e r' m. (Load r' ix e, Mutable r ix e, MonadIO m) + 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 <- unsafeNewUninitialized arr - unsafeLoadIntoM marr arr + marr <- unsafeNewUpper arr + unsafeLoadIntoIO marr arr {-# INLINE loadArray #-} @@ -312,16 +409,16 @@ -- -- @since 0.1.3 computeInto :: - (Load r' ix' e, Mutable r ix e, MonadIO m) + (Size r', 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 - unless (totalElem (msize mArr) == totalElem (size arr)) $ - throwM $ SizeElementsMismatchException (msize mArr) (size arr) + unless (totalElem (sizeOfMArray mArr) == totalElem (size arr)) $ + throwM $ SizeElementsMismatchException (sizeOfMArray mArr) (size arr) withMassivScheduler_ (getComp arr) $ \scheduler -> - loadArrayM scheduler arr (unsafeLinearWrite mArr) + stToPrim $ iterArrayLinearST_ scheduler arr (unsafeLinearWrite mArr) {-# INLINE computeInto #-} @@ -329,7 +426,7 @@ -- -- @since 0.3.0 makeMArrayS :: - forall r ix e m. (Mutable r ix e, PrimMonad m) + 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) @@ -341,13 +438,13 @@ -- -- @since 0.3.0 makeMArrayLinearS :: - forall r ix e m. (Mutable r ix e, PrimMonad m) + 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 - loopM_ 0 (< totalElem (msize marr)) (+ 1) (\ !i -> f i >>= unsafeLinearWrite marr i) + loopM_ 0 (< totalElem (sizeOfMArray marr)) (+ 1) (\ !i -> f i >>= unsafeLinearWrite marr i) return marr {-# INLINE makeMArrayLinearS #-} @@ -355,11 +452,11 @@ -- -- @since 0.3.0 makeMArray :: - forall r ix e m. (PrimMonad m, MonadUnliftIO m, Mutable r ix e) + forall r ix e m. (MonadUnliftIO m, Manifest r e, Index ix) => Comp -> Sz ix -> (ix -> m e) - -> m (MArray (PrimState m) r ix e) + -> m (MArray RealWorld r ix e) makeMArray comp sz f = makeMArrayLinear comp sz (f . fromLinearIndex sz) {-# INLINE makeMArray #-} @@ -368,15 +465,16 @@ -- -- @since 0.3.0 makeMArrayLinear :: - forall r ix e m. (PrimMonad m, MonadUnliftIO m, Mutable r ix e) + forall r ix e m. (MonadUnliftIO m, Manifest r e, Index ix) => Comp -> Sz ix -> (Int -> m e) - -> m (MArray (PrimState m) r ix e) + -> m (MArray RealWorld r ix e) makeMArrayLinear comp sz f = do - marr <- unsafeNew sz + marr <- liftIO $ unsafeNew sz withScheduler_ comp $ \scheduler -> - splitLinearlyWithM_ scheduler (totalElem sz) f (unsafeLinearWrite marr) + withRunInIO $ \run -> + splitLinearlyWithM_ scheduler (totalElem sz) (run . f) (unsafeLinearWrite marr) return marr {-# INLINE makeMArrayLinear #-} @@ -397,16 +495,16 @@ -- @since 0.3.0 -- createArray_ :: - forall r ix e a m. (Mutable r ix e, PrimMonad m, MonadUnliftIO m) + 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 m () -> MArray (PrimState m) r ix e -> m a) + -> (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 <- new sz + marr <- liftIO $ newMArray' sz withScheduler_ comp (`action` marr) - unsafeFreeze comp marr + liftIO $ unsafeFreeze comp marr {-# INLINE createArray_ #-} -- | Just like `createArray_`, but together with `Array` it returns results of scheduled filling @@ -415,16 +513,16 @@ -- @since 0.3.0 -- createArray :: - forall r ix e a m b. (Mutable r ix e, PrimMonad m, MonadUnliftIO m) + 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 m a -> MArray (PrimState m) r ix e -> m b) + -> (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 <- new sz + marr <- liftIO $ newMArray' sz a <- withScheduler comp (`action` marr) - arr <- unsafeFreeze comp marr + arr <- liftIO $ unsafeFreeze comp marr return (a, arr) {-# INLINE createArray #-} @@ -442,7 +540,7 @@ -- -- @since 0.3.0 createArrayS_ :: - forall r ix e a m. (Mutable r ix e, PrimMonad m) + 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 @@ -454,13 +552,13 @@ -- -- @since 0.3.0 createArrayS :: - forall r ix e a m. (Mutable r ix e, PrimMonad m) + 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 <- new sz + marr <- newMArray' sz a <- action marr arr <- unsafeFreeze Seq marr return (a, arr) @@ -470,7 +568,7 @@ -- -- @since 0.3.0 createArrayST_ :: - forall r ix e a. Mutable r ix e + 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 @@ -482,7 +580,7 @@ -- -- @since 0.2.6 createArrayST :: - forall r ix e a. Mutable r ix e + 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) @@ -491,7 +589,7 @@ -- | 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 +-- 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. @@ -515,7 +613,7 @@ -- -- @since 0.2.6 generateArrayS :: - forall r ix e m. (Mutable r ix e, PrimMonad m) + forall r ix e m. (Manifest r e, Index ix, PrimMonad m) => Sz ix -- ^ Resulting size of the array -> (ix -> m e) -- ^ Element producing generator -> m (Array r ix e) @@ -526,13 +624,13 @@ -- -- @since 0.3.0 generateArrayLinearS :: - forall r ix e m. (Mutable r ix e, PrimMonad m) + 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 - loopM_ 0 (< totalElem (msize marr)) (+ 1) $ \i -> gen i >>= unsafeLinearWrite marr i + loopM_ 0 (< totalElem (sizeOfMArray marr)) (+ 1) $ \i -> gen i >>= unsafeLinearWrite marr i unsafeFreeze Seq marr {-# INLINE generateArrayLinearS #-} @@ -542,7 +640,7 @@ -- -- @since 0.2.6 generateArray :: - forall r ix e m. (MonadUnliftIO m, PrimMonad m, Mutable r ix e) + forall r ix e m. (MonadUnliftIO m, Manifest r e, Index ix) => Comp -> Sz ix -> (ix -> m e) @@ -555,12 +653,12 @@ -- -- @since 0.3.0 generateArrayLinear :: - forall r ix e m. (MonadUnliftIO m, PrimMonad m, Mutable r ix e) + forall r ix e m. (MonadUnliftIO m, Manifest r e, Index ix) => Comp -> Sz ix -> (Int -> m e) -> m (Array r ix e) -generateArrayLinear comp sz f = makeMArrayLinear comp sz f >>= unsafeFreeze comp +generateArrayLinear comp sz f = makeMArrayLinear comp sz f >>= liftIO . unsafeFreeze comp {-# INLINE generateArrayLinear #-} @@ -568,7 +666,7 @@ -- -- @since 0.3.4 generateArrayLinearWS :: - forall r ix e s m. (Mutable r ix e, MonadUnliftIO m, PrimMonad m) + forall r ix e s m. (Manifest r e, Index ix, MonadUnliftIO m, PrimMonad m) => WorkerStates s -> Sz ix -> (Int -> s -> m e) @@ -589,7 +687,7 @@ -- -- @since 0.3.4 generateArrayWS :: - forall r ix e s m. (Mutable r ix e, MonadUnliftIO m, PrimMonad m) + forall r ix e s m. (Manifest r e, Index ix, MonadUnliftIO m, PrimMonad m) => WorkerStates s -> Sz ix -> (ix -> s -> m e) @@ -622,7 +720,7 @@ -- -- @since 0.3.0 unfoldrPrimM_ :: - forall r ix e a m. (Mutable r ix e, PrimMonad m) + 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 @@ -634,7 +732,7 @@ -- -- @since 0.3.0 iunfoldrPrimM_ :: - forall r ix e a m. (Mutable r ix e, PrimMonad m) + 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 @@ -647,14 +745,14 @@ -- -- @since 0.3.0 iunfoldrPrimM :: - forall r ix e a m. (Mutable r ix e, PrimMonad m) + 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' = msize 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 @@ -665,14 +763,14 @@ -- -- @since 0.3.0 unfoldrPrimM :: - forall r ix e a m. (Mutable r ix e, PrimMonad m) + 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' = msize marr + let sz' = sizeOfMArray marr in loopM 0 (< totalElem sz') (+ 1) acc0 $ \ !i !acc -> do (e, acc') <- gen acc unsafeLinearWrite marr i e @@ -703,7 +801,7 @@ -- -- @since 0.3.0 unfoldlPrimM_ :: - forall r ix e a m. (Mutable r ix e, PrimMonad m) + 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 @@ -715,7 +813,7 @@ -- -- @since 0.3.0 iunfoldlPrimM_ :: - forall r ix e a m. (Mutable r ix e, PrimMonad m) + 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 @@ -728,14 +826,14 @@ -- -- @since 0.3.0 iunfoldlPrimM :: - forall r ix e a m. (Mutable r ix e, PrimMonad m) + 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' = msize 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 @@ -746,14 +844,14 @@ -- -- @since 0.3.0 unfoldlPrimM :: - forall r ix e a m. (Mutable r ix e, PrimMonad m) + 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' = msize marr + let sz' = sizeOfMArray marr in loopDeepM 0 (< totalElem sz') (+1) acc0 $ \ !i !acc -> do (acc', e) <- gen acc unsafeLinearWrite marr i e @@ -764,17 +862,17 @@ -- 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_ :: (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> (e -> m ()) -> m () forPrimM_ marr f = - loopM_ 0 (< totalElem (msize marr)) (+1) (unsafeLinearRead marr >=> f) + loopM_ 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 :: (Mutable r ix e, PrimMonad m) => MArray (PrimState m) r ix e -> (e -> m e) -> m () +forPrimM :: (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> (e -> m e) -> m () forPrimM marr f = - loopM_ 0 (< totalElem (msize marr)) (+1) (unsafeLinearModify marr f) + loopM_ 0 (< totalElem (sizeOfMArray marr)) (+1) (unsafeLinearModify marr f) {-# INLINE forPrimM #-} @@ -784,16 +882,16 @@ -- -- @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)) + (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 :: - (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)) + (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 #-} @@ -803,27 +901,61 @@ -- -- @since 0.4.0 iforLinearPrimM_ :: - (Mutable r ix e, PrimMonad m) => MArray (PrimState m) r ix e -> (Int -> e -> m ()) -> m () + (Manifest r e, Index ix, 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) + loopM_ 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 :: - (Mutable r ix e, PrimMonad m) => MArray (PrimState m) r ix e -> (Int -> e -> m e) -> m () + (Manifest r e, Index ix, 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) + loopM_ 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)) + iterM_ 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 :: - (Mutable r ix e, MonadUnliftIO m) + (Manifest r e, Index ix, MonadUnliftIO m) => Array r ix e - -> (Scheduler m a -> MArray RealWorld r ix e -> m b) + -> (Scheduler RealWorld a -> MArray RealWorld r ix e -> m b) -> m ([a], Array r ix e) withMArray arr action = do marr <- thaw arr @@ -845,9 +977,9 @@ -- -- @since 0.5.0 withMArray_ :: - (Mutable r ix e, MonadUnliftIO m) + (Manifest r e, Index ix, MonadUnliftIO m) => Array r ix e - -> (Scheduler m () -> MArray RealWorld r ix e -> m a) + -> (Scheduler RealWorld () -> MArray RealWorld r ix e -> m a) -> m (Array r ix e) withMArray_ arr action = do marr <- thaw arr @@ -861,15 +993,13 @@ -- -- @since 0.6.1 withLoadMArray_ :: - forall r ix e r' m b. (Load r' ix e, Mutable r ix e, MonadUnliftIO m) + forall r ix e r' m b. (Load r' ix e, Manifest r e, MonadUnliftIO m) => Array r' ix e - -> (Scheduler m () -> MArray RealWorld r ix e -> m b) + -> (Scheduler RealWorld () -> MArray RealWorld r ix e -> m b) -> m (Array r ix e) withLoadMArray_ arr action = do - marr <- liftIO $ unsafeNew (size arr) - withScheduler_ (getComp arr) $ \scheduler -> do - runBatch_ scheduler $ \_ -> loadArrayM scheduler arr (\i -> liftIO . unsafeLinearWrite marr i) - action scheduler marr + marr <- loadArray arr + withScheduler_ (getComp arr) (`action` marr) liftIO $ unsafeFreeze (getComp arr) marr {-# INLINE[2] withLoadMArray_ #-} {-# RULES @@ -884,7 +1014,7 @@ -- -- @since 0.5.0 withMArrayS :: - (Mutable r ix e, PrimMonad m) + (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) @@ -899,7 +1029,7 @@ -- -- @since 0.5.0 withMArrayS_ :: - (Mutable r ix e, PrimMonad m) + (Manifest r e, Index ix, PrimMonad m) => Array r ix e -> (MArray (PrimState m) r ix e -> m a) -> m (Array r ix e) @@ -911,13 +1041,12 @@ -- -- @since 0.6.1 withLoadMArrayS :: - forall r ix e r' m a. (Load r' ix e, Mutable r ix e, PrimMonad m) + 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 <- unsafeNew (size arr) - loadArrayM trivialScheduler_ arr (unsafeLinearWrite marr) + marr <- loadArrayS arr a <- action marr (,) a <$> unsafeFreeze (getComp arr) marr {-# INLINE[2] withLoadMArrayS #-} @@ -926,7 +1055,7 @@ -- -- @since 0.6.1 withLoadMArrayS_ :: - forall r ix e r' m a. (Load r' ix e, Mutable r ix e, PrimMonad m) + 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) @@ -939,7 +1068,7 @@ -- -- @since 0.5.0 withMArrayST :: - Mutable r ix e + (Manifest r e, Index ix) => Array r ix e -> (forall s . MArray s r ix e -> ST s a) -> (a, Array r ix e) @@ -951,7 +1080,7 @@ -- -- @since 0.5.0 withMArrayST_ :: - Mutable r ix e => Array r ix e -> (forall s. MArray s r ix e -> ST s a) -> Array r ix e + (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_ #-} @@ -960,7 +1089,7 @@ -- -- @since 0.6.1 withLoadMArrayST :: - forall r ix e r' a. (Load r' ix e, Mutable r ix e) + 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) @@ -971,7 +1100,7 @@ -- -- @since 0.6.1 withLoadMArrayST_ :: - forall r ix e r' a. (Load r' ix e, Mutable r ix e) + 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 @@ -982,10 +1111,10 @@ -- | /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) => +read :: (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> ix -> m (Maybe e) read marr ix = - if isSafeIndex (msize marr) ix + if isSafeIndex (sizeOfMArray marr) ix then Just <$> unsafeRead marr ix else return Nothing {-# INLINE read #-} @@ -994,34 +1123,22 @@ -- | /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) => +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 (msize marr) ix + Nothing -> throwM $ IndexOutOfBoundsException (sizeOfMArray 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 :: (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> ix -> e -> m Bool write marr ix e = - if isSafeIndex (msize marr) ix + if isSafeIndex (sizeOfMArray marr) ix then unsafeWrite marr ix e >> pure True else pure False {-# INLINE write #-} @@ -1032,42 +1149,32 @@ -- 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 +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 :: - (Mutable r ix e, PrimMonad m, MonadThrow m) => MArray (PrimState m) r ix e -> ix -> e -> m () + (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 (msize marr) ix)) + write marr ix e >>= (`unless` throwM (IndexOutOfBoundsException (sizeOfMArray 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) + (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 (msize marr) ix + if isSafeIndex (sizeOfMArray marr) ix then Just <$> unsafeModify marr f ix else return Nothing {-# INLINE modify #-} @@ -1078,12 +1185,12 @@ -- -- @since 0.4.4 modify_ :: - (Mutable r ix e, PrimMonad m) + (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 (msize marr) ix) $ void $ unsafeModify marr f ix +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 @@ -1092,14 +1199,14 @@ -- -- @since 0.4.0 modifyM :: - (Mutable r ix e, PrimMonad m, MonadThrow m) + (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 (msize marr) ix = unsafeModify marr f ix - | otherwise = throwM (IndexOutOfBoundsException (msize marr) 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 @@ -1115,7 +1222,7 @@ -- -- @since 0.4.0 modifyM_ :: - (Mutable r ix e, PrimMonad m, MonadThrow m) + (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. @@ -1124,27 +1231,14 @@ {-# 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 :: (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 = msize marr + let !sz = sizeOfMArray marr in if isSafeIndex sz ix1 && isSafeIndex sz ix2 then Just <$> unsafeSwap marr ix1 ix2 else pure Nothing @@ -1155,9 +1249,9 @@ -- 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_ :: (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> ix -> ix -> m () swap_ marr ix1 ix2 = - let !sz = msize marr + let !sz = sizeOfMArray marr in when (isSafeIndex sz ix1 && isSafeIndex sz ix2) $ void $ unsafeSwap marr ix1 ix2 {-# INLINE swap_ #-} @@ -1167,7 +1261,7 @@ -- -- @since 0.4.0 swapM :: - (Mutable r ix e, PrimMonad m, MonadThrow m) + (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. @@ -1175,11 +1269,11 @@ -- 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 + | 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 = msize marr + !sz = sizeOfMArray marr {-# INLINE swapM #-} @@ -1187,22 +1281,40 @@ -- -- @since 0.4.0 swapM_ :: - (Mutable r ix e, PrimMonad m, MonadThrow m) => MArray (PrimState m) r ix e -> ix -> ix -> m () + (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) + iterM_ 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_ #-} --- | /O(1)/ - Same as `swap`, but throws an `IndexOutOfBoundsException` on invalid indices. +-- | Get the size of a mutable array. -- -- @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`" #-} +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 @@ -28,16 +28,16 @@ -- >>> 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 +-- @since 1.0.0 unstablePartitionM :: - forall r e m. (Mutable r Ix1 e, PrimMonad m) + forall r e m. (Manifest r e, PrimMonad m) => MVector (PrimState m) r e - -> (e -> Bool) -- ^ Predicate + -> (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 @@ -37,7 +37,7 @@ 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 #-} @@ -49,7 +49,7 @@ 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 #-} @@ -65,7 +65,7 @@ -> 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 #-} @@ -77,7 +77,7 @@ 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 #-} @@ -88,7 +88,7 @@ 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 #-} @@ -99,7 +99,7 @@ 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 #-} @@ -110,7 +110,7 @@ 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 #-} @@ -121,7 +121,7 @@ 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 #-} @@ -132,7 +132,7 @@ 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 #-} @@ -143,6 +143,6 @@ 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 --- a/src/Data/Massiv/Array/Mutable/Internal.hs +++ b/src/Data/Massiv/Array/Mutable/Internal.hs @@ -21,7 +21,7 @@ -- -- @since 0.5.0 unsafeCreateArrayS :: - forall r ix e a m. (Mutable r ix e, PrimMonad m) + 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 @@ -38,16 +38,16 @@ -- -- @since 0.5.0 unsafeCreateArray :: - forall r ix e a m b. (Mutable r ix e, PrimMonad m, MonadUnliftIO m) + 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 m a -> MArray (PrimState m) r ix e -> m b) + -> (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 <- unsafeNew sz + marr <- liftIO $ unsafeNew sz a <- withScheduler comp (`action` marr) - arr <- unsafeFreeze comp marr + arr <- liftIO $ unsafeFreeze comp marr return (a, arr) {-# INLINE unsafeCreateArray #-} @@ -56,15 +56,15 @@ -- -- @since 0.5.0 unsafeCreateArray_ :: - forall r ix e a m b. (Mutable r ix e, PrimMonad m, MonadUnliftIO m) + 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 m a -> MArray (PrimState m) r ix e -> m b) + -> (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 <- unsafeNew sz + marr <- liftIO $ unsafeNew sz withScheduler_ comp (`action` marr) - arr <- unsafeFreeze comp 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 @@ -15,11 +15,14 @@ ( -- * Numeric Numeric , NumericFloat + , liftNumArray2M -- ** Pointwise addition , (.+) , (+.) , (.+.) , (!+!) + , sumArraysM + , sumArrays' -- ** Pointwise subtraction , (.-) , (-.) @@ -31,6 +34,8 @@ , (.*.) , (!*!) , (.^) + , productArraysM + , productArrays' -- ** Dot product , (!.!) , dotM @@ -96,14 +101,13 @@ ) where import Data.Massiv.Array.Mutable -import Data.Massiv.Array.Manifest import Data.Massiv.Array.Delayed.Pull import Data.Massiv.Array.Delayed.Push import Data.Massiv.Array.Manifest.Internal import Data.Massiv.Array.Ops.Map as A import Data.Massiv.Array.Ops.Construct import Data.Massiv.Core -import Data.Massiv.Core.Common +import Data.Massiv.Core.Common as A import Data.Massiv.Core.Operations import Prelude as P import System.IO.Unsafe @@ -116,28 +120,40 @@ infixl 7 !*!, .*., .*, *., !/!, ./., ./, /., `quotA`, `remA`, `divA`, `modA` infixl 6 !+!, .+., .+, +., !-!, .-., .-, -. -liftArray2M :: - (Load r ix e, Numeric r e, MonadThrow m) +-- | Similar to `liftArray2M`, except it can be applied only to representations +-- with `Numeric` instance and result representation stays the same. +-- +-- @since 1.0.0 +liftNumArray2M :: + (Index ix, 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 +liftNumArray2M f a1 a2 | size a1 == size a2 = pure $ unsafeLiftArray2 f a1 a2 - | otherwise = throwM $ SizeMismatchException (size a1) (size a2) -{-# INLINE liftArray2M #-} + | isZeroSz sz1 && isZeroSz sz2 = pure $ unsafeResize zeroSz a1 + | otherwise = throwM $ SizeMismatchException sz1 sz2 + where + !sz1 = size a1 + !sz2 = size a2 +{-# INLINE liftNumArray2M #-} -liftNumericArray2M :: - (Load r ix e, MonadThrow m) +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) -liftNumericArray2M f a1 a2 +applyExactSize2M f a1 a2 | size a1 == size a2 = pure $ f a1 a2 - | otherwise = throwM $ SizeMismatchException (size a1) (size a2) -{-# INLINE liftNumericArray2M #-} + | 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 @@ -146,9 +162,8 @@ -- /__Throws Exception__/: `SizeMismatchException` when array 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 +(.+.) :: (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 @@ -165,7 +180,7 @@ -- [ 20, 22, 24, 26, 28, 30, 32, 34, 36, 38, 40 ] -- -- @since 0.5.6 -(!+!) :: (Load r ix e, Numeric r e) => Array r ix e -> Array r ix e -> Array r ix e +(!+!) :: (Index ix, Numeric r e) => Array r ix e -> Array r ix e -> Array r ix e (!+!) a1 a2 = throwEither (a1 .+. a2) {-# INLINE (!+!) #-} @@ -190,8 +205,8 @@ -- -- @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 + (Index ix, Numeric r e, MonadThrow m) => Array r ix e -> Array r ix e -> m (Array r ix e) +(.-.) = applyExactSize2M subtractionPointwise {-# INLINE (.-.) #-} @@ -209,7 +224,7 @@ -- [ -20, -20, -20, -20, -20, -20, -20, -20, -20, -20, -20 ] -- -- @since 0.5.6 -(!-!) :: (Load r ix e, Numeric r e) => Array r ix e -> Array r ix e -> Array r ix e +(!-!) :: (Index ix, Numeric r e) => Array r ix e -> Array r ix e -> Array r ix e (!-!) a1 a2 = throwEither (a1 .-. a2) {-# INLINE (!-!) #-} @@ -235,8 +250,8 @@ -- -- @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 + (Index ix, Numeric r e, MonadThrow m) => Array r ix e -> Array r ix e -> m (Array r ix e) +(.*.) = applyExactSize2M multiplicationPointwise {-# INLINE (.*.) #-} @@ -256,7 +271,7 @@ -- [ 0, 21, 44, 69, 96, 125, 156, 189, 224, 261, 300 ] -- -- @since 0.5.6 -(!*!) :: (Load r ix e, Numeric r e) => Array r ix e -> Array r ix e -> Array r ix e +(!*!) :: (Index ix, Numeric r e) => Array r ix e -> Array r ix e -> Array r ix e (!*!) a1 a2 = throwEither (a1 .*. a2) {-# INLINE (!*!) #-} @@ -320,7 +335,7 @@ -- [Partial] Throws an impure exception when lengths of vectors do not match -- -- @since 0.5.6 -(!.!) :: (Numeric r e, Source r Ix1 e) => Vector r e -> Vector r e -> e +(!.!) :: (Numeric r e, Source r e) => Vector r e -> Vector r e -> e (!.!) v1 v2 = throwEither $ dotM v1 v2 {-# INLINE (!.!) #-} @@ -329,7 +344,7 @@ -- /__Throws Exception__/: `SizeMismatchException` when lengths of vectors do not match -- -- @since 0.5.6 -dotM :: (FoldNumeric r e, Source r Ix1 e, MonadThrow m) => Vector r e -> Vector r e -> m e +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 @@ -340,14 +355,14 @@ unsafeDotProductIO :: - (MonadUnliftIO m, FoldNumeric r b, Source r ix b) + (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 -> do + splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> liftIO $ do let n = SafeSz chunkLength loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start -> scheduleWork scheduler $ @@ -367,21 +382,21 @@ -- | Compute L2 norm of an array. -- -- @since 0.5.6 -normL2 :: (Floating e, FoldNumeric r e, Source r ix e) => Array r ix e -> e +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, FoldNumeric r b, Source r ix b) + (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 -> do + splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> liftIO $ do let n = SafeSz chunkLength loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start -> scheduleWork scheduler $ pure $! powerSumArray (unsafeLinearSlice start n v) p @@ -401,7 +416,7 @@ -- -- @since 0.5.6 (.><) :: - (MonadThrow m, FoldNumeric r e, Source r Ix1 e, Source r Ix2 e) + (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) @@ -422,7 +437,7 @@ -- -- @since 0.5.7 multiplyMatrixByVector :: - (MonadThrow m, Numeric r e, Mutable r Ix1 e, Mutable r Ix2 e) + (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) @@ -436,7 +451,7 @@ -- -- @since 0.5.6 (!><) :: - (Numeric r e, Source r Ix1 e, Source r Ix2 e) + (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 @@ -450,7 +465,7 @@ -- /__Throws Exception__/: `SizeMismatchException` when inner dimensions of arrays do not match. -- -- @since 0.5.6 -(><.) :: (MonadThrow m, Numeric r e, Mutable r Ix1 e, Mutable r Ix2 e) => +(><.) :: (MonadThrow m, Numeric r e, Manifest r e) => Vector r e -- ^ Row vector -> Matrix r e -- ^ Matrix -> m (Vector r e) @@ -464,13 +479,13 @@ -- -- @since 0.5.7 multiplyVectorByMatrix :: - (MonadThrow m, Numeric r e, Mutable r Ix1 e, Mutable r Ix2 e) + (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 $ setComp comp empty + | mRows == 0 || mCols == 0 = pure $ runST (unsafeFreeze comp =<< unsafeNew zeroSz) | otherwise = pure $! unsafePerformIO $ do @@ -501,7 +516,7 @@ -- -- @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 @@ -516,6 +531,7 @@ -- -- ====__Examples__ -- +-- >>> 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 +(!> Matrix r e -> Matrix r e -> Matrix r e (!><.) :: (Numeric r e, Mutable r Ix2 e, MonadThrow m) => Matrix r e -> Matrix r e -> m (Matrix r e) +(.><.) :: (Numeric r e, Manifest r e, MonadThrow m) => Matrix r e -> Matrix r e -> m (Matrix r e) (.><.) = multiplyMatrices {-# INLINE (.><.) #-} @@ -547,12 +563,12 @@ -- -- @since 0.5.6 multiplyMatrices :: - (Numeric r e, Mutable r Ix2 e, MonadThrow m) => Matrix r e -> Matrix r e -> m (Matrix r e) + (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 $ setComp comp empty + | isEmpty arrA || isEmpty arrB = pure $ runST (unsafeFreeze comp =<< unsafeNew zeroSz) | otherwise = pure $! unsafePerformIO $ do marrC <- newMArray (SafeSz (mA :. nB)) 0 withScheduler_ comp $ \scheduler -> do @@ -693,7 +709,7 @@ -- -- @since 0.5.6 multiplyMatricesTransposed :: - (Numeric r e, Manifest r Ix2 e, MonadThrow m) + (Numeric r e, Manifest r e, MonadThrow m) => Matrix r e -> Matrix r e -> m (Matrix D e) @@ -711,8 +727,6 @@ SafeSz (n2 :. m2) = size arr2 {-# INLINE multiplyMatricesTransposed #-} - - -- | Create an indentity matrix. -- -- ==== __Example__ @@ -813,11 +827,11 @@ -- -- @since 0.4.0 (./.) :: - (Load r ix e, NumericFloat r e, MonadThrow m) + (Index ix, NumericFloat r e, MonadThrow m) => Array r ix e -> Array r ix e -> m (Array r ix e) -(./.) = liftNumericArray2M divisionPointwise +(./.) = applyExactSize2M divisionPointwise {-# INLINE (./.) #-} @@ -835,7 +849,7 @@ -- [ 0.2, 0.20792079, 0.21568628, 0.22330096, 0.23076923 ] -- -- @since 0.5.6 -(!/!) :: (Load r ix e, NumericFloat r e) => Array r ix e -> Array r ix e -> Array r ix e +(!/!) :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e -> Array r ix e (!/!) a1 a2 = throwEither (a1 ./. a2) {-# INLINE (!/!) #-} @@ -928,9 +942,9 @@ -- -- @since 0.4.0 logBaseA - :: (Source r1 ix e, Source r2 ix e, Floating e) + :: (Index ix, Source r1 e, Source r2 e, Floating e) => Array r1 ix e -> Array r2 ix e -> Array D ix e -logBaseA = liftArray2Matching logBase +logBaseA = liftArray2' logBase {-# INLINE logBaseA #-} -- TODO: siwtch to -- (breaking) logBaseA :: Array r ix e -> e -> Array D ix e @@ -948,9 +962,9 @@ -- -- @since 0.4.0 (.**) - :: (Source r1 ix e, Source r2 ix e, Floating e) + :: (Index ix, Source r1 e, Source r2 e, Floating e) => Array r1 ix e -> Array r2 ix e -> Array D ix e -(.**) = liftArray2Matching (**) +(.**) = liftArray2' (**) {-# INLINE (.**) #-} -- TODO: -- !**! :: Array r1 ix e -> Array r2 ix e -> Array D ix e @@ -1077,9 +1091,9 @@ -- -- @since 0.1.0 quotA - :: (Source r1 ix e, Source r2 ix e, Integral e) + :: (HasCallStack, Index ix, Source r1 e, Source r2 e, Integral e) => Array r1 ix e -> Array r2 ix e -> Array D ix e -quotA = liftArray2Matching quot +quotA = liftArray2' quot {-# INLINE quotA #-} @@ -1091,9 +1105,9 @@ -- -- @since 0.1.0 remA - :: (Source r1 ix e, Source r2 ix e, Integral e) + :: (HasCallStack, Index ix, Source r1 e, Source r2 e, Integral e) => Array r1 ix e -> Array r2 ix e -> Array D ix e -remA = liftArray2Matching rem +remA = liftArray2' rem {-# INLINE remA #-} -- | Perform a pointwise integer division where first array contains numerators and the @@ -1105,9 +1119,9 @@ -- -- @since 0.1.0 divA - :: (Source r1 ix e, Source r2 ix e, Integral e) + :: (HasCallStack, Index ix, Source r1 e, Source r2 e, Integral e) => Array r1 ix e -> Array r2 ix e -> Array D ix e -divA = liftArray2Matching div +divA = liftArray2' div {-# INLINE divA #-} -- TODO: -- * Array r ix e -> Array r ix e -> m (Array r ix e) @@ -1122,9 +1136,9 @@ -- -- @since 0.1.0 modA - :: (Source r1 ix e, Source r2 ix e, Integral e) + :: (HasCallStack, Index ix, Source r1 e, Source r2 e, Integral e) => Array r1 ix e -> Array r2 ix e -> Array D ix e -modA = liftArray2Matching mod +modA = liftArray2' mod {-# INLINE modA #-} @@ -1138,9 +1152,9 @@ -- -- @since 0.1.0 quotRemA - :: (Source r1 ix e, Source r2 ix e, Integral e) + :: (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 . liftArray2Matching quotRem arr1 +quotRemA arr1 = A.unzip . liftArray2' quotRem arr1 {-# INLINE quotRemA #-} @@ -1153,9 +1167,9 @@ -- -- @since 0.1.0 divModA - :: (Source r1 ix e, Source r2 ix e, Integral e) + :: (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 . liftArray2Matching divMod arr1 +divModA arr1 = A.unzip . liftArray2' divMod arr1 {-# INLINE divModA #-} @@ -1165,9 +1179,7 @@ -- > truncateA arr == map truncate arr -- -- @since 0.1.0 -truncateA - :: (Source r ix a, RealFrac a, Integral e) - => Array r ix a -> Array D ix e +truncateA :: (Index ix, Source r a, RealFrac a, Integral e) => Array r ix a -> Array D ix e truncateA = A.map truncate {-# INLINE truncateA #-} @@ -1177,7 +1189,7 @@ -- > truncateA arr == map truncate arr -- -- @since 0.1.0 -roundA :: (Source r ix a, RealFrac a, Integral e) => Array r ix a -> Array D ix e +roundA :: (Index ix, Source r a, RealFrac a, Integral e) => Array r ix a -> Array D ix e roundA = A.map round {-# INLINE roundA #-} @@ -1187,7 +1199,7 @@ -- > truncateA arr == map truncate arr -- -- @since 0.1.0 -ceilingA :: (Source r ix a, RealFrac a, Integral e) => Array r ix a -> Array D ix e +ceilingA :: (Index ix, Source r a, RealFrac a, Integral e) => Array r ix a -> Array D ix e ceilingA = A.map ceiling {-# INLINE ceilingA #-} @@ -1197,7 +1209,7 @@ -- > truncateA arr == map truncate arr -- -- @since 0.1.0 -floorA :: (Source r ix a, RealFrac a, Integral e) => Array r ix a -> Array D ix e +floorA :: (Index ix, Source r a, RealFrac a, Integral e) => Array r ix a -> Array D ix e floorA = A.map floor {-# INLINE floorA #-} @@ -1209,9 +1221,151 @@ -- -- @since 0.1.0 atan2A :: - (Load r ix e, Numeric r e, RealFloat e, MonadThrow m) + (Index ix, Numeric r e, RealFloat e, MonadThrow m) => Array r ix e -> Array r ix e -> m (Array r ix e) -atan2A = liftArray2M atan2 +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 @@ -111,7 +111,7 @@ -- | Integrate with a stencil along a particular dimension. integrateWith :: - (Fractional e, StrideLoad DW ix e, Mutable r ix e) + (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 @@ -126,15 +126,15 @@ -- | 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) + (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 M ix e + -> 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 @@ -144,7 +144,7 @@ -- | Use midpoint rule to approximate an integral. midpointRule :: - (Fractional e, StrideLoad DW ix e, Mutable r ix e) + (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 @@ -152,7 +152,7 @@ -> 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 + -> 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 #-} @@ -160,7 +160,7 @@ -- | Use trapezoid rule to approximate an integral. trapezoidRule :: - (Fractional e, StrideLoad DW ix e, Mutable r ix e) + (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 @@ -168,14 +168,14 @@ -> 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 + -> 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) + (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 @@ -184,7 +184,7 @@ -> 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 + -> 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 #-} @@ -260,8 +260,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: -- @@ -282,7 +283,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 @@ -305,7 +306,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 @@ -318,14 +319,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 @@ -37,6 +37,8 @@ , iunfoldrS_ --, iunfoldrS -- *** Random + , uniformArray + , uniformRangeArray , randomArray , randomArrayS , randomArrayWS @@ -74,6 +76,7 @@ import Data.Massiv.Array.Mutable import Data.Massiv.Core.Common import Prelude hiding (enumFromTo, replicate) +import System.Random.Stateful -- | 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,28 +97,28 @@ -- ] -- -- @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 #-} 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 #-} @@ -127,7 +130,7 @@ -- -- @since 0.2.6 makeArrayA :: - forall r ix e f. (Mutable r ix e, Applicative f) + forall r ix e f. (Manifest r e, Index ix, Applicative f) => Sz ix -> (ix -> f e) -> f (Array r ix e) @@ -147,7 +150,7 @@ -- -- @since 0.4.5 makeArrayLinearA :: - forall r ix e f. (Mutable r ix e, Applicative f) + forall r ix e f. (Manifest r e, Index ix, Applicative f) => Sz ix -> (Int -> f e) -> f (Array r ix e) @@ -165,7 +168,7 @@ -- -- @since 0.2.6 makeArrayAR :: - forall r ix e f. (Mutable r ix e, Applicative f) + forall r ix e f. (Manifest r e, Index ix, Applicative f) => r -> Sz ix -> (ix -> f e) @@ -208,7 +211,12 @@ -- [ 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_ #-} @@ -216,15 +224,14 @@ -- -- @since 0.3.0 iunfoldrS_ :: - forall ix e a. Construct DL ix e + 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 :: Monad m => - Scheduler m () -> Ix1 -> (Ix1 -> e -> m ()) -> (Ix1 -> Sz1 -> e -> m ()) -> m () + load :: Loader e load _ startAt dlWrite _ = void $ loopM startAt (< totalElem sz + startAt) (+ 1) acc0 $ \ !i !acc -> @@ -239,7 +246,7 @@ -- `Data.Massiv.Array.Mutable.unfoldlPrimM` to achive 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_ #-} @@ -247,15 +254,14 @@ -- -- @since 0.3.0 iunfoldlS_ :: - forall ix e a. Construct DL ix e + 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 :: Monad m => - Scheduler m () -> Ix1 -> (Ix1 -> e -> m ()) -> (Ix1 -> Sz1 -> e -> m ()) -> m () + load :: Loader e load _ startAt dlWrite _ = void $ loopDeepM startAt (< totalElem sz + startAt) (+ 1) acc0 $ \ !i !acc -> @@ -286,20 +292,22 @@ -- ] -- -- >>> 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 +-- @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. @@ -308,7 +316,7 @@ randomArray gen splitGen nextRandom comp sz = unsafeMakeLoadArray comp sz Nothing load where !totalLength = totalElem sz - load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m () + 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 @@ -329,6 +337,31 @@ void $ loopM slackStartAt (< totalLength + startAt) (+ 1) genForSlack writeRandom {-# INLINE randomArray #-} + +-- | 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 + +-- | 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) + + -- | 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. @@ -358,13 +391,13 @@ -- >>> 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 + 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)) @@ -385,27 +418,32 @@ -- -- ==== __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 + 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) @@ -474,12 +512,13 @@ -- *** 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 (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 = @@ -499,8 +538,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. @@ -512,7 +551,7 @@ {-# 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) @@ -522,8 +561,8 @@ -- | 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' #-} @@ -568,7 +607,7 @@ -- __/Similar/__: -- -- [@Prelude.`Prelude.enumFromTo`@] Very similar to @[i .. i + n - 1]@, except that --- `senumFromN` is faster, but it only works for `Num` and not for `Enum` elements +-- `enumFromN` is faster, but it only works for `Num` and not for `Enum` elements -- -- [@Data.Vector.Generic.`Data.Vector.Generic.enumFromN`@] -- @@ -662,7 +701,7 @@ -- -- @since 0.2.6 expandWithin :: - forall ix e r n a. (IsIndexDimension ix n, Manifest r (Lower ix) a) + forall n ix e r a. (IsIndexDimension ix n, Index (Lower ix), Manifest r a) => Dimension n -> Sz1 -> (a -> Ix1 -> e) @@ -681,22 +720,22 @@ -- will throw an exception on an invalid dimension. -- -- @since 0.2.6 -expandWithin' - :: (Index ix, Manifest r (Lower ix) a) +expandWithin' :: + 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 -- will throw an exception on an invalid dimension. -- -- @since 0.4.0 -expandWithinM - :: (Index ix, Manifest r (Lower ix) a, MonadThrow m) +expandWithinM :: + forall r ix a b m. (Index ix, Index (Lower ix), Manifest r a, MonadThrow m) => Dim -> Sz1 -> (a -> Ix1 -> b) @@ -713,8 +752,8 @@ -- | Similar to `expandWithin`, except it uses the outermost dimension. -- -- @since 0.2.6 -expandOuter - :: (Index ix, Manifest r (Lower ix) a) +expandOuter :: + forall r ix a b. (Index ix, Index (Lower ix), Manifest r a) => Sz1 -> (a -> Ix1 -> b) -> Array r (Lower ix) a @@ -731,8 +770,8 @@ -- | Similar to `expandWithin`, except it uses the innermost dimension. -- -- @since 0.2.6 -expandInner - :: (Index ix, Manifest r (Lower ix) a) +expandInner :: + 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 @@ -112,7 +112,7 @@ -- -- @since 0.2.4 ifoldMono :: - (Source r ix e, Monoid m) + (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 @@ -124,7 +124,7 @@ -- -- @since 0.2.4 ifoldSemi :: - (Source r ix e, Semigroup m) + (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 @@ -137,7 +137,7 @@ -- -- @since 0.1.6 foldSemi :: - (Source r ix e, Semigroup m) + (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 @@ -149,7 +149,7 @@ -- | 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) => +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 #-} @@ -175,7 +175,7 @@ -- [ [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) => +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 #-} @@ -184,7 +184,7 @@ -- | 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) => +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 #-} @@ -193,7 +193,7 @@ -- | Right fold along a specified dimension. -- -- @since 0.2.4 -foldrWithin :: (Index (Lower ix), IsIndexDimension ix n, Source r ix e) => +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 #-} @@ -203,7 +203,7 @@ -- will throw an exception on an invalid dimension. -- -- @since 0.2.4 -ifoldlWithin' :: (Index (Lower ix), Source r ix e) => +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 -> @@ -224,7 +224,7 @@ -- throw an exception on an invalid dimension. -- -- @since 0.2.4 -foldlWithin' :: (Index (Lower ix), Source r ix e) => +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' #-} @@ -235,7 +235,7 @@ -- -- -- @since 0.2.4 -ifoldrWithin' :: (Index (Lower ix), Source r ix e) => +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 -> @@ -255,7 +255,7 @@ -- will throw an exception on an invalid dimension. -- -- @since 0.2.4 -foldrWithin' :: (Index (Lower ix), Source r ix e) => +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' #-} @@ -264,7 +264,7 @@ -- | Left fold over the inner most dimension with index aware function. -- -- @since 0.2.4 -ifoldlInner :: (Index (Lower ix), Source r ix e) => +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 #-} @@ -272,7 +272,7 @@ -- | Left fold over the inner most dimension. -- -- @since 0.2.4 -foldlInner :: (Index (Lower ix), Source r ix e) => +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 #-} @@ -280,7 +280,7 @@ -- | Right fold over the inner most dimension with index aware function. -- -- @since 0.2.4 -ifoldrInner :: (Index (Lower ix), Source r ix e) => +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 #-} @@ -288,7 +288,7 @@ -- | Right fold over the inner most dimension. -- -- @since 0.2.4 -foldrInner :: (Index (Lower ix), Source r ix e) => +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 #-} @@ -296,7 +296,7 @@ -- | 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 #-} @@ -304,7 +304,7 @@ -- -- @since 0.4.3 foldWithin :: - (Source r ix a, Monoid a, Index (Lower ix), IsIndexDimension ix n) + (Source r a, Monoid a, Index (Lower ix), IsIndexDimension ix n) => Dimension n -> Array r ix a -> Array D (Lower ix) a @@ -316,7 +316,7 @@ -- -- @since 0.4.3 foldWithin' :: - (Source r ix a, Monoid a, Index (Lower ix)) + (HasCallStack, Index ix, Source r a, Monoid a, Index (Lower ix)) => Dim -> Array r ix a -> Array D (Lower ix) a @@ -342,7 +342,11 @@ -- 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 #-} @@ -351,10 +355,15 @@ -- 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 #-} @@ -377,7 +386,8 @@ -- 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 #-} @@ -386,50 +396,54 @@ -- 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' #-} @@ -437,7 +451,7 @@ -- -- -- -- @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 @@ -446,7 +460,7 @@ -- | /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 #-} @@ -454,7 +468,7 @@ -- | /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 #-} @@ -462,7 +476,7 @@ -- | /O(n)/ - Compute conjunction of all elements. -- -- @since 0.1.0 -and :: Source r ix Bool => Array r ix Bool -> Bool +and :: (Index ix, Source r Bool) => Array r ix Bool -> Bool and = all id {-# INLINE and #-} @@ -470,7 +484,7 @@ -- | /O(n)/ - Compute disjunction of all elements. -- -- @since 0.1.0 -or :: Source r ix Bool => Array r ix Bool -> Bool +or :: (Index ix, Source r Bool) => Array r ix Bool -> Bool or = any id {-# INLINE or #-} @@ -478,14 +492,14 @@ -- | /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 :: (Index ix, Source r e) => (e -> Bool) -> Array r ix e -> Bool all f = not . any (not . f) {-# INLINE all #-} -- | /O(n)/ - Determines whether an element is present in the array. -- -- @since 0.5.5 -elem :: (Eq e, Source r ix e) => e -> Array r ix e -> Bool +elem :: (Eq e, Index ix, Source r e) => e -> Array r ix e -> Bool elem e = any (e ==) {-# INLINE elem #-} 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,6 +2,7 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UndecidableInstances #-} -- | -- Module : Data.Massiv.Array.Ops.Fold.Internal @@ -41,7 +42,7 @@ , ifoldrP , ifoldlIO , ifoldrIO - -- , splitReduce + , splitReduce , any , anySu , anyPu @@ -60,7 +61,7 @@ -- -- @since 0.3.0 fold :: - (Monoid e, Source r ix e) + (Monoid e, Index ix, Source r e) => Array r ix e -- ^ Source array -> e fold = foldlInternal mappend mempty mappend mempty @@ -73,7 +74,7 @@ -- -- @since 0.1.4 foldMono :: - (Source r ix e, Monoid m) + (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 @@ -84,7 +85,7 @@ -- | /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 :: (Index ix, Source r e, Monad m) => (a -> e -> m a) -> a -> Array r ix e -> m a foldlM f = ifoldlM (\ a _ b -> f a b) {-# INLINE foldlM #-} @@ -92,7 +93,7 @@ -- | /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_ :: (Index ix, Source r e, Monad m) => (a -> e -> m a) -> a -> Array r ix e -> m () foldlM_ f = ifoldlM_ (\ a _ b -> f a b) {-# INLINE foldlM_ #-} @@ -100,7 +101,7 @@ -- | /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) {-# INLINE ifoldlM #-} @@ -109,7 +110,7 @@ -- | /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_ #-} @@ -117,7 +118,7 @@ -- | /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 :: (Index ix, Source r e, Monad m) => (e -> a -> m a) -> a -> Array r ix e -> m a foldrM f = ifoldrM (\_ e a -> f e a) {-# INLINE foldrM #-} @@ -125,7 +126,7 @@ -- | /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_ #-} @@ -133,7 +134,7 @@ -- | /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 @@ -143,7 +144,7 @@ -- | /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_ #-} @@ -152,7 +153,7 @@ -- | /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 :: (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) @@ -165,7 +166,7 @@ -- | /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 #-} @@ -173,7 +174,7 @@ -- | /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 :: (Index ix, Source r e) => (a -> e -> a) -> a -> Array r ix e -> a foldlS f = ifoldlS (\ a _ e -> f a e) {-# INLINE foldlS #-} @@ -181,7 +182,7 @@ -- | /O(n)/ - Left fold with an index aware function, computed sequentially. -- -- @since 0.1.0 -ifoldlS :: Source r ix e +ifoldlS :: (Index ix, Source r 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 {-# INLINE ifoldlS #-} @@ -190,7 +191,7 @@ -- | /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 :: (Index ix, Source r e) => (e -> a -> a) -> a -> Array r ix e -> a foldrS f = ifoldrS (\_ e a -> f e a) {-# INLINE foldrS #-} @@ -198,7 +199,7 @@ -- | /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 :: (Index ix, Source r 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 {-# INLINE ifoldrS #-} @@ -206,7 +207,7 @@ -- | 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) @@ -235,7 +236,7 @@ -- [1,0,3,2,5,4] -- -- @since 0.1.0 -foldlP :: (MonadIO m, Source r ix e) => +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@. @@ -249,7 +250,7 @@ -- element it is being applied to. -- -- @since 0.1.0 -ifoldlP :: (MonadIO m, Source r ix e) => +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 @@ -270,7 +271,7 @@ -- [[0,1],[2,3],[4,5]] -- -- @since 0.1.0 -foldrP :: (MonadIO m, Source r ix e) => +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 #-} @@ -282,7 +283,7 @@ -- -- @since 0.1.0 ifoldrP :: - (MonadIO m, Source r ix e) + (MonadIO m, Index ix, Source r e) => (ix -> e -> a -> a) -> a -> (a -> b -> b) @@ -295,12 +296,12 @@ -- | 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 #-} @@ -309,7 +310,7 @@ -- -- @since 0.1.0 ifoldlIO :: - (MonadUnliftIO m, Source r ix e) + (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 @@ -322,53 +323,54 @@ 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 + loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start -> + scheduleWork scheduler $ run $ + iterLinearM sz start (start + chunkLength) 1 (<) initAcc $ \ !i ix !acc -> + f acc ix (unsafeLinearIndex arr i) + when (slackStart < totalLength) $ + scheduleWork scheduler $ run $ + iterLinearM sz slackStart totalLength 1 (<) initAcc $ \ !i ix !acc -> + f acc ix (unsafeLinearIndex arr i) 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.6.0 --- splitReduce :: --- (MonadUnliftIO m, Source r ix e) --- => (Scheduler m a -> BatchId -> 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 -> do --- batchId <- getCurrentBatchId scheduler --- splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do --- loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start -> --- scheduleWork scheduler $ f scheduler batchId $ --- unsafeLinearSlice start (SafeSz chunkLength) arr --- when (slackStart < totalLength) $ --- scheduleWork scheduler $ f scheduler batchId $ --- unsafeLinearSlice slackStart (SafeSz (totalLength - slackStart)) arr --- F.foldlM g tAcc results --- {-# INLINE splitReduce #-} +-- | 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 -> do + withRunInIO $ \run -> do + splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do + loopM_ 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 splitReduce #-} -- | Similar to `ifoldrP`, except that folding functions themselves do live in IO -- -- @since 0.1.0 -ifoldrIO :: (MonadUnliftIO m, Source r ix e) => +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) @@ -376,21 +378,22 @@ let !sz = size arr !totalLength = totalElem sz results <- - withScheduler (getComp arr) $ \ scheduler -> - splitLinearly (numWorkers scheduler) totalLength $ \ chunkLength slackStart -> do - 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 -> + 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 !acc -> f ix (unsafeLinearIndex arr i) acc + loopM_ slackStart (> 0) (subtract chunkLength) $ \ !start -> + scheduleWork scheduler $ run $ + 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 #-} -- | Sequential implementation of `any` with unrolling -anySu :: Source r ix a => (a -> Bool) -> Array r ix a -> Bool +anySu :: (Index ix, Source r e) => (e -> Bool) -> Array r ix e -> Bool anySu f arr = go 0 where !k = elemsCount arr @@ -409,14 +412,14 @@ -- | Implementaton of `any` on a slice of an array with short-circuiting using batch cancellation. anySliceSuM :: - Source r ix a - => Batch IO Bool + (Index ix, Source r e) + => Batch RealWorld Bool -> Ix1 -> Sz1 - -> (a -> Bool) - -> Array r ix a + -> (e -> Bool) + -> Array r ix e -> IO Bool -anySliceSuM batch ix0 (Sz k) f arr = go ix0 +anySliceSuM batch ix0 (Sz1 k) f arr = go ix0 where !k' = k - ix0 !k4 = ix0 + (k' - (k' `rem` 4)) @@ -444,7 +447,7 @@ -- | Parallelizable implementation of `any` with unrolling -anyPu :: Source r ix e => (e -> Bool) -> Array r ix e -> IO Bool +anyPu :: (Index ix, Source r e) => (e -> Bool) -> Array r ix e -> IO Bool anyPu f arr = do let !sz = size arr !totalLength = totalElem sz @@ -464,7 +467,7 @@ -- | /O(n)/ - Determines whether any element of the array satisfies a predicate. -- -- @since 0.1.0 -any :: Source r ix e => (e -> Bool) -> Array r ix e -> Bool +any :: (Index ix, Source r e) => (e -> Bool) -> Array r ix e -> Bool any f arr = case getComp arr of Seq -> anySu f arr 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 @@ -2,6 +2,7 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE MonoLocalBinds #-} -- | -- Module : Data.Massiv.Array.Ops.Map -- Copyright : (c) Alexey Kuleshevich 2018-2021 @@ -49,6 +50,9 @@ , iforIO_ , imapSchedulerM_ , iforSchedulerM_ + , iterArrayLinearM_ + , iterArrayLinearWithSetM_ + , iterArrayLinearWithStrideM_ -- ** Zipping , zip , zip3 @@ -62,7 +66,6 @@ , izipWith , izipWith3 , izipWith4 - , liftArray2 -- *** Applicative , zipWithA , izipWithA @@ -71,6 +74,7 @@ ) where import Control.Monad (void) +import Control.Monad.Primitive import Control.Scheduler import Data.Coerce import Data.Massiv.Array.Delayed.Pull @@ -85,7 +89,9 @@ -------------------------------------------------------------------------------- -- | Map a function over an array -map :: Source r ix e' => (e' -> e) -> Array r ix e' -> Array D ix e +-- +-- @since 0.1.0 +map :: (Index ix, Source r e') => (e' -> e) -> Array r ix e' -> Array D ix e map f = imap (const f) {-# INLINE map #-} @@ -95,13 +101,17 @@ -------------------------------------------------------------------------------- -- | Zip two arrays -zip :: (Source r1 ix e1, Source r2 ix 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) +-- +-- @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 #-} @@ -110,7 +120,7 @@ -- -- @since 0.5.4 zip4 :: - (Source r1 ix e1, Source r2 ix e2, Source r3 ix e3, Source r4 ix e4) + (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 @@ -120,12 +130,16 @@ {-# 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) +-- +-- @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 #-} @@ -133,7 +147,7 @@ -- | Unzip four arrays -- -- @since 0.5.4 -unzip4 :: Source r ix (e1, e2, e3, e4) +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 @@ -148,14 +162,14 @@ -- | 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) +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 = izipWith (\ _ e1 e2 -> f e1 e2) {-# INLINE zipWith #-} -- | Just like `zipWith`, except with an index aware function. -izipWith :: (Source r1 ix e1, Source r2 ix e2) +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 @@ -166,7 +180,7 @@ -- | Just like `zipWith`, except zip three arrays with a function. -zipWith3 :: (Source r1 ix e1, Source r2 ix e2, Source r3 ix 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 = izipWith3 (\ _ e1 e2 e3 -> f e1 e2 e3) {-# INLINE zipWith3 #-} @@ -174,7 +188,7 @@ -- | 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 @@ -197,7 +211,7 @@ -- -- @since 0.5.4 zipWith4 :: - (Source r1 ix e1, Source r2 ix e2, Source r3 ix e3, Source r4 ix e4) + (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 @@ -212,7 +226,7 @@ -- -- @since 0.5.4 izipWith4 - :: (Source r1 ix e1, Source r2 ix e2, Source r3 ix e3, Source r4 ix e4) + :: (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 @@ -235,11 +249,11 @@ -- | Similar to `zipWith`, except does it sequentially and using the `Applicative`. Note that --- resulting array has Mutable representation. +-- resulting array has Manifest representation. -- -- @since 0.3.0 zipWithA :: - (Source r1 ix e1, Source r2 ix e2, Applicative f, Mutable r ix e) + (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 @@ -248,11 +262,11 @@ {-# INLINE zipWithA #-} -- | Similar to `zipWith`, except does it sequentiall and using the `Applicative`. Note that --- resulting array has Mutable representation. +-- resulting array has Manifest representation. -- -- @since 0.3.0 izipWithA :: - (Source r1 ix e1, Source r2 ix e2, Applicative f, Mutable r ix e) + (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 @@ -268,7 +282,7 @@ -- -- @since 0.3.0 zipWith3A :: - (Source r1 ix e1, Source r2 ix e2, Source r3 ix e3, Applicative f, Mutable r ix e) + (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 @@ -281,7 +295,7 @@ -- -- @since 0.3.0 izipWith3A :: - (Source r1 ix e1, Source r2 ix e2, Source r3 ix e3, Applicative f, Mutable r ix e) + (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 @@ -297,27 +311,6 @@ {-# 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 -------------------------------------------------------------------- -------------------------------------------------------------------------------- @@ -329,7 +322,7 @@ -- @since 0.2.6 -- traverseA :: - forall r ix e r' a f . (Source r' ix a, Mutable r ix e, Applicative f) + 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) @@ -340,7 +333,11 @@ -- -- @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_ :: + forall r ix e a f. (Index ix, Source r e, Applicative f) + => (e -> f a) + -> Array r ix e + -> f () traverseA_ f arr = loopA_ 0 (< totalElem (size arr)) (+ 1) (f . unsafeLinearIndex arr) {-# INLINE traverseA_ #-} @@ -349,7 +346,7 @@ -- @since 0.3.0 -- sequenceA :: - forall r ix e r' f. (Source r' ix (f e), Mutable r ix e, Applicative f) + 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 @@ -359,7 +356,10 @@ -- -- @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_ #-} @@ -369,7 +369,7 @@ -- @since 0.2.6 -- itraverseA :: - forall r ix e r' a f . (Source r' ix a, Mutable r ix e, Applicative f) + 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) @@ -383,7 +383,7 @@ -- @since 0.2.6 -- itraverseA_ :: - forall r ix e a f. (Source r ix a, Applicative f) + forall r ix e a f. (Source r a, Index ix, Applicative f) => (ix -> a -> f e) -> Array r ix a -> f () @@ -399,7 +399,7 @@ -- @since 0.3.0 -- traversePrim :: - forall r ix b r' a m . (Source r' ix a, Mutable r ix b, PrimMonad m) + 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) @@ -411,7 +411,7 @@ -- @since 0.3.0 -- itraversePrim :: - forall r ix b r' a m . (Source r' ix a, Mutable r ix b, PrimMonad m) + 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) @@ -432,7 +432,7 @@ -- -- @since 0.2.6 mapM :: - forall r ix b r' a m. (Source r' ix a, Mutable r ix b, Monad m) + 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) @@ -444,7 +444,7 @@ -- -- @since 0.2.6 forM :: - forall r ix b r' a m. (Source r' ix a, Mutable r ix b, Monad m) + 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) @@ -456,7 +456,7 @@ -- -- @since 0.2.6 imapM :: - forall r ix b r' a m. (Source r' ix a, Mutable r ix b, Monad m) + 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) @@ -468,7 +468,7 @@ -- -- @since 0.5.1 iforM :: - forall r ix b r' a m. (Source r' ix a, Mutable r ix b, Monad m) + 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) -> m (Array r ix b) @@ -489,7 +489,7 @@ -- 58 -- -- @since 0.1.0 -mapM_ :: (Source r ix a, Monad m) => (a -> m b) -> Array r ix a -> m () +mapM_ :: (Source r a, Index ix, Monad m) => (a -> m b) -> Array r ix a -> m () mapM_ f !arr = iterM_ zeroIndex (unSz (size arr)) (pureIndex 1) (<) (f . unsafeIndex arr) {-# INLINE mapM_ #-} @@ -508,13 +508,13 @@ -- >>> readIORef ref -- 499500 -- -forM_ :: (Source r ix a, Monad m) => Array r ix a -> (a -> m b) -> m () +forM_ :: (Source r a, Index ix, Monad m) => Array r ix a -> (a -> m b) -> m () forM_ = flip mapM_ {-# INLINE forM_ #-} -- | Just like `imapM_`, except with flipped arguments. -iforM_ :: (Source r ix a, Monad m) => Array r ix a -> (ix -> a -> m b) -> m () +iforM_ :: (Source r a, Index ix, Monad m) => Array r ix a -> (ix -> a -> m b) -> m () iforM_ = flip imapM_ {-# INLINE iforM_ #-} @@ -525,69 +525,72 @@ -- -- @since 0.2.6 mapIO :: - forall r ix b r' a m. (Source r' ix a, Mutable r ix b, MonadUnliftIO m, PrimMonad m) + 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 + -- 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. + in 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. 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) + 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) + 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) @@ -596,13 +599,13 @@ --- | 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) + 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 @@ -614,7 +617,7 @@ -- -- @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) + 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 @@ -627,7 +630,7 @@ -- -- @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) + 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) @@ -639,7 +642,7 @@ -- -- @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) + 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) @@ -664,7 +667,7 @@ -- 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_ #-} @@ -672,7 +675,7 @@ -- -- @since 0.2.6 iforIO :: - forall r ix b r' a m. (Source r' ix a, Mutable r ix b, MonadUnliftIO m, PrimMonad m) + 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) @@ -682,6 +685,127 @@ -- | 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,4 +1,5 @@ {-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} -- | @@ -62,25 +63,28 @@ -- ] -- ] -- >>> 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 (!>) #-} @@ -88,12 +92,15 @@ -- `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 (!?>) #-} @@ -104,15 +111,19 @@ -- -- >>> 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 (??>) #-} @@ -120,30 +131,38 @@ -- | /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 @@ -161,7 +184,12 @@ internalInnerSlice :: - (MonadThrow m, Slice r ix e) => Dim -> Sz ix -> Array r ix e -> Int -> m (Elt r ix e) + (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 @@ -176,18 +204,23 @@ -- 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 () #-} @@ -205,15 +238,13 @@ -- [ 2 :. 0, 2 :. 1 ] -- -- @since 0.5.4 -outerSlices :: OuterSlice r ix e => Array r ix e -> Array D Ix1 (Elt r ix e) -outerSlices arr = makeArray Seq k (unsafeOuterSlice arr) +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, _) = unconsSz $ size arr --- TODO: move setComp to Load --- outerSlices arr = makeArray (getComp arr) k (unsafeOuterSlice arr') --- where --- arr' = setComp Seq arr --- (k, _) = unconsSz $ size arr + (k, szL) = unconsSz $ size arr {-# INLINE outerSlices #-} @@ -229,15 +260,13 @@ -- [ 0 :. 1, 1 :. 1, 2 :. 1 ] -- -- @since 0.5.4 -innerSlices :: InnerSlice r ix e => Array r ix e -> Array D Ix1 (Elt r ix e) -innerSlices arr = makeArray Seq k (unsafeInnerSlice arr sz) +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 - sz@(_, k) = unsnocSz $ size arr --- TODO: move setComp to Load --- innerSlices arr = makeArray (getComp arr) k (unsafeInnerSlice arr' sz) --- where --- arr' = setComp Seq arr --- sz@(_, k) = unsnocSz $ size arr + (szL, k) = unsnocSz $ size arr {-# INLINE innerSlices #-} -- | Create a delayed array of slices from within. Checks dimension at compile time. @@ -287,10 +316,10 @@ -- -- @since 0.5.4 withinSlices :: - (IsIndexDimension ix n, Slice r ix e) + forall n r ix e. (IsIndexDimension ix n, Index (Lower ix), Source r e) => Dimension n -> Array r ix e - -> Array D Ix1 (Elt r ix e) + -> Array D Ix1 (Array D (Lower ix) e) withinSlices dim = either throwImpossible id . withinSlicesM (fromDimension dim) {-# INLINE withinSlices #-} @@ -301,7 +330,11 @@ -- /__Throws Exceptions__/: `IndexDimensionException` -- -- @since 0.5.4 -withinSlicesM :: (MonadThrow m, Slice r ix e) => Dim -> Array r ix e -> m (Array D Ix1 (Elt r ix e)) +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 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,6 +1,7 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE ExplicitForAll #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MonoLocalBinds #-} -- | -- Module : Data.Massiv.Array.Ops.Sort -- Copyright : (c) Alexey Kuleshevich 2018-2021 @@ -21,6 +22,7 @@ import Control.Monad.IO.Unlift import Control.Monad (when) +import Control.Monad.Primitive import Control.Scheduler import Data.Massiv.Array.Delayed.Stream import Data.Massiv.Array.Mutable @@ -44,10 +46,10 @@ -- [ (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 -> Vector DS (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 = scatMaybes $ sunfoldrN (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) @@ -62,14 +64,18 @@ {-# INLINE tally #-} -unsafeUnstablePartitionRegionM' :: - forall r e m. (Mutable r Ix1 e, PrimMonad m) - => MArray (PrimState m) r Ix1 e + +-- | Partition a segment of a vector. Starting and ending indices are unchecked. +-- +-- @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) +unsafeUnstablePartitionRegionM marr f start end = fromLeft start (end + 1) where fromLeft i j | i == j = pure i @@ -89,21 +95,6 @@ unsafeLinearWrite marr i x fromLeft (i + 1) j else fromRight i (j - 1) -{-# INLINE unsafeUnstablePartitionRegionM' #-} - - --- TODO: Replace `unsafeUnstablePartitionRegionM` with `unsafeUnstablePartitionRegionM'` --- | 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) - => MVector (PrimState m) r e - -> (e -> Bool) - -> Ix1 -- ^ Start index of the region - -> Ix1 -- ^ End index of the region - -> m Ix1 -unsafeUnstablePartitionRegionM marr f = unsafeUnstablePartitionRegionM' marr (pure . f) {-# INLINE unsafeUnstablePartitionRegionM #-} @@ -115,7 +106,7 @@ -- -- @since 0.3.2 quicksort :: - (Mutable r Ix1 e, Ord e) => Array r Ix1 e -> Array r Ix1 e + (Manifest r e, Ord e) => Vector r e -> Vector r e quicksort arr = unsafePerformIO $ withMArray_ arr quicksortM_ {-# INLINE quicksort #-} @@ -124,26 +115,25 @@ -- -- @since 0.6.1 quicksortByM :: - (Mutable r Ix1 e, MonadUnliftIO m) => (e -> e -> m Ordering) -> Vector r e -> m (Vector r e) + (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 :: - (Mutable r Ix1 e) => (e -> e -> Ordering) -> Vector r e -> Vector r e +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 () - -> MVector (PrimState m) r e + (Ord e, Manifest r e, MonadPrimBase s m) + => Scheduler s () + -> MVector s r e -> m () quicksortM_ = quicksortInternalM_ (\e1 e2 -> pure $ e1 < e2) (\e1 e2 -> pure $ e1 == e2) {-# INLINE quicksortM_ #-} @@ -153,10 +143,10 @@ -- -- @since 0.6.1 quicksortByM_ :: - (Mutable r Ix1 e, PrimMonad m) + (Manifest r e, MonadPrimBase s m) => (e -> e -> m Ordering) - -> Scheduler m () - -> MVector (PrimState m) r e + -> Scheduler s () + -> MVector s r e -> m () quicksortByM_ compareM = quicksortInternalM_ (\x y -> (LT ==) <$> compareM x y) (\x y -> (EQ ==) <$> compareM x y) @@ -164,14 +154,14 @@ quicksortInternalM_ :: - (Mutable r Ix1 e, PrimMonad m) + (Manifest r e, MonadPrimBase s m) => (e -> e -> m Bool) -> (e -> e -> m Bool) - -> Scheduler m () - -> MVector (PrimState m) r e + -> Scheduler s () + -> MVector s r e -> m () quicksortInternalM_ fLT fEQ scheduler marr = - scheduleWork scheduler $ qsort (numWorkers scheduler) 0 (unSz (msize marr) - 1) + scheduleWork scheduler $ qsort (numWorkers scheduler) 0 (unSz (sizeOfMArray marr) - 1) where ltSwap i j = do ei <- unsafeLinearRead marr i @@ -193,8 +183,8 @@ qsort !n !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 + l <- unsafeUnstablePartitionRegionM marr (`fLT` p) lo (hi - 1) + h <- unsafeUnstablePartitionRegionM marr (`fEQ` p) l hi if n > 0 then do let !n' = n - 1 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 @@ -64,27 +64,73 @@ , transform2' ) where +import Control.Monad as M (foldM_, forM_, unless) +import Control.Monad.ST import Control.Scheduler (traverse_) -import Control.Monad as M (foldM_, unless, forM_) import Data.Bifunctor (bimap) -import Data.Foldable as F (foldl', foldrM, toList, length) +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.Common -import Prelude as P hiding (concat, splitAt, traverse, mapM_, reverse, take, drop) +import Data.Massiv.Core +import Data.Massiv.Core.Index.Internal +import Data.Massiv.Core.Common (size, unsafeIndex, unsafeResize, evaluate', evaluateM) +import Data.Proxy +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. -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) +-- +-- ====__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 @@ -99,45 +145,74 @@ -- 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 +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 :: (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 :: + 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 `extractFromTo`, but throws an error on invalid indices. +-- | Same as `extractFromToM`, 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' :: + 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)/ - Changes the shape of an array. Returns `Nothing` if total --- number of elements does not match the source array. +-- | /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 :: - (MonadThrow m, Index ix', Load r ix e, Resize r ix) + 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) @@ -147,14 +222,18 @@ -- | 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 +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 :: (Load r ix e, Resize r ix) => Array r ix e -> Array r Ix1 e +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 #-} @@ -178,7 +257,7 @@ -- ] -- -- @since 0.1.0 -transpose :: Source r Ix2 e => Array r Ix2 e -> Array D Ix2 e +transpose :: forall r e. Source r e => Matrix r e -> Matrix D e transpose = transposeInner {-# INLINE [1] transpose #-} @@ -220,8 +299,10 @@ -- ] -- -- @since 0.1.0 -transposeInner :: (Index (Lower ix), Source r' ix e) - => Array r' ix e -> Array D ix e +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 = @@ -271,8 +352,10 @@ -- -- -- @since 0.1.0 -transposeOuter :: (Index (Lower ix), Source r' ix e) - => Array r' ix e -> Array D ix e +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 = @@ -316,7 +399,11 @@ -- ] -- -- @since 0.4.1 -reverse :: (IsIndexDimension ix n, Source r ix e) => Dimension n -> Array r ix e -> Array D ix e +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 #-} @@ -324,7 +411,11 @@ -- `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 :: + 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 @@ -336,8 +427,12 @@ -- `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 +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 @@ -370,7 +465,7 @@ -- @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) + (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. @@ -385,11 +480,12 @@ -- * 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' :: + 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' #-} @@ -431,7 +527,7 @@ -- -- @since 0.3.0 appendM :: - forall r1 r2 ix e m. (MonadThrow m, Source r1 ix e, Source r2 ix e) + 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 @@ -444,8 +540,7 @@ unless (szl1 == szl2) $ throwM $ SizeMismatchException sz1 sz2 let !k1' = unSz k1 newSz <- insertSzM szl1 n (SafeSz (k1' + unSz k2)) - let load :: Monad n => - Scheduler n () -> Ix1 -> (Ix1 -> e -> n ()) -> (Ix1 -> Sz1 -> e -> n ()) -> n () + let load :: Loader e load scheduler !startAt dlWrite _dlSet = do scheduleWork scheduler $ iterM_ zeroIndex (unSz sz1) (pureIndex 1) (<) $ \ix -> @@ -465,16 +560,24 @@ -- | 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 +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' :: (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 +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 @@ -484,11 +587,11 @@ -- -- @since 0.3.0 concatM :: - forall r ix e f m. (MonadThrow m, Foldable f, Source r ix e) + 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 = +concatM n arrsF = case L.uncons (F.toList arrsF) of Nothing -> pure empty Just (a, arrs) -> do @@ -504,20 +607,20 @@ (dropWhile ((== szl) . snd) $ P.zip szs szls) let kTotal = SafeSz $ F.foldl' (+) k ks newSz <- insertSzM (SafeSz szl) n kTotal - let load :: Monad n => - Scheduler n () -> Ix1 -> (Ix1 -> e -> n ()) -> (Ix1 -> Sz1 -> e -> n ()) -> n () + let load :: Loader e load scheduler startAt dlWrite _dlSet = - let arrayLoader !kAcc (kCur, arr) = do + let arrayLoader !kAcc (!kCur, arr) = do scheduleWork scheduler $ - iforM_ arr $ \ix e -> - let i = getDim' ix n - ix' = setDim' ix n (i + kAcc) - in dlWrite (startAt + toLinearIndex newSz ix') e - pure (kAcc + kCur) + 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 = foldMap getComp arrsF, dlSize = newSz, dlLoad = load} + DLArray {dlComp = getComp a <> foldMap getComp arrs, dlSize = newSz, dlLoad = load} {-# INLINE concatM #-} @@ -585,7 +688,7 @@ -- -- @since 0.5.4 stackSlicesM :: - forall r ix e f m. (Foldable f, MonadThrow m, Source r (Lower ix) e, Index ix) + 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) @@ -599,8 +702,7 @@ M.forM_ arrsF $ \arr -> unless (sz == size arr) $ throwM (SizeMismatchException sz (size arr)) newSz <- insertSzM sz dim len - let load :: Monad n => - Scheduler n () -> Ix1 -> (Ix1 -> e -> n ()) -> (Ix1 -> Sz1 -> e -> n ()) -> n () + 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) @@ -630,11 +732,11 @@ -- ] -- >>> rows = outerSlices x -- >>> A.mapM_ print rows --- Array M Seq (Sz1 3) +-- Array P Seq (Sz1 3) -- [ 1, 2, 3 ] --- Array M Seq (Sz1 3) +-- Array P Seq (Sz1 3) -- [ 4, 5, 6 ] --- Array M Seq (Sz1 3) +-- Array P Seq (Sz1 3) -- [ 7, 8, 9 ] -- >>> stackOuterSlicesM rows :: IO (Matrix DL Int) -- Array DL Seq (Sz (3 :. 3)) @@ -645,7 +747,7 @@ -- -- @since 0.5.4 stackOuterSlicesM :: - forall r ix e f m. (Foldable f, MonadThrow m, Source r (Lower ix) e, Index ix) + 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)) @@ -670,11 +772,11 @@ -- ] -- >>> columns = innerSlices x -- >>> A.mapM_ print columns --- Array M Seq (Sz1 3) +-- Array D Seq (Sz1 3) -- [ 1, 4, 7 ] --- Array M Seq (Sz1 3) +-- Array D Seq (Sz1 3) -- [ 2, 5, 8 ] --- Array M Seq (Sz1 3) +-- Array D Seq (Sz1 3) -- [ 3, 6, 9 ] -- >>> stackInnerSlicesM columns :: IO (Matrix DL Int) -- Array DL Seq (Sz (3 :. 3)) @@ -685,7 +787,7 @@ -- -- @since 0.5.4 stackInnerSlicesM :: - forall r ix e f m. (Foldable f, MonadThrow m, Source r (Lower ix) e, Index ix) + 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 @@ -700,11 +802,11 @@ -- -- @since 0.3.0 splitAtM :: - (MonadThrow m, Extract r ix e) + 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 (R r) ix e, Array (R r) ix e) + -> m (Array D ix e, Array D ix e) splitAtM dim i arr = do let Sz sz = size arr eIx <- setDimM sz dim i @@ -723,9 +825,13 @@ -- -- -- @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 +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' #-} @@ -733,12 +839,12 @@ -- -- @since 0.3.5 splitExtractM :: - (MonadThrow m, Extract r ix e, Source (R r) ix e) + 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 (R r) ix e, Array (R r) ix e, Array (R 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 @@ -781,21 +887,16 @@ -- -- @since 0.6.1 replaceSlice :: - ( MonadThrow m - , Extract r ix e - , Source (R r) ix e - , Load (R r) (Lower ix) e - , Resize (R r) (Lower ix) - ) + forall r r' ix e m. (MonadThrow m, Source r e, Source r' e, Index ix, Index (Lower ix)) => Dim -> Ix1 - -> Array (R r) (Lower ix) e + -> 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, m', r] + concatM dim [l, delay m', r] {-# INLINE replaceSlice #-} @@ -827,14 +928,9 @@ -- -- @since 0.6.1 replaceOuterSlice :: - ( MonadThrow m - , Extract r ix e - , Source (R r) ix e - , Load (R r) (Lower ix) e - , Resize (R r) (Lower ix) - ) + forall r ix e m. (MonadThrow m, Index ix, Source r e, Load r (Lower ix) e) => Ix1 - -> Array (R r) (Lower ix) e + -> 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 @@ -866,7 +962,7 @@ -- -- @since 0.3.5 deleteRegionM :: - (MonadThrow m, Extract r ix e, Source (R r) ix e) + 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 @@ -898,7 +994,7 @@ -- -- @since 0.3.5 deleteRowsM :: - (MonadThrow m, Extract r ix e, Source (R r) ix e, Index (Lower ix)) + forall r ix e m. (MonadThrow m, Index ix, Index (Lower ix), Source r e) => Ix1 -> Sz Ix1 -> Array r ix e @@ -927,7 +1023,7 @@ -- -- @since 0.3.5 deleteColumnsM :: - (MonadThrow m, Extract r ix e, Source (R r) ix e) + forall r ix e m. (MonadThrow m, Index ix, Source r e) => Ix1 -> Sz Ix1 -> Array r ix e @@ -940,7 +1036,7 @@ -- -- @since 0.3.0 downsample :: - forall r ix e. Source r ix e + forall r ix e. (Source r e, Load r ix e) => Stride ix -> Array r ix e -> Array DL ix e @@ -952,8 +1048,7 @@ unsafeLinearWriteWithStride = unsafeIndex arr . liftIndex2 (*) strideIx . fromLinearIndex resultSize {-# INLINE unsafeLinearWriteWithStride #-} - load :: Monad m => - Scheduler m () -> Ix1 -> (Ix1 -> e -> m ()) -> (Ix1 -> Sz1 -> e -> m ()) -> m () + load :: Loader e load scheduler startAt dlWrite _ = splitLinearlyWithStartAtM_ scheduler @@ -1007,18 +1102,17 @@ , dlLoad = load } where - load :: Monad m => - Scheduler m () -> Ix1 -> (Ix1 -> e -> m ()) -> (Ix1 -> Sz1 -> e -> m ()) -> m () + load :: Loader e load scheduler startAt uWrite uSet = do uSet startAt (toLinearSz newsz) fillWith - loadArrayM scheduler arr (\i -> uWrite (adjustLinearStride (i + startAt))) + 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 = size arr + ~sz = outerSize arr -- intentionally lazy in case it is used with DS !newsz = SafeSz (timesStride $ unSz sz) {-# INLINE upsample #-} @@ -1028,7 +1122,7 @@ -- @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) + (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' @@ -1043,7 +1137,8 @@ -- -- @since 0.3.0 transform' :: - (Source r' ix' e', Index ix) + 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' @@ -1057,7 +1152,16 @@ -- -- @since 0.3.0 transform2M :: - (Mutable r ix e, Source r1 ix1 e1, Source r2 ix2 e2, MonadUnliftIO m, PrimMonad m, MonadThrow m) + ( 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 @@ -1073,7 +1177,7 @@ -- -- @since 0.3.0 transform2' :: - (Source r1 ix1 e1, Source r2 ix2 e2, Index ix) + (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 @@ -1118,7 +1222,7 @@ -- -- @since 0.3.1 zoomWithGrid :: - forall r ix e. Source r ix e + 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 @@ -1128,7 +1232,7 @@ !kx = liftIndex (+ 1) zoomFactor !lastNewIx = liftIndex2 (*) kx $ unSz (size arr) !newSz = Sz (liftIndex (+ 1) lastNewIx) - load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m () + 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 @@ -1172,7 +1276,7 @@ -- -- @since 0.4.4 zoom :: - forall r ix e. Source r ix e + forall r ix e. (Index ix, Source r e) => Stride ix -- ^ Scaling factor -> Array r ix e -- ^ Source array -> Array DL ix e @@ -1180,7 +1284,7 @@ where !lastNewIx = liftIndex2 (*) zoomFactor $ unSz (size arr) !newSz = Sz lastNewIx - load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m () + 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 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 @@ -15,7 +15,6 @@ ( -- * Stencil Stencil , makeStencil - , makeStencilDef , getStencilSize , getStencilCenter -- ** Padding @@ -70,7 +69,7 @@ -- -- @since 0.1.0 mapStencil :: - (Source r ix e, Manifest r ix e) + (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 @@ -179,7 +178,7 @@ -- -- @since 0.4.3 applyStencil :: - (Source r ix e, Manifest r ix e) + (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 @@ -252,26 +251,6 @@ inline relStencil $ \ !ixD -> getVal (liftIndex2 (+) ix ixD) {-# INLINE stencil #-} {-# 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 -> e) -> a) - -- ^ Stencil function. - -> Stencil ix e a -makeStencilDef _defVal !sSz !sCenter relStencil = - Stencil sSz sCenter stencil - where - stencil _ getVal !ix = - inline relStencil $ \ !ixD -> getVal (liftIndex2 (+) ix ixD) - {-# INLINE stencil #-} -{-# INLINE makeStencilDef #-} -{-# DEPRECATED makeStencilDef "In favor of `makeStencil`. Validation is no longer possible" #-} -- | Identity stencil that does not change the elements of the source array. -- 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 @@ -61,7 +61,7 @@ -- -- @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 @@ -104,7 +104,7 @@ -- -- @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 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,4 +1,5 @@ {-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} @@ -102,6 +103,7 @@ !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) @@ -109,6 +111,7 @@ !newSz = unionStencilSizes maxCenter s1 s2 !maxCenter = unionStencilCenters s1 s2 {-# INLINE liftA2 #-} +#endif instance (Index ix, Num a) => Num (Stencil ix e a) where (+) = liftA2 (+) 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 @@ -17,45 +17,11 @@ , makeUnsafeConvolutionStencil , makeUnsafeCorrelationStencil , unsafeTransformStencil - -- ** Deprecated - , unsafeMapStencil ) 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) - - --- | This is an unsafe version of `Data.Massiv.Array.Stencil.mapStencil`, which does not --- take a `Stencil`, but instead accepts all necessary information as separate arguments. --- --- @since 0.5.0 -unsafeMapStencil :: - Manifest r ix e - => Border e - -> Sz ix - -> ix - -> (ix -> (ix -> e) -> a) - -> Array r ix e - -> Array DW ix a -unsafeMapStencil 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 ix) $ \ !ixD -> getVal (liftIndex2 (+) ix ixD) - {-# INLINE stencil #-} -{-# INLINE unsafeMapStencil #-} -{-# DEPRECATED unsafeMapStencil "In favor of `Data.Massiv.Array.mapStencil` that is applied to stencil created with `makeUnsafeStencil`" #-} -- | Similar to `Data.Massiv.Array.Stencil.makeStencil`, but there are no guarantees that the 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,4 +1,5 @@ {-# LANGUAGE BangPatterns #-} +{-# LANGUAGE ExplicitForAll #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} @@ -32,9 +33,13 @@ , unsafeInnerSlice , unsafeLinearSlice -- * Mutable interface + , unsafeResizeMArray + , unsafeLinearSliceMArray , unsafeThaw , unsafeFreeze , unsafeNew + , unsafeLoadIntoST + , unsafeLoadIntoIO , unsafeLoadIntoS , unsafeLoadIntoM , unsafeCreateArray @@ -85,22 +90,32 @@ , 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.Pull (D) -import Data.Massiv.Array.Delayed.Push (unsafeMakeLoadArray, unsafeMakeLoadArrayAdjusted) +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.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 -import Data.Massiv.Array.Stencil.Unsafe -unsafeBackpermute :: (Source r' ix' e, Index 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 $ \ !ix -> unsafeIndex arr (ixF ix) @@ -111,7 +126,7 @@ -- -- @since 0.3.0 unsafeTransform :: - (Source r' ix' e', Index ix) + (Index ix', Source r' e', Index ix) => (Sz ix' -> (Sz ix, a)) -> (a -> (ix' -> e') -> ix -> e) -> Array r' ix' e' @@ -126,7 +141,7 @@ -- -- @since 0.3.0 unsafeTransform2 :: - (Source r1 ix1 e1, Source r2 ix2 e2, Index ix) + (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 @@ -7,36 +7,36 @@ -- Portability : non-portable -- module Data.Massiv.Core - ( Array(List, unList) + ( Array(LArray) + , List(..) , Vector , MVector , Matrix , MMatrix - , Elt - , Construct - , Load(R, loadArrayM, loadArrayWithSetM) + , Load(iterArrayLinearST_, iterArrayLinearWithSetST_) , Stream(..) , Source - , Resize - , Extract + , Size + , Shape(..) + , LengthHint(..) , StrideLoad(..) - , Slice - , OuterSlice - , InnerSlice , Manifest , Mutable , Ragged - , Nested(..) - , NestedStruct , L(..) - , LN , ListItem , Scheduler , SchedulerWS + , Strategy , Comp(Seq, Par, Par', ParOn, ParN) + , getComp + , setComp , appComp , WorkerStates , initWorkerStates + , scheduleWork + , scheduleWork_ + , withMassivScheduler_ , module Data.Massiv.Core.Index -- * Numeric , FoldNumeric @@ -44,7 +44,6 @@ , NumericFloat -- * Exceptions , MonadThrow(..) - , throw , IndexException(..) , SizeException(..) , ShapeException(..) @@ -66,6 +65,6 @@ -- | Append computation strategy using `Comp`'s `Monoid` instance. -- -- @since 0.6.0 -appComp :: (Construct r ix e, Load r ix e) => Comp -> Array r ix e -> Array r ix e +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,7 +1,9 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE CPP #-} +{-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE DefaultSignatures #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} @@ -16,23 +18,20 @@ module Data.Massiv.Core.Common ( Array , Vector - , MVector , Matrix + , MArray + , MVector , MMatrix - , Elt , Steps(..) , Stream(..) - , Construct(..) + , Strategy(..) , Source(..) , Load(..) , StrideLoad(..) - , Resize(..) - , Extract(..) - , Slice(..) - , OuterSlice(..) - , InnerSlice(..) + , Size(..) + , Shape(..) , Manifest(..) - , Mutable(..) + , Mutable , Comp(..) , Scheduler , numWorkers @@ -48,16 +47,15 @@ , unsafeLinearSwap , unsafeDefaultLinearShrink , Ragged(..) - , Nested(..) - , NestedStruct , empty , singleton -- * Size , elemsCount - , linearSize + , isNotNull + , isEmpty , isNotEmpty , Sz(SafeSz) - , Size(..) + , LengthHint(..) -- * Indexing , (!?) , index @@ -69,13 +67,15 @@ , borderIndex , evaluateM , evaluate' + , inline0 + , inline1 + , inline2 , module Data.Massiv.Core.Index -- * Common Operations , imapM_ , Semigroup((<>)) -- * Exceptions , MonadThrow(..) - , throw , IndexException(..) , SizeException(..) , ShapeException(..) @@ -85,27 +85,29 @@ -- * Stateful Monads , runST , ST - , MonadUnliftIO + , MonadUnliftIO(..) , MonadIO(liftIO) , PrimMonad(PrimState) + , RealWorld ) 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.IO.Unlift (MonadIO(liftIO), MonadUnliftIO(..)) import Control.Monad.Primitive import Control.Monad.ST import Control.Scheduler (Comp(..), Scheduler, WorkerStates, numWorkers, - scheduleWork, scheduleWork_, withScheduler_, trivialScheduler_) + scheduleWork, scheduleWork_, trivialScheduler_, + withScheduler_) import Control.Scheduler.Global +import GHC.Exts (IsList) 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 Data.Kind import qualified Data.Vector.Fusion.Stream.Monadic as S (Stream) import Data.Vector.Fusion.Util @@ -113,43 +115,41 @@ -- | 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, even if that element does not yet exist in memory and the arry has to be --- computed in order to get access to that element. Data is always arranged in a nested --- row-major fashion, depth of which is controlled by @`Rank` ix@. -data family Array r ix e :: * +-- 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 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 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 - -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 Load r ix e => Stream r ix e where toStream :: Array r ix e -> Steps Id e @@ -157,14 +157,10 @@ data Steps m e = Steps { stepsStream :: S.Stream m e - , stepsSize :: Size + , stepsSize :: LengthHint } - --- | Array types that can be constructed. -class (Typeable r, Index ix) => Construct r ix e where - {-# MINIMAL setComp,(makeArray|makeArrayLinear) #-} - +class Typeable r => Strategy r where -- | Set computation strategy for this array -- -- ==== __Example__ @@ -181,78 +177,120 @@ -- 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. + -- | Get computation strategy of this array -- - -- >>> 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 ] - -- ] + -- @since 0.1.0 + getComp :: Array r ix e -> Comp + + +-- | Size hint +-- +-- @since 1.0.0 +data LengthHint + = LengthExact Sz1 -- ^ Exact known size + | LengthMax Sz1 -- ^ Upper bound on the size + | LengthUnknown -- ^ Unknown size + 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 -- - -- Instead of restricting the full type manually we can use `TypeApplications` as convenience: + -- @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 -- - -- >>> :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.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 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 #-} + -- @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 #-} - -- | Same as `makeArray`, but produce elements using linear row-major index. + -- | /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 - -- >>> makeArrayLinear Seq (Sz (2 :. 4)) id :: Array D Ix2 Int - -- Array D Seq (Sz (2 :. 4)) - -- [ [ 0, 1, 2, 3 ] - -- , [ 4, 5, 6, 7 ] - -- ] + -- >>> 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 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 #-} + -- @since 1.0.0 + isNull :: Array r ix e -> Bool + isNull = (zeroSz ==) . linearSize + {-# INLINE isNull #-} - replicate :: Comp -> Sz ix -> e -> Array r ix e - replicate comp sz !e = makeArray comp sz (const e) - {-# INLINE replicate #-} -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 +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 -class Load r ix e => Extract r ix e where - -- | /O(1)/ - Extract a portion of an array. Staring index and new size are + -- | /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. - unsafeExtract :: ix -> Sz ix -> Array r ix e -> Array (R r) ix e + -- + -- @since 0.1.0 + unsafeResize :: (Index ix, Index ix') => Sz ix' -> Array r ix e -> Array r ix' e + -- | Arrays that can be used as source to practically any manipulation function. -class (Resize r ix, Load r ix e) => Source r ix e where +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 :: Array r ix e -> ix -> e + unsafeIndex :: Index ix => Array r ix e -> ix -> e unsafeIndex = - INDEX_CHECK("(Source r ix e).unsafeIndex", + INDEX_CHECK("(Source r e).unsafeIndex", size, \ !arr -> unsafeLinearIndex arr . toLinearIndex (size arr)) {-# INLINE unsafeIndex #-} @@ -260,216 +298,262 @@ -- bounds check is performed -- -- @since 0.1.0 - unsafeLinearIndex :: Array r ix e -> Int -> e + unsafeLinearIndex :: Index ix => Array r ix e -> Int -> e unsafeLinearIndex !arr = unsafeIndex arr . fromLinearIndex (size arr) {-# INLINE unsafeLinearIndex #-} + + -- | /O(1)/ - Take a slice out of an array from the outside + -- + -- @since 0.1.0 + unsafeOuterSlice :: (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 :: Ix1 -> Sz1 -> Array r ix e -> Array r Ix1 e + unsafeLinearSlice :: Index ix => 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 - {-# MINIMAL getComp, size, (loadArrayM | loadArrayWithSetM) #-} +class (Strategy r, Shape r ix) => Load r ix e where + {-# MINIMAL (makeArray | makeArrayLinear), (iterArrayLinearST_ | iterArrayLinearWithSetST_)#-} - -- | Get computation strategy of this array + -- | 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. -- - -- @since 0.1.0 - getComp :: Array r ix e -> Comp - - -- | Get the exact size of an immutabe array. Most of the time will produce the size in - -- constant time, except for `DS` representation, which could result in evaluation of - -- the whole stream. See `maxSize` and `Data.Massiv.Vector.slength` for more info. + -- >>> 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 - size :: Array r ix e -> Sz ix + 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 #-} - -- | Load an array into memory. + -- | 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 - 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 () - loadArrayM scheduler arr uWrite = - loadArrayWithSetM scheduler arr uWrite $ \offset sz e -> - loopM_ offset (< (offset + unSz sz)) (+1) (\i -> uWrite i e) - {-# INLINE loadArrayM #-} + makeArrayLinear :: Comp -> Sz ix -> (Int -> e) -> Array r ix e + makeArrayLinear comp sz f = makeArray comp sz (f . toLinearIndex sz) + {-# INLINE makeArrayLinear #-} - -- | 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 - :: Monad m => - Scheduler m () - -> Array r ix e -- ^ Array that is being loaded - -> (Ix1 -> 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 () - loadArrayWithSetM scheduler arr uWrite _ = loadArrayM scheduler arr uWrite - {-# INLINE loadArrayWithSetM #-} - -- | /O(1)/ - Get the possible maximum 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. + -- | Construct an array of the specified size that contains the same element in all of + -- the cells. -- - -- @since 0.5.0 - maxSize :: Array r ix e -> Maybe (Sz ix) - maxSize = Just . size - {-# INLINE maxSize #-} + -- @since 0.3.0 + replicate :: Comp -> Sz ix -> e -> Array r ix e + replicate comp sz !e = makeArrayLinear comp sz (const e) + {-# INLINE replicate #-} - -- | /O(1)/ - Check if an 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 + -- | Iterate over an array with a ST action that is applied to each element and its index. -- - -- @since 0.1.0 - isEmpty :: Array r ix e -> Bool - isEmpty !arr = 0 == elemsCount arr - {-# INLINE isEmpty #-} + -- @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 -> + loopM_ 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. There is no guarantees, but some array representations, might + -- utilize this region setting function, in which case for such regions index + -- aware action will not be called. + -- + -- @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 + -- the same. -- - -- @since 0.5.7 - unsafeLoadIntoS :: - (Mutable r' ix e, PrimMonad m) - => MArray (PrimState m) r' ix e + -- @since 1.0.0 + unsafeLoadIntoST :: + Manifest r' e + => MVector s r' e -> Array r ix e - -> m (MArray (PrimState m) r' ix e) - unsafeLoadIntoS marr arr = - marr <$ loadArrayWithSetM trivialScheduler_ arr (unsafeLinearWrite marr) (unsafeLinearSet marr) - {-# INLINE unsafeLoadIntoS #-} + -> 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 `unsafeLoadIntoS`, but respecting computation strategy. + -- | Same as `unsafeLoadIntoST`, but respecting computation strategy. -- - -- @since 0.5.7 - unsafeLoadIntoM :: - (Mutable r' ix e, MonadIO m) - => MArray RealWorld r' ix e + -- @since 1.0.0 + unsafeLoadIntoIO :: + Manifest r' e + => MVector RealWorld r' e -> Array r ix e - -> m (MArray RealWorld r' ix e) - unsafeLoadIntoM marr arr = do - liftIO $ withMassivScheduler_ (getComp arr) $ \scheduler -> - loadArrayWithSetM scheduler arr (unsafeLinearWrite marr) (unsafeLinearSet marr) - pure marr - {-# INLINE unsafeLoadIntoM #-} - + -> 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 #-} --- | Selects an optimal scheduler for the supplied strategy, but it works only in `IO` -withMassivScheduler_ :: Comp -> (Scheduler IO () -> IO ()) -> IO () -withMassivScheduler_ comp f = - case comp of - Par -> withGlobalScheduler_ globalScheduler f - Seq -> f trivialScheduler_ - _ -> withScheduler_ comp f +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`. - loadArrayWithStrideM - :: Monad m => - Scheduler m () + 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 -> m ()) -- ^ Function that writes an element into target array - -> m () - default loadArrayWithStrideM - :: (Source r ix e, Monad m) => - Scheduler m () + -> (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 -> m ()) - -> m () - loadArrayWithStrideM scheduler stride resultSize arr = + -> (Int -> e -> ST s ()) + -> ST s () + iterArrayLinearWithStrideST_ 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 + {-# INLINE iterArrayLinearWithStrideST_ #-} -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) +-- 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. However, +-- this type class synonym will be deprecated in the next major version. +type Mutable r e = Manifest r 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 Source r ix e => Manifest r ix e where +-- 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 :: Array r ix e -> Int -> e + 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 -class (Construct r ix e, Manifest r ix e) => Mutable r ix e where - data MArray s r ix e :: * + -- | /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 - -- | Get the size of a mutable array. + -- | /O(1)/ - Take a linear slice out of a mutable array. -- - -- @since 0.1.0 - msize :: MArray s r ix e -> Sz ix + -- @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 :: PrimMonad m => Array r ix e -> m (MArray (PrimState m) r ix e) + 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 :: PrimMonad m => Comp -> MArray (PrimState m) r ix e -> m (Array r ix e) + 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 :: PrimMonad m => Sz ix -> m (MArray (PrimState m) r ix e) + 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 :: PrimMonad m => MArray (PrimState m) r ix e -> Int -> m e + unsafeLinearRead :: (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 :: PrimMonad m => MArray (PrimState m) r ix e -> Int -> e -> m () + unsafeLinearWrite :: (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 :: PrimMonad m => MArray (PrimState m) r ix e -> m () + 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 :: PrimMonad m => Maybe e -> Sz ix -> m (MArray (PrimState m) r ix e) + 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 #-} @@ -477,7 +561,7 @@ -- | Create new mutable array while initializing all elements to the specified value. -- -- @since 0.6.0 - newMArray :: PrimMonad m => Sz ix -> e -> m (MArray (PrimState m) r ix e) + 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 @@ -486,7 +570,7 @@ -- | Set all cells in the mutable array within the range to a specified value. -- -- @since 0.3.0 - unsafeLinearSet :: PrimMonad m => + unsafeLinearSet :: (Index ix, 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) @@ -495,7 +579,7 @@ -- | Copy part of one mutable array into another -- -- @since 0.3.6 - unsafeLinearCopy :: (Mutable r ix' e, PrimMonad m) => + unsafeLinearCopy :: (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 @@ -511,7 +595,7 @@ -- | Copy a part of a pure array into a mutable array -- -- @since 0.3.6 - unsafeArrayLinearCopy :: (Mutable r ix' e, PrimMonad m) => + unsafeArrayLinearCopy :: (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 @@ -529,7 +613,7 @@ -- no longer be used. -- -- @since 0.3.6 - unsafeLinearShrink :: PrimMonad m => + unsafeLinearShrink :: (Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> Sz ix -> m (MArray (PrimState m) r ix e) unsafeLinearShrink = unsafeDefaultLinearShrink {-# INLINE unsafeLinearShrink #-} @@ -539,17 +623,17 @@ -- should no longer be used. -- -- @since 0.3.6 - unsafeLinearGrow :: PrimMonad m => + unsafeLinearGrow :: (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 (msize marr)) + unsafeLinearCopy marr 0 marr' 0 $ SafeSz (totalElem (sizeOfMArray marr)) pure marr' {-# INLINE unsafeLinearGrow #-} unsafeDefaultLinearShrink :: - (Mutable r ix e, PrimMonad m) + (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> Sz ix -> m (MArray (PrimState m) r ix e) @@ -560,27 +644,39 @@ {-# INLINE unsafeDefaultLinearShrink #-} +-- | 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_ #-} + + -- | Read an array element -- -- @since 0.1.0 -unsafeRead :: (Mutable r ix e, PrimMonad m) => +unsafeRead :: (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> ix -> m e -unsafeRead marr = unsafeLinearRead marr . toLinearIndex (msize marr) +unsafeRead marr = unsafeLinearRead marr . toLinearIndex (sizeOfMArray marr) {-# INLINE unsafeRead #-} -- | Write an element into array -- -- @since 0.1.0 -unsafeWrite :: (Mutable r ix e, PrimMonad m) => +unsafeWrite :: (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> ix -> e -> m () -unsafeWrite marr = unsafeLinearWrite marr . toLinearIndex (msize marr) +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 :: (Mutable r ix e, PrimMonad m) => +unsafeLinearModify :: (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 @@ -592,19 +688,19 @@ -- | 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) => +unsafeModify :: (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 (msize marr) ix) +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 :: (Mutable r ix e, PrimMonad m) => - MArray (PrimState m) r ix e -> ix -> ix -> m (e, e) +unsafeSwap :: (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 = msize marr + where sz = sizeOfMArray marr {-# INLINE unsafeSwap #-} @@ -612,7 +708,7 @@ -- previous values. -- -- @since 0.4.0 -unsafeLinearSwap :: (Mutable r ix e, PrimMonad m) => +unsafeLinearSwap :: (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 @@ -623,32 +719,15 @@ {-# 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) +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) - edgeSize :: Array r ix e -> Sz ix - - flattenRagged :: Array r ix e -> Array r Ix1 e + flattenRagged :: Array r ix e -> Vector r e - loadRagged :: - Monad m => (m () -> m ()) -> (Int -> e -> m a) -> Int -> Int -> Sz ix -> Array r ix e -> m () + loadRaggedST :: + Scheduler s () -> Array r ix e -> (Ix1 -> e -> ST s ()) -> Ix1 -> Ix1 -> Sz ix -> ST s () - -- TODO: test property: - -- (read $ raggedFormat show "\n" (ls :: Array L (IxN n) Int)) == ls raggedFormat :: (e -> String) -> String -> Array r ix e -> String @@ -667,7 +746,7 @@ -- -- @since 0.3.0 empty :: - forall r ix e. Construct r ix e + forall r ix e. Load r ix e => Array r ix e empty = makeArray Seq zeroSz (const (throwImpossible Uninitialized)) {-# INLINE empty #-} @@ -696,7 +775,7 @@ -- -- @since 0.1.0 singleton :: - forall r ix e. Construct r ix e + forall r ix e. Load r ix e => e -- ^ The only element -> Array r ix e singleton = makeArray Seq oneSz . const @@ -718,12 +797,14 @@ -- ] -- >>> 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' +(!) :: + forall r ix e. (HasCallStack, Manifest r e, Index ix) + => Array r ix e + -> ix + -> e +(!) arr = throwEither . evaluateM arr {-# INLINE (!) #-} @@ -749,7 +830,11 @@ -- Nothing -- -- @since 0.1.0 -(!?) :: (Manifest r ix e, MonadThrow m) => Array r ix e -> ix -> m e +(!?) :: + forall r ix e m. (Index ix, Manifest r e, MonadThrow m) + => Array r ix e + -> ix + -> m e (!?) = indexM {-# INLINE (!?) #-} @@ -774,7 +859,7 @@ -- ] -- ) -- >>> ma ??> 1 --- Just (Array M Seq (Sz (1 :. 3)) +-- Just (Array U Seq (Sz (1 :. 3)) -- [ [ 4, 5, 6 ] -- ] -- ) @@ -784,7 +869,7 @@ -- Just 6 -- -- @since 0.1.0 -(??) :: (Manifest r ix e, MonadThrow m) => m (Array r ix e) -> ix -> m e +(??) :: (Index ix, Manifest r e, MonadThrow m) => m (Array r ix e) -> ix -> m e (??) marr ix = marr >>= (!? ix) {-# INLINE (??) #-} @@ -793,7 +878,7 @@ -- general and it 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 :: (Index ix, Manifest r e) => Array r ix e -> ix -> Maybe e index = indexM {-# INLINE index #-} @@ -802,7 +887,7 @@ -- /__Exceptions__/: `IndexOutOfBoundsException` -- -- @since 0.3.0 -indexM :: (Manifest r ix e, MonadThrow m) => Array r ix e -> ix -> m e +indexM :: (Index ix, Manifest r e, MonadThrow m) => Array r ix e -> ix -> m e indexM = evaluateM {-# INLINE indexM #-} @@ -820,7 +905,7 @@ -- 999 -- -- @since 0.1.0 -defaultIndex :: Manifest r ix e => e -> Array r ix e -> ix -> e +defaultIndex :: (Index ix, Manifest r e) => e -> Array r ix e -> ix -> e defaultIndex defVal = borderIndex (Fill defVal) {-# INLINE defaultIndex #-} @@ -837,12 +922,12 @@ -- [ 99, 100, 0, 1, 2 ] -- -- @since 0.1.0 -borderIndex :: Manifest r ix e => Border e -> Array r ix e -> ix -> e +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 can throw --- `IndexOutOfBoundsException` inside pure code. It is safer to use `index` instead. +-- | /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__ -- @@ -851,12 +936,10 @@ -- >>> 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' +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 @@ -874,33 +957,23 @@ -- 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 +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 exception in pure code. +-- | 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 --- >>> 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 +evaluate' :: (HasCallStack, Index ix, Source r e) => Array r ix e -> ix -> e +evaluate' arr ix = throwEither (evaluateM arr ix) {-# INLINE evaluate' #-} @@ -917,35 +990,44 @@ -- (4,14) -- -- @since 0.1.0 -imapM_ :: (Source r ix a, Monad m) => (ix -> a -> m b) -> Array r ix a -> m () +imapM_ :: (Index ix, Source r 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. --- --- /Note/ - It is always a constant time operation except for some arrays with --- `Data.Massiv.Array.DS` representation. See `Data.Massiv.Vector.slength` for more info. + +-- | /O(1)/ - Check if array has elements. -- -- ==== __Examples__ -- -- >>> import Data.Massiv.Array --- >>> elemsCount $ range Seq (Ix1 10) 15 --- 5 +-- >>> isNotNull (singleton 1 :: Array D Ix2 Int) +-- True +-- >>> isNotNull (empty :: Array D Ix2 Int) +-- False -- --- @since 0.1.0 -elemsCount :: Load r ix e => Array r ix e -> Int -elemsCount = totalElem . size -{-# INLINE elemsCount #-} +-- @since 0.5.1 +isNotNull :: Shape r ix => Array r ix e -> Bool +isNotNull = not . isNull +{-# INLINE isNotNull #-} --- | Get the number of elements in the array + +-- | /O(1)/ - Check if array has elements. -- --- @since 0.5.8 -linearSize :: Load r ix e => Array r ix e -> Sz1 -linearSize = toLinearSz . size -{-# INLINE linearSize #-} +-- ==== __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. @@ -958,7 +1040,34 @@ -- >>> isNotEmpty (empty :: Array D Ix2 Int) -- False -- --- @since 0.5.1 -isNotEmpty :: Load r ix e => Array r ix e -> Bool +-- @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,6 +1,7 @@ {-# LANGUAGE CPP #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE LambdaCase #-} +{-# LANGUAGE ImplicitParams #-} {-# OPTIONS_GHC -fno-warn-orphans #-} -- | -- Module : Data.Massiv.Core.Exception @@ -11,19 +12,21 @@ -- Portability : non-portable -- module Data.Massiv.Core.Exception - ( ImpossibleException(..) - , throwImpossible + ( 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.Stack +import GHC.Exception #if !MIN_VERSION_exceptions(0, 10, 3) import Control.Monad.ST (ST) @@ -34,31 +37,34 @@ throwM = unsafeIOToST . throwIO #endif - -newtype ImpossibleException = - ImpossibleException SomeException - deriving (Show) +-- | 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." -throwImpossible :: Exception e => e -> a -throwImpossible = throw . ImpossibleException . toException {-# NOINLINE throwImpossible #-} -throwEither :: Either SomeException a -> a +-- | 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 exc + Left exc -> throw (errorCallWithCallStackException (displayException exc) ?callStack) Right res -> res {-# INLINE throwEither #-} -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." - -- | An error that gets thrown when an unitialized element of a boxed array gets accessed. Can only --- happen when array was constructed with `unsafeNew`. +-- happen when array was constructed with `Data.Massiv.Array.Unsafe.unsafeNew`. data Uninitialized = Uninitialized deriving Show instance Exception Uninitialized where 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 @@ -17,6 +17,7 @@ , pattern Ix1 , type Ix2(Ix2, (:.)) , IxN((:>), Ix3, Ix4, Ix5) + , HighIxN , type Ix3 , type Ix4 , type Ix5 @@ -63,7 +64,8 @@ , Index(..) , zeroIndex , oneIndex - , isNonEmpty + , isZeroSz + , isNotZeroSz , headDim , tailDim , lastDim @@ -97,11 +99,10 @@ ) where import Control.DeepSeq -import Control.Exception (throw) 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.Ix import Data.Massiv.Core.Index.Stride @@ -109,6 +110,33 @@ import Data.Massiv.Core.Iterator import GHC.TypeLits + +-- | 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, @@ -212,22 +240,35 @@ 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 @@ -283,8 +324,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__ -- @@ -292,26 +333,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__ @@ -320,17 +359,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 @@ -344,41 +383,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. 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 @@ -26,7 +26,6 @@ ( Sz(SafeSz) , pattern Sz , pattern Sz1 - , type Sz1 , unSz , zeroSz , oneSz @@ -66,14 +65,42 @@ import Control.Monad (when) import Control.Monad.Catch (MonadThrow(..)) import Data.Coerce +import Data.Kind import Data.Massiv.Core.Iterator 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 @@ -92,20 +119,25 @@ Sz ix = SafeSz (liftIndex (max 0) ix) {-# COMPLETE Sz #-} --- | 1-dimensional type synonym for size. --- --- @since 0.3.0 -type Sz1 = Sz Ix1 - -- | 1-dimensional size constructor. Especially useful with literals: @(Sz1 5) == Sz (5 :: Int)@. -- -- @since 0.3.0 -pattern Sz1 :: Ix1 -> Sz1 +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 #-} + +instance (UniformRange ix, Index ix) => Random (Sz ix) + instance Index ix => Show (Sz ix) where showsPrec n sz@(SafeSz usz) = showsPrecWrapped n (str ++) where @@ -115,16 +147,20 @@ 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 #-} @@ -222,7 +258,7 @@ -- 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 #-} @@ -236,7 +272,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 #-} @@ -279,7 +315,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 #-} @@ -292,7 +328,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 #-} @@ -318,6 +354,14 @@ 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 + +instance Random Dim + -- | A way to select Array dimension at a type level. -- -- @since 0.2.4 @@ -366,7 +410,7 @@ -- 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 @@ -387,6 +431,7 @@ , Ord ix , Show ix , NFData ix + , Typeable ix , Eq (Lower ix) , Ord (Lower ix) , Show (Lower ix) @@ -708,7 +753,7 @@ -- | Index contains a zero value along one of the dimensions. IndexZeroException :: Index ix => !ix -> IndexException -- | Dimension is out of reach. - IndexDimensionException :: (NFData ix, 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 @@ -723,11 +768,13 @@ 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 @@ -763,15 +810,22 @@ 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 - (SizeOverflowException sz1, SizeOverflowException sz2) -> show sz1 == show sz2 - (SizeNegativeException sz1, SizeNegativeException 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 NFData SizeException where @@ -808,16 +862,27 @@ -- -- @since 0.3.0 data ShapeException - = DimTooShortException !Sz1 !Sz1 - | DimTooLongException + = DimTooShortException !Dim !(Sz Ix1) !(Sz Ix1) + -- ^ Across a specific dimension there was not enough elements for the supplied size + | DimTooLongException !Dim !(Sz Ix1) !(Sz Ix1) + -- ^ Across a specific dimension there was too many elements for the supplied size + | ShapeNonEmpty + -- ^ Expected an empty size, but the shape was not empty. 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 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 @@ -25,22 +25,17 @@ , 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 , HighIxN ) where @@ -49,10 +44,12 @@ import Control.DeepSeq import Data.Massiv.Core.Index.Internal import Data.Proxy +import qualified GHC.Arr as I 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 +import System.Random.Stateful #if !MIN_VERSION_base(4,11,0) import Data.Semigroup #endif @@ -73,15 +70,10 @@ pattern Ix2 i2 i1 = i2 :. i1 {-# COMPLETE Ix2 #-} --- | 2-dimensional size type synonym. --- --- @since 0.3.0 -type Sz2 = Sz 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 #-} @@ -98,15 +90,10 @@ pattern Ix3 i3 i2 i1 = i3 :> i2 :. i1 {-# COMPLETE Ix3 #-} --- | 3-dimensional size type synonym. --- --- @since 0.3.0 -type Sz3 = Sz 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 #-} @@ -122,15 +109,10 @@ 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 - -- | 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 #-} @@ -146,15 +128,10 @@ 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 - -- | 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 #-} @@ -183,7 +160,57 @@ 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 #-} + +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 #-} + +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] (+) #-} @@ -374,9 +401,9 @@ -- | Constraint synonym that encapsulates all constraints needed for dimension 4 and higher. -- --- @since 0.6.0 +-- @since 1.0.0 type HighIxN n - = (4 <= n, KnownNat n, KnownNat (n - 1), Index (Ix (n - 1)), IxN (n - 1) ~ Ix (n - 1)) + = (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 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 @@ -21,6 +21,7 @@ import Control.DeepSeq import Data.Massiv.Core.Index.Internal +import System.Random.Stateful -- | 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 @@ -63,6 +64,17 @@ 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 #-} + +instance (UniformRange ix, Index ix) => Random (Stride ix) -- | Just a helper function for unwrapping `Stride`. diff --git a/src/Data/Massiv/Core/Iterator.hs b/src/Data/Massiv/Core/Iterator.hs --- a/src/Data/Massiv/Core/Iterator.hs +++ b/src/Data/Massiv/Core/Iterator.hs @@ -1,4 +1,6 @@ {-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MonoLocalBinds #-} -- | -- Module : Data.Massiv.Core.Iterator -- Copyright : (c) Alexey Kuleshevich 2018-2021 @@ -23,6 +25,8 @@ import Control.Scheduler import Control.Monad +import Control.Monad.Primitive +import Control.Monad.IO.Unlift -- | Efficient loop with an accumulator -- @@ -126,7 +130,7 @@ -- -- @since 0.5.7 splitLinearlyM_ :: - Monad m => Scheduler m () -> Int -> (Int -> Int -> m ()) -> m () + MonadPrimBase s m => Scheduler s () -> Int -> (Int -> Int -> m ()) -> m () splitLinearlyM_ scheduler totalLength action = splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do loopNextM_ 0 (< slackStart) (+ chunkLength) $ \ start next -> @@ -140,7 +144,7 @@ -- -- @since 0.2.1 splitLinearlyWith_ :: - Monad m => Scheduler m () -> Int -> (Int -> b) -> (Int -> b -> m ()) -> m () + MonadPrimBase s m => Scheduler s () -> Int -> (Int -> b) -> (Int -> b -> m ()) -> m () splitLinearlyWith_ scheduler totalLength index = splitLinearlyWithM_ scheduler totalLength (pure . index) {-# INLINE splitLinearlyWith_ #-} @@ -150,7 +154,7 @@ -- -- @since 0.2.6 splitLinearlyWithM_ :: - Monad m => Scheduler m () -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m () + MonadPrimBase s m => Scheduler s () -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m () splitLinearlyWithM_ scheduler totalLength make write = splitLinearlyM_ scheduler totalLength go where @@ -163,7 +167,7 @@ -- -- @since 0.3.0 splitLinearlyWithStartAtM_ :: - Monad m => Scheduler m () -> Int -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m () + 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 -> @@ -180,20 +184,21 @@ -- -- @since 0.3.4 splitLinearlyWithStatefulM_ :: - Monad m - => SchedulerWS s m () + MonadUnliftIO m + => SchedulerWS ws () -> Int -- ^ Total linear length - -> (Int -> s -> m b) -- ^ Element producing action + -> (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 splitLinearly nWorkers totalLength $ \chunkLength slackStart -> do + in withRunInIO $ \run -> + 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 + run (make k s >>= store k) scheduleWorkState_ schedulerWS $ \s -> loopM_ slackStart (< totalLength) (+ 1) $ \ !k -> - make k s >>= store k + run (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,11 +1,14 @@ {-# 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 #-} -- | @@ -17,217 +20,203 @@ -- Portability : non-portable -- module Data.Massiv.Core.List - ( LN - , L(..) + ( 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.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.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 (L.map f [0 .. n - 1]) - {-# INLINE makeArray #-} - makeArrayLinear _ (Sz n) f = coerce (L.map f [0 .. n - 1]) - {-# INLINE makeArrayLinear #-} - -instance {-# OVERLAPPING #-} Nested LN Ix1 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 -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 #-} +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 #-} + , 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 + + +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 (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 !sz = outerSize arr + {-# INLINE iterArrayLinearST_ #-} -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 Ragged L Ix2 e where + generateRaggedM = unsafeGenerateParM + {-# INLINE generateRaggedM #-} + flattenRagged arr = LArray {lComp = lComp arr, lData = coerce xs} where - !arrL = LArray Seq arr - !sz = size arrL - {-# INLINE loadArrayM #-} - - - -outerLength :: Array L ix e -> Sz Int -outerLength = SafeSz . length . unList . lData + 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 ( 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] +instance ( Shape L (IxN n) + , Ragged L (Ix (n - 1)) e + , Coercible (Elt (Ix (n - 1)) e) (ListItem (Ix (n - 1)) 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 - 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) + Ragged L (IxN n) e where generateRaggedM = unsafeGenerateParM {-# INLINE generateRaggedM #-} 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 + 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) 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 #-} + 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) + (Elt ix e ~ List (Lower ix) e, Index ix, Monad m, Ragged L (Lower ix) e) => Comp -> Sz ix -> (ix -> m e) @@ -251,62 +240,45 @@ 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 +instance Strategy L where setComp c arr = arr {lComp = c} {-# INLINE setComp #-} - makeArray = unsafeGenerateN - {-# INLINE makeArray #-} + getComp = lComp + {-# INLINE getComp #-} - -- 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 #-} +-- -- 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) + showsPrec n arr = showsArrayLAsPrec (Proxy :: Proxy L) (outerSize arr) n arr -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 +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 @@ -320,35 +292,37 @@ 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 = +showsArrayLAsPrec pr sz n arr = opp . ("Array " ++) . showsTypeRep (typeRep pr) . (' ':) . - showsPrec 1 (getComp arr) . (" (" ++) . shows (size arr) . (")\n" ++) . shows lnarr . clp + showsPrec 1 (getComp arr) . (" (" ++) . shows sz . (")\n" ++) . shows lnarr . clp where (opp, clp) = if n == 0 then (id, id) else (('(':), ("\n)" ++)) - lnarr = toNested arr + 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 + 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 @@ -363,32 +337,8 @@ 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 #-} - - toStreamIx = S.indexed . S.fromList . coerce - {-# INLINE toStreamIx #-} - instance Stream L Ix1 e where - toStream = toStream . lData + toStream = S.fromList . unList . lData {-# INLINE toStream #-} - - toStreamIx = toStreamIx . lData + toStreamIx = S.indexed . S.fromList . unList . lData {-# INLINE toStreamIx #-} 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 @@ -24,8 +24,11 @@ import Data.Massiv.Core.Common -class Num e => FoldNumeric r e where +class (Size r, Num e) => FoldNumeric r e where + + {-# MINIMAL foldArray, powerSumArray, unsafeDotProduct #-} + -- | Compute sum of all elements in the array -- -- @since 0.5.6 @@ -57,7 +60,7 @@ defaultUnsafeDotProduct :: - (Num e, Source r ix e) => Array r ix e -> Array r ix e -> e + (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) @@ -66,7 +69,7 @@ | otherwise = acc {-# INLINE defaultUnsafeDotProduct #-} -defaultPowerSumArray :: (Source r ix e, Num e) => Array r ix e -> Int -> e +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) @@ -75,7 +78,7 @@ | otherwise = acc {-# INLINE defaultPowerSumArray #-} -defaultFoldArray :: Source r ix e => (e -> e -> e) -> e -> Array r ix e -> e +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) @@ -134,13 +137,13 @@ defaultUnsafeLiftArray :: - (Construct r ix e, Source r ix e) => (e -> e) -> Array r ix e -> Array r ix e + (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 :: - (Construct r ix e, Source r ix e) + (Load r ix e, Source r e) => (e -> e -> e) -> Array r ix e -> Array r ix e diff --git a/src/Data/Massiv/Vector.hs b/src/Data/Massiv/Vector.hs --- a/src/Data/Massiv/Vector.hs +++ b/src/Data/Massiv/Vector.hs @@ -1,4 +1,5 @@ {-# LANGUAGE BangPatterns #-} +{-# LANGUAGE ExplicitForAll #-} {-# LANGUAGE FlexibleContexts #-} {-# OPTIONS_GHC -fno-warn-duplicate-exports #-} -- | @@ -15,9 +16,10 @@ -- * Accessors -- *** Size , slength - , maxSize + , maxLinearSize , size - , snull + , isNull + , isNotNull -- *** Indexing , (!?) , (!) @@ -116,7 +118,7 @@ -- -- ** Permutations -- , reverse -- , backpermute - -- -- ** Mutable updates + -- -- ** Manifest updates -- , modify -- -- * Elementwise -- -- ** Mapping @@ -264,30 +266,6 @@ , convert , convertAs , convertProxy - -- -- ** Other vector types - -- , convert - -- -- ** Mutable vectors - -- , freeze - -- , thaw - -- , copy - -- , unsafeFreeze - -- , unsafeThaw - -- , unsafeCopy - -- * Deprecated - , takeS - , dropS - , unfoldr - , unfoldrN - , filterS - , ifilterS - , filterM - , ifilterM - , mapMaybeS - , imapMaybeS - , mapMaybeM - , imapMaybeM - , catMaybesS - , traverseS -- ** Re-exports , module Data.Massiv.Core , module Data.Massiv.Array.Delayed @@ -306,13 +284,14 @@ 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 (makeArrayR, replicate) +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, init, length, null, replicate, splitAt, tail, take, takeWhile, dropWhile) +import Prelude hiding (drop, dropWhile, init, length, null, replicate, splitAt, + tail, take, takeWhile) -- ========= -- -- Accessors -- @@ -326,7 +305,7 @@ -- | /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/: `maxSize`, `size`, `elemsCount` and `totalElem` +-- /Related/: `maxLinearSize`, `size`, `elemsCount` and `totalElem` -- -- ==== __Examples__ -- @@ -355,45 +334,16 @@ -- the vector. -- -- @since 0.5.0 -slength :: Stream r ix e => Array r ix e -> Maybe Sz1 +slength :: + forall r ix e. Stream r ix e + => Array r ix e + -> Maybe Sz1 slength v = case stepsSize (toStream v) of - Exact sz -> Just (SafeSz sz) - _ -> Nothing + LengthExact sz -> Just sz + _ -> Nothing {-# INLINE slength #-} --- | /O(1)/ - Check whether a `Stream` array is empty or not. It only looks at the exact size --- (i.e. `slength`), if it is available, otherwise checks if there is at least one element --- in a stream. --- --- /Related/: `isEmpty`, `isNotEmpty` --- --- ==== __Examples__ --- --- >>> snull sempty --- True --- >>> snull (empty :: Array D Ix5 Int) --- True --- >>> snull $ ssingleton "A Vector with a single String element" --- False --- >>> snull $ sfromList [] --- True --- >>> snull $ sfromList [1 :: Int ..] --- False --- --- /__Similar__/: --- --- [@Data.Foldable.`Data.Foldable.null`@] List fusion is also broken with a check for --- emptiness, unless there are no other consumers of the list. --- --- [@Data.Vector.Generic.`Data.Vector.Generic.null`@] Same as with --- `Data.Vector.Generic.length`, unless it is the only operation applied to the vector it --- will break fusion and will result in the vector being fully materialized in memory. --- --- @since 0.5.0 -snull :: Load r ix e => Array r ix e -> Bool -snull = isEmpty -{-# INLINE snull #-} -------------- -- Indexing -- @@ -408,8 +358,6 @@ -- -- >>> head' (Ix1 10 ..: 10000000000000) -- 10 --- >>> head' (Ix1 10 ..: 10) --- *** Exception: SizeEmptyException: (Sz1 0) corresponds to an empty array -- -- /__Similar__/: -- @@ -420,8 +368,11 @@ -- cause materialization of the full vector if any other function is applied to the vector. -- -- @since 0.5.0 -head' :: Source r Ix1 e => Vector r e -> e -head' = either throw id . headM +head' :: + forall r e. (HasCallStack, Source r e) + => Vector r e + -> e +head' = throwEither . headM {-# INLINE head' #-} @@ -448,9 +399,12 @@ -- except it is restricted to `Maybe` -- -- @since 0.5.0 -headM :: (Source r Ix1 e, MonadThrow m) => Vector r e -> m e +headM :: + forall r e m. (Source r e, MonadThrow m) + => Vector r e + -> m e headM v - | isEmpty v = throwM $ SizeEmptyException (size v) + | elemsCount v == 0 = throwM $ SizeEmptyException (size v) | otherwise = pure $ unsafeLinearIndex v 0 {-# INLINE headM #-} @@ -463,13 +417,15 @@ -- -- >>> shead' $ sunfoldr (\x -> Just (x, x)) (0 :: Int) -- 0 --- >>> x = shead' $ sunfoldr (\_ -> Nothing) (0 :: Int) --- >>> print x --- *** Exception: SizeEmptyException: (Sz1 0) corresponds to an empty array +-- >>> shead' (Ix1 3 ... 5) +-- 3 -- -- @since 0.5.0 -shead' :: Stream r Ix1 e => Vector r e -> e -shead' = either throw id . sheadM +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. @@ -490,10 +446,13 @@ -- *** Exception: SizeEmptyException: (Sz1 0) corresponds to an empty array -- -- @since 0.5.0 -sheadM :: (Stream r Ix1 e, MonadThrow m) => Vector r e -> m e +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 (size v) + Nothing -> throwM $ SizeEmptyException (zeroSz :: Sz1) Just e -> pure e {-# INLINE sheadM #-} @@ -517,7 +476,10 @@ -- the more general `MonadThrow` -- -- @since 0.3.0 -unconsM :: (MonadThrow m, Source r Ix1 e) => Vector r e -> m (e, Vector r e) +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) @@ -539,7 +501,10 @@ -- [ 1, 2 ],3) -- -- @since 0.3.0 -unsnocM :: (MonadThrow m, Source r Ix1 e) => Vector r e -> m (Vector r e, e) +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) @@ -557,20 +522,18 @@ -- -- >>> last' (Ix1 10 ... 10000000000000) -- 10000000000000 --- >>> last' (fromList Seq [] :: Array P Ix1 Int) --- *** Exception: SizeEmptyException: (Sz1 0) corresponds to an empty array -- -- /__Similar__/: -- --- [@Data.List.`Data.List.last`@] Also partial, but it has /O(n)/ complixity. Fusion is +-- [@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' :: Source r Ix1 e => Vector r e -> e -last' = either throw id . lastM +last' :: forall r e. (HasCallStack, Source r e) => Vector r e -> e +last' = throwEither . lastM {-# INLINE last' #-} @@ -590,7 +553,7 @@ -- "SizeEmptyException: (Sz1 0) corresponds to an empty array" -- -- @since 0.5.0 -lastM :: (Source r Ix1 e, MonadThrow m) => Vector r e -> m e +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) @@ -613,7 +576,7 @@ -- [ 9999999999998, 9999999999999, 10000000000000 ] -- -- @since 0.5.0 -slice :: Source r Ix1 e => Ix1 -> Sz1 -> Vector r e -> Vector r e +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) @@ -628,19 +591,13 @@ -- >>> slice' 10 5 (Ix1 0 ... 100) -- Array D Seq (Sz1 5) -- [ 10, 11, 12, 13, 14 ] --- >>> slice' (-10) 5 (Ix1 0 ... 100) --- Array D *** Exception: SizeSubregionException: (Sz1 101) is to small for -10 (Sz1 5) --- >>> slice' 98 50 (Ix1 0 ... 100) --- Array D *** Exception: SizeSubregionException: (Sz1 101) is to small for 98 (Sz1 50) --- >>> slice' 9999999999998 50 (Ix1 0 ... 10000000000000) --- Array D *** Exception: SizeSubregionException: (Sz1 10000000000001) is to small for 9999999999998 (Sz1 50) -- >>> slice' 9999999999998 3 (Ix1 0 ... 10000000000000) -- Array D Seq (Sz1 3) -- [ 9999999999998, 9999999999999, 10000000000000 ] -- -- @since 0.5.0 -slice' :: Source r Ix1 e => Ix1 -> Sz1 -> Vector r e -> Vector r e -slice' i k = either throw id . sliceM i k +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' #-} @@ -650,9 +607,27 @@ -- -- ==== __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 :: (Source r Ix1 e, MonadThrow m) => Ix1 -> Sz1 -> Vector r e -> m (Vector r e) +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 @@ -685,8 +660,16 @@ -- Nothing -- -- @since 0.5.0 -sslice :: Stream r Ix1 e => Ix1 -> Sz1 -> Vector r e -> Vector DS e -sslice !i (Sz k) = fromSteps . S.slice i k . S.toStream +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 #-} @@ -694,15 +677,16 @@ -- -- ==== __Examples__ -- --- >>> init (0 ..: 10) +-- >>> import Data.Massiv.Array as A +-- >>> A.init (0 ..: 10) -- Array D Seq (Sz1 9) -- [ 0, 1, 2, 3, 4, 5, 6, 7, 8 ] --- >>> init (empty :: Array D Ix1 Int) +-- >>> A.init (empty :: Array D Ix1 Int) -- Array D Seq (Sz1 0) -- [ ] -- -- @since 0.5.0 -init :: Source r Ix1 e => Vector r e -> Vector r e +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 #-} @@ -713,30 +697,29 @@ -- >>> init' (0 ..: 10) -- Array D Seq (Sz1 9) -- [ 0, 1, 2, 3, 4, 5, 6, 7, 8 ] --- >>> init' (empty :: Array D Ix1 Int) --- Array D *** Exception: SizeEmptyException: (Sz1 0) corresponds to an empty array -- -- @since 0.5.0 -init' :: Source r Ix1 e => Vector r e -> Vector r e -init' = either throw id . initM +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 sum $ initM (0 ..: 10) +-- >>> maybe 0 A.sum $ initM (0 ..: 10) -- 36 --- >>> maybe 0 sum $ initM (empty :: Array D Ix1 Int) +-- >>> maybe 0 A.sum $ initM (empty :: Array D Ix1 Int) -- 0 -- -- @since 0.5.0 -initM :: (Source r Ix1 e, MonadThrow m) => Vector r e -> m (Vector r e) +initM :: forall r e m. (Source r e, MonadThrow m) => Vector r e -> m (Vector r e) initM v = do - when (isEmpty v) $ throwM $ SizeEmptyException $ size v + when (elemsCount v == 0) $ throwM $ SizeEmptyException $ size v pure $ unsafeInit v {-# INLINE initM #-} @@ -746,16 +729,17 @@ -- -- ==== __Examples__ -- --- >>> tail (0 ..: 10) +-- >>> import Data.Massiv.Array as A +-- >>> A.tail (0 ..: 10) -- Array D Seq (Sz1 9) -- [ 1, 2, 3, 4, 5, 6, 7, 8, 9 ] --- >>> tail (empty :: Array D Ix1 Int) +-- >>> A.tail (empty :: Array D Ix1 Int) -- Array D Seq (Sz1 0) -- [ ] -- -- @since 0.5.0 -tail :: Source r Ix1 e => Vector r e -> Vector r e -tail = drop 1 +tail :: forall r e. Source r e => Vector r e -> Vector r e +tail = drop oneSz {-# INLINE tail #-} @@ -770,8 +754,8 @@ -- Array D *** Exception: SizeEmptyException: (Sz1 0) corresponds to an empty array -- -- @since 0.5.0 -tail' :: Source r Ix1 e => Vector r e -> Vector r e -tail' = either throw id . tailM +tail' :: forall r e. (HasCallStack, Source r e) => Vector r e -> Vector r e +tail' = throwEither . tailM {-# INLINE tail' #-} @@ -779,18 +763,19 @@ -- -- ==== __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 sum $ tailM (0 ..: 10) +-- >>> maybe 0 A.sum $ tailM (0 ..: 10) -- 45 --- >>> maybe 0 sum $ tailM (empty :: Array D Ix1 Int) +-- >>> maybe 0 A.sum $ tailM (empty :: Array D Ix1 Int) -- 0 -- -- @since 0.5.0 -tailM :: (Source r Ix1 e, MonadThrow m) => Vector r e -> m (Vector r e) +tailM :: forall r e m. (Source r e, MonadThrow m) => Vector r e -> m (Vector r e) tailM v = do - when (isEmpty v) $ throwM $ SizeEmptyException $ size v + when (elemsCount v == 0) $ throwM $ SizeEmptyException $ size v pure $ unsafeTail v {-# INLINE tailM #-} @@ -800,19 +785,19 @@ -- -- ==== __Examples__ -- --- >>> take 5 (0 ..: 10) +-- >>> import Data.Massiv.Array as A +-- >>> A.take 5 (0 ..: 10) -- Array D Seq (Sz1 5) -- [ 0, 1, 2, 3, 4 ] --- >>> take (-5) (0 ..: 10) +-- >>> A.take 0 (0 ..: 10) -- Array D Seq (Sz1 0) -- [ ] --- >>> take 100 (0 ..: 10) +-- >>> 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 Ix1 e => Sz1 -> Vector r e -> Vector r e +take :: Source r e => Sz1 -> Vector r e -> Vector r e take k = fst . sliceAt k {-# INLINE take #-} @@ -821,7 +806,7 @@ -- satisfy the supplied predicate. -- -- @since 0.5.5 -takeWhile :: Manifest r Ix1 e => (e -> Bool) -> Vector r e -> Vector r e +takeWhile :: Manifest r e => (e -> Bool) -> Vector r e -> Vector r e takeWhile f v = take (go 0) v where !k = elemsCount v @@ -831,8 +816,6 @@ {-# INLINE takeWhile #-} - - -- | /O(1)/ - Get the vector with the first @n@ elements. Throws an error size is less -- than @n@. -- @@ -844,28 +827,29 @@ -- >>> take' 5 (0 ..: 10) -- Array D Seq (Sz1 5) -- [ 0, 1, 2, 3, 4 ] --- >>> take' 15 (0 ..: 10) --- Array D *** Exception: SizeSubregionException: (Sz1 10) is to small for 0 (Sz1 15) -- -- @since 0.5.0 -take' :: Source r Ix1 e => Sz1 -> Vector r e -> Vector r e -take' k = either throw id . takeM k +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 sum $ takeM 5 (0 ..: 10) +-- >>> maybe 0 A.sum $ takeM 5 (0 ..: 10) -- 10 --- >>> maybe (-1) sum $ takeM 15 (0 ..: 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 :: (Source r Ix1 e, MonadThrow m) => Sz1 -> Vector r e -> m (Vector r e) +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 @@ -877,8 +861,8 @@ -- ==== __Examples__ -- -- @since 0.5.0 -stake :: Stream r Ix1 e => Sz1 -> Vector r e -> Vector DS e -stake n = fromSteps . S.take (unSz n) . S.toStream +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 #-} -- | @@ -886,7 +870,7 @@ -- ==== __Examples__ -- -- @since 0.5.0 -drop :: Source r Ix1 e => Sz1 -> Vector r e -> Vector r e +drop :: forall r e. Source r e => Sz1 -> Vector r e -> Vector r e drop k = snd . sliceAt k {-# INLINE drop #-} @@ -895,7 +879,7 @@ -- that satisfy the supplied predicate. -- -- @since 0.5.5 -dropWhile :: Manifest r Ix1 e => (e -> Bool) -> Vector r e -> Vector r e +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 @@ -910,8 +894,8 @@ -- ==== __Examples__ -- -- @since 0.5.0 -sdrop :: Stream r Ix1 e => Sz1 -> Vector r e -> Vector DS e -sdrop n = fromSteps . S.drop (unSz n) . S.toStream +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 #-} -- | @@ -919,8 +903,8 @@ -- ==== __Examples__ -- -- @since 0.5.0 -drop' :: Source r Ix1 e => Sz1 -> Vector r e -> Vector r e -drop' k = either throw id . dropM k +drop' :: forall r e. (HasCallStack, Source r e) => Sz1 -> Vector r e -> Vector r e +drop' k = throwEither . dropM k {-# INLINE drop' #-} -- | @@ -928,10 +912,10 @@ -- ==== __Examples__ -- -- @since 0.5.0 -dropM :: (Source r Ix1 e, MonadThrow m) => Sz1 -> Vector r e -> m (Vector r e) +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 (sz - k) + when (k > sz) $ throwM $ SizeSubregionException sz d (SafeSz (n - d)) pure $ unsafeLinearSlice d (SafeSz (n - d)) v {-# INLINE dropM #-} @@ -942,7 +926,7 @@ -- ==== __Examples__ -- -- @since 0.5.0 -sliceAt :: Source r Ix1 e => Sz1 -> Vector r e -> (Vector r e, Vector r e) +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) @@ -954,8 +938,8 @@ -- ==== __Examples__ -- -- @since 0.5.0 -sliceAt' :: Source r Ix1 e => Sz1 -> Vector r e -> (Vector r e, Vector r e) -sliceAt' k = either throw id . sliceAtM k +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. @@ -963,7 +947,7 @@ -- ==== __Examples__ -- -- @since 0.5.0 -sliceAtM :: (Source r Ix1 e, MonadThrow m) => Sz1 -> Vector r e -> m (Vector r e, Vector r e) +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) @@ -991,7 +975,7 @@ -- | /O(1)/ - Add an element to the vector from the left side -- -- @since 0.3.0 -cons :: Load r Ix1 e => e -> Vector r e -> Vector DL e +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 = @@ -1003,7 +987,7 @@ -- | /O(1)/ - Add an element to the vector from the right side -- -- @since 0.3.0 -snoc :: Load r Ix1 e => Vector r e -> e -> Vector DL e +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) @@ -1021,7 +1005,7 @@ -- -- @since 0.5.0 sreplicate :: Sz1 -> e -> Vector DS e -sreplicate (Sz n) = DSArray . S.replicate n +sreplicate n = DSArray . S.replicate n {-# INLINE sreplicate #-} -- | Create a delayed vector of length @n@ with a function that maps an index to an @@ -1041,7 +1025,7 @@ -- -- @since 0.5.0 sgenerate :: Sz1 -> (Ix1 -> e) -> Vector DS e -sgenerate (Sz n) = DSArray . S.generate n +sgenerate n = DSArray . S.generate n {-# INLINE sgenerate #-} @@ -1070,7 +1054,7 @@ -- -- @since 0.5.0 siterateN :: Sz1 -> (e -> e) -> e -> Vector DS e -siterateN n f a = fromSteps $ S.iterateN (unSz n) f a +siterateN n f a = fromSteps $ S.iterateN n f a {-# INLINE siterateN #-} @@ -1079,8 +1063,8 @@ -- ==== __Examples__ -- -- @since 0.5.0 -sreplicateM :: Monad m => Sz1 -> m e -> m (Vector DS e) -sreplicateM n f = fromStepsM $ S.replicateM (unSz n) f +sreplicateM :: forall e m. Monad m => Sz1 -> m e -> m (Vector DS e) +sreplicateM n f = fromStepsM $ S.replicateM n f {-# INLINE sreplicateM #-} @@ -1090,8 +1074,8 @@ -- ==== __Examples__ -- -- @since 0.5.0 -sgenerateM :: Monad m => Sz1 -> (Ix1 -> m e) -> m (Vector DS e) -sgenerateM n f = fromStepsM $ S.generateM (unSz n) f +sgenerateM :: forall e m. Monad m => Sz1 -> (Ix1 -> m e) -> m (Vector DS e) +sgenerateM n f = fromStepsM $ S.generateM n f {-# INLINE sgenerateM #-} @@ -1101,8 +1085,8 @@ -- ==== __Examples__ -- -- @since 0.5.0 -siterateNM :: Monad m => Sz1 -> (e -> m e) -> e -> m (Vector DS e) -siterateNM n f a = fromStepsM $ S.iterateNM (unSz n) f a +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 #-} @@ -1119,7 +1103,7 @@ -- [ 0, 1, 4, 9, 16, 25, 36, 49, 64 ] -- -- @since 0.5.0 -sunfoldr :: (s -> Maybe (e, s)) -> s -> Vector DS e +sunfoldr :: forall e s. (s -> Maybe (e, s)) -> s -> Vector DS e sunfoldr f = DSArray . S.unfoldr f {-# INLINE sunfoldr #-} @@ -1136,6 +1120,7 @@ -- -- @since 0.5.0 sunfoldrN :: + forall e s. Sz1 -- ^ @n@ - maximum number of elements that the vector will have -> (s -> Maybe (e, s)) @@ -1143,7 +1128,7 @@ -- is reached. -> s -- ^ Inititial element. -> Vector DS e -sunfoldrN (Sz n) f = DSArray . S.unfoldrN n f +sunfoldrN n f = DSArray . S.unfoldrN n f {-# INLINE sunfoldrN #-} -- | /O(n)/ - Same as `unfoldr`, but with monadic generating function. @@ -1159,7 +1144,7 @@ -- ) -- -- @since 0.5.0 -sunfoldrM :: Monad m => (s -> m (Maybe (e, s))) -> s -> m (Vector DS e) +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 #-} @@ -1184,7 +1169,7 @@ -- -- -- @since 0.5.0 -sunfoldrNM :: Monad m => Sz1 -> (s -> m (Maybe (e, s))) -> s -> m (Vector DS e) +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 #-} @@ -1198,8 +1183,8 @@ -- [ 100, 121, 144, 169, 196, 225, 256, 289, 324, 361 ] -- -- @since 0.5.0 -sunfoldrExactN :: Sz1 -> (s -> (e, s)) -> s -> Vector DS e -sunfoldrExactN (Sz n) f = fromSteps . S.unfoldrExactN n f +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 `unfoldrNM`, except the length of the resulting vector will be exactly @n@ @@ -1216,15 +1201,15 @@ -- ) -- -- @since 0.5.0 -sunfoldrExactNM :: Monad m => Sz1 -> (s -> m (e, s)) -> s -> m (Vector DS e) -sunfoldrExactNM (Sz n) f = fromStepsM . S.unfoldrExactNM n f +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` +-- `rangeStepSize`, `range`, `rangeStep'` -- -- ==== __Examples__ -- @@ -1246,7 +1231,7 @@ => e -- ^ @x@ - starting number -> Sz1 -- ^ @n@ - length of resulting vector -> Vector DS e -senumFromN x (Sz n) = DSArray $ S.enumFromStepN x 1 n +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@ @@ -1272,7 +1257,7 @@ -> e -- ^ @dx@ - Step -> Sz1 -- ^ @n@ - length of resulting vector -> Vector DS e -senumFromStepN x step (Sz n) = DSArray $ S.enumFromStepN x step n +senumFromStepN x step n = DSArray $ S.enumFromStepN x step n {-# INLINE senumFromStepN #-} @@ -1300,7 +1285,11 @@ -- memory representations. -- -- @since 0.5.0 -sappend :: (Stream r1 Ix1 e, Stream r2 Ix1 e) => Vector r1 e -> Vector r2 e -> Vector DS e +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 #-} @@ -1330,7 +1319,7 @@ -- implementation underneath as `sconcat`. -- -- @since 0.5.0 -sconcat :: Stream r Ix1 e => [Vector r e] -> Vector DS e +sconcat :: forall r e. Stream r Ix1 e => [Vector r e] -> Vector DS e sconcat = DSArray . foldMap toStream {-# INLINE sconcat #-} @@ -1379,15 +1368,12 @@ -- ==== __Examples__ -- -- @since 0.5.0 -stoList :: Stream r ix e => Array r ix e -> [e] +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__ @@ -1405,7 +1391,7 @@ -- [ (0,0), (0,1), (0,2), (0,3), (2,0), (2,1), (2,2), (2,3) ] -- -- @since 0.5.0 -sfilter :: S.Stream r ix e => (e -> Bool) -> Array r ix e -> Vector DS e +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 #-} @@ -1415,7 +1401,7 @@ -- ==== __Examples__ -- -- @since 0.5.0 -sifilter :: Stream r ix a => (ix -> a -> Bool) -> Array r ix a -> Vector DS a +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 @@ -1457,7 +1443,11 @@ -- [ (0,0), (0,2), (1,0), (1,2), (2,0), (2,2) ] -- -- @since 0.5.0 -sfilterM :: (S.Stream r ix e, Applicative f) => (e -> f Bool) -> Array r ix e -> f (Vector DS e) +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 #-} @@ -1468,7 +1458,10 @@ -- -- @since 0.5.0 sifilterM :: - (Stream r ix a, Applicative f) => (ix -> a -> f Bool) -> Array r ix a -> f (Vector DS a) + 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 -> @@ -1485,7 +1478,7 @@ -- ==== __Examples__ -- -- @since 0.5.0 -smapMaybe :: S.Stream r ix a => (a -> Maybe b) -> Array r ix a -> Vector DS b +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 #-} @@ -1495,7 +1488,11 @@ -- ==== __Examples__ -- -- @since 0.5.0 -simapMaybe :: Stream r ix a => (ix -> a -> Maybe b) -> Array r ix a -> Vector DS b +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 #-} @@ -1505,7 +1502,10 @@ -- -- @since 0.5.0 simapMaybeM :: - (Stream r ix a, Applicative f) => (ix -> a -> f (Maybe b)) -> Array r ix a -> f (Vector DS b) + 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 #-} @@ -1515,7 +1515,7 @@ -- ==== __Examples__ -- -- @since 0.5.0 -scatMaybes :: S.Stream r ix (Maybe a) => Array r ix (Maybe a) -> Vector DS a +scatMaybes :: forall r ix a. S.Stream r ix (Maybe a) => Array r ix (Maybe a) -> Vector DS a scatMaybes = smapMaybe id {-# INLINE scatMaybes #-} @@ -1528,7 +1528,10 @@ -- -- @since 0.5.0 smapMaybeM :: - (S.Stream r ix a, Applicative f) => (a -> f (Maybe b)) -> Array r ix a -> f (Vector DS b) + 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 #-} @@ -1539,7 +1542,11 @@ -- ==== __Examples__ -- -- @since 0.5.0 -smap :: S.Stream r ix a => (a -> b) -> Array r ix a -> Vector DS b +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 #-} @@ -1548,7 +1555,11 @@ -- ==== __Examples__ -- -- @since 0.5.0 -simap :: S.Stream r ix a => (ix -> a -> b) -> Array r ix a -> Vector DS b +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 #-} @@ -1558,7 +1569,11 @@ -- ==== __Examples__ -- -- @since 0.5.0 -straverse :: (S.Stream r ix a, Applicative f) => (a -> f b) -> Array r ix a -> f (Vector DS b) +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 #-} @@ -1568,7 +1583,11 @@ -- ==== __Examples__ -- -- @since 0.5.0 -sitraverse :: (S.Stream r ix a, Applicative f) => (ix -> a -> f b) -> Array r ix a -> f (Vector DS b) +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 #-} @@ -1578,7 +1597,11 @@ -- ==== __Examples__ -- -- @since 0.5.0 -smapM :: (S.Stream r ix a, Monad m) => (a -> m b) -> Array r ix a -> m (Vector DS b) +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 #-} @@ -1589,7 +1612,11 @@ -- ==== __Examples__ -- -- @since 0.5.0 -simapM :: (S.Stream r ix a, Monad m) => (ix -> a -> m b) -> Array r ix a -> m (Vector DS b) +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 #-} @@ -1598,7 +1625,11 @@ -- ==== __Examples__ -- -- @since 0.5.0 -smapM_ :: (S.Stream r ix a, Monad m) => (a -> m b) -> Array r ix a -> m () +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_ #-} @@ -1607,7 +1638,11 @@ -- ==== __Examples__ -- -- @since 0.5.0 -simapM_ :: (S.Stream r ix a, Monad m) => (ix -> a -> m b) -> Array r ix a -> m () +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_ #-} @@ -1617,7 +1652,11 @@ -- ==== __Examples__ -- -- @since 0.5.0 -sforM :: (S.Stream r ix a, Monad m) => Array r ix a -> (a -> m b) -> m (Vector DS b) +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 #-} @@ -1626,7 +1665,11 @@ -- ==== __Examples__ -- -- @since 0.5.0 -siforM :: (S.Stream r ix a, Monad m) => Array r ix a -> (ix -> a -> m b) -> m (Vector DS b) +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 #-} @@ -1644,7 +1687,11 @@ -- ==== __Examples__ -- -- @since 0.5.0 -siforM_ :: (S.Stream r ix a, Monad m) => Array r ix a -> (ix -> a -> m b) -> m () +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_ #-} @@ -1657,7 +1704,10 @@ -- -- @since 0.5.0 szip :: - (S.Stream ra Ix1 a, S.Stream rb Ix1 b) => Vector ra a -> Vector rb b -> Vector DS (a, b) + 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 #-} @@ -1665,7 +1715,7 @@ -- -- @since 0.5.0 szip3 :: - (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c) + 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 @@ -1677,6 +1727,7 @@ -- -- @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 @@ -1690,6 +1741,7 @@ -- -- @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 @@ -1704,6 +1756,7 @@ -- -- @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 @@ -1732,6 +1785,7 @@ -- -- @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 @@ -1744,6 +1798,7 @@ -- -- @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 @@ -1757,6 +1812,7 @@ -- -- @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 @@ -1772,6 +1828,7 @@ -- -- @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 @@ -1789,6 +1846,7 @@ -- -- @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 @@ -1822,6 +1880,7 @@ -- -- @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 @@ -1834,6 +1893,7 @@ -- -- @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 @@ -1848,6 +1908,7 @@ -- -- @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 @@ -1864,6 +1925,7 @@ -- -- @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 @@ -1887,6 +1949,7 @@ -- -- @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 @@ -1921,6 +1984,7 @@ -- -- @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 @@ -1933,6 +1997,7 @@ -- -- @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 @@ -1946,6 +2011,7 @@ -- -- @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 @@ -1961,6 +2027,7 @@ -- -- @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 @@ -1984,6 +2051,7 @@ -- -- @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 @@ -2019,6 +2087,7 @@ -- -- @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 @@ -2032,6 +2101,7 @@ -- -- @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 @@ -2046,6 +2116,7 @@ -- -- @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 @@ -2062,6 +2133,7 @@ -- -- @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 @@ -2093,6 +2165,7 @@ -- -- @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 @@ -2128,7 +2201,7 @@ -- -- @since 0.5.0 szipWithM_ :: - (S.Stream ra Ix1 a, S.Stream rb Ix1 b, Monad m) + 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 @@ -2140,6 +2213,7 @@ -- -- @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 @@ -2153,6 +2227,7 @@ -- -- @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 @@ -2168,6 +2243,7 @@ -- -- @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 @@ -2190,6 +2266,7 @@ -- -- @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 @@ -2226,7 +2303,7 @@ -- -- @since 0.5.0 sizipWithM_ :: - (S.Stream ra Ix1 a, S.Stream rb Ix1 b, Monad m) + 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 @@ -2239,6 +2316,7 @@ -- -- @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 @@ -2252,6 +2330,7 @@ -- -- @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 @@ -2267,6 +2346,7 @@ -- -- @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 @@ -2295,6 +2375,7 @@ -- -- @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 @@ -2326,12 +2407,17 @@ --- | +-- | Strict left fold sequentially over a streamed array. -- -- ==== __Examples__ -- -- @since 0.5.0 -sfoldl :: Stream r ix e => (a -> e -> a) -> a -> Array r ix e -> a +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 #-} @@ -2340,7 +2426,12 @@ -- ==== __Examples__ -- -- @since 0.5.0 -sfoldlM :: (Stream r ix e, Monad m) => (a -> e -> m a) -> a -> Array r ix e -> m a +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 #-} @@ -2349,7 +2440,12 @@ -- ==== __Examples__ -- -- @since 0.5.0 -sfoldlM_ :: (Stream r ix e, Monad m) => (a -> e -> m a) -> a -> Array r ix e -> m () +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_ #-} @@ -2359,8 +2455,12 @@ -- ==== __Examples__ -- -- @since 0.5.0 -sfoldl1' :: Stream r ix e => (e -> e -> e) -> Array r ix e -> e -sfoldl1' f = either throw id . sfoldl1M (\e -> pure . f e) +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' #-} -- | @@ -2368,11 +2468,15 @@ -- ==== __Examples__ -- -- @since 0.5.0 -sfoldl1M :: (Stream r ix e, MonadThrow m) => (e -> e -> m e) -> Array r ix e -> m e +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 - nullStream <- S.null str - when nullStream $ throwM $ SizeEmptyException (size arr) + isNullStream <- S.null str + when isNullStream $ throwM $ SizeEmptyException (outerSize arr) S.foldl1M f str {-# INLINE sfoldl1M #-} @@ -2381,7 +2485,11 @@ -- ==== __Examples__ -- -- @since 0.5.0 -sfoldl1M_ :: (Stream r ix e, MonadThrow m) => (e -> e -> m e) -> Array r ix e -> m () +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_ #-} @@ -2392,7 +2500,12 @@ -- ==== __Examples__ -- -- @since 0.5.0 -sifoldl :: Stream r ix e => (a -> ix -> e -> a) -> a -> Array r ix e -> a +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 #-} @@ -2401,7 +2514,12 @@ -- ==== __Examples__ -- -- @since 0.5.0 -sifoldlM :: (Stream r ix e, Monad m) => (a -> ix -> e -> m a) -> a -> Array r ix e -> m a +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 #-} @@ -2410,7 +2528,12 @@ -- ==== __Examples__ -- -- @since 0.5.0 -sifoldlM_ :: (Stream r ix e, Monad m) => (a -> ix -> e -> m a) -> a -> Array r ix e -> m () +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_ #-} @@ -2420,7 +2543,10 @@ -- ==== __Examples__ -- -- @since 0.5.0 -sor :: Stream r ix Bool => Array r ix Bool -> Bool +sor :: + forall r ix. Stream r ix Bool + => Array r ix Bool + -> Bool sor = S.unId . S.or . toStream {-# INLINE sor #-} @@ -2430,7 +2556,7 @@ -- ==== __Examples__ -- -- @since 0.5.0 -sand :: Stream r ix Bool => Array r ix Bool -> Bool +sand :: forall r ix. Stream r ix Bool => Array r ix Bool -> Bool sand = S.unId . S.and . toStream {-# INLINE sand #-} @@ -2440,7 +2566,7 @@ -- ==== __Examples__ -- -- @since 0.5.0 -sany :: Stream r ix e => (e -> Bool) -> Array r ix e -> Bool +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 #-} @@ -2450,7 +2576,7 @@ -- ==== __Examples__ -- -- @since 0.5.0 -sall :: Stream r ix e => (e -> Bool) -> Array r ix e -> Bool +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 #-} @@ -2467,7 +2593,7 @@ -- 88 -- -- @since 0.5.0 -ssum :: (Num e, Stream r ix e) => Array r ix e -> e +ssum :: forall r ix e. (Num e, Stream r ix e) => Array r ix e -> e ssum = sfoldl (+) 0 {-# INLINE ssum #-} @@ -2482,7 +2608,7 @@ -- 10500 -- -- @since 0.5.0 -sproduct :: (Num e, Stream r ix e) => Array r ix e -> e +sproduct :: forall r ix e. (Num e, Stream r ix e) => Array r ix e -> e sproduct = sfoldl (*) 1 {-# INLINE sproduct #-} @@ -2496,11 +2622,9 @@ -- >>> import Data.Massiv.Vector as V -- >>> V.smaximum' $ V.sfromList [10, 3, 70, 5 :: Int] -- 70 --- >>> V.smaximum' (V.empty :: Vector D Int) --- *** Exception: SizeEmptyException: (Sz1 0) corresponds to an empty array -- -- @since 0.5.0 -smaximum' :: (Ord e, Stream r ix e) => Array r ix e -> e +smaximum' :: forall r ix e. (HasCallStack, Ord e, Stream r ix e) => Array r ix e -> e smaximum' = sfoldl1' max {-# INLINE smaximum' #-} @@ -2519,7 +2643,7 @@ -- Nothing -- -- @since 0.5.0 -smaximumM :: (Ord e, Stream r ix e, MonadThrow m) => Array r ix e -> m e +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 #-} @@ -2533,11 +2657,9 @@ -- >>> import Data.Massiv.Vector as V -- >>> V.sminimum' $ V.sfromList [10, 3, 70, 5 :: Int] -- 3 --- >>> V.sminimum' (V.empty :: Array D Ix2 Int) --- *** Exception: SizeEmptyException: (Sz (0 :. 0)) corresponds to an empty array -- -- @since 0.5.0 -sminimum' :: (Ord e, Stream r ix e) => Array r ix e -> e +sminimum' :: forall r ix e. (HasCallStack, Ord e, Stream r ix e) => Array r ix e -> e sminimum' = sfoldl1' min {-# INLINE sminimum' #-} @@ -2556,138 +2678,6 @@ -- Nothing -- -- @since 0.5.0 -sminimumM :: (Ord e, Stream r ix e, MonadThrow m) => Array r ix e -> m e +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 #-} - - --- | See `stake`. --- --- @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 #-} -{-# DEPRECATED takeS "In favor of `stake`" #-} - --- | See `sdrop`. --- --- @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 #-} -{-# DEPRECATED dropS "In favor of `sdrop`" #-} - --- | See `sunfoldr` --- --- @since 0.4.1 -unfoldr :: (s -> Maybe (e, s)) -> s -> Vector DS e -unfoldr = sunfoldr -{-# INLINE unfoldr #-} -{-# DEPRECATED unfoldr "In favor of `sunfoldr`" #-} - - --- | See `sunfoldrN` --- --- @since 0.4.1 -unfoldrN :: Sz1 -> (s -> Maybe (e, s)) -> s -> Vector DS e -unfoldrN = unfoldrN -{-# INLINE unfoldrN #-} -{-# DEPRECATED unfoldrN "In favor of `sunfoldrN`" #-} - - --- | See `sfilterM` --- --- @since 0.4.1 -filterM :: (S.Stream r ix e, Applicative f) => (e -> f Bool) -> Array r ix e -> f (Vector DS e) -filterM f arr = DSArray <$> S.filterA f (S.toStream arr) -{-# INLINE filterM #-} -{-# DEPRECATED filterM "In favor of `sfilterM`" #-} - --- | See `sfilter` --- --- @since 0.4.1 -filterS :: S.Stream r ix e => (e -> Bool) -> Array r ix e -> Array DS Ix1 e -filterS = sfilter -{-# INLINE filterS #-} -{-# DEPRECATED filterS "In favor of `sfilter`" #-} - - --- | See `smapMaybe` --- --- @since 0.4.1 -mapMaybeS :: S.Stream r ix a => (a -> Maybe b) -> Array r ix a -> Vector DS b -mapMaybeS = smapMaybe -{-# INLINE mapMaybeS #-} -{-# DEPRECATED mapMaybeS "In favor of `smapMaybe`" #-} - --- | See `scatMaybes` --- --- @since 0.4.4 -catMaybesS :: S.Stream r ix (Maybe a) => Array r ix (Maybe a) -> Vector DS a -catMaybesS = scatMaybes -{-# INLINE catMaybesS #-} -{-# DEPRECATED catMaybesS "In favor of `scatMaybes`" #-} - --- | See `smapMaybeM` --- --- @since 0.4.1 -mapMaybeM :: - (S.Stream r ix a, Applicative f) => (a -> f (Maybe b)) -> Array r ix a -> f (Vector DS b) -mapMaybeM = smapMaybeM -{-# INLINE mapMaybeM #-} -{-# DEPRECATED mapMaybeM "In favor of `smapMaybeM`" #-} - --- | See `traverseS` --- --- @since 0.4.5 -traverseS :: (S.Stream r ix a, Applicative f) => (a -> f b) -> Array r ix a -> f (Vector DS b) -traverseS = straverse -{-# INLINE traverseS #-} -{-# DEPRECATED traverseS "In favor of `straverse`" #-} - --- | See `simapMaybe` --- --- @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) $ A.makeArrayR D (getComp arr) (size arr) $ \ix -> (ix, unsafeIndex arr ix) -{-# INLINE imapMaybeS #-} -{-# DEPRECATED imapMaybeS "In favor of `simapMaybe`" #-} - --- | See `simapMaybeM` --- --- @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) $ A.makeArrayR D (getComp arr) (size arr) $ \ix -> (ix, unsafeIndex arr ix) -{-# INLINE imapMaybeM #-} -{-# DEPRECATED imapMaybeM "In favor of `simapMaybeM`" #-} - --- | 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 #-} -{-# DEPRECATED ifilterS "In favor of `sifilter`" #-} - - --- | Similar to `filterM`, but map with an index aware function. --- --- @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 #-} -{-# DEPRECATED ifilterM "In favor of `sifilterM`" #-} diff --git a/src/Data/Massiv/Vector/Stream.hs b/src/Data/Massiv/Vector/Stream.hs --- a/src/Data/Massiv/Vector/Stream.hs +++ b/src/Data/Massiv/Vector/Stream.hs @@ -1,11 +1,14 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE LambdaCase #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} -{-# OPTIONS_HADDOCK hide, not-home #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE TypeFamilies #-} {-# OPTIONS_GHC -fno-warn-orphans #-} +{-# OPTIONS_HADDOCK hide, not-home #-} -- | -- Module : Data.Massiv.Vector.Stream -- Copyright : (c) Alexey Kuleshevich 2019-2021 @@ -123,7 +126,6 @@ , transSteps , transStepsId -- * Useful re-exports - , module Data.Vector.Fusion.Bundle.Size , module Data.Vector.Fusion.Util , Id(..) ) where @@ -132,24 +134,25 @@ import Control.Monad.ST import qualified Data.Foldable as F import Data.Massiv.Core.Common hiding (empty, singleton, replicate) +import Data.Coerce import Data.Maybe (catMaybes) 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.Bundle.Size as B import qualified Data.Vector.Fusion.Stream.Monadic as S import Data.Vector.Fusion.Util import Prelude hiding (and, concatMap, drop, filter, foldl, foldl1, foldr, foldr1, length, map, mapM, mapM_, null, or, replicate, take, traverse, zipWith, zipWith3) - +import qualified GHC.Exts (IsList(..)) 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 - Exact n -> str {stepsStream = S.replicate n e} - _ -> fmap (const e) str + LengthExact n -> str {stepsStream = S.replicate (coerce n) e} + _ -> fmap (const e) str {-# INLINE (<$) #-} instance Monad m => Semigroup (Steps m e) where @@ -163,6 +166,15 @@ {-# INLINE mappend #-} +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 #-} @@ -191,8 +203,8 @@ -- 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) +steps :: forall r ix e m . (Monad m, Index ix, Source r e) => Array r ix e -> Steps m e +steps arr = k `seq` arr `seq` Steps (S.Stream step 0) (LengthExact (coerce k)) where k = totalElem $ size arr step i @@ -204,8 +216,8 @@ {-# 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) +isteps :: forall r ix e m . (Monad m, Index ix, Source r e) => Array r ix e -> Steps m (ix, e) +isteps arr = k `seq` arr `seq` Steps (S.Stream step 0) (LengthExact (coerce k)) where sz = size arr k = totalElem sz @@ -217,39 +229,39 @@ {-# INLINE step #-} {-# INLINE isteps #-} -toBundle :: (Monad m, Source r ix e) => Array r ix e -> B.Bundle m v e +toBundle :: (Monad m, Index ix, Source r e) => Array r ix e -> B.Bundle m v e toBundle arr = let Steps str k = steps arr - in B.fromStream str k + in B.fromStream str (sizeHintToBundleSize k) {-# INLINE toBundle #-} -fromBundle :: Mutable r Ix1 e => B.Bundle Id v e -> Array r Ix1 e +fromBundle :: Manifest r e => B.Bundle Id v e -> Vector r 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 :: (Monad m, Manifest r e) => B.Bundle m v e -> m (Vector r 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 :: forall r e . Manifest r e => B.Size -> S.Stream Id e -> Vector r e fromStream sz str = - case upperBound sz of + case B.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 :: 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 upperBound sz of + 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, Mutable r ix e) + forall r ix e m. (Monad m, Manifest r e, Index ix) => Sz ix -> S.Stream m e -> m (Array r ix e) @@ -260,25 +272,25 @@ unstreamIntoM :: - (Mutable r Ix1 a, PrimMonad m) - => MArray (PrimState m) r Ix1 a - -> Size + (Manifest r a, PrimMonad m) + => MVector (PrimState m) r a + -> LengthHint -> S.Stream Id a - -> m (MArray (PrimState m) r Ix1 a) + -> m (MVector (PrimState m) r a) unstreamIntoM marr sz str = case sz of - Exact _ -> marr <$ unstreamMaxM marr str - Max _ -> unsafeLinearShrink marr . SafeSz =<< unstreamMaxM marr str - Unknown -> unstreamUnknownM marr str + LengthExact _ -> marr <$ unstreamMaxM marr str + LengthMax _ -> unsafeLinearShrink marr . SafeSz =<< unstreamMaxM marr str + LengthUnknown -> unstreamUnknownM marr str {-# INLINE unstreamIntoM #-} unstreamMax :: - forall r e. (Mutable r Ix1 e) + forall r e. (Manifest r e) => Int -> S.Stream Id e - -> Array r Ix1 e + -> Vector r e unstreamMax kMax str = runST $ do marr <- unsafeNew (SafeSz kMax) @@ -288,7 +300,7 @@ unstreamMaxM :: - (Mutable r ix a, PrimMonad m) => MArray (PrimState m) r ix a -> S.Stream Id a -> m Int + (Manifest r a, Index ix, 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 = @@ -302,7 +314,7 @@ {-# INLINE unstreamMaxM #-} -unstreamUnknown :: Mutable r Ix1 a => S.Stream Id a -> Array r Ix1 a +unstreamUnknown :: Manifest r a => S.Stream Id a -> Vector r a unstreamUnknown str = runST $ do marr <- unsafeNew zeroSz @@ -311,11 +323,11 @@ unstreamUnknownM :: - (Mutable r Ix1 a, PrimMonad m) - => MArray (PrimState m) r Ix1 a + (Manifest r a, PrimMonad m) + => MVector (PrimState m) r 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 + -> m (MVector (PrimState m) r a) +unstreamUnknownM marrInit (S.Stream step s) = stepLoad s 0 (unSz (sizeOfMArray marrInit)) marrInit where stepLoad t i kMax marr | i < kMax = @@ -334,7 +346,7 @@ unstreamExact :: - forall r ix e. (Mutable r ix e) + forall r ix e. (Manifest r e, Index ix) => Sz ix -> S.Stream Id e -> Array r ix e @@ -348,28 +360,28 @@ length :: Monad m => Steps m a -> m Int length (Steps str sz) = case sz of - Exact k -> pure k - _ -> S.length str + 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 - Exact k -> pure (k == 0) - _ -> S.null str + LengthExact k -> pure (k == zeroSz) + _ -> S.null str {-# INLINE null #-} empty :: Monad m => Steps m e -empty = Steps S.empty (Exact 0) +empty = Steps S.empty (LengthExact zeroSz) {-# INLINE empty #-} singleton :: Monad m => e -> Steps m e -singleton e = Steps (S.singleton e) (Exact 1) +singleton e = Steps (S.singleton e) (LengthExact oneSz) {-# INLINE singleton #-} -generate :: Monad m => Int -> (Int -> e) -> Steps m e -generate k f = Steps (S.generate k f) (Exact k) +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 @@ -387,16 +399,16 @@ cons :: Monad m => e -> Steps m e -> Steps m e -cons e (Steps str k) = Steps (S.cons e str) (k + 1) +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 -> (\x -> (x, drop 1 sts)) <$> mx) <$> headMaybe sts +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 + 1) +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) @@ -404,7 +416,7 @@ {-# 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) +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 @@ -424,18 +436,18 @@ {-# 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) (smaller ka kb) +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) (smaller ka (smaller kb 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) (smaller ka (smaller kb (smaller kc kd))) + Steps (S.zipWith4 f sa sb sc sd) (minLengthHint ka (minLengthHint kb (minLengthHint kc kd))) {-# INLINE zipWith4 #-} zipWith5 :: @@ -448,7 +460,7 @@ -> 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) (smaller ka (smaller kb (smaller kc (smaller kd ke)))) + Steps (S.zipWith5 f sa sb sc sd se) (minLengthHint ka (minLengthHint kb (minLengthHint kc (minLengthHint kd ke)))) {-# INLINE zipWith5 #-} zipWith6 :: @@ -464,17 +476,17 @@ 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) - (smaller ka (smaller kb (smaller kc (smaller kd (smaller ke kf))))) + (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) (smaller ka kb) +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) (smaller ka (smaller kb kc)) + Steps (S.zipWith3M f sa sb sc) (minLengthHint ka (minLengthHint kb kc)) {-# INLINE zipWith3M #-} zipWith4M :: @@ -486,7 +498,7 @@ -> 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) (smaller ka (smaller kb (smaller kc kd))) + Steps (S.zipWith4M f sa sb sc sd) (minLengthHint ka (minLengthHint kb (minLengthHint kc kd))) {-# INLINE zipWith4M #-} zipWith5M :: @@ -499,7 +511,7 @@ -> 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) (smaller ka (smaller kb (smaller kc (smaller kd ke)))) + Steps (S.zipWith5M f sa sb sc sd se) (minLengthHint ka (minLengthHint kb (minLengthHint kc (minLengthHint kd ke)))) {-# INLINE zipWith5M #-} zipWith6M :: @@ -515,7 +527,7 @@ 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) - (smaller ka (smaller kb (smaller kc (smaller kd (smaller ke kf))))) + (minLengthHint ka (minLengthHint kb (minLengthHint kc (minLengthHint kd (minLengthHint ke kf))))) {-# INLINE zipWith6M #-} @@ -582,10 +594,10 @@ {-# INLINE transStepsId #-} transSteps :: (Monad m, Monad n) => Steps m e -> m (Steps n e) -transSteps (Steps strM sz@(Exact _)) = (`Steps` sz) <$> transListM strM +transSteps (Steps strM sz@(LengthExact _)) = (`Steps` sz) <$> transListM strM transSteps (Steps strM _) = do (n, strN) <- transListNM strM - pure (Steps strN (Exact n)) + pure (Steps strN (LengthExact n)) {-# INLINE transSteps #-} @@ -655,20 +667,20 @@ mapMaybe :: Monad m => (a -> Maybe e) -> Steps m a -> Steps m e -mapMaybe f (Steps str k) = Steps (S.mapMaybe f str) (toMax k) +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) Unknown +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` toMax k) <$> liftListA (mapMaybeListA f) str +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) (toMax k) +mapMaybeM f (Steps str k) = Steps (mapMaybeStreamM f str) (toLengthMax k) {-# INLINE mapMaybeM #-} mapMaybeListA :: Applicative f => (a -> f (Maybe b)) -> [a] -> f [b] @@ -693,87 +705,87 @@ {-# 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) +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` toMax k) <$> liftListA (M.filterM f) str +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) (toMax k) +filterM f (Steps str k) = Steps (S.filterM f str) (toLengthMax k) {-# INLINE filterM #-} -take :: Monad m => Int -> Steps m a -> Steps m a +take :: Monad m => Sz1 -> Steps m a -> Steps m a take n (Steps str sz) = - Steps (S.take n str) $! + Steps (S.take (coerce n) str) $! case sz of - Exact k -> Exact (min n k) - Max k -> Max (min n k) - Unknown -> Unknown + LengthExact k -> LengthExact (inline0 min n k) + LengthMax k -> LengthMax (inline0 min n k) + LengthUnknown -> LengthUnknown {-# 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) +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 -> Int -> Steps m a -> Steps m a -slice i k (Steps str _) = Steps (S.slice i k str) (Max k) +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 => Int -> (a -> a) -> a -> Steps m a -iterateN n f a = Steps (S.iterateN n f a) (Exact n) +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 => Int -> (a -> m a) -> a -> Steps m a -iterateNM n f a = Steps (S.iterateNM n f a) (Exact n) +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 => Int -> a -> Steps m a -replicate n a = Steps (S.replicate n a) (Exact n) +replicate :: Monad m => Sz1 -> a -> Steps m a +replicate n a = Steps (S.replicate (coerce n) a) (LengthExact n) {-# INLINE replicate #-} -replicateM :: Monad m => Int -> m a -> Steps m a -replicateM n f = Steps (S.replicateM n f) (Exact n) +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 => Int -> (Int -> m a) -> Steps m a -generateM n f = Steps (S.generateM n f) (Exact n) +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) Unknown +unfoldr f e0 = Steps (S.unfoldr f e0) LengthUnknown {-# INLINE unfoldr #-} -unfoldrN :: Monad m => Int -> (s -> Maybe (e, s)) -> s -> Steps m e -unfoldrN n f e0 = Steps (S.unfoldrN n f e0) Unknown +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 => Int -> (s -> Maybe (e, s)) -> s -> Steps m e -unsafeUnfoldrN n f e0 = Steps (S.unfoldrN n f e0) (Max n) +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) Unknown +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) Unknown +unfoldrNM n f e0 = Steps (S.unfoldrNM n f e0) LengthUnknown {-# INLINE unfoldrNM #-} -unsafeUnfoldrNM :: Monad m => Int -> (s -> m (Maybe (e, s))) -> s -> Steps m e -unsafeUnfoldrNM n f e0 = Steps (S.unfoldrNM n f e0) (Max n) +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 => Int -> (s -> (a, s)) -> s -> Steps m a +unfoldrExactN :: Monad m => Sz1 -> (s -> (a, s)) -> s -> Steps m a unfoldrExactN n f = unfoldrExactNM n (pure . f) {-# INLINE unfoldrExactN #-} -unfoldrExactNM :: Monad m => Int -> (s -> m (a, s)) -> s -> Steps m a -unfoldrExactNM n f t = Steps (S.Stream step (t, n)) (Exact n) +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 @@ -782,8 +794,8 @@ {-# INLINE unfoldrExactNM #-} -enumFromStepN :: (Num a, Monad m) => a -> a -> Int -> Steps m a -enumFromStepN x step k = Steps (S.enumFromStepN x step k) (Exact k) +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 #-} @@ -794,15 +806,15 @@ {-# INLINE toList #-} fromList :: Monad m => [e] -> Steps m e -fromList = (`Steps` Unknown) . S.fromList +fromList = (`Steps` LengthUnknown) . S.fromList {-# INLINE fromList #-} fromListN :: Monad m => Int -> [e] -> Steps m e -fromListN n = (`Steps` Unknown) . S.fromListN n +fromListN n = (`Steps` LengthUnknown) . S.fromListN n {-# INLINE fromListN #-} -unsafeFromListN :: Monad m => Int -> [e] -> Steps m e -unsafeFromListN n = (`Steps` Max n) . S.fromListN n +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) @@ -816,14 +828,76 @@ pure $ S.fromList xs {-# INLINE transListM #-} -transListNM :: (Monad m, Monad n) => S.Stream m a -> m (Int, S.Stream n a) +transListNM :: (Monad m, Monad n) => S.Stream m a -> m (Sz1, S.Stream n a) transListNM str = do (n, xs) <- toListN str - pure (n, S.fromList xs) + 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 #-} + + +sizeHintToBundleSize :: LengthHint -> B.Size +sizeHintToBundleSize = + \case + LengthExact k -> B.Exact (coerce k) + LengthMax k -> B.Max (coerce k) + LengthUnknown -> B.Unknown +{-# INLINE sizeHintToBundleSize #-} + +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 #-} + diff --git a/src/Data/Massiv/Vector/Unsafe.hs b/src/Data/Massiv/Vector/Unsafe.hs --- a/src/Data/Massiv/Vector/Unsafe.hs +++ b/src/Data/Massiv/Vector/Unsafe.hs @@ -56,14 +56,14 @@ -- | -- -- @since 0.5.0 -unsafeHead :: Source r Ix1 e => Vector r e -> e +unsafeHead :: Source r e => Vector r e -> e unsafeHead = (`unsafeLinearIndex` 0) {-# INLINE unsafeHead #-} -- | -- -- @since 0.5.0 -unsafeLast :: Source r Ix1 e => Vector r e -> e +unsafeLast :: Source r e => Vector r e -> e unsafeLast v = unsafeLinearIndex v (max 0 (unSz (size v) - 1)) {-# INLINE unsafeLast #-} @@ -74,7 +74,7 @@ -- | -- -- @since 0.5.0 -unsafeIndexM :: (Source r Ix1 e, Monad m) => Vector r e -> Ix1 -> m e +unsafeIndexM :: (Source r e, Monad m) => Vector r e -> Ix1 -> m e unsafeIndexM v i = pure $! unsafeLinearIndex v i {-# INLINE unsafeIndexM #-} @@ -82,14 +82,14 @@ -- | -- -- @since 0.5.0 -unsafeHeadM :: Monad m => Source r Ix1 e => Vector r e -> m e +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 Ix1 e => Vector r e -> m e +unsafeLastM :: (Monad m, Source r e) => Vector r e -> m e unsafeLastM v = pure $! unsafeLast v {-# INLINE unsafeLastM #-} @@ -102,7 +102,7 @@ -- | -- -- @since 0.5.0 -unsafeInit :: Source r Ix1 e => Vector r e -> Vector r e +unsafeInit :: Source r e => Vector r e -> Vector r e unsafeInit v = unsafeLinearSlice 0 (SafeSz (coerce (size v) - 1)) v {-# INLINE unsafeInit #-} @@ -110,22 +110,22 @@ -- | -- -- @since 0.5.0 -unsafeTail :: Source r Ix1 e => Vector r e -> Vector r e -unsafeTail = unsafeDrop 1 +unsafeTail :: Source r e => Vector r e -> Vector r e +unsafeTail = unsafeDrop oneSz {-# INLINE unsafeTail #-} -- | -- -- @since 0.5.0 -unsafeTake :: Source r Ix1 e => Sz1 -> Vector r e -> Vector r e +unsafeTake :: Source r e => Sz1 -> Vector r e -> Vector r e unsafeTake = unsafeLinearSlice 0 {-# INLINE unsafeTake #-} -- | -- -- @since 0.5.0 -unsafeDrop :: Source r Ix1 e => Sz1 -> Vector r e -> Vector r e +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 #-} @@ -139,7 +139,7 @@ -- -- @since 0.5.1 unsafeFromListN :: Sz1 -> [e] -> Vector DS e -unsafeFromListN (Sz n) = fromSteps . S.unsafeFromListN n +unsafeFromListN n = fromSteps . S.unsafeFromListN n {-# INLINE unsafeFromListN #-} -- | /O(n)/ - Right unfolding function with at most @n@ number of elements. @@ -158,7 +158,7 @@ -- is reached. -> s -- ^ Inititial element. -> Vector DS e -unsafeUnfoldrN (Sz n) f = DSArray . S.unsafeUnfoldrN n f +unsafeUnfoldrN n f = DSArray . S.unsafeUnfoldrN n f {-# INLINE unsafeUnfoldrN #-} @@ -172,5 +172,5 @@ -- -- @since 0.5.1 unsafeUnfoldrNM :: Monad m => Sz1 -> (s -> m (Maybe (e, s))) -> s -> m (Vector DS e) -unsafeUnfoldrNM (Sz n) f = fromStepsM . S.unsafeUnfoldrNM n f +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,7 +1,7 @@ {-# LANGUAGE CPP #-} module Main where -#if __GLASGOW_HASKELL__ >= 802 && __GLASGOW_HASKELL__ < 810 +#if __GLASGOW_HASKELL__ >= 802 && __GLASGOW_HASKELL__ != 810 import Test.DocTest (doctest) @@ -12,6 +12,6 @@ -- TODO: fix doctest support main :: IO () -main = putStrLn "\nDoctests are not supported for older ghc version\n" +main = putStrLn "\nDoctests are not supported for ghc version 8.2 and prior as well as 8.10\n" #endif