diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,9 @@ +# 1.0.3 + +* Deprecated `indexWith` in favor of `indexAssert` +* Addition of scans: `sscanl`, `sscanl1`, `sprescanl`, `spostscanl` and `spostscanlAcc` +* Expose `unsafePrefIndex` + # 1.0.2 * Addition of `Iterator` type class and related fucntions: diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -13,9 +13,9 @@ | Package | Hackage | Nightly | LTS | |:-------------------|:-------:|:-------:|:---:| -| [`massiv`](https://github.com/lehins/massiv/tree/master/massiv)| [![Hackage](https://img.shields.io/hackage/v/massiv.svg)](https://hackage.haskell.org/package/massiv)| [![Nightly](https://www.stackage.org/package/massiv/badge/nightly)](https://www.stackage.org/nightly/package/massiv)| ](https://www.stackage.org/package/massiv-io/badge/lts)](https://www.stackage.org/lts/package/massiv-io)| -| [`massiv-test`](https://github.com/lehins/massiv/tree/master/massiv-test)| [![Hackage](https://img.shields.io/hackage/v/massiv-test.svg)](https://hackage.haskell.org/package/massiv-test)| [![Nightly](https://www.stackage.org/package/massiv-test/badge/nightly)](https://www.stackage.org/nightly/package/massiv-test)| [![Nightly](https://www.stackage.org/package/massiv-test/badge/lts)](https://www.stackage.org/lts/package/massiv-test)| -| [`haskell-scheduler`](https://github.com/lehins/haskell-scheduler)| [![Hackage](https://img.shields.io/hackage/v/scheduler.svg)](https://hackage.haskell.org/package/scheduler)| [![Nightly](https://www.stackage.org/package/scheduler/badge/nightly)](https://www.stackage.org/nightly/package/scheduler)| [![Nightly](https://www.stackage.org/package/scheduler/badge/lts)](https://www.stackage.org/lts/package/scheduler)| +| [`massiv`](https://github.com/lehins/massiv/tree/master/massiv)| [![Hackage](https://img.shields.io/hackage/v/massiv.svg)](https://hackage.haskell.org/package/massiv)| [![Nightly](https://www.stackage.org/package/massiv/badge/nightly)](https://www.stackage.org/nightly/package/massiv)| [![LTS](https://www.stackage.org/package/massiv-io/badge/lts)](https://www.stackage.org/lts/package/massiv-io)| +| [`massiv-test`](https://github.com/lehins/massiv/tree/master/massiv-test)| [![Hackage](https://img.shields.io/hackage/v/massiv-test.svg)](https://hackage.haskell.org/package/massiv-test)| [![Nightly](https://www.stackage.org/package/massiv-test/badge/nightly)](https://www.stackage.org/nightly/package/massiv-test)| [![LTS](https://www.stackage.org/package/massiv-test/badge/lts)](https://www.stackage.org/lts/package/massiv-test)| +| [`haskell-scheduler`](https://github.com/lehins/haskell-scheduler)| [![Hackage](https://img.shields.io/hackage/v/scheduler.svg)](https://hackage.haskell.org/package/scheduler)| [![Nightly](https://www.stackage.org/package/scheduler/badge/nightly)](https://www.stackage.org/nightly/package/scheduler)| [![LTS](https://www.stackage.org/package/scheduler/badge/lts)](https://www.stackage.org/lts/package/scheduler)| ## Introduction @@ -94,7 +94,7 @@ go ahead and square its elements ```haskell -λ> makeVectorR D Seq 10 id +λ> vec = makeVectorR D Seq 10 id λ> evaluateM vec 4 4 λ> vec2 = A.map (^ (2 :: Int)) vec diff --git a/include/massiv.h b/include/massiv.h --- a/include/massiv.h +++ b/include/massiv.h @@ -3,9 +3,9 @@ #define MASSIV_INCLUDE #if MASSIV_UNSAFE_CHECKS -#define INDEX_CHECK(name, s, f) (indexWith __FILE__ __LINE__ (name) (s) (f)) +#define HAS_CALL_STACK (HasCallStack) #else -#define INDEX_CHECK(name, s, f) ((f)) +#define HAS_CALL_STACK () #endif #endif diff --git a/massiv.cabal b/massiv.cabal --- a/massiv.cabal +++ b/massiv.cabal @@ -1,5 +1,5 @@ name: massiv -version: 1.0.2.0 +version: 1.0.3.0 synopsis: Massiv (Массив) is an Array Library. description: Multi-dimensional Arrays with fusion, stencils and parallel computation. homepage: https://github.com/lehins/massiv @@ -20,7 +20,8 @@ , GHC == 8.8.4 , GHC == 8.10.7 , GHC == 9.0.2 - , GHC == 9.2.4 + , GHC == 9.2.5 + , GHC == 9.4.3 flag unsafe-checks description: Enable all the bounds checks for unsafe functions at the cost of @@ -88,6 +89,7 @@ , random >= 1.2.0 , unliftio-core , vector >= 0.12 + , vector-stream include-dirs: include install-includes: massiv.h diff --git a/src/Data/Massiv/Array.hs b/src/Data/Massiv/Array.hs --- a/src/Data/Massiv/Array.hs +++ b/src/Data/Massiv/Array.hs @@ -1,4 +1,5 @@ {-# OPTIONS_GHC -fno-warn-duplicate-exports #-} + -- | -- Module : Data.Massiv.Array -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -76,85 +77,98 @@ -- import Prelude as P -- import Data.Massiv.Array as A -- @ --- -module Data.Massiv.Array - ( -- * Construct - module Data.Massiv.Array.Ops.Construct +module Data.Massiv.Array ( + -- * Construct + module Data.Massiv.Array.Ops.Construct, + -- * Compute - , getComp - , setComp - , appComp - , compute - , computeS - , computeP - , computeIO - , computePrimM - , computeAs - , computeProxy - , computeSource - , computeWithStride - , computeWithStrideAs - , clone - , convert - , convertAs - , convertProxy - , fromRaggedArrayM - , fromRaggedArray' + getComp, + setComp, + appComp, + compute, + computeS, + computeP, + computeIO, + computePrimM, + computeAs, + computeProxy, + computeSource, + computeWithStride, + computeWithStrideAs, + clone, + convert, + convertAs, + convertProxy, + fromRaggedArrayM, + fromRaggedArray', + -- * Vector - , module Data.Massiv.Vector + module Data.Massiv.Vector, + -- * Size - , size - , elemsCount - , isEmpty - , isNotEmpty - , isNull - , isNotNull + size, + elemsCount, + isEmpty, + isNotEmpty, + isNull, + isNotNull, + -- * Indexing - , (!?) - , (!) - , (??) - , indexM - , index - , index' - , defaultIndex - , borderIndex - , evaluateM - , evaluate' + (!?), + (!), + (??), + indexM, + index, + index', + defaultIndex, + borderIndex, + evaluateM, + evaluate', + -- * Mapping - , module Data.Massiv.Array.Ops.Map - -- * Folding + module Data.Massiv.Array.Ops.Map, + -- * Folding -- $folding + module Data.Massiv.Array.Ops.Fold, - , module Data.Massiv.Array.Ops.Fold -- * Transforming - , module Data.Massiv.Array.Ops.Transform + module Data.Massiv.Array.Ops.Transform, + -- * Slicing - , module Data.Massiv.Array.Ops.Slice + module Data.Massiv.Array.Ops.Slice, + -- * Algorithms + -- ** Sorting - , quicksort - , quicksortBy - , quicksortByM - , tally + quicksort, + quicksortBy, + quicksortByM, + tally, + -- ** Iterations - , iterateUntil + iterateUntil, + -- * Conversion - , module Data.Massiv.Array.Manifest.List + module Data.Massiv.Array.Manifest.List, + -- * Mutable - , module Data.Massiv.Array.Mutable + module Data.Massiv.Array.Mutable, + -- * Core - , module Data.Massiv.Core + module Data.Massiv.Core, + -- * Representations - , module Data.Massiv.Array.Delayed - , module Data.Massiv.Array.Manifest + module Data.Massiv.Array.Delayed, + module Data.Massiv.Array.Manifest, + -- * Stencil - , module Data.Massiv.Array.Stencil + module Data.Massiv.Array.Stencil, + -- * Numeric Operations - , module Data.Massiv.Array.Numeric - ) where + module Data.Massiv.Array.Numeric, +) where -import Data.Massiv.Vector import Data.Massiv.Array.Delayed import Data.Massiv.Array.Manifest import Data.Massiv.Array.Manifest.Internal @@ -170,12 +184,26 @@ import Data.Massiv.Array.Stencil import Data.Massiv.Core import Data.Massiv.Core.Common -import Prelude as P hiding (all, and, any, enumFromTo, foldl, foldr, mapM, - mapM_, maximum, minimum, or, product, replicate, splitAt, - sum, zip) - -{- $folding - -All folding is done in a row-major order. +import Data.Massiv.Vector +import Prelude as P hiding ( + all, + and, + any, + enumFromTo, + foldl, + foldr, + mapM, + mapM_, + maximum, + minimum, + or, + product, + replicate, + splitAt, + sum, + zip, + ) --} +-- $folding +-- +-- All folding is done in a row-major order. diff --git a/src/Data/Massiv/Array/Delayed.hs b/src/Data/Massiv/Array/Delayed.hs --- a/src/Data/Massiv/Array/Delayed.hs +++ b/src/Data/Massiv/Array/Delayed.hs @@ -5,37 +5,41 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Delayed - ( -- * Delayed +module Data.Massiv.Array.Delayed ( + -- * Delayed + -- ** Delayed Pull Array - D(..) - , delay - , liftArray2' - , liftArray2M + D (..), + delay, + liftArray2', + liftArray2M, + -- ** Delayed Push Array - , DL(..) - , toLoadArray - , makeLoadArrayS - , makeLoadArray - , fromStrideLoad + DL (..), + toLoadArray, + makeLoadArrayS, + makeLoadArray, + fromStrideLoad, + -- ** Delayed Stream Array - , DS(..) - , toStreamArray - , toSteps - , fromSteps + DS (..), + toStreamArray, + toSteps, + fromSteps, + -- ** Delayed Interleaved Array - , DI(..) - , toInterleaved - , fromInterleaved + DI (..), + toInterleaved, + fromInterleaved, + -- ** Delayed Windowed Array - , DW(..) - , Window(..) - , insertWindow - , getWindow - , dropWindow - , makeWindowedArray - ) where + DW (..), + Window (..), + insertWindow, + getWindow, + dropWindow, + makeWindowedArray, +) where import Data.Massiv.Array.Delayed.Interleaved import Data.Massiv.Array.Delayed.Pull diff --git a/src/Data/Massiv/Array/Delayed/Interleaved.hs b/src/Data/Massiv/Array/Delayed/Interleaved.hs --- a/src/Data/Massiv/Array/Delayed/Interleaved.hs +++ b/src/Data/Massiv/Array/Delayed/Interleaved.hs @@ -5,6 +5,7 @@ {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UndecidableInstances #-} + -- | -- Module : Data.Massiv.Array.Delayed.Interleaved -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -12,19 +13,17 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Delayed.Interleaved - ( DI(..) - , Array(..) - , toInterleaved - , fromInterleaved - ) where +module Data.Massiv.Array.Delayed.Interleaved ( + DI (..), + Array (..), + toInterleaved, + fromInterleaved, +) where import Data.Massiv.Array.Delayed.Pull import Data.Massiv.Core.Common import Data.Massiv.Core.List (L, showArrayList, showsArrayPrec) - -- | Delayed array that will be loaded in an interleaved fashion during parallel -- computation. -- @@ -33,14 +32,15 @@ newtype instance Array DI ix e = DIArray { diArray :: Array D ix e - } deriving (Eq, Ord, Functor, Applicative, Foldable) + } + 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 Strategy DI where - setComp c arr = arr { diArray = (diArray arr) { dComp = c } } + setComp c arr = arr{diArray = (diArray arr){dComp = c}} {-# INLINE setComp #-} getComp = dComp . diArray {-# INLINE getComp #-} @@ -50,14 +50,12 @@ maxLinearSize = Just . SafeSz . elemsCount {-# INLINE maxLinearSize #-} - instance Size DI where size (DIArray arr) = size arr {-# INLINE size #-} unsafeResize sz = DIArray . unsafeResize sz . diArray {-# INLINE unsafeResize #-} - instance Index ix => Load DI ix e where makeArray c sz = DIArray . makeArray c sz {-# INLINE makeArray #-} @@ -78,7 +76,6 @@ -- | 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 :: (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 @@ -17,20 +17,20 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable -module Data.Massiv.Array.Delayed.Pull - ( D(..) - , Array(..) - , delay - , eqArrays - , compareArrays - , imap - , liftArray2' - , liftArray2M - , unsafeExtract - , unsafeSlice - , unsafeInnerSlice - , zipWithInternal - ) where +module Data.Massiv.Array.Delayed.Pull ( + D (..), + Array (..), + delay, + eqArrays, + compareArrays, + imap, + liftArray2', + liftArray2M, + unsafeExtract, + unsafeSlice, + unsafeInnerSlice, + zipWithInternal, +) where import Control.Applicative import qualified Data.Foldable as F @@ -45,16 +45,15 @@ #include "massiv.h" -- | Delayed representation. -data D = - D +data D + = D deriving (Show) - -data instance Array D ix e = - DArray { dComp :: !Comp - , dSize :: !(Sz ix) - , dPrefIndex :: !(PrefIndex ix e) - } +data instance Array D ix e = DArray + { dComp :: !Comp + , dSize :: !(Sz ix) + , dPrefIndex :: !(PrefIndex ix e) + } instance (Ragged L ix e, Show e) => Show (Array D ix e) where showsPrec = showsArrayPrec id @@ -72,7 +71,7 @@ {-# INLINE unsafeResize #-} instance Strategy D where - setComp c arr = arr {dComp = c} + setComp c arr = arr{dComp = c} {-# INLINE setComp #-} getComp = dComp {-# INLINE getComp #-} @@ -106,8 +105,8 @@ {-# INLINE unsafeExtract #-} -- | /O(1)/ - Take a slice out of an array from within -unsafeSlice :: - (Source r e, Index ix, Index (Lower ix), MonadThrow m) +unsafeSlice + :: (Source r e, Index ix, Index (Lower ix), MonadThrow m) => Array r ix e -> ix -> Sz ix @@ -119,8 +118,8 @@ {-# INLINE unsafeSlice #-} -- | /O(1)/ - Take a slice out of an array from the inside -unsafeInnerSlice :: - (Source r e, Index ix) => Array r ix e -> Sz (Lower ix) -> Int -> Array D (Lower ix) e +unsafeInnerSlice + :: (Source r e, Index ix) => Array r ix e -> Sz (Lower ix) -> Int -> Array D (Lower ix) e unsafeInnerSlice !arr szL !i = DArray (getComp arr) szL $ PrefIndex (unsafeIndex arr . (`snocDim` i)) {-# INLINE unsafeInnerSlice #-} @@ -177,7 +176,7 @@ {-# INLINE makeArray #-} makeArrayLinear comp sz = DArray comp sz . PrefIndexLinear {-# INLINE makeArrayLinear #-} - iterArrayLinearST_ !scheduler DArray {..} uWrite = + iterArrayLinearST_ !scheduler DArray{..} uWrite = case dPrefIndex of PrefIndex f -> iterTargetFullST_ defRowMajor scheduler 0 dSize $ \ !i -> uWrite i . f @@ -186,7 +185,7 @@ {-# INLINE iterArrayLinearST_ #-} instance Index ix => StrideLoad D ix e where - iterArrayLinearWithStrideST_ !scheduler !stride sz DArray {..} uWrite = + iterArrayLinearWithStrideST_ !scheduler !stride sz DArray{..} uWrite = case dPrefIndex of PrefIndex f -> iterTargetFullWithStrideST_ defRowMajor scheduler 0 sz stride $ \i -> @@ -205,8 +204,9 @@ -- | Map an index aware function over an array -- -- @since 0.1.0 -imap :: - forall r ix e a. (Index ix, Source r e) +imap + :: forall r ix e a + . (Index ix, Source r e) => (ix -> e -> a) -> Array r ix e -> Array D ix a @@ -228,7 +228,7 @@ {-# INLINE foldArray #-} instance Num e => Numeric D e where - unsafeLiftArray f arr = arr {dPrefIndex = f <$> dPrefIndex arr} + unsafeLiftArray f arr = arr{dPrefIndex = f <$> dPrefIndex arr} {-# INLINE unsafeLiftArray #-} unsafeLiftArray2 f a1 a2 = zipWithInternal (size a1) f a1 a2 {-# INLINE unsafeLiftArray2 #-} @@ -242,16 +242,17 @@ PrefIndex gix -> makeArray (getComp arr) (size arr) gix PrefIndexLinear gi -> makeArrayLinear (getComp arr) (size arr) gi {-# INLINE [1] delay #-} + {-# RULES -"delay" [~1] forall (arr :: Array D ix e) . delay arr = arr - #-} +"delay" [~1] forall (arr :: Array D ix e). delay arr = arr + #-} -- | Compute array equality by applying a comparing function to each -- element. Empty arrays are always equal, regardless of their size. -- -- @since 0.5.7 -eqArrays :: - (Index ix, Source r1 e1, Source r2 e2) +eqArrays + :: (Index ix, Source r1 e1, Source r2 e2) => (e1 -> e2 -> Bool) -> Array r1 ix e1 -> Array r2 ix e2 @@ -259,13 +260,16 @@ eqArrays f arr1 arr2 = let sz1 = size arr1 sz2 = size arr2 - in (sz1 == sz2 && - not - (A.any - not - (makeArray @D (getComp arr1 <> getComp arr2) (size arr1) $ \ix -> - f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)))) || - (isZeroSz sz1 && isZeroSz sz2) + in ( sz1 == sz2 + && not + ( A.any + not + ( makeArray @D (getComp arr1 <> getComp arr2) (size arr1) $ \ix -> + f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix) + ) + ) + ) + || (isZeroSz sz1 && isZeroSz sz2) {-# INLINE eqArrays #-} -- | Compute array ordering by applying a comparing function to each element. @@ -273,25 +277,26 @@ -- you need an ordering but do not care about which one is used. -- -- @since 0.5.7 -compareArrays :: - (Index ix, Source r1 e1, Source r2 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) <> - A.fold - (makeArray @D (getComp arr1 <> getComp arr2) (size arr1) $ \ix -> - f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)) + compare (size arr1) (size arr2) + <> A.fold + ( makeArray @D (getComp arr1 <> getComp arr2) (size arr1) $ \ix -> + f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix) + ) {-# INLINE compareArrays #-} -- | 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) +liftArray2' + :: (HasCallStack, Index ix, Source r1 a, Source r2 b) => (a -> b -> e) -> Array r1 ix a -> Array r2 ix b @@ -303,8 +308,8 @@ -- have to be the same, otherwise it throws `SizeMismatchException`. -- -- @since 1.0.0 -liftArray2M :: - (Index ix, Source r1 a, Source r2 b, MonadThrow m) +liftArray2M + :: (Index ix, Source r1 a, Source r2 b, MonadThrow m) => (a -> b -> e) -> Array r1 ix a -> Array r2 ix b @@ -318,8 +323,8 @@ sz2 = size arr2 {-# INLINE liftArray2M #-} -zipWithInternal :: - (Index ix, Source r1 e1, Source r2 e2) +zipWithInternal + :: (Index ix, Source r1 e1, Source r2 e2) => Sz ix -> (e1 -> e2 -> e3) -> Array r1 ix e1 @@ -329,7 +334,7 @@ case unsafePrefIndex arr1 of PrefIndexLinear gi1 | PrefIndexLinear gi2 <- unsafePrefIndex arr2 -> - makeArrayLinear comp sz (\ !i -> f (gi1 i) (gi2 i)) + makeArrayLinear comp sz (\ !i -> f (gi1 i) (gi2 i)) _ -> makeArray comp sz (\ !ix -> f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)) where comp = getComp arr1 <> getComp arr2 diff --git a/src/Data/Massiv/Array/Delayed/Stream.hs b/src/Data/Massiv/Array/Delayed/Stream.hs --- a/src/Data/Massiv/Array/Delayed/Stream.hs +++ b/src/Data/Massiv/Array/Delayed/Stream.hs @@ -3,6 +3,7 @@ {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE TypeFamilies #-} + -- | -- Module : Data.Massiv.Array.Delayed.Stream -- Copyright : (c) Alexey Kuleshevich 2019-2022 @@ -10,27 +11,26 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Delayed.Stream - ( DS(..) - , Array (..) - , toStreamArray - , toStreamM - , toStreamIxM - , toSteps - , fromSteps - , fromStepsM - ) where +module Data.Massiv.Array.Delayed.Stream ( + DS (..), + Array (..), + toStreamArray, + toStreamM, + toStreamIxM, + toSteps, + fromSteps, + fromStepsM, +) where import Control.Applicative import Control.Monad.ST import Data.Coerce import Data.Foldable import Data.Massiv.Array.Delayed.Pull -import qualified Data.Massiv.Vector.Stream as S import Data.Massiv.Core.Common +import qualified Data.Massiv.Vector.Stream as S import GHC.Exts -import Prelude hiding (take, drop) +import Prelude hiding (drop, take) -- | 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 @@ -62,7 +62,6 @@ fromStepsM = fmap DSArray . S.transSteps {-# INLINE fromStepsM #-} - instance Shape DS Ix1 where linearSizeHint = stepsSize . dsArray {-# INLINE linearSizeHint #-} @@ -82,7 +81,6 @@ setComp _ = id repr = DS - instance Functor (Array DS Ix1) where fmap f = coerce . S.map f . dsArray {-# INLINE fmap #-} @@ -104,7 +102,6 @@ (>>=) arr f = coerce (S.concatMap (coerce . f) (dsArray arr)) {-# INLINE (>>=) #-} - instance Foldable (Array DS Ix1) where foldr f acc = S.unId . S.foldrLazy f acc . toSteps {-# INLINE foldr #-} @@ -135,7 +132,6 @@ (<>) a1 a2 = fromSteps (coerce a1 `S.append` coerce a2) {-# INLINE (<>) #-} - instance Monoid (Array DS Ix1 e) where mempty = DSArray S.empty {-# INLINE mempty #-} @@ -153,20 +149,19 @@ toList = S.toList . coerce {-# INLINE toList #-} - instance S.Stream DS Ix1 e where toStream = coerce {-# INLINE toStream #-} toStreamIx = S.indexed . coerce {-# INLINE toStreamIx #-} - -- | Flatten an array into a stream of values. -- -- @since 0.4.1 toStreamArray :: (Index ix, Source r e) => Array r ix e -> Vector DS e toStreamArray = DSArray . S.steps -{-# INLINE[1] toStreamArray #-} +{-# INLINE [1] toStreamArray #-} + {-# RULES "toStreamArray/id" toStreamArray = id #-} -- | /O(1)/ - Convert an array into monadic `Steps` @@ -183,10 +178,8 @@ toStreamIxM = S.transStepsId . toStreamIx {-# INLINE toStreamIxM #-} - -- | /O(n)/ - `size` implementation. instance Load DS Ix1 e where - makeArrayLinear _ k = fromSteps . S.generate k {-# INLINE makeArrayLinear #-} replicate _ k = fromSteps . S.replicate k @@ -203,7 +196,6 @@ unsafeLoadIntoIO marr arr = stToIO $ unsafeLoadIntoST marr arr {-# INLINE unsafeLoadIntoIO #-} - -- cons :: e -> Array DS Ix1 e -> Array DS Ix1 e -- cons e = coerce . S.cons e . dsArray -- {-# INLINE cons #-} @@ -216,7 +208,6 @@ -- snoc (DSArray sts) e = DSArray (S.snoc sts e) -- {-# INLINE snoc #-} - -- TODO: skip the stride while loading -- instance StrideLoad DS Ix1 e where -- iterArrayLinearWithStrideST_ scheduler stride resultSize arr uWrite = @@ -227,5 +218,3 @@ -- iterLinearM_ resultSize start (totalElem resultSize) (numWorkers scheduler) (<) $ -- \ !i ix -> uWrite i (f (liftIndex2 (*) strideIx ix)) -- {-# INLINE iterArrayLinearWithStrideST_ #-} - - diff --git a/src/Data/Massiv/Array/Delayed/Windowed.hs b/src/Data/Massiv/Array/Delayed/Windowed.hs --- a/src/Data/Massiv/Array/Delayed/Windowed.hs +++ b/src/Data/Massiv/Array/Delayed/Windowed.hs @@ -8,6 +8,7 @@ {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE TypeOperators #-} {-# LANGUAGE UndecidableInstances #-} + -- | -- Module : Data.Massiv.Array.Delayed.Windowed -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -15,16 +16,15 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Delayed.Windowed - ( DW(..) - , Array(..) - , Window(..) - , insertWindow - , getWindow - , dropWindow - , makeWindowedArray - ) where +module Data.Massiv.Array.Delayed.Windowed ( + DW (..), + Array (..), + Window (..), + insertWindow, + getWindow, + dropWindow, + makeWindowedArray, +) where import Control.Monad (when) import Data.Massiv.Array.Delayed.Pull @@ -39,46 +39,46 @@ -- | Delayed Windowed Array representation. data DW = DW -data Window ix e = Window { windowStart :: !ix - -- ^ Index of where window will start at. - , windowSize :: !(Sz ix) - -- ^ Size of the window - , windowIndex :: ix -> e - -- ^ Indexing function for the window - , windowUnrollIx2 :: !(Maybe Int) - -- ^ Setting this value during stencil application improves cache - -- utilization by unrolling the loop for Ix2 and higher dimensions. - -- Has no affect on arrays with one dimension. - } +data Window ix e = Window + { windowStart :: !ix + -- ^ Index of where window will start at. + , windowSize :: !(Sz ix) + -- ^ Size of the window + , windowIndex :: ix -> e + -- ^ Indexing function for the window + , windowUnrollIx2 :: !(Maybe Int) + -- ^ Setting this value during stencil application improves cache + -- utilization by unrolling the loop for Ix2 and higher dimensions. + -- Has no affect on arrays with one dimension. + } instance Functor (Window ix) where - fmap f arr@Window{windowIndex} = arr { windowIndex = f . windowIndex } + fmap f arr@Window{windowIndex} = arr{windowIndex = f . windowIndex} -data instance Array DW ix e = DWArray { dwArray :: !(Array D ix e) - , dwWindow :: !(Maybe (Window ix e)) - } +data instance Array DW ix e = DWArray + { dwArray :: !(Array D ix e) + , dwWindow :: !(Maybe (Window ix e)) + } instance (Ragged L ix e, Load DW ix e, Show e) => Show (Array DW ix e) where showsPrec = showsArrayPrec (computeAs B) showList = showArrayList instance Strategy DW where - setComp c arr = arr { dwArray = (dwArray arr) { dComp = c } } + setComp c arr = arr{dwArray = (dwArray arr){dComp = c}} {-# INLINE setComp #-} getComp = dComp . dwArray {-# INLINE getComp #-} repr = DW - instance Functor (Array DW ix) where fmap f arr@DWArray{dwArray, dwWindow} = arr - { dwArray = fmap f dwArray - , dwWindow = fmap f <$> dwWindow - } + { dwArray = fmap f dwArray + , dwWindow = fmap f <$> dwWindow + } {-# INLINE fmap #-} - -- -- -- @since 0.3.0 @@ -120,14 +120,18 @@ -- @since 0.1.3 makeWindowedArray :: (Index ix, Source r e) - => Array r ix e -- ^ Source array that will have a window inserted into it - -> ix -- ^ Start index for the window - -> Sz ix -- ^ Size of the window - -> (ix -> e) -- ^ Indexing function foto use inside window + => Array r ix e + -- ^ Source array that will have a window inserted into it + -> ix + -- ^ Start index for the window + -> Sz ix + -- ^ Size of the window + -> (ix -> e) + -- ^ Indexing function foto use inside window -> Array DW ix e makeWindowedArray !arr wStart wSize wIndex = insertWindow (delay arr) $ - Window {windowStart = wStart, windowSize = wSize, windowIndex = wIndex, windowUnrollIx2 = Nothing} + Window{windowStart = wStart, windowSize = wSize, windowIndex = wIndex, windowUnrollIx2 = Nothing} {-# INLINE makeWindowedArray #-} -- | Inserts a `Window` into a delayed array while scaling the window down if it doesn't fit inside @@ -136,32 +140,34 @@ -- @since 0.3.0 insertWindow :: 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 D ix e + -- ^ Source array that will have a window inserted into it + -> Window ix e + -- ^ Window to place inside the delayed array -> Array DW ix e insertWindow !arr !window = DWArray { dwArray = delay arr , dwWindow = Just $! - Window - { windowStart = wStart' - , windowSize = Sz (liftIndex2 min wSize (liftIndex2 (-) sz wStart')) - , windowIndex = wIndex - , windowUnrollIx2 = wUnrollIx2 - } + Window + { windowStart = wStart' + , windowSize = Sz (liftIndex2 min wSize (liftIndex2 (-) sz wStart')) + , windowIndex = wIndex + , windowUnrollIx2 = wUnrollIx2 + } } where wStart' = unSz (Sz (liftIndex2 min wStart (liftIndex (subtract 1) sz))) Sz sz = size arr - Window { windowStart = wStart - , windowSize = Sz wSize - , windowIndex = wIndex - , windowUnrollIx2 = wUnrollIx2 - } = window + Window + { windowStart = wStart + , windowSize = Sz wSize + , windowIndex = wIndex + , windowUnrollIx2 = wUnrollIx2 + } = window {-# INLINE insertWindow #-} - -- | Get the `Window` from a windowed array. -- -- @since 0.2.1 @@ -176,7 +182,6 @@ dropWindow = dwArray {-# INLINE dropWindow #-} - zeroWindow :: Index ix => Window ix e zeroWindow = Window zeroIndex zeroSz windowError Nothing {-# INLINE zeroWindow #-} @@ -184,16 +189,14 @@ data EmptyWindowException = EmptyWindowException deriving (Eq, Show) instance Exception EmptyWindowException where - displayException _ = "Index of zero size Window" windowError :: a windowError = throwImpossible EmptyWindowException {-# NOINLINE windowError #-} - -loadWithIx1 :: - (Monad m) +loadWithIx1 + :: (Monad m) => (m () -> m ()) -> Array DW Ix1 e -> (Ix1 -> e -> m a) @@ -230,18 +233,18 @@ instance StrideLoad DW Ix1 e where iterArrayLinearWithStrideST_ scheduler stride sz arr uWrite = do - (loadWindow, (wStart, wEnd)) <- loadArrayWithIx1 (scheduleWork scheduler) arr stride sz uWrite - let (chunkWidth, slackWidth) = (wEnd - wStart) `quotRem` numWorkers scheduler - loopA_ 0 (< numWorkers scheduler) (+ 1) $ \ !wid -> - let !it' = wid * chunkWidth + wStart - in loadWindow (it', it' + chunkWidth) - when (slackWidth > 0) $ - let !itSlack = numWorkers scheduler * chunkWidth + wStart - in loadWindow (itSlack, itSlack + slackWidth) + (loadWindow, (wStart, wEnd)) <- loadArrayWithIx1 (scheduleWork scheduler) arr stride sz uWrite + let (chunkWidth, slackWidth) = (wEnd - wStart) `quotRem` numWorkers scheduler + loopA_ 0 (< numWorkers scheduler) (+ 1) $ \ !wid -> + let !it' = wid * chunkWidth + wStart + in loadWindow (it', it' + chunkWidth) + when (slackWidth > 0) $ + let !itSlack = numWorkers scheduler * chunkWidth + wStart + in loadWindow (itSlack, itSlack + slackWidth) {-# INLINE iterArrayLinearWithStrideST_ #-} -loadArrayWithIx1 :: - (Monad m) +loadArrayWithIx1 + :: (Monad m) => (m () -> m ()) -> Array DW Ix1 e -> Stride Ix1 @@ -259,15 +262,14 @@ return ( \(from, to) -> with $ - iterA_ (strideStart stride from) to strideIx (<) $ \ !i -> - uWrite (i `div` strideIx) (indexW i) - , (it, wEnd)) + iterA_ (strideStart stride from) to strideIx (<) $ \ !i -> + uWrite (i `div` strideIx) (indexW i) + , (it, wEnd) + ) {-# INLINE loadArrayWithIx1 #-} - - -loadWithIx2 :: - Monad m +loadWithIx2 + :: Monad m => (m () -> m ()) -> Array DW Ix2 t1 -> (Int -> t1 -> m ()) @@ -293,8 +295,8 @@ return (f, it :. ib) {-# INLINE loadWithIx2 #-} -loadArrayWithIx2 :: - Monad m +loadArrayWithIx2 + :: Monad m => (m () -> m ()) -> Array DW Ix2 e -> Stride Ix2 @@ -317,15 +319,15 @@ with $ iterA_ (strideStart stride (it :. 0)) (ib :. jt) strideIx (<) writeB with $ iterA_ (strideStart stride (it :. jb)) (ib :. n) strideIx (<) writeB let f (it' :. ib') - | is > 1 || blockHeight <= 1 -- Turn off unrolling for vertical strides - = iterA_ (strideStart stride (it' :. jt)) (ib' :. jb) strideIx (<) writeW + | is > 1 || blockHeight <= 1 = + -- Turn off unrolling for vertical strides + iterA_ (strideStart stride (it' :. jt)) (ib' :. jb) strideIx (<) writeW | otherwise = - unrollAndJam blockHeight (strideStart stride (it' :. jt)) (ib' :. jb) js writeW + unrollAndJam blockHeight (strideStart stride (it' :. jt)) (ib' :. jb) js writeW {-# INLINE f #-} return (with . f, it :. ib) {-# INLINE loadArrayWithIx2 #-} - loadWindowIx2 :: Monad m => Int -> (Ix2 -> m ()) -> Ix2 -> m () loadWindowIx2 nWorkers loadWindow (it :. ib) = do let !(chunkHeight, slackHeight) = (ib - it) `quotRem` nWorkers @@ -337,22 +339,20 @@ in loadWindow (itSlack :. (itSlack + slackHeight)) {-# INLINE loadWindowIx2 #-} - instance Load DW Ix2 e where 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)) + loadWithIx2 (scheduleWork scheduler) arr uWrite + >>= uncurry (loadWindowIx2 (numWorkers scheduler)) {-# INLINE iterArrayLinearST_ #-} instance StrideLoad DW Ix2 e where iterArrayLinearWithStrideST_ scheduler stride sz arr uWrite = - loadArrayWithIx2 (scheduleWork scheduler) arr stride sz uWrite >>= - uncurry (loadWindowIx2 (numWorkers scheduler)) + loadArrayWithIx2 (scheduleWork scheduler) arr stride sz uWrite + >>= uncurry (loadWindowIx2 (numWorkers scheduler)) {-# INLINE iterArrayLinearWithStrideST_ #-} - instance (Index (IxN n), Load DW (Ix (n - 1)) e) => Load DW (IxN n) e where makeArray c sz f = DWArray (makeArray c sz f) Nothing {-# INLINE makeArray #-} @@ -363,8 +363,8 @@ iterArrayLinearWithStrideST_ = loadArrayWithIxN {-# INLINE iterArrayLinearWithStrideST_ #-} -loadArrayWithIxN :: - (Index ix, StrideLoad DW (Lower ix) e) +loadArrayWithIxN + :: (Index ix, StrideLoad DW (Lower ix) e) => Scheduler s () -> Stride ix -> Sz ix @@ -373,7 +373,7 @@ -> ST s () loadArrayWithIxN scheduler stride szResult arr uWrite = do let DWArray darr window = arr - Window {windowStart, windowSize, windowIndex, windowUnrollIx2} = fromMaybe zeroWindow window + Window{windowStart, windowSize, windowIndex, windowUnrollIx2} = fromMaybe zeroWindow window !(headSourceSize, lowerSourceSize) = unconsSz (dSize darr) !lowerSize = snd $ unconsSz szResult !(s, lowerStrideIx) = unconsDim $ unStride stride @@ -390,10 +390,12 @@ } mkLowerArray mw i = DWArray - { dwArray = darr { dComp = Seq - , dSize = lowerSourceSize - , dPrefIndex = PrefIndex (unsafeIndex darr . consDim i) - } + { dwArray = + darr + { dComp = Seq + , dSize = lowerSourceSize + , dPrefIndex = PrefIndex (unsafeIndex darr . consDim i) + } , dwWindow = ($ i) <$> mw } loadLower mw !i = @@ -413,17 +415,15 @@ loopA_ (strideStart (Stride s) curWindowEnd) (< unSz headSourceSize) (+ s) (loadLower Nothing) {-# INLINE loadArrayWithIxN #-} - - -loadWithIxN :: - (Index ix, Load DW (Lower ix) e) +loadWithIxN + :: (Index ix, Load DW (Lower ix) e) => Scheduler s () -> Array DW ix e -> (Int -> e -> ST s ()) -> ST s () loadWithIxN scheduler arr uWrite = do let DWArray darr window = arr - Window {windowStart, windowSize, windowIndex, windowUnrollIx2} = fromMaybe zeroWindow window + Window{windowStart, windowSize, windowIndex, windowUnrollIx2} = fromMaybe zeroWindow window !(si, szL) = unconsSz (dSize darr) !windowEnd = liftIndex2 (+) windowStart (unSz windowSize) !(t, windowStartL) = unconsDim windowStart @@ -438,42 +438,46 @@ mkLowerArray mw i = DWArray { dwArray = - darr {dComp = Seq, dSize = szL, dPrefIndex = PrefIndex (unsafeIndex darr . consDim i)} + darr{dComp = Seq, dSize = szL, dPrefIndex = PrefIndex (unsafeIndex darr . consDim i)} , dwWindow = ($ i) <$> mw } loadLower mw !i = scheduleWork_ scheduler $ - iterArrayLinearST_ scheduler (mkLowerArray mw i) (\k -> uWrite (k + pageElements * i)) + iterArrayLinearST_ scheduler (mkLowerArray mw i) (\k -> uWrite (k + pageElements * i)) {-# NOINLINE loadLower #-} loopA_ 0 (< headDim windowStart) (+ 1) (loadLower Nothing) loopA_ t (< headDim windowEnd) (+ 1) (loadLower (Just mkLowerWindow)) loopA_ (headDim windowEnd) (< unSz si) (+ 1) (loadLower Nothing) {-# INLINE loadWithIxN #-} - - -unrollAndJam :: Monad m => - Int -- ^ Block height - -> Ix2 -- ^ Top corner - -> Ix2 -- ^ Bottom corner - -> Int -- ^ Column Stride - -> (Ix2 -> m ()) -- ^ Writing function - -> m () +unrollAndJam + :: Monad m + => Int + -- ^ Block height + -> Ix2 + -- ^ Top corner + -> Ix2 + -- ^ Bottom corner + -> Int + -- ^ Column Stride + -> (Ix2 -> m ()) + -- ^ Writing function + -> m () unrollAndJam !bH (it :. jt) (ib :. jb) js f = do - let f2 (i :. j) = f (i :. j) >> f ((i + 1) :. j) + let f2 (i :. j) = f (i :. j) >> f ((i + 1) :. j) let f3 (i :. j) = f (i :. j) >> f2 ((i + 1) :. j) let f4 (i :. j) = f (i :. j) >> f3 ((i + 1) :. j) let f5 (i :. j) = f (i :. j) >> f4 ((i + 1) :. j) let f6 (i :. j) = f (i :. j) >> f5 ((i + 1) :. j) let f7 (i :. j) = f (i :. j) >> f6 ((i + 1) :. j) let f' = case bH of - 1 -> f - 2 -> f2 - 3 -> f3 - 4 -> f4 - 5 -> f5 - 6 -> f6 - _ -> f7 + 1 -> f + 2 -> f2 + 3 -> f3 + 4 -> f4 + 5 -> f5 + 6 -> f6 + _ -> f7 let !ibS = ib - ((ib - it) `mod` bH) loopA_ it (< ibS) (+ bH) $ \ !i -> loopA_ jt (< jb) (+ js) $ \ !j -> @@ -482,7 +486,5 @@ loopA_ jt (< jb) (+ js) $ \ !j -> f (i :. j) {-# INLINE unrollAndJam #-} - - -- TODO: Implement Hilbert curve 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 @@ -171,7 +171,7 @@ toByteString = castToByteString . #if __GLASGOW_HASKELL__ >= 820 convert - {- For ghc-8.0 `covert` results in "internal error: ARR_WORDS object entered!" -} + {- For ghc-8.0 `convert` results in "internal error: ARR_WORDS object entered!" -} #else compute #endif 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 @@ -158,8 +158,7 @@ instance Source BL e where unsafeLinearIndex (BLArray _ _sz o a) i = - INDEX_CHECK("(Source BL ix e).unsafeLinearIndex", - SafeSz . A.sizeofArray, A.indexArray) a (i + o) + indexAssert "BL.unsafeLinearIndex" (SafeSz . A.sizeofArray) A.indexArray a (i + o) {-# INLINE unsafeLinearIndex #-} unsafeOuterSlice (BLArray c _ o a) szL i = BLArray c szL (i * totalElem szL + o) a @@ -171,8 +170,7 @@ instance Manifest BL e where unsafeLinearIndexM (BLArray _ _sz o a) i = - INDEX_CHECK("(Manifest BL ix e).unsafeLinearIndexM", - SafeSz . A.sizeofArray, A.indexArray) a (i + o) + indexAssert "BL.unsafeLinearIndexM" (SafeSz . A.sizeofArray) A.indexArray a (i + o) {-# INLINE unsafeLinearIndexM #-} sizeOfMArray (MBLArray sz _ _) = sz @@ -200,13 +198,11 @@ {-# INLINE newMArray #-} unsafeLinearRead (MBLArray _ o ma) i = - INDEX_CHECK("(Manifest BL ix e).unsafeLinearRead", - SafeSz . A.sizeofMutableArray, A.readArray) ma (i + o) + indexAssert "B.unsafeLinearRead" (SafeSz . A.sizeofMutableArray) A.readArray ma (i + o) {-# INLINE unsafeLinearRead #-} - unsafeLinearWrite (MBLArray _sz o ma) i e = - INDEX_CHECK("(Manifest BL ix e).unsafeLinearWrite", - SafeSz . A.sizeofMutableArray, A.writeArray) ma (i + o) e + unsafeLinearWrite (MBLArray _sz o ma) i = + indexAssert "B.unsafeLinearWrite" (SafeSz . A.sizeofMutableArray) A.writeArray ma (i + o) {-# INLINE unsafeLinearWrite #-} instance Size BL where 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 @@ -8,6 +8,7 @@ {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE TypeOperators #-} {-# LANGUAGE UndecidableInstances #-} + -- | -- Module : Data.Massiv.Array.Manifest.Internal -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -15,38 +16,37 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Manifest.Internal - ( Manifest(..) - , Array(..) - , flattenMArray - , compute - , computeS - , computeP - , computeIO - , computePrimM - , computeAs - , computeProxy - , computeSource - , computeWithStride - , computeWithStrideAs - , clone - , convert - , convertAs - , convertProxy - , gcastArr - , fromRaggedArrayM - , fromRaggedArray' - , unsafeLoadIntoS - , unsafeLoadIntoM - , iterateUntil - , iterateUntilM - ) where +module Data.Massiv.Array.Manifest.Internal ( + Manifest (..), + Array (..), + flattenMArray, + compute, + computeS, + computeP, + computeIO, + computePrimM, + computeAs, + computeProxy, + computeSource, + computeWithStride, + computeWithStrideAs, + clone, + convert, + convertAs, + convertProxy, + gcastArr, + fromRaggedArrayM, + fromRaggedArray', + unsafeLoadIntoS, + unsafeLoadIntoM, + iterateUntil, + iterateUntilM, +) where -import Control.Exception (try) import Control.DeepSeq -import Control.Monad.ST +import Control.Exception (try) import Control.Monad.Primitive +import Control.Monad.ST import Control.Scheduler import Data.Massiv.Array.Delayed.Pull import Data.Massiv.Array.Mutable @@ -57,13 +57,12 @@ import Data.Typeable import System.IO.Unsafe (unsafePerformIO) - -- | 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' . (Manifest r 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 #-} @@ -71,18 +70,18 @@ -- the same as `computePrimM`, but executed in `ST`, thus pure. -- -- @since 0.1.0 -computeS :: forall r ix e r' . (Manifest r e, Load r' ix e) => Array r' ix e -> Array r ix e +computeS :: forall r ix e r'. (Manifest r e, Load r' ix e) => Array r' ix e -> Array r ix e computeS !arr = runST $ computePrimM arr {-# INLINE computeS #-} - -- | Compute array in parallel using all cores disregarding predefined computation -- strategy. Computation stategy of the resulting array will match the source, despite -- that it is diregarded. -- -- @since 0.5.4 -computeP :: - forall r ix e r'. (Manifest r e, Load r' ix e) +computeP + :: forall r ix e r' + . (Manifest r e, Load r' ix e) => Array r' ix e -> Array r ix e computeP arr = setComp (getComp arr) $ compute (setComp Par arr) @@ -94,8 +93,9 @@ -- computing an array during benchmarking. -- -- @since 0.4.5 -computeIO :: - forall r ix e r' m. (Manifest r e, Load r' ix e, MonadIO m) +computeIO + :: forall r ix e r' m + . (Manifest r e, Load r' ix e, MonadIO m) => Array r' ix e -> m (Array r ix e) computeIO arr = liftIO (loadArray arr >>= unsafeFreeze (getComp arr)) @@ -105,14 +105,14 @@ -- strategy. -- -- @since 0.4.5 -computePrimM :: - forall r ix e r' m. (Manifest r e, Load r' ix e, PrimMonad m) +computePrimM + :: forall r ix e r' m + . (Manifest r e, Load r' ix e, PrimMonad m) => Array r' ix e -> m (Array r ix e) computePrimM arr = loadArrayS arr >>= unsafeFreeze (getComp arr) {-# INLINE computePrimM #-} - -- | Just as `compute`, but let's you supply resulting representation type as an argument. -- -- ====__Examples__ @@ -121,12 +121,10 @@ -- >>> computeAs P $ range Seq (Ix1 0) 10 -- Array P Seq (Sz1 10) -- [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9 ] --- computeAs :: (Manifest r e, Load r' ix e) => r -> Array r' ix e -> Array r ix e computeAs _ = compute {-# INLINE computeAs #-} - -- | Same as `compute` and `computeAs`, but let's you supply resulting representation type as a proxy -- argument. -- @@ -146,17 +144,18 @@ computeProxy _ = compute {-# INLINE computeProxy #-} - -- | This is just like `convert`, but restricted to `Source` arrays. Will be a noop if -- resulting type is the same as the input. -- -- @since 0.1.0 -computeSource :: forall r ix e r' . (Manifest r e, Source r' e, Index ix) - => Array r' ix e -> Array 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 $ delay arr) (\Refl -> arr) (eqT :: Maybe (r' :~: r)) {-# INLINE computeSource #-} - -- | /O(n)/ - Make an exact immutable copy of an Array. -- -- @since 0.1.0 @@ -164,80 +163,91 @@ clone arr = unsafePerformIO $ thaw arr >>= unsafeFreeze (getComp arr) {-# INLINE clone #-} - -- | /O(1)/ - Cast over Array representation -gcastArr :: forall r ix e r' . (Typeable r, Typeable r') - => Array r' ix e -> Maybe (Array r ix e) +gcastArr + :: forall r ix e r' + . (Typeable r, Typeable r') + => Array r' ix e + -> Maybe (Array r ix e) gcastArr arr = fmap (\Refl -> arr) (eqT :: Maybe (r :~: r')) - -- | /O(n)/ - conversion between array types. A full copy will occur, unless when the source and -- result arrays are of the same representation, in which case it is an /O(1)/ operation. -- -- @since 0.1.0 -convert :: forall r ix e r' . (Manifest r e, Load r' ix e) - => Array r' ix e -> Array r ix e +convert + :: forall r ix e r' + . (Manifest r e, Load r' ix e) + => Array r' ix e + -> Array r ix e convert arr = fromMaybe (compute arr) (gcastArr arr) {-# INLINE convert #-} -- | Same as `convert`, but let's you supply resulting representation type as an argument. -- -- @since 0.1.0 -convertAs :: (Manifest r e, Load r' ix e) - => r -> Array r' ix e -> Array r ix e +convertAs + :: (Manifest r e, Load r' ix e) + => r + -> Array r' ix e + -> Array r ix e convertAs _ = convert {-# INLINE convertAs #-} - -- | Same as `convert` and `convertAs`, but let's you supply resulting representation type as a -- proxy argument. -- -- @since 0.1.1 -convertProxy :: (Manifest r e, Load r' ix e) - => proxy r -> Array r' ix e -> Array r ix e +convertProxy + :: (Manifest r e, Load r' ix e) + => proxy r + -> Array r' ix e + -> Array r ix e convertProxy _ = convert {-# INLINE convertProxy #-} - -- | Convert a ragged array into a common array with rectangular shape. Throws `ShapeException` -- whenever supplied ragged array does not have a rectangular shape. -- -- @since 0.4.0 -fromRaggedArrayM :: - forall r ix e r' m . (Manifest r e, Ragged r' ix e, MonadThrow m) +fromRaggedArrayM + :: forall r ix e r' m + . (Manifest r e, Ragged r' ix e, MonadThrow m) => Array r' ix e -> m (Array r ix e) fromRaggedArrayM arr = let sz = outerSize arr in either (\(e :: ShapeException) -> throwM e) pure $ - unsafePerformIO $ do - marr <- unsafeNew sz - traverse (\_ -> unsafeFreeze (getComp arr) marr) =<< - try (withMassivScheduler_ (getComp arr) $ \scheduler -> - stToIO $ loadRaggedST scheduler arr (unsafeLinearWrite marr) 0 (totalElem sz) sz) + unsafePerformIO $ do + marr <- unsafeNew sz + traverse (\_ -> unsafeFreeze (getComp arr) marr) + =<< try + ( withMassivScheduler_ (getComp arr) $ \scheduler -> + stToIO $ loadRaggedST scheduler arr (unsafeLinearWrite marr) 0 (totalElem sz) sz + ) {-# INLINE fromRaggedArrayM #-} - -- | Same as `fromRaggedArrayM`, but will throw an impure exception if its shape is not -- rectangular. -- -- @since 0.1.1 -fromRaggedArray' :: - forall r ix e r'. (HasCallStack, Manifest r e, Ragged r' ix e) +fromRaggedArray' + :: forall r ix e r' + . (HasCallStack, Manifest r e, Ragged r' ix e) => Array r' ix e -> Array r ix e fromRaggedArray' = throwEither . fromRaggedArrayM {-# INLINE fromRaggedArray' #-} - -- | Same as `compute`, but with `Stride`. -- -- /O(n div k)/ - Where @n@ is number of elements in the source array and @k@ is number of -- elements in the stride. -- -- @since 0.3.0 -computeWithStride :: - forall r ix e r'. (Manifest r e, StrideLoad r' ix e) +computeWithStride + :: forall r ix e r' + . (Manifest r e, StrideLoad r' ix e) => Stride ix -> Array r' ix e -> Array r ix e @@ -248,23 +258,22 @@ stToIO $ iterArrayLinearWithStrideST_ scheduler stride sz arr (unsafeLinearWrite marr) {-# INLINE computeWithStride #-} - -- | Same as `computeWithStride`, but with ability to specify resulting array representation. -- -- @since 0.3.0 -computeWithStrideAs :: - (Manifest r e, StrideLoad r' ix e) => r -> Stride ix -> Array r' ix e -> Array r ix e +computeWithStrideAs + :: (Manifest r e, StrideLoad r' ix e) => r -> Stride ix -> Array r' ix e -> Array r ix e computeWithStrideAs _ = computeWithStride {-# INLINE computeWithStrideAs #-} - -- | Load into a supplied mutable vector sequentially. Returned array is not -- necesserally the same vector as the one that was supplied. It will be the -- same only if it had enough space to load all the elements in. -- -- @since 0.5.7 -unsafeLoadIntoS :: - forall r r' ix e m s. (Load r ix e, Manifest r' e, MonadPrim s m) +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) @@ -274,15 +283,15 @@ -- | 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) +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 @@ -319,16 +328,17 @@ -- ] -- -- @since 0.3.6 -iterateUntil :: - (Load r' ix e, Manifest r e, NFData (Array r ix e)) +iterateUntil + :: (Load r' ix e, Manifest r e, NFData (Array r ix e)) => (Int -> Array r ix e -> Array r ix e -> Bool) -- ^ Convergence condition. Accepts current iteration counter, array at the previous -- state and at the current state. -> (Int -> Array r ix e -> Array r' ix e) -- ^ A modifying function to apply at each iteration. The size of resulting array may -- differ if necessary - -> Array r ix e -- ^ Initial source array -> Array r ix e + -- ^ Initial source array + -> Array r ix e iterateUntil convergence iteration initArr0 = unsafePerformIO $ do let loadArr0 = iteration 0 initArr0 initMVec1 <- unsafeNew (fromMaybe zeroSz (maxLinearSize loadArr0)) @@ -344,8 +354,8 @@ -- iteration. -- -- @since 0.3.6 -iterateUntilM :: - (Load r' ix e, Manifest r e, MonadIO m) +iterateUntilM + :: (Load r' ix e, Manifest r e, MonadIO m) => (Int -> Array r ix e -> MArray RealWorld r ix e -> m Bool) -- ^ Convergence condition. Accepts current iteration counter, pure array at previous -- state and a mutable at the current state, therefore after each iteration its contents @@ -353,7 +363,8 @@ -> (Int -> Array r ix e -> m (Array r' ix e)) -- ^ A modifying function to apply at each iteration. The size of resulting array may -- differ if necessary. - -> Array r ix e -- ^ Initial source array + -> Array r ix e + -- ^ Initial source array -> m (Array r ix e) iterateUntilM convergence iteration initArr0 = do loadArr0 <- iteration 0 initArr0 @@ -362,9 +373,8 @@ iterateLoop conv iteration 0 initArr0 loadArr0 initMVec1 {-# INLINE iterateUntilM #-} - -iterateLoop :: - (Load r' ix e, Manifest r e, MonadIO m) +iterateLoop + :: (Load r' ix e, Manifest r e, MonadIO m) => (Int -> Array r ix e -> Comp -> MArray RealWorld r ix e -> m (Bool, Array r ix e)) -> (Int -> Array r ix e -> m (Array r' ix e)) -> Int diff --git a/src/Data/Massiv/Array/Manifest/List.hs b/src/Data/Massiv/Array/Manifest/List.hs --- a/src/Data/Massiv/Array/Manifest/List.hs +++ b/src/Data/Massiv/Array/Manifest/List.hs @@ -4,6 +4,7 @@ {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} + -- | -- Module : Data.Massiv.Array.Manifest.List -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -11,38 +12,39 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Manifest.List - ( -- ** List - fromList - , fromListsM - , fromLists' - , toList - , toLists - , toLists2 - , toLists3 - , toLists4 - ) where +module Data.Massiv.Array.Manifest.List ( + -- ** List + fromList, + fromListsM, + fromLists', + toList, + toLists, + toLists2, + toLists3, + toLists4, +) where import Data.Massiv.Array.Manifest.Internal import Data.Massiv.Array.Ops.Fold (foldrInner) import Data.Massiv.Array.Ops.Fold.Internal (foldrFB) import Data.Massiv.Core.Common import Data.Massiv.Core.List -import qualified GHC.Exts as GHC (build, IsList(..)) +import qualified GHC.Exts as GHC (IsList (..), build) -- | Convert a flat list into a vector -- -- @since 0.1.0 -fromList :: - forall r e. Manifest r e - => Comp -- ^ Computation startegy to use - -> [e] -- ^ Flat list +fromList + :: forall r e + . Manifest r e + => Comp + -- ^ Computation startegy to use + -> [e] + -- ^ Flat list -> Vector r e fromList = fromLists' {-# INLINE fromList #-} - -- | /O(n)/ - Convert a nested list into an array. Nested list must be of a rectangular shape, -- otherwise a runtime error will occur. Also, nestedness must match the rank of resulting array, -- which should be specified through an explicit type signature. @@ -82,26 +84,26 @@ -- *** Exception: DimTooShortException for (Dim 1): expected (Sz1 4), got (Sz1 3) -- -- @since 0.3.0 -fromListsM :: - forall r ix e m. (Ragged L ix e, Manifest r e, MonadThrow m) +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 #-} - -fromListToListArray :: - forall ix e. GHC.IsList (Array L ix e) +fromListToListArray + :: forall ix e + . GHC.IsList (Array L ix e) => [ListItem ix e] -> Array L ix e fromListToListArray = GHC.fromList {-# INLINE fromListToListArray #-} - -- | Same as `fromListsM`, but will throw an error on irregular shaped lists. -- --- __Note__: This function is the same as if you would turn on @{-\# LANGUAGE OverloadedLists #-}@ +-- __Note__: This function is the same as if you would turn on @{\-\# LANGUAGE OverloadedLists #-\}@ -- extension. For that reason you can also use `GHC.Exts.fromList`. -- -- prop> \xs -> fromLists' Seq xs == (fromList Seq xs :: Vector P Int) @@ -127,16 +129,17 @@ -- ] -- -- @since 0.1.0 -fromLists' :: - forall r ix e. (HasCallStack, Ragged L ix e, Manifest r e) - => Comp -- ^ Computation startegy to use - -> [ListItem ix e] -- ^ Nested list +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' #-} - - -- | Convert any array to a flat list. -- -- ==== __Examples__ @@ -147,10 +150,9 @@ -- -- @since 0.1.0 toList :: (Index ix, Source r e) => Array r ix e -> [e] -toList !arr = GHC.build (\ c n -> foldrFB c n arr) +toList !arr = GHC.build (\c n -> foldrFB c n arr) {-# INLINE toList #-} - -- | /O(n)/ - Convert an array into a nested list. Number of array dimensions and list nestedness -- will always match, but you can use `toList`, `toLists2`, etc. if flattening of inner dimensions -- is desired. @@ -172,15 +174,14 @@ -- [[[0 :> 0 :. 0,0 :> 0 :. 1,0 :> 0 :. 2]],[[1 :> 0 :. 0,1 :> 0 :. 1,1 :> 0 :. 2]]] -- -- @since 0.1.0 -toLists :: - (Ragged L ix e, Shape r ix, Source r e) - => Array r ix e -- ^ Array to be converted to nested lists +toLists + :: (Ragged L ix e, Shape r ix, Source r e) + => Array r ix e + -- ^ Array to be converted to nested lists -> [ListItem ix e] toLists = GHC.toList . toListArray {-# INLINE toLists #-} - - -- | Convert an array with at least 2 dimensions into a list of lists. Inner dimensions will get -- flattened. -- @@ -197,13 +198,12 @@ toLists2 = toList . foldrInner (:) [] {-# INLINE toLists2 #-} - -- | Convert an array with at least 3 dimensions into a 3 deep nested list. Inner dimensions will -- get flattened. -- -- @since 0.1.0 -toLists3 :: - (Source r e, Index ix, Index (Lower ix), Index (Lower (Lower ix))) => 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 #-} @@ -211,8 +211,8 @@ -- get flattened. -- -- @since 0.1.0 -toLists4 :: - ( Source r e +toLists4 + :: ( Source r e , Index ix , Index (Lower ix) , Index (Lower (Lower ix)) @@ -222,7 +222,6 @@ -> [[[[e]]]] toLists4 = toList . foldrInner (:) [] . foldrInner (:) [] . foldrInner (:) [] {-# INLINE toLists4 #-} - -- $setup -- >>> import Data.Massiv.Array as A diff --git a/src/Data/Massiv/Array/Manifest/Primitive.hs b/src/Data/Massiv/Array/Manifest/Primitive.hs --- a/src/Data/Massiv/Array/Manifest/Primitive.hs +++ b/src/Data/Massiv/Array/Manifest/Primitive.hs @@ -1,5 +1,4 @@ {-# LANGUAGE BangPatterns #-} -{-# LANGUAGE CPP #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MagicHash #-} @@ -9,6 +8,7 @@ {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UnboxedTuples #-} {-# LANGUAGE UndecidableInstances #-} + -- | -- Module : Data.Massiv.Array.Manifest.Primitive -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -16,76 +16,74 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Manifest.Primitive - ( P(..) - , Array(..) - , MArray(..) - , Prim - , toPrimitiveVector - , toPrimitiveMVector - , fromPrimitiveVector - , fromPrimitiveMVector - , toByteArray - , toByteArrayM - , unwrapByteArray - , unwrapByteArrayOffset - , unwrapMutableByteArray - , unwrapMutableByteArrayOffset - , fromByteArray - , fromByteArrayM - , fromByteArrayOffsetM - , toMutableByteArray - , toMutableByteArrayM - , fromMutableByteArrayM - , fromMutableByteArrayOffsetM - , fromMutableByteArray - , shrinkMutableByteArray - , unsafeAtomicReadIntArray - , unsafeAtomicWriteIntArray - , unsafeCasIntArray - , unsafeAtomicModifyIntArray - , unsafeAtomicAddIntArray - , unsafeAtomicSubIntArray - , unsafeAtomicAndIntArray - , unsafeAtomicNandIntArray - , unsafeAtomicOrIntArray - , unsafeAtomicXorIntArray - ) where +module Data.Massiv.Array.Manifest.Primitive ( + P (..), + Array (..), + MArray (..), + Prim, + toPrimitiveVector, + toPrimitiveMVector, + fromPrimitiveVector, + fromPrimitiveMVector, + toByteArray, + toByteArrayM, + unwrapByteArray, + unwrapByteArrayOffset, + unwrapMutableByteArray, + unwrapMutableByteArrayOffset, + fromByteArray, + fromByteArrayM, + fromByteArrayOffsetM, + toMutableByteArray, + toMutableByteArrayM, + fromMutableByteArrayM, + fromMutableByteArrayOffsetM, + fromMutableByteArray, + shrinkMutableByteArray, + unsafeAtomicReadIntArray, + unsafeAtomicWriteIntArray, + unsafeCasIntArray, + unsafeAtomicModifyIntArray, + unsafeAtomicAddIntArray, + unsafeAtomicSubIntArray, + unsafeAtomicAndIntArray, + unsafeAtomicNandIntArray, + unsafeAtomicOrIntArray, + unsafeAtomicXorIntArray, +) where +import Control.DeepSeq (NFData (..), deepseq) import Control.Monad -import Control.DeepSeq (NFData(..), deepseq) -import Control.Monad.Primitive (PrimMonad(..), primitive_) +import Control.Monad.Primitive (PrimMonad (..), primitive_) import Data.Massiv.Array.Delayed.Pull -- (eq, ord) import Data.Massiv.Array.Manifest.Internal import Data.Massiv.Array.Manifest.List as A import Data.Massiv.Array.Mutable import Data.Massiv.Core.Common -import Data.Massiv.Core.Operations import Data.Massiv.Core.List -import Data.Massiv.Vector.Stream as S (steps, isteps) +import Data.Massiv.Core.Operations +import Data.Massiv.Vector.Stream as S (isteps, steps) import Data.Maybe (fromMaybe) -import Data.Primitive (sizeOf, Prim) +import Data.Primitive (Prim, sizeOf) import Data.Primitive.ByteArray import qualified Data.Vector.Primitive as VP import qualified Data.Vector.Primitive.Mutable as MVP import GHC.Exts as GHC -import Prelude hiding (mapM) import System.IO.Unsafe (unsafePerformIO) - -#include "massiv.h" +import Prelude hiding (mapM) -- | Representation for `Prim`itive elements -data P = P deriving Show +data P = P deriving (Show) -data instance Array P ix e = PArray { pComp :: !Comp - , pSize :: !(Sz ix) - , pOffset :: {-# UNPACK #-} !Int - , pData :: {-# UNPACK #-} !ByteArray - } +data instance Array P ix e = PArray + { pComp :: !Comp + , pSize :: !(Sz ix) + , pOffset :: {-# UNPACK #-} !Int + , pData :: {-# UNPACK #-} !ByteArray + } -data instance MArray s P ix e = - MPArray !(Sz ix) {-# UNPACK #-} !Int {-# UNPACK #-} !(MutableByteArray s) +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 @@ -110,11 +108,10 @@ instance Strategy P where getComp = pComp {-# INLINE getComp #-} - setComp c arr = arr { pComp = c } + setComp c arr = arr{pComp = c} {-# INLINE setComp #-} repr = P - instance Index ix => Shape P ix where maxLinearSize = Just . SafeSz . elemsCount {-# INLINE maxLinearSize #-} @@ -122,13 +119,12 @@ instance Size P where size = pSize {-# INLINE size #-} - unsafeResize !sz !arr = arr { pSize = sz } + 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) + indexAssert "P.unsafeLinearIndex" (SafeSz . elemsBA _arr) indexByteArray a (i + o) {-# INLINE unsafeLinearIndex #-} unsafeOuterSlice (PArray c _ o a) szL i = @@ -138,12 +134,9 @@ unsafeLinearSlice i k (PArray c _ o a) = PArray c k (i + o) a {-# INLINE unsafeLinearSlice #-} - 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) + indexAssert "P.unsafeLinearIndexM" (const (Sz (totalElem _sz))) indexByteArray a (i + o) {-# INLINE unsafeLinearIndexM #-} sizeOfMArray (MPArray sz _ _) = sz @@ -164,23 +157,22 @@ unsafeNew sz | n <= (maxBound :: Int) `div` eSize = MPArray sz 0 <$> newByteArray (n * eSize) | otherwise = error $ "Array size is too big: " ++ show sz - where !n = totalElem sz - !eSize = sizeOf (undefined :: e) + where + !n = totalElem sz + !eSize = sizeOf (undefined :: e) {-# INLINE unsafeNew #-} initialize (MPArray sz o mba) = let k = totalElem sz * sizeOf (undefined :: e) - in when (k > 0) $ fillByteArray mba o k 0 + in when (k > 0) $ fillByteArray mba o k 0 {-# INLINE initialize #-} unsafeLinearRead _mpa@(MPArray _sz o ma) i = - INDEX_CHECK("(Manifest P ix e).unsafeLinearRead", - const (Sz (totalElem _sz)), readByteArray) ma (i + o) + indexAssert "P.unsafeLinearRead" (const (Sz (totalElem _sz))) readByteArray ma (i + o) {-# INLINE unsafeLinearRead #-} unsafeLinearWrite _mpa@(MPArray _sz o ma) i = - INDEX_CHECK("(Manifest P ix e).unsafeLinearWrite", - const (Sz (totalElem _sz)), writeByteArray) ma (i + o) + indexAssert "P.unsafeLinearWrite" (const (Sz (totalElem _sz))) writeByteArray ma (i + o) {-# INLINE unsafeLinearWrite #-} unsafeLinearSet (MPArray _ o ma) offset (SafeSz sz) = setByteArray ma (offset + o) sz @@ -188,12 +180,14 @@ unsafeLinearCopy (MPArray _ oFrom maFrom) iFrom (MPArray _ oTo maTo) iTo (Sz k) = copyMutableByteArray maTo ((oTo + iTo) * esz) maFrom ((oFrom + iFrom) * esz) (k * esz) - where esz = sizeOf (undefined :: e) + where + esz = sizeOf (undefined :: e) {-# INLINE unsafeLinearCopy #-} unsafeArrayLinearCopy (PArray _ _ oFrom aFrom) iFrom (MPArray _ oTo maTo) iTo (Sz k) = copyByteArray maTo ((oTo + iTo) * esz) aFrom ((oFrom + iFrom) * esz) (k * esz) - where esz = sizeOf (undefined :: e) + where + esz = sizeOf (undefined :: e) {-# INLINE unsafeArrayLinearCopy #-} unsafeLinearShrink (MPArray _ o ma) sz = do @@ -205,7 +199,6 @@ MPArray sz o <$> resizeMutableByteArray ma ((o + totalElem sz) * sizeOf (undefined :: e)) {-# INLINE unsafeLinearGrow #-} - instance (Prim e, Index ix) => Load P ix e where makeArray comp sz f = compute (makeArray comp sz f :: Array D ix e) {-# INLINE makeArray #-} @@ -227,7 +220,6 @@ toStreamIx = S.isteps {-# INLINE toStreamIx #-} - instance (Prim e, Num e) => FoldNumeric P e where unsafeDotProduct = defaultUnsafeDotProduct {-# INLINE unsafeDotProduct #-} @@ -242,10 +234,8 @@ unsafeLiftArray2 = defaultUnsafeLiftArray2 {-# INLINE unsafeLiftArray2 #-} - instance (Prim e, Floating e) => NumericFloat P e - instance (Prim e, IsList (Array L ix e), Ragged L ix e) => IsList (Array P ix e) where type Item (Array P ix e) = Item (Array L ix e) fromList = A.fromLists' Seq @@ -253,17 +243,14 @@ toList = GHC.toList . toListArray {-# INLINE toList #-} - -elemsBA :: forall proxy e . Prim e => proxy e -> ByteArray -> Int +elemsBA :: forall proxy e. Prim e => proxy e -> ByteArray -> Int elemsBA _ a = sizeofByteArray a `div` sizeOf (undefined :: e) {-# INLINE elemsBA #-} - -elemsMBA :: forall proxy e s . Prim e => proxy e -> MutableByteArray s -> Int +elemsMBA :: forall proxy e s. Prim e => proxy e -> MutableByteArray s -> Int elemsMBA _ a = sizeofMutableByteArray a `div` sizeOf (undefined :: e) {-# INLINE elemsMBA #-} - -- | /O(n)/ - Ensure that the size matches the internal `ByteArray`. If not make a copy of -- the slice and return it as `ByteArray` -- @@ -281,7 +268,6 @@ unwrapByteArray = pData {-# INLINE unwrapByteArray #-} - -- | /O(1)/ - Extract potential linear offset into the underlying `ByteArray`, which can -- also be extracted with `unwrapByteArray`. -- @@ -290,16 +276,14 @@ unwrapByteArrayOffset = pOffset {-# INLINE unwrapByteArrayOffset #-} - -- | /O(1)/ - Unwrap Ensure that the size matches the internal `ByteArray`. -- -- @since 0.5.0 toByteArrayM :: (Prim e, Index ix, MonadThrow m) => Array P ix e -> m ByteArray -toByteArrayM arr@PArray {pSize, pData} = do +toByteArrayM arr@PArray{pSize, pData} = do pData <$ guardNumberOfElements pSize (Sz (elemsBA arr pData)) {-# INLINE toByteArrayM #-} - -- | /O(1)/ - Construct a primitive array from the `ByteArray`. Will return `Nothing` if -- number of elements doesn't match. -- @@ -312,8 +296,8 @@ -- number of elements doesn't match. -- -- @since 0.5.9 -fromByteArrayOffsetM :: - (MonadThrow m, Index ix, Prim e) => Comp -> Sz ix -> Int -> ByteArray -> m (Array P ix e) +fromByteArrayOffsetM + :: (MonadThrow m, Index ix, Prim e) => Comp -> Sz ix -> Int -> ByteArray -> m (Array P ix e) fromByteArrayOffsetM comp sz off ba = arr <$ guardNumberOfElements sz (SafeSz (elemsBA arr ba - off)) where @@ -323,11 +307,10 @@ -- | /O(1)/ - Construct a flat Array from `ByteArray` -- -- @since 0.4.0 -fromByteArray :: forall e . Prim e => Comp -> ByteArray -> Array P Ix1 e +fromByteArray :: forall e. Prim e => Comp -> ByteArray -> Array P Ix1 e fromByteArray comp ba = PArray comp (SafeSz (elemsBA (Proxy :: Proxy e) ba)) 0 ba {-# INLINE fromByteArray #-} - -- | /O(1)/ - Extract the internal `MutableByteArray`. This will discard any possible -- slicing that has been applied to the array. -- @@ -349,8 +332,9 @@ -- that the source at the resulting array are still pointing to the same location in memory. -- -- @since 0.5.0 -toMutableByteArray :: - forall ix e m. (Prim e, Index ix, PrimMonad m) +toMutableByteArray + :: forall ix e m + . (Prim e, Index ix, PrimMonad m) => MArray (PrimState m) P ix e -> m (Bool, MutableByteArray (PrimState m)) toMutableByteArray marr@(MPArray sz offset mbas) = @@ -364,7 +348,6 @@ pure (False, mbad) {-# INLINE toMutableByteArray #-} - -- | /O(1)/ - Extract the internal `MutableByteArray`. -- -- @since 0.2.1 @@ -373,13 +356,12 @@ mba <$ guardNumberOfElements sz (Sz (elemsMBA marr mba)) {-# INLINE toMutableByteArrayM #-} - -- | /O(1)/ - Construct a primitive mutable array from the `MutableByteArray`. Will throw -- `SizeElementsMismatchException` if number of elements doesn't match. -- -- @since 0.3.0 -fromMutableByteArrayM :: - (MonadThrow m, Index ix, Prim e) => Sz ix -> MutableByteArray s -> m (MArray s P ix e) +fromMutableByteArrayM + :: (MonadThrow m, Index ix, Prim e) => Sz ix -> MutableByteArray s -> m (MArray s P ix e) fromMutableByteArrayM sz = fromMutableByteArrayOffsetM sz 0 {-# INLINE fromMutableByteArrayM #-} @@ -387,33 +369,28 @@ -- `SizeElementsMismatchException` if number of elements doesn't match. -- -- @since 0.5.9 -fromMutableByteArrayOffsetM :: - (MonadThrow m, Index ix, Prim e) => Sz ix -> Ix1 -> MutableByteArray s -> m (MArray s P ix e) +fromMutableByteArrayOffsetM + :: (MonadThrow m, Index ix, Prim e) => Sz ix -> Ix1 -> MutableByteArray s -> m (MArray s P ix e) fromMutableByteArrayOffsetM sz off mba = marr <$ guardNumberOfElements sz (SafeSz (elemsMBA marr mba - off)) where marr = MPArray sz off mba {-# INLINE fromMutableByteArrayOffsetM #-} - -- | /O(1)/ - Construct a flat Array from `MutableByteArray` -- -- @since 0.4.0 -fromMutableByteArray :: forall e s . Prim e => MutableByteArray s -> MArray s P Ix1 e +fromMutableByteArray :: forall e s. Prim e => MutableByteArray s -> MArray s P Ix1 e fromMutableByteArray mba = MPArray (SafeSz (elemsMBA (Proxy :: Proxy e) mba)) 0 mba {-# INLINE fromMutableByteArray #-} - - - -- | /O(1)/ - Cast a primitive array to a primitive vector. -- -- @since 0.5.0 toPrimitiveVector :: Index ix => Array P ix e -> VP.Vector e -toPrimitiveVector PArray {pSize, pOffset, pData} = VP.Vector pOffset (totalElem pSize) pData +toPrimitiveVector PArray{pSize, pOffset, pData} = VP.Vector pOffset (totalElem pSize) pData {-# INLINE toPrimitiveVector #-} - -- | /O(1)/ - Cast a mutable primitive array to a mutable primitive vector. -- -- @since 0.5.0 @@ -421,13 +398,12 @@ toPrimitiveMVector (MPArray sz offset mba) = MVP.MVector offset (totalElem sz) mba {-# INLINE toPrimitiveMVector #-} - -- | /O(1)/ - Cast a primitive vector to a primitive array. -- -- @since 0.5.0 fromPrimitiveVector :: VP.Vector e -> Array P Ix1 e fromPrimitiveVector (VP.Vector offset len ba) = - PArray {pComp = Seq, pSize = SafeSz len, pOffset = offset, pData = ba} + PArray{pComp = Seq, pSize = SafeSz len, pOffset = offset, pData = ba} {-# INLINE fromPrimitiveVector #-} -- | /O(1)/ - Cast a mutable primitive vector to a mutable primitive array. @@ -440,59 +416,67 @@ -- | Atomically read an `Int` element from the array -- -- @since 0.3.0 -unsafeAtomicReadIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> m Int +unsafeAtomicReadIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> m Int unsafeAtomicReadIntArray _mpa@(MPArray sz o mba) ix = - INDEX_CHECK( "unsafeAtomicReadIntArray" - , SafeSz . elemsMBA _mpa - , \(MutableByteArray mba#) (I# i#) -> - primitive $ \s# -> - case atomicReadIntArray# mba# i# s# of - (# s'#, e# #) -> (# s'#, I# e# #)) - mba - (o + toLinearIndex sz ix) + indexAssert + "P.unsafeAtomicReadIntArray" + (SafeSz . elemsMBA _mpa) + ( \(MutableByteArray mba#) (I# i#) -> + primitive $ \s# -> + case atomicReadIntArray# mba# i# s# of + (# s'#, e# #) -> (# s'#, I# e# #) + ) + mba + (o + toLinearIndex sz ix) {-# INLINE unsafeAtomicReadIntArray #-} -- | Atomically write an `Int` element int the array -- -- @since 0.3.0 -unsafeAtomicWriteIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m () +unsafeAtomicWriteIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m () unsafeAtomicWriteIntArray _mpa@(MPArray sz o mba) ix (I# e#) = - INDEX_CHECK( "unsafeAtomicWriteIntArray" - , SafeSz . elemsMBA _mpa - , \(MutableByteArray mba#) (I# i#) -> - primitive_ (atomicWriteIntArray# mba# i# e#)) - mba - (o + toLinearIndex sz ix) + indexAssert + "P.unsafeAtomicWriteIntArray" + (SafeSz . elemsMBA _mpa) + ( \(MutableByteArray mba#) (I# i#) -> + primitive_ (atomicWriteIntArray# mba# i# e#) + ) + mba + (o + toLinearIndex sz ix) {-# INLINE unsafeAtomicWriteIntArray #-} -- | Atomically CAS an `Int` in the array. Returns the old value. -- -- @since 0.3.0 -unsafeCasIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> Int -> m Int +unsafeCasIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> Int -> m Int unsafeCasIntArray _mpa@(MPArray sz o mba) ix (I# e#) (I# n#) = - INDEX_CHECK( "unsafeCasIntArray" - , SafeSz . elemsMBA _mpa - , \(MutableByteArray mba#) (I# i#) -> - primitive $ \s# -> - case casIntArray# mba# i# e# n# s# of - (# s'#, o# #) -> (# s'#, I# o# #)) - mba - (o + toLinearIndex sz ix) + indexAssert + "P.unsafeCasIntArray" + (SafeSz . elemsMBA _mpa) + ( \(MutableByteArray mba#) (I# i#) -> + primitive $ \s# -> + case casIntArray# mba# i# e# n# s# of + (# s'#, o# #) -> (# s'#, I# o# #) + ) + mba + (o + toLinearIndex sz ix) {-# INLINE unsafeCasIntArray #-} - -- | Atomically modify an `Int` element of the array. Returns the old value. -- -- @since 0.3.0 -unsafeAtomicModifyIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> (Int -> Int) -> m Int +unsafeAtomicModifyIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> (Int -> Int) -> m Int unsafeAtomicModifyIntArray _mpa@(MPArray sz o mba) ix f = - INDEX_CHECK("unsafeAtomicModifyIntArray", SafeSz . elemsMBA _mpa, atomicModify) - mba - (o + toLinearIndex sz ix) + indexAssert + "P.unsafeAtomicModifyIntArray" + (SafeSz . elemsMBA _mpa) + atomicModify + mba + (o + toLinearIndex sz ix) where atomicModify (MutableByteArray mba#) (I# i#) = let go s# o# = @@ -501,111 +485,117 @@ (# s'#, o'# #) -> case o# ==# o'# of 0# -> go s# o'# - _ -> (# s'#, I# o# #) + _ -> (# s'#, I# o# #) in primitive $ \s# -> case atomicReadIntArray# mba# i# s# of (# s'#, o# #) -> go s'# o# {-# INLINE atomicModify #-} {-# INLINE unsafeAtomicModifyIntArray #-} - -- | Atomically add to an `Int` element in the array. Returns the old value. -- -- @since 0.3.0 -unsafeAtomicAddIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int +unsafeAtomicAddIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int unsafeAtomicAddIntArray _mpa@(MPArray sz o mba) ix (I# e#) = - INDEX_CHECK( "unsafeAtomicAddIntArray" - , SafeSz . elemsMBA _mpa - , \(MutableByteArray mba#) (I# i#) -> - primitive $ \s# -> - case fetchAddIntArray# mba# i# e# s# of - (# s'#, p# #) -> (# s'#, I# p# #)) - mba - (o + toLinearIndex sz ix) + indexAssert + "P.unsafeAtomicAddIntArray" + (SafeSz . elemsMBA _mpa) + ( \(MutableByteArray mba#) (I# i#) -> + primitive $ \s# -> + case fetchAddIntArray# mba# i# e# s# of + (# s'#, p# #) -> (# s'#, I# p# #) + ) + mba + (o + toLinearIndex sz ix) {-# INLINE unsafeAtomicAddIntArray #-} - -- | Atomically subtract from an `Int` element in the array. Returns the old value. -- -- @since 0.3.0 -unsafeAtomicSubIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int +unsafeAtomicSubIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int unsafeAtomicSubIntArray _mpa@(MPArray sz o mba) ix (I# e#) = - INDEX_CHECK( "unsafeAtomicSubIntArray" - , SafeSz . elemsMBA _mpa - , \(MutableByteArray mba#) (I# i#) -> - primitive $ \s# -> - case fetchSubIntArray# mba# i# e# s# of - (# s'#, p# #) -> (# s'#, I# p# #)) - mba - (o + toLinearIndex sz ix) + indexAssert + "P.unsafeAtomicSubIntArray" + (SafeSz . elemsMBA _mpa) + ( \(MutableByteArray mba#) (I# i#) -> + primitive $ \s# -> + case fetchSubIntArray# mba# i# e# s# of + (# s'#, p# #) -> (# s'#, I# p# #) + ) + mba + (o + toLinearIndex sz ix) {-# INLINE unsafeAtomicSubIntArray #-} - -- | Atomically AND an `Int` element in the array. Returns the old value. -- -- @since 0.3.0 -unsafeAtomicAndIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int +unsafeAtomicAndIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int unsafeAtomicAndIntArray _mpa@(MPArray sz o mba) ix (I# e#) = - INDEX_CHECK( "unsafeAtomicAndIntArray" - , SafeSz . elemsMBA _mpa - , \(MutableByteArray mba#) (I# i#) -> - primitive $ \s# -> - case fetchAndIntArray# mba# i# e# s# of - (# s'#, p# #) -> (# s'#, I# p# #)) - mba - (o + toLinearIndex sz ix) + indexAssert + "P.unsafeAtomicAndIntArray" + (SafeSz . elemsMBA _mpa) + ( \(MutableByteArray mba#) (I# i#) -> + primitive $ \s# -> + case fetchAndIntArray# mba# i# e# s# of + (# s'#, p# #) -> (# s'#, I# p# #) + ) + mba + (o + toLinearIndex sz ix) {-# INLINE unsafeAtomicAndIntArray #-} - -- | Atomically NAND an `Int` element in the array. Returns the old value. -- -- @since 0.3.0 -unsafeAtomicNandIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int +unsafeAtomicNandIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int unsafeAtomicNandIntArray _mpa@(MPArray sz o mba) ix (I# e#) = - INDEX_CHECK( "unsafeAtomicNandIntArray" - , SafeSz . elemsMBA _mpa - , \(MutableByteArray mba#) (I# i#) -> - primitive $ \s# -> - case fetchNandIntArray# mba# i# e# s# of - (# s'#, p# #) -> (# s'#, I# p# #)) - mba - (o + toLinearIndex sz ix) + indexAssert + "P.unsafeAtomicNandIntArray" + (SafeSz . elemsMBA _mpa) + ( \(MutableByteArray mba#) (I# i#) -> + primitive $ \s# -> + case fetchNandIntArray# mba# i# e# s# of + (# s'#, p# #) -> (# s'#, I# p# #) + ) + mba + (o + toLinearIndex sz ix) {-# INLINE unsafeAtomicNandIntArray #-} - -- | Atomically OR an `Int` element in the array. Returns the old value. -- -- @since 0.3.0 -unsafeAtomicOrIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int +unsafeAtomicOrIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int unsafeAtomicOrIntArray _mpa@(MPArray sz o mba) ix (I# e#) = - INDEX_CHECK( "unsafeAtomicOrIntArray" - , SafeSz . elemsMBA _mpa - , \(MutableByteArray mba#) (I# i#) -> - primitive $ \s# -> - case fetchOrIntArray# mba# i# e# s# of - (# s'#, p# #) -> (# s'#, I# p# #)) - mba - (o + toLinearIndex sz ix) + indexAssert + "P.unsafeAtomicOrIntArray" + (SafeSz . elemsMBA _mpa) + ( \(MutableByteArray mba#) (I# i#) -> + primitive $ \s# -> + case fetchOrIntArray# mba# i# e# s# of + (# s'#, p# #) -> (# s'#, I# p# #) + ) + mba + (o + toLinearIndex sz ix) {-# INLINE unsafeAtomicOrIntArray #-} - -- | Atomically XOR an `Int` element in the array. Returns the old value. -- -- @since 0.3.0 -unsafeAtomicXorIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int +unsafeAtomicXorIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int unsafeAtomicXorIntArray _mpa@(MPArray sz o mba) ix (I# e#) = - INDEX_CHECK( "unsafeAtomicXorIntArray" - , SafeSz . elemsMBA _mpa - , \(MutableByteArray mba#) (I# i#) -> - primitive $ \s# -> - case fetchXorIntArray# mba# i# e# s# of - (# s'#, p# #) -> (# s'#, I# p# #)) - mba - (o + toLinearIndex sz ix) + indexAssert + "P.unsafeAtomicXorIntArray" + (SafeSz . elemsMBA _mpa) + ( \(MutableByteArray mba#) (I# i#) -> + primitive $ \s# -> + case fetchXorIntArray# mba# i# e# s# of + (# s'#, p# #) -> (# s'#, I# p# #) + ) + mba + (o + toLinearIndex sz ix) {-# INLINE unsafeAtomicXorIntArray #-} diff --git a/src/Data/Massiv/Array/Manifest/Storable.hs b/src/Data/Massiv/Array/Manifest/Storable.hs --- a/src/Data/Massiv/Array/Manifest/Storable.hs +++ b/src/Data/Massiv/Array/Manifest/Storable.hs @@ -7,6 +7,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UndecidableInstances #-} + -- | -- Module : Data.Massiv.Array.Manifest.Storable -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -14,28 +15,27 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Manifest.Storable - ( S (..) - , Array(..) - , MArray(..) - , Storable - , toStorableVector - , toStorableMVector - , fromStorableVector - , fromStorableMVector - , withPtr - , unsafeWithPtr - , unsafeMallocMArray - , unsafeArrayToForeignPtr - , unsafeMArrayToForeignPtr - , unsafeArrayFromForeignPtr - , unsafeArrayFromForeignPtr0 - , unsafeMArrayFromForeignPtr - , unsafeMArrayFromForeignPtr0 - ) where +module Data.Massiv.Array.Manifest.Storable ( + S (..), + Array (..), + MArray (..), + Storable, + toStorableVector, + toStorableMVector, + fromStorableVector, + fromStorableMVector, + withPtr, + unsafeWithPtr, + unsafeMallocMArray, + unsafeArrayToForeignPtr, + unsafeMArrayToForeignPtr, + unsafeArrayFromForeignPtr, + unsafeArrayFromForeignPtr0, + unsafeMArrayFromForeignPtr, + unsafeMArrayFromForeignPtr0, +) where -import Control.DeepSeq (NFData(..), deepseq) +import Control.DeepSeq (NFData (..), deepseq) import Control.Exception import Control.Monad import Control.Monad.IO.Unlift @@ -58,23 +58,22 @@ import Foreign.Storable import GHC.Exts as GHC import GHC.ForeignPtr -import Prelude hiding (mapM) import System.IO.Unsafe (unsafePerformIO) import Unsafe.Coerce +import Prelude hiding (mapM) 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 S = S deriving (Show) -data instance Array S ix e = SArray { sComp :: !Comp - , sSize :: !(Sz ix) - , sData :: {-# UNPACK #-} !(ForeignPtr 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) @@ -101,14 +100,14 @@ instance Strategy S where getComp = sComp {-# INLINE getComp #-} - setComp c arr = arr { sComp = c } + setComp c arr = arr{sComp = c} {-# INLINE setComp #-} repr = S plusFp :: ForeignPtr a -> Int -> ForeignPtr b plusFp (ForeignPtr addr c) (I# d) = ForeignPtr (plusAddr# addr d) c -advanceForeignPtr :: forall e . Storable e => ForeignPtr e -> Int -> ForeignPtr e +advanceForeignPtr :: forall e. Storable e => ForeignPtr e -> Int -> ForeignPtr e advanceForeignPtr fp i = plusFp fp (i * sizeOf (undefined :: e)) {-# INLINE advanceForeignPtr #-} @@ -118,12 +117,12 @@ instance Storable e => Source S e where unsafeLinearIndex (SArray _ _sz fp) = - INDEX_CHECK("(Source S ix e).unsafeLinearIndex", const (toLinearSz _sz), indexForeignPtr) fp + indexAssert "S.unsafeLinearIndex" (const (toLinearSz _sz)) indexForeignPtr fp {-# INLINE unsafeLinearIndex #-} unsafeOuterSlice (SArray c _ fp) szL i = let k = totalElem szL - in SArray c szL $ advanceForeignPtr fp (i * k) + in SArray c szL $ advanceForeignPtr fp (i * k) {-# INLINE unsafeOuterSlice #-} unsafeLinearSlice i k (SArray c _ fp) = @@ -137,14 +136,12 @@ instance Size S where size = sSize {-# INLINE size #-} - unsafeResize !sz !arr = arr { sSize = sz } + unsafeResize !sz !arr = arr{sSize = sz} {-# INLINE unsafeResize #-} - instance Storable e => Manifest S e where - unsafeLinearIndexM (SArray _ _sz fp) = - INDEX_CHECK("(Source S ix e).unsafeLinearIndex", const (toLinearSz _sz), indexForeignPtr) fp + indexAssert "S.unsafeLinearIndex" (const (toLinearSz _sz)) indexForeignPtr fp {-# INLINE unsafeLinearIndexM #-} sizeOfMArray (MSArray sz _) = sz @@ -178,12 +175,24 @@ setPtr (castPtr p) (totalElem sz * sizeOf (undefined :: e)) (0 :: Word8) {-# INLINE initialize #-} - unsafeLinearRead (MSArray _sz fp) o = unsafeIOToPrim $ - INDEX_CHECK("(Manifest S ix e).unsafeLinearRead", const (toLinearSz _sz), (\_ _ -> unsafeWithForeignPtr fp (`peekElemOff` o))) fp o + unsafeLinearRead (MSArray _sz fp) o = + unsafeIOToPrim $ + indexAssert + "S.unsafeLinearRead" + (const (toLinearSz _sz)) + (\_ _ -> unsafeWithForeignPtr fp (`peekElemOff` o)) + fp + o {-# INLINE unsafeLinearRead #-} - 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 + unsafeLinearWrite (MSArray _sz fp) o e = + unsafeIOToPrim $ + indexAssert + "S.unsafeLinearWrite" + (const (toLinearSz _sz)) + (\_ _ -> unsafeWithForeignPtr fp (\p -> pokeElemOff p o e)) + fp + o {-# INLINE unsafeLinearWrite #-} unsafeLinearSet (MSArray _ fp) i k e = @@ -192,8 +201,8 @@ unsafeLinearCopy (MSArray _ fpFrom) iFrom (MSArray _ fpTo) iTo (Sz k) = do unsafePrimToPrim $ - withForeignPtr fpFrom $ \ ptrFrom -> - withForeignPtr fpTo $ \ ptrTo -> do + withForeignPtr fpFrom $ \ptrFrom -> + withForeignPtr fpTo $ \ptrTo -> do let ptrFrom' = advancePtr ptrFrom iFrom ptrTo' = advancePtr ptrTo iTo copyArray ptrTo' ptrFrom' k @@ -240,7 +249,6 @@ toStreamIx = S.isteps {-# INLINE toStreamIx #-} - instance (Storable e, Num e) => FoldNumeric S e where unsafeDotProduct = defaultUnsafeDotProduct {-# INLINE unsafeDotProduct #-} @@ -257,7 +265,6 @@ instance (Storable e, Floating e) => NumericFloat S e - instance (Storable e, IsList (Array L ix e), Ragged L ix e) => IsList (Array S ix e) where type Item (Array S ix e) = Item (Array L ix e) fromList = A.fromLists' Seq @@ -273,7 +280,6 @@ 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 @@ -281,18 +287,17 @@ 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 :: 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)) + -- this hack is needed to workaround the redundant Storable constraint + -- see haskell/vector#394 + unsafeCoerce $ + VS.unsafeFromForeignPtr0 (castForeignPtr (sData arr) :: ForeignPtr Word) (totalElem (sSize arr)) {-# INLINE toStorableVector #-} - -- | /O(1)/ - Unwrap storable mutable array and pull out the underlying storable mutable vector. -- -- @since 0.2.1 @@ -308,7 +313,7 @@ -- 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} + (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. @@ -318,7 +323,6 @@ 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 @@ -337,7 +341,7 @@ -- -- @since 0.3.0 unsafeArrayFromForeignPtr0 :: Comp -> ForeignPtr e -> Sz1 -> Vector S e -unsafeArrayFromForeignPtr0 comp fp sz = SArray {sComp = comp, sSize = sz, sData = fp} +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. @@ -345,10 +349,9 @@ -- @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 = advanceForeignPtr ptr offset} + SArray{sComp = comp, sSize = sz, sData = advanceForeignPtr ptr offset} {-# INLINE unsafeArrayFromForeignPtr #-} - -- | /O(1)/ - Wrap a `ForeignPtr` and it's size into a mutable storable array. It is still safe to -- modify the pointer, unless the array gets frozen prior to modification. -- @@ -357,7 +360,6 @@ unsafeMArrayFromForeignPtr0 fp sz = MSArray sz fp {-# INLINE unsafeMArrayFromForeignPtr0 #-} - -- | /O(1)/ - Wrap a `ForeignPtr`, an offset and it's size into a mutable storable array. It is -- still safe to modify the pointer, unless the array gets frozen prior to modification. -- @@ -366,13 +368,13 @@ unsafeMArrayFromForeignPtr fp offset sz = MSArray sz (advanceForeignPtr fp offset) {-# INLINE unsafeMArrayFromForeignPtr #-} - -- | Allocate memory using @malloc@ on C heap, instead of on Haskell heap. Memory is left -- uninitialized -- -- @since 0.5.9 -unsafeMallocMArray :: - forall ix e m. (Index ix, Storable e, PrimMonad m) +unsafeMallocMArray + :: forall ix e m + . (Index ix, Storable e, PrimMonad m) => Sz ix -> m (MArray (PrimState m) S ix e) unsafeMallocMArray sz = unsafePrimToPrim $ do @@ -382,7 +384,6 @@ newForeignPtr finalizerFree ptr pure $ MSArray sz foreignPtr {-# INLINE unsafeMallocMArray #-} - #if !MIN_VERSION_base(4,15,0) -- | A compatibility wrapper for 'GHC.ForeignPtr.unsafeWithForeignPtr' provided diff --git a/src/Data/Massiv/Array/Manifest/Unboxed.hs b/src/Data/Massiv/Array/Manifest/Unboxed.hs --- a/src/Data/Massiv/Array/Manifest/Unboxed.hs +++ b/src/Data/Massiv/Array/Manifest/Unboxed.hs @@ -1,11 +1,11 @@ {-# LANGUAGE BangPatterns #-} -{-# LANGUAGE CPP #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UndecidableInstances #-} + -- | -- Module : Data.Massiv.Array.Manifest.Unboxed -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -13,19 +13,18 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Manifest.Unboxed - ( U (..) - , Unbox - , Array(..) - , MArray(..) - , toUnboxedVector - , toUnboxedMVector - , fromUnboxedVector - , fromUnboxedMVector - ) where +module Data.Massiv.Array.Manifest.Unboxed ( + U (..), + Unbox, + Array (..), + MArray (..), + toUnboxedVector, + toUnboxedMVector, + fromUnboxedVector, + fromUnboxedMVector, +) where -import Control.DeepSeq (NFData(..), deepseq) +import Control.DeepSeq (NFData (..), deepseq) import Control.Monad.Primitive (stToPrim) import Data.Massiv.Array.Delayed.Pull (D, compareArrays, eqArrays) import Data.Massiv.Array.Manifest.Internal @@ -39,19 +38,18 @@ import Data.Vector.Unboxed (Unbox) import qualified Data.Vector.Unboxed as VU import qualified Data.Vector.Unboxed.Mutable as MVU -import GHC.Exts as GHC (IsList(..)) -import Prelude hiding (mapM) +import GHC.Exts as GHC (IsList (..)) import System.IO.Unsafe (unsafePerformIO) - -#include "massiv.h" +import Prelude hiding (mapM) -- | Representation for `Unbox`ed elements -data U = U deriving Show +data U = U deriving (Show) -data instance Array U ix e = UArray { uComp :: !Comp - , uSize :: !(Sz ix) - , uData :: !(VU.Vector e) - } +data instance Array U ix e = UArray + { uComp :: !Comp + , uSize :: !(Sz ix) + , uData :: !(VU.Vector e) + } data instance MArray s U ix e = MUArray !(Sz ix) !(VU.MVector s e) instance (Ragged L ix e, Show e, Unbox e) => Show (Array U ix e) where @@ -69,11 +67,10 @@ instance Strategy U where getComp = uComp {-# INLINE getComp #-} - setComp c arr = arr { uComp = c } + setComp c arr = arr{uComp = c} {-# INLINE setComp #-} repr = U - instance (Unbox e, Eq e, Index ix) => Eq (Array U ix e) where (==) = eqArrays (==) {-# INLINE (==) #-} @@ -82,15 +79,14 @@ compare = compareArrays compare {-# INLINE compare #-} - instance Unbox e => Source U e where unsafeLinearIndex (UArray _ _ v) = - INDEX_CHECK("(Source U ix e).unsafeLinearIndex", Sz . VU.length, VU.unsafeIndex) v + indexAssert "U.unsafeLinearIndex" (SafeSz . VU.length) VU.unsafeIndex v {-# INLINE unsafeLinearIndex #-} unsafeOuterSlice (UArray c _ v) szL i = let k = totalElem szL - in UArray c szL $ VU.unsafeSlice (i * k) k v + 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 @@ -103,7 +99,7 @@ instance Size U where size = uSize {-# INLINE size #-} - unsafeResize !sz !arr = arr { uSize = sz } + unsafeResize !sz !arr = arr{uSize = sz} {-# INLINE unsafeResize #-} instance (Unbox e, Index ix) => Load U ix e where @@ -122,14 +118,9 @@ instance (Unbox e, Index ix) => StrideLoad U ix e - - - - instance Unbox e => Manifest U e where - unsafeLinearIndexM (UArray _ _ v) = - INDEX_CHECK("(Manifest U ix e).unsafeLinearIndexM", Sz . VU.length, VU.unsafeIndex) v + indexAssert "S.unsafeLinearIndexM" (SafeSz . VU.length) VU.unsafeIndex v {-# INLINE unsafeLinearIndexM #-} sizeOfMArray (MUArray sz _) = sz @@ -158,24 +149,22 @@ {-# INLINE unsafeLinearCopy #-} unsafeLinearRead (MUArray _ mv) = - INDEX_CHECK("(Manifest U ix e).unsafeLinearRead", Sz . MVU.length, MVU.unsafeRead) mv + indexAssert "U.unsafeLinearRead" (Sz . MVU.length) MVU.unsafeRead mv {-# INLINE unsafeLinearRead #-} unsafeLinearWrite (MUArray _ mv) = - INDEX_CHECK("(Manifest U ix e).unsafeLinearWrite", Sz . MVU.length, MVU.unsafeWrite) mv + indexAssert "U.unsafeLinearWrite" (Sz . MVU.length) MVU.unsafeWrite mv {-# INLINE unsafeLinearWrite #-} unsafeLinearGrow (MUArray _ mv) sz = MUArray sz <$> MVU.unsafeGrow mv (totalElem sz) {-# INLINE unsafeLinearGrow #-} - instance (Index ix, Unbox e) => Stream U ix e where toStream = S.steps {-# INLINE toStream #-} toStreamIx = S.isteps {-# INLINE toStreamIx #-} - 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 @@ -183,7 +172,6 @@ toList = GHC.toList . toListArray {-# INLINE toList #-} - instance (VU.Unbox e, Num e) => FoldNumeric U e where unsafeDotProduct = defaultUnsafeDotProduct {-# INLINE unsafeDotProduct #-} @@ -198,7 +186,6 @@ unsafeLiftArray2 = defaultUnsafeLiftArray2 {-# INLINE unsafeLiftArray2 #-} - -- | /O(1)/ - Unwrap unboxed array and pull out the underlying unboxed vector. -- -- @since 0.2.1 @@ -206,7 +193,6 @@ toUnboxedVector = uData {-# INLINE toUnboxedVector #-} - -- | /O(1)/ - Unwrap unboxed mutable array and pull out the underlying unboxed mutable vector. -- -- @since 0.2.1 @@ -214,15 +200,12 @@ toUnboxedMVector (MUArray _ mv) = mv {-# INLINE toUnboxedMVector #-} - - -- | /O(1)/ - Wrap an unboxed vector and produce an unboxed flat array. -- -- @since 0.6.0 fromUnboxedVector :: VU.Unbox e => Comp -> VU.Vector e -> Vector U e fromUnboxedVector comp v = UArray comp (SafeSz (VU.length v)) v {-# INLINE fromUnboxedVector #-} - -- | /O(1)/ - Wrap an unboxed mutable vector and produce a mutable unboxed flat array. -- diff --git a/src/Data/Massiv/Array/Manifest/Vector.hs b/src/Data/Massiv/Array/Manifest/Vector.hs --- a/src/Data/Massiv/Array/Manifest/Vector.hs +++ b/src/Data/Massiv/Array/Manifest/Vector.hs @@ -4,6 +4,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE TypeOperators #-} + -- | -- Module : Data.Massiv.Array.Manifest.Vector -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -11,16 +12,15 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Manifest.Vector - ( fromVectorM - , fromVector' - , castFromVector - , toVector - , castToVector - , ARepr - , VRepr - ) where +module Data.Massiv.Array.Manifest.Vector ( + fromVectorM, + fromVector', + castFromVector, + toVector, + castToVector, + ARepr, + VRepr, +) where import Control.Monad (guard, join, msum) import Data.Kind @@ -54,46 +54,54 @@ VRepr BN = VB.Vector VRepr BL = VB.Vector - -- | /O(1)/ - conversion from vector to an array with a corresponding representation. Will -- return `Nothing` if there is a size mismatch or if some non-standard vector type is -- supplied. Is suppplied is the boxed `Data.Vector.Vector` then it's all elements will be -- evaluated toWHNF, therefore complexity will be /O(n)/ -castFromVector :: forall v r ix e. (VG.Vector v e, Typeable v, Index ix, ARepr v ~ r) - => Comp - -> Sz ix -- ^ Size of the result Array - -> v e -- ^ Source Vector - -> Maybe (Array r ix e) +castFromVector + :: forall v r ix e + . (VG.Vector v e, Typeable v, Index ix, ARepr v ~ r) + => Comp + -> Sz ix + -- ^ Size of the result Array + -> v e + -- ^ Source Vector + -> Maybe (Array r ix e) castFromVector comp sz vector = do guard (totalElem sz == VG.length vector) msum - [ do Refl <- eqT :: Maybe (v :~: VU.Vector) - uVector <- join $ gcast1 (Just vector) - return $ UArray {uComp = comp, uSize = sz, uData = uVector} - , do Refl <- eqT :: Maybe (v :~: VS.Vector) - sVector <- join $ gcast1 (Just vector) - return $ unsafeResize sz $ fromStorableVector comp sVector - , do Refl <- eqT :: Maybe (v :~: VP.Vector) - VP.Vector o _ ba <- join $ gcast1 (Just vector) - return $ PArray {pComp = comp, pSize = sz, pOffset = o, pData = ba} - , do Refl <- eqT :: Maybe (v :~: VB.Vector) - bVector <- join $ gcast1 (Just vector) - pure $ unsafeResize sz $ setComp comp $ fromBoxedVector bVector + [ do + Refl <- eqT :: Maybe (v :~: VU.Vector) + uVector <- join $ gcast1 (Just vector) + return $ UArray{uComp = comp, uSize = sz, uData = uVector} + , do + Refl <- eqT :: Maybe (v :~: VS.Vector) + sVector <- join $ gcast1 (Just vector) + return $ unsafeResize sz $ fromStorableVector comp sVector + , do + Refl <- eqT :: Maybe (v :~: VP.Vector) + VP.Vector o _ ba <- join $ gcast1 (Just vector) + return $ PArray{pComp = comp, pSize = sz, pOffset = o, pData = ba} + , do + Refl <- eqT :: Maybe (v :~: VB.Vector) + bVector <- join $ gcast1 (Just vector) + pure $ unsafeResize sz $ setComp comp $ fromBoxedVector bVector ] {-# NOINLINE castFromVector #-} - -- | In case when resulting array representation matches the one of vector's it -- will do a /O(1)/ - conversion using `castFromVector`, otherwise Vector elements -- will be copied into a new array. Will throw an error if length of resulting -- array doesn't match the source vector length. -- -- @since 0.3.0 -fromVectorM :: - (MonadThrow m, Typeable v, VG.Vector v a, Manifest r a, Load (ARepr v) ix a, Load r ix a) +fromVectorM + :: (MonadThrow m, Typeable v, VG.Vector v a, Manifest r a, Load (ARepr v) ix a, Load r ix a) => Comp - -> Sz ix -- ^ Resulting size of the array - -> v a -- ^ Source Vector + -> Sz ix + -- ^ Resulting size of the array + -> v a + -- ^ Source Vector -> m (Array r ix a) fromVectorM comp sz v = case castFromVector comp sz v of @@ -103,15 +111,16 @@ pure (makeArrayLinear comp sz (VG.unsafeIndex v)) {-# NOINLINE fromVectorM #-} - -- | Just like `fromVectorM`, but will throw an exception on a mismatched size. -- -- @since 0.3.0 -fromVector' :: - (HasCallStack, Typeable v, VG.Vector v a, Load (ARepr v) ix a, Load r ix a, Manifest r a) +fromVector' + :: (HasCallStack, Typeable v, VG.Vector v a, Load (ARepr v) ix a, Load r ix a, Manifest r a) => Comp - -> Sz ix -- ^ Resulting size of the array - -> v a -- ^ Source Vector + -> Sz ix + -- ^ Resulting size of the array + -> v a + -- ^ Source Vector -> Array r ix a fromVector' comp sz = throwEither . fromVectorM comp sz {-# INLINE fromVector' #-} @@ -119,34 +128,40 @@ -- | /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. (Manifest r e, Index ix, VRepr r ~ v) +castToVector + :: forall v r ix e + . (Manifest r e, Index ix, VRepr r ~ v) => Array r ix e -> Maybe (v e) castToVector arr = msum - [ do Refl <- eqT :: Maybe (r :~: U) - uArr <- gcastArr arr - return $ uData uArr - , do Refl <- eqT :: Maybe (r :~: S) - sArr <- gcastArr arr - return $ toStorableVector sArr - , do Refl <- eqT :: Maybe (r :~: P) - pArr <- gcastArr arr - return $ VP.Vector (pOffset pArr) (totalElem (size arr)) $ pData pArr - , do Refl <- eqT :: Maybe (r :~: B) - bArr <- gcastArr arr - return $ toBoxedVector $ toLazyArray bArr - , do Refl <- eqT :: Maybe (r :~: BN) - bArr <- gcastArr arr - return $ toBoxedVector $ toLazyArray $ unwrapNormalForm bArr - , do Refl <- eqT :: Maybe (r :~: BL) - bArr <- gcastArr arr - return $ toBoxedVector bArr + [ do + Refl <- eqT :: Maybe (r :~: U) + uArr <- gcastArr arr + return $ uData uArr + , do + Refl <- eqT :: Maybe (r :~: S) + sArr <- gcastArr arr + return $ toStorableVector sArr + , do + Refl <- eqT :: Maybe (r :~: P) + pArr <- gcastArr arr + return $ VP.Vector (pOffset pArr) (totalElem (size arr)) $ pData pArr + , do + Refl <- eqT :: Maybe (r :~: B) + bArr <- gcastArr arr + return $ toBoxedVector $ toLazyArray bArr + , do + Refl <- eqT :: Maybe (r :~: BN) + bArr <- gcastArr arr + return $ toBoxedVector $ toLazyArray $ unwrapNormalForm bArr + , do + Refl <- eqT :: Maybe (r :~: BL) + bArr <- gcastArr arr + return $ toBoxedVector bArr ] {-# NOINLINE castToVector #-} - -- | Convert an array into a vector. Will perform a cast if resulting vector is -- of compatible representation, otherwise memory copy will occur. -- @@ -168,10 +183,9 @@ -- >>> import qualified Data.Vector.Unboxed as VU -- >>> toVector (makeArrayR S Par (Sz2 5 6) (\(i :. j) -> i + j)) :: VU.Vector Int -- [0,1,2,3,4,5,1,2,3,4,5,6,2,3,4,5,6,7,3,4,5,6,7,8,4,5,6,7,8,9] --- -toVector :: - forall r ix e v. - ( Manifest r e +toVector + :: forall r ix e v + . ( Manifest r e , Load r ix e , Manifest (ARepr v) e , VG.Vector v e @@ -184,4 +198,3 @@ (VG.generate (totalElem (size arr)) (unsafeLinearIndex arr)) (castToVector (convert arr :: Array (ARepr v) ix e)) {-# NOINLINE toVector #-} - diff --git a/src/Data/Massiv/Array/Mutable.hs b/src/Data/Massiv/Array/Mutable.hs --- a/src/Data/Massiv/Array/Mutable.hs +++ b/src/Data/Massiv/Array/Mutable.hs @@ -5,6 +5,7 @@ {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} + -- | -- Module : Data.Massiv.Array.Mutable -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -12,100 +13,111 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Mutable - ( -- ** Size - sizeOfMArray - , msize - , resizeMArrayM - , flattenMArray - , outerSliceMArrayM - , outerSlicesMArray - -- ** Element-wise mutation - , read - , readM - , write - , write_ - , writeM - , modify - , modify_ - , modifyM - , modifyM_ - , swap - , swap_ - , swapM - , swapM_ - , zipSwapM_ +module Data.Massiv.Array.Mutable ( + -- ** Size + sizeOfMArray, + msize, + resizeMArrayM, + flattenMArray, + outerSliceMArrayM, + outerSlicesMArray, + + -- ** Element-wise mutation + read, + readM, + write, + write_, + writeM, + modify, + modify_, + modifyM, + modifyM_, + swap, + swap_, + swapM, + swapM_, + zipSwapM_, + -- ** Operations on @MArray@ + -- *** Immutable conversion - , thaw - , thawS - , freeze - , freezeS + thaw, + thawS, + freeze, + freezeS, + -- *** Create mutable - , newMArray - , newMArray' - , makeMArray - , makeMArrayLinear - , makeMArrayS - , makeMArrayLinearS + newMArray, + newMArray', + makeMArray, + makeMArrayLinear, + makeMArrayS, + makeMArrayLinearS, + -- *** Create pure - , createArray_ - , createArray - , createArrayS_ - , createArrayS - , createArrayST_ - , createArrayST + createArray_, + createArray, + createArrayS_, + createArrayS, + createArrayST_, + createArrayST, + -- *** Generate - , generateArray - , generateArrayLinear - , generateArrayS - , generateArrayLinearS - , generateSplitSeedArray + generateArray, + generateArrayLinear, + generateArrayS, + generateArrayLinearS, + generateSplitSeedArray, + -- *** Stateful worker threads - , generateArrayWS - , generateArrayLinearWS + generateArrayWS, + generateArrayLinearWS, + -- *** Unfold - , unfoldrPrimM_ - , iunfoldrPrimM_ - , unfoldrPrimM - , iunfoldrPrimM - , unfoldlPrimM_ - , iunfoldlPrimM_ - , unfoldlPrimM - , iunfoldlPrimM + unfoldrPrimM_, + iunfoldrPrimM_, + unfoldrPrimM, + iunfoldrPrimM, + unfoldlPrimM_, + iunfoldlPrimM_, + unfoldlPrimM, + iunfoldlPrimM, + -- *** Mapping - , forPrimM - , forPrimM_ - , iforPrimM - , iforPrimM_ - , iforLinearPrimM - , iforLinearPrimM_ - , for2PrimM_ - , ifor2PrimM_ + forPrimM, + forPrimM_, + iforPrimM, + iforPrimM_, + iforLinearPrimM, + iforLinearPrimM_, + for2PrimM_, + ifor2PrimM_, + -- *** Modify - , withMArray - , withMArray_ - , withLoadMArray_ - , withMArrayS - , withLoadMArrayS - , withMArrayS_ - , withLoadMArrayS_ - , withMArrayST - , withLoadMArrayST - , withMArrayST_ - , withLoadMArrayST_ + withMArray, + withMArray_, + withLoadMArray_, + withMArrayS, + withLoadMArrayS, + withMArrayS_, + withLoadMArrayS_, + withMArrayST, + withLoadMArrayST, + withMArrayST_, + withLoadMArrayST_, + -- *** Initialize - , initialize - , initializeNew + initialize, + initializeNew, + -- ** Computation - , Manifest - , MArray - , RealWorld - , computeInto - , loadArray - , loadArrayS - ) where + Manifest, + MArray, + RealWorld, + computeInto, + loadArray, + loadArrayS, +) where -- TODO: add fromListM, et al. @@ -118,16 +130,16 @@ import Data.Massiv.Array.Mutable.Internal import Data.Massiv.Core.Common import Data.Maybe (fromMaybe) -import Prelude hiding (mapM, read) import System.IO.Unsafe (unsafePerformIO) +import Prelude hiding (mapM, read) -- | /O(1)/ - Change the size of a mutable array. Throws -- `SizeElementsMismatchException` if total number of elements does not match -- the supplied array. -- -- @since 1.0.0 -resizeMArrayM :: - (Manifest r e, Index ix', Index ix, MonadThrow m) +resizeMArrayM + :: (Manifest r e, Index ix', Index ix, MonadThrow m) => Sz ix' -> MArray s r ix e -> m (MArray s r ix' e) @@ -135,7 +147,6 @@ unsafeResizeMArray sz marr <$ guardNumberOfElements (sizeOfMArray marr) sz {-# INLINE resizeMArrayM #-} - -- | /O(1)/ - Change a mutable array to a mutable vector. -- -- @since 1.0.0 @@ -143,14 +154,14 @@ 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) +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) @@ -215,8 +226,9 @@ -- ] -- -- @since 1.0.0 -outerSlicesMArray :: - forall r ix e s. (Index (Lower ix), Index ix, Manifest r e) +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) @@ -227,7 +239,6 @@ (k, szL) = unconsSz $ sizeOfMArray marr {-# INLINE outerSlicesMArray #-} - -- | /O(n)/ - Initialize a new mutable array. All elements will be set to some default value. For -- boxed arrays it will be a thunk with `Uninitialized` exception, while for others it will be -- simply zeros. @@ -254,14 +265,14 @@ -- *** Exception: Uninitialized -- -- @since 0.6.0 -newMArray' :: - forall r ix e m. (Manifest r e, Index ix, PrimMonad m) +newMArray' + :: forall r ix e m + . (Manifest r e, Index ix, PrimMonad m) => Sz ix -> m (MArray (PrimState m) r ix e) newMArray' sz = unsafeNew sz >>= \ma -> ma <$ initialize ma {-# INLINE newMArray' #-} - -- | /O(n)/ - Make a mutable copy of a pure array. Keep in mind that both `freeze` and `thaw` trigger a -- copy of the full array. -- @@ -293,7 +304,7 @@ let slackLength = totalLength - slackStart when (slackLength > 0) $ scheduleWork_ scheduler $ - unsafeArrayLinearCopy arr slackStart marr slackStart (SafeSz slackLength) + unsafeArrayLinearCopy arr slackStart marr slackStart (SafeSz slackLength) pure marr {-# INLINE thaw #-} @@ -311,8 +322,9 @@ -- [ 1.0, 2.0, 3.0, 4.0, 5.0, 100.0, 7.0, 8.0, 9.0, 10.0 ] -- -- @since 0.3.0 -thawS :: - forall r ix e m. (Manifest r e, Index ix, PrimMonad m) +thawS + :: forall r ix e m + . (Manifest r e, Index ix, PrimMonad m) => Array r ix e -> m (MArray (PrimState m) r ix e) thawS arr = do @@ -321,7 +333,6 @@ pure tmarr {-# INLINE thawS #-} - -- | /O(n)/ - Yield an immutable copy of the mutable array. Note that mutable representations -- have to be the same. -- @@ -337,8 +348,9 @@ -- ] -- -- @since 0.1.0 -freeze :: - forall r ix e m. (Manifest r e, Index ix, MonadIO m) +freeze + :: forall r ix e m + . (Manifest r e, Index ix, MonadIO m) => Comp -> MArray RealWorld r ix e -> m (Array r ix e) @@ -354,17 +366,17 @@ let slackLength = totalLength - slackStart when (slackLength > 0) $ scheduleWork_ scheduler $ - unsafeLinearCopy smarr slackStart tmarr slackStart (SafeSz slackLength) + unsafeLinearCopy smarr slackStart tmarr slackStart (SafeSz slackLength) unsafeFreeze comp tmarr {-# INLINE freeze #-} - -- | Same as `freeze`, but do the copy of supplied muable array sequentially. Also, unlike `freeze` -- that has to be done in `IO`, `freezeS` can be used with `ST`. -- -- @since 0.3.0 -freezeS :: - forall r ix e m. (Manifest r e, Index ix, PrimMonad m) +freezeS + :: forall r ix e m + . (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> m (Array r ix e) freezeS smarr = do @@ -374,16 +386,17 @@ unsafeFreeze Seq tmarr {-# INLINE freezeS #-} -unsafeNewUpper :: - (Load r' ix e, Manifest r e, PrimMonad m) => Array r' ix e -> m (MArray (PrimState m) r Ix1 e) +unsafeNewUpper + :: (Load r' ix e, Manifest r e, PrimMonad m) => Array r' ix e -> m (MArray (PrimState m) r Ix1 e) unsafeNewUpper !arr = unsafeNew (fromMaybe zeroSz (maxLinearSize arr)) {-# INLINE unsafeNewUpper #-} -- | Load sequentially a pure array into the newly created mutable array. -- -- @since 0.3.0 -loadArrayS :: - forall r ix e r' m. (Load r' ix e, Manifest r e, PrimMonad m) +loadArrayS + :: forall r ix e r' m + . (Load r' ix e, Manifest r e, PrimMonad m) => Array r' ix e -> m (MArray (PrimState m) r ix e) loadArrayS arr = do @@ -391,12 +404,12 @@ stToPrim $ unsafeLoadIntoST marr arr {-# INLINE loadArrayS #-} - -- | Load a pure array into the newly created mutable array, while respecting computation startegy. -- -- @since 0.3.0 -loadArray :: - forall r ix e r' m. (Load r' ix e, Manifest r e, MonadIO m) +loadArray + :: forall r ix e r' m + . (Load r' ix e, Manifest r e, MonadIO m) => Array r' ix e -> m (MArray RealWorld r ix e) loadArray arr = @@ -405,44 +418,47 @@ unsafeLoadIntoIO marr arr {-# INLINE loadArray #-} - - -- | Compute an Array while loading the results into the supplied mutable target array. Number of -- elements for arrays must agree, otherwise `SizeElementsMismatchException` exception is thrown. -- -- @since 0.1.3 -computeInto :: - (Load r' ix' e, Manifest r e, Index ix, MonadIO m) - => MArray RealWorld r ix e -- ^ Target Array - -> Array r' ix' e -- ^ Array to load +computeInto + :: (Load r' ix' e, Manifest r e, Index ix, MonadIO m) + => MArray RealWorld r ix e + -- ^ Target Array + -> Array r' ix' e + -- ^ Array to load -> m () computeInto !mArr !arr = liftIO $ do let sz = outerSize arr unless (totalElem (sizeOfMArray mArr) == totalElem sz) $ - throwM $ SizeElementsMismatchException (sizeOfMArray mArr) sz + throwM $ + SizeElementsMismatchException (sizeOfMArray mArr) sz withMassivScheduler_ (getComp arr) $ \scheduler -> stToPrim $ iterArrayLinearST_ scheduler arr (unsafeLinearWrite mArr) {-# INLINE computeInto #-} - -- | Create a mutable array using an index aware generating action. -- -- @since 0.3.0 -makeMArrayS :: - forall r ix e m. (Manifest r e, Index ix, PrimMonad m) - => Sz ix -- ^ Size of the create array - -> (ix -> m e) -- ^ Element generating action +makeMArrayS + :: forall r ix e m + . (Manifest r e, Index ix, PrimMonad m) + => Sz ix + -- ^ Size of the create array + -> (ix -> m e) + -- ^ Element generating action -> m (MArray (PrimState m) r ix e) makeMArrayS sz f = makeMArrayLinearS sz (f . fromLinearIndex sz) {-# INLINE makeMArrayS #-} - -- | Same as `makeMArrayS`, but index supplied to the action is row-major linear index. -- -- @since 0.3.0 -makeMArrayLinearS :: - forall r ix e m. (Manifest r e, Index ix, PrimMonad m) +makeMArrayLinearS + :: forall r ix e m + . (Manifest r e, Index ix, PrimMonad m) => Sz ix -> (Int -> m e) -> m (MArray (PrimState m) r ix e) @@ -455,8 +471,9 @@ -- | Just like `makeMArrayS`, but also accepts computation strategy and runs in `IO`. -- -- @since 0.3.0 -makeMArray :: - forall r ix e m. (MonadUnliftIO m, Manifest r e, Index ix) +makeMArray + :: forall r ix e m + . (MonadUnliftIO m, Manifest r e, Index ix) => Comp -> Sz ix -> (ix -> m e) @@ -464,12 +481,12 @@ makeMArray comp sz f = makeMArrayLinear comp sz (f . fromLinearIndex sz) {-# INLINE makeMArray #-} - -- | Just like `makeMArrayLinearS`, but also accepts computation strategy and runs in `IO`. -- -- @since 0.3.0 -makeMArrayLinear :: - forall r ix e m. (MonadUnliftIO m, Manifest r e, Index ix) +makeMArrayLinear + :: forall r ix e m + . (MonadUnliftIO m, Manifest r e, Index ix) => Comp -> Sz ix -> (Int -> m e) @@ -478,13 +495,10 @@ marr <- liftIO $ unsafeNew sz withScheduler_ comp $ \scheduler -> withRunInIO $ \run -> - splitLinearlyWithM_ scheduler (totalElem sz) (run . f) (unsafeLinearWrite marr) + splitLinearlyWithM_ scheduler (totalElem sz) (run . f) (unsafeLinearWrite marr) return marr {-# INLINE makeMArrayLinear #-} - - - -- | Create a new array by supplying an action that will fill the new blank mutable array. Use -- `createArray` if you'd like to keep the result of the filling function. -- @@ -497,11 +511,13 @@ -- [ 10, 11 ] -- -- @since 0.3.0 --- -createArray_ :: - forall r ix e a m. (Manifest r e, Index ix, MonadUnliftIO m) - => Comp -- ^ Computation strategy to use after `MArray` gets frozen and onward. - -> Sz ix -- ^ Size of the newly created array +createArray_ + :: forall r ix e a m + . (Manifest r e, Index ix, MonadUnliftIO m) + => Comp + -- ^ Computation strategy to use after `MArray` gets frozen and onward. + -> Sz ix + -- ^ Size of the newly created array -> (Scheduler RealWorld () -> MArray RealWorld r ix e -> m a) -- ^ An action that should fill all elements of the brand new mutable array -> m (Array r ix e) @@ -515,11 +531,13 @@ -- actions. -- -- @since 0.3.0 --- -createArray :: - forall r ix e a m b. (Manifest r e, Index ix, MonadUnliftIO m) - => Comp -- ^ Computation strategy to use after `MArray` gets frozen and onward. - -> Sz ix -- ^ Size of the newly created array +createArray + :: forall r ix e a m b + . (Manifest r e, Index ix, MonadUnliftIO m) + => Comp + -- ^ Computation strategy to use after `MArray` gets frozen and onward. + -> Sz ix + -- ^ Size of the newly created array -> (Scheduler RealWorld a -> MArray RealWorld r ix e -> m b) -- ^ An action that should fill all elements of the brand new mutable array -> m ([a], Array r ix e) @@ -530,7 +548,6 @@ return (a, arr) {-# INLINE createArray #-} - -- | Create a new array by supplying an action that will fill the new blank mutable array. Use -- `createArrayS` if you'd like to keep the result of the filling function. -- @@ -543,9 +560,11 @@ -- [ 10, 12 ] -- -- @since 0.3.0 -createArrayS_ :: - forall r ix e a m. (Manifest r e, Index ix, PrimMonad m) - => Sz ix -- ^ Size of the newly created array +createArrayS_ + :: forall r ix e a m + . (Manifest r e, Index ix, PrimMonad m) + => Sz ix + -- ^ Size of the newly created array -> (MArray (PrimState m) r ix e -> m a) -- ^ An action that should fill all elements of the brand new mutable array -> m (Array r ix e) @@ -555,9 +574,11 @@ -- | Just like `createArray_`, but together with `Array` it returns the result of the filling action. -- -- @since 0.3.0 -createArrayS :: - forall r ix e a m. (Manifest r e, Index ix, PrimMonad m) - => Sz ix -- ^ Size of the newly created array +createArrayS + :: forall r ix e a m + . (Manifest r e, Index ix, PrimMonad m) + => Sz ix + -- ^ Size of the newly created array -> (MArray (PrimState m) r ix e -> m a) -- ^ An action that should fill all elements of the brand new mutable array -> m (a, Array r ix e) @@ -571,27 +592,27 @@ -- | Just like `createArrayS_`, but restricted to `ST`. -- -- @since 0.3.0 -createArrayST_ :: - forall r ix e a. (Manifest r e, Index ix) +createArrayST_ + :: forall r ix e a + . (Manifest r e, Index ix) => Sz ix -> (forall s. MArray s r ix e -> ST s a) -> Array r ix e createArrayST_ sz action = runST $ createArrayS_ sz action {-# INLINE createArrayST_ #-} - -- | Just like `createArrayS`, but restricted to `ST`. -- -- @since 0.2.6 -createArrayST :: - forall r ix e a. (Manifest r e, Index ix) +createArrayST + :: forall r ix e a + . (Manifest r e, Index ix) => Sz ix -> (forall s. MArray s r ix e -> ST s a) -> (a, Array r ix e) createArrayST sz action = runST $ createArrayS sz action {-# INLINE createArrayST #-} - -- | Sequentially generate a pure array. Much like `makeArray` creates a pure array this -- function will use `Manifest` interface to generate a pure `Array` in the end, except that -- computation strategy is set to `Seq`. Element producing function no longer has to be pure @@ -616,10 +637,13 @@ -- 15 -- -- @since 0.2.6 -generateArrayS :: - forall r ix e m. (Manifest r e, Index ix, PrimMonad m) - => Sz ix -- ^ Size of the array - -> (ix -> m e) -- ^ Element producing action +generateArrayS + :: forall r ix e m + . (Manifest r e, Index ix, PrimMonad m) + => Sz ix + -- ^ Size of the array + -> (ix -> m e) + -- ^ Element producing action -> m (Array r ix e) generateArrayS sz gen = generateArrayLinearS sz (gen . fromLinearIndex sz) {-# INLINE generateArrayS #-} @@ -627,10 +651,13 @@ -- | Same as `generateArray` but with action that accepts row-major linear index. -- -- @since 0.3.0 -generateArrayLinearS :: - forall r ix e m. (Manifest r e, Index ix, PrimMonad m) - => Sz ix -- ^ Resulting size of the array - -> (Int -> m e) -- ^ Element producing generator +generateArrayLinearS + :: forall r ix e m + . (Manifest r e, Index ix, PrimMonad m) + => Sz ix + -- ^ Resulting size of the array + -> (Int -> m e) + -- ^ Element producing generator -> m (Array r ix e) generateArrayLinearS sz gen = do marr <- unsafeNew sz @@ -638,13 +665,13 @@ unsafeFreeze Seq marr {-# INLINE generateArrayLinearS #-} - -- | Just like `generateArrayS`, except this generator __will__ respect the supplied computation -- strategy, and for that reason it is restricted to `IO`. -- -- @since 0.2.6 -generateArray :: - forall r ix e m. (MonadUnliftIO m, Manifest r e, Index ix) +generateArray + :: forall r ix e m + . (MonadUnliftIO m, Manifest r e, Index ix) => Comp -> Sz ix -> (ix -> m e) @@ -656,8 +683,9 @@ -- index. -- -- @since 0.3.0 -generateArrayLinear :: - forall r ix e m. (MonadUnliftIO m, Manifest r e, Index ix) +generateArrayLinear + :: forall r ix e m + . (MonadUnliftIO m, Manifest r e, Index ix) => Comp -> Sz ix -> (Ix1 -> m e) @@ -665,29 +693,33 @@ generateArrayLinear comp sz f = makeMArrayLinear comp sz f >>= liftIO . unsafeFreeze comp {-# INLINE generateArrayLinear #-} - -- | Similar to `Data.Massiv.Array.makeSplitSeedArray`, except it will produce a -- Manifest array and will return back the last unused seed together with all -- final seeds produced by each scheduled job. This function can be thought of -- as an unfolding done in parallel while iterating in a customizable manner. -- -- @since 1.0.2 -generateSplitSeedArray :: - forall r ix e g it. (Iterator it, Manifest r e, Index ix) - => it -- ^ Iterator - -> g -- ^ Initial seed +generateSplitSeedArray + :: forall r ix e g it + . (Iterator it, Manifest r e, Index ix) + => it + -- ^ Iterator + -> g + -- ^ Initial seed -> (forall s. g -> ST s (g, g)) - -- ^ An ST action that can split a seed into two independent seeds. It will - -- be called the same number of times as the number of jobs that will get - -- scheduled during parallelization. Eg. only once for the sequential case. - -> Comp -- ^ Computation strategy. - -> Sz ix -- ^ Resulting size of the array. + -- ^ An ST action that can split a seed into two independent seeds. It will + -- be called the same number of times as the number of jobs that will get + -- scheduled during parallelization. Eg. only once for the sequential case. + -> Comp + -- ^ Computation strategy. + -> Sz ix + -- ^ Resulting size of the array. -> (forall s. Ix1 -> ix -> g -> ST s (e, g)) - -- ^ An ST action that produces a value and the next seed. It takes both - -- versions of the index, in linear and in multi-dimensional forms, as well - -- as the current seeding value. Returns the element for the array cell - -- together with the new seed that will be used for the next element - -- generation + -- ^ An ST action that produces a value and the next seed. It takes both + -- versions of the index, in linear and in multi-dimensional forms, as well + -- as the current seeding value. Returns the element for the array cell + -- together with the new seed that will be used for the next element + -- generation -> (g, [g], Array r ix e) -- ^ Returned values are: -- @@ -700,30 +732,32 @@ unsafePerformIO $ do marr <- unsafeNew sz ref <- newIORef Nothing - res <- withSchedulerR comp $ \ scheduler -> do + res <- withSchedulerR comp $ \scheduler -> do fin <- stToIO $ iterTargetFullAccST it scheduler 0 sz seed splitSeed $ \ !i ix !g -> - genFunc i ix g >>= \ (x, g') -> g' <$ unsafeLinearWrite marr i x + genFunc i ix g >>= \(x, g') -> g' <$ unsafeLinearWrite marr i x writeIORef ref $ Just fin mFin <- readIORef ref case res of - Finished gs | - Just fin <- mFin -> do - arr <- unsafeFreeze comp marr - pure (fin, gs, arr) + Finished gs + | Just fin <- mFin -> do + arr <- unsafeFreeze comp marr + pure (fin, gs, arr) -- This case does not make much sence for array filling and can only -- happen with a custom 'Iterator' defined outside massiv, therefore it is -- ok to not support it. - _ -> error $ "Parallelized array filling finished prematurely. " ++ - "This feature is not supported by the 'generateSplitSeedArray' function." + _ -> + error $ + "Parallelized array filling finished prematurely. " + ++ "This feature is not supported by the 'generateSplitSeedArray' function." {-# INLINE generateSplitSeedArray #-} - -- | Same as `generateArrayWS`, but use linear indexing instead. -- -- @since 0.3.4 -generateArrayLinearWS :: - forall r ix e s m. (Manifest r e, Index ix, MonadUnliftIO m, PrimMonad m) +generateArrayLinearWS + :: forall r ix e s m + . (Manifest r e, Index ix, MonadUnliftIO m, PrimMonad m) => WorkerStates s -> Sz ix -> (Int -> s -> m e) @@ -743,8 +777,9 @@ -- things that are not thread safe. -- -- @since 0.3.4 -generateArrayWS :: - forall r ix e s m. (Manifest r e, Index ix, MonadUnliftIO m, PrimMonad m) +generateArrayWS + :: forall r ix e s m + . (Manifest r e, Index ix, MonadUnliftIO m, PrimMonad m) => WorkerStates s -> Sz ix -> (ix -> s -> m e) @@ -752,7 +787,6 @@ generateArrayWS states sz make = generateArrayLinearWS states sz (make . fromLinearIndex sz) {-# INLINE generateArrayWS #-} - -- | Sequentially unfold an array from the left. -- -- ====__Examples__ @@ -775,11 +809,15 @@ -- [ 0, 1, 1, 2, 3, 5, 8, 13, 21, 34 ] -- -- @since 0.3.0 -unfoldrPrimM_ :: - forall r ix e a m. (Manifest r e, Index ix, PrimMonad m) - => Sz ix -- ^ Size of the desired array - -> (a -> m (e, a)) -- ^ Unfolding action - -> a -- ^ Initial accumulator +unfoldrPrimM_ + :: forall r ix e a m + . (Manifest r e, Index ix, PrimMonad m) + => Sz ix + -- ^ Size of the desired array + -> (a -> m (e, a)) + -- ^ Unfolding action + -> a + -- ^ Initial accumulator -> m (Array r ix e) unfoldrPrimM_ sz gen acc0 = snd <$> unfoldrPrimM sz gen acc0 {-# INLINE unfoldrPrimM_ #-} @@ -787,24 +825,31 @@ -- | Same as `unfoldrPrimM_` but do the unfolding with index aware function. -- -- @since 0.3.0 -iunfoldrPrimM_ :: - forall r ix e a m. (Manifest r e, Index ix, PrimMonad m) - => Sz ix -- ^ Size of the desired array - -> (a -> ix -> m (e, a)) -- ^ Unfolding action - -> a -- ^ Initial accumulator +iunfoldrPrimM_ + :: forall r ix e a m + . (Manifest r e, Index ix, PrimMonad m) + => Sz ix + -- ^ Size of the desired array + -> (a -> ix -> m (e, a)) + -- ^ Unfolding action + -> a + -- ^ Initial accumulator -> m (Array r ix e) iunfoldrPrimM_ sz gen acc0 = snd <$> iunfoldrPrimM sz gen acc0 {-# INLINE iunfoldrPrimM_ #-} - -- | Just like `iunfoldrPrimM_`, but also returns the final value of the accumulator. -- -- @since 0.3.0 -iunfoldrPrimM :: - forall r ix e a m. (Manifest r e, Index ix, PrimMonad m) - => Sz ix -- ^ Size of the desired array - -> (a -> ix -> m (e, a)) -- ^ Unfolding action - -> a -- ^ Initial accumulator +iunfoldrPrimM + :: forall r ix e a m + . (Manifest r e, Index ix, PrimMonad m) + => Sz ix + -- ^ Size of the desired array + -> (a -> ix -> m (e, a)) + -- ^ Unfolding action + -> a + -- ^ Initial accumulator -> m (a, Array r ix e) iunfoldrPrimM sz gen acc0 = unsafeCreateArrayS sz $ \marr -> @@ -818,11 +863,15 @@ -- | Just like `iunfoldrPrimM`, but do the unfolding with index aware function. -- -- @since 0.3.0 -unfoldrPrimM :: - forall r ix e a m. (Manifest r e, Index ix, PrimMonad m) - => Sz ix -- ^ Size of the desired array - -> (a -> m (e, a)) -- ^ Unfolding action - -> a -- ^ Initial accumulator +unfoldrPrimM + :: forall r ix e a m + . (Manifest r e, Index ix, PrimMonad m) + => Sz ix + -- ^ Size of the desired array + -> (a -> m (e, a)) + -- ^ Unfolding action + -> a + -- ^ Initial accumulator -> m (a, Array r ix e) unfoldrPrimM sz gen acc0 = unsafeCreateArrayS sz $ \marr -> @@ -856,11 +905,15 @@ -- [ 34, 21, 13, 8, 5, 3, 2, 1, 1, 0 ] -- -- @since 0.3.0 -unfoldlPrimM_ :: - forall r ix e a m. (Manifest r e, Index ix, PrimMonad m) - => Sz ix -- ^ Size of the desired array - -> (a -> m (a, e)) -- ^ Unfolding action - -> a -- ^ Initial accumulator +unfoldlPrimM_ + :: forall r ix e a m + . (Manifest r e, Index ix, PrimMonad m) + => Sz ix + -- ^ Size of the desired array + -> (a -> m (a, e)) + -- ^ Unfolding action + -> a + -- ^ Initial accumulator -> m (Array r ix e) unfoldlPrimM_ sz gen acc0 = snd <$> unfoldlPrimM sz gen acc0 {-# INLINE unfoldlPrimM_ #-} @@ -868,24 +921,31 @@ -- | Same as `unfoldlPrimM_` but do the unfolding with index aware function. -- -- @since 0.3.0 -iunfoldlPrimM_ :: - forall r ix e a m. (Manifest r e, Index ix, PrimMonad m) - => Sz ix -- ^ Size of the desired array - -> (a -> ix -> m (a, e)) -- ^ Unfolding action - -> a -- ^ Initial accumulator +iunfoldlPrimM_ + :: forall r ix e a m + . (Manifest r e, Index ix, PrimMonad m) + => Sz ix + -- ^ Size of the desired array + -> (a -> ix -> m (a, e)) + -- ^ Unfolding action + -> a + -- ^ Initial accumulator -> m (Array r ix e) iunfoldlPrimM_ sz gen acc0 = snd <$> iunfoldlPrimM sz gen acc0 {-# INLINE iunfoldlPrimM_ #-} - -- | Just like `iunfoldlPrimM_`, but also returns the final value of the accumulator. -- -- @since 0.3.0 -iunfoldlPrimM :: - forall r ix e a m. (Manifest r e, Index ix, PrimMonad m) - => Sz ix -- ^ Size of the desired array - -> (a -> ix -> m (a, e)) -- ^ Unfolding action - -> a -- ^ Initial accumulator +iunfoldlPrimM + :: forall r ix e a m + . (Manifest r e, Index ix, PrimMonad m) + => Sz ix + -- ^ Size of the desired array + -> (a -> ix -> m (a, e)) + -- ^ Unfolding action + -> a + -- ^ Initial accumulator -> m (a, Array r ix e) iunfoldlPrimM sz gen acc0 = unsafeCreateArrayS sz $ \marr -> @@ -899,16 +959,20 @@ -- | Just like `iunfoldlPrimM`, but do the unfolding with index aware function. -- -- @since 0.3.0 -unfoldlPrimM :: - forall r ix e a m. (Manifest r e, Index ix, PrimMonad m) - => Sz ix -- ^ Size of the desired array - -> (a -> m (a, e)) -- ^ Unfolding action - -> a -- ^ Initial accumulator +unfoldlPrimM + :: forall r ix e a m + . (Manifest r e, Index ix, PrimMonad m) + => Sz ix + -- ^ Size of the desired array + -> (a -> m (a, e)) + -- ^ Unfolding action + -> a + -- ^ Initial accumulator -> m (a, Array r ix e) unfoldlPrimM sz gen acc0 = unsafeCreateArrayS sz $ \marr -> let sz' = sizeOfMArray marr - in loopDeepM 0 (< totalElem sz') (+1) acc0 $ \ !i !acc -> do + in loopDeepM 0 (< totalElem sz') (+ 1) acc0 $ \ !i !acc -> do (acc', e) <- gen acc unsafeLinearWrite marr i e pure acc' @@ -920,7 +984,7 @@ -- @since 0.4.0 forPrimM_ :: (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> (e -> m ()) -> m () forPrimM_ marr f = - loopA_ 0 (< totalElem (sizeOfMArray marr)) (+1) (unsafeLinearRead marr >=> f) + loopA_ 0 (< totalElem (sizeOfMArray marr)) (+ 1) (unsafeLinearRead marr >=> f) {-# INLINE forPrimM_ #-} -- | Sequentially loop over a mutable array while modifying each element with an action. @@ -928,36 +992,34 @@ -- @since 0.4.0 forPrimM :: (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> (e -> m e) -> m () forPrimM marr f = - loopA_ 0 (< totalElem (sizeOfMArray marr)) (+1) (unsafeLinearModify marr f) + loopA_ 0 (< totalElem (sizeOfMArray marr)) (+ 1) (unsafeLinearModify marr f) {-# INLINE forPrimM #-} - -- | Sequentially loop over a mutable array while reading each element and applying an -- index aware action to it. There is no mutation to the array, unless the -- action itself modifies it. -- -- @since 0.4.0 -iforPrimM_ :: - (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> (ix -> e -> m ()) -> m () +iforPrimM_ + :: (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> (ix -> e -> m ()) -> m () iforPrimM_ marr f = iforLinearPrimM_ marr (f . fromLinearIndex (sizeOfMArray marr)) {-# INLINE iforPrimM_ #-} -- | Sequentially loop over a mutable array while modifying each element with an index aware action. -- -- @since 0.4.0 -iforPrimM :: - (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> (ix -> e -> m e) -> m () +iforPrimM + :: (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> (ix -> e -> m e) -> m () iforPrimM marr f = iforLinearPrimM marr (f . fromLinearIndex (sizeOfMArray marr)) {-# INLINE iforPrimM #-} - -- | Sequentially loop over a mutable array while reading each element and applying a -- linear index aware action to it. There is no mutation to the array, unless the action -- itself modifies it. -- -- @since 0.4.0 -iforLinearPrimM_ :: - (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> (Int -> e -> m ()) -> m () +iforLinearPrimM_ + :: (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> (Int -> e -> m ()) -> m () iforLinearPrimM_ marr f = loopA_ 0 (< totalElem (sizeOfMArray marr)) (+ 1) (\i -> unsafeLinearRead marr i >>= f i) {-# INLINE iforLinearPrimM_ #-} @@ -965,21 +1027,20 @@ -- | Sequentially loop over a mutable array while modifying each element with an index aware action. -- -- @since 0.4.0 -iforLinearPrimM :: - (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> (Int -> e -> m e) -> m () +iforLinearPrimM + :: (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> (Int -> e -> m e) -> m () iforLinearPrimM marr f = loopA_ 0 (< totalElem (sizeOfMArray marr)) (+ 1) (\i -> unsafeLinearModify marr (f i) i) {-# INLINE iforLinearPrimM #-} - - -- | Sequentially loop over the intersection of two mutable arrays while reading -- elements from both and applying an action to it. There is no mutation to the -- actual arrays, unless the action itself modifies either one of them. -- -- @since 1.0.0 -for2PrimM_ :: - forall r1 r2 e1 e2 ix m. (PrimMonad m, Index ix, Manifest r1 e1, Manifest r2 e2) +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 ()) @@ -990,8 +1051,9 @@ -- | 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) +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 ()) @@ -1004,12 +1066,11 @@ f ix e1 e2 {-# INLINE ifor2PrimM_ #-} - -- | Same as `withMArray_`, but allows to keep artifacts of scheduled tasks. -- -- @since 0.5.0 -withMArray :: - (Manifest r e, Index ix, MonadUnliftIO m) +withMArray + :: (Manifest r e, Index ix, MonadUnliftIO m) => Array r ix e -> (Scheduler RealWorld a -> MArray RealWorld r ix e -> m b) -> m ([a], Array r ix e) @@ -1032,8 +1093,8 @@ -- * And, of course, the mutable array itself. -- -- @since 0.5.0 -withMArray_ :: - (Manifest r e, Index ix, MonadUnliftIO m) +withMArray_ + :: (Manifest r e, Index ix, MonadUnliftIO m) => Array r ix e -> (Scheduler RealWorld () -> MArray RealWorld r ix e -> m a) -> m (Array r ix e) @@ -1043,13 +1104,13 @@ liftIO $ unsafeFreeze (getComp arr) marr {-# INLINE withMArray_ #-} - -- | Same as `withMArray_`, but the array supplied to this function can be any loadable -- array. For that reason it will be faster if supplied array is delayed. -- -- @since 0.6.1 -withLoadMArray_ :: - forall r ix e r' m b. (Load r' ix e, Manifest r e, MonadUnliftIO m) +withLoadMArray_ + :: forall r ix e r' m b + . (Load r' ix e, Manifest r e, MonadUnliftIO m) => Array r' ix e -> (Scheduler RealWorld () -> MArray RealWorld r ix e -> m b) -> m (Array r ix e) @@ -1057,20 +1118,21 @@ marr <- loadArray arr withScheduler_ (getComp arr) (`action` marr) liftIO $ unsafeFreeze (getComp arr) marr -{-# INLINE[2] withLoadMArray_ #-} +{-# INLINE [2] withLoadMArray_ #-} + {-# RULES "withLoadMArray_/withMArray_" [~2] withLoadMArray_ = withMArray_ "withLoadMArrayS/withMArrayS" [~2] withLoadMArrayS = withMArrayS "withLoadMArrayS_/withMArrayS_" [~2] withLoadMArrayS_ = withMArrayS_ -#-} + #-} -- | Create a copy of a pure array, mutate it in place and return its frozen version. The important -- benefit over doing a manual `thawS` followed by a `freezeS` is that an array will only be copied -- once. -- -- @since 0.5.0 -withMArrayS :: - (Manifest r e, Index ix, PrimMonad m) +withMArrayS + :: (Manifest r e, Index ix, PrimMonad m) => Array r ix e -> (MArray (PrimState m) r ix e -> m a) -> m (a, Array r ix e) @@ -1080,24 +1142,23 @@ (,) a <$> unsafeFreeze (getComp arr) marr {-# INLINE withMArrayS #-} - -- | Same as `withMArrayS`, except it discards the value produced by the supplied action -- -- @since 0.5.0 -withMArrayS_ :: - (Manifest r e, Index ix, PrimMonad m) +withMArrayS_ + :: (Manifest r e, Index ix, PrimMonad m) => Array r ix e -> (MArray (PrimState m) r ix e -> m a) -> m (Array r ix e) withMArrayS_ arr action = snd <$> withMArrayS arr action {-# INLINE withMArrayS_ #-} - -- | Same as `withMArrayS`, but will work with any loadable array. -- -- @since 0.6.1 -withLoadMArrayS :: - forall r ix e r' m a. (Load r' ix e, Manifest r e, PrimMonad m) +withLoadMArrayS + :: forall r ix e r' m a + . (Load r' ix e, Manifest r e, PrimMonad m) => Array r' ix e -> (MArray (PrimState m) r ix e -> m a) -> m (a, Array r ix e) @@ -1105,28 +1166,28 @@ marr <- loadArrayS arr a <- action marr (,) a <$> unsafeFreeze (getComp arr) marr -{-# INLINE[2] withLoadMArrayS #-} +{-# INLINE [2] withLoadMArrayS #-} -- | Same as `withMArrayS_`, but will work with any loadable array. -- -- @since 0.6.1 -withLoadMArrayS_ :: - forall r ix e r' m a. (Load r' ix e, Manifest r e, PrimMonad m) +withLoadMArrayS_ + :: forall r ix e r' m a + . (Load r' ix e, Manifest r e, PrimMonad m) => Array r' ix e -> (MArray (PrimState m) r ix e -> m a) -> m (Array r ix e) withLoadMArrayS_ arr action = snd <$> withLoadMArrayS arr action -{-# INLINE[2] withLoadMArrayS_ #-} - +{-# INLINE [2] withLoadMArrayS_ #-} -- | Same as `withMArrayS` but in `ST`. This is not only pure, but also the safest way to do -- mutation to the array. -- -- @since 0.5.0 -withMArrayST :: - (Manifest r e, Index ix) +withMArrayST + :: (Manifest r e, Index ix) => Array r ix e - -> (forall s . MArray s r ix e -> ST s a) + -> (forall s. MArray s r ix e -> ST s a) -> (a, Array r ix e) withMArrayST arr f = runST $ withMArrayS arr f {-# INLINE withMArrayST #-} @@ -1135,59 +1196,63 @@ -- mutation to the array. -- -- @since 0.5.0 -withMArrayST_ :: - (Manifest r e, Index ix) => Array r ix e -> (forall s. MArray s r ix e -> ST s a) -> Array r ix e +withMArrayST_ + :: (Manifest r e, Index ix) => Array r ix e -> (forall s. MArray s r ix e -> ST s a) -> Array r ix e withMArrayST_ arr f = runST $ withMArrayS_ arr f {-# INLINE withMArrayST_ #-} - -- | Same as `withMArrayST`, but works with any loadable array. -- -- @since 0.6.1 -withLoadMArrayST :: - forall r ix e r' a. (Load r' ix e, Manifest r e) +withLoadMArrayST + :: forall r ix e r' a + . (Load r' ix e, Manifest r e) => Array r' ix e -> (forall s. MArray s r ix e -> ST s a) -> (a, Array r ix e) withLoadMArrayST arr f = runST $ withLoadMArrayS arr f -{-# INLINE[2] withLoadMArrayST #-} +{-# INLINE [2] withLoadMArrayST #-} -- | Same as `withMArrayST_`, but works with any loadable array. -- -- @since 0.6.1 -withLoadMArrayST_ :: - forall r ix e r' a. (Load r' ix e, Manifest r e) +withLoadMArrayST_ + :: forall r ix e r' a + . (Load r' ix e, Manifest r e) => Array r' ix e -> (forall s. MArray s r ix e -> ST s a) -> Array r ix e withLoadMArrayST_ arr f = runST $ withLoadMArrayS_ arr f -{-# INLINE[2] withLoadMArrayST_ #-} - +{-# INLINE [2] withLoadMArrayST_ #-} -- | /O(1)/ - Lookup an element in the mutable array. Returns `Nothing` when index is out of bounds. -- -- @since 0.1.0 -read :: (Manifest r e, Index ix, PrimMonad m) => - MArray (PrimState m) r ix e -> ix -> m (Maybe e) +read + :: (Manifest r e, Index ix, PrimMonad m) + => MArray (PrimState m) r ix e + -> ix + -> m (Maybe e) read marr ix = if isSafeIndex (sizeOfMArray marr) ix then Just <$> unsafeRead marr ix else return Nothing {-# INLINE read #-} - -- | /O(1)/ - Same as `read`, but throws `IndexOutOfBoundsException` on an invalid index. -- -- @since 0.4.0 -readM :: (Manifest r e, Index ix, PrimMonad m, MonadThrow m) => - MArray (PrimState m) r ix e -> ix -> m e +readM + :: (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 + Just e -> pure e Nothing -> throwM $ IndexOutOfBoundsException (sizeOfMArray marr) ix {-# INLINE readM #-} - -- | /O(1)/ - Write an element into the cell of a mutable array. Returns `False` when index is out -- of bounds. -- @@ -1195,8 +1260,8 @@ write :: (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> ix -> e -> m Bool write marr ix e = if isSafeIndex (sizeOfMArray marr) ix - then unsafeWrite marr ix e >> pure True - else pure False + then unsafeWrite marr ix e >> pure True + else pure False {-# INLINE write #-} -- | /O(1)/ - Write an element into the cell of a mutable array. Same as `write` function @@ -1212,22 +1277,24 @@ -- | /O(1)/ - Same as `write`, but throws `IndexOutOfBoundsException` on an invalid index. -- -- @since 0.4.0 -writeM :: - (Manifest r e, Index ix, PrimMonad m, MonadThrow m) => MArray (PrimState m) r ix e -> ix -> e -> m () +writeM + :: (Manifest r e, Index ix, PrimMonad m, MonadThrow m) => MArray (PrimState m) r ix e -> ix -> e -> m () writeM marr ix e = write marr ix e >>= (`unless` throwM (IndexOutOfBoundsException (sizeOfMArray marr) ix)) {-# INLINE writeM #-} - -- | /O(1)/ - Modify an element in the cell of a mutable array with a supplied -- action. Returns the previous value, if index was not out of bounds. -- -- @since 0.1.0 -modify :: - (Manifest r e, Index ix, PrimMonad m) - => MArray (PrimState m) r ix e -- ^ Array to mutate. - -> (e -> m e) -- ^ Monadic action that modifies the element - -> ix -- ^ Index at which to perform modification. +modify + :: (Manifest r e, Index ix, PrimMonad m) + => MArray (PrimState m) r ix e + -- ^ Array to mutate. + -> (e -> m e) + -- ^ Monadic action that modifies the element + -> ix + -- ^ Index at which to perform modification. -> m (Maybe e) modify marr f ix = if isSafeIndex (sizeOfMArray marr) ix @@ -1240,11 +1307,14 @@ -- just like `modifyM_`, but doesn't throw an exception. -- -- @since 0.4.4 -modify_ :: - (Manifest r e, Index ix, PrimMonad m) - => MArray (PrimState m) r ix e -- ^ Array to mutate. - -> (e -> m e) -- ^ Monadic action that modifies the element - -> ix -- ^ Index at which to perform modification. +modify_ + :: (Manifest r e, Index ix, PrimMonad m) + => MArray (PrimState m) r ix e + -- ^ Array to mutate. + -> (e -> m e) + -- ^ Monadic action that modifies the element + -> ix + -- ^ Index at which to perform modification. -> m () modify_ marr f ix = when (isSafeIndex (sizeOfMArray marr) ix) $ void $ unsafeModify marr f ix {-# INLINE modify_ #-} @@ -1254,11 +1324,14 @@ -- the previous value otherwise. -- -- @since 0.4.0 -modifyM :: - (Manifest r e, Index ix, PrimMonad m, MonadThrow m) - => MArray (PrimState m) r ix e -- ^ Array to mutate. - -> (e -> m e) -- ^ Monadic action that modifies the element - -> ix -- ^ Index at which to perform modification. +modifyM + :: (Manifest r e, Index ix, PrimMonad m, MonadThrow m) + => MArray (PrimState m) r ix e + -- ^ Array to mutate. + -> (e -> m e) + -- ^ Monadic action that modifies the element + -> ix + -- ^ Index at which to perform modification. -> m e modifyM marr f ix | isSafeIndex (sizeOfMArray marr) ix = unsafeModify marr f ix @@ -1277,16 +1350,18 @@ -- [ 0, 10, 0 ] -- -- @since 0.4.0 -modifyM_ :: - (Manifest r e, Index ix, PrimMonad m, MonadThrow m) - => MArray (PrimState m) r ix e -- ^ Array to mutate. - -> (e -> m e) -- ^ Monadic action that modifies the element - -> ix -- ^ Index at which to perform modification. +modifyM_ + :: (Manifest r e, Index ix, PrimMonad m, MonadThrow m) + => MArray (PrimState m) r ix e + -- ^ Array to mutate. + -> (e -> m e) + -- ^ Monadic action that modifies the element + -> ix + -- ^ Index at which to perform modification. -> m () modifyM_ marr f ix = void $ modifyM marr f ix {-# INLINE modifyM_ #-} - -- | /O(1)/ - Same as `swapM`, but instead of throwing an exception returns `Nothing` when -- either one of the indices is out of bounds and `Just` elements under those indices -- otherwise. @@ -1300,7 +1375,6 @@ else pure Nothing {-# INLINE swap #-} - -- | /O(1)/ - Same as `swap`, but instead of returning `Nothing` it does nothing. In other -- words, it is similar to `swapM_`, but does not throw any exceptions. -- @@ -1316,13 +1390,15 @@ -- elements under those indices otherwise. -- -- @since 0.4.0 -swapM :: - (Manifest r e, Index ix, PrimMonad m, MonadThrow m) +swapM + :: (Manifest r e, Index ix, PrimMonad m, MonadThrow m) => MArray (PrimState m) r ix e - -> ix -- ^ Index for the first element, which will be returned as the first element in the - -- tuple. - -> ix -- ^ Index for the second element, which will be returned as the second element in - -- the tuple. + -> ix + -- ^ Index for the first element, which will be returned as the first element in the + -- tuple. + -> ix + -- ^ Index for the second element, which will be returned as the second element in + -- the tuple. -> m (e, e) swapM marr ix1 ix2 | not (isSafeIndex sz ix1) = throwM $ IndexOutOfBoundsException (sizeOfMArray marr) ix1 @@ -1332,12 +1408,11 @@ !sz = sizeOfMArray marr {-# INLINE swapM #-} - -- | /O(1)/ - Same as `swapM`, but discard the returned elements -- -- @since 0.4.0 -swapM_ :: - (Manifest r e, Index ix, PrimMonad m, MonadThrow m) +swapM_ + :: (Manifest r e, Index ix, PrimMonad m, MonadThrow m) => MArray (PrimState m) r ix e -> ix -> ix @@ -1349,8 +1424,9 @@ -- 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) +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 diff --git a/src/Data/Massiv/Array/Mutable/Algorithms.hs b/src/Data/Massiv/Array/Mutable/Algorithms.hs --- a/src/Data/Massiv/Array/Mutable/Algorithms.hs +++ b/src/Data/Massiv/Array/Mutable/Algorithms.hs @@ -1,5 +1,6 @@ {-# LANGUAGE ExplicitForAll #-} {-# LANGUAGE FlexibleContexts #-} + -- | -- Module : Data.Massiv.Array.Mutable.Algorithms -- Copyright : (c) Alexey Kuleshevich 2019-2022 @@ -7,19 +8,17 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Mutable.Algorithms - ( quicksortM_ - , quicksortByM_ - , unstablePartitionM - , iterateUntilM - ) where +module Data.Massiv.Array.Mutable.Algorithms ( + quicksortM_, + quicksortByM_, + unstablePartitionM, + iterateUntilM, +) where -import Data.Massiv.Array.Ops.Sort import Data.Massiv.Array.Manifest.Internal (iterateUntilM) +import Data.Massiv.Array.Ops.Sort import Data.Massiv.Core.Common - -- | Partition elements of the supplied mutable vector according to the predicate. -- -- ==== __Example__ @@ -35,9 +34,11 @@ -- [ 2, 1, 8, 10, 20, 50 ] -- -- @since 1.0.0 -unstablePartitionM :: - forall r e m. (Manifest r e, PrimMonad m) +unstablePartitionM + :: forall r e m + . (Manifest r e, PrimMonad m) => MVector (PrimState m) r e - -> (e -> m Bool) -- ^ Predicate + -> (e -> m Bool) + -- ^ Predicate -> m Ix1 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 @@ -3,6 +3,7 @@ {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} + -- | -- Module : Data.Massiv.Array.Mutable.Atomic -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -10,20 +11,19 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Mutable.Atomic - ( -- * Atomic element-wise mutation - atomicReadIntArray - , atomicWriteIntArray - , atomicModifyIntArray - , atomicAddIntArray - , atomicSubIntArray - , atomicAndIntArray - , atomicNandIntArray - , atomicOrIntArray - , atomicXorIntArray - , casIntArray - ) where +module Data.Massiv.Array.Mutable.Atomic ( + -- * Atomic element-wise mutation + atomicReadIntArray, + atomicWriteIntArray, + atomicModifyIntArray, + atomicAddIntArray, + atomicSubIntArray, + atomicAndIntArray, + atomicNandIntArray, + atomicOrIntArray, + atomicXorIntArray, + casIntArray, +) where import Control.Monad.Primitive import Data.Massiv.Array.Manifest.Primitive @@ -34,114 +34,109 @@ -- | Atomically read an `Int` element from the array -- -- @since 0.3.0 -atomicReadIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> m (Maybe Int) +atomicReadIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> m (Maybe Int) atomicReadIntArray marr ix | isSafeIndex (sizeOfMArray marr) ix = Just <$> unsafeAtomicReadIntArray marr ix | otherwise = pure Nothing {-# INLINE atomicReadIntArray #-} - -- | Atomically write an `Int` element int the array. Returns `True` if supplied index was correct -- and write was successfull. -- -- @since 0.3.0 -atomicWriteIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Bool +atomicWriteIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Bool atomicWriteIntArray marr ix f | isSafeIndex (sizeOfMArray marr) ix = unsafeAtomicWriteIntArray marr ix f >> pure True | otherwise = pure False {-# INLINE atomicWriteIntArray #-} - -- | Atomically CAS (Compare-and-Swap) an `Int` in the array. Returns the old value. -- -- @since 0.3.0 -casIntArray :: - (Index ix, PrimMonad m) - => MArray (PrimState m) P ix Int -- ^ Array to mutate - -> ix -- ^ Index at which to mutate - -> Int -- ^ Expected value - -> Int -- ^ New value +casIntArray + :: (Index ix, PrimMonad m) + => MArray (PrimState m) P ix Int + -- ^ Array to mutate + -> ix + -- ^ Index at which to mutate + -> Int + -- ^ Expected value + -> Int + -- ^ New value -> m (Maybe Int) casIntArray marr ix e n | isSafeIndex (sizeOfMArray marr) ix = Just <$> unsafeCasIntArray marr ix e n | otherwise = pure Nothing {-# INLINE casIntArray #-} - -- | Atomically modify an `Int` element of the array. Returns the old value, unless the -- supplied index was out of bounds. -- -- @since 0.3.0 -atomicModifyIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> (Int -> Int) -> m (Maybe Int) +atomicModifyIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> (Int -> Int) -> m (Maybe Int) atomicModifyIntArray marr ix f | isSafeIndex (sizeOfMArray marr) ix = Just <$> unsafeAtomicModifyIntArray marr ix f | otherwise = pure Nothing {-# INLINE atomicModifyIntArray #-} - -- | Atomically add to an `Int` element in the array. Returns the old value. -- -- @since 0.3.0 -atomicAddIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m (Maybe Int) +atomicAddIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m (Maybe Int) atomicAddIntArray marr ix e | isSafeIndex (sizeOfMArray marr) ix = Just <$> unsafeAtomicAddIntArray marr ix e | otherwise = pure Nothing {-# INLINE atomicAddIntArray #-} - -- | Atomically subtract from an `Int` element in the array. Returns the old value. -- -- @since 0.3.0 -atomicSubIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m (Maybe Int) +atomicSubIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m (Maybe Int) atomicSubIntArray marr ix e | isSafeIndex (sizeOfMArray marr) ix = Just <$> unsafeAtomicSubIntArray marr ix e | otherwise = pure Nothing {-# INLINE atomicSubIntArray #-} - -- | Atomically AND an `Int` element in the array. Returns the old value. -- -- @since 0.3.0 -atomicAndIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m (Maybe Int) +atomicAndIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m (Maybe Int) atomicAndIntArray marr ix e | isSafeIndex (sizeOfMArray marr) ix = Just <$> unsafeAtomicAndIntArray marr ix e | otherwise = pure Nothing {-# INLINE atomicAndIntArray #-} - -- | Atomically NAND an `Int` element in the array. Returns the old value. -- -- @since 0.3.0 -atomicNandIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m (Maybe Int) +atomicNandIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m (Maybe Int) atomicNandIntArray marr ix e | isSafeIndex (sizeOfMArray marr) ix = Just <$> unsafeAtomicNandIntArray marr ix e | otherwise = pure Nothing {-# INLINE atomicNandIntArray #-} - -- | Atomically OR an `Int` element in the array. Returns the old value. -- -- @since 0.3.0 -atomicOrIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m (Maybe Int) +atomicOrIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m (Maybe Int) atomicOrIntArray marr ix e | isSafeIndex (sizeOfMArray marr) ix = Just <$> unsafeAtomicOrIntArray marr ix e | otherwise = pure Nothing {-# INLINE atomicOrIntArray #-} - -- | Atomically XOR an `Int` element in the array. Returns the old value. -- -- @since 0.3.0 -atomicXorIntArray :: - (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m (Maybe Int) +atomicXorIntArray + :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m (Maybe Int) atomicXorIntArray marr ix e | isSafeIndex (sizeOfMArray marr) ix = Just <$> unsafeAtomicXorIntArray marr ix e | otherwise = pure Nothing 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 @@ -1,4 +1,5 @@ {-# LANGUAGE ExplicitForAll #-} + -- | -- Module : Data.Massiv.Array.Mutable.Internal -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -6,12 +7,11 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Mutable.Internal - ( unsafeCreateArray - , unsafeCreateArray_ - , unsafeCreateArrayS - ) where +module Data.Massiv.Array.Mutable.Internal ( + unsafeCreateArray, + unsafeCreateArray_, + unsafeCreateArrayS, +) where import Control.Scheduler import Data.Massiv.Core.Common @@ -20,9 +20,11 @@ -- and for unboxed types might contain garbage. -- -- @since 0.5.0 -unsafeCreateArrayS :: - forall r ix e a m. (Manifest r e, Index ix, PrimMonad m) - => Sz ix -- ^ Size of the newly created array +unsafeCreateArrayS + :: forall r ix e a m + . (Manifest r e, Index ix, PrimMonad m) + => Sz ix + -- ^ Size of the newly created array -> (MArray (PrimState m) r ix e -> m a) -- ^ An action that should fill all elements of the brand new mutable array -> m (a, Array r ix e) @@ -37,10 +39,13 @@ -- and for unboxed types might contain garbage. -- -- @since 0.5.0 -unsafeCreateArray :: - forall r ix e a m b. (Manifest r e, Index ix, MonadUnliftIO m) - => Comp -- ^ Computation strategy to use after `MArray` gets frozen and onward. - -> Sz ix -- ^ Size of the newly created array +unsafeCreateArray + :: forall r ix e a m b + . (Manifest r e, Index ix, MonadUnliftIO m) + => Comp + -- ^ Computation strategy to use after `MArray` gets frozen and onward. + -> Sz ix + -- ^ Size of the newly created array -> (Scheduler RealWorld a -> MArray RealWorld r ix e -> m b) -- ^ An action that should fill all elements of the brand new mutable array -> m ([a], Array r ix e) @@ -55,10 +60,13 @@ -- and for unboxed types might contain garbage. -- -- @since 0.5.0 -unsafeCreateArray_ :: - forall r ix e a m b. (Manifest r e, Index ix, MonadUnliftIO m) - => Comp -- ^ Computation strategy to use after `MArray` gets frozen and onward. - -> Sz ix -- ^ Size of the newly created array +unsafeCreateArray_ + :: forall r ix e a m b + . (Manifest r e, Index ix, MonadUnliftIO m) + => Comp + -- ^ Computation strategy to use after `MArray` gets frozen and onward. + -> Sz ix + -- ^ Size of the newly created array -> (Scheduler RealWorld a -> MArray RealWorld r ix e -> m b) -- ^ An action that should fill all elements of the brand new mutable array -> m (Array r ix e) diff --git a/src/Data/Massiv/Array/Numeric.hs b/src/Data/Massiv/Array/Numeric.hs --- a/src/Data/Massiv/Array/Numeric.hs +++ b/src/Data/Massiv/Array/Numeric.hs @@ -3,6 +3,7 @@ {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} + -- | -- Module : Data.Massiv.Array.Numeric -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -10,122 +11,132 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Numeric - ( -- * Numeric - Numeric - , NumericFloat - , liftNumArray2M - -- ** Pointwise addition - , (.+) - , (+.) - , (.+.) - , (!+!) - , sumArraysM - , sumArrays' +module Data.Massiv.Array.Numeric ( + -- * Numeric + Numeric, + NumericFloat, + liftNumArray2M, + + -- ** Pointwise addition + (.+), + (+.), + (.+.), + (!+!), + sumArraysM, + sumArrays', + -- ** Pointwise subtraction - , (.-) - , (-.) - , (.-.) - , (!-!) + (.-), + (-.), + (.-.), + (!-!), + -- ** Pointwise multiplication - , (.*) - , (*.) - , (.*.) - , (!*!) - , (.^) - , productArraysM - , productArrays' + (.*), + (*.), + (.*.), + (!*!), + (.^), + productArraysM, + productArrays', + -- ** Dot product - , (!.!) - , dotM + (!.!), + dotM, + -- ** Matrix multiplication - , (.><) - , (!><) - , multiplyMatrixByVector - , (><.) - , (><.) - , (!><), + (!><), + multiplyMatrixByVector, + (><.), + (><.), + (!> (e -> e -> e) -> Array r ix e -> Array r ix e @@ -139,9 +150,8 @@ !sz2 = size a2 {-# INLINE liftNumArray2M #-} - -applyExactSize2M :: - (Index ix, Size r, 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 @@ -155,7 +165,6 @@ !sz2 = size a2 {-# INLINE applyExactSize2M #-} - -- | Add two arrays together pointwise. Same as `!+!` but produces monadic computation -- that allows for handling failure. -- @@ -204,12 +213,11 @@ -- /__Throws Exception__/: `SizeMismatchException` when array sizes do not match. -- -- @since 0.4.0 -(.-.) :: - (Index ix, Numeric r e, MonadThrow m) => Array r ix e -> Array r ix e -> m (Array r ix e) +(.-.) + :: (Index ix, Numeric r e, MonadThrow m) => Array r ix e -> Array r ix e -> m (Array r ix e) (.-.) = applyExactSize2M subtractionPointwise {-# INLINE (.-.) #-} - -- | Subtract one array from another pointwise. Prefer to use monadic version of this -- function `.-.` whenever possible, because it is better to avoid partial functions. -- @@ -242,19 +250,17 @@ (-.) = scalarMinus {-# INLINE (-.) #-} - -- | Multiply two arrays together pointwise. Same as `!*!` but produces monadic -- computation that allows for handling failure. -- -- /__Throws Exception__/: `SizeMismatchException` when array sizes do not match. -- -- @since 0.4.0 -(.*.) :: - (Index ix, Numeric r e, MonadThrow m) => Array r ix e -> Array r ix e -> m (Array r ix e) +(.*.) + :: (Index ix, Numeric r e, MonadThrow m) => Array r ix e -> Array r ix e -> m (Array r ix e) (.*.) = applyExactSize2M multiplicationPointwise {-# INLINE (.*.) #-} - -- | Multiplication of two arrays pointwise, -- i.e. . -- Prefer to use monadic version of this function `.*.` whenever possible, @@ -275,7 +281,6 @@ (!*!) a1 a2 = throwEither (a1 .*. a2) {-# INLINE (!*!) #-} - -- | Multiply each element of the array by a scalar value. Scalar is on the right. -- -- ====__Example__ @@ -293,7 +298,6 @@ (.*) = multiplyScalar {-# INLINE (.*) #-} - -- | Multiply each element of the array by a scalar value. Scalar is on the left. -- -- ====__Example__ @@ -311,7 +315,6 @@ (*.) = flip multiplyScalar {-# INLINE (*.) #-} - -- | Raise each element of the array to a power. -- -- ====__Example__ @@ -329,7 +332,6 @@ (.^) = powerPointwise {-# INLINE (.^) #-} - -- | Dot product of two vectors. -- -- [Partial] Throws an impure exception when lengths of vectors do not match @@ -353,9 +355,8 @@ comp = getComp v1 <> getComp v2 {-# INLINE dotM #-} - -unsafeDotProductIO :: - (MonadUnliftIO m, Index ix, FoldNumeric r b, Source r b) +unsafeDotProductIO + :: (MonadUnliftIO m, Index ix, FoldNumeric r b, Source r b) => Array r ix b -> Array r ix b -> m b @@ -366,19 +367,19 @@ let n = SafeSz chunkLength loopA_ 0 (< slackStart) (+ chunkLength) $ \ !start -> scheduleWork scheduler $ - pure $! unsafeDotProduct (unsafeLinearSlice start n v1) (unsafeLinearSlice start n v2) + pure $! + unsafeDotProduct (unsafeLinearSlice start n v1) (unsafeLinearSlice start n v2) when (slackStart < totalLength) $ do let k = SafeSz (totalLength - slackStart) scheduleWork scheduler $ pure $! - unsafeDotProduct (unsafeLinearSlice slackStart k v1) (unsafeLinearSlice slackStart k v2) + unsafeDotProduct (unsafeLinearSlice slackStart k v1) (unsafeLinearSlice slackStart k v2) pure $! F.foldl' (+) 0 results where totalLength = totalElem (size v1) comp = getComp v1 <> getComp v2 {-# INLINE unsafeDotProductIO #-} - -- | Compute L2 norm of an array. -- -- @since 0.5.6 @@ -388,8 +389,8 @@ | otherwise = sqrt $! unsafePerformIO $ powerSumArrayIO v 2 {-# INLINE normL2 #-} -powerSumArrayIO :: - (MonadUnliftIO m, Index ix, FoldNumeric r b, Source r b) +powerSumArrayIO + :: (MonadUnliftIO m, Index ix, FoldNumeric r b, Source r b) => Array r ix b -> Int -> m b @@ -408,17 +409,18 @@ totalLength = totalElem (size v) {-# INLINE powerSumArrayIO #-} - -- | Multiply a matrix by a column vector. Same as `!><` but produces monadic -- computation that allows for handling failure. -- -- /__Throws Exception__/: `SizeMismatchException` when inner dimensions of arrays do not match. -- -- @since 0.5.6 -(.><) :: - (MonadThrow m, FoldNumeric r e, Source r e) - => Matrix r e -- ^ Matrix - -> Vector r e -- ^ Column vector (Used many times, so make sure it is computed) +(.><) + :: (MonadThrow m, FoldNumeric r e, Source r e) + => Matrix r e + -- ^ Matrix + -> Vector r e + -- ^ Column vector (Used many times, so make sure it is computed) -> m (Vector D e) (.><) mm v | mCols /= n = throwM $ SizeMismatchException (size mm) (Sz2 n 1) @@ -436,39 +438,44 @@ -- /__Throws Exception__/: `SizeMismatchException` when inner dimensions of arrays do not match. -- -- @since 0.5.7 -multiplyMatrixByVector :: - (MonadThrow m, Numeric r e, Manifest r e) - => Matrix r e -- ^ Matrix - -> Vector r e -- ^ Column vector (Used many times, so make sure it is computed) +multiplyMatrixByVector + :: (MonadThrow m, Numeric r e, Manifest r e) + => Matrix r e + -- ^ Matrix + -> Vector r e + -- ^ Column vector (Used many times, so make sure it is computed) -> m (Vector r e) multiplyMatrixByVector mm v = compute <$> mm .>< v {-# INLINE multiplyMatrixByVector #-} - -- | Multiply a matrix by a column vector -- -- [Partial] Throws impure exception when inner dimensions do not agree -- -- @since 0.5.6 -(!><) :: - (Numeric r e, Source r e) - => Matrix r e -- ^ Matrix - -> Vector r e -- ^ Column vector (Used many times, so make sure it is computed) +(!><) + :: (Numeric r e, Source r e) + => Matrix r e + -- ^ Matrix + -> Vector r e + -- ^ Column vector (Used many times, so make sure it is computed) -> Vector D e (!><) mm v = throwEither (mm .>< v) {-# INLINE (!><) #-} - -- | Multiply a row vector by a matrix. Same as `><.) :: (MonadThrow m, Numeric r e, Manifest r e) => - Vector r e -- ^ Row vector - -> Matrix r e -- ^ Matrix - -> m (Vector r e) +(><.) + :: (MonadThrow m, Numeric r e, Manifest r e) + => Vector r e + -- ^ Row vector + -> Matrix r e + -- ^ Matrix + -> m (Vector r e) (><.) = multiplyVectorByMatrix {-# INLINE (><.) #-} @@ -478,53 +485,54 @@ -- /__Throws Exception__/: `SizeMismatchException` when inner dimensions of arrays do not match. -- -- @since 0.5.7 -multiplyVectorByMatrix :: - (MonadThrow m, Numeric r e, Manifest r e) - => Vector r e -- ^ Row vector - -> Matrix r e -- ^ Matrix +multiplyVectorByMatrix + :: (MonadThrow m, Numeric r e, Manifest r e) + => Vector r e + -- ^ Row vector + -> Matrix r e + -- ^ Matrix -> m (Vector r e) multiplyVectorByMatrix v mm | mRows /= n = throwM $ SizeMismatchException (Sz2 1 n) (size mm) | mRows == 0 || mCols == 0 = pure $ runST (unsafeFreeze comp =<< unsafeNew zeroSz) | otherwise = - pure $! - unsafePerformIO $ do - mv <- newMArray (Sz mCols) 0 - withMassivScheduler_ comp $ \scheduler -> do - let loopCols x ivto = - fix $ \go im iv -> - when (iv < ivto) $ do - _ <- unsafeLinearModify mv (\a -> pure $ a + unsafeLinearIndex mm im * x) iv - go (im + 1) (iv + 1) - loopRows i0 from to = - flip fix i0 $ \go i -> - when (i < mRows) $ do - loopCols (unsafeLinearIndex v i) to (i * mCols + from) from - go (i + 1) - splitLinearlyM_ scheduler mCols (loopRows 0) - unsafeFreeze comp mv + pure $! + unsafePerformIO $ do + mv <- newMArray (Sz mCols) 0 + withMassivScheduler_ comp $ \scheduler -> do + let loopCols x ivto = + fix $ \go im iv -> + when (iv < ivto) $ do + _ <- unsafeLinearModify mv (\a -> pure $ a + unsafeLinearIndex mm im * x) iv + go (im + 1) (iv + 1) + loopRows i0 from to = + flip fix i0 $ \go i -> + when (i < mRows) $ do + loopCols (unsafeLinearIndex v i) to (i * mCols + from) from + go (i + 1) + splitLinearlyM_ scheduler mCols (loopRows 0) + unsafeFreeze comp mv where comp = getComp mm <> getComp v Sz2 mRows mCols = size mm Sz1 n = size v {-# INLINE multiplyVectorByMatrix #-} - -- | Multiply a row vector by a matrix. -- -- [Partial] Throws impure exception when inner dimensions do not agree -- -- @since 0.5.6 -(> Vector r e -- ^ Row vector (Used many times, so make sure it is computed) - -> Matrix r e -- ^ Matrix +(> Vector r e + -- ^ Row vector (Used many times, so make sure it is computed) + -> Matrix r e + -- ^ Matrix -> Vector r e (><. mm) {-# INLINE (><.) = multiplyMatrices {-# INLINE (.><.) #-} - -- | Synonym for `.><.` -- -- @since 0.5.6 -multiplyMatrices :: - (Numeric r e, Manifest r e, MonadThrow m) => Matrix r e -> Matrix r e -> m (Matrix r e) +multiplyMatrices + :: (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 + -- mA == 1 = -- TODO: call multiplyVectorByMatrix + -- nA == 1 = -- TODO: call multiplyMatrixByVector | nA /= mB = throwM $ SizeMismatchException (size arrA) (size arrB) | isEmpty arrA || isEmpty arrB = pure $ runST (unsafeFreeze comp =<< unsafeNew zeroSz) | otherwise = pure $! unsafePerformIO $ do - marrC <- newMArray (SafeSz (mA :. nB)) 0 - withScheduler_ comp $ \scheduler -> do - let withC00 iA jB f = let !ixC00 = iA * nB + jB - in f ixC00 =<< unsafeLinearRead marrC ixC00 - withC01 ixC00 f = let !ixC01 = ixC00 + 1 - in f ixC01 =<< unsafeLinearRead marrC ixC01 - withC10 ixC00 f = let !ixC10 = ixC00 + nB - in f ixC10 =<< unsafeLinearRead marrC ixC10 - withC11 ixC01 f = let !ixC11 = ixC01 + nB - in f ixC11 =<< unsafeLinearRead marrC ixC11 - withB00 iB jB f = let !ixB00 = iB * nB + jB - in f ixB00 $! unsafeLinearIndex arrB ixB00 - withB00B10 iB jB f = - withB00 iB jB $ \ixB00 b00 -> let !ixB10 = ixB00 + nB - in f ixB00 b00 ixB10 $! unsafeLinearIndex arrB ixB10 - withA00 iA jA f = let !ixA00 = iA * nA + jA - in f ixA00 $! unsafeLinearIndex arrA ixA00 - withA00A10 iA jA f = - withA00 iA jA $ \ixA00 a00 -> let !ixA10 = ixA00 + nA - in f ixA00 a00 ixA10 $! unsafeLinearIndex arrA ixA10 - let loopColsB_UnRowBColA_UnRowA a00 a01 a10 a11 iA iB jB - | jB < n2B = do - withB00B10 iB jB $ \ixB00 b00 ixB10 b10 -> do - let !b01 = unsafeLinearIndex arrB (ixB00 + 1) - !b11 = unsafeLinearIndex arrB (ixB10 + 1) - withC00 iA jB $ \ixC00 c00 -> do - unsafeLinearWrite marrC ixC00 (c00 + a00 * b00 + a01 * b10) - withC01 ixC00 $ \ixC01 c01 -> do - unsafeLinearWrite marrC ixC01 (c01 + a00 * b01 + a01 * b11) + marrC <- newMArray (SafeSz (mA :. nB)) 0 + withScheduler_ comp $ \scheduler -> do + let withC00 iA jB f = + let !ixC00 = iA * nB + jB + in f ixC00 =<< unsafeLinearRead marrC ixC00 + withC01 ixC00 f = + let !ixC01 = ixC00 + 1 + in f ixC01 =<< unsafeLinearRead marrC ixC01 + withC10 ixC00 f = + let !ixC10 = ixC00 + nB + in f ixC10 =<< unsafeLinearRead marrC ixC10 + withC11 ixC01 f = + let !ixC11 = ixC01 + nB + in f ixC11 =<< unsafeLinearRead marrC ixC11 + withB00 iB jB f = + let !ixB00 = iB * nB + jB + in f ixB00 $! unsafeLinearIndex arrB ixB00 + withB00B10 iB jB f = + withB00 iB jB $ \ixB00 b00 -> + let !ixB10 = ixB00 + nB + in f ixB00 b00 ixB10 $! unsafeLinearIndex arrB ixB10 + withA00 iA jA f = + let !ixA00 = iA * nA + jA + in f ixA00 $! unsafeLinearIndex arrA ixA00 + withA00A10 iA jA f = + withA00 iA jA $ \ixA00 a00 -> + let !ixA10 = ixA00 + nA + in f ixA00 a00 ixA10 $! unsafeLinearIndex arrA ixA10 + let loopColsB_UnRowBColA_UnRowA a00 a01 a10 a11 iA iB jB + | jB < n2B = do + withB00B10 iB jB $ \ixB00 b00 ixB10 b10 -> do + let !b01 = unsafeLinearIndex arrB (ixB00 + 1) + !b11 = unsafeLinearIndex arrB (ixB10 + 1) + withC00 iA jB $ \ixC00 c00 -> do + unsafeLinearWrite marrC ixC00 (c00 + a00 * b00 + a01 * b10) + withC01 ixC00 $ \ixC01 c01 -> do + unsafeLinearWrite marrC ixC01 (c01 + a00 * b01 + a01 * b11) + withC10 ixC00 $ \ixC10 c10 -> + unsafeLinearWrite marrC ixC10 (c10 + a10 * b00 + a11 * b10) + withC11 ixC01 $ \ixC11 c11 -> + unsafeLinearWrite marrC ixC11 (c11 + a10 * b01 + a11 * b11) + loopColsB_UnRowBColA_UnRowA a00 a01 a10 a11 iA iB (jB + 2) + | jB < nB = withB00B10 iB jB $ \_ b00 _ b10 -> + withC00 iA jB $ \ixC00 c00 -> do + unsafeLinearWrite marrC ixC00 (c00 + a00 * b00 + a01 * b10) withC10 ixC00 $ \ixC10 c10 -> unsafeLinearWrite marrC ixC10 (c10 + a10 * b00 + a11 * b10) - withC11 ixC01 $ \ixC11 c11 -> - unsafeLinearWrite marrC ixC11 (c11 + a10 * b01 + a11 * b11) - loopColsB_UnRowBColA_UnRowA a00 a01 a10 a11 iA iB (jB + 2) - - | jB < nB = withB00B10 iB jB $ \_ b00 _ b10 -> - withC00 iA jB $ \ixC00 c00 -> do - unsafeLinearWrite marrC ixC00 (c00 + a00 * b00 + a01 * b10) - withC10 ixC00 $ \ixC10 c10 -> - unsafeLinearWrite marrC ixC10 (c10 + a10 * b00 + a11 * b10) - | otherwise = pure () - - loopColsB_UnRowBColA_RowA a00 a01 iA iB jB - | jB < n2B = do - withB00B10 iB jB $ \ixB00 b00 ixB10 b10 -> do - let !b01 = unsafeLinearIndex arrB (ixB00 + 1) - !b11 = unsafeLinearIndex arrB (ixB10 + 1) - withC00 iA jB $ \ixC00 c00 -> do - unsafeLinearWrite marrC ixC00 (c00 + a00 * b00 + a01 * b10) - withC01 ixC00 $ \ixC01 c01 -> - unsafeLinearWrite marrC ixC01 (c01 + a00 * b01 + a01 * b11) - loopColsB_UnRowBColA_RowA a00 a01 iA iB (jB + 2) + | otherwise = pure () - | jB < nB = withB00B10 iB jB $ \_ b00 _ b10 -> - withC00 iA jB $ \ixC00 c00 -> - unsafeLinearWrite marrC ixC00 (c00 + a00 * b00 + a01 * b10) - | otherwise = pure () + loopColsB_UnRowBColA_RowA a00 a01 iA iB jB + | jB < n2B = do + withB00B10 iB jB $ \ixB00 b00 ixB10 b10 -> do + let !b01 = unsafeLinearIndex arrB (ixB00 + 1) + !b11 = unsafeLinearIndex arrB (ixB10 + 1) + withC00 iA jB $ \ixC00 c00 -> do + unsafeLinearWrite marrC ixC00 (c00 + a00 * b00 + a01 * b10) + withC01 ixC00 $ \ixC01 c01 -> + unsafeLinearWrite marrC ixC01 (c01 + a00 * b01 + a01 * b11) + loopColsB_UnRowBColA_RowA a00 a01 iA iB (jB + 2) + | jB < nB = withB00B10 iB jB $ \_ b00 _ b10 -> + withC00 iA jB $ \ixC00 c00 -> + unsafeLinearWrite marrC ixC00 (c00 + a00 * b00 + a01 * b10) + | otherwise = pure () - loopColsB_RowBColA_UnRowA a00 a10 iA iB jB - | jB < n2B = do - withB00 iB jB $ \ixB00 b00 -> do - let !b01 = unsafeLinearIndex arrB (ixB00 + 1) - withC00 iA jB $ \ixC00 c00 -> do - unsafeLinearWrite marrC ixC00 (c00 + a00 * b00) - withC01 ixC00 $ \ixC01 c01 -> do - unsafeLinearWrite marrC ixC01 (c01 + a00 * b01) + loopColsB_RowBColA_UnRowA a00 a10 iA iB jB + | jB < n2B = do + withB00 iB jB $ \ixB00 b00 -> do + let !b01 = unsafeLinearIndex arrB (ixB00 + 1) + withC00 iA jB $ \ixC00 c00 -> do + unsafeLinearWrite marrC ixC00 (c00 + a00 * b00) + withC01 ixC00 $ \ixC01 c01 -> do + unsafeLinearWrite marrC ixC01 (c01 + a00 * b01) + withC10 ixC00 $ \ixC10 c10 -> + unsafeLinearWrite marrC ixC10 (c10 + a10 * b00) + withC11 ixC01 $ \ixC11 c11 -> + unsafeLinearWrite marrC ixC11 (c11 + a10 * b01) + loopColsB_RowBColA_UnRowA a00 a10 iA iB (jB + 2) + | jB < nB = withB00 iB jB $ \_ b00 -> + withC00 iA jB $ \ixC00 c00 -> do + unsafeLinearWrite marrC ixC00 (c00 + a00 * b00) withC10 ixC00 $ \ixC10 c10 -> unsafeLinearWrite marrC ixC10 (c10 + a10 * b00) - withC11 ixC01 $ \ixC11 c11 -> - unsafeLinearWrite marrC ixC11 (c11 + a10 * b01) - loopColsB_RowBColA_UnRowA a00 a10 iA iB (jB + 2) - - | jB < nB = withB00 iB jB $ \_ b00 -> - withC00 iA jB $ \ixC00 c00 -> do - unsafeLinearWrite marrC ixC00 (c00 + a00 * b00) - withC10 ixC00 $ \ixC10 c10 -> - unsafeLinearWrite marrC ixC10 (c10 + a10 * b00) - | otherwise = pure () - - loopColsB_RowBColA_RowA a00 iA iB jB - | jB < n2B = do - withB00 iB jB $ \ixB00 b00 -> do - let !b01 = unsafeLinearIndex arrB (ixB00 + 1) - withC00 iA jB $ \ixC00 c00 -> do - unsafeLinearWrite marrC ixC00 (c00 + a00 * b00) - withC01 ixC00 $ \ixC01 c01 -> do - unsafeLinearWrite marrC ixC01 (c01 + a00 * b01) - loopColsB_RowBColA_RowA a00 iA iB (jB + 2) - | jB < nB = withB00 iB jB $ \_ b00 -> - withC00 iA jB $ \ixC00 c00 -> - unsafeLinearWrite marrC ixC00 (c00 + a00 * b00) + | otherwise = pure () - | otherwise = pure () + loopColsB_RowBColA_RowA a00 iA iB jB + | jB < n2B = do + withB00 iB jB $ \ixB00 b00 -> do + let !b01 = unsafeLinearIndex arrB (ixB00 + 1) + withC00 iA jB $ \ixC00 c00 -> do + unsafeLinearWrite marrC ixC00 (c00 + a00 * b00) + withC01 ixC00 $ \ixC01 c01 -> do + unsafeLinearWrite marrC ixC01 (c01 + a00 * b01) + loopColsB_RowBColA_RowA a00 iA iB (jB + 2) + | jB < nB = withB00 iB jB $ \_ b00 -> + withC00 iA jB $ \ixC00 c00 -> + unsafeLinearWrite marrC ixC00 (c00 + a00 * b00) + | otherwise = pure () - loopRowsB_UnRowA iA iB - | iB < m2B = do - withA00A10 iA iB $ \ixA00 a00 ixA10 a10 -> do - let !a01 = unsafeLinearIndex arrA (ixA00 + 1) - !a11 = unsafeLinearIndex arrA (ixA10 + 1) - loopColsB_UnRowBColA_UnRowA a00 a01 a10 a11 iA iB 0 - loopRowsB_UnRowA iA (iB + 2) - | iB < mB = - withA00A10 iA iB $ \_ a00 _ a10 -> loopColsB_RowBColA_UnRowA a00 a10 iA iB 0 - | otherwise = pure () + loopRowsB_UnRowA iA iB + | iB < m2B = do + withA00A10 iA iB $ \ixA00 a00 ixA10 a10 -> do + let !a01 = unsafeLinearIndex arrA (ixA00 + 1) + !a11 = unsafeLinearIndex arrA (ixA10 + 1) + loopColsB_UnRowBColA_UnRowA a00 a01 a10 a11 iA iB 0 + loopRowsB_UnRowA iA (iB + 2) + | iB < mB = + withA00A10 iA iB $ \_ a00 _ a10 -> loopColsB_RowBColA_UnRowA a00 a10 iA iB 0 + | otherwise = pure () - loopRowsB_RowA iA iB - | iB < m2B = do - withA00 iA iB $ \ixA00 a00 -> do - let !a01 = unsafeLinearIndex arrA (ixA00 + 1) - loopColsB_UnRowBColA_RowA a00 a01 iA iB 0 - loopRowsB_RowA iA (iB + 2) - | iB < mB = withA00 iA iB $ \_ a00 -> loopColsB_RowBColA_RowA a00 iA iB 0 - | otherwise = pure () + loopRowsB_RowA iA iB + | iB < m2B = do + withA00 iA iB $ \ixA00 a00 -> do + let !a01 = unsafeLinearIndex arrA (ixA00 + 1) + loopColsB_UnRowBColA_RowA a00 a01 iA iB 0 + loopRowsB_RowA iA (iB + 2) + | iB < mB = withA00 iA iB $ \_ a00 -> loopColsB_RowBColA_RowA a00 iA iB 0 + | otherwise = pure () - loopRowsA iA - | iA < m2A = do - scheduleWork_ scheduler $ loopRowsB_UnRowA iA 0 - loopRowsA (iA + 2) - | iA < mA = scheduleWork_ scheduler $ loopRowsB_RowA iA 0 - | otherwise = pure () - loopRowsA 0 + loopRowsA iA + | iA < m2A = do + scheduleWork_ scheduler $ loopRowsB_UnRowA iA 0 + loopRowsA (iA + 2) + | iA < mA = scheduleWork_ scheduler $ loopRowsB_RowA iA 0 + | otherwise = pure () + loopRowsA 0 - unsafeFreeze comp marrC + unsafeFreeze comp marrC where comp = getComp arrA <> getComp arrB m2A = mA - mA `rem` 2 @@ -708,8 +719,8 @@ -- > m1 .><. transpose m2 == multiplyMatricesTransposed m1 m2 -- -- @since 0.5.6 -multiplyMatricesTransposed :: - (Numeric r e, Manifest r e, MonadThrow m) +multiplyMatricesTransposed + :: (Numeric r e, Manifest r e, MonadThrow m) => Matrix r e -> Matrix r e -> m (Matrix D e) @@ -717,9 +728,9 @@ | n1 /= m2 = throwM $ SizeMismatchException (size arr1) (Sz2 m2 n2) | isEmpty arr1 || isEmpty arr2 = pure $ setComp comp empty | otherwise = - pure $ - makeArray comp (SafeSz (m1 :. n2)) $ \(i :. j) -> - unsafeDotProduct (unsafeLinearSlice (i * n1) n arr1) (unsafeLinearSlice (j * n1) n arr2) + pure $ + makeArray comp (SafeSz (m1 :. n2)) $ \(i :. j) -> + unsafeDotProduct (unsafeLinearSlice (i * n1) n arr1) (unsafeLinearSlice (j * n1) n arr2) where comp = getComp arr1 <> getComp arr2 n = SafeSz n1 @@ -744,7 +755,7 @@ -- @since 0.3.6 identityMatrix :: Num e => Sz1 -> Matrix DL e identityMatrix (Sz n) = - makeLoadArrayS (Sz2 n n) 0 $ \ w -> loopA_ 0 (< n) (+1) $ \ i -> w (i :. i) 1 + makeLoadArrayS (Sz2 n n) 0 $ \w -> loopA_ 0 (< n) (+ 1) $ \i -> w (i :. i) 1 {-# INLINE identityMatrix #-} -- | Create a lower triangular (L in LU decomposition) matrix of size @NxN@ @@ -828,15 +839,14 @@ -- /__Throws Exception__/: `SizeMismatchException` when array sizes do not match. -- -- @since 0.4.0 -(./.) :: - (Index ix, 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) (./.) = applyExactSize2M divisionPointwise {-# INLINE (./.) #-} - -- | Divide two arrays pointwise. Prefer to use monadic version of this function `./.` -- whenever possible, because it is better to avoid partial functions. -- @@ -870,7 +880,7 @@ -- [ 5.0, 4.7619047, 4.5454545, 4.347826, 4.1666665 ] -- -- @since 0.5.6 -(/.) ::(Index ix, NumericFloat r e) => e -> Array r ix e -> Array r ix e +(/.) :: (Index ix, NumericFloat r e) => e -> Array r ix e -> Array r ix e (/.) = scalarDivide {-# INLINE (/.) #-} @@ -887,13 +897,15 @@ -- [ 0.2, 0.21, 0.22, 0.23, 0.24 ] -- -- @since 0.4.0 -(./) ::(Index ix, NumericFloat r e) => Array r ix e -> e -> Array r ix e +(./) :: (Index ix, NumericFloat r e) => Array r ix e -> e -> Array r ix e (./) = divideScalar {-# INLINE (./) #-} (.^^) :: (Index ix, Numeric r e, Fractional e, Integral b) - => Array r ix e -> b -> Array r ix e + => Array r ix e + -> b + -> Array r ix e (.^^) arr n = unsafeLiftArray (^^ n) arr {-# INLINE (.^^) #-} @@ -906,7 +918,6 @@ recipA = recipPointwise {-# INLINE recipA #-} - -- | Apply exponent to each element of the array. -- -- > expA arr == map exp arr @@ -934,7 +945,6 @@ logA = unsafeLiftArray log {-# INLINE logA #-} - -- | Apply logarithm to each element of the array where the base is in the same cell in -- the second array. -- @@ -945,16 +955,16 @@ -- @since 0.4.0 logBaseA :: (Index ix, Source r1 e, Source r2 e, Floating e) - => Array r1 ix e -> Array r2 ix e -> Array D ix e + => Array r1 ix e + -> Array r2 ix e + -> Array D ix e logBaseA = liftArray2' logBase {-# INLINE logBaseA #-} + -- TODO: siwtch to -- (breaking) logBaseA :: Array r ix e -> e -> Array D ix e -- logBasesM :: Array r ix e -> Array r ix e -> m (Array D ix e) - - - -- | Apply power to each element of the array where the power value is in the same cell -- in the second array. -- @@ -965,16 +975,17 @@ -- @since 0.4.0 (.**) :: (Index ix, Source r1 e, Source r2 e, Floating e) - => Array r1 ix e -> Array r2 ix e -> Array D ix e + => Array r1 ix e + -> Array r2 ix e + -> Array D ix e (.**) = liftArray2' (**) {-# INLINE (.**) #-} + -- TODO: -- !**! :: Array r1 ix e -> Array r2 ix e -> Array D ix e -- .**. :: Array r1 ix e -> Array r2 ix e -> m (Array D ix e) -- (breaking) .** :: Array r1 ix e -> e -> Array D ix e - - -- | Apply sine function to each element of the array. -- -- > sinA arr == map sin arr @@ -1083,7 +1094,6 @@ atanhA = unsafeLiftArray atanh {-# INLINE atanhA #-} - -- | Perform a pointwise quotient where first array contains numerators and the second -- one denominators -- @@ -1094,11 +1104,12 @@ -- @since 0.1.0 quotA :: (HasCallStack, Index ix, Source r1 e, Source r2 e, Integral e) - => Array r1 ix e -> Array r2 ix e -> Array D ix e + => Array r1 ix e + -> Array r2 ix e + -> Array D ix e quotA = liftArray2' quot {-# INLINE quotA #-} - -- | Perform a pointwise remainder computation -- -- > remA arr1 arr2 == zipWith rem arr1 arr2 @@ -1108,7 +1119,9 @@ -- @since 0.1.0 remA :: (HasCallStack, Index ix, Source r1 e, Source r2 e, Integral e) - => Array r1 ix e -> Array r2 ix e -> Array D ix e + => Array r1 ix e + -> Array r2 ix e + -> Array D ix e remA = liftArray2' rem {-# INLINE remA #-} @@ -1122,9 +1135,12 @@ -- @since 0.1.0 divA :: (HasCallStack, Index ix, Source r1 e, Source r2 e, Integral e) - => Array r1 ix e -> Array r2 ix e -> Array D ix e + => Array r1 ix e + -> Array r2 ix e + -> Array D ix e divA = liftArray2' div {-# INLINE divA #-} + -- TODO: -- * Array r ix e -> Array r ix e -> m (Array r ix e) -- * Array r ix e -> e -> Array r ix e @@ -1139,12 +1155,12 @@ -- @since 0.1.0 modA :: (HasCallStack, Index ix, Source r1 e, Source r2 e, Integral e) - => Array r1 ix e -> Array r2 ix e -> Array D ix e + => Array r1 ix e + -> Array r2 ix e + -> Array D ix e modA = liftArray2' mod {-# INLINE modA #-} - - -- | Perform a pointwise quotient with remainder where first array contains numerators -- and the second one denominators -- @@ -1155,11 +1171,12 @@ -- @since 0.1.0 quotRemA :: (HasCallStack, Index ix, Source r1 e, Source r2 e, Integral e) - => Array r1 ix e -> Array r2 ix e -> (Array D ix e, Array D ix e) + => Array r1 ix e + -> Array r2 ix e + -> (Array D ix e, Array D ix e) quotRemA arr1 = A.unzip . liftArray2' quotRem arr1 {-# INLINE quotRemA #-} - -- | Perform a pointwise integer division with modulo where first array contains -- numerators and the second one denominators -- @@ -1170,12 +1187,12 @@ -- @since 0.1.0 divModA :: (HasCallStack, Index ix, Source r1 e, Source r2 e, Integral e) - => Array r1 ix e -> Array r2 ix e -> (Array D ix e, Array D ix e) + => Array r1 ix e + -> Array r2 ix e + -> (Array D ix e, Array D ix e) divModA arr1 = A.unzip . liftArray2' divMod arr1 {-# INLINE divModA #-} - - -- | Truncate each element of the array. -- -- > truncateA arr == map truncate arr @@ -1185,7 +1202,6 @@ truncateA = A.map truncate {-# INLINE truncateA #-} - -- | Round each element of the array. -- -- > truncateA arr == map truncate arr @@ -1195,7 +1211,6 @@ roundA = A.map round {-# INLINE roundA #-} - -- | Ceiling of each element of the array. -- -- > truncateA arr == map truncate arr @@ -1205,7 +1220,6 @@ ceilingA = A.map ceiling {-# INLINE ceilingA #-} - -- | Floor each element of the array. -- -- > truncateA arr == map truncate arr @@ -1222,8 +1236,8 @@ -- /__Throws Exception__/: `SizeMismatchException` when array sizes do not match. -- -- @since 0.1.0 -atan2A :: - (Index ix, Numeric r e, RealFloat e, MonadThrow m) +atan2A + :: (Index ix, Numeric r e, RealFloat e, MonadThrow m) => Array r ix e -> Array r ix e -> m (Array r ix e) @@ -1291,13 +1305,14 @@ -- ] -- -- @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 + :: (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 + (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 @@ -1305,12 +1320,11 @@ -- 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' + :: (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__ @@ -1364,10 +1378,10 @@ -- ] -- -- @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 + :: (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 + (x : xs) -> F.foldlM (.*.) x xs {-# INLINE productArraysM #-} diff --git a/src/Data/Massiv/Array/Numeric/Integral.hs b/src/Data/Massiv/Array/Numeric/Integral.hs --- a/src/Data/Massiv/Array/Numeric/Integral.hs +++ b/src/Data/Massiv/Array/Numeric/Integral.hs @@ -1,5 +1,6 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE FlexibleContexts #-} + -- | -- Module : Data.Massiv.Array.Numeric.Integral -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -7,30 +8,31 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Numeric.Integral - ( +module Data.Massiv.Array.Numeric.Integral ( -- $integral_intro - -- - -- * Integral Approximation - -- ** Midpoint Rule - midpointRule - , midpointStencil + midpointRule, + midpointStencil, + -- ** Trapezoid Rule - , trapezoidRule - , trapezoidStencil + trapezoidRule, + trapezoidStencil, + -- ** Simpson's Rule - , simpsonsRule - , simpsonsStencil + simpsonsRule, + simpsonsStencil, + -- * General Integral approximation - , integrateWith - , integralApprox + integrateWith, + integralApprox, + -- * From functions + -- ** Sampled at the edge - , fromFunction + fromFunction, + -- ** Sampled at the midpoint - , fromFunctionMidpoint - ) where + fromFunctionMidpoint, +) where import Data.Coerce import Data.Massiv.Array.Delayed.Pull (D) @@ -42,7 +44,6 @@ import Data.Massiv.Array.Unsafe import Data.Massiv.Core.Common - -- | -- -- __Midpoint Rule__ @@ -50,18 +51,20 @@ -- \[ -- \int_{{\,a}}^{{\,b}}{{f\left( x \right)\,dx}} \approx \Delta x \cdot \,f\left( {x_1 + \frac{\Delta x}{2}} \right) + \Delta x \cdot \,f\left( {x_2 + \frac{\Delta x}{2}} \right) + \cdots + \Delta x \cdot \,f\left( {x_n + \frac{\Delta x}{2}} \right) -- \] -midpointStencil :: - (Fractional e, Index ix) - => e -- ^ @Δx@ - distance between sample points - -> Dim -- ^ Dimension along which to integrate - -> Int -- ^ @n@ - number of sample points. +midpointStencil + :: (Fractional e, Index ix) + => e + -- ^ @Δx@ - distance between sample points + -> Dim + -- ^ Dimension along which to integrate + -> Int + -- ^ @n@ - number of sample points. -> Stencil ix e e midpointStencil dx dim k = makeUnsafeStencil (Sz (setDim' (pureIndex 1) dim k)) zeroIndex $ \_ g -> dx * loop 0 (< k) (+ 1) 0 (\i -> (+ g (setDim' zeroIndex dim i))) {-# INLINE midpointStencil #-} - -- | -- -- __Trapezoid Rule__ @@ -69,20 +72,23 @@ -- \[ -- \int_{{\,a}}^{{\,b}}{{f\left( x \right)\,dx}} \approx \frac{{\Delta x}}{2}\cdot\left( {f\left( {{x_0}} \right) + f\left( {{x_1}} \right)} \right) + \frac{{\Delta x}}{2}\cdot\left( {f\left( {{x_1}} \right) + f\left( {{x_2}} \right)} \right) + \cdots + \frac{{\Delta x}}{2}\cdot\left( {f\left( {{x_{n - 1}}} \right) + f\left( {{x_n}} \right)} \right) -- \] -trapezoidStencil :: - (Fractional e, Index ix) - => e -- ^ @Δx@ - distance between sample points - -> Dim -- ^ Dimension along which to integrate - -> Int -- ^ @n@ - number of sample points. +trapezoidStencil + :: (Fractional e, Index ix) + => e + -- ^ @Δx@ - distance between sample points + -> Dim + -- ^ Dimension along which to integrate + -> Int + -- ^ @n@ - number of sample points. -> Stencil ix e e trapezoidStencil dx dim n = makeUnsafeStencil (Sz (setDim' (pureIndex 1) dim (n + 1))) zeroIndex $ \_ g -> - dx / 2 * - (loop 1 (< n) (+ 1) (g zeroIndex) (\i -> (+ 2 * g (setDim' zeroIndex dim i))) + - g (setDim' zeroIndex dim n)) + (dx / 2) + * ( loop 1 (< n) (+ 1) (g zeroIndex) (\i -> (+ 2 * g (setDim' zeroIndex dim i))) + + g (setDim' zeroIndex dim n) + ) {-# INLINE trapezoidStencil #-} - -- | -- -- __Simpson's Rule__ @@ -90,31 +96,36 @@ -- \[ -- \int_{{\,a}}^{{\,b}}{{f\left( x \right)\,dx}} \approx \frac{{\Delta x}}{3}\cdot\left( {f\left( {{x_0}} \right) + 4\cdot f\left( {{x_1}} \right) + f\left( {{x_2}} \right)} \right) + \frac{{\Delta x}}{3}\cdot\left( {f\left( {{x_2}} \right) + 4\cdot f\left( {{x_3}} \right) + f\left( {{x_4}} \right)} \right) + \cdots + \frac{{\Delta x}}{3}\cdot\left( {f\left( {{x_{n - 2}}} \right) + 4\cdot f\left( {{x_{n - 1}}} \right) + f\left( {{x_n}} \right)} \right) -- \] -simpsonsStencil :: - (Fractional e, Index ix) - => e -- ^ @Δx@ - distance between sample points - -> Dim -- ^ Dimension along which to integrate - -> Int -- ^ @n@ - Number of sample points. This value should be even, otherwise error. +simpsonsStencil + :: (Fractional e, Index ix) + => e + -- ^ @Δx@ - distance between sample points + -> Dim + -- ^ Dimension along which to integrate + -> Int + -- ^ @n@ - Number of sample points. This value should be even, otherwise error. -> Stencil ix e e simpsonsStencil dx dim n | odd n = - error $ - "Number of sample points for Simpson's rule stencil should be even, but received: " ++ show n + error $ + "Number of sample points for Simpson's rule stencil should be even, but received: " ++ show n | otherwise = - makeUnsafeStencil (Sz (setDim' (pureIndex 1) dim (n + 1))) zeroIndex $ \_ g -> - let simAcc i (prev, acc) = - let !fx3 = g (setDim' zeroIndex dim (i + 2)) - !newAcc = acc + prev + 4 * g (setDim' zeroIndex dim (i + 1)) + fx3 - in (fx3, newAcc) - in dx / 3 * snd (loop 2 (< n - 1) (+ 2) (simAcc 0 (g zeroIndex, 0)) simAcc) + makeUnsafeStencil (Sz (setDim' (pureIndex 1) dim (n + 1))) zeroIndex $ \_ g -> + let simAcc i (prev, acc) = + let !fx3 = g (setDim' zeroIndex dim (i + 2)) + !newAcc = acc + prev + 4 * g (setDim' zeroIndex dim (i + 1)) + fx3 + in (fx3, newAcc) + in dx / 3 * snd (loop 2 (< n - 1) (+ 2) (simAcc 0 (g zeroIndex, 0)) simAcc) {-# INLINE simpsonsStencil #-} -- | Integrate with a stencil along a particular dimension. -integrateWith :: - (Fractional e, StrideLoad DW ix e, Manifest r e) +integrateWith + :: (Fractional e, StrideLoad DW ix e, Manifest r e) => (Dim -> Int -> Stencil ix e e) - -> Dim -- ^ Dimension along which integration should be estimated. - -> Int -- ^ @n@ - Number of samples + -> Dim + -- ^ Dimension along which integration should be estimated. + -> Int + -- ^ @n@ - Number of samples -> Array r ix e -> Array r ix e integrateWith stencil dim n arr = @@ -123,15 +134,19 @@ !nsz = setDim' (pureIndex 1) dim n {-# INLINE integrateWith #-} - -- | Compute an approximation of integral using a supplied rule in a form of `Stencil`. -integralApprox :: - (Fractional e, StrideLoad DW ix e, 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. +integralApprox + :: (Fractional e, StrideLoad DW ix e, Manifest r e) + => (e -> Dim -> Int -> Stencil ix e e) + -- ^ Integration Stencil + -> e + -- ^ @d@ - Length of interval per cell + -> Sz ix + -- ^ @sz@ - Result size of the matrix + -> Int + -- ^ @n@ - Number of samples + -> Array r ix e + -- ^ Array with values of @f(x,y,..)@ that will be used as source for integration. -> Array D ix e integralApprox stencil d sz n arr = extract' zeroIndex sz $ loop 1 (<= coerce (dimensions sz)) (+ 1) arr integrateAlong @@ -141,55 +156,73 @@ {-# INLINE integrateAlong #-} {-# INLINE integralApprox #-} - -- | Use midpoint rule to approximate an integral. -midpointRule :: - (Fractional e, StrideLoad DW ix e, Manifest r e) - => Comp -- ^ Computation strategy. - -> r -- ^ Intermediate array representation. - -> ((Int -> e) -> ix -> e) -- ^ @f(x,y,...)@ - Function to integrate - -> e -- ^ @a@ - Starting value point. - -> e -- ^ @d@ - Distance per matrix cell. - -> Sz ix -- ^ @sz@ - Result matrix size. - -> Int -- ^ @n@ - Number of sample points per cell in each direction. +midpointRule + :: (Fractional e, StrideLoad DW ix e, Manifest r e) + => Comp + -- ^ Computation strategy. + -> r + -- ^ Intermediate array representation. + -> ((Int -> e) -> ix -> e) + -- ^ @f(x,y,...)@ - Function to integrate + -> e + -- ^ @a@ - Starting value point. + -> e + -- ^ @d@ - Distance per matrix cell. + -> Sz ix + -- ^ @sz@ - Result matrix size. + -> Int + -- ^ @n@ - Number of sample points per cell in each direction. -> Array D ix e midpointRule comp r f a d sz n = integralApprox midpointStencil d sz n $ computeAs r $ fromFunctionMidpoint comp f a d sz n {-# INLINE midpointRule #-} - -- | Use trapezoid rule to approximate an integral. -trapezoidRule :: - (Fractional e, StrideLoad DW ix e, Manifest r e) - => Comp -- ^ Computation strategy - -> r -- ^ Intermediate array representation - -> ((Int -> e) -> ix -> e) -- ^ @f(x,y,...)@ - function to integrate - -> e -- ^ @a@ - Starting value point. - -> e -- ^ @d@ - Distance per matrix cell. - -> Sz ix -- ^ @sz@ - Result matrix size. - -> Int -- ^ @n@ - Number of sample points per cell in each direction. +trapezoidRule + :: (Fractional e, StrideLoad DW ix e, Manifest r e) + => Comp + -- ^ Computation strategy + -> r + -- ^ Intermediate array representation + -> ((Int -> e) -> ix -> e) + -- ^ @f(x,y,...)@ - function to integrate + -> e + -- ^ @a@ - Starting value point. + -> e + -- ^ @d@ - Distance per matrix cell. + -> Sz ix + -- ^ @sz@ - Result matrix size. + -> Int + -- ^ @n@ - Number of sample points per cell in each direction. -> Array D ix e trapezoidRule comp r f a d sz n = integralApprox trapezoidStencil d sz n $ computeAs r $ fromFunction comp f a d sz n {-# INLINE trapezoidRule #-} -- | Use Simpson's rule to approximate an integral. -simpsonsRule :: - (Fractional e, StrideLoad DW ix e, Manifest r e) - => Comp -- ^ Computation strategy - -> r -- ^ Intermediate array representation - -> ((Int -> e) -> ix -> e) -- ^ @f(x,y,...)@ - Function to integrate - -> e -- ^ @a@ - Starting value point. - -> e -- ^ @d@ - Distance per matrix cell. - -> Sz ix -- ^ @sz@ - Result matrix size. - -> Int -- ^ @n@ - Number of sample points per cell in each direction. This value must be even, - -- otherwise error. +simpsonsRule + :: (Fractional e, StrideLoad DW ix e, Manifest r e) + => Comp + -- ^ Computation strategy + -> r + -- ^ Intermediate array representation + -> ((Int -> e) -> ix -> e) + -- ^ @f(x,y,...)@ - Function to integrate + -> e + -- ^ @a@ - Starting value point. + -> e + -- ^ @d@ - Distance per matrix cell. + -> Sz ix + -- ^ @sz@ - Result matrix size. + -> Int + -- ^ @n@ - Number of sample points per cell in each direction. This value must be even, + -- otherwise error. -> Array D ix e simpsonsRule comp r f a d sz n = integralApprox simpsonsStencil d sz n $ computeAs r $ fromFunction comp f a d sz n {-# INLINE simpsonsRule #-} - -- | Create an array from a function with sample points at the edges -- -- >>> fromFunction Seq (\ scale (i :. j) -> scale i + scale j :: Double) (-2) 1 (Sz 4) 2 @@ -204,17 +237,21 @@ -- , [ -0.5, 0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5 ] -- , [ 0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0 ] -- ] --- -fromFunction :: - (Index ix, Fractional a) - => Comp -- ^ Computation strategy +fromFunction + :: (Index ix, Fractional a) + => Comp + -- ^ Computation strategy -> ((Int -> a) -> ix -> e) -- ^ A function that will produce elements of scaled up array. First argument is a scaling -- function that should be applied to individual indicies. - -> a -- ^ @a@ - Starting point - -> a -- ^ @d@ - Distance per cell - -> Sz ix -- ^ @sz@ - Size of the desired array - -> Int -- ^ @n@ - Scaling factor, i.e. number of sample points per cell. + -> a + -- ^ @a@ - Starting point + -> a + -- ^ @d@ - Distance per cell + -> Sz ix + -- ^ @sz@ - Size of the desired array + -> Int + -- ^ @n@ - Scaling factor, i.e. number of sample points per cell. -> Array D ix e fromFunction comp f a d (Sz sz) n = f scale <$> rangeInclusive comp zeroIndex (liftIndex (n *) sz) @@ -224,7 +261,6 @@ {-# INLINE scale #-} {-# INLINE fromFunction #-} - -- | Similar to `fromFunction`, but will create an array from a function with sample points in the -- middle of cells. -- @@ -239,10 +275,15 @@ -- , [ -0.5, 0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0 ] -- , [ 0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5 ] -- ] --- fromFunctionMidpoint - :: (Index ix, Fractional a) => - Comp -> ((Int -> a) -> ix -> e) -> a -> a -> Sz ix -> Int -> Array D ix e + :: (Index ix, Fractional a) + => Comp + -> ((Int -> a) -> ix -> e) + -> a + -> a + -> Sz ix + -> Int + -> Array D ix e fromFunctionMidpoint comp f a d (Sz sz) n = f scale <$> rangeInclusive comp zeroIndex (liftIndex (\i -> n * i - 1) sz) where @@ -251,7 +292,6 @@ scale i = dx2 + a + d * fromIntegral i / nFrac {-# INLINE scale #-} {-# INLINE fromFunctionMidpoint #-} - -- $integral_intro -- 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 @@ -5,6 +5,7 @@ {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} + -- | -- Module : Data.Massiv.Array.Ops.Construct -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -12,68 +13,75 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Ops.Construct - ( -- ** With constant value - empty - , singleton - , replicate - -- ** With a function - , makeArray - , makeArrayLinear - , makeArrayR - , makeArrayLinearR - , makeVectorR - -- *** Iterating - , iterateN - , iiterateN - -- *** Unfolding - , unfoldlS_ - , iunfoldlS_ - , unfoldrS_ - , iunfoldrS_ - , makeSplitSeedArray - -- *** Random - , uniformArray - , uniformRangeArray - , randomArray - , randomArrayS - , randomArrayWS - -- *** Applicative - , makeArrayA - , makeArrayAR - , makeArrayLinearA - -- ** Enumeration - , (...) - , (..:) - , range - , rangeStepM - , rangeStep' - , rangeInclusive - , rangeStepInclusiveM - , rangeStepInclusive' - , rangeSize - , rangeStepSize - , enumFromN - , enumFromStepN - -- ** Expansion - , expandWithin - , expandWithinM - , expandWithin' - , expandOuter - , expandInner - ) where +module Data.Massiv.Array.Ops.Construct ( + -- ** With constant value + empty, + singleton, + replicate, + -- ** With a function + makeArray, + makeArrayLinear, + makeArrayR, + makeArrayLinearR, + makeVectorR, + + -- *** Iterating + iterateN, + iiterateN, + + -- *** Unfolding + unfoldlS_, + iunfoldlS_, + unfoldrS_, + iunfoldrS_, + makeSplitSeedArray, + + -- *** Random + uniformArray, + uniformRangeArray, + randomArray, + randomArrayS, + randomArrayWS, + + -- *** Applicative + makeArrayA, + makeArrayAR, + makeArrayLinearA, + + -- ** Enumeration + (...), + (..:), + range, + rangeStepM, + rangeStep', + rangeInclusive, + rangeStepInclusiveM, + rangeStepInclusive', + rangeSize, + rangeStepSize, + enumFromN, + enumFromStepN, + + -- ** Expansion + expandWithin, + expandWithinM, + expandWithin', + expandOuter, + expandInner, +) where + import Control.Applicative hiding (empty) -import Control.Monad (when, void) +import Control.Monad (void, when) import Control.Monad.ST import Data.Massiv.Array.Delayed.Pull import Data.Massiv.Array.Delayed.Push ---import Data.Massiv.Array.Delayed.Stream (unfoldr, unfoldrN) + +-- import Data.Massiv.Array.Delayed.Stream (unfoldr, unfoldrN) import Data.Massiv.Array.Mutable import Data.Massiv.Core.Common -import Prelude hiding (enumFromTo, replicate) import System.Random.Stateful +import Prelude hiding (enumFromTo, replicate) -- | Just like `makeArray` but with ability to specify the result representation as an -- argument. Note the `Data.Massiv.Array.U`nboxed type constructor in the below example. @@ -112,12 +120,11 @@ 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 :: (Manifest r e, Index ix) => Sz ix -> STA r ix e -> Array r ix e runSTA !sz (STA m) = runST (unsafeNew sz >>= m) -{-# INLINE runSTA #-} +{-# INLINE runSTA #-} -- | Similar to `makeArray`, but construct the array sequentially using an `Applicative` interface. -- @@ -126,37 +133,39 @@ -- -- -- @since 0.2.6 -makeArrayA :: - forall r ix e f. (Manifest r e, Index ix, Applicative f) +makeArrayA + :: forall r ix e f + . (Manifest r e, Index ix, Applicative f) => Sz ix -> (ix -> f e) -> f (Array r ix e) makeArrayA sz@(Sz n) f = fmap (runSTA sz) $ - iterF zeroIndex n oneIndex (<) (pure (STA (unsafeFreeze Seq))) $ \ix g -> - liftA2 (\e (STA st) -> STA (\ma -> unsafeWrite ma ix e >> st ma)) (f ix) g -{-# INLINE makeArrayA #-} + iterF zeroIndex n oneIndex (<) (pure (STA (unsafeFreeze Seq))) $ \ix g -> + liftA2 (\e (STA st) -> STA (\ma -> unsafeWrite ma ix e >> st ma)) (f ix) g +{-# INLINE makeArrayA #-} -- | Same as `makeArrayA`, but with linear index. -- -- @since 0.4.5 -makeArrayLinearA :: - forall r ix e f. (Manifest r e, Index ix, Applicative f) +makeArrayLinearA + :: forall r ix e f + . (Manifest r e, Index ix, Applicative f) => Sz ix -> (Int -> f e) -> f (Array r ix e) makeArrayLinearA !sz f = fmap (runSTA sz) $ - loopF 0 (< totalElem sz) (+ 1) (pure (STA (unsafeFreeze Seq))) $ \i -> - liftA2 (\e (STA st) -> STA (\ma -> unsafeLinearWrite ma i e >> st ma)) (f i) -{-# INLINE makeArrayLinearA #-} - + loopF 0 (< totalElem sz) (+ 1) (pure (STA (unsafeFreeze Seq))) $ \i -> + liftA2 (\e (STA st) -> STA (\ma -> unsafeLinearWrite ma i e >> st ma)) (f i) +{-# INLINE makeArrayLinearA #-} -- | Same as `makeArrayA`, but with ability to supply result array representation. -- -- @since 0.2.6 -makeArrayAR :: - forall r ix e f. (Manifest r e, Index ix, Applicative f) +makeArrayAR + :: forall r ix e f + . (Manifest r e, Index ix, Applicative f) => r -> Sz ix -> (ix -> f e) @@ -164,7 +173,6 @@ makeArrayAR _ = makeArrayA {-# INLINE makeArrayAR #-} - -- | Sequentially iterate over each cell in the array in the row-major order while continuously -- aplying the accumulator at each step. -- @@ -178,14 +186,14 @@ -- ] -- -- @since 0.3.0 -iterateN :: forall ix e . Index ix => Sz ix -> (e -> e) -> e -> Array DL ix e +iterateN :: forall ix e. Index ix => Sz ix -> (e -> e) -> e -> Array DL ix e iterateN sz f = unfoldrS_ sz $ \a -> let !a' = f a in (a', a') {-# INLINE iterateN #-} -- | Same as `iterateN`, but with index aware function. -- -- @since 0.3.0 -iiterateN :: forall ix e . Index ix => Sz ix -> (e -> ix -> e) -> e -> Array DL ix e +iiterateN :: forall ix e. Index ix => Sz ix -> (e -> ix -> e) -> e -> Array DL ix e iiterateN sz f = iunfoldrS_ sz $ \a ix -> let !a' = f a ix in (a', a') {-# INLINE iiterateN #-} @@ -199,8 +207,9 @@ -- [ 10, 11, 12, 13, 14, 15, 16, 17, 18, 19 ] -- -- @since 0.3.0 -unfoldrS_ :: - forall ix e a. Index ix +unfoldrS_ + :: forall ix e a + . Index ix => Sz ix -> (a -> (e, a)) -> a @@ -211,24 +220,24 @@ -- | Right unfold of a delayed load array with index aware function -- -- @since 0.3.0 -iunfoldrS_ :: - forall ix e a. Index ix +iunfoldrS_ + :: 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} +iunfoldrS_ sz f acc0 = DLArray{dlComp = Seq, dlSize = sz, dlLoad = load} where load :: Loader e load _ startAt dlWrite _ = void $ - iterTargetM defRowMajor startAt sz zeroIndex oneStride acc0 $ \ !i !ix !acc -> - case f acc ix of - (e, !acc') -> acc' <$ dlWrite i e + iterTargetM defRowMajor startAt sz zeroIndex oneStride acc0 $ \ !i !ix !acc -> + case f acc ix of + (e, !acc') -> acc' <$ dlWrite i e {-# INLINE load #-} {-# INLINE iunfoldrS_ #-} - -- | Unfold sequentially from the end. There is no way to save the accumulator after -- unfolding is done, since resulting array is delayed, but it's possible to use -- `Data.Massiv.Array.Mutable.unfoldlPrimM` to achieve such effect. @@ -241,24 +250,24 @@ -- | Unfold sequentially from the right with an index aware function. -- -- @since 0.3.0 -iunfoldlS_ :: - forall ix e a. Index ix +iunfoldlS_ + :: 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} +iunfoldlS_ sz f acc0 = DLArray{dlComp = Seq, dlSize = sz, dlLoad = load} where load :: Loader e load _ startAt dlWrite _ = void $ - loopDeepM startAt (< totalElem sz + startAt) (+ 1) acc0 $ \ !i !acc -> - let (acc', e) = f (fromLinearIndex sz (i - startAt)) acc - in acc' <$ dlWrite i e + loopDeepM startAt (< totalElem sz + startAt) (+ 1) acc0 $ \ !i !acc -> + let (acc', e) = f (fromLinearIndex sz (i - startAt)) acc + in acc' <$ dlWrite i e {-# INLINE load #-} {-# INLINE iunfoldlS_ #-} - -- | Create an array with random values by using a pure splittable random number generator -- such as one provided by either [splitmix](https://www.stackage.org/package/splitmix) or -- [random](https://www.stackage.org/package/random) packages. If you don't have a @@ -292,17 +301,21 @@ -- ] -- -- @since 1.0.0 -randomArray :: - forall ix e g. Index ix - => g -- ^ Initial random value generator +randomArray + :: forall ix e g + . Index ix + => g + -- ^ Initial random value generator -> (g -> (g, g)) - -- ^ 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. + -- ^ A function that can split a generator into two independent + -- generators. It will only be called if supplied computation strategy + -- needs more than one worker threads. -> (g -> (e, g)) - -- ^ A function that produces a random value and the next generator - -> Comp -- ^ Computation strategy. - -> Sz ix -- ^ Resulting size of the array. + -- ^ A function that produces a random value and the next generator + -> Comp + -- ^ Computation strategy. + -> Sz ix + -- ^ Resulting size of the array. -> Array DL ix e randomArray gen splitGen nextRandom comp sz = unsafeMakeLoadArray comp sz Nothing load where @@ -321,16 +334,15 @@ then (genI, genI) else splitGen genI scheduleWork_ scheduler $ - void $ loopM start (< start + chunkLength) (+ 1) genI0 writeRandom + void $ + loopM start (< start + chunkLength) (+ 1) genI0 writeRandom pure genI1 when (slackStart < totalLength) $ scheduleWork_ scheduler $ - void $ loopM slackStartAt (< totalLength + startAt) (+ 1) genForSlack writeRandom + void $ + loopM slackStartAt (< totalLength + startAt) (+ 1) genForSlack writeRandom {-# INLINE randomArray #-} - - - -- | Create a delayed array with an initial seed and a splitting function. It is -- somewhat similar to `iunfoldlS_` function, but it is capable of parallelizing -- computation and iterating over the array accoriding to the supplied @@ -340,41 +352,49 @@ -- `Data.Massiv.Array.Manifest.generateSplitSeedArray` -- -- @since 1.0.2 -makeSplitSeedArray :: - forall ix e g it. (Iterator it, Index ix) - => it -- ^ Iterator - -> g -- ^ Initial seed +makeSplitSeedArray + :: forall ix e g it + . (Iterator it, Index ix) + => it + -- ^ Iterator + -> g + -- ^ Initial seed -> (g -> (g, g)) - -- ^ A function that can split a seed into two independent seeds. It will - -- be called the same number of times as the number of jobs that will get - -- scheduled during parallelization. Eg. only once for the sequential case. - -> Comp -- ^ Computation strategy. - -> Sz ix -- ^ Resulting size of the array. + -- ^ A function that can split a seed into two independent seeds. It will + -- be called the same number of times as the number of jobs that will get + -- scheduled during parallelization. Eg. only once for the sequential case. + -> Comp + -- ^ Computation strategy. + -> Sz ix + -- ^ Resulting size of the array. -> (Ix1 -> ix -> g -> (e, g)) - -- ^ A function that produces a value and the next seed. It takes both - -- versions of the index, in linear and in multi-dimensional forms, as well as - -- the current seeding value. + -- ^ A function that produces a value and the next seed. It takes both + -- versions of the index, in linear and in multi-dimensional forms, as well as + -- the current seeding value. -> Array DL ix e makeSplitSeedArray it seed splitSeed comp sz genFunc = - DLArray {dlComp = comp, dlSize = sz, dlLoad = load} + DLArray{dlComp = comp, dlSize = sz, dlLoad = load} where load :: Loader e load scheduler startAt writeAt _ = - iterTargetFullAccST_ it scheduler startAt sz seed (pure . splitSeed) $ \ i ix g -> + iterTargetFullAccST_ it scheduler startAt sz seed (pure . splitSeed) $ \i ix g -> case genFunc (i - startAt) ix g of (x, g') -> g' <$ writeAt i x {-# INLINE load #-} {-# INLINE makeSplitSeedArray #-} - -- | Generate a random array where all elements are sampled from a uniform distribution. -- -- @since 1.0.0 -uniformArray :: - forall ix e g. (Index ix, RandomGen g, Uniform e) - => g -- ^ Initial random value generator. - -> Comp -- ^ Computation strategy. - -> Sz ix -- ^ Resulting size of the array. +uniformArray + :: forall ix e g + . (Index ix, RandomGen g, Uniform e) + => g + -- ^ Initial random value generator. + -> Comp + -- ^ Computation strategy. + -> Sz ix + -- ^ Resulting size of the array. -> Array DL ix e uniformArray gen = randomArray gen split uniform {-# INLINE uniformArray #-} @@ -382,17 +402,21 @@ -- | 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. +uniformRangeArray + :: forall ix e g + . (Index ix, RandomGen g, UniformRange e) + => g + -- ^ Initial random value generator. + -> (e, e) + -- ^ Inclusive range in which values will be generated in. + -> Comp + -- ^ Computation strategy. + -> Sz ix + -- ^ Resulting size of the array. -> Array DL ix e uniformRangeArray gen r = randomArray gen split (uniformR r) {-# INLINE uniformRangeArray #-} - -- | Similar to `randomArray` but performs generation sequentially, which means it doesn't -- require splitability property. Another consequence is that it returns the new generator -- together with /manifest/ array of random values. @@ -427,12 +451,15 @@ -- ] -- -- @since 0.3.4 -randomArrayS :: - forall r ix e g. (Manifest r e, Index ix) - => g -- ^ Initial random value generator - -> Sz ix -- ^ Resulting size of the array. +randomArrayS + :: forall r ix e g + . (Manifest r e, Index ix) + => g + -- ^ Initial random value generator + -> Sz ix + -- ^ Resulting size of the array. -> (g -> (e, g)) - -- ^ A function that produces a random value and the next generator + -- ^ A function that produces a random value and the next generator -> (g, Array r ix e) randomArrayS gen sz nextRandom = runST $ unfoldrPrimM sz (pure . nextRandom) gen @@ -471,12 +498,15 @@ -- [ 8, 8, 7, 1, 1, 2 ] -- -- @since 0.3.4 -randomArrayWS :: - forall r ix e g m. (Manifest r e, Index ix, MonadUnliftIO m, PrimMonad m) +randomArrayWS + :: forall r ix e g m + . (Manifest r e, Index ix, MonadUnliftIO m, PrimMonad m) => WorkerStates g -- ^ Use `Control.Scheduler.initWorkerStates` to initialize you per thread generators - -> Sz ix -- ^ Resulting size of the array - -> (g -> m e) -- ^ Generate the value using the per thread generator. + -> Sz ix + -- ^ Resulting size of the array + -> (g -> m e) + -- ^ Generate the value using the per thread generator. -> m (Array r ix e) randomArrayWS states sz genRandom = generateArrayLinearWS states sz (const genRandom) {-# INLINE randomArrayWS #-} @@ -505,9 +535,9 @@ (..:) = range Seq {-# INLINE (..:) #-} - -- prop> range comp from to == rangeStep comp from 1 to -- + -- | Create an array of indices with a range from start to finish (not-including), where indices are -- incremeted by one. -- @@ -543,20 +573,25 @@ -- *** Exception: IndexZeroException: 0 -- -- @since 0.3.0 -rangeStepM :: - forall ix m. (Index ix, MonadThrow m) - => Comp -- ^ Computation strategy - -> ix -- ^ Start - -> ix -- ^ Step. Negative and positive values are ok, but can't have zeros - -> ix -- ^ End +rangeStepM + :: forall ix m + . (Index ix, MonadThrow m) + => Comp + -- ^ Computation strategy + -> ix + -- ^ Start + -> ix + -- ^ Step. Negative and positive values are ok, but can't have zeros + -> ix + -- ^ End -> m (Array D ix ix) rangeStepM comp !from !step !to | foldlIndex (\acc i -> acc || i == 0) False step = throwM $ IndexZeroException step | otherwise = - let dist = liftIndex2 (-) to from - sz = liftIndex2 div dist step - r = liftIndex signum $ liftIndex2 mod dist step - in pure $ rangeStepSize comp from step (Sz (liftIndex2 (+) sz r)) + let dist = liftIndex2 (-) to from + sz = liftIndex2 div dist step + r = liftIndex signum $ liftIndex2 mod dist step + in pure $ rangeStepSize comp from step (Sz (liftIndex2 (+) sz r)) {-# INLINE rangeStepM #-} -- | Same as `rangeStepM`, but will throw an error whenever @step@ contains zeros. @@ -581,7 +616,6 @@ rangeSize comp ixFrom (Sz (liftIndex2 (-) (liftIndex (+ 1) ixTo) ixFrom)) {-# INLINE rangeInclusive #-} - -- | Just like `rangeStepM`, except the finish index is included. -- -- @since 0.3.0 @@ -596,33 +630,40 @@ rangeStepInclusive' comp ixFrom step = throwEither . rangeStepInclusiveM comp ixFrom step {-# INLINE rangeStepInclusive' #-} - -- | Create an array of specified size with indices starting with some index at position @0@ and -- incremented by @1@ until the end of the array is reached -- -- @since 0.3.0 -rangeSize :: Index ix => - Comp -- ^ Computation strategy - -> ix -- ^ @x@ - start value - -> Sz ix -- ^ @sz@ - Size of resulting array - -> Array D ix ix +rangeSize + :: Index ix + => Comp + -- ^ Computation strategy + -> ix + -- ^ @x@ - start value + -> Sz ix + -- ^ @sz@ - Size of resulting array + -> Array D ix ix rangeSize comp !from !sz = makeArray comp sz (liftIndex2 (+) from) {-# INLINE rangeSize #-} -- | Same as `rangeSize`, but with ability to specify the step. -- -- @since 0.3.0 -rangeStepSize :: Index ix => - Comp -- ^ Computation strategy - -> ix -- ^ @x@ - start value - -> ix -- ^ @delta@ - step value - -> Sz ix -- ^ @sz@ - Size of resulting array - -> Array D ix ix +rangeStepSize + :: Index ix + => Comp + -- ^ Computation strategy + -> ix + -- ^ @x@ - start value + -> ix + -- ^ @delta@ - step value + -> Sz ix + -- ^ @sz@ - Size of resulting array + -> Array D ix ix rangeStepSize comp !from !step !sz = makeArray comp sz (liftIndex2 (+) from . liftIndex2 (*) step) {-# INLINE rangeStepSize #-} - -- | Same as `enumFromStepN` with step @dx = 1@. -- -- /Related/: `Data.Massiv.Vector.senumFromN`, `Data.Massiv.Vector.senumFromStepN`, @@ -643,15 +684,17 @@ -- [@Data.Vector.Generic.`Data.Vector.Generic.enumFromN`@] -- -- @since 0.1.0 -enumFromN :: Num e => - Comp - -> e -- ^ @x@ - start value - -> Sz1 -- ^ @n@ - length of resulting vector. - -> Vector D e -enumFromN comp !from !sz = makeArrayLinear comp sz $ \ i -> from + fromIntegral i +enumFromN + :: Num e + => Comp + -> e + -- ^ @x@ - start value + -> Sz1 + -- ^ @n@ - length of resulting vector. + -> Vector D e +enumFromN comp !from !sz = makeArrayLinear comp sz $ \i -> from + fromIntegral i {-# INLINE enumFromN #-} - -- | Enumerate from a starting number @x@ exactly @n@ times with a custom step value -- @dx@. Unlike `Data.Massiv.Vector.senumFromStepN`, there is no dependency on neigboring -- elements therefore `enumFromStepN` is parallelizable. @@ -680,16 +723,19 @@ -- -- -- @since 0.1.0 -enumFromStepN :: Num e => - Comp - -> e -- ^ @x@ - start number - -> e -- ^ @dx@ - step number - -> Sz1 -- ^ @n@ - length of resulting vector - -> Vector D e -enumFromStepN comp !from !step !sz = makeArrayLinear comp sz $ \ i -> from + fromIntegral i * step +enumFromStepN + :: Num e + => Comp + -> e + -- ^ @x@ - start number + -> e + -- ^ @dx@ - step number + -> Sz1 + -- ^ @n@ - length of resulting vector + -> Vector D e +enumFromStepN comp !from !step !sz = makeArrayLinear comp sz $ \i -> from + fromIntegral i * step {-# INLINE enumFromStepN #-} - -- | Function that expands an array to one with a higher dimension. -- -- This is useful for constructing arrays where there is shared computation @@ -731,8 +777,9 @@ -- ] -- -- @since 0.2.6 -expandWithin :: - forall n ix e r a. (IsIndexDimension ix n, Index (Lower ix), Manifest r a) +expandWithin + :: forall n ix e r a + . (IsIndexDimension ix n, Index (Lower ix), Manifest r a) => Dimension n -> Sz1 -> (a -> Ix1 -> e) @@ -751,8 +798,9 @@ -- will throw an exception on an invalid dimension. -- -- @since 0.2.6 -expandWithin' :: - forall r ix a b. (HasCallStack, Index ix, Index (Lower ix), Manifest r a) +expandWithin' + :: forall r ix a b + . (HasCallStack, Index ix, Index (Lower ix), Manifest r a) => Dim -> Sz1 -> (a -> Ix1 -> b) @@ -765,8 +813,9 @@ -- will throw an exception on an invalid dimension. -- -- @since 0.4.0 -expandWithinM :: - forall r ix a b m. (Index ix, Index (Lower ix), Manifest r 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) @@ -783,8 +832,9 @@ -- | Similar to `expandWithin`, except it uses the outermost dimension. -- -- @since 0.2.6 -expandOuter :: - forall r ix a b. (Index ix, Index (Lower ix), Manifest r a) +expandOuter + :: forall r ix a b + . (Index ix, Index (Lower ix), Manifest r a) => Sz1 -> (a -> Ix1 -> b) -> Array r (Lower ix) a @@ -801,8 +851,9 @@ -- | Similar to `expandWithin`, except it uses the innermost dimension. -- -- @since 0.2.6 -expandInner :: - forall r ix a b. (Index ix, Index (Lower ix), Manifest r 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 @@ -3,6 +3,7 @@ {-# LANGUAGE GADTs #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE ScopedTypeVariables #-} + -- | -- Module : Data.Massiv.Array.Ops.Fold -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -10,151 +11,156 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Ops.Fold - ( +module Data.Massiv.Array.Ops.Fold ( -- ** Unstructured folds - -- $unstruct_folds - fold - , ifoldMono - , foldMono - , ifoldSemi - , foldSemi - , foldOuterSlice - , ifoldOuterSlice - , foldInnerSlice - , ifoldInnerSlice - , minimumM - , minimum' - , maximumM - , maximum' - , sum - , product - , and - , or - , all - , any - , elem - , eqArrays - , compareArrays + fold, + ifoldMono, + foldMono, + ifoldSemi, + foldSemi, + foldOuterSlice, + ifoldOuterSlice, + foldInnerSlice, + ifoldInnerSlice, + minimumM, + minimum', + maximumM, + maximum', + sum, + product, + and, + or, + all, + any, + elem, + eqArrays, + compareArrays, -- ** Single dimension folds + -- *** Safe inner most + -- -- Folding along the inner most dimension will always be faster when compared to doing the same -- operation along any other dimension, this is due to the fact that inner most folds follow the -- memory layout of data. - , ifoldlInner - , foldlInner - , ifoldrInner - , foldrInner - , foldInner + ifoldlInner, + foldlInner, + ifoldrInner, + foldrInner, + foldInner, + -- *** Type safe within - , ifoldlWithin - , foldlWithin - , ifoldrWithin - , foldrWithin - , foldWithin + ifoldlWithin, + foldlWithin, + ifoldrWithin, + foldrWithin, + foldWithin, + -- *** Partial within - , ifoldlWithin' - , foldlWithin' - , ifoldrWithin' - , foldrWithin' - , foldWithin' + ifoldlWithin', + foldlWithin', + ifoldrWithin', + foldrWithin', + foldWithin', -- ** Sequential folds - -- $seq_folds - - , foldlS - , foldrS - , ifoldlS - , ifoldrS + foldlS, + foldrS, + ifoldlS, + ifoldrS, -- *** Monadic - , foldlM - , foldrM - , foldlM_ - , foldrM_ - , ifoldlM - , ifoldrM - , ifoldlM_ - , ifoldrM_ + foldlM, + foldrM, + foldlM_, + foldrM_, + ifoldlM, + ifoldrM, + ifoldlM_, + ifoldrM_, -- *** Special folds - , foldrFB - , lazyFoldlS - , lazyFoldrS + foldrFB, + lazyFoldlS, + lazyFoldrS, -- ** Parallel folds - -- $par_folds - - , foldlP - , foldrP - , ifoldlP - , ifoldrP - , ifoldlIO - , ifoldrIO + foldlP, + foldrP, + ifoldlP, + ifoldrP, + ifoldlIO, + ifoldrIO, -- , splitReduce - ) where +) where import Data.Massiv.Array.Delayed.Pull import Data.Massiv.Array.Ops.Construct import Data.Massiv.Array.Ops.Fold.Internal import Data.Massiv.Core import Data.Massiv.Core.Common -import Prelude hiding (all, and, any, foldl, foldr, map, maximum, minimum, or, product, sum, elem) - +import Prelude hiding (all, and, any, elem, foldl, foldr, map, maximum, minimum, or, product, sum) -- | /O(n)/ - Monoidal fold over an array with an index aware function. Also known as reduce. -- -- @since 0.2.4 -ifoldMono :: - (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 +ifoldMono + :: (Index ix, Source r e, Monoid m) + => (ix -> e -> m) + -- ^ Convert each element of an array to an appropriate `Monoid`. + -> Array r ix e + -- ^ Source array -> m ifoldMono f = ifoldlInternal (\a ix e -> a `mappend` f ix e) mempty mappend mempty {-# INLINE ifoldMono #-} - -- | /O(n)/ - Semigroup fold over an array with an index aware function. -- -- @since 0.2.4 -ifoldSemi :: - (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 +ifoldSemi + :: (Index ix, Source r e, Semigroup m) + => (ix -> e -> m) + -- ^ Convert each element of an array to an appropriate `Semigroup`. -> m + -- ^ Initial element that must be neutral to the (`<>`) function. + -> Array r ix e + -- ^ Source array + -> m ifoldSemi f m = ifoldlInternal (\a ix e -> a <> f ix e) m (<>) m {-# INLINE ifoldSemi #-} - -- | /O(n)/ - Semigroup fold over an array. -- -- @since 0.1.6 -foldSemi :: - (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 +foldSemi + :: (Index ix, Source r e, Semigroup m) + => (e -> m) + -- ^ Convert each element of an array to an appropriate `Semigroup`. -> m + -- ^ Initial element that must be neutral to the (`<>`) function. + -> Array r ix e + -- ^ Source array + -> m foldSemi f m = foldlInternal (\a e -> a <> f e) m (<>) m {-# INLINE foldSemi #-} - -- | Left fold along a specified dimension with an index aware function. -- -- @since 0.2.4 -ifoldlWithin :: (Index (Lower ix), IsIndexDimension ix n, Source r e) => - Dimension n -> (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a +ifoldlWithin + :: (Index (Lower ix), IsIndexDimension ix n, Source r e) + => Dimension n + -> (ix -> a -> e -> a) + -> a + -> Array r ix e + -> Array D (Lower ix) a ifoldlWithin dim = ifoldlWithin' (fromDimension dim) {-# INLINE ifoldlWithin #-} - -- | Left fold along a specified dimension. -- -- ====__Example__ @@ -175,36 +181,53 @@ -- [ [5,0], [6,1], [7,2], [8,3], [9,4] ] -- -- @since 0.2.4 -foldlWithin :: (Index (Lower ix), IsIndexDimension ix n, Source r e) => - Dimension n -> (a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a +foldlWithin + :: (Index (Lower ix), IsIndexDimension ix n, Source r e) + => Dimension n + -> (a -> e -> a) + -> a + -> Array r ix e + -> Array D (Lower ix) a foldlWithin dim f = ifoldlWithin dim (const f) {-# INLINE foldlWithin #-} - -- | Right fold along a specified dimension with an index aware function. -- -- @since 0.2.4 -ifoldrWithin :: (Index (Lower ix), IsIndexDimension ix n, Source r e) => - Dimension n -> (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a +ifoldrWithin + :: (Index (Lower ix), IsIndexDimension ix n, Source r e) + => Dimension n + -> (ix -> e -> a -> a) + -> a + -> Array r ix e + -> Array D (Lower ix) a ifoldrWithin dim = ifoldrWithin' (fromDimension dim) {-# INLINE ifoldrWithin #-} - -- | Right fold along a specified dimension. -- -- @since 0.2.4 -foldrWithin :: (Index (Lower ix), IsIndexDimension ix n, Source r e) => - Dimension n -> (e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a +foldrWithin + :: (Index (Lower ix), IsIndexDimension ix n, Source r e) + => Dimension n + -> (e -> a -> a) + -> a + -> Array r ix e + -> Array D (Lower ix) a foldrWithin dim f = ifoldrWithin dim (const f) {-# INLINE foldrWithin #-} - -- | Similar to `ifoldlWithin`, except that dimension is specified at a value level, which means it -- will throw an exception on an invalid dimension. -- -- @since 0.2.4 -ifoldlWithin' :: (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' + :: (HasCallStack, Index (Lower ix), Index ix, Source r e) + => Dim + -> (ix -> a -> e -> a) + -> a + -> Array r ix e + -> Array D (Lower ix) a ifoldlWithin' dim f acc0 arr = makeArray (getComp arr) (SafeSz szl) $ \ixl -> iter @@ -219,24 +242,32 @@ (k, szl) = pullOutDim' sz dim {-# INLINE ifoldlWithin' #-} - -- | Similar to `foldlWithin`, except that dimension is specified at a value level, which means it will -- throw an exception on an invalid dimension. -- -- @since 0.2.4 -foldlWithin' :: (HasCallStack, Index (Lower ix), Index ix, Source r e) => - Dim -> (a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a +foldlWithin' + :: (HasCallStack, Index (Lower ix), Index ix, Source r e) + => Dim + -> (a -> e -> a) + -> a + -> Array r ix e + -> Array D (Lower ix) a foldlWithin' dim f = ifoldlWithin' dim (const f) {-# INLINE foldlWithin' #-} - -- | Similar to `ifoldrWithin`, except that dimension is specified at a value level, which means it -- will throw an exception on an invalid dimension. -- -- -- @since 0.2.4 -ifoldrWithin' :: (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' + :: (HasCallStack, Index (Lower ix), Index ix, Source r e) + => Dim + -> (ix -> e -> a -> a) + -> a + -> Array r ix e + -> Array D (Lower ix) a ifoldrWithin' dim f acc0 arr = makeArray (getComp arr) (SafeSz szl) $ \ixl -> iter @@ -255,41 +286,61 @@ -- will throw an exception on an invalid dimension. -- -- @since 0.2.4 -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' + :: (HasCallStack, Index (Lower ix), Index ix, Source r e) + => Dim + -> (e -> a -> a) + -> a + -> Array r ix e + -> Array D (Lower ix) a foldrWithin' dim f = ifoldrWithin' dim (const f) {-# INLINE foldrWithin' #-} - -- | Left fold over the inner most dimension with index aware function. -- -- @since 0.2.4 -ifoldlInner :: (Index (Lower ix), Index ix, Source r e) => - (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a +ifoldlInner + :: (Index (Lower ix), Index ix, Source r e) + => (ix -> a -> e -> a) + -> a + -> Array r ix e + -> Array D (Lower ix) a ifoldlInner = ifoldlWithin' 1 {-# INLINE ifoldlInner #-} -- | Left fold over the inner most dimension. -- -- @since 0.2.4 -foldlInner :: (Index (Lower ix), Index ix, Source r e) => - (a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a +foldlInner + :: (Index (Lower ix), Index ix, Source r e) + => (a -> e -> a) + -> a + -> Array r ix e + -> Array D (Lower ix) a foldlInner = foldlWithin' 1 {-# INLINE foldlInner #-} -- | Right fold over the inner most dimension with index aware function. -- -- @since 0.2.4 -ifoldrInner :: (Index (Lower ix), Index ix, Source r e) => - (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a +ifoldrInner + :: (Index (Lower ix), Index ix, Source r e) + => (ix -> e -> a -> a) + -> a + -> Array r ix e + -> Array D (Lower ix) a ifoldrInner = ifoldrWithin' 1 {-# INLINE ifoldrInner #-} -- | Right fold over the inner most dimension. -- -- @since 0.2.4 -foldrInner :: (Index (Lower ix), Index ix, Source r e) => - (e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a +foldrInner + :: (Index (Lower ix), Index ix, Source r e) + => (e -> a -> a) + -> a + -> Array r ix e + -> Array D (Lower ix) a foldrInner = foldrWithin' 1 {-# INLINE foldrInner #-} @@ -303,8 +354,8 @@ -- | Monoidal fold over some internal dimension. -- -- @since 0.4.3 -foldWithin :: - (Source r a, Monoid a, Index (Lower ix), IsIndexDimension ix n) +foldWithin + :: (Source r a, Monoid a, Index (Lower ix), IsIndexDimension ix n) => Dimension n -> Array r ix a -> Array D (Lower ix) a @@ -315,15 +366,14 @@ -- result in `IndexDimensionException` if supplied dimension is invalid. -- -- @since 0.4.3 -foldWithin' :: - (HasCallStack, Index ix, Source r a, Monoid a, Index (Lower ix)) +foldWithin' + :: (HasCallStack, Index ix, Source r a, Monoid a, Index (Lower ix)) => Dim -> Array r ix a -> Array D (Lower ix) a foldWithin' dim = foldlWithin' dim mappend mempty {-# INLINE foldWithin' #-} - -- | Reduce each outer slice into a monoid and mappend results together -- -- ==== __Example__ @@ -342,21 +392,20 @@ -- 1620 -- -- @since 0.4.3 -foldOuterSlice :: - (Index ix, Index (Lower ix), Source r e, Monoid m) +foldOuterSlice + :: (Index ix, Index (Lower ix), Source r e, Monoid m) => (Array r (Lower ix) e -> m) -> Array r ix e -> m foldOuterSlice f = ifoldOuterSlice (const f) {-# INLINE foldOuterSlice #-} - -- | Reduce each outer slice into a monoid with an index aware function and mappend results -- together -- -- @since 0.4.3 -ifoldOuterSlice :: - (Index ix, Index (Lower ix), Source r e, Monoid m) +ifoldOuterSlice + :: (Index ix, Index (Lower ix), Source r e, Monoid m) => (Ix1 -> Array r (Lower ix) e -> m) -> Array r ix e -> m @@ -367,7 +416,6 @@ {-# INLINE g #-} {-# INLINE ifoldOuterSlice #-} - -- | Reduce each inner slice into a monoid and mappend results together -- -- ==== __Example__ @@ -386,18 +434,17 @@ -- 19575 -- -- @since 0.4.3 -foldInnerSlice :: - (Source r e, Index ix, Monoid m) => (Array D (Lower ix) e -> m) -> Array r ix e -> m +foldInnerSlice + :: (Source r e, Index ix, Monoid m) => (Array D (Lower ix) e -> m) -> Array r ix e -> m foldInnerSlice f = ifoldInnerSlice (const f) {-# INLINE foldInnerSlice #-} - -- | Reduce each inner slice into a monoid with an index aware function and mappend -- results together -- -- @since 0.4.3 -ifoldInnerSlice :: - (Source r e, Index ix, Monoid m) => (Ix1 -> Array D (Lower 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 (szL, !k) = unsnocSz (size arr) @@ -412,22 +459,22 @@ maximumM arr = if isNull arr then throwM (SizeEmptyException (size arr)) - else let !e0 = unsafeIndex arr zeroIndex - in pure $ foldlInternal max e0 max e0 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' :: - forall r ix e. (HasCallStack, Shape r ix, Source r e, Ord e) +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 @@ -435,8 +482,9 @@ minimumM arr = if isNull arr then throwM (SizeEmptyException (size arr)) - else let !e0 = unsafeIndex arr zeroIndex - in pure $ foldlInternal min e0 min e0 arr + else + let !e0 = unsafeIndex arr zeroIndex + in pure $ foldlInternal min e0 min e0 arr {-# INLINE minimumM #-} -- | /O(n)/ - Compute minimum of all elements. @@ -446,7 +494,6 @@ minimum' = throwEither . minimumM {-# INLINE minimum' #-} - -- -- | /O(n)/ - Compute sum of all elements. -- -- -- -- @since 0.1.0 @@ -464,7 +511,6 @@ sum = foldlInternal (+) 0 (+) 0 {-# INLINE sum #-} - -- | /O(n)/ - Compute product of all elements. -- -- @since 0.1.0 @@ -472,7 +518,6 @@ product = foldlInternal (*) 1 (*) 1 {-# INLINE product #-} - -- | /O(n)/ - Compute conjunction of all elements. -- -- @since 0.1.0 @@ -480,7 +525,6 @@ and = all id {-# INLINE and #-} - -- | /O(n)/ - Compute disjunction of all elements. -- -- @since 0.1.0 @@ -488,7 +532,6 @@ or = any id {-# INLINE or #-} - -- | /O(n)/ - Determines whether all elements of the array satisfy a predicate. -- -- @since 0.1.0 @@ -503,36 +546,28 @@ elem e = any (e ==) {-# INLINE elem #-} - -{- $unstruct_folds - -Functions in this section will fold any `Source` array with respect to the inner -`Comp`utation strategy setting. - --} - - -{- $seq_folds - -Functions in this section will fold any `Source` array sequentially, regardless of the inner -`Comp`utation strategy setting. - --} - - -{- $par_folds - -__Note__ It is important to compile with @-threaded -with-rtsopts=-N@ flags, otherwise there will be -no parallelization. - -Functions in this section will fold any `Source` array in parallel, regardless of the inner -`Comp`utation strategy setting. All of the parallel structured folds are performed inside `IO` -monad, because referential transparency can't generally be preserved and results will depend on the -number of cores/capabilities that computation is being performed on. +-- $unstruct_folds +-- +-- Functions in this section will fold any `Source` array with respect to the inner +-- `Comp`utation strategy setting. -In contrast to sequential folds, each parallel folding function accepts two functions and two -initial elements as arguments. This is necessary because an array is first split into chunks, which -folded individually on separate cores with the first function, and the results of those folds are -further folded with the second function. +-- $seq_folds +-- +-- Functions in this section will fold any `Source` array sequentially, regardless of the inner +-- `Comp`utation strategy setting. --} +-- $par_folds +-- +-- __Note__ It is important to compile with @-threaded -with-rtsopts=-N@ flags, otherwise +-- there will be no parallelization. +-- +-- Functions in this section will fold any `Source` array in parallel, regardless of the +-- inner `Comp`utation strategy setting. All of the parallel structured folds are +-- performed inside `IO` monad, because referential transparency can't generally be +-- preserved and results will depend on the number of cores/capabilities that computation +-- is being performed on. +-- +-- In contrast to sequential folds, each parallel folding function accepts two functions +-- and two initial elements as arguments. This is necessary because an array is first +-- split into chunks, which folded individually on separate cores with the first function, +-- and the results of those folds are further folded with the second function. diff --git a/src/Data/Massiv/Array/Ops/Fold/Internal.hs b/src/Data/Massiv/Array/Ops/Fold/Internal.hs --- a/src/Data/Massiv/Array/Ops/Fold/Internal.hs +++ b/src/Data/Massiv/Array/Ops/Fold/Internal.hs @@ -4,6 +4,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UndecidableInstances #-} + -- | -- Module : Data.Massiv.Array.Ops.Fold.Internal -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -11,43 +12,41 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Ops.Fold.Internal - ( - foldlS - , foldrS - , ifoldlS - , ifoldrS - --Monadic - , foldlM - , foldrM - , foldlM_ - , foldrM_ - , ifoldlM - , ifoldrM - , ifoldlM_ - , ifoldrM_ - --Special folds - , fold - , foldMono - , foldlInternal - , ifoldlInternal - , foldrFB - , lazyFoldlS - , lazyFoldrS +module Data.Massiv.Array.Ops.Fold.Internal ( + foldlS, + foldrS, + ifoldlS, + ifoldrS, + -- Monadic + foldlM, + foldrM, + foldlM_, + foldrM_, + ifoldlM, + ifoldrM, + ifoldlM_, + ifoldrM_, + -- Special folds + fold, + foldMono, + foldlInternal, + ifoldlInternal, + foldrFB, + lazyFoldlS, + lazyFoldrS, -- Parallel folds - , foldlP - , foldrP - , ifoldlP - , ifoldrP - , foldlIO - , ifoldlIO - , ifoldrIO - , splitReduce - , any - , anySu - , anyPu - ) where + foldlP, + foldrP, + ifoldlP, + ifoldrP, + foldlIO, + ifoldlIO, + ifoldrIO, + splitReduce, + any, + anySu, + anyPu, +) where import Control.Monad (void, when) import Control.Monad.Primitive @@ -55,35 +54,35 @@ import qualified Data.Foldable as F import Data.Functor.Identity (runIdentity) import Data.Massiv.Core.Common -import Prelude hiding (any, foldl, foldr) import System.IO.Unsafe (unsafePerformIO) - +import Prelude hiding (any, foldl, foldr) -- | /O(n)/ - Unstructured fold of an array. -- -- @since 0.3.0 -fold :: - (Monoid e, Index ix, Source r e) - => Array r ix e -- ^ Source array +fold + :: (Monoid e, Index ix, Source r e) + => Array r ix e + -- ^ Source array -> e fold = foldlInternal mappend mempty mappend mempty {-# INLINE fold #-} - -- | /O(n)/ - This is exactly like `Data.Foldable.foldMap`, but for arrays. Fold over an array, -- while converting each element into a `Monoid`. Also known as map-reduce. If elements of the array -- are already a `Monoid` you can use `fold` instead. -- -- @since 0.1.4 -foldMono :: - (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 +foldMono + :: (Index ix, Source r e, Monoid m) + => (e -> m) + -- ^ Convert each element of an array to an appropriate `Monoid`. + -> Array r ix e + -- ^ Source array -> m foldMono f = foldlInternal (\a e -> a `mappend` f e) mempty mappend mempty {-# INLINE foldMono #-} - -- | /O(n)/ - Monadic left fold. -- -- @since 0.1.0 @@ -98,7 +97,6 @@ sz = size arr {-# INLINE foldlM #-} - -- | /O(n)/ - Monadic left fold, that discards the result. -- -- @since 0.1.0 @@ -106,7 +104,6 @@ foldlM_ f acc = void . foldlM f acc {-# INLINE foldlM_ #-} - -- | /O(n)/ - Monadic left fold with an index aware function. -- -- @since 0.1.0 @@ -119,7 +116,6 @@ iterTargetM defRowMajor 0 (size arr) zeroIndex oneStride acc $ \i ix !a -> f a ix (gi i) {-# INLINE ifoldlM #-} - -- | /O(n)/ - Monadic left fold with an index aware function, that discards the result. -- -- @since 0.1.0 @@ -127,7 +123,6 @@ ifoldlM_ f acc = void . ifoldlM f acc {-# INLINE ifoldlM_ #-} - -- | /O(n)/ - Monadic right fold. -- -- @since 0.1.0 @@ -142,7 +137,6 @@ !sz = size arr {-# INLINE foldrM #-} - -- | /O(n)/ - Monadic right fold, that discards the result. -- -- @since 0.1.0 @@ -150,7 +144,6 @@ foldrM_ f = ifoldrM_ (\_ e a -> f e a) {-# INLINE foldrM_ #-} - -- | /O(n)/ - Monadic right fold with an index aware function. -- -- @since 0.1.0 @@ -160,7 +153,6 @@ f ix (unsafeIndex arr ix) {-# INLINE ifoldrM #-} - -- | /O(n)/ - Monadic right fold with an index aware function, that discards the result. -- -- @since 0.1.0 @@ -168,8 +160,6 @@ ifoldrM_ f !acc !arr = void $ ifoldrM f acc arr {-# INLINE ifoldrM_ #-} - - -- | /O(n)/ - Left fold, computed sequentially with lazy accumulator. -- -- @since 0.1.0 @@ -182,7 +172,6 @@ | otherwise = acc {-# INLINE lazyFoldlS #-} - -- | /O(n)/ - Right fold, computed sequentially with lazy accumulator. -- -- @since 0.1.0 @@ -190,40 +179,39 @@ lazyFoldrS = foldrFB {-# INLINE lazyFoldrS #-} - -- | /O(n)/ - Left fold, computed sequentially. -- -- @since 0.1.0 foldlS :: (Index ix, Source r e) => (a -> e -> a) -> a -> Array r ix e -> a -foldlS f acc = runIdentity . foldlM (\ a e -> pure $! f a e) acc +foldlS f acc = runIdentity . foldlM (\a e -> pure $! f a e) acc {-# INLINE foldlS #-} - -- | /O(n)/ - Left fold with an index aware function, computed sequentially. -- -- @since 0.1.0 -ifoldlS :: (Index ix, Source r e) - => (a -> ix -> e -> a) -> a -> Array r ix e -> a -ifoldlS f acc = runIdentity . ifoldlM (\ a ix e -> pure $! f a ix e) acc +ifoldlS + :: (Index ix, Source r e) + => (a -> ix -> e -> a) + -> a + -> Array r ix e + -> a +ifoldlS f acc = runIdentity . ifoldlM (\a ix e -> pure $! f a ix e) acc {-# INLINE ifoldlS #-} - -- | /O(n)/ - Right fold, computed sequentially. -- -- @since 0.1.0 foldrS :: (Index ix, Source r e) => (e -> a -> a) -> a -> Array r ix e -> a -foldrS f acc = runIdentity . foldrM (\ e a -> pure $! f e a) acc +foldrS f acc = runIdentity . foldrM (\e a -> pure $! f e a) acc {-# INLINE foldrS #-} - -- | /O(n)/ - Right fold with an index aware function, computed sequentially. -- -- @since 0.1.0 ifoldrS :: (Index ix, Source r e) => (ix -> e -> a -> a) -> a -> Array r ix e -> a -ifoldrS f acc = runIdentity . ifoldrM (\ ix e a -> pure $! f ix e a) acc +ifoldrS f acc = runIdentity . ifoldrM (\ix e a -> pure $! f ix e a) acc {-# INLINE ifoldrS #-} - -- | Version of foldr that supports @foldr/build@ list fusion implemented by GHC. -- -- @since 0.1.0 @@ -236,8 +224,6 @@ | otherwise = let v = unsafeLinearIndex arr i in v `c` go (i + 1) {-# INLINE [0] foldrFB #-} - - -- | /O(n)/ - Left fold, computed with respect of array's computation strategy. Because we do -- potentially split the folding among many threads, we also need a combining function and an -- accumulator for the results. Depending on the number of threads being used, results can be @@ -256,12 +242,18 @@ -- [1,0,3,2,5,4] -- -- @since 0.1.0 -foldlP :: (MonadIO m, Index ix, Source r e) => - (a -> e -> a) -- ^ Folding function @g@. - -> a -- ^ Accumulator. Will be applied to @g@ multiple times, thus must be neutral. - -> (b -> a -> b) -- ^ Chunk results folding function @f@. - -> b -- ^ Accumulator for results of chunks folding. - -> Array r ix e -> m b +foldlP + :: (MonadIO m, Index ix, Source r e) + => (a -> e -> a) + -- ^ Folding function @g@. + -> a + -- ^ Accumulator. Will be applied to @g@ multiple times, thus must be neutral. + -> (b -> a -> b) + -- ^ Chunk results folding function @f@. + -> b + -- ^ Accumulator for results of chunks folding. + -> Array r ix e + -> m b foldlP f fAcc g gAcc = liftIO . foldlIO (\acc -> pure . f acc) fAcc (\acc -> pure . g acc) gAcc {-# INLINE foldlP #-} @@ -271,13 +263,18 @@ -- element it is being applied to. -- -- @since 0.1.0 -ifoldlP :: (MonadIO m, Index ix, Source r e) => - (a -> ix -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> m b +ifoldlP + :: (MonadIO m, Index ix, Source r e) + => (a -> ix -> e -> a) + -> a + -> (b -> a -> b) + -> b + -> Array r ix e + -> m b ifoldlP f fAcc g gAcc = liftIO . ifoldlIO (\acc ix -> pure . f acc ix) fAcc (\acc -> pure . g acc) gAcc {-# INLINE ifoldlP #-} - -- | /O(n)/ - Right fold, computed with respect to computation strategy. Same as `foldlP`, except -- directed from the last element in the array towards beginning. -- @@ -292,19 +289,24 @@ -- [[0,1],[2,3],[4,5]] -- -- @since 0.1.0 -foldrP :: (MonadIO m, Index ix, Source r e) => - (e -> a -> a) -> a -> (a -> b -> b) -> b -> Array r ix e -> m b +foldrP + :: (MonadIO m, Index ix, Source r e) + => (e -> a -> a) + -> a + -> (a -> b -> b) + -> b + -> Array r ix e + -> m b foldrP f fAcc g gAcc = liftIO . ifoldrP (const f) fAcc g gAcc {-# INLINE foldrP #-} - -- | /O(n)/ - Right fold with an index aware function, while respecting the computation strategy. -- Same as `ifoldlP`, except directed from the last element in the array towards -- beginning, but also row-major. -- -- @since 0.1.0 -ifoldrP :: - (MonadIO m, Index ix, Source r e) +ifoldrP + :: (MonadIO m, Index ix, Source r e) => (ix -> e -> a -> a) -> a -> (a -> b -> b) @@ -314,79 +316,84 @@ ifoldrP f fAcc g gAcc = liftIO . ifoldrIO (\ix e -> pure . f ix e) fAcc (\e -> pure . g e) gAcc {-# INLINE ifoldrP #-} - -- | This folding function breaks referential transparency on some functions -- @f@, therefore it is kept here for internal use only. -foldlInternal :: - (Index ix, Source r 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 :: - (Index ix, Source r e) => (a -> ix -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> b +ifoldlInternal + :: (Index ix, Source r e) => (a -> ix -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> b ifoldlInternal g initAcc f resAcc = unsafePerformIO . ifoldlP g initAcc f resAcc {-# INLINE ifoldlInternal #-} - -- | Similar to `foldlP`, except that folding functions themselves do live in IO -- -- @since 0.1.0 -foldlIO :: - (MonadUnliftIO m, Index ix, Source r e) - => (a -> e -> m a) -- ^ Index aware folding IO action - -> a -- ^ Accumulator - -> (b -> a -> m b) -- ^ Folding action that is applied to the results of a parallel fold - -> b -- ^ Accumulator for chunks folding +foldlIO + :: (MonadUnliftIO m, Index ix, Source r e) + => (a -> e -> m a) + -- ^ Index aware folding IO action + -> a + -- ^ Accumulator + -> (b -> a -> m b) + -- ^ Folding action that is applied to the results of a parallel fold + -> b + -- ^ Accumulator for chunks folding -> Array r ix e -> m b foldlIO f !initAcc g !tAcc !arr | getComp arr == Seq = foldlM f initAcc arr >>= g tAcc | otherwise = do - let splitAcc _ = pure (initAcc, initAcc) - !sz = size arr - results <- - withScheduler (getComp arr) $ \scheduler -> - withRunInIO $ \run -> - stToPrim $ - case unsafePrefIndex arr of - PrefIndex gix -> - iterFullAccST defRowMajor scheduler zeroIndex sz initAcc splitAcc $ \ !ix !acc -> - ioToPrim (run (f acc (gix ix))) - PrefIndexLinear gi -> - iterFullAccST defRowMajor scheduler 0 (toLinearSz sz) initAcc splitAcc $ \ !i !acc -> - ioToPrim (run (f acc (gi i))) - F.foldlM g tAcc results + let splitAcc _ = pure (initAcc, initAcc) + !sz = size arr + results <- + withScheduler (getComp arr) $ \scheduler -> + withRunInIO $ \run -> + stToPrim $ + case unsafePrefIndex arr of + PrefIndex gix -> + iterFullAccST defRowMajor scheduler zeroIndex sz initAcc splitAcc $ \ !ix !acc -> + ioToPrim (run (f acc (gix ix))) + PrefIndexLinear gi -> + iterFullAccST defRowMajor scheduler 0 (toLinearSz sz) initAcc splitAcc $ \ !i !acc -> + ioToPrim (run (f acc (gi i))) + F.foldlM g tAcc results {-# INLINE foldlIO #-} -- | Similar to `ifoldlP`, except that folding functions themselves do live in IO -- -- @since 0.1.0 -ifoldlIO :: - (MonadUnliftIO m, Index ix, Source r e) - => (a -> ix -> e -> m a) -- ^ Index aware folding IO action - -> a -- ^ Accumulator - -> (b -> a -> m b) -- ^ Folding action that is applied to the results of a parallel fold - -> b -- ^ Accumulator for chunks folding +ifoldlIO + :: (MonadUnliftIO m, Index ix, Source r e) + => (a -> ix -> e -> m a) + -- ^ Index aware folding IO action + -> a + -- ^ Accumulator + -> (b -> a -> m b) + -- ^ Folding action that is applied to the results of a parallel fold + -> b + -- ^ Accumulator for chunks folding -> Array r ix e -> m b ifoldlIO f !initAcc g !tAcc !arr | getComp arr == Seq = ifoldlM f initAcc arr >>= g tAcc | otherwise = do - let !sz = size arr - splitAcc _ = pure (initAcc, initAcc) - results <- - withScheduler (getComp arr) $ \scheduler -> - withRunInIO $ \run -> - stToPrim $ - case unsafePrefIndex arr of - PrefIndex gix -> - iterFullAccST defRowMajor scheduler zeroIndex sz initAcc splitAcc $ \ !ix !acc -> - ioToPrim (run (f acc ix (gix ix))) - PrefIndexLinear gi -> - iterTargetFullAccST defRowMajor scheduler 0 sz initAcc splitAcc $ \ !i !ix !acc -> - ioToPrim (run (f acc ix (gi i))) - F.foldlM g tAcc results + let !sz = size arr + splitAcc _ = pure (initAcc, initAcc) + results <- + withScheduler (getComp arr) $ \scheduler -> + withRunInIO $ \run -> + stToPrim $ + case unsafePrefIndex arr of + PrefIndex gix -> + iterFullAccST defRowMajor scheduler zeroIndex sz initAcc splitAcc $ \ !ix !acc -> + ioToPrim (run (f acc ix (gix ix))) + PrefIndexLinear gi -> + iterTargetFullAccST defRowMajor scheduler 0 sz initAcc splitAcc $ \ !i !ix !acc -> + ioToPrim (run (f acc ix (gi i))) + F.foldlM g tAcc results {-# INLINE ifoldlIO #-} -- | Slice an array into linear row-major vector chunks and apply an action to each of @@ -394,11 +401,13 @@ -- will be combined with a folding function. -- -- @since 1.0.0 -splitReduce :: - (MonadUnliftIO m, Index ix, Source r e) +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 + -> (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 @@ -409,39 +418,49 @@ withRunInIO $ \run -> do splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do loopA_ 0 (< slackStart) (+ chunkLength) $ \ !start -> - scheduleWork scheduler $ run $ f scheduler $ - unsafeLinearSlice start (SafeSz chunkLength) arr + scheduleWork scheduler $ + run $ + f scheduler $ + unsafeLinearSlice start (SafeSz chunkLength) arr when (slackStart < totalLength) $ - scheduleWork scheduler $ run $ f scheduler $ - unsafeLinearSlice slackStart (SafeSz (totalLength - slackStart)) arr + 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, Index ix, Source r e) => - (ix -> e -> a -> m a) -> a -> (a -> b -> m b) -> b -> Array r ix e -> m b +ifoldrIO + :: (MonadUnliftIO m, Index ix, Source r e) + => (ix -> e -> a -> m a) + -> a + -> (a -> b -> m b) + -> b + -> Array r ix e + -> m b ifoldrIO f !initAcc g !tAcc !arr | getComp arr == Seq = ifoldrM f initAcc arr >>= (`g` tAcc) | otherwise = do - let !sz = size arr - !totalLength = totalElem sz - results <- - withRunInIO $ \run -> do - withScheduler (getComp arr) $ \ scheduler -> - splitLinearly (numWorkers scheduler) totalLength $ \ chunkLength slackStart -> do - when (slackStart < totalLength) $ - scheduleWork scheduler $ run $ - iterLinearM sz (totalLength - 1) slackStart (-1) (>=) initAcc $ \ !i ix -> - f ix (unsafeLinearIndex arr i) - loopA_ slackStart (> 0) (subtract chunkLength) $ \ !start -> - scheduleWork scheduler $ run $ - iterLinearM sz (start - 1) (start - chunkLength) (-1) (>=) initAcc $ \ !i ix -> - f ix (unsafeLinearIndex arr i) - F.foldlM (flip g) tAcc results + let !sz = size arr + !totalLength = totalElem sz + results <- + withRunInIO $ \run -> do + withScheduler (getComp arr) $ \scheduler -> + splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do + when (slackStart < totalLength) $ + scheduleWork scheduler $ + run $ + iterLinearM sz (totalLength - 1) slackStart (-1) (>=) initAcc $ \ !i ix -> + f ix (unsafeLinearIndex arr i) + loopA_ slackStart (> 0) (subtract chunkLength) $ \ !start -> + scheduleWork scheduler $ + run $ + iterLinearM sz (start - 1) (start - chunkLength) (-1) (>=) initAcc $ \ !i ix -> + f ix (unsafeLinearIndex arr i) + F.foldlM (flip g) tAcc results {-# INLINE ifoldrIO #-} -- | Sequential implementation of `any` with unrolling @@ -452,19 +471,18 @@ !k4 = k - (k `rem` 4) go !i | i < k4 = - f (unsafeLinearIndex arr i ) || - f (unsafeLinearIndex arr (i + 1)) || - f (unsafeLinearIndex arr (i + 2)) || - f (unsafeLinearIndex arr (i + 3)) || - go (i + 4) + f (unsafeLinearIndex arr i) + || f (unsafeLinearIndex arr (i + 1)) + || f (unsafeLinearIndex arr (i + 2)) + || f (unsafeLinearIndex arr (i + 3)) + || go (i + 4) | i < k = f (unsafeLinearIndex arr i) || go (i + 1) | otherwise = False {-# INLINE anySu #-} - -- | Implementaton of `any` on a slice of an array with short-circuiting using batch cancellation. -anySliceSuM :: - (Index ix, Source r e) +anySliceSuM + :: (Index ix, Source r e) => Batch RealWorld Bool -> Ix1 -> Sz1 @@ -477,27 +495,25 @@ !k4 = ix0 + (k' - (k' `rem` 4)) go !i | i < k4 = do - let r = - f (unsafeLinearIndex arr i) || - f (unsafeLinearIndex arr (i + 1)) || - f (unsafeLinearIndex arr (i + 2)) || - f (unsafeLinearIndex arr (i + 3)) - in if r - then cancelBatchWith batch True - else do - done <- hasBatchFinished batch - if done - then pure True - else go (i + 4) + let r = + f (unsafeLinearIndex arr i) + || f (unsafeLinearIndex arr (i + 1)) + || f (unsafeLinearIndex arr (i + 2)) + || f (unsafeLinearIndex arr (i + 3)) + in if r + then cancelBatchWith batch True + else do + done <- hasBatchFinished batch + if done + then pure True + else go (i + 4) | i < k = - if f (unsafeLinearIndex arr i) - then cancelBatchWith batch True - else go (i + 1) + if f (unsafeLinearIndex arr i) + then cancelBatchWith batch True + else go (i + 1) | otherwise = pure False {-# INLINE anySliceSuM #-} - - -- | Parallelizable implementation of `any` with unrolling anyPu :: (Index ix, Source r e) => (e -> Bool) -> Array r ix e -> IO Bool -- TODO: switch to splitReduce @@ -513,12 +529,11 @@ loopA_ 0 (< slackStart) (+ chunkLength) $ \ !start -> scheduleWork scheduler $ anySliceSuM batch start (Sz (start + chunkLength)) f arr when (slackStart < totalLength) $ - scheduleWork scheduler $ anySliceSuM batch slackStart (Sz totalLength) f arr + scheduleWork scheduler $ + anySliceSuM batch slackStart (Sz totalLength) f arr pure $ F.foldl' (||) False results {-# INLINE anyPu #-} - - -- | /O(n)/ - Determines whether any element of the array satisfies a predicate. -- -- @since 0.1.0 @@ -526,5 +541,5 @@ any f arr = case getComp arr of Seq -> anySu f arr - _ -> unsafePerformIO $ anyPu f arr + _ -> unsafePerformIO $ anyPu f arr {-# INLINE any #-} diff --git a/src/Data/Massiv/Array/Ops/Map.hs b/src/Data/Massiv/Array/Ops/Map.hs --- a/src/Data/Massiv/Array/Ops/Map.hs +++ b/src/Data/Massiv/Array/Ops/Map.hs @@ -1,8 +1,9 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MonoLocalBinds #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE MonoLocalBinds #-} + -- | -- Module : Data.Massiv.Array.Ops.Map -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -10,81 +11,99 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Ops.Map - ( map - , imap +module Data.Massiv.Array.Ops.Map ( + map, + imap, + -- ** Traversing + -- *** Applicative - , traverseA - , traverseA_ - , itraverseA - , itraverseA_ - , sequenceA - , sequenceA_ + traverseA, + traverseA_, + itraverseA, + itraverseA_, + sequenceA, + sequenceA_, + -- *** PrimMonad - , traversePrim - , itraversePrim + traversePrim, + itraversePrim, + -- ** Monadic mapping + -- *** Sequential - , mapM - , forM - , imapM - , iforM - , mapM_ - , forM_ - , imapM_ - , iforM_ + mapM, + forM, + imapM, + iforM, + mapM_, + forM_, + imapM_, + iforM_, + -- *** Parallelizable - , mapIO - , mapWS - , mapIO_ - , imapIO - , imapWS - , imapIO_ - , forIO - , forWS - , forIO_ - , iforIO - , iforWS - , iforIO_ - , imapSchedulerM_ - , iforSchedulerM_ - , iterArrayLinearM_ - , iterArrayLinearWithSetM_ - , iterArrayLinearWithStrideM_ + mapIO, + mapWS, + mapIO_, + imapIO, + imapWS, + imapIO_, + forIO, + forWS, + forIO_, + iforIO, + iforWS, + iforIO_, + imapSchedulerM_, + iforSchedulerM_, + iterArrayLinearM_, + iterArrayLinearWithSetM_, + iterArrayLinearWithStrideM_, + -- ** Zipping - , zip - , zip3 - , zip4 - , unzip - , unzip3 - , unzip4 - , zipWith - , zipWith3 - , zipWith4 - , izipWith - , izipWith3 - , izipWith4 + zip, + zip3, + zip4, + unzip, + unzip3, + unzip4, + zipWith, + zipWith3, + zipWith4, + izipWith, + izipWith3, + izipWith4, + -- *** Applicative - , zipWithA - , izipWithA - , zipWith3A - , izipWith3A - ) where + zipWithA, + izipWithA, + zipWith3A, + izipWith3A, +) where -import Data.Traversable (traverse) -import Data.Massiv.Array.Manifest.List import Control.Monad (void) import Control.Monad.Primitive import Control.Scheduler import Data.Coerce import Data.Massiv.Array.Delayed.Pull +import Data.Massiv.Array.Manifest.List import Data.Massiv.Array.Mutable import Data.Massiv.Array.Ops.Construct (makeArrayA, makeArrayLinearA) import Data.Massiv.Core.Common -import Prelude hiding (map, mapM, mapM_, sequenceA, traverse, unzip, unzip3, - zip, zip3, zipWith, zipWith3) +import Data.Traversable (traverse) +import Prelude hiding ( + map, + mapM, + mapM_, + sequenceA, + traverse, + unzip, + unzip3, + zip, + zip3, + zipWith, + zipWith3, + ) -------------------------------------------------------------------------------- -- map ------------------------------------------------------------------------- @@ -97,7 +116,6 @@ map f = fmap f . delay {-# INLINE map #-} - -------------------------------------------------------------------------------- -- zip ------------------------------------------------------------------------- -------------------------------------------------------------------------------- @@ -105,24 +123,31 @@ -- | Zip two arrays -- -- @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 + :: (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 -- -- @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 + :: (Index ix, Source r1 e1, Source r2 e2, Source r3 e3) + => Array r1 ix e1 + -> Array r2 ix e2 + -> Array r3 ix e3 + -> Array D ix (e1, e2, e3) zip3 = zipWith3 (,,) {-# INLINE zip3 #-} -- | Zip four arrays -- -- @since 0.5.4 -zip4 :: - (Index ix, Source r1 e1, Source r2 e2, Source r3 e3, Source r4 e4) +zip4 + :: (Index ix, Source r1 e1, Source r2 e2, Source r3 e3, Source r4 e4) => Array r1 ix e1 -> Array r2 ix e2 -> Array r3 ix e3 @@ -141,21 +166,26 @@ -- | Unzip three arrays -- -- @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) +unzip3 + :: (Index ix, Source r (e1, e2, e3)) + => Array r ix (e1, e2, e3) + -> (Array D ix e1, Array D ix e2, Array D ix e3) +unzip3 arr = (map (\(e, _, _) -> e) arr, map (\(_, e, _) -> e) arr, map (\(_, _, e) -> e) arr) {-# INLINE unzip3 #-} -- | Unzip four arrays -- -- @since 0.5.4 -unzip4 :: (Index ix, Source r (e1, e2, e3, e4)) - => Array r ix (e1, e2, e3, e4) -> (Array D ix e1, Array D ix e2, Array D ix e3, Array D ix e4) +unzip4 + :: (Index ix, Source r (e1, e2, e3, e4)) + => Array r ix (e1, e2, e3, e4) + -> (Array D ix e1, Array D ix e2, Array D ix e3, Array D ix e4) unzip4 arr = ( map (\(e, _, _, _) -> e) arr , map (\(_, e, _, _) -> e) arr , map (\(_, _, e, _) -> e) arr - , map (\(_, _, _, e) -> e) arr) + , map (\(_, _, _, e) -> e) arr + ) {-# INLINE unzip4 #-} -------------------------------------------------------------------------------- @@ -164,25 +194,32 @@ -- | Zip two arrays with a function. Resulting array will be an intersection of -- source arrays in case their dimensions do not match. -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 + :: (Index ix, Source r1 e1, Source r2 e2) + => (e1 -> e2 -> e) + -> Array r1 ix e1 + -> Array r2 ix e2 + -> Array D ix e zipWith f arr1 arr2 = DArray comp sz prefIndex where sz = SafeSz (liftIndex2 min (coerce (size arr1)) (coerce (size arr2))) comp = getComp arr1 <> getComp arr2 prefIndex = PrefIndex (\ix -> f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)) - -- Somehow checking for size equality destroys performance - -- | PrefIndexLinear gi1 <- unsafePrefIndex arr1, - -- PrefIndexLinear gi2 <- unsafePrefIndex arr2, - -- size arr1 == size arr2 = - -- PrefIndexLinear (\i -> f (gi1 i) (gi2 i)) - -- | otherwise = PrefIndex (\ix -> f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)) +-- Somehow checking for size equality destroys performance +-- | PrefIndexLinear gi1 <- unsafePrefIndex arr1, +-- PrefIndexLinear gi2 <- unsafePrefIndex arr2, +-- size arr1 == size arr2 = +-- PrefIndexLinear (\i -> f (gi1 i) (gi2 i)) +-- | otherwise = PrefIndex (\ix -> f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)) {-# INLINE zipWith #-} - -- | Just like `zipWith`, except with an index aware function. -izipWith :: (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 + :: (Index ix, Source r1 e1, Source r2 e2) + => (ix -> e1 -> e2 -> e) + -> Array r1 ix e1 + -> Array r2 ix e2 + -> Array D ix e izipWith f arr1 arr2 = DArray (getComp arr1 <> getComp arr2) @@ -190,29 +227,27 @@ (PrefIndex (\ix -> f ix (unsafeIndex arr1 ix) (unsafeIndex arr2 ix))) {-# INLINE izipWith #-} - -- | Just like `zipWith`, except zip three arrays with a function. -zipWith3 :: - (Index ix, Source r1 e1, Source r2 e2, Source r3 e3) +zipWith3 + :: (Index ix, Source r1 e1, Source r2 e2, Source r3 e3) => (e1 -> e2 -> e3 -> e) -> Array r1 ix e1 -> Array r2 ix e2 -> Array r3 ix e3 -> Array D ix e zipWith3 f arr1 arr2 arr3 = izipWith3 (\_ e1 e2 e3 -> f e1 e2 e3) arr1 arr2 arr3 - -- See note on zipWith - -- | sz1 == size arr2 && sz1 == size arr3 - -- , PrefIndexLinear gi1 <- unsafePrefIndex arr1 - -- , PrefIndexLinear gi2 <- unsafePrefIndex arr2 - -- , PrefIndexLinear gi3 <- unsafePrefIndex arr3 = - -- makeArrayLinear comp sz1 (\ !i -> f (gi1 i) (gi2 i) (gi3 i)) - -- | otherwise = izipWith3 (\_ e1 e2 e3 -> f e1 e2 e3) arr1 arr2 arr3 - -- where - -- comp = getComp arr1 <> getComp arr2 <> getComp arr3 - -- sz1 = size arr1 +-- See note on zipWith +-- | sz1 == size arr2 && sz1 == size arr3 +-- , PrefIndexLinear gi1 <- unsafePrefIndex arr1 +-- , PrefIndexLinear gi2 <- unsafePrefIndex arr2 +-- , PrefIndexLinear gi3 <- unsafePrefIndex arr3 = +-- makeArrayLinear comp sz1 (\ !i -> f (gi1 i) (gi2 i) (gi3 i)) +-- | otherwise = izipWith3 (\_ e1 e2 e3 -> f e1 e2 e3) arr1 arr2 arr3 +-- where +-- comp = getComp arr1 <> getComp arr2 <> getComp arr3 +-- sz1 = size arr1 {-# INLINE zipWith3 #-} - -- | Just like `zipWith3`, except with an index aware function. izipWith3 :: (Index ix, Source r1 e1, Source r2 e2, Source r3 e3) @@ -224,21 +259,21 @@ izipWith3 f arr1 arr2 arr3 = DArray (getComp arr1 <> getComp arr2 <> getComp arr3) - (SafeSz - (liftIndex2 - min - (liftIndex2 min (coerce (size arr1)) (coerce (size arr2))) - (coerce (size arr3)))) + ( SafeSz + ( liftIndex2 + min + (liftIndex2 min (coerce (size arr1)) (coerce (size arr2))) + (coerce (size arr3)) + ) + ) (PrefIndex $ \ !ix -> f ix (unsafeIndex arr1 ix) (unsafeIndex arr2 ix) (unsafeIndex arr3 ix)) {-# INLINE izipWith3 #-} - - -- | Just like `zipWith`, except zip four arrays with a function. -- -- @since 0.5.4 -zipWith4 :: - (Index ix, Source r1 e1, Source r2 e2, Source r3 e3, Source r4 e4) +zipWith4 + :: (Index ix, Source r1 e1, Source r2 e2, Source r3 e3, Source r4 e4) => (e1 -> e2 -> e3 -> e4 -> e) -> Array r1 ix e1 -> Array r2 ix e2 @@ -246,21 +281,20 @@ -> Array r4 ix e4 -> Array D ix e zipWith4 f arr1 arr2 arr3 arr4 = - izipWith4 (\ _ e1 e2 e3 e4 -> f e1 e2 e3 e4) arr1 arr2 arr3 arr4 - -- See note on zipWith - -- | sz1 == size arr2 && sz1 == size arr3 && sz1 == size arr4 - -- , PrefIndexLinear gi1 <- unsafePrefIndex arr1 - -- , PrefIndexLinear gi2 <- unsafePrefIndex arr2 - -- , PrefIndexLinear gi3 <- unsafePrefIndex arr3 - -- , PrefIndexLinear gi4 <- unsafePrefIndex arr4 = - -- makeArrayLinear comp sz1 (\ !i -> f (gi1 i) (gi2 i) (gi3 i) (gi4 i)) - -- | otherwise = izipWith4 (\ _ e1 e2 e3 e4 -> f e1 e2 e3 e4) arr1 arr2 arr3 arr4 - -- where - -- comp = getComp arr1 <> getComp arr2 <> getComp arr3 <> getComp arr4 - -- sz1 = size arr1 + izipWith4 (\_ e1 e2 e3 e4 -> f e1 e2 e3 e4) arr1 arr2 arr3 arr4 +-- See note on zipWith +-- | sz1 == size arr2 && sz1 == size arr3 && sz1 == size arr4 +-- , PrefIndexLinear gi1 <- unsafePrefIndex arr1 +-- , PrefIndexLinear gi2 <- unsafePrefIndex arr2 +-- , PrefIndexLinear gi3 <- unsafePrefIndex arr3 +-- , PrefIndexLinear gi4 <- unsafePrefIndex arr4 = +-- makeArrayLinear comp sz1 (\ !i -> f (gi1 i) (gi2 i) (gi3 i) (gi4 i)) +-- | otherwise = izipWith4 (\ _ e1 e2 e3 e4 -> f e1 e2 e3 e4) arr1 arr2 arr3 arr4 +-- where +-- comp = getComp arr1 <> getComp arr2 <> getComp arr3 <> getComp arr4 +-- sz1 = size arr1 {-# INLINE zipWith4 #-} - -- | Just like `zipWith4`, except with an index aware function. -- -- @since 0.5.4 @@ -275,61 +309,64 @@ izipWith4 f arr1 arr2 arr3 arr4 = makeArray (getComp arr1 <> getComp arr2 <> getComp arr3 <> getComp arr4) - (SafeSz - (liftIndex2 - min - (liftIndex2 - min - (liftIndex2 min (coerce (size arr1)) (coerce (size arr2))) - (coerce (size arr3))) - (coerce (size arr4)))) - (\ !ix -> - f ix (unsafeIndex arr1 ix) (unsafeIndex arr2 ix) (unsafeIndex arr3 ix) (unsafeIndex arr4 ix)) + ( SafeSz + ( liftIndex2 + min + ( liftIndex2 + min + (liftIndex2 min (coerce (size arr1)) (coerce (size arr2))) + (coerce (size arr3)) + ) + (coerce (size arr4)) + ) + ) + ( \ !ix -> + f ix (unsafeIndex arr1 ix) (unsafeIndex arr2 ix) (unsafeIndex arr3 ix) (unsafeIndex arr4 ix) + ) {-# INLINE izipWith4 #-} - -- | Similar to `zipWith`, except does it sequentially and using the `Applicative`. Note that -- resulting array has Manifest representation. -- -- @since 0.3.0 -zipWithA :: - (Source r1 e1, Source r2 e2, Applicative f, Manifest r e, Index ix) => - (e1 -> e2 -> f e) -> - Array r1 ix e1 -> - Array r2 ix e2 -> - f (Array r ix e) +zipWithA + :: (Source r1 e1, Source r2 e2, Applicative f, Manifest r e, Index ix) + => (e1 -> e2 -> f e) + -> Array r1 ix e1 + -> Array r2 ix e2 + -> f (Array r ix e) zipWithA f arr1 arr2 | sz1 == size arr2 , PrefIndexLinear gi1 <- unsafePrefIndex arr1 , PrefIndexLinear gi2 <- unsafePrefIndex arr2 = - setComp (getComp arr1 <> getComp arr2) <$> makeArrayLinearA sz1 (\ !i -> f (gi1 i) (gi2 i)) + setComp (getComp arr1 <> getComp arr2) <$> makeArrayLinearA sz1 (\ !i -> f (gi1 i) (gi2 i)) | otherwise = izipWithA (const f) arr1 arr2 where !sz1 = size arr1 {-# INLINE zipWithA #-} --- | Similar to `zipWith`, except does it sequentiall and using the `Applicative`. Note that +-- | Similar to `zipWith`, except does it sequentially and using the `Applicative`. Note that -- resulting array has Manifest representation. -- -- @since 0.3.0 -izipWithA :: - (Source r1 e1, Source r2 e2, Applicative f, Manifest r e, Index ix) +izipWithA + :: (Source r1 e1, Source r2 e2, Applicative f, Manifest r e, Index ix) => (ix -> e1 -> e2 -> f e) -> Array r1 ix e1 -> Array r2 ix e2 -> f (Array r ix e) izipWithA f arr1 arr2 = - setComp (getComp arr1 <> getComp arr2) <$> - makeArrayA - (SafeSz (liftIndex2 min (coerce (size arr1)) (coerce (size arr2)))) - (\ !ix -> f ix (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)) + setComp (getComp arr1 <> getComp arr2) + <$> makeArrayA + (SafeSz (liftIndex2 min (coerce (size arr1)) (coerce (size arr2)))) + (\ !ix -> f ix (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)) {-# INLINE izipWithA #-} -- | Same as `zipWithA`, but for three arrays. -- -- @since 0.3.0 -zipWith3A :: - (Source r1 e1, Source r2 e2, Source r3 e3, Applicative f, Manifest r e, Index ix) +zipWith3A + :: (Source r1 e1, Source r2 e2, Source r3 e3, Applicative f, Manifest r e, Index ix) => (e1 -> e2 -> e3 -> f e) -> Array r1 ix e1 -> Array r2 ix e2 @@ -341,35 +378,35 @@ -- | Same as `izipWithA`, but for three arrays. -- -- @since 0.3.0 -izipWith3A :: - (Source r1 e1, Source r2 e2, Source r3 e3, Applicative f, Manifest r e, Index ix) +izipWith3A + :: (Source r1 e1, Source r2 e2, Source r3 e3, Applicative f, Manifest r e, Index ix) => (ix -> e1 -> e2 -> e3 -> f e) -> Array r1 ix e1 -> Array r2 ix e2 -> Array r3 ix e3 -> f (Array r ix e) izipWith3A f arr1 arr2 arr3 = - setComp (getComp arr1 <> getComp arr2 <> getComp arr3) <$> - makeArrayA sz (\ !ix -> f ix (unsafeIndex arr1 ix) (unsafeIndex arr2 ix) (unsafeIndex arr3 ix)) + setComp (getComp arr1 <> getComp arr2 <> getComp arr3) + <$> makeArrayA sz (\ !ix -> f ix (unsafeIndex arr1 ix) (unsafeIndex arr2 ix) (unsafeIndex arr3 ix)) where sz = SafeSz $ - liftIndex2 min (liftIndex2 min (coerce (size arr1)) (coerce (size arr2))) (coerce (size arr3)) + liftIndex2 min (liftIndex2 min (coerce (size arr1)) (coerce (size arr2))) (coerce (size arr3)) {-# INLINE izipWith3A #-} - -------------------------------------------------------------------------------- -- traverse -------------------------------------------------------------------- -------------------------------------------------------------------------------- -- | Traverse with an `Applicative` action over an array sequentially. -- --- /Note/ - using `traversePrim` will always be faster, althought not always possible. +-- /Note/ - using `traversePrim` instead will always be significantly faster, roughly +-- about 30 times faster in practice. -- -- @since 0.2.6 --- -traverseA :: - forall r ix e r' a f . (Source r' a, Manifest r e, Index ix, Applicative f) +traverseA + :: forall r ix e r' a f + . (Source r' a, Manifest r e, Index ix, Applicative f) => (a -> f e) -> Array r' ix a -> f (Array r ix e) @@ -380,9 +417,9 @@ -- | Traverse sequentially over a source array, while discarding the result. -- -- @since 0.3.0 --- -traverseA_ :: - forall r ix e a f. (Index ix, Source r e, Applicative f) +traverseA_ + :: forall r ix e a f + . (Index ix, Source r e, Applicative f) => (e -> f a) -> Array r ix e -> f () @@ -397,9 +434,9 @@ -- | Sequence actions in a source array. -- -- @since 0.3.0 --- -sequenceA :: - forall r ix e r' f. (Source r' (f e), Manifest r e, Index ix, Applicative f) +sequenceA + :: forall r ix e r' f + . (Source r' (f e), Manifest r e, Index ix, Applicative f) => Array r' ix (f e) -> f (Array r ix e) sequenceA = traverseA id @@ -408,21 +445,20 @@ -- | Sequence actions in a source array, while discarding the result. -- -- @since 0.3.0 --- -sequenceA_ :: - forall r ix e f. (Index ix, Source r (f e), Applicative f) +sequenceA_ + :: forall r ix e f + . (Index ix, Source r (f e), Applicative f) => Array r ix (f e) -> f () sequenceA_ = traverseA_ id {-# INLINE sequenceA_ #-} - -- | Traverse with an `Applicative` index aware action over an array sequentially. -- -- @since 0.2.6 --- -itraverseA :: - forall r ix e r' a f . (Source r' a, Manifest r e, Index ix, Applicative f) +itraverseA + :: forall r ix e r' a f + . (Source r' a, Manifest r e, Index ix, Applicative f) => (ix -> a -> f e) -> Array r' ix a -> f (Array r ix e) @@ -430,13 +466,12 @@ setComp (getComp arr) <$> makeArrayA (size arr) (\ !ix -> f ix (unsafeIndex arr ix)) {-# INLINE itraverseA #-} - -- | Traverse with an `Applicative` index aware action over an array sequentially. -- -- @since 0.2.6 --- -itraverseA_ :: - forall r ix e a f. (Source r a, Index ix, Applicative f) +itraverseA_ + :: forall r ix e a f + . (Source r a, Index ix, Applicative f) => (ix -> a -> f e) -> Array r ix a -> f () @@ -450,13 +485,12 @@ sz = size arr {-# INLINE itraverseA_ #-} - -- | Traverse sequentially within `PrimMonad` over an array with an action. -- -- @since 0.3.0 --- -traversePrim :: - forall r ix b r' a m . (Source r' a, Manifest r b, Index ix, PrimMonad m) +traversePrim + :: forall r ix b r' a m + . (Source r' a, Manifest r b, Index ix, PrimMonad m) => (a -> m b) -> Array r' ix a -> m (Array r ix b) @@ -476,9 +510,9 @@ -- | Same as `traversePrim`, but traverse with index aware action. -- -- @since 0.3.0 --- -itraversePrim :: - forall r ix b r' a m . (Source r' a, Manifest r b, Index ix, PrimMonad m) +itraversePrim + :: forall r ix b r' a m + . (Source r' a, Manifest r b, Index ix, PrimMonad m) => (ix -> a -> m b) -> Array r' ix a -> m (Array r ix b) @@ -502,51 +536,53 @@ -- | Map a monadic action over an array sequentially. -- -- @since 0.2.6 -mapM :: - forall r ix b r' a m. (Source r' a, Manifest r b, Index ix, Monad m) - => (a -> m b) -- ^ Mapping action - -> Array r' ix a -- ^ Source array +mapM + :: forall r ix b r' a m + . (Source r' a, Manifest r b, Index ix, Monad m) + => (a -> m b) + -- ^ Mapping action + -> Array r' ix a + -- ^ Source array -> m (Array r ix b) mapM = traverseA {-# INLINE mapM #-} - -- | Same as `mapM` except with arguments flipped. -- -- @since 0.2.6 -forM :: - forall r ix b r' a m. (Source r' a, Manifest r b, Index ix, Monad m) +forM + :: forall r ix b r' a m + . (Source r' a, Manifest r b, Index ix, Monad m) => Array r' ix a -> (a -> m b) -> m (Array r ix b) forM = flip traverseA {-# INLINE forM #-} - -- | Map an index aware monadic action over an array sequentially. -- -- @since 0.2.6 -imapM :: - forall r ix b r' a m. (Source r' a, Manifest r b, Index ix, Monad m) +imapM + :: forall r ix b r' a m + . (Source r' a, Manifest r b, Index ix, Monad m) => (ix -> a -> m b) -> Array r' ix a -> m (Array r ix b) imapM = itraverseA {-# INLINE imapM #-} - -- | Same as `forM`, except with an index aware action. -- -- @since 0.5.1 -iforM :: - forall r ix b r' a m. (Source r' a, Manifest r b, Index ix, Monad m) +iforM + :: forall r ix b r' a m + . (Source r' a, Manifest r b, Index ix, Monad m) => Array r' ix a -> (ix -> a -> m b) -> m (Array r ix b) iforM = flip itraverseA {-# INLINE iforM #-} - -- | Map a monadic function over an array sequentially, while discarding the result. -- -- ==== __Examples__ @@ -564,7 +600,6 @@ mapM_ = traverseA_ {-# INLINE mapM_ #-} - -- | Just like `mapM_`, except with flipped arguments. -- -- ==== __Examples__ @@ -578,12 +613,10 @@ -- >>> A.forM_ (range Seq (Ix1 0) 1000) $ \ i -> modifyIORef' ref (+i) -- >>> readIORef ref -- 499500 --- forM_ :: (Source r a, Index ix, Monad m) => Array r ix a -> (a -> m b) -> m () forM_ = flip traverseA_ {-# INLINE forM_ #-} - -- | Map a monadic index aware function over an array sequentially, while discarding the result. -- -- ==== __Examples__ @@ -601,20 +634,19 @@ imapM_ = itraverseA_ {-# INLINE imapM_ #-} - -- | Just like `imapM_`, except with flipped arguments. iforM_ :: (Source r a, Index ix, Monad m) => Array r ix a -> (ix -> a -> m b) -> m () iforM_ = flip itraverseA_ {-# INLINE iforM_ #-} - -- | Map an `IO` action over an `Array`. Underlying computation strategy is respected and will be -- parallelized when requested. Unfortunately no fusion is possible and new array will be create -- upon each call. -- -- @since 0.2.6 -mapIO :: - forall r ix b r' a m. (Size r', Load r' ix a, Manifest r b, MonadUnliftIO m) +mapIO + :: forall r ix b r' a m + . (Size r', Load r' ix a, Manifest r b, MonadUnliftIO m) => (a -> m b) -> Array r' ix a -> m (Array r ix b) @@ -627,8 +659,9 @@ -- logic outlined by the supplied array. -- -- @since 0.2.6 -mapIO_ :: - forall r ix e a m. (Load r ix e, MonadUnliftIO m) +mapIO_ + :: forall r ix e a m + . (Load r ix e, MonadUnliftIO m) => (e -> m a) -> Array r ix e -> m () @@ -641,8 +674,9 @@ -- | Same as `mapIO_`, but map an index aware action instead. -- -- @since 0.2.6 -imapIO_ :: - forall r ix e a m. (Load r ix e, MonadUnliftIO m) +imapIO_ + :: forall r ix e a m + . (Load r ix e, MonadUnliftIO m) => (ix -> e -> m a) -> Array r ix e -> m () @@ -650,18 +684,18 @@ 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, + in -- It is ok to use outerSize in context of DS and L. Former is 1-dim, -- so sz is never evaluated and for the latter outerSize has to be -- called regardless how this function is implemented. - in iterArrayLinearM_ scheduler arr (\i -> void . run . action (fromLinearIndex sz i)) + iterArrayLinearM_ scheduler arr (\i -> void . run . action (fromLinearIndex sz i)) {-# INLINE imapIO_ #-} - -- | 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. (Size r', Load r' ix a, Manifest r b, MonadUnliftIO m) +imapIO + :: forall r ix b r' a m + . (Size r', Load r' ix a, Manifest r b, MonadUnliftIO m) => (ix -> a -> m b) -> Array r' ix a -> m (Array r ix b) @@ -678,23 +712,23 @@ -- | Same as `mapIO` but with arguments flipped. -- -- @since 0.2.6 -forIO :: - forall r ix b r' a m. (Size r', Load r' ix a, Manifest r b, MonadUnliftIO m) +forIO + :: forall r ix b r' a m + . (Size r', Load r' ix a, Manifest r b, MonadUnliftIO m) => Array r' ix a -> (a -> m b) -> m (Array r ix b) forIO = flip mapIO {-# INLINE forIO #-} - - -- | Same as `imapIO`, but ignores the inner computation strategy and uses -- stateful workers during computation instead. Use -- `Control.Scheduler.initWorkerStates` for the `WorkerStates` initialization. -- -- @since 0.3.4 -imapWS :: - forall r ix b r' a s m. (Source r' a, Manifest r b, Index ix, MonadUnliftIO m, PrimMonad m) +imapWS + :: forall r ix b r' a s m + . (Source r' a, Manifest r b, Index ix, MonadUnliftIO m, PrimMonad m) => WorkerStates s -> (ix -> a -> s -> m b) -> Array r' ix a @@ -705,21 +739,22 @@ -- | Same as `imapWS`, but without the index. -- -- @since 0.3.4 -mapWS :: - forall r ix b r' a s m. (Source r' a, Manifest r b, Index ix, MonadUnliftIO m, PrimMonad m) +mapWS + :: forall r ix b r' a s m + . (Source r' a, Manifest r b, Index ix, MonadUnliftIO m, PrimMonad m) => WorkerStates s -> (a -> s -> m b) -> Array r' ix a -> m (Array r ix b) -mapWS states f = imapWS states (\ _ -> f) +mapWS states f = imapWS states (\_ -> f) {-# INLINE mapWS #-} - -- | Same as `imapWS`, but with source array and mapping action arguments flipped. -- -- @since 0.3.4 -iforWS :: - forall r ix b r' a s m. (Source r' a, Manifest r b, Index ix, MonadUnliftIO m, PrimMonad m) +iforWS + :: forall r ix b r' a s m + . (Source r' a, Manifest r b, Index ix, MonadUnliftIO m, PrimMonad m) => WorkerStates s -> Array r' ix a -> (ix -> a -> s -> m b) @@ -730,17 +765,16 @@ -- | Same as `iforWS`, but without the index. -- -- @since 0.3.4 -forWS :: - forall r ix b r' a s m. (Source r' a, Manifest r b, Index ix, MonadUnliftIO m, PrimMonad m) +forWS + :: forall r ix b r' a s m + . (Source r' a, Manifest r b, Index ix, MonadUnliftIO m, PrimMonad m) => WorkerStates s -> Array r' ix a -> (a -> s -> m b) -> m (Array r ix b) -forWS states arr f = imapWS states (\ _ -> f) arr +forWS states arr f = imapWS states (\_ -> f) arr {-# INLINE forWS #-} - - -- | Same as `mapIO_` but with arguments flipped. -- -- ==== __Example__ @@ -763,8 +797,9 @@ -- | Same as `imapIO` but with arguments flipped. -- -- @since 0.2.6 -iforIO :: - forall r ix b r' a m. (Size r', Load r' ix a, Manifest r b, MonadUnliftIO m) +iforIO + :: forall r ix b r' a m + . (Size r', Load r' ix a, Manifest r b, MonadUnliftIO m) => Array r' ix a -> (ix -> a -> m b) -> m (Array r ix b) @@ -774,53 +809,62 @@ -- | Same as `imapIO_` but with arguments flipped. -- -- @since 0.2.6 -iforIO_ :: - forall r ix e a m. (Load r ix e, MonadUnliftIO 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) +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 + -> 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) +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 + -> 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) + 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) +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 + -> 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 @@ -830,12 +874,11 @@ -- Deprecated - -- | Same as `imapM_`, but will use the supplied scheduler. -- -- @since 0.3.1 -imapSchedulerM_ :: - (Index ix, Source r e, MonadPrimBase s m) +imapSchedulerM_ + :: (Index ix, Source r e, MonadPrimBase s m) => Scheduler s () -> (ix -> e -> m a) -> Array r ix e @@ -849,12 +892,11 @@ (\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) +iforSchedulerM_ + :: (Index ix, Source r e, MonadPrimBase s m) => Scheduler s () -> Array r ix e -> (ix -> e -> m a) @@ -862,7 +904,6 @@ iforSchedulerM_ scheduler arr action = imapSchedulerM_ scheduler action arr {-# INLINE iforSchedulerM_ #-} - -- -- | Load an array into memory. -- -- -- -- @since 0.3.0 @@ -891,10 +932,10 @@ -- 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 () +-- 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 @@ -2,6 +2,7 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} + -- | -- Module : Data.Massiv.Array.Ops.Slice -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -9,36 +10,35 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Ops.Slice - ( +module Data.Massiv.Array.Ops.Slice ( -- ** From the outside - (!>) - , (!?>) - , (??>) + (!>), + (!?>), + (??>), + -- ** From the inside - , () - , () - , () + (), + (), + (), + -- ** Many slices - , outerSlices - , innerSlices - , withinSlices - , withinSlicesM - ) where + outerSlices, + innerSlices, + withinSlices, + withinSlicesM, +) where import Control.Monad (unless) import Data.Massiv.Array.Delayed.Pull import Data.Massiv.Core.Common - infixl 4 !>, !?>, ??>, , , - -- | /O(1)/ - Slices the array from the outside. For 2-dimensional array this will -- be equivalent of taking a row. Throws an error when index is out of bounds. -- @@ -79,21 +79,22 @@ -- -- -- @since 0.1.0 -(!>) :: - forall r ix e. (HasCallStack, Index ix, Index (Lower ix), Source r e) +(!>) + :: forall r ix e + . (HasCallStack, Index ix, Index (Lower ix), Source r e) => Array r ix e -> Int -> Array r (Lower ix) e (!>) !arr !ix = throwEither (arr !?> ix) {-# INLINE (!>) #-} - -- | /O(1)/ - Just like `!>` slices the array from the outside, but returns -- `Nothing` when index is out of bounds. -- -- @since 0.1.0 -(!?>) :: - forall r ix e m. (MonadThrow m, Index ix, Index (Lower ix), Source r e) +(!?>) + :: 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) @@ -103,7 +104,6 @@ pure $ unsafeOuterSlice arr szL i {-# INLINE (!?>) #-} - -- | /O(1)/ - Safe slicing continuation from the outside. Similarly to (`!>`) slices the array from -- the outside, but takes `Maybe` array as input and returns `Nothing` when index is out of bounds. -- @@ -119,20 +119,21 @@ -- *** Exception: IndexOutOfBoundsException: -10 is not safe for (Sz1 2) -- -- @since 0.1.0 -(??>) :: - forall r ix e m. (MonadThrow m, Index ix, Index (Lower ix), Source r e) +(??>) + :: forall r ix e m + . (MonadThrow m, Index ix, Index (Lower ix), Source r e) => m (Array r ix e) -> Int -> m (Array r (Lower ix) e) (??>) marr !ix = marr >>= (!?> ix) {-# INLINE (??>) #-} - -- | /O(1)/ - Safe slice from the inside -- -- @since 0.1.0 -( Array r ix e -> Int -> m (Array D (Lower ix) e) @@ -142,36 +143,36 @@ pure $ unsafeInnerSlice arr szL i {-# INLINE (`) slice an array from an opposite direction. -- -- @since 0.1.0 -( Array r ix e -> Int -> Array D (Lower 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 -() :: - forall r ix e m. (MonadThrow m, Index ix, Index (Lower ix), Source r 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) @@ -182,9 +183,8 @@ internalInnerSlice dim cutSz arr i {-# INLINE () #-} - -internalInnerSlice :: - (MonadThrow m, Index ix, Index (Lower ix), Source r e) +internalInnerSlice + :: (MonadThrow m, Index ix, Index (Lower ix), Source r e) => Dim -> Sz ix -> Array r ix e @@ -195,29 +195,29 @@ unsafeSlice arr start cutSz dim {-# INLINE internalInnerSlice #-} - - -- prop> arr !> i == arr (dimensions (size arr), i) -- prop> arr (1,i) -- + -- | /O(1)/ - Slices the array in any available dimension. Throws an error when -- index is out of bounds or dimensions is invalid. -- -- @since 0.1.0 -() :: - forall r ix e. (HasCallStack, Index ix, Index (Lower ix), Source r e) +() + :: 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 -() :: - forall r ix e m. (MonadThrow m, Index ix, Index (Lower ix), Source r 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) @@ -238,8 +238,9 @@ -- [ 2 :. 0, 2 :. 1 ] -- -- @since 0.5.4 -outerSlices :: - forall r ix e. (Index ix, Index (Lower ix), Source r e) +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) @@ -247,7 +248,6 @@ (k, szL) = unconsSz $ size arr {-# INLINE outerSlices #-} - -- | Create a delayed array of inner slices. -- -- ====__Examples__ @@ -260,8 +260,9 @@ -- [ 0 :. 1, 1 :. 1, 2 :. 1 ] -- -- @since 0.5.4 -innerSlices :: - forall r ix e. (Index ix, Source r e) +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) @@ -315,23 +316,24 @@ -- ] -- -- @since 0.5.4 -withinSlices :: - forall n r ix e. (IsIndexDimension ix n, Index (Lower ix), Source r e) +withinSlices + :: forall n r ix e + . (IsIndexDimension ix n, Index (Lower ix), Source r e) => Dimension n -> Array r ix e -> Array D Ix1 (Array D (Lower ix) e) withinSlices dim = either throwImpossible id . withinSlicesM (fromDimension dim) {-# INLINE withinSlices #-} - -- | Create a delayed array of slices from within. Same as `withinSlices`, but throws an -- error on invalid dimension. -- -- /__Throws Exceptions__/: `IndexDimensionException` -- -- @since 0.5.4 -withinSlicesM :: - forall r ix e m. (MonadThrow m, Index ix, Index (Lower ix), Source r e) +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)) 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 @@ -2,6 +2,7 @@ {-# LANGUAGE ExplicitForAll #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE MonoLocalBinds #-} + -- | -- Module : Data.Massiv.Array.Ops.Sort -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -9,22 +10,21 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Ops.Sort - ( tally - , quicksort - , quicksortBy - , quicksortByM - , quicksortAs - , quicksortAsBy - , quicksortAsByM - , quicksortM_ - , quicksortByM_ - , unsafeUnstablePartitionRegionM - ) where +module Data.Massiv.Array.Ops.Sort ( + tally, + quicksort, + quicksortBy, + quicksortByM, + quicksortAs, + quicksortAsBy, + quicksortAsByM, + quicksortM_, + quicksortByM_, + unsafeUnstablePartitionRegionM, +) where -import Control.Monad.IO.Unlift import Control.Monad (when) +import Control.Monad.IO.Unlift import Control.Monad.Primitive import Control.Scheduler import Data.Massiv.Array.Delayed.Stream @@ -57,73 +57,72 @@ sz@(Sz k) = size sorted count (!i, !n, !prev) | i < k = - let !e' = unsafeLinearIndex sorted i - in if prev == e' - then Just (Nothing, (i + 1, n + 1, prev)) - else Just (Just (prev, n), (i + 1, 1, e')) + let !e' = unsafeLinearIndex sorted i + in if prev == e' + then Just (Nothing, (i + 1, n + 1, prev)) + else Just (Just (prev, n), (i + 1, 1, e')) | otherwise = Just (Just (prev, n), (i + 1, n, prev)) {-# INLINE count #-} sorted = quicksort $ flatten arr {-# INLINE tally #-} - - -- | Partition a segment of a vector. Starting and ending indices are unchecked. -- -- @since 1.0.0 -unsafeUnstablePartitionRegionM :: - forall r e m. (Manifest r e, PrimMonad m) +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 + -> Ix1 + -- ^ Start index of the region + -> Ix1 + -- ^ End index of the region -> m Ix1 unsafeUnstablePartitionRegionM marr f start end = fromLeft start (end + 1) where fromLeft i j | i == j = pure i | otherwise = do - e <- f =<< unsafeLinearRead marr i - if e - then fromLeft (i + 1) j - else fromRight i (j - 1) + e <- f =<< unsafeLinearRead marr i + if e + then fromLeft (i + 1) j + else fromRight i (j - 1) fromRight i j | i == j = pure i | otherwise = do - x <- unsafeLinearRead marr j - e <- f x - if e - then do - unsafeLinearWrite marr j =<< unsafeLinearRead marr i - unsafeLinearWrite marr i x - fromLeft (i + 1) j - else fromRight i (j - 1) + x <- unsafeLinearRead marr j + e <- f x + if e + then do + unsafeLinearWrite marr j =<< unsafeLinearRead marr i + unsafeLinearWrite marr i x + fromLeft (i + 1) j + else fromRight i (j - 1) {-# INLINE unsafeUnstablePartitionRegionM #-} - -- | Same as `quicksort` except it accepts any array that is computable. -- -- @since 1.0.2 -quicksortAs :: - (Load r Ix1 e, Manifest r' e, Ord e) => r' -> Vector r e -> Vector r' e +quicksortAs + :: (Load r Ix1 e, Manifest r' e, Ord e) => r' -> Vector r e -> Vector r' e quicksortAs _ arr = unsafePerformIO $ withLoadMArray_ arr quicksortM_ {-# INLINE quicksortAs #-} -- | Same as `quicksortBy` except it accepts any array that is computable. -- -- @since 1.0.2 -quicksortAsBy :: - (Load r Ix1 e, Manifest r' e) => r' -> (e -> e -> Ordering) -> Vector r e -> Vector r' e +quicksortAsBy + :: (Load r Ix1 e, Manifest r' e) => r' -> (e -> e -> Ordering) -> Vector r e -> Vector r' e quicksortAsBy _ f arr = unsafePerformIO $ withLoadMArray_ arr (quicksortByM_ (\x y -> pure $ f x y)) {-# INLINE quicksortAsBy #-} - -- | Same as `quicksortByM` except it accepts any array that is computable. -- -- @since 1.0.2 -quicksortAsByM :: - (Load r Ix1 e, Manifest r' e, MonadUnliftIO m) +quicksortAsByM + :: (Load r Ix1 e, Manifest r' e, MonadUnliftIO m) => r' -> (e -> e -> m Ordering) -> Vector r e @@ -132,7 +131,6 @@ withRunInIO $ \run -> withLoadMArray_ arr (quicksortByM_ (\x y -> run (f x y))) {-# INLINE quicksortAsByM #-} - -- | This is an implementation of -- [Quicksort](https://en.wikipedia.org/wiki/Quicksort), which is an efficient, -- but unstable sort. This implementation uses Median-of-three for pivot @@ -142,17 +140,16 @@ -- will result in a nice speed up for systems with multiple CPUs. -- -- @since 0.3.2 -quicksort :: - (Manifest r e, Ord e) => Vector r e -> Vector r e +quicksort + :: (Manifest r e, Ord e) => Vector r e -> Vector r e quicksort arr = unsafePerformIO $ withMArray_ arr quicksortM_ {-# INLINE quicksort #-} - -- | Same as `quicksortBy`, but instead of `Ord` constraint expects a custom `Ordering`. -- -- @since 0.6.1 -quicksortByM :: - (Manifest r e, MonadUnliftIO m) => (e -> e -> m Ordering) -> Vector r e -> m (Vector r e) +quicksortByM + :: (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 #-} @@ -167,20 +164,19 @@ -- | Manifest version of `quicksort` -- -- @since 0.3.2 -quicksortM_ :: - (Ord e, Manifest r e, MonadPrimBase s m) +quicksortM_ + :: (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_ #-} - -- | Same as `quicksortM_`, but instead of `Ord` constraint expects a custom `Ordering`. -- -- @since 0.6.1 -quicksortByM_ :: - (Manifest r e, MonadPrimBase s m) +quicksortByM_ + :: (Manifest r e, MonadPrimBase s m) => (e -> e -> m Ordering) -> Scheduler s () -> MVector s r e @@ -189,9 +185,8 @@ quicksortInternalM_ (\x y -> (LT ==) <$> compareM x y) (\x y -> (EQ ==) <$> compareM x y) {-# INLINE quicksortByM_ #-} - -quicksortInternalM_ :: - (Manifest r e, MonadPrimBase s m) +quicksortInternalM_ + :: (Manifest r e, MonadPrimBase s m) => (e -> e -> m Bool) -> (e -> e -> m Bool) -> Scheduler s () 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 @@ -4,6 +4,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} {-# OPTIONS_GHC -fno-warn-redundant-constraints #-} + -- | -- Module : Data.Massiv.Array.Ops.Transform -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -11,58 +12,65 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Ops.Transform - ( -- ** Transpose - transpose - , transposeInner - , transposeOuter +module Data.Massiv.Array.Ops.Transform ( + -- ** Transpose + transpose, + transposeInner, + transposeOuter, + -- ** Reverse - , reverse - , reverse' - , reverseM + reverse, + reverse', + reverseM, + -- ** Backpermute - , backpermuteM - , backpermute' + backpermuteM, + backpermute', + -- ** Resize - , resizeM - , resize' - , flatten + resizeM, + resize', + flatten, + -- ** Extract - , extractM - , extract' - , extractFromToM - , extractFromTo' - , deleteRowsM - , deleteColumnsM - , deleteRegionM + extractM, + extract', + extractFromToM, + extractFromTo', + deleteRowsM, + deleteColumnsM, + deleteRegionM, + -- ** Append/Split - , appendOuterM - , appendM - , append' - , concatOuterM - , concatM - , concat' - , stackSlicesM - , stackOuterSlicesM - , stackInnerSlicesM - , splitAtM - , splitAt' - , splitExtractM - , replaceSlice - , replaceOuterSlice + appendOuterM, + appendM, + append', + concatOuterM, + concatM, + concat', + stackSlicesM, + stackOuterSlicesM, + stackInnerSlicesM, + splitAtM, + splitAt', + splitExtractM, + replaceSlice, + replaceOuterSlice, + -- ** Upsample/Downsample - , upsample - , downsample + upsample, + downsample, + -- ** Zoom - , zoom - , zoomWithGrid + zoom, + zoomWithGrid, + -- ** Transform - , transformM - , transform' - , transform2M - , transform2' - ) where + transformM, + transform', + transform2M, + transform2', +) where import Control.Monad as M (foldM_, forM_, unless) import Control.Monad.ST @@ -76,11 +84,16 @@ import Data.Massiv.Array.Ops.Construct import Data.Massiv.Array.Ops.Map import Data.Massiv.Core -import Data.Massiv.Core.Index.Internal -import Data.Massiv.Core.Common --(size, unsafeIndex, unsafeResize, evaluate', evaluateM) -import Prelude as P hiding (concat, drop, mapM_, reverse, splitAt, take, - traverse) - +import Data.Massiv.Core.Common -- (size, unsafeIndex, unsafeResize, evaluate', evaluateM) +import Prelude as P hiding ( + concat, + drop, + mapM_, + reverse, + splitAt, + take, + traverse, + ) -- | Extract a sub-array from within a larger source array. Array that is being extracted must be -- fully encapsulated in a source array, otherwise `SizeSubregionException` will be thrown. @@ -124,19 +137,23 @@ -- ] -- -- @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 +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 + pure $ unsafeExtract sIx newSz arr | otherwise = throwM $ SizeSubregionException (size arr) sIx newSz where - sz1 = Sz (liftIndex (+1) (unSz (size arr))) - eIx1 = Sz (liftIndex (+1) eIx) + sz1 = Sz (liftIndex (+ 1) (unSz (size arr))) + eIx1 = Sz (liftIndex (+ 1) eIx) eIx = liftIndex2 (+) sIx $ unSz newSz {-# INLINE extractM #-} @@ -144,16 +161,19 @@ -- are incorrect. -- -- @since 0.1.0 -extract' :: - forall r ix e. (HasCallStack, Index ix, Source r e) - => ix -- ^ Starting index - -> Sz ix -- ^ Size of the resulting array - -> Array r ix e -- ^ Source array +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. -- @@ -183,11 +203,15 @@ -- ] -- -- @since 0.3.0 -extractFromToM :: - forall r ix e m. (MonadThrow m, Index ix, Source r e) - => ix -- ^ Starting index - -> ix -- ^ Index up to which elements should be extracted. - -> Array r ix e -- ^ Source array. +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 #-} @@ -195,23 +219,27 @@ -- | Same as `extractFromToM`, but throws an error on invalid indices. -- -- @since 0.2.4 -extractFromTo' :: - forall r ix e. (HasCallStack, Index ix, Source r e) - => ix -- ^ Starting index - -> ix -- ^ Index up to which elmenets should be extracted. - -> Array r ix e -- ^ Source array. +extractFromTo' + :: forall r ix e + . (HasCallStack, Index ix, Source r e) + => ix + -- ^ Starting index + -> ix + -- ^ Index up to which elmenets should be extracted. + -> Array r ix e + -- ^ Source array. -> Array D ix e extractFromTo' sIx eIx = extract' sIx $ Sz (liftIndex2 (-) eIx sIx) {-# INLINE extractFromTo' #-} - -- | /O(1)/ - Change the size of an array. Throws -- `SizeElementsMismatchException` if total number of elements does not match -- the supplied array. -- -- @since 0.3.0 -resizeM :: - forall r ix ix' e m. (MonadThrow m, Index ix', Index ix, Size r) +resizeM + :: forall r ix ix' e m + . (MonadThrow m, Index ix', Index ix, Size r) => Sz ix' -> Array r ix e -> m (Array r ix' e) @@ -221,8 +249,9 @@ -- | Same as `resizeM`, but will throw an error if supplied dimensions are incorrect. -- -- @since 0.1.0 -resize' :: - forall r ix ix' e. (HasCallStack, Index ix', Index ix, Size r) +resize' + :: forall r ix ix' e + . (HasCallStack, Index ix', Index ix, Size r) => Sz ix' -> Array r ix e -> Array r ix' e @@ -236,7 +265,6 @@ flatten arr = unsafeResize (SafeSz (totalElem (size arr))) arr {-# INLINE flatten #-} - -- | Transpose a 2-dimensional array -- -- ==== __Examples__ @@ -261,11 +289,10 @@ {-# INLINE [1] transpose #-} {-# RULES -"transpose . transpose" [~1] forall arr . transpose (transpose arr) = delay arr -"transposeInner . transposeInner" [~1] forall arr . transposeInner (transposeInner arr) = delay arr -"transposeOuter . transposeOuter" [~1] forall arr . transposeOuter (transposeOuter arr) = delay arr - #-} - +"transpose . transpose" [~1] forall arr. transpose (transpose arr) = delay arr +"transposeInner . transposeInner" [~1] forall arr. transposeInner (transposeInner arr) = delay arr +"transposeOuter . transposeOuter" [~1] forall arr. transposeOuter (transposeOuter arr) = delay arr + #-} -- | Transpose inner two dimensions of at least rank-2 array. -- @@ -298,8 +325,9 @@ -- ] -- -- @since 0.1.0 -transposeInner :: - forall r ix e. (Index (Lower ix), Index ix, Source r 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 @@ -351,8 +379,9 @@ -- -- -- @since 0.1.0 -transposeOuter :: - forall r ix e. (Index (Lower ix), Index ix, Source r 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 @@ -398,8 +427,9 @@ -- ] -- -- @since 0.4.1 -reverse :: - forall n r ix e. (IsIndexDimension ix n, Index ix, Source r 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 @@ -410,15 +440,16 @@ -- `IndexDimensionException` for an incorrect dimension. -- -- @since 0.4.1 -reverseM :: - forall r ix e m. (MonadThrow m, Index ix, Source r 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 - pure $ makeArray (getComp arr) sz $ \ ix -> + pure $ makeArray (getComp arr) sz $ \ix -> unsafeIndex arr (snd $ modifyDim' ix dim (\i -> k - i - 1)) {-# INLINE reverseM #-} @@ -426,8 +457,9 @@ -- `IndexDimensionException` from pure code. -- -- @since 0.4.1 -reverse' :: - forall r ix e. (HasCallStack, Index ix, Source r e) +reverse' + :: forall r ix e + . (HasCallStack, Index ix, Source r e) => Dim -> Array r ix e -> Array D ix e @@ -462,12 +494,15 @@ -- ] -- -- @since 0.3.0 -backpermuteM :: - forall r ix e r' ix' m. - (Manifest r e, Index ix, Source r' e, Index ix', MonadUnliftIO m, PrimMonad m, MonadThrow m) - => Sz ix -- ^ Size of the result array - -> (ix -> ix') -- ^ A function that maps indices of the new array into the source one. - -> Array r' ix' e -- ^ Source array. +backpermuteM + :: forall r ix e r' ix' m + . (Manifest r e, Index ix, Source r' e, Index ix', MonadUnliftIO m, PrimMonad m, MonadThrow m) + => Sz ix + -- ^ Size of the result array + -> (ix -> ix') + -- ^ A function that maps indices of the new array into the source one. + -> Array r' ix' e + -- ^ Source array. -> m (Array r ix e) backpermuteM sz ixF !arr = generateArray (getComp arr) sz (evaluateM arr . ixF) {-# INLINE backpermuteM #-} @@ -479,16 +514,19 @@ -- * Throws a runtime `IndexOutOfBoundsException` from pure code. -- -- @since 0.3.0 -backpermute' :: - forall r ix ix' e. (HasCallStack, Source r e, Index ix, Index ix') - => Sz ix' -- ^ Size of the result array - -> (ix' -> ix) -- ^ A function that maps indices of the new array into the source one. - -> Array r ix e -- ^ Source array. +backpermute' + :: forall r ix ix' e + . (HasCallStack, Source r e, Index ix, Index ix') + => Sz ix' + -- ^ Size of the result array + -> (ix' -> ix) + -- ^ A function that maps indices of the new array into the source one. + -> Array r ix e + -- ^ Source array. -> Array D ix' e backpermute' sz ixF !arr = makeArray (getComp arr) sz (evaluate' arr . ixF) {-# INLINE backpermute' #-} - -- | Append two arrays together along a particular dimension. Sizes of both arrays must match, with -- an allowed exception of the dimension they are being appended along, otherwise `Nothing` is -- returned. @@ -525,8 +563,9 @@ -- *** Exception: SizeMismatchException: (Sz (2 :. 3)) vs (Sz (2 :. 4)) -- -- @since 0.3.0 -appendM :: - forall r1 r2 ix e m. (MonadThrow m, Index ix, Source r1 e, Source r2 e) +appendM + :: forall r1 r2 ix e m + . (MonadThrow m, Index ix, Source r1 e, Source r2 e) => Dim -> Array r1 ix e -> Array r2 ix e @@ -552,15 +591,18 @@ {-# INLINE load #-} return $ DLArray - {dlComp = getComp arr1 <> getComp arr2, dlSize = newSz, dlLoad = load} + { dlComp = getComp arr1 <> getComp arr2 + , dlSize = newSz + , dlLoad = load + } {-# INLINE appendM #-} - -- | Same as `appendM`, but will throw an exception in pure code on mismatched sizes. -- -- @since 0.3.0 -append' :: - forall r1 r2 ix e. (HasCallStack, Index ix, Source r1 e, Source r2 e) +append' + :: forall r1 r2 ix e + . (HasCallStack, Index ix, Source r1 e, Source r2 e) => Dim -> Array r1 ix e -> Array r2 ix e @@ -571,8 +613,9 @@ -- | Concat many arrays together along some dimension. -- -- @since 0.3.0 -concat' :: - forall f r ix e. (HasCallStack, Foldable f, Index ix, Source r e) +concat' + :: forall f r ix e + . (HasCallStack, Foldable f, Index ix, Source r e) => Dim -> f (Array r ix e) -> Array DL ix e @@ -585,8 +628,9 @@ -- /__Exceptions__/: `IndexDimensionException`, `SizeMismatchException` -- -- @since 0.3.0 -concatM :: - forall r ix e f m. (MonadThrow m, Foldable f, Index ix, Source r e) +concatM + :: forall r ix e f m + . (MonadThrow m, Foldable f, Index ix, Source r e) => Dim -> f (Array r ix e) -> m (Array DL ix e) @@ -619,10 +663,9 @@ in M.foldM_ arrayLoader 0 $ (k, a) : P.zip ks arrs {-# INLINE load #-} return $ - DLArray {dlComp = getComp a <> foldMap getComp arrs, dlSize = newSz, dlLoad = load} + DLArray{dlComp = getComp a <> foldMap getComp arrs, dlSize = newSz, dlLoad = load} {-# INLINE concatM #-} - -- | Stack slices on top of each other along the specified dimension. -- -- /__Exceptions__/: `IndexDimensionException`, `SizeMismatchException` @@ -686,8 +729,9 @@ -- ] -- -- @since 0.5.4 -stackSlicesM :: - forall r ix e f m. (Foldable f, MonadThrow m, Index (Lower ix), Source r e, Index ix) +stackSlicesM + :: forall r ix e f m + . (Foldable f, MonadThrow m, Index (Lower ix), Source r e, Index ix) => Dim -> f (Array r (Lower ix) e) -> m (Array DL ix e) @@ -699,7 +743,7 @@ len = SafeSz (F.length arrsF) -- / make sure all arrays have the same size M.forM_ arrsF $ \arr -> - unless (sz == size arr) $ throwM (SizeMismatchException sz (size arr)) + unless (sz == size arr) $ throwM (SizeMismatchException sz (size arr)) newSz <- insertSzM sz dim len let load :: Loader e load scheduler startAt dlWrite _dlSet = @@ -709,7 +753,7 @@ in M.foldM_ arrayLoader 0 arrsF {-# INLINE load #-} return $ - DLArray {dlComp = foldMap getComp arrs, dlSize = newSz, dlLoad = load} + DLArray{dlComp = foldMap getComp arrs, dlSize = newSz, dlLoad = load} {-# INLINE stackSlicesM #-} -- | Specialized `stackSlicesM` to handling stacking from the outside. It is the inverse of @@ -745,8 +789,9 @@ -- ] -- -- @since 0.5.4 -stackOuterSlicesM :: - forall r ix e f m. (Foldable f, MonadThrow m, Index (Lower ix), Source r e, Index ix) +stackOuterSlicesM + :: forall r ix e f m + . (Foldable f, MonadThrow m, Index (Lower ix), Source r e, Index ix) => f (Array r (Lower ix) e) -> m (Array DL ix e) stackOuterSlicesM = stackSlicesM (dimensions (Proxy :: Proxy ix)) @@ -785,14 +830,14 @@ -- ] -- -- @since 0.5.4 -stackInnerSlicesM :: - forall r ix e f m. (Foldable f, MonadThrow m, Index (Lower ix), Source r e, Index ix) +stackInnerSlicesM + :: forall r ix e f m + . (Foldable f, MonadThrow m, Index (Lower ix), Source r e, Index ix) => f (Array r (Lower ix) e) -> m (Array DL ix e) stackInnerSlicesM = stackSlicesM 1 {-# INLINE stackInnerSlicesM #-} - -- | /O(1)/ - Split an array into two at an index along a specified dimension. -- -- /Related/: 'splitAt'', `splitExtractM`, 'Data.Massiv.Vector.sliceAt'', `Data.Massiv.Vector.sliceAtM` @@ -800,11 +845,15 @@ -- /__Exceptions__/: `IndexDimensionException`, `SizeSubregionException` -- -- @since 0.3.0 -splitAtM :: - forall r ix e m. (MonadThrow m, Index ix, Source r e) - => Dim -- ^ Dimension along which to split - -> Int -- ^ Index along the dimension to split at - -> Array r ix e -- ^ Source array +splitAtM + :: forall r ix e m + . (MonadThrow m, Index ix, Source r e) + => Dim + -- ^ Dimension along which to split + -> Int + -- ^ Index along the dimension to split at + -> Array r ix e + -- ^ Source array -> m (Array D ix e, Array D ix e) splitAtM dim i arr = do let Sz sz = size arr @@ -824,8 +873,9 @@ -- -- -- @since 0.1.0 -splitAt' :: - forall r ix e. (HasCallStack, Index ix, Source r e) +splitAt' + :: forall r ix e + . (HasCallStack, Index ix, Source r e) => Dim -> Int -> Array r ix e @@ -833,15 +883,18 @@ splitAt' dim i = throwEither . splitAtM dim i {-# INLINE splitAt' #-} - -- | Split an array in three parts across some dimension -- -- @since 0.3.5 -splitExtractM :: - forall r ix e m. (MonadThrow m, Index ix, Source r e) - => Dim -- ^ Dimension along which to do the extraction - -> Ix1 -- ^ Start index along the dimension that needs to be extracted - -> Sz Ix1 -- ^ Size of the extracted array along the dimension that it will be extracted +splitExtractM + :: forall r ix e m + . (MonadThrow m, Index ix, Source r e) + => Dim + -- ^ Dimension along which to do the extraction + -> Ix1 + -- ^ Start index along the dimension that needs to be extracted + -> Sz Ix1 + -- ^ Size of the extracted array along the dimension that it will be extracted -> Array r ix e -> m (Array D ix e, Array D ix e, Array D ix e) splitExtractM dim startIx1 (Sz extractSzIx1) arr = do @@ -856,8 +909,6 @@ pure (leftArr, midArr, rightArr) {-# INLINE splitExtractM #-} - - -- | Replace a slice of an array with another one -- -- ==== __Example__ @@ -885,8 +936,9 @@ -- ] -- -- @since 0.6.1 -replaceSlice :: - forall r r' ix e m. (MonadThrow m, Source r e, Source r' e, Index ix, Index (Lower ix)) +replaceSlice + :: forall r r' ix e m + . (MonadThrow m, Source r e, Source r' e, Index ix, Index (Lower ix)) => Dim -> Ix1 -> Array r' (Lower ix) e @@ -898,7 +950,6 @@ concatM dim [l, delay m', r] {-# INLINE replaceSlice #-} - -- | Replace an outer slice of an array with another one -- -- ==== __Example__ @@ -926,8 +977,9 @@ -- ] -- -- @since 0.6.1 -replaceOuterSlice :: - forall r ix e m. (MonadThrow m, Index ix, Source r e, Load r (Lower ix) e) +replaceOuterSlice + :: forall r ix e m + . (MonadThrow m, Index ix, Source r e, Load r (Lower ix) e) => Ix1 -> Array r (Lower ix) e -> Array r ix e @@ -935,7 +987,6 @@ replaceOuterSlice i sl arr = replaceSlice (dimensions (size arr)) i sl arr {-# INLINE replaceOuterSlice #-} - -- | Delete a region from an array along the specified dimension. -- -- ==== __Examples__ @@ -960,12 +1011,17 @@ -- [ 0, 1, 2, 8, 9, 10 ] -- -- @since 0.3.5 -deleteRegionM :: - forall r ix e m. (MonadThrow m, Index ix, Source r e) - => Dim -- ^ Along which axis should the removal happen - -> Ix1 -- ^ At which index to start dropping slices - -> Sz Ix1 -- ^ Number of slices to drop - -> Array r ix e -- ^ Array that will have it's subarray removed +deleteRegionM + :: forall r ix e m + . (MonadThrow m, Index ix, Source r e) + => Dim + -- ^ Along which axis should the removal happen + -> Ix1 + -- ^ At which index to start dropping slices + -> Sz Ix1 + -- ^ Number of slices to drop + -> Array r ix e + -- ^ Array that will have it's subarray removed -> m (Array DL ix e) deleteRegionM dim ix sz arr = do (leftArr, _, rightArr) <- splitExtractM dim ix sz arr @@ -992,8 +1048,9 @@ -- ] -- -- @since 0.3.5 -deleteRowsM :: - forall r ix e m. (MonadThrow m, Index ix, Index (Lower ix), Source r e) +deleteRowsM + :: forall r ix e m + . (MonadThrow m, Index ix, Index (Lower ix), Source r e) => Ix1 -> Sz Ix1 -> Array r ix e @@ -1021,8 +1078,9 @@ -- ] -- -- @since 0.3.5 -deleteColumnsM :: - forall r ix e m. (MonadThrow m, Index ix, Source r e) +deleteColumnsM + :: forall r ix e m + . (MonadThrow m, Index ix, Source r e) => Ix1 -> Sz Ix1 -> Array r ix e @@ -1030,17 +1088,17 @@ deleteColumnsM = deleteRegionM 1 {-# INLINE deleteColumnsM #-} - -- | Discard elements from the source array according to the stride. -- -- @since 0.3.0 -downsample :: - forall r ix e. (Source r e, Load r ix e) +downsample + :: forall r ix e + . (Source r e, Load r ix e) => Stride ix -> Array r ix e -> Array DL ix e downsample stride arr = - DLArray {dlComp = getComp arr, dlSize = resultSize, dlLoad = load} + DLArray{dlComp = getComp arr, dlSize = resultSize, dlLoad = load} where resultSize = strideSize stride (size arr) strideIx = unStride stride @@ -1058,7 +1116,6 @@ {-# INLINE load #-} {-# INLINE downsample #-} - -- | Insert the same element into a `Load`able array according to the supplied stride. -- -- ====__Examples__ @@ -1088,11 +1145,15 @@ -- ] -- -- @since 0.3.0 -upsample :: - forall r ix e. Load r ix e - => e -- ^ Element to use for filling the newly added cells - -> Stride ix -- ^ Fill cells according to this stride - -> Array r ix e -- ^ Array that will have cells added to +upsample + :: forall r ix e + . Load r ix e + => e + -- ^ Element to use for filling the newly added cells + -> Stride ix + -- ^ Fill cells according to this stride + -> Array r ix e + -- ^ Array that will have cells added to -> Array DL ix e upsample !fillWith safeStride arr = DLArray @@ -1115,13 +1176,12 @@ !newsz = SafeSz (timesStride $ unSz sz) {-# INLINE upsample #-} - -- | General array transformation, that forces computation and produces a manifest array. -- -- @since 0.3.0 -transformM :: - forall r ix e r' ix' e' a m. - (Manifest r e, Index ix, Source r' e', Index ix', MonadUnliftIO m, PrimMonad m, MonadThrow m) +transformM + :: forall r ix e r' ix' e' a m + . (Manifest r e, Index ix, Source r' e', Index ix', MonadUnliftIO m, PrimMonad m, MonadThrow m) => (Sz ix' -> m (Sz ix, a)) -> (a -> (ix' -> m e') -> ix -> m e) -> Array r' ix' e' @@ -1131,13 +1191,12 @@ generateArray (getComp arr) sz (getM a (evaluateM arr)) {-# INLINE transformM #-} - -- | General array transformation -- -- @since 0.3.0 -transform' :: - forall ix e r' ix' e' a. - (HasCallStack, Source r' e', Index ix', Index ix) +transform' + :: forall ix e r' ix' e' a + . (HasCallStack, Source r' e', Index ix', Index ix) => (Sz ix' -> (Sz ix, a)) -> (a -> (ix' -> e') -> ix -> e) -> Array r' ix' e' @@ -1150,8 +1209,8 @@ -- | Same as `transformM`, but operates on two arrays -- -- @since 0.3.0 -transform2M :: - ( Manifest r e +transform2M + :: ( Manifest r e , Index ix , Source r1 e1 , Source r2 e2 @@ -1171,12 +1230,11 @@ generateArray (getComp arr1 <> getComp arr2) sz (getM a (evaluateM arr1) (evaluateM arr2)) {-# INLINE transform2M #-} - -- | Same as 'transform'', but operates on two arrays -- -- @since 0.3.0 -transform2' :: - (HasCallStack, Source r1 e1, Source r2 e2, Index ix, Index ix1, Index ix2) +transform2' + :: (HasCallStack, Source r1 e1, Source r2 e2, Index ix, Index ix1, Index ix2) => (Sz ix1 -> Sz ix2 -> (Sz ix, a)) -> (a -> (ix1 -> e1) -> (ix2 -> e2) -> ix -> e) -> Array r1 ix1 e1 @@ -1188,8 +1246,6 @@ (sz, a) = getSz (size arr1) (size arr2) {-# INLINE transform2' #-} - - -- | Replicate each element of the array by a factor in stride along each dimension and surround each -- such group with a box of supplied grid value. It will essentially zoom up an array and create a -- grid around each element from the original array. Very useful for zooming up images to inspect @@ -1220,11 +1276,15 @@ -- ] -- -- @since 0.3.1 -zoomWithGrid :: - forall r ix e. (Index ix, Source r e) - => e -- ^ Value to use for the grid - -> Stride ix -- ^ Scaling factor - -> Array r ix e -- ^ Source array +zoomWithGrid + :: forall r ix e + . (Index ix, Source r e) + => e + -- ^ Value to use for the grid + -> Stride ix + -- ^ Scaling factor + -> Array r ix e + -- ^ Source array -> Array DL ix e zoomWithGrid gridVal (Stride zoomFactor) arr = unsafeMakeLoadArray Seq newSz (Just gridVal) load where @@ -1236,7 +1296,7 @@ iforSchedulerM_ scheduler arr $ \ !ix !e -> let !kix = liftIndex2 (*) ix kx in mapM_ (\ !ix' -> writeElement (toLinearIndex newSz ix') e) $ - range Seq (liftIndex (+ 1) kix) (liftIndex2 (+) kix kx) + range Seq (liftIndex (+ 1) kix) (liftIndex2 (+) kix kx) {-# INLINE load #-} {-# INLINE zoomWithGrid #-} @@ -1274,10 +1334,13 @@ -- ] -- -- @since 0.4.4 -zoom :: - forall r ix e. (Index ix, Source r e) - => Stride ix -- ^ Scaling factor - -> Array r ix e -- ^ Source array +zoom + :: forall r ix e + . (Index ix, Source r e) + => Stride ix + -- ^ Scaling factor + -> Array r ix e + -- ^ Source array -> Array DL ix e zoom (Stride zoomFactor) arr = unsafeMakeLoadArray Seq newSz Nothing load where @@ -1288,6 +1351,6 @@ iforSchedulerM_ scheduler arr $ \ !ix !e -> let !kix = liftIndex2 (*) ix zoomFactor in mapM_ (\ !ix' -> writeElement (toLinearIndex newSz ix') e) $ - range Seq kix (liftIndex2 (+) kix zoomFactor) + range Seq kix (liftIndex2 (+) kix zoomFactor) {-# INLINE load #-} {-# INLINE zoom #-} diff --git a/src/Data/Massiv/Array/Stencil.hs b/src/Data/Massiv/Array/Stencil.hs --- a/src/Data/Massiv/Array/Stencil.hs +++ b/src/Data/Massiv/Array/Stencil.hs @@ -3,6 +3,7 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE TypeFamilies #-} + -- | -- Module : Data.Massiv.Array.Stencil -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -10,37 +11,41 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Stencil - ( -- * Stencil - Stencil - , makeStencil - , getStencilSize - , getStencilCenter +module Data.Massiv.Array.Stencil ( + -- * Stencil + Stencil, + makeStencil, + getStencilSize, + getStencilCenter, + -- ** Padding - , Padding(..) - , noPadding - , samePadding + Padding (..), + noPadding, + samePadding, + -- ** Application - , mapStencil - , applyStencil + mapStencil, + applyStencil, + -- ** Common stencils - , idStencil - , sumStencil - , productStencil - , avgStencil - , maxStencil - , minStencil - , foldlStencil - , foldrStencil - , foldStencil + idStencil, + sumStencil, + productStencil, + avgStencil, + maxStencil, + minStencil, + foldlStencil, + foldrStencil, + foldStencil, + -- ** Profunctor - , dimapStencil - , lmapStencil - , rmapStencil + dimapStencil, + lmapStencil, + rmapStencil, + -- * Convolution - , module Data.Massiv.Array.Stencil.Convolution - ) where + module Data.Massiv.Array.Stencil.Convolution, +) where import Data.Coerce import Data.Massiv.Array.Delayed.Windowed @@ -68,16 +73,18 @@ -- `Data.Massiv.Array.compute`d in order to be useful. -- -- @since 0.1.0 -mapStencil :: - (Index ix, Manifest r e) - => Border e -- ^ Border resolution technique - -> Stencil ix e a -- ^ Stencil to map over the array - -> Array r ix e -- ^ Source array +mapStencil + :: (Index ix, Manifest r e) + => Border e + -- ^ Border resolution technique + -> Stencil ix e a + -- ^ Stencil to map over the array + -> Array r ix e + -- ^ Source array -> Array DW ix a mapStencil b stencil = applyStencil (samePadding stencil b) stencil {-# INLINE mapStencil #-} - -- | Padding of the source array before stencil application. -- -- ==== __Examples__ @@ -143,11 +150,12 @@ -- -- @since 0.4.3 data Padding ix e = Padding - { paddingFromOrigin :: !(Sz ix) - , paddingFromBottom :: !(Sz ix) + { paddingFromOrigin :: !(Sz ix) + , paddingFromBottom :: !(Sz ix) , paddingWithElement :: !(Border e) -- ^ Element to do padding with - } deriving (Eq, Show) + } + deriving (Eq, Show) -- | Also known as "valid" padding. When stencil is applied to an array, that array will -- shrink, unless the stencil is of size 1. @@ -166,24 +174,25 @@ samePadding (Stencil (Sz sSz) sCenter _) border = Padding { paddingFromOrigin = Sz sCenter - , paddingFromBottom = Sz (liftIndex2 (-) sSz (liftIndex (+1) sCenter)) + , paddingFromBottom = Sz (liftIndex2 (-) sSz (liftIndex (+ 1) sCenter)) , paddingWithElement = border } - -- | Apply a constructed stencil over an array. Resulting array must be -- `Data.Massiv.Array.compute`d in order to be useful. Unlike `mapStencil`, the size of -- the resulting array will not necesserally be the same as the source array, which will -- depend on the padding. -- -- @since 0.4.3 -applyStencil :: - (Index ix, Manifest r e) +applyStencil + :: (Index ix, Manifest r e) => Padding ix e -- ^ Padding to be applied to the source array. This will dictate the resulting size of -- the array. No padding will cause it to shrink by the size of the stencil - -> Stencil ix e a -- ^ Stencil to apply to the array - -> Array r ix e -- ^ Source array + -> Stencil ix e a + -- ^ Stencil to apply to the array + -> Array r ix e + -- ^ Source array -> Array DW ix a applyStencil (Padding (Sz po) (Sz pb) border) (Stencil sSz sCenter stencilF) !arr = insertWindow warr window @@ -207,7 +216,6 @@ } {-# INLINE applyStencil #-} - -- | Construct a stencil from a function, which describes how to calculate the -- value at a point while having access to neighboring elements with a function -- that accepts idices relative to the center of stencil. Trying to index @@ -225,18 +233,22 @@ -- -- /Note/ - Make sure to add an @INLINE@ pragma, otherwise performance will be terrible. -- --- > average3x3Stencil :: Fractional a => Stencil Ix2 a a --- > average3x3Stencil = makeStencil (Sz (3 :. 3)) (1 :. 1) $ \ get -> --- > ( get (-1 :. -1) + get (-1 :. 0) + get (-1 :. 1) + --- > get ( 0 :. -1) + get ( 0 :. 0) + get ( 0 :. 1) + --- > get ( 1 :. -1) + get ( 1 :. 0) + get ( 1 :. 1) ) / 9 --- > {-# INLINE average3x3Stencil #-} +-- @ +-- average3x3Stencil :: Fractional a => Stencil Ix2 a a +-- average3x3Stencil = makeStencil (Sz (3 :. 3)) (1 :. 1) $ \ get -> +-- ( get (-1 :. -1) + get (-1 :. 0) + get (-1 :. 1) + +-- get ( 0 :. -1) + get ( 0 :. 0) + get ( 0 :. 1) + +-- get ( 1 :. -1) + get ( 1 :. 0) + get ( 1 :. 1) ) / 9 +-- {\-# INLINE average3x3Stencil #-\} +-- @ -- -- @since 0.1.0 makeStencil :: Index ix - => Sz ix -- ^ Size of the stencil - -> ix -- ^ Center of the stencil + => Sz ix + -- ^ Size of the stencil + -> ix + -- ^ Center of the stencil -> ((ix -> e) -> a) -- ^ Stencil function that receives a "get" function as it's argument that can -- retrieve values of cells in the source array with respect to the center of @@ -255,10 +267,9 @@ -- -- @since 0.4.3 idStencil :: Index ix => Stencil ix e e -idStencil = makeUnsafeStencil oneSz zeroIndex $ \ _ get -> get zeroIndex +idStencil = makeUnsafeStencil oneSz zeroIndex $ \_ get -> get zeroIndex {-# INLINE idStencil #-} - -- | Stencil that does a left fold in a row-major order. Regardless of the supplied size -- resulting stencil will be centered at zero, although by using `Padding` it is possible -- to overcome this limitation. @@ -318,7 +329,6 @@ iter ixStart zeroIndex (pureIndex (-1)) (>=) acc0 $ \ix -> f (get ix) {-# INLINE foldrStencil #-} - -- | Create a stencil that will fold all elements in the region monoidally. -- -- @since 0.4.3 @@ -358,7 +368,6 @@ maxStencil :: (Bounded e, Ord e, Index ix) => Sz ix -> Stencil ix e e maxStencil = dimapStencil coerce getMax . foldStencil {-# INLINE maxStencil #-} - -- | Create a stencil centered at 0 that will extract the maximum value in the region of -- supplied size. diff --git a/src/Data/Massiv/Array/Stencil/Convolution.hs b/src/Data/Massiv/Array/Stencil/Convolution.hs --- a/src/Data/Massiv/Array/Stencil/Convolution.hs +++ b/src/Data/Massiv/Array/Stencil/Convolution.hs @@ -1,5 +1,6 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE FlexibleContexts #-} + -- | -- Module : Data.Massiv.Array.Stencil.Convolution -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -7,13 +8,12 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Stencil.Convolution - ( makeConvolutionStencil - , makeConvolutionStencilFromKernel - , makeCorrelationStencil - , makeCorrelationStencilFromKernel - ) where +module Data.Massiv.Array.Stencil.Convolution ( + makeConvolutionStencil, + makeConvolutionStencilFromKernel, + makeCorrelationStencil, + makeCorrelationStencilFromKernel, +) where import Data.Massiv.Array.Ops.Fold (ifoldlS) import Data.Massiv.Array.Stencil.Internal @@ -24,19 +24,21 @@ -- an accumulator function. -- -- /Note/ - Using `Data.Massiv.Array.Stencil.Unsafe.makeUnsafeConvolutionStencil` will be --- much faster, therefore it is recommended to switch from this function, after manual +-- slightly faster, therefore it is recommended to switch from this function, after manual -- verification that the created stencil behaves as expected. -- -- ==== __Examples__ -- -- Here is how to create a 2D horizontal Sobel Stencil: -- --- > sobelX :: Num e => Stencil Ix2 e e --- > sobelX = makeConvolutionStencil (Sz2 3 3) (1 :. 1) $ --- > \f -> f (-1 :. -1) (-1) . f (-1 :. 1) 1 . --- > f ( 0 :. -1) (-2) . f ( 0 :. 1) 2 . --- > f ( 1 :. -1) (-1) . f ( 1 :. 1) 1 --- > {-# INLINE sobelX #-} +-- @ +-- sobelX :: Num e => Stencil Ix2 e e +-- sobelX = makeConvolutionStencil (Sz2 3 3) (1 :. 1) +-- $ \f -> f (-1 :. -1) (-1) . f (-1 :. 1) 1 . +-- f ( 0 :. -1) (-2) . f ( 0 :. 1) 2 . +-- f ( 1 :. -1) (-1) . f ( 1 :. 1) 1 +-- {\-# INLINE sobelX #-\} +-- @ -- -- @since 0.1.0 makeConvolutionStencil @@ -54,7 +56,6 @@ {-# INLINE stencil #-} {-# INLINE makeConvolutionStencil #-} - -- | Make a stencil out of a Kernel Array. This `Stencil` will be slower than if -- `makeConvolutionStencil` is used, but sometimes we just really don't know the -- kernel at compile time. @@ -70,14 +71,14 @@ !szi1 = liftIndex (subtract 1) szi !sInvertCenter = liftIndex2 (-) szi1 sCenter !sCenter = liftIndex (`quot` 2) szi - stencil uget _ !ix = ifoldlS accum 0 kArr where - !ixOff = liftIndex2 (+) ix sCenter - accum !acc !kIx !kVal = uget (liftIndex2 (-) ixOff kIx) * kVal + acc - {-# INLINE accum #-} + stencil uget _ !ix = ifoldlS accum 0 kArr + where + !ixOff = liftIndex2 (+) ix sCenter + accum !acc !kIx !kVal = uget (liftIndex2 (-) ixOff kIx) * kVal + acc + {-# INLINE accum #-} {-# INLINE stencil #-} {-# INLINE makeConvolutionStencilFromKernel #-} - -- | Make a stencil -- -- /Note/ - Using `Data.Massiv.Array.Stencil.Unsafe.makeUnsafeCorrelationStencil` will be @@ -111,9 +112,10 @@ where !sz = size kArr !sCenter = liftIndex (`div` 2) $ unSz sz - stencil uget _ !ix = ifoldlS accum 0 kArr where - !ixOff = liftIndex2 (-) ix sCenter - accum !acc !kIx !kVal = uget (liftIndex2 (+) ixOff kIx) * kVal + acc - {-# INLINE accum #-} + stencil uget _ !ix = ifoldlS accum 0 kArr + where + !ixOff = liftIndex2 (-) ix sCenter + accum !acc !kIx !kVal = uget (liftIndex2 (+) ixOff kIx) * kVal + acc + {-# INLINE accum #-} {-# INLINE stencil #-} {-# INLINE makeCorrelationStencilFromKernel #-} diff --git a/src/Data/Massiv/Array/Stencil/Internal.hs b/src/Data/Massiv/Array/Stencil/Internal.hs --- a/src/Data/Massiv/Array/Stencil/Internal.hs +++ b/src/Data/Massiv/Array/Stencil/Internal.hs @@ -4,6 +4,7 @@ {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE ScopedTypeVariables #-} + -- | -- Module : Data.Massiv.Array.Stencil.Internal -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -11,13 +12,12 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Stencil.Internal - ( Stencil(..) - , dimapStencil - , lmapStencil - , rmapStencil - ) where +module Data.Massiv.Array.Stencil.Internal ( + Stencil (..), + dimapStencil, + lmapStencil, + rmapStencil, +) where import Control.Applicative import Control.DeepSeq @@ -28,12 +28,11 @@ -- `Data.Massiv.Array.makeStencil` and `Data.Massiv.Array.makeConvolutionStencil` in order -- to create a stencil. data Stencil ix e a = Stencil - { stencilSize :: !(Sz ix) + { stencilSize :: !(Sz ix) , stencilCenter :: !ix - , stencilFunc :: (ix -> e) -> (ix -> e) -> ix -> a + , stencilFunc :: (ix -> e) -> (ix -> e) -> ix -> a } - instance Index ix => NFData (Stencil ix e a) where rnf (Stencil sz ix f) = sz `deepseq` ix `deepseq` f `seq` () @@ -41,14 +40,13 @@ fmap = rmapStencil {-# INLINE fmap #-} - -- Profunctor -- | A Profunctor dimap. Same caviat applies as in `lmapStencil` -- -- @since 0.2.3 dimapStencil :: (c -> d) -> (a -> b) -> Stencil ix d a -> Stencil ix c b -dimapStencil f g stencil@Stencil {stencilFunc = sf} = stencil {stencilFunc = sf'} +dimapStencil f g stencil@Stencil{stencilFunc = sf} = stencil{stencilFunc = sf'} where sf' us s = g . sf (f . us) (f . s) {-# INLINE sf' #-} @@ -63,7 +61,7 @@ -- -- @since 0.2.3 lmapStencil :: (c -> d) -> Stencil ix d a -> Stencil ix c a -lmapStencil f stencil@Stencil {stencilFunc = sf} = stencil {stencilFunc = sf'} +lmapStencil f stencil@Stencil{stencilFunc = sf} = stencil{stencilFunc = sf'} where sf' us s = sf (f . us) (f . s) {-# INLINE sf' #-} @@ -76,7 +74,7 @@ -- -- @since 0.2.3 rmapStencil :: (a -> b) -> Stencil ix e a -> Stencil ix e b -rmapStencil f stencil@Stencil {stencilFunc = sf} = stencil {stencilFunc = sf'} +rmapStencil f stencil@Stencil{stencilFunc = sf} = stencil{stencilFunc = sf'} where sf' us s = f . sf us s {-# INLINE sf' #-} 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 @@ -3,6 +3,7 @@ {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE RecordWildCards #-} + -- | -- Module : Data.Massiv.Array.Stencil.Unsafe -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -10,20 +11,18 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Stencil.Unsafe - ( -- * Stencil - makeUnsafeStencil - , makeUnsafeConvolutionStencil - , makeUnsafeCorrelationStencil - , unsafeTransformStencil - ) where +module Data.Massiv.Array.Stencil.Unsafe ( + -- * Stencil + makeUnsafeStencil, + makeUnsafeConvolutionStencil, + makeUnsafeCorrelationStencil, + unsafeTransformStencil, +) where import Data.Massiv.Array.Stencil.Internal import Data.Massiv.Core.Common import GHC.Exts (inline) - -- | Similar to `Data.Massiv.Array.Stencil.makeStencil`, but there are no guarantees that the -- stencil will not read out of bounds memory. This stencil is also a bit more powerful in sense it -- gets an extra peice of information, namely the exact index for the element it is constructing. @@ -31,8 +30,10 @@ -- @since 0.3.0 makeUnsafeStencil :: Index ix - => Sz ix -- ^ Size of the stencil - -> ix -- ^ Center of the stencil + => Sz ix + -- ^ Size of the stencil + -> ix + -- ^ Center of the stencil -> (ix -> (ix -> e) -> a) -- ^ Stencil function. -> Stencil ix e a @@ -80,7 +81,6 @@ {-# INLINE stencil #-} {-# INLINE makeUnsafeCorrelationStencil #-} - -- | Perform an arbitrary transformation of a stencil. This stencil modifier can be used for -- example to turn a vector stencil into a matrix stencil implement, or transpose a matrix -- stencil. It is really easy to get this wrong, so be extremely careful. @@ -112,18 +112,22 @@ -- ] -- -- @since 0.5.4 -unsafeTransformStencil :: - (Sz ix' -> Sz ix) +unsafeTransformStencil + :: (Sz ix' -> Sz ix) -- ^ Forward modifier for the size -> (ix' -> ix) -- ^ Forward index modifier - -> (((ix' -> e) -> (ix' -> e) -> ix' -> a) - -> (ix -> e) -> (ix -> e) -> ix -> a) + -> ( ((ix' -> e) -> (ix' -> e) -> ix' -> a) + -> (ix -> e) + -> (ix -> e) + -> ix + -> a + ) -- ^ Inverse stencil function modifier -> Stencil ix' e a -- ^ Original stencil. -> Stencil ix e a -unsafeTransformStencil transformSize transformIndex transformFunc Stencil {..} = +unsafeTransformStencil transformSize transformIndex transformFunc Stencil{..} = Stencil { stencilSize = transformSize stencilSize , stencilCenter = transformIndex stencilCenter @@ -131,15 +135,12 @@ } {-# INLINE unsafeTransformStencil #-} - - {- Invalid stencil transformer function. TODO: figure out if there is a safe way to do stencil index trnasformation. - transformStencil :: (Default e, Index ix) => (Sz ix' -> Sz ix) @@ -154,6 +155,5 @@ transformStencil transformSize transformIndex transformIndex' stencil = validateStencil def $! unsafeTransformStencil transformSize transformIndex transformIndex' stencil {-# INLINE transformStencil #-} - -} diff --git a/src/Data/Massiv/Array/Unsafe.hs b/src/Data/Massiv/Array/Unsafe.hs --- a/src/Data/Massiv/Array/Unsafe.hs +++ b/src/Data/Massiv/Array/Unsafe.hs @@ -3,6 +3,7 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} + -- | -- Module : Data.Massiv.Array.Unsafe -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -10,98 +11,119 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Array.Unsafe - ( -- * Creation - unsafeMakeLoadArray - , unsafeMakeLoadArrayAdjusted - -- * Indexing - , Sz(SafeSz) - , Stride(SafeStride) - , unsafeIndex - , unsafeLinearIndex - , unsafeLinearIndexM - -- * Manipulations - , unsafeBackpermute - , unsafeResize - , unsafeExtract - , unsafeTransform - , unsafeTransform2 - -- * Slicing - , unsafeSlice - , unsafeOuterSlice - , unsafeInnerSlice - , unsafeLinearSlice - -- * Mutable interface - , unsafeResizeMArray - , unsafeLinearSliceMArray - , unsafeThaw - , unsafeFreeze - , unsafeNew - , unsafeLoadIntoST - , unsafeLoadIntoIO - , unsafeLoadIntoS - , unsafeLoadIntoM - , unsafeCreateArray - , unsafeCreateArray_ - , unsafeCreateArrayS - -- ** Read - , unsafeRead - , unsafeLinearRead - -- ** Write - , unsafeWrite - , unsafeLinearWrite - -- ** Modify - , unsafeModify - , unsafeLinearModify - -- ** Swap - , unsafeSwap - , unsafeLinearSwap - -- ** Range modification - , unsafeLinearSet - , unsafeLinearCopy - , unsafeArrayLinearCopy - -- ** Resizing - , unsafeLinearShrink - , unsafeLinearGrow - -- * Pointer access - , unsafeMallocMArray - , unsafeWithPtr - , unsafeArrayToForeignPtr - , unsafeMArrayToForeignPtr - , unsafeArrayFromForeignPtr - , unsafeArrayFromForeignPtr0 - , unsafeMArrayFromForeignPtr - , unsafeMArrayFromForeignPtr0 - -- ** Atomic Operations - , unsafeAtomicReadIntArray - , unsafeAtomicWriteIntArray - , unsafeAtomicModifyIntArray - , unsafeAtomicAddIntArray - , unsafeAtomicSubIntArray - , unsafeAtomicAndIntArray - , unsafeAtomicNandIntArray - , unsafeAtomicOrIntArray - , unsafeAtomicXorIntArray - , unsafeCasIntArray - -- ** Other operations - , coerceBoxedArray - , coerceNormalBoxedArray - , unsafeUnstablePartitionRegionM - , module Data.Massiv.Vector.Unsafe - , module Data.Massiv.Array.Stencil.Unsafe - -- * Constructors - , Array(PArray, SArray, UArray, BArray, BLArray, BNArray, DArray, DLArray, DSArray, DIArray, DWArray) - , MArray(MPArray, MSArray, MUArray, MBArray, MBLArray, MBNArray) - ) where +module Data.Massiv.Array.Unsafe ( + -- * Creation + 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)) + -- * Indexing + Sz (SafeSz), + Stride (SafeStride), + unsafeIndex, + unsafePrefIndex, + unsafeLinearIndex, + unsafeLinearIndexM, + + -- * Manipulations + unsafeBackpermute, + unsafeResize, + unsafeExtract, + unsafeTransform, + unsafeTransform2, + + -- * Slicing + unsafeSlice, + unsafeOuterSlice, + unsafeInnerSlice, + unsafeLinearSlice, + + -- * Mutable interface + unsafeResizeMArray, + unsafeLinearSliceMArray, + unsafeThaw, + unsafeFreeze, + unsafeNew, + unsafeLoadIntoST, + unsafeLoadIntoIO, + unsafeLoadIntoS, + unsafeLoadIntoM, + unsafeCreateArray, + unsafeCreateArray_, + unsafeCreateArrayS, + + -- ** Read + unsafeRead, + unsafeLinearRead, + + -- ** Write + unsafeWrite, + unsafeLinearWrite, + + -- ** Modify + unsafeModify, + unsafeLinearModify, + + -- ** Swap + unsafeSwap, + unsafeLinearSwap, + + -- ** Range modification + unsafeLinearSet, + unsafeLinearCopy, + unsafeArrayLinearCopy, + + -- ** Resizing + unsafeLinearShrink, + unsafeLinearGrow, + + -- * Pointer access + unsafeMallocMArray, + unsafeWithPtr, + unsafeArrayToForeignPtr, + unsafeMArrayToForeignPtr, + unsafeArrayFromForeignPtr, + unsafeArrayFromForeignPtr0, + unsafeMArrayFromForeignPtr, + unsafeMArrayFromForeignPtr0, + + -- ** Atomic Operations + unsafeAtomicReadIntArray, + unsafeAtomicWriteIntArray, + unsafeAtomicModifyIntArray, + unsafeAtomicAddIntArray, + unsafeAtomicSubIntArray, + unsafeAtomicAndIntArray, + unsafeAtomicNandIntArray, + unsafeAtomicOrIntArray, + unsafeAtomicXorIntArray, + unsafeCasIntArray, + + -- ** Other operations + coerceBoxedArray, + coerceNormalBoxedArray, + unsafeUnstablePartitionRegionM, + module Data.Massiv.Vector.Unsafe, + module Data.Massiv.Array.Stencil.Unsafe, + + -- * Constructors + Array (PArray, SArray, UArray, BArray, BLArray, BNArray, DArray, DLArray, DSArray, DIArray, DWArray), + MArray (MPArray, MSArray, MUArray, MBArray, MBLArray, MBNArray), +) where + +import Data.Massiv.Array.Delayed.Interleaved (Array (DIArray)) +import Data.Massiv.Array.Delayed.Pull ( + D, + unsafeExtract, + unsafeInnerSlice, + unsafeSlice, + ) +import Data.Massiv.Array.Delayed.Push ( + Array (DLArray), + unsafeMakeLoadArray, + unsafeMakeLoadArrayAdjusted, + ) +import Data.Massiv.Array.Delayed.Stream (Array (DSArray)) +import Data.Massiv.Array.Delayed.Windowed (Array (DWArray)) import Data.Massiv.Array.Manifest.Boxed import Data.Massiv.Array.Manifest.Internal import Data.Massiv.Array.Manifest.Primitive @@ -111,12 +133,15 @@ 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.Core.Index.Stride (Stride (SafeStride)) import Data.Massiv.Vector.Unsafe - -unsafeBackpermute :: (Index ix', Source r' e, Index ix) => - Sz ix -> (ix -> ix') -> Array r' ix' e -> Array D ix e +unsafeBackpermute + :: (Index ix', Source r' e, Index ix) + => Sz ix + -> (ix -> ix') + -> Array r' ix' e + -> Array D ix e unsafeBackpermute !sz ixF !arr = makeArray (getComp arr) sz (unsafeIndex arr . ixF) {-# INLINE unsafeBackpermute #-} @@ -124,8 +149,8 @@ -- but unsafe. -- -- @since 0.3.0 -unsafeTransform :: - (Index ix', Source r' e', Index ix) +unsafeTransform + :: (Index ix', Source r' e', Index ix) => (Sz ix' -> (Sz ix, a)) -> (a -> (ix' -> e') -> ix -> e) -> Array r' ix' e' @@ -139,8 +164,8 @@ -- but unsafe. -- -- @since 0.3.0 -unsafeTransform2 :: - (Index ix1, Source r1 e1, Index ix2, Source r2 e2, Index ix) +unsafeTransform2 + :: (Index ix1, Source r1 e1, Index ix2, Source r2 e2, Index ix) => (Sz ix1 -> Sz ix2 -> (Sz ix, a)) -> (a -> (ix1 -> e1) -> (ix2 -> e2) -> ix -> e) -> Array r1 ix1 e1 diff --git a/src/Data/Massiv/Core.hs b/src/Data/Massiv/Core.hs --- a/src/Data/Massiv/Core.hs +++ b/src/Data/Massiv/Core.hs @@ -5,62 +5,63 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Core - ( Array(LArray) - , List(..) - , Vector - , MVector - , Matrix - , MMatrix - , Load(iterArrayLinearST_, iterArrayLinearWithSetST_) - , Stream(..) - , Source - , PrefIndex(..) - , Size - , Shape(..) - , LengthHint(..) - , StrideLoad(..) - , Manifest - , Mutable - , Ragged - , L(..) - , ListItem - , Scheduler - , SchedulerWS - , Strategy - , Comp(Seq, Par, Par', ParOn, ParN) - , getComp - , setComp - , appComp - , WorkerStates - , initWorkerStates - , scheduleWork - , scheduleWork_ - , module Data.Massiv.Core.Index +module Data.Massiv.Core ( + Array (LArray), + List (..), + Vector, + MVector, + Matrix, + MMatrix, + Load (iterArrayLinearST_, iterArrayLinearWithSetST_), + Stream (..), + Source, + PrefIndex (..), + Size, + Shape (..), + LengthHint (..), + StrideLoad (..), + Manifest, + Mutable, + Ragged, + L (..), + ListItem, + Scheduler, + SchedulerWS, + Strategy, + Comp (Seq, Par, Par', ParOn, ParN), + getComp, + setComp, + appComp, + WorkerStates, + initWorkerStates, + scheduleWork, + scheduleWork_, + module Data.Massiv.Core.Index, + -- * Numeric - , FoldNumeric - , Numeric - , NumericFloat + FoldNumeric, + Numeric, + NumericFloat, + -- * Exceptions - , MonadThrow(..) - , IndexException(..) - , SizeException(..) - , ShapeException(..) - , module Data.Massiv.Core.Exception + MonadThrow (..), + IndexException (..), + SizeException (..), + ShapeException (..), + module Data.Massiv.Core.Exception, + -- * Stateful Monads - , MonadUnliftIO - , MonadIO(liftIO) - , PrimMonad(PrimState) - ) where + MonadUnliftIO, + MonadIO (liftIO), + PrimMonad (PrimState), +) where import Control.Scheduler (SchedulerWS, initWorkerStates) import Data.Massiv.Core.Common +import Data.Massiv.Core.Exception import Data.Massiv.Core.Index import Data.Massiv.Core.List -import Data.Massiv.Core.Exception import Data.Massiv.Core.Operations (FoldNumeric, Numeric, NumericFloat) - -- | Append computation strategy using `Comp`'s `Monoid` instance. -- 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 @@ -8,6 +8,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UndecidableInstances #-} + -- | -- Module : Data.Massiv.Core.Common -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -15,99 +16,113 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable -module Data.Massiv.Core.Common - ( Array - , Vector - , Matrix - , MArray - , MVector - , MMatrix - , Steps(..) - , Stream(..) - , Strategy(..) - , Source(..) - , PrefIndex(..) - , Load(..) - , StrideLoad(..) - , Size(..) - , Shape(..) - , Manifest(..) - , Mutable - , Comp(..) - , Scheduler - , numWorkers - , scheduleWork - , scheduleWork_ - , WorkerStates - , unsafeRead - , unsafeWrite - , unsafeModify - , unsafeLinearModify - , unsafeSwap - , unsafeLinearSwap - , unsafeDefaultLinearShrink - , Ragged(..) - , empty - , singleton +module Data.Massiv.Core.Common ( + Array, + Vector, + Matrix, + MArray, + MVector, + MMatrix, + Steps (..), + Stream (..), + Strategy (..), + Source (..), + PrefIndex (..), + Load (..), + StrideLoad (..), + Size (..), + Shape (..), + Manifest (..), + Mutable, + Comp (..), + Scheduler, + numWorkers, + scheduleWork, + scheduleWork_, + WorkerStates, + unsafeRead, + unsafeWrite, + unsafeModify, + unsafeLinearModify, + unsafeSwap, + unsafeLinearSwap, + unsafeDefaultLinearShrink, + Ragged (..), + empty, + singleton, + -- * Size - , elemsCount - , isNotNull - , isEmpty - , isNotEmpty - , Sz(SafeSz) - , LengthHint(..) + elemsCount, + isNotNull, + isEmpty, + isNotEmpty, + Sz (SafeSz), + LengthHint (..), + -- * Indexing - , (!?) - , index - , indexM - , (!) - , index' - , (??) - , defaultIndex - , borderIndex - , evaluateM - , evaluate' - , inline0 - , inline1 - , inline2 - , module Data.Massiv.Core.Index + (!?), + index, + indexM, + (!), + index', + (??), + defaultIndex, + borderIndex, + evaluateM, + evaluate', + inline0, + inline1, + inline2, + module Data.Massiv.Core.Index, + -- * Common Operations - , Semigroup((<>)) + Semigroup ((<>)), + -- * Exceptions - , MonadThrow(..) - , IndexException(..) - , SizeException(..) - , ShapeException(..) - , module Data.Massiv.Core.Exception - , Proxy(..) - , Id(..) + MonadThrow (..), + IndexException (..), + SizeException (..), + ShapeException (..), + module Data.Massiv.Core.Exception, + Proxy (..), + Id (..), + -- * Stateful Monads - , runST - , ST - , MonadUnliftIO(..) - , MonadIO(liftIO) - , PrimMonad(PrimState) - , RealWorld - ) where + runST, + ST, + MonadUnliftIO (..), + MonadIO (liftIO), + PrimMonad (PrimState), + RealWorld, +) where #if !MIN_VERSION_base(4,11,0) import Data.Semigroup (Semigroup((<>))) #endif -import Control.Monad.Catch (MonadThrow(..)) -import Control.Monad.IO.Unlift (MonadIO(liftIO), MonadUnliftIO(..)) +import Control.Monad.Catch (MonadThrow (..)) +import Control.Monad.IO.Unlift (MonadIO (liftIO), MonadUnliftIO (..)) import Control.Monad.Primitive import Control.Monad.ST -import Control.Scheduler (Comp(..), Scheduler, WorkerStates, numWorkers, - scheduleWork, scheduleWork_, trivialScheduler_) -import GHC.Exts (IsList) +import Control.Scheduler ( + Comp (..), + Scheduler, + WorkerStates, + numWorkers, + scheduleWork, + scheduleWork_, + trivialScheduler_, + ) +import Data.Kind import Data.Massiv.Core.Exception import Data.Massiv.Core.Index -import Data.Massiv.Core.Index.Internal (Sz(SafeSz)) +import Data.Massiv.Core.Index.Internal (Sz (SafeSz)) import Data.Typeable -import Data.Kind -import qualified Data.Vector.Fusion.Stream.Monadic as S (Stream) +import qualified Data.Stream.Monadic as S (Stream) import Data.Vector.Fusion.Util +import GHC.Exts (IsList) +#include "massiv.h" + -- | The array family. Representations @r@ describe how data is arranged or computed. All -- arrays have a common property that each index @ix@ always maps to the same unique -- element @e@, even if that element does not yet exist in memory and the array has to be @@ -127,7 +142,6 @@ -- @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`. -- @@ -144,7 +158,6 @@ -- @since 0.5.0 type MMatrix s r e = MArray s r Ix2 e - class Load r ix e => Stream r ix e where toStream :: Array r ix e -> Steps Id e @@ -152,7 +165,7 @@ data Steps m e = Steps { stepsStream :: S.Stream m e - , stepsSize :: LengthHint + , stepsSize :: LengthHint } class Typeable r => Strategy r where @@ -169,7 +182,6 @@ -- >>> setComp (ParN 6) a -- use 6 capabilities -- Array DL (ParN 6) (Sz1 1) -- [ 0 ] - -- setComp :: Comp -> Array r ix e -> Array r ix e -- | Get computation strategy of this array @@ -184,26 +196,27 @@ -- @since 1.0.2 repr :: r repr = - error $ "Array representation should never be evaluated: " ++ - show (typeRep (Proxy :: Proxy r)) - + error $ + "Array representation should never be evaluated: " + ++ show (typeRep (Proxy :: Proxy r)) -- | Size hint -- -- @since 1.0.0 data LengthHint - = LengthExact Sz1 -- ^ Exact known size - | LengthMax Sz1 -- ^ Upper bound on the size - | LengthUnknown -- ^ Unknown size + = -- | Exact known size + LengthExact Sz1 + | -- | Upper bound on the size + LengthMax Sz1 + | -- | Unknown size + LengthUnknown deriving (Eq, Show) - -- | The shape of an array. It is different from `Size` in that it can be applicable to -- non-square matrices and might not be available in constant time. -- -- @since 1.0.0 class Index ix => Shape r ix where - -- | /O(1)/ - Check what do we know about the number of elements without doing any work -- -- @since 1.0.0 @@ -257,18 +270,16 @@ isNull = (zeroSz ==) . linearSize {-# INLINE isNull #-} - lengthHintUpperBound :: LengthHint -> Maybe Sz1 lengthHintUpperBound = \case - LengthExact sz -> Just sz - LengthMax sz -> Just sz - LengthUnknown -> Nothing + LengthExact sz -> Just sz + LengthMax sz -> Just sz + LengthUnknown -> Nothing {-# INLINE lengthHintUpperBound #-} -- | Arrays that have information about their size availible in constant -- time. class Size r where - -- | /O(1)/ - Get the exact size of an immutabe array. Most of the time will -- produce the size in constant time, except for `Data.Massiv.Array.DS` -- representation, which could result in evaluation of the whole stream. See @@ -298,13 +309,13 @@ -- | Arrays that can be used as source to practically any manipulation function. class (Strategy r, Size r) => Source r e where - {-# MINIMAL (unsafeIndex|unsafeLinearIndex), unsafeLinearSlice #-} + {-# 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 :: Index ix => Array r ix e -> ix -> e + unsafeIndex :: HAS_CALL_STACK => Index ix => Array r ix e -> ix -> e unsafeIndex !arr = unsafeLinearIndex arr . toLinearIndex (size arr) {-# INLINE unsafeIndex #-} @@ -312,7 +323,7 @@ -- bounds check is performed -- -- @since 0.1.0 - unsafeLinearIndex :: Index ix => Array r ix e -> Int -> e + unsafeLinearIndex :: HAS_CALL_STACK => Index ix => Array r ix e -> Int -> e unsafeLinearIndex !arr = unsafeIndex arr . fromLinearIndex (size arr) {-# INLINE unsafeLinearIndex #-} @@ -320,16 +331,20 @@ -- efficient for underlying representation -- -- @since 1.0.2 - unsafePrefIndex :: Index ix => Array r ix e -> PrefIndex ix e + unsafePrefIndex :: HAS_CALL_STACK => Index ix => Array r ix e -> PrefIndex ix e unsafePrefIndex !arr = PrefIndexLinear (unsafeLinearIndex arr) {-# INLINE unsafePrefIndex #-} - -- | /O(1)/ - Take a slice out of an array from the outside -- -- @since 0.1.0 - unsafeOuterSlice :: (Index ix, Index (Lower ix)) => - Array r ix e -> Sz (Lower ix) -> Int -> Array r (Lower ix) e + unsafeOuterSlice + :: HAS_CALL_STACK + => (Index ix, Index (Lower ix)) + => Array r ix e + -> Sz (Lower ix) + -> Int + -> Array r (Lower ix) e unsafeOuterSlice arr sz i = unsafeResize sz $ unsafeLinearSlice i (toLinearSz sz) arr {-# INLINE unsafeOuterSlice #-} @@ -366,10 +381,13 @@ -- ] -- -- @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 + 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 #-} @@ -388,7 +406,6 @@ makeArrayLinear comp sz f = makeArray comp sz (f . toLinearIndex sz) {-# INLINE makeArrayLinear #-} - -- | Construct an array of the specified size that contains the same element in all of -- the cells. -- @@ -397,18 +414,19 @@ replicate comp sz !e = makeArrayLinear comp sz (const e) {-# INLINE replicate #-} - -- | Iterate over an array with a ST action that is applied to each element and its index. -- -- @since 1.0.0 iterArrayLinearST_ :: Scheduler s () - -> Array r ix e -- ^ Array that is being loaded - -> (Int -> e -> ST s ()) -- ^ Function that writes an element into target array + -> Array r ix e + -- ^ Array that is being loaded + -> (Int -> e -> ST s ()) + -- ^ Function that writes an element into target array -> ST s () iterArrayLinearST_ scheduler arr uWrite = iterArrayLinearWithSetST_ scheduler arr uWrite $ \offset sz e -> - loopA_ offset (< (offset + unSz sz)) (+1) (`uWrite` e) + loopA_ offset (< (offset + unSz sz)) (+ 1) (`uWrite` e) {-# INLINE iterArrayLinearST_ #-} -- | Similar to `iterArrayLinearST_`. Except it also accepts a function that is @@ -418,10 +436,13 @@ -- @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 + -> 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_ #-} @@ -430,8 +451,8 @@ -- the same. -- -- @since 1.0.0 - unsafeLoadIntoST :: - Manifest r' e + unsafeLoadIntoST + :: Manifest r' e => MVector s r' e -> Array r ix e -> ST s (MArray s r' ix e) @@ -445,21 +466,22 @@ -- | Same as `unsafeLoadIntoST`, but respecting computation strategy. -- -- @since 1.0.0 - unsafeLoadIntoIO :: - Manifest r' e + unsafeLoadIntoIO + :: Manifest r' e => MVector RealWorld r' e -> Array r ix e -> IO (MArray RealWorld r' ix e) unsafeLoadIntoIO mvec arr = do let sz = outerSize arr mvec' <- resizeMVector mvec $ toLinearSz sz - withMassivScheduler_ (getComp arr) $ \scheduler -> stToIO $ - iterArrayLinearWithSetST_ scheduler arr (unsafeLinearWrite mvec') (unsafeLinearSet mvec') + withMassivScheduler_ (getComp arr) $ \scheduler -> + stToIO $ + iterArrayLinearWithSetST_ scheduler arr (unsafeLinearWrite mvec') (unsafeLinearSet mvec') pure $ unsafeResizeMArray sz mvec' {-# INLINE unsafeLoadIntoIO #-} -resizeMVector :: - (Manifest r e, PrimMonad f) +resizeMVector + :: (Manifest r e, PrimMonad f) => MVector (PrimState f) r e -> Sz1 -> f (MVector (PrimState f) r e) @@ -467,25 +489,29 @@ let mk = sizeOfMArray mvec in if k == mk then pure mvec - else if k < mk - then unsafeLinearShrink mvec k - else unsafeLinearGrow mvec k + else + if k < mk + then unsafeLinearShrink mvec k + else unsafeLinearGrow mvec k {-# INLINE resizeMVector #-} - class Load r ix e => StrideLoad r ix e where -- | Load an array into memory with stride. Default implementation requires an instance of -- `Source`. iterArrayLinearWithStrideST_ :: Scheduler s () - -> Stride ix -- ^ Stride to use - -> Sz ix -- ^ Size of the target array affected by the stride. - -> Array r ix e -- ^ Array that is being loaded - -> (Int -> e -> ST s ()) -- ^ Function that writes an element into target array + -> Stride ix + -- ^ Stride to use + -> Sz ix + -- ^ Size of the target array affected by the stride. + -> Array r ix e + -- ^ Array that is being loaded + -> (Int -> e -> ST s ()) + -- ^ Function that writes an element into target array -> ST s () default iterArrayLinearWithStrideST_ - :: Source r e => - Scheduler s () + :: Source r e + => Scheduler s () -> Stride ix -> Sz ix -> Array r ix e @@ -502,17 +528,18 @@ -- 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) +-- +-- unsafeLoadIntoWithStrideST :: -- TODO: this would remove Size constraint and allow DS and LN instances for vectors. +-- Manifest r' ix e +-- => Array r ix e +-- -> Stride ix -- ^ Stride to use +-- -> MArray RealWorld r' ix e +-- -> m (MArray RealWorld r' ix e) -- | Starting with massiv-1.0 `Mutable` and `Manifest` are synonymous. Since massiv-1.1 -- it is deprecated and will be removed in massiv-1.2 type Mutable r e = Manifest r e + {-# DEPRECATED Mutable "In favor of `Manifest`" #-} -- | Manifest arrays are backed by actual memory and values are looked up versus @@ -521,7 +548,6 @@ -- of changed a mutable `MArray` back into an immutable `Array` is called -- freezing. class Source r e => Manifest r e where - unsafeLinearIndexM :: Index ix => Array r ix e -> Int -> e -- | /O(1)/ - Get the size of a mutable array. @@ -540,7 +566,6 @@ -- @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 @@ -559,12 +584,23 @@ -- | Read an element at linear row-major index -- -- @since 0.1.0 - unsafeLinearRead :: (Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> Int -> m e + unsafeLinearRead + :: HAS_CALL_STACK + => (Index ix, PrimMonad m) + => MArray (PrimState m) r ix e + -> Int + -> m e -- | Write an element into mutable array with linear row-major index -- -- @since 0.1.0 - unsafeLinearWrite :: (Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> Int -> e -> m () + unsafeLinearWrite + :: HAS_CALL_STACK + => (Index ix, PrimMonad m) + => MArray (PrimState m) r ix e + -> Int + -> e + -> m () -- | Initialize mutable array to some default value. -- @@ -591,41 +627,61 @@ -- | Set all cells in the mutable array within the range to a specified value. -- -- @since 0.3.0 - unsafeLinearSet :: (Index ix, PrimMonad m) => - MArray (PrimState m) r ix e -> Ix1 -> Sz1 -> e -> m () + unsafeLinearSet + :: HAS_CALL_STACK + => (Index ix, PrimMonad m) + => MArray (PrimState m) r ix e + -> Ix1 + -> Sz1 + -> e + -> m () unsafeLinearSet marr offset len e = - loopA_ offset (< (offset + unSz len)) (+1) (\i -> unsafeLinearWrite marr i e) + loopA_ offset (< (offset + unSz len)) (+ 1) (\i -> unsafeLinearWrite marr i e) {-# INLINE unsafeLinearSet #-} -- | Copy part of one mutable array into another -- -- @since 0.3.6 - unsafeLinearCopy :: (Index ix', Index ix, PrimMonad m) => - MArray (PrimState m) r ix' e -- ^ Source mutable array - -> Ix1 -- ^ Starting index at source array - -> MArray (PrimState m) r ix e -- ^ Target mutable array - -> Ix1 -- ^ Starting index at target array - -> Sz1 -- ^ Number of elements to copy - -> m () + unsafeLinearCopy + :: HAS_CALL_STACK + => (Index ix', Index ix, PrimMonad m) + => MArray (PrimState m) r ix' e + -- ^ Source mutable array + -> Ix1 + -- ^ Starting index at source array + -> MArray (PrimState m) r ix e + -- ^ Target mutable array + -> Ix1 + -- ^ Starting index at target array + -> Sz1 + -- ^ Number of elements to copy + -> m () unsafeLinearCopy marrFrom iFrom marrTo iTo (SafeSz k) = do let delta = iTo - iFrom - loopA_ iFrom (< k + iFrom) (+1) $ \i -> + loopA_ iFrom (< k + iFrom) (+ 1) $ \i -> unsafeLinearRead marrFrom i >>= unsafeLinearWrite marrTo (i + delta) {-# INLINE unsafeLinearCopy #-} -- | Copy a part of a pure array into a mutable array -- -- @since 0.3.6 - unsafeArrayLinearCopy :: (Index ix', Index ix, PrimMonad m) => - Array r ix' e -- ^ Source pure array - -> Ix1 -- ^ Starting index at source array - -> MArray (PrimState m) r ix e -- ^ Target mutable array - -> Ix1 -- ^ Starting index at target array - -> Sz1 -- ^ Number of elements to copy - -> m () + unsafeArrayLinearCopy + :: HAS_CALL_STACK + => (Index ix', Index ix, PrimMonad m) + => Array r ix' e + -- ^ Source pure array + -> Ix1 + -- ^ Starting index at source array + -> MArray (PrimState m) r ix e + -- ^ Target mutable array + -> Ix1 + -- ^ Starting index at target array + -> Sz1 + -- ^ Number of elements to copy + -> m () unsafeArrayLinearCopy arrFrom iFrom marrTo iTo (SafeSz k) = do let delta = iTo - iFrom - loopA_ iFrom (< k + iFrom) (+1) $ \i -> + loopA_ iFrom (< k + iFrom) (+ 1) $ \i -> unsafeLinearWrite marrTo (i + delta) (unsafeLinearIndex arrFrom i) {-# INLINE unsafeArrayLinearCopy #-} @@ -634,8 +690,12 @@ -- no longer be used. -- -- @since 0.3.6 - unsafeLinearShrink :: (Index ix, PrimMonad m) => - MArray (PrimState m) r ix e -> Sz ix -> m (MArray (PrimState m) r ix e) + unsafeLinearShrink + :: HAS_CALL_STACK + => (Index ix, PrimMonad m) + => MArray (PrimState m) r ix e + -> Sz ix + -> m (MArray (PrimState m) r ix e) unsafeLinearShrink = unsafeDefaultLinearShrink {-# INLINE unsafeLinearShrink #-} @@ -644,17 +704,21 @@ -- should no longer be used. -- -- @since 0.3.6 - unsafeLinearGrow :: (Index ix, PrimMonad m) => - MArray (PrimState m) r ix e -> Sz ix -> m (MArray (PrimState m) r ix e) + unsafeLinearGrow + :: HAS_CALL_STACK + => (Index ix, PrimMonad m) + => MArray (PrimState m) r ix e + -> Sz ix + -> m (MArray (PrimState m) r ix e) unsafeLinearGrow marr sz = do marr' <- unsafeNew sz unsafeLinearCopy marr 0 marr' 0 $ SafeSz (totalElem (sizeOfMArray marr)) pure marr' {-# INLINE unsafeLinearGrow #-} - -unsafeDefaultLinearShrink :: - (Manifest r e, Index ix, PrimMonad m) +unsafeDefaultLinearShrink + :: HAS_CALL_STACK + => (Manifest r e, Index ix, PrimMonad m) => MArray (PrimState m) r ix e -> Sz ix -> m (MArray (PrimState m) r ix e) @@ -667,25 +731,38 @@ -- | Read an array element -- -- @since 0.1.0 -unsafeRead :: (Manifest r e, Index ix, PrimMonad m) => - MArray (PrimState m) r ix e -> ix -> m e +unsafeRead + :: HAS_CALL_STACK + => (Manifest r e, Index ix, PrimMonad m) + => MArray (PrimState m) r ix e + -> ix + -> m e unsafeRead marr = unsafeLinearRead marr . toLinearIndex (sizeOfMArray marr) {-# INLINE unsafeRead #-} -- | Write an element into array -- -- @since 0.1.0 -unsafeWrite :: (Manifest r e, Index ix, PrimMonad m) => - MArray (PrimState m) r ix e -> ix -> e -> m () +unsafeWrite + :: HAS_CALL_STACK + => (Manifest r e, Index ix, PrimMonad m) + => MArray (PrimState m) r ix e + -> ix + -> e + -> m () unsafeWrite marr = unsafeLinearWrite marr . toLinearIndex (sizeOfMArray marr) {-# INLINE unsafeWrite #-} - -- | Modify an element in the array with a monadic action. Returns the previous value. -- -- @since 0.4.0 -unsafeLinearModify :: (Manifest r e, Index ix, PrimMonad m) => - MArray (PrimState m) r ix e -> (e -> m e) -> Int -> m e +unsafeLinearModify + :: HAS_CALL_STACK + => (Manifest r e, Index ix, PrimMonad m) + => MArray (PrimState m) r ix e + -> (e -> m e) + -> Int + -> m e unsafeLinearModify !marr f !i = do v <- unsafeLinearRead marr i v' <- f v @@ -696,8 +773,13 @@ -- | Modify an element in the array with a monadic action. Returns the previous value. -- -- @since 0.4.0 -unsafeModify :: (Manifest r e, Index ix, PrimMonad m) => - MArray (PrimState m) r ix e -> (e -> m e) -> ix -> m e +unsafeModify + :: HAS_CALL_STACK + => (Manifest r e, Index ix, PrimMonad m) + => MArray (PrimState m) r ix e + -> (e -> m e) + -> ix + -> m e unsafeModify marr f ix = unsafeLinearModify marr f (toLinearIndex (sizeOfMArray marr) ix) {-# INLINE unsafeModify #-} @@ -705,19 +787,29 @@ -- values. -- -- @since 0.4.0 -unsafeSwap :: (Manifest r e, Index ix, PrimMonad m) => - MArray (PrimState m) r ix e -> ix -> ix -> m (e, e) +unsafeSwap + :: HAS_CALL_STACK + => (Manifest r e, Index ix, PrimMonad m) + => MArray (PrimState m) r ix e + -> ix + -> ix + -> m (e, e) unsafeSwap !marr !ix1 !ix2 = unsafeLinearSwap marr (toLinearIndex sz ix1) (toLinearIndex sz ix2) - where sz = sizeOfMArray marr + where + sz = sizeOfMArray marr {-# INLINE unsafeSwap #-} - -- | Swap two elements in a mutable array under the supplied linear indices. Returns the -- previous values. -- -- @since 0.4.0 -unsafeLinearSwap :: (Manifest r e, Index ix, PrimMonad m) => - MArray (PrimState m) r ix e -> Int -> Int -> m (e, e) +unsafeLinearSwap + :: HAS_CALL_STACK + => (Manifest r e, Index ix, PrimMonad m) + => MArray (PrimState m) r ix e + -> Int + -> Int + -> m (e, e) unsafeLinearSwap !marr !i1 !i2 = do val1 <- unsafeLinearRead marr i1 val2 <- unsafeLinearRead marr i2 @@ -726,20 +818,16 @@ return (val1, val2) {-# INLINE unsafeLinearSwap #-} - class (IsList (Array r ix e), Load r ix e) => Ragged r ix e where - generateRaggedM :: Monad m => Comp -> Sz ix -> (ix -> m e) -> m (Array r ix e) flattenRagged :: Array r ix e -> Vector r e - loadRaggedST :: - Scheduler s () -> Array r ix e -> (Ix1 -> e -> ST s ()) -> Ix1 -> Ix1 -> Sz ix -> ST s () + loadRaggedST + :: Scheduler s () -> Array r ix e -> (Ix1 -> e -> ST s ()) -> Ix1 -> Ix1 -> Sz ix -> ST s () raggedFormat :: (e -> String) -> String -> Array r ix e -> String - - -- | Create an Array with no elements. By itself it is not particularly useful, but it serves as a -- nice base for constructing larger arrays. -- @@ -753,8 +841,9 @@ -- [ 4.0, 5.0, 22.0 ] -- -- @since 0.3.0 -empty :: - forall r ix e. Load r ix e +empty + :: forall r ix e + . Load r ix e => Array r ix e empty = makeArray Seq zeroSz (const (throwImpossible Uninitialized)) {-# INLINE empty #-} @@ -782,14 +871,15 @@ -- ] -- -- @since 0.1.0 -singleton :: - forall r ix e. Load r ix e - => e -- ^ The only element +singleton + :: forall r ix e + . Load r ix e + => e + -- ^ The only element -> Array r ix e singleton = makeArray Seq oneSz . const {-# INLINE singleton #-} - infixl 4 !, !?, ?? -- | /O(1)/ - Infix version of 'index''. @@ -807,15 +897,15 @@ -- 3 -- -- @since 0.1.0 -(!) :: - forall r ix e. (HasCallStack, Manifest r e, Index ix) +(!) + :: forall r ix e + . (HasCallStack, Manifest r e, Index ix) => Array r ix e -> ix -> e (!) arr = throwEither . evaluateM arr {-# INLINE (!) #-} - -- | /O(1)/ - Infix version of `indexM`. -- -- /__Exceptions__/: `IndexOutOfBoundsException` @@ -838,15 +928,15 @@ -- Nothing -- -- @since 0.1.0 -(!?) :: - forall r ix e m. (Index ix, Manifest r e, MonadThrow m) +(!?) + :: forall r ix e m + . (Index ix, Manifest r e, MonadThrow m) => Array r ix e -> ix -> m e (!?) = indexM {-# INLINE (!?) #-} - -- | /O(1)/ - Lookup an element in the array, where array itself is wrapped with -- `MonadThrow`. This operator is useful when used together with slicing or other -- functions that can fail. @@ -984,8 +1074,6 @@ evaluate' arr ix = throwEither (evaluateM arr ix) {-# INLINE evaluate' #-} - - -- | /O(1)/ - Check if array has elements. -- -- ==== __Examples__ @@ -1001,8 +1089,6 @@ isNotNull = not . isNull {-# INLINE isNotNull #-} - - -- | /O(1)/ - Check if array has elements. -- -- ==== __Examples__ @@ -1015,10 +1101,9 @@ -- -- @since 1.0.0 isEmpty :: (Index ix, Size r) => Array r ix e -> Bool -isEmpty = (==0) . elemsCount +isEmpty = (== 0) . elemsCount {-# INLINE isEmpty #-} - -- | /O(1)/ - Check if array has elements. -- -- ==== __Examples__ @@ -1034,7 +1119,6 @@ isNotEmpty = not . isEmpty {-# INLINE isNotEmpty #-} - -- | /O(1)/ - Get the number of elements in the array. -- -- ==== __Examples__ @@ -1047,7 +1131,6 @@ elemsCount :: (Index ix, Size r) => Array r ix e -> Int elemsCount = totalElem . size {-# INLINE elemsCount #-} - inline0 :: (a -> b) -> a -> b inline0 f = f 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,8 +1,9 @@ {-# LANGUAGE CPP #-} {-# LANGUAGE GADTs #-} -{-# LANGUAGE LambdaCase #-} {-# LANGUAGE ImplicitParams #-} +{-# LANGUAGE LambdaCase #-} {-# OPTIONS_GHC -fno-warn-orphans #-} + -- | -- Module : Data.Massiv.Core.Exception -- Copyright : (c) Alexey Kuleshevich 2019-2022 @@ -10,23 +11,22 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Core.Exception - ( throwImpossible - , throwEither - , Uninitialized(..) - , guardNumberOfElements - , Exception(..) - , SomeException - , HasCallStack - ) where +module Data.Massiv.Core.Exception ( + throwImpossible, + throwEither, + Uninitialized (..), + guardNumberOfElements, + Exception (..), + SomeException, + HasCallStack, +) where import Control.Exception import Control.Monad import Control.Monad.Catch import Data.Massiv.Core.Index.Internal -import GHC.Stack import GHC.Exception +import GHC.Stack #if !MIN_VERSION_exceptions(0, 10, 3) import Control.Monad.ST (ST) @@ -44,11 +44,10 @@ 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." - + " ImpossibleException (" + ++ displayException exc + ++ "): Either one of the unsafe functions was used or it is a bug in the library. " + ++ "In latter case please report this error." {-# NOINLINE throwImpossible #-} -- | Throw an error on `Left` or produce the result on `Right`. Exception type is lost, so @@ -65,12 +64,11 @@ -- | An error that gets thrown when an unitialized element of a boxed array gets accessed. Can only -- happen when array was constructed with `Data.Massiv.Array.Unsafe.unsafeNew`. -data Uninitialized = Uninitialized deriving Show +data Uninitialized = Uninitialized deriving (Show) instance Exception Uninitialized where displayException Uninitialized = "Array element is uninitialized" - -- | Throw `SizeElementsMismatchException` whenever number of elements in both sizes do -- not match. -- @@ -79,4 +77,3 @@ guardNumberOfElements sz sz' = unless (totalElem sz == totalElem sz') $ throwM $ SizeElementsMismatchException sz sz' {-# INLINE guardNumberOfElements #-} - diff --git a/src/Data/Massiv/Core/Index.hs b/src/Data/Massiv/Core/Index.hs --- a/src/Data/Massiv/Core/Index.hs +++ b/src/Data/Massiv/Core/Index.hs @@ -1,8 +1,10 @@ {-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} {-# LANGUAGE DataKinds #-} +{-# LANGUAGE ExplicitNamespaces #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE ExplicitNamespaces #-} + -- | -- Module : Data.Massiv.Core.Index -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -10,110 +12,118 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Core.Index - ( Ix0(..) - , type Ix1 - , pattern Ix1 - , type Ix2(Ix2, (:.)) - , IxN((:>), Ix3, Ix4, Ix5) - , HighIxN - , type Ix3 - , type Ix4 - , type Ix5 - , Ix +module Data.Massiv.Core.Index ( + Ix0 (..), + type Ix1, + pattern Ix1, + type Ix2 (Ix2, (:.)), + IxN ((:>), Ix3, Ix4, Ix5), + HighIxN, + type Ix3, + type Ix4, + type Ix5, + Ix, + -- ** Size - , type Sz1 - , type Sz2 - , type Sz3 - , type Sz4 - , type Sz5 - , Sz(Sz, Sz1, Sz2, Sz3, Sz4, Sz5) - , unSz - , zeroSz - , oneSz - , liftSz - , liftSz2 - , consSz - , unconsSz - , snocSz - , unsnocSz - , setSzM - , insertSzM - , pullOutSzM - , toLinearSz - , mkSzM + type Sz1, + type Sz2, + type Sz3, + type Sz4, + type Sz5, + Sz (Sz, Sz1, Sz2, Sz3, Sz4, Sz5), + unSz, + zeroSz, + oneSz, + liftSz, + liftSz2, + consSz, + unconsSz, + snocSz, + unsnocSz, + setSzM, + insertSzM, + pullOutSzM, + toLinearSz, + mkSzM, + -- ** Dimension - , Dim(..) - , Dimension(Dim1, Dim2, Dim3, Dim4, Dim5, DimN) - , IsIndexDimension - , IsDimValid - , ReportInvalidDim + Dim (..), + Dimension (Dim1, Dim2, Dim3, Dim4, Dim5, DimN), + IsIndexDimension, + IsDimValid, + ReportInvalidDim, + -- ** Stride - , Stride(Stride) - , unStride - , toLinearIndexStride - , strideStart - , strideSize - , oneStride + Stride (Stride), + unStride, + toLinearIndexStride, + strideStart, + strideSize, + oneStride, + -- ** Border - , Border(..) - , handleBorderIndex + Border (..), + handleBorderIndex, + -- ** Index functions - , Lower - , Index(..) - , zeroIndex - , oneIndex - , isZeroSz - , isNotZeroSz - , headDim - , tailDim - , lastDim - , initDim - , getDim' - , setDim' - , modifyDim' - , dropDimM - , dropDim' - , pullOutDim' - , insertDim' - , fromDimension - , getDimension - , setDimension - , modifyDimension - , dropDimension - , pullOutDimension - , insertDimension + Lower, + Index (..), + zeroIndex, + oneIndex, + isZeroSz, + isNotZeroSz, + headDim, + tailDim, + lastDim, + initDim, + getDim', + setDim', + modifyDim', + dropDimM, + dropDim', + pullOutDim', + insertDim', + fromDimension, + getDimension, + setDimension, + modifyDimension, + dropDimension, + pullOutDimension, + insertDimension, + -- * Iterators - , iter - , iterA_ - , iterM_ - , iterLinearM - , iterLinearM_ - , module Data.Massiv.Core.Loop - , module Data.Massiv.Core.Index.Iterator - , module Data.Massiv.Core.Index.Tuple + iter, + iterA_, + iterM_, + iterLinearM, + iterLinearM_, + module Data.Massiv.Core.Loop, + module Data.Massiv.Core.Index.Iterator, + module Data.Massiv.Core.Index.Tuple, + -- * Exceptions - , IndexException(..) - , SizeException(..) - , ShapeException(..) - , guardNumberOfElements - , indexWith - ) where + IndexException (..), + SizeException (..), + ShapeException (..), + guardNumberOfElements, + indexAssert, + indexWith, +) where import Control.DeepSeq -import Control.Monad.Catch (MonadThrow(..)) +import Control.Monad.Catch (MonadThrow (..)) import Data.Coerce import Data.Functor.Identity (runIdentity) import Data.Massiv.Core.Exception import Data.Massiv.Core.Index.Internal +import Data.Massiv.Core.Index.Iterator import Data.Massiv.Core.Index.Ix import Data.Massiv.Core.Index.Stride import Data.Massiv.Core.Index.Tuple -import Data.Massiv.Core.Index.Iterator import Data.Massiv.Core.Loop import GHC.TypeLits +#include "massiv.h" -- | 1-dimensional type synonym for size. -- @@ -140,57 +150,55 @@ -- @since 0.3.0 type Sz5 = Sz Ix5 - -- | Approach to be used near the borders during various transformations. -- Whenever a function needs information not only about an element of interest, but -- also about it's neighbors, it will go out of bounds near the array edges, -- hence is this set of approaches that specify how to handle such situation. -data Border e = - Fill e -- ^ Fill in a constant element. - -- - -- @ - -- outside | Array | outside - -- ('Fill' 0) : 0 0 0 0 | 1 2 3 4 | 0 0 0 0 - -- @ - -- - | Wrap -- ^ Wrap around from the opposite border of the array. - -- - -- @ - -- outside | Array | outside - -- 'Wrap' : 1 2 3 4 | 1 2 3 4 | 1 2 3 4 - -- @ - -- - | Edge -- ^ Replicate the element at the edge. - -- - -- @ - -- outside | Array | outside - -- 'Edge' : 1 1 1 1 | 1 2 3 4 | 4 4 4 4 - -- @ - -- - | Reflect -- ^ Mirror like reflection. - -- - -- @ - -- outside | Array | outside - -- 'Reflect' : 4 3 2 1 | 1 2 3 4 | 4 3 2 1 - -- @ - -- - | Continue -- ^ Also mirror like reflection, but without repeating the edge element. - -- - -- @ - -- outside | Array | outside - -- 'Continue' : 1 4 3 2 | 1 2 3 4 | 3 2 1 4 - -- @ - -- +data Border e + = -- | Fill in a constant element. + -- + -- @ + -- outside | Array | outside + -- ('Fill' 0) : 0 0 0 0 | 1 2 3 4 | 0 0 0 0 + -- @ + Fill e + | -- | Wrap around from the opposite border of the array. + -- + -- @ + -- outside | Array | outside + -- 'Wrap' : 1 2 3 4 | 1 2 3 4 | 1 2 3 4 + -- @ + Wrap + | -- | Replicate the element at the edge. + -- + -- @ + -- outside | Array | outside + -- 'Edge' : 1 1 1 1 | 1 2 3 4 | 4 4 4 4 + -- @ + Edge + | -- | Mirror like reflection. + -- + -- @ + -- outside | Array | outside + -- 'Reflect' : 4 3 2 1 | 1 2 3 4 | 4 3 2 1 + -- @ + Reflect + | -- | Also mirror like reflection, but without repeating the edge element. + -- + -- @ + -- outside | Array | outside + -- 'Continue' : 1 4 3 2 | 1 2 3 4 | 3 2 1 4 + -- @ + Continue deriving (Eq, Show) instance NFData e => NFData (Border e) where rnf b = case b of - Fill e -> rnf e - Wrap -> () - Edge -> () - Reflect -> () - Continue -> () - + Fill e -> rnf e + Wrap -> () + Edge -> () + Reflect -> () + Continue -> () -- | Apply a border resolution technique to an index -- @@ -204,25 +212,41 @@ -- 1 :. 2 -- -- @since 0.1.0 -handleBorderIndex :: - Index ix - => Border e -- ^ Broder resolution technique - -> Sz ix -- ^ Size - -> (ix -> e) -- ^ Index function that produces an element - -> ix -- ^ Index +handleBorderIndex + :: Index ix + => Border e + -- ^ Broder resolution technique + -> Sz ix + -- ^ Size + -> (ix -> e) + -- ^ Index function that produces an element + -> ix + -- ^ Index -> e handleBorderIndex border !sz getVal !ix = case border of Fill val -> if isSafeIndex sz ix then getVal ix else val - Wrap -> getVal (repairIndex sz ix wrap wrap) - Edge -> getVal (repairIndex sz ix (const (const 0)) (\ (SafeSz k) _ -> k - 1)) - Reflect -> getVal (repairIndex sz ix (\ (SafeSz k) !i -> (abs i - 1) `mod` k) - (\ (SafeSz k) !i -> (-i - 1) `mod` k)) - Continue -> getVal (repairIndex sz ix (\ (SafeSz k) !i -> abs i `mod` k) - (\ (SafeSz k) !i -> (-i - 2) `mod` k)) - - where wrap (SafeSz k) i = i `mod` k - {-# INLINE [1] wrap #-} + Wrap -> getVal (repairIndex sz ix wrap wrap) + Edge -> getVal (repairIndex sz ix (const (const 0)) (\(SafeSz k) _ -> k - 1)) + Reflect -> + getVal + ( repairIndex + sz + ix + (\(SafeSz k) !i -> (abs i - 1) `mod` k) + (\(SafeSz k) !i -> (-i - 1) `mod` k) + ) + Continue -> + getVal + ( repairIndex + sz + ix + (\(SafeSz k) !i -> abs i `mod` k) + (\(SafeSz k) !i -> (-i - 2) `mod` k) + ) + where + wrap (SafeSz k) i = i `mod` k + {-# INLINE [1] wrap #-} {-# INLINE [1] handleBorderIndex #-} -- | Index with all zeros @@ -255,6 +279,7 @@ 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) @@ -272,7 +297,6 @@ isZeroSz = not . isNotZeroSz {-# INLINE [1] isZeroSz #-} - -- | Convert a size to a linear size. -- -- @since 0.5.8 @@ -470,7 +494,6 @@ getDimension ix = getDim' ix . fromDimension {-# INLINE [1] getDimension #-} - -- | Type safe way of dropping a particular dimension, thus lowering index -- dimensionality. -- @@ -527,19 +550,25 @@ -- 3615 -- -- @since 0.1.0 -iter :: Index ix - => ix -- ^ Start index - -> ix -- ^ End index - -> ix -- ^ Increment - -> (Int -> Int -> Bool) -- ^ Continuation condition - -> a -- ^ Accumulator - -> (ix -> a -> a) -- ^ Iterating function +iter + :: Index ix + => ix + -- ^ Start index + -> ix + -- ^ End index + -> ix + -- ^ Increment + -> (Int -> Int -> Bool) + -- ^ Continuation condition -> a + -- ^ Accumulator + -> (ix -> a -> a) + -- ^ Iterating function + -> a iter sIx eIx incIx cond acc f = runIdentity $ iterM sIx eIx incIx cond acc (\ix -> return . f ix) {-# INLINE iter #-} - -- | Iterate over N-dimensional space linearly from start to end in row-major fashion with an -- accumulator -- @@ -553,15 +582,22 @@ -- 103 -- -- @since 0.1.0 -iterLinearM :: (Index ix, Monad m) - => Sz ix -- ^ Size - -> Int -- ^ Linear start (must be non-negative) - -> Int -- ^ Linear end (must be less than or equal to @`totalElem` sz@) - -> Int -- ^ Increment (must not be zero) - -> (Int -> Int -> Bool) -- ^ Continuation condition (continue if @True@) - -> a -- ^ Accumulator - -> (Int -> ix -> a -> m a) - -> m a +iterLinearM + :: (Index ix, Monad m) + => Sz ix + -- ^ Size + -> Int + -- ^ Linear start (must be non-negative) + -> Int + -- ^ Linear end (must be less than or equal to @`totalElem` sz@) + -> Int + -- ^ Increment (must not be zero) + -> (Int -> Int -> Bool) + -- ^ Continuation condition (continue if @True@) + -> a + -- ^ Accumulator + -> (Int -> ix -> a -> m a) + -> m a iterLinearM !sz !k0 !k1 !inc cond !acc f = loopM k0 (`cond` k1) (+ inc) acc $ \ !i !acc0 -> f i (fromLinearIndex sz i) acc0 {-# INLINE iterLinearM #-} @@ -577,43 +613,76 @@ -- True -- -- @since 0.1.0 -iterLinearM_ :: (Index ix, Monad m) => - Sz ix -- ^ Size - -> Int -- ^ Start (must be non-negative) - -> Int -- ^ End - -> Int -- ^ Increment (must not be zero) - -> (Int -> Int -> Bool) -- ^ Continuation condition (continue if @True@) - -> (Int -> ix -> m ()) -- ^ Monadic action that takes index in both forms - -> m () +iterLinearM_ + :: (Index ix, Monad m) + => Sz ix + -- ^ Size + -> Int + -- ^ Start (must be non-negative) + -> Int + -- ^ End + -> Int + -- ^ Increment (must not be zero) + -> (Int -> Int -> Bool) + -- ^ Continuation condition (continue if @True@) + -> (Int -> ix -> m ()) + -- ^ Monadic action that takes index in both forms + -> m () iterLinearM_ sz !k0 !k1 !inc cond f = loopA_ k0 (`cond` k1) (+ inc) $ \ !i -> f i (fromLinearIndex sz i) {-# INLINE iterLinearM_ #-} +-- | This is used by the @unsafe-checks@ cabal flag. +-- +-- @since 1.1.0 +#ifdef MASSIV_UNSAFE_CHECKS +indexAssert :: (HasCallStack, Index ix) => String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e +indexAssert funName getSize f arr ix + | isSafeIndex sz ix = f arr ix + | otherwise = _errorIx ("<" ++ funName ++ ">") sz ix + where + sz = getSize arr +#else +indexAssert :: String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e +indexAssert _funName _getSize f arr ix = f arr ix +#endif +{-# INLINE indexAssert #-} + -- | This is used by @INDEX_CHECK@ macro and thus used whenever the @unsafe-checks@ cabal -- flag is on. -- -- @since 0.4.0 -indexWith :: - Index ix - => String -- ^ Source file name, eg. __FILE__ - -> Int -- ^ Line number in th source file, eg. __LINE__ +indexWith + :: Index ix + => String + -- ^ Source file name, eg. __FILE__ + -> Int + -- ^ Line number in th source file, eg. __LINE__ -> String - -> (arr -> Sz ix) -- ^ Get size of the array - -> (arr -> ix -> e) -- ^ Indexing function - -> arr -- ^ Array - -> ix -- ^ Index + -> (arr -> Sz ix) + -- ^ Get size of the array + -> (arr -> ix -> e) + -- ^ Indexing function + -> arr + -- ^ Array + -> ix + -- ^ Index -> e indexWith fileName lineNo funName getSize f arr ix | isSafeIndex sz ix = f arr ix - | otherwise = errorIx ("<" ++ fileName ++ ":" ++ show lineNo ++ "> " ++ funName) sz ix + | otherwise = _errorIx ("<" ++ fileName ++ ":" ++ show lineNo ++ "> " ++ funName) sz ix where sz = getSize arr +{-# DEPRECATED indexWith "In favor of `indexAssert` that uses HasCallStack" #-} --- | Helper function for throwing out of bounds error. Used by `indexWith` -errorIx :: (Show ix, Show ix') => String -> ix -> ix' -> a -errorIx fName sz ix = +-- | Helper function for throwing out of bounds error. Used by `indexAssert` +_errorIx :: (HasCallStack, Show ix, Show ix') => String -> ix -> ix' -> a +_errorIx fName sz ix = error $ - fName ++ - ": Index out of bounds: (" ++ show ix ++ ") for Array of size: (" ++ show sz ++ ")" -{-# NOINLINE errorIx #-} - + fName + ++ ": Index out of bounds: (" + ++ show ix + ++ ") for Array of size: (" + ++ show sz + ++ ")" +{-# NOINLINE _errorIx #-} diff --git a/src/Data/Massiv/Core/Index/Internal.hs b/src/Data/Massiv/Core/Index/Internal.hs --- a/src/Data/Massiv/Core/Index/Internal.hs +++ b/src/Data/Massiv/Core/Index/Internal.hs @@ -15,6 +15,7 @@ {-# LANGUAGE TypeOperators #-} {-# LANGUAGE UndecidableInstances #-} {-# OPTIONS_GHC -Wno-unticked-promoted-constructors #-} + -- | -- Module : Data.Massiv.Core.Index.Internal -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -22,53 +23,52 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Core.Index.Internal - ( Sz(SafeSz) - , pattern Sz - , pattern Sz1 - , unSz - , zeroSz - , oneSz - , liftSz - , liftSz2 - , consSz - , unconsSz - , snocSz - , unsnocSz - , setSzM - , insertSzM - , pullOutSzM - , mkSzM - , Dim(..) - , Dimension(DimN) - , pattern Dim1 - , pattern Dim2 - , pattern Dim3 - , pattern Dim4 - , pattern Dim5 - , IsIndexDimension - , IsDimValid - , ReportInvalidDim - , Lower - , Index(..) - , iterA_ - , iterM_ - , Ix0(..) - , type Ix1 - , pattern Ix1 - , IndexException(..) - , SizeException(..) - , ShapeException(..) - , showsPrecWrapped - ) where +module Data.Massiv.Core.Index.Internal ( + Sz (SafeSz), + pattern Sz, + pattern Sz1, + unSz, + zeroSz, + oneSz, + liftSz, + liftSz2, + consSz, + unconsSz, + snocSz, + unsnocSz, + setSzM, + insertSzM, + pullOutSzM, + mkSzM, + Dim (..), + Dimension (DimN), + pattern Dim1, + pattern Dim2, + pattern Dim3, + pattern Dim4, + pattern Dim5, + IsIndexDimension, + IsDimValid, + ReportInvalidDim, + Lower, + Index (..), + iterA_, + iterM_, + Ix0 (..), + type Ix1, + pattern Ix1, + IndexException (..), + SizeException (..), + ShapeException (..), + showsPrecWrapped, +) where +import Control.DeepSeq +import Control.Exception (Exception (..), throw) +import Control.Monad (void, when) +import Control.Monad.Catch (MonadThrow (..)) import Control.Monad.ST import Control.Scheduler -import Control.DeepSeq -import Control.Exception (Exception(..), throw) -import Control.Monad (when, void) -import Control.Monad.Catch (MonadThrow(..)) import Data.Coerce import Data.Kind import Data.Massiv.Core.Loop @@ -107,12 +107,12 @@ -- [#114](https://github.com/lehins/massiv/issues/114). -- -- @since 0.3.0 -newtype Sz ix = - SafeSz ix - -- ^ Safe size constructor. It is unsafe to use it without making sure that it does not contain - -- negative components. Use `Data.Massiv.Core.Index.Sz` pattern instead. - -- - -- @since 0.3.0 +newtype Sz ix + = -- | Safe size constructor. It is unsafe to use it without making sure that it does not contain + -- negative components. Use `Data.Massiv.Core.Index.Sz` pattern instead. + -- + -- @since 0.3.0 + SafeSz ix deriving (Eq, Ord, NFData) -- | A safe bidirectional pattern synonym for `Sz` construction that will make sure that none of @@ -120,18 +120,21 @@ -- -- @since 0.3.0 pattern Sz :: Index ix => ix -> Sz ix -pattern Sz ix <- SafeSz ix where - Sz ix = SafeSz (liftIndex (max 0) ix) +pattern Sz ix <- SafeSz ix + where + Sz ix = SafeSz (liftIndex (max 0) ix) + {-# COMPLETE Sz #-} -- | 1-dimensional size constructor. Especially useful with literals: @(Sz1 5) == Sz (5 :: Int)@. -- -- @since 0.3.0 pattern Sz1 :: Ix1 -> Sz Ix1 -pattern Sz1 ix <- SafeSz ix where - Sz1 ix = SafeSz (max 0 ix) -{-# COMPLETE Sz1 #-} +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 @@ -147,10 +150,10 @@ showsPrec n sz@(SafeSz usz) = showsPrecWrapped n (str ++) where str = - "Sz" ++ - case unDim (dimensions sz) of - 1 -> "1 " ++ show usz - _ -> " (" ++ shows usz ")" + "Sz" + ++ case unDim (dimensions sz) of + 1 -> "1 " ++ show usz + _ -> " (" ++ shows usz ")" -- | Calling `negate` is an error. instance (Num ix, Index ix) => Num (Sz ix) where @@ -165,9 +168,10 @@ 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" + 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 #-} @@ -188,8 +192,6 @@ Sz ix <$ foldlIndex (\acc i -> acc >>= guardNegativeOverflow i) (pure 1) ix {-# INLINE mkSzM #-} - - -- | Function for unwrapping `Sz`. -- -- ==== __Example__ @@ -229,7 +231,6 @@ oneSz = SafeSz (pureIndex 1) {-# INLINE oneSz #-} - -- | Same as `liftIndex`, but for `Sz` -- -- ==== __Example__ @@ -256,7 +257,6 @@ liftSz2 f sz1 sz2 = Sz (liftIndex2 f (coerce sz1) (coerce sz2)) {-# INLINE liftSz2 #-} - -- | Same as `consDim`, but for `Sz` -- -- ==== __Example__ @@ -270,7 +270,6 @@ consSz (SafeSz i) (SafeSz ix) = SafeSz (consDim i ix) {-# INLINE consSz #-} - -- | Same as `snocDim`, but for `Sz` -- -- ==== __Example__ @@ -353,11 +352,10 @@ pullOutSzM (SafeSz sz) = fmap coerce . pullOutDimM sz {-# INLINE pullOutSzM #-} - -- | A way to select Array dimension at a value level. -- -- @since 0.1.0 -newtype Dim = Dim { unDim :: Int } deriving (Eq, Ord, Num, Real, Integral, Enum, NFData) +newtype Dim = Dim {unDim :: Int} deriving (Eq, Ord, Num, Real, Integral, Enum, NFData) instance Show Dim where show (Dim d) = "(Dim " ++ show d ++ ")" @@ -406,26 +404,28 @@ pattern Dim5 :: Dimension 5 pattern Dim5 = DimN - -- | A type level constraint that ensures index is indeed valid and that supplied dimension can be -- safely used with it. -- -- @since 0.2.4 type IsIndexDimension ix n = (1 <= n, n <= Dimensions ix, Index ix, KnownNat n) - -- | This type family will always point to a type for a dimension that is one lower than the type -- argument. -- -- @since 0.1.0 type family Lower ix :: Type - type family ReportInvalidDim (dims :: Nat) (n :: Nat) isNotZero isLess :: Bool where ReportInvalidDim dims n True True = True ReportInvalidDim dims n True False = - TypeError (Text "Dimension " :<>: ShowType n :<>: Text " is higher than " :<>: - Text "the maximum expected " :<>: ShowType dims) + TypeError + ( Text "Dimension " + :<>: ShowType n + :<>: Text " is higher than " + :<>: Text "the maximum expected " + :<>: ShowType dims + ) ReportInvalidDim dims n False isLess = TypeError (Text "Zero dimensional indices are not supported") @@ -435,18 +435,19 @@ -- | This is bread and butter of multi-dimensional array indexing. It is unlikely that any of the -- functions in this class will be useful to a regular user, unless general algorithms are being -- implemented that do span multiple dimensions. -class ( Eq ix - , Ord ix - , Show ix - , NFData ix - , Typeable ix - , Eq (Lower ix) - , Ord (Lower ix) - , Show (Lower ix) - , NFData (Lower ix) - , KnownNat (Dimensions ix) - ) => - Index ix +class + ( Eq ix + , Ord ix + , Show ix + , NFData ix + , Typeable ix + , Eq (Lower ix) + , Ord (Lower ix) + , Show (Lower ix) + , NFData (Lower ix) + , KnownNat (Dimensions ix) + ) => + Index ix where -- | Type level information on how many dimensions this index has. -- @@ -535,23 +536,33 @@ -- | Perform a left fold over the index foldlIndex :: (a -> Int -> a) -> a -> ix -> a - default foldlIndex :: Index (Lower ix) => - (a -> Int -> a) -> a -> ix -> a + default foldlIndex + :: Index (Lower ix) + => (a -> Int -> a) + -> a + -> ix + -> a foldlIndex f !acc !ix = foldlIndex f (f acc i0) ixL where !(i0, ixL) = unconsDim ix {-# INLINE [1] foldlIndex #-} -- TODO: implement in terms of foldlIndex and pull out of the class + -- | Check whether index is positive and is within the size. -- -- @since 0.1.0 - isSafeIndex :: - Sz ix -- ^ Size - -> ix -- ^ Index + isSafeIndex + :: Sz ix + -- ^ Size + -> ix + -- ^ Index -> Bool - default isSafeIndex :: Index (Lower ix) => - Sz ix -> ix -> Bool + default isSafeIndex + :: Index (Lower ix) + => Sz ix + -> ix + -> Bool isSafeIndex sz !ix = isSafeIndex n0 i0 && isSafeIndex szL ixL where !(n0, szL) = unconsSz sz @@ -561,9 +572,11 @@ -- | Convert linear index from size and index -- -- @since 0.1.0 - toLinearIndex :: - Sz ix -- ^ Size - -> ix -- ^ Index + toLinearIndex + :: Sz ix + -- ^ Size + -> ix + -- ^ Index -> Ix1 default toLinearIndex :: Index (Lower ix) => Sz ix -> ix -> Ix1 toLinearIndex (SafeSz sz) !ix = toLinearIndex (SafeSz szL) ixL * n + i @@ -611,14 +624,23 @@ -- that will be invoked whenever index (2nd arg) is outsize the supplied size (1st arg) -- -- @since 0.1.0 - repairIndex :: - Sz ix -- ^ Size - -> ix -- ^ Index - -> (Sz Int -> Int -> Int) -- ^ Repair when below zero - -> (Sz Int -> Int -> Int) -- ^ Repair when higher than size + repairIndex + :: Sz ix + -- ^ Size -> ix - default repairIndex :: Index (Lower ix) => - Sz ix -> ix -> (Sz Int -> Int -> Int) -> (Sz Int -> Int -> Int) -> ix + -- ^ Index + -> (Sz Int -> Int -> Int) + -- ^ Repair when below zero + -> (Sz Int -> Int -> Int) + -- ^ Repair when higher than size + -> ix + default repairIndex + :: Index (Lower ix) + => Sz ix + -> ix + -> (Sz Int -> Int -> Int) + -> (Sz Int -> Int -> Int) + -> ix repairIndex sz !ix rBelow rOver = consDim (repairIndex n i rBelow rOver) (repairIndex szL ixL rBelow rOver) where @@ -629,17 +651,30 @@ -- | This function is what makes it possible to iterate over an array of any dimension. -- -- @since 0.1.0 - iterM :: - Monad m - => ix -- ^ Start index - -> ix -- ^ End index - -> ix -- ^ Increment - -> (Int -> Int -> Bool) -- ^ Continue iterating while predicate is True (eg. until end of row) - -> a -- ^ Initial value for an accumulator - -> (ix -> a -> m a) -- ^ Accumulator function + iterM + :: Monad m + => ix + -- ^ Start index + -> ix + -- ^ End index + -> ix + -- ^ Increment + -> (Int -> Int -> Bool) + -- ^ Continue iterating while predicate is True (eg. until end of row) + -> a + -- ^ Initial value for an accumulator + -> (ix -> a -> m a) + -- ^ Accumulator function -> m a - default iterM :: (Index (Lower ix), Monad m) => - ix -> ix -> ix -> (Int -> Int -> Bool) -> a -> (ix -> a -> m a) -> m a + default iterM + :: (Index (Lower ix), Monad m) + => ix + -> ix + -> ix + -> (Int -> Int -> Bool) + -> a + -> (ix -> a -> m a) + -> m a iterM !sIx eIx !incIx cond !acc f = loopM s (`cond` e) (+ inc) acc $ \ !i !acc0 -> iterM sIxL eIxL incIxL cond acc0 $ \ !ix -> f (consDim i ix) @@ -649,25 +684,35 @@ !(inc, incIxL) = unconsDim incIx {-# INLINE iterM #-} - iterRowMajorST :: Int -- ^ Scheduler multiplying factor. Must be positive - -> Scheduler s a -- ^ Scheduler to use - -> ix -- ^ Start index - -> ix -- ^ Stride - -> Sz ix -- ^ Size - -> a -- ^ Initial accumulator - -> (a -> ST s (a, a)) -- ^ Function that splits accumulator for each scheduled job. - -> (ix -> a -> ST s a) -- ^ Action - -> ST s a - default iterRowMajorST :: Index (Lower ix) - => Int - -> Scheduler s a - -> ix - -> ix - -> Sz ix - -> a - -> (a -> ST s (a, a)) - -> (ix -> a -> ST s a) - -> ST s a + iterRowMajorST + :: Int + -- ^ Scheduler multiplying factor. Must be positive + -> Scheduler s a + -- ^ Scheduler to use + -> ix + -- ^ Start index + -> ix + -- ^ Stride + -> Sz ix + -- ^ Size + -> a + -- ^ Initial accumulator + -> (a -> ST s (a, a)) + -- ^ Function that splits accumulator for each scheduled job. + -> (ix -> a -> ST s a) + -- ^ Action + -> ST s a + default iterRowMajorST + :: Index (Lower ix) + => Int + -> Scheduler s a + -> ix + -> ix + -> Sz ix + -> a + -> (a -> ST s (a, a)) + -> (ix -> a -> ST s a) + -> ST s a iterRowMajorST !fact scheduler ixStart ixStride sz initAcc splitAcc f = do let !(SafeSz n, szL@(SafeSz nL)) = unconsSz sz if n > 0 @@ -679,11 +724,10 @@ let !newFact = 1 + (fact `quot` n) loopM start (< start + n * stride) (+ stride) initAcc $ \j acc -> iterRowMajorST newFact scheduler ixL sL szL acc splitAcc (f . consDim j) - else - splitWorkWithFactorST fact scheduler start stride n initAcc splitAcc $ - \ _ _ chunkStartAdj chunkStopAdj acc -> - loopM chunkStartAdj (< chunkStopAdj) (+ stride) acc $ \j a -> - iterM ixL nL sL (<) a (f . consDim j) + else splitWorkWithFactorST fact scheduler start stride n initAcc splitAcc $ + \_ _ chunkStartAdj chunkStopAdj acc -> + loopM chunkStartAdj (< chunkStopAdj) (+ stride) acc $ \j a -> + iterM ixL nL sL (<) a (f . consDim j) else pure initAcc {-# INLINE iterRowMajorST #-} @@ -691,8 +735,15 @@ -- -- @since 1.0.2 iterF :: ix -> ix -> ix -> (Int -> Int -> Bool) -> f a -> (ix -> f a -> f a) -> f a - default iterF :: (Index (Lower ix)) => - ix -> ix -> ix -> (Int -> Int -> Bool) -> f a -> (ix -> f a -> f a) -> f a + default iterF + :: (Index (Lower ix)) + => ix + -> ix + -> ix + -> (Int -> Int -> Bool) + -> f a + -> (ix -> f a -> f a) + -> f a iterF !sIx !eIx !incIx cond initAct f = loopF s (`cond` e) (+ inc) initAct $ \ !i g -> iterF sIxL eIxL incIxL cond g (\ !ix -> f (consDim i ix)) @@ -706,11 +757,17 @@ -- -- @since 0.1.0 stepNextMF :: ix -> ix -> ix -> (Int -> Int -> Bool) -> (Maybe ix -> f a) -> f a - default stepNextMF :: (Index (Lower ix)) => - ix -> ix -> ix -> (Int -> Int -> Bool) -> (Maybe ix -> f a) -> f a + default stepNextMF + :: (Index (Lower ix)) + => ix + -> ix + -> ix + -> (Int -> Int -> Bool) + -> (Maybe ix -> f a) + -> f a stepNextMF !sIx !eIx !incIx cond f = nextMaybeF s (`cond` e) (+ inc) $ \ !mni -> - stepNextMF sIxL eIxL incIxL cond $ \ mIxN -> + stepNextMF sIxL eIxL incIxL cond $ \mIxN -> f $! case mIxN of Just ixN -> Just $! consDim s ixN @@ -724,24 +781,31 @@ !(inc, incIxL) = unconsDim incIx {-# INLINE stepNextMF #-} - iterTargetRowMajorA_ :: Applicative f - => Int -- ^ Target linear index accumulator - -> Int -- ^ Target linear index start - -> Sz ix -- ^ Target size - -> ix -- ^ Source start index - -> ix -- ^ Source stride - -> (Ix1 -> ix -> f a) - -- ^ Action that accepts a linear index of the target, - -- multi-dimensional index of the source and accumulator - -> f () - default iterTargetRowMajorA_ :: (Applicative f, Index (Lower ix)) - => Int - -> Int - -> Sz ix - -> ix - -> ix - -> (Ix1 -> ix -> f a) - -> f () + iterTargetRowMajorA_ + :: Applicative f + => Int + -- ^ Target linear index accumulator + -> Int + -- ^ Target linear index start + -> Sz ix + -- ^ Target size + -> ix + -- ^ Source start index + -> ix + -- ^ Source stride + -> (Ix1 -> ix -> f a) + -- ^ Action that accepts a linear index of the target, + -- multi-dimensional index of the source and accumulator + -> f () + default iterTargetRowMajorA_ + :: (Applicative f, Index (Lower ix)) + => Int + -> Int + -> Sz ix + -> ix + -> ix + -> (Ix1 -> ix -> f a) + -> f () iterTargetRowMajorA_ !iAcc !iStart szRes ixStart ixStride f = do let !(SafeSz nRes, !szL) = unconsSz szRes !(!start, !ixL) = unconsDim ixStart @@ -750,26 +814,34 @@ iterTargetRowMajorA_ k iStart szL ixL sL $ \i jl -> f i (consDim j jl) {-# INLINE iterTargetRowMajorA_ #-} - iterTargetRowMajorAccM :: Monad m => - Int -- ^ Target linear index accumulator - -> Int -- ^ Target linear index start - -> Sz ix -- ^ Target size - -> ix -- ^ Source start index - -> ix -- ^ Source stride - -> a -- ^ Accumulator - -> (Ix1 -> ix -> a -> m a) - -- ^ Action that accepts a linear index of the target, - -- multi-dimensional index of the source and accumulator - -> m a - default iterTargetRowMajorAccM :: (Monad m, Index (Lower ix)) - => Int - -> Int - -> Sz ix - -> ix - -> ix - -> a - -> (Ix1 -> ix -> a -> m a) - -> m a + iterTargetRowMajorAccM + :: Monad m + => Int + -- ^ Target linear index accumulator + -> Int + -- ^ Target linear index start + -> Sz ix + -- ^ Target size + -> ix + -- ^ Source start index + -> ix + -- ^ Source stride + -> a + -- ^ Accumulator + -> (Ix1 -> ix -> a -> m a) + -- ^ Action that accepts a linear index of the target, + -- multi-dimensional index of the source and accumulator + -> m a + default iterTargetRowMajorAccM + :: (Monad m, Index (Lower ix)) + => Int + -> Int + -> Sz ix + -> ix + -> ix + -> a + -> (Ix1 -> ix -> a -> m a) + -> m a iterTargetRowMajorAccM !iAcc !iStart szRes ixStart ixStride initAcc f = do let !(SafeSz nRes, !szL) = unconsSz szRes !(!start, !ixL) = unconsDim ixStart @@ -778,30 +850,41 @@ iterTargetRowMajorAccM k iStart szL ixL sL acc $ \i jl -> f i (consDim j jl) {-# INLINE iterTargetRowMajorAccM #-} - iterTargetRowMajorAccST :: - Int -- ^ Linear index accumulator - -> Int -- ^ Scheduler multiplying factor. Must be positive - -> Scheduler s a -- ^ Scheduler to use - -> Int -- ^ Target linear index start - -> Sz ix -- ^ Target size - -> ix -- ^ Source start index - -> ix -- ^ Source stride - -> a -- ^ Initial accumulator - -> (a -> ST s (a, a)) -- ^ Function that splits accumulator for each scheduled job. - -> (Ix1 -> ix -> a -> ST s a) -- ^ Action + iterTargetRowMajorAccST + :: Int + -- ^ Linear index accumulator + -> Int + -- ^ Scheduler multiplying factor. Must be positive + -> Scheduler s a + -- ^ Scheduler to use + -> Int + -- ^ Target linear index start + -> Sz ix + -- ^ Target size + -> ix + -- ^ Source start index + -> ix + -- ^ Source stride + -> a + -- ^ Initial accumulator + -> (a -> ST s (a, a)) + -- ^ Function that splits accumulator for each scheduled job. + -> (Ix1 -> ix -> a -> ST s a) + -- ^ Action -> ST s a - default iterTargetRowMajorAccST :: Index (Lower ix) - => Int - -> Int - -> Scheduler s a - -> Int - -> Sz ix - -> ix - -> ix - -> a - -> (a -> ST s (a, a)) - -> (Ix1 -> ix -> a -> ST s a) - -> ST s a + default iterTargetRowMajorAccST + :: Index (Lower ix) + => Int + -> Int + -> Scheduler s a + -> Int + -> Sz ix + -> ix + -> ix + -> a + -> (a -> ST s (a, a)) + -> (Ix1 -> ix -> a -> ST s a) + -> ST s a iterTargetRowMajorAccST !iAcc !fact scheduler iStart sz ixStart ixStride initAcc splitAcc f = do let !(SafeSz n, nL) = unconsSz sz if n > 0 @@ -815,26 +898,34 @@ iloopM iAccL start (< start + n * stride) (+ stride) initAcc $ \k j acc -> do iterTargetRowMajorAccST k newFact scheduler iStart nL ixL sL acc splitAcc $ \i -> f i . consDim j - else - splitWorkWithFactorST fact scheduler start stride n initAcc splitAcc $ - \ chunkStart _ chunkStartAdj chunkStopAdj acc -> - iloopM (iAccL + chunkStart) chunkStartAdj (< chunkStopAdj) (+ stride) acc $ \k j a -> - iterTargetRowMajorAccM k iStart nL ixL sL a $ \i -> f i . consDim j + else splitWorkWithFactorST fact scheduler start stride n initAcc splitAcc $ + \chunkStart _ chunkStartAdj chunkStopAdj acc -> + iloopM (iAccL + chunkStart) chunkStartAdj (< chunkStopAdj) (+ stride) acc $ \k j a -> + iterTargetRowMajorAccM k iStart nL ixL sL a $ \i -> f i . consDim j else pure initAcc {-# INLINE iterTargetRowMajorAccST #-} - iterTargetRowMajorAccST_ - :: Int -- ^ Index accumulator - -> Int -- ^ Scheduler multiplying factor. Must be positive - -> Scheduler s () -- ^ Scheduler to use - -> Int -- ^ Target linear start index - -> Sz ix -- ^ Target size - -> ix -- ^ Source start index - -> ix -- ^ Source stride - -> a -- ^ Initial accumulator - -> (a -> ST s (a, a)) -- ^ Function that splits accumulator for each scheduled job. - -> (Ix1 -> ix -> a -> ST s a) -- ^ Action + :: Int + -- ^ Index accumulator + -> Int + -- ^ Scheduler multiplying factor. Must be positive + -> Scheduler s () + -- ^ Scheduler to use + -> Int + -- ^ Target linear start index + -> Sz ix + -- ^ Target size + -> ix + -- ^ Source start index + -> ix + -- ^ Source stride + -> a + -- ^ Initial accumulator + -> (a -> ST s (a, a)) + -- ^ Function that splits accumulator for each scheduled job. + -> (Ix1 -> ix -> a -> ST s a) + -- ^ Action -> ST s () default iterTargetRowMajorAccST_ :: Index (Lower ix) @@ -864,12 +955,12 @@ iterTargetRowMajorAccST_ k newFact scheduler iStart szL ixL sL accCur splitAcc $ \i -> f i . consDim j pure accNext - else - void $ splitWorkWithFactorST fact scheduler start stride n initAcc splitAcc $ - \ chunkStart _ chunkStartAdj chunkStopAdj acc -> - void $ - iloopM (iAccL + chunkStart) chunkStartAdj (< chunkStopAdj) (+ stride) acc $ \k j a -> - iterTargetRowMajorAccM k iStart szL ixL sL a $ \i -> f i . consDim j + else void $ + splitWorkWithFactorST fact scheduler start stride n initAcc splitAcc $ + \chunkStart _ chunkStartAdj chunkStopAdj acc -> + void $ + iloopM (iAccL + chunkStart) chunkStartAdj (< chunkStopAdj) (+ stride) acc $ \k j a -> + iterTargetRowMajorAccM k iStart szL ixL sL a $ \i -> f i . consDim j {-# INLINE iterTargetRowMajorAccST_ #-} -- | Zero-dimension, i.e. a scalar. Can't really be used directly as there is no instance of @@ -898,11 +989,11 @@ -- @since 0.1.0 pattern Ix1 :: Int -> Ix1 pattern Ix1 i = i + {-# COMPLETE Ix1 #-} type instance Lower Int = Ix0 - instance Index Ix1 where type Dimensions Ix1 = 1 dimensions _ = 1 @@ -979,19 +1070,19 @@ let !n = unSz sz !iAccL = iStart + iAcc * n splitWorkWithFactorST fact scheduler start stride n initAcc splitAcc $ - \ chunkStart _ chunkStartAdj chunkStopAdj acc -> + \chunkStart _ chunkStartAdj chunkStopAdj acc -> iloopM (iAccL + chunkStart) chunkStartAdj (< chunkStopAdj) (+ stride) acc action {-# INLINE iterTargetRowMajorAccST #-} iterTargetRowMajorAccST_ iAcc fact scheduler iStart sz start stride initAcc splitAcc action = do let !n = unSz sz !iAccL = iStart + iAcc * n - void $ splitWorkWithFactorST fact scheduler start stride n initAcc splitAcc $ - \ chunkStart _ chunkStartAdj chunkStopAdj acc -> - void $ iloopM (iAccL + chunkStart) chunkStartAdj (< chunkStopAdj) (+ stride) acc action + void $ + splitWorkWithFactorST fact scheduler start stride n initAcc splitAcc $ + \chunkStart _ chunkStartAdj chunkStopAdj acc -> + void $ iloopM (iAccL + chunkStart) chunkStartAdj (< chunkStopAdj) (+ stride) acc action {-# INLINE iterTargetRowMajorAccST_ #-} - -- | Same as `iterM`, but don't bother with accumulator and return value. -- -- @since 0.1.0 @@ -1004,20 +1095,24 @@ -- `Applicative` instead of a `Monad` and don't bother with accumulator or return value. -- -- @since 1.0.2 -iterA_ :: - forall ix f a. (Index ix, Applicative f) - => ix -- ^ Starting index - -> ix -- ^ Ending index (not included) - -> ix -- ^ Stepping index - -> (Int -> Int -> Bool) -- ^ Continuation function. Loop will stop on `False` - -> (ix -> f a) -- ^ Action applied to an index. Result is ignored. +iterA_ + :: forall ix f a + . (Index ix, Applicative f) + => ix + -- ^ Starting index + -> ix + -- ^ Ending index (not included) + -> ix + -- ^ Stepping index + -> (Int -> Int -> Bool) + -- ^ Continuation function. Loop will stop on `False` + -> (ix -> f a) + -- ^ Action applied to an index. Result is ignored. -> f () iterA_ sIx eIx incIx cond f = iterF sIx eIx incIx cond (pure ()) $ \ix go -> f ix *> go {-# INLINE iterA_ #-} - - -- | Exceptions that get thrown when there is a problem with an index, size or dimension. -- -- @since 0.3.0 @@ -1046,14 +1141,15 @@ | 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 + , Just sz2 <- cast sz2t -> + sz1 == sz2 && i1 == i2 _ -> False instance NFData IndexException where rnf = \case - IndexZeroException i -> rnf i - IndexDimensionException i d -> i `deepseq` rnf d + IndexZeroException i -> rnf i + IndexDimensionException i d -> i `deepseq` rnf d IndexOutOfBoundsException sz i -> sz `deepseq` rnf i instance Exception IndexException @@ -1084,14 +1180,17 @@ case (e1, e2) of (SizeMismatchException sz1 sz1', SizeMismatchException sz2t sz2t') | Just sz2 <- cast sz2t - , Just sz2' <- cast sz2t' -> sz1 == sz2 && sz1' == sz2' + , 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' + , 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' + , Just sz2' <- cast sz2t' -> + sz1 == sz2 && i1 == i2 && sz1' == sz2' (SizeEmptyException sz1, SizeEmptyException sz2t) | Just sz2 <- cast sz2t -> sz1 == sz2 (SizeOverflowException sz1, SizeOverflowException sz2t) @@ -1103,24 +1202,28 @@ instance NFData SizeException where rnf = \case - SizeMismatchException sz sz' -> sz `deepseq` rnf sz' + SizeMismatchException sz sz' -> sz `deepseq` rnf sz' SizeElementsMismatchException sz sz' -> sz `deepseq` rnf sz' - SizeSubregionException sz i sz' -> sz `deepseq` i `deepseq` rnf sz' - SizeEmptyException sz -> rnf sz - SizeOverflowException sz -> rnf sz - SizeNegativeException sz -> rnf sz + SizeSubregionException sz i sz' -> sz `deepseq` i `deepseq` rnf sz' + SizeEmptyException sz -> rnf sz + SizeOverflowException sz -> rnf sz + SizeNegativeException sz -> rnf sz instance Exception SizeException - instance Show SizeException where show (SizeMismatchException sz sz') = "SizeMismatchException: (" ++ show sz ++ ") vs (" ++ show sz' ++ ")" show (SizeElementsMismatchException sz sz') = "SizeElementsMismatchException: (" ++ show sz ++ ") vs (" ++ show sz' ++ ")" show (SizeSubregionException sz' ix sz) = - "SizeSubregionException: (" ++ - show sz' ++ ") is to small for " ++ show ix ++ " (" ++ show sz ++ ")" + "SizeSubregionException: (" + ++ show sz' + ++ ") is to small for " + ++ show ix + ++ " (" + ++ show sz + ++ ")" show (SizeEmptyException sz) = "SizeEmptyException: (" ++ show sz ++ ") corresponds to an empty array" show (SizeOverflowException sz) = @@ -1134,13 +1237,13 @@ -- -- @since 0.3.0 data ShapeException - = 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 + = -- | Across a specific dimension there was not enough elements for the supplied size + DimTooShortException !Dim !(Sz Ix1) !(Sz Ix1) + | -- | Across a specific dimension there was too many elements for the supplied size + DimTooLongException !Dim !(Sz Ix1) !(Sz Ix1) + | -- | Expected an empty size, but the shape was not empty. + ShapeNonEmpty + deriving (Eq) instance Show ShapeException where showsPrec n = @@ -1152,14 +1255,19 @@ showsShapeExc tyName d sz sz' = showsPrecWrapped n - ((tyName ++) . - (" for " ++) . - shows d . (": expected (" ++) . shows sz . ("), got (" ++) . shows sz' . (")" ++)) + ( (tyName ++) + . (" for " ++) + . shows d + . (": expected (" ++) + . shows sz + . ("), got (" ++) + . shows sz' + . (")" ++) + ) instance Exception ShapeException - showsPrecWrapped :: Int -> ShowS -> ShowS showsPrecWrapped n inner | n < 1 = inner - | otherwise = ('(':) . inner . (")" ++) + | otherwise = ('(' :) . inner . (")" ++) diff --git a/src/Data/Massiv/Core/Index/Iterator.hs b/src/Data/Massiv/Core/Index/Iterator.hs --- a/src/Data/Massiv/Core/Index/Iterator.hs +++ b/src/Data/Massiv/Core/Index/Iterator.hs @@ -3,6 +3,7 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE MonoLocalBinds #-} {-# LANGUAGE PatternSynonyms #-} + -- | -- Module : Data.Massiv.Core.Index.Iterator -- Copyright : (c) Alexey Kuleshevich 2021-2022 @@ -10,24 +11,25 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Core.Index.Iterator - ( Iterator(..) +module Data.Massiv.Core.Index.Iterator ( + Iterator (..), + -- * Extra iterator functions - , iterTargetAccST - , iterTargetAccST_ - , iterTargetFullWithStrideAccST - , iterTargetFullWithStrideAccST_ - , iterTargetST_ - , iterTargetFullWithStrideST_ + iterTargetAccST, + iterTargetAccST_, + iterTargetFullWithStrideAccST, + iterTargetFullWithStrideAccST_, + iterTargetST_, + iterTargetFullWithStrideST_, + -- * Iterator implementations - , RowMajor(RowMajor) - , defRowMajor - , RowMajorLinear(RowMajorLinear) - , defRowMajorLinear - , RowMajorUnbalanced(RowMajorUnbalanced) - , defRowMajorUnbalanced - ) where + RowMajor (RowMajor), + defRowMajor, + RowMajorLinear (RowMajorLinear), + defRowMajorLinear, + RowMajorUnbalanced (RowMajorUnbalanced), + defRowMajorUnbalanced, +) where import Control.Monad import Control.Monad.ST @@ -36,23 +38,25 @@ import Data.Massiv.Core.Index.Stride import Data.Massiv.Core.Loop - class Iterator it where {-# MINIMAL (iterTargetM, iterTargetA_, iterTargetWithStrideAccST, iterTargetWithStrideAccST_) #-} - -- | Iterate over a target region using linear index with access to the source -- index, which adjusted according to the stride. Use `iterTargetM` if you -- need an accumulator. -- -- @since 1.0.2 - iterTargetA_ :: - (Index ix, Applicative f) + iterTargetA_ + :: (Index ix, Applicative f) => it - -> Int -- ^ Target linear index start - -> Sz ix -- ^ Target size - -> ix -- ^ Source start index - -> Stride ix -- ^ Source stride + -> Int + -- ^ Target linear index start + -> Sz ix + -- ^ Target size + -> ix + -- ^ Source start index + -> Stride ix + -- ^ Source stride -> (Ix1 -> ix -> f a) -- ^ Action that accepts a linear index of the target and multi-dimensional -- index of the source. @@ -62,56 +66,78 @@ -- index, which adjusted according to the stride. -- -- @since 1.0.2 - iterTargetM :: - (Index ix, Monad m) + iterTargetM + :: (Index ix, Monad m) => it - -> Ix1 -- ^ Target linear index start - -> Sz ix -- ^ Target size - -> ix -- ^ Source start index - -> Stride ix -- ^ Source stride - -> a -- ^ Accumulator + -> Ix1 + -- ^ Target linear index start + -> Sz ix + -- ^ Target size + -> ix + -- ^ Source start index + -> Stride ix + -- ^ Source stride + -> a + -- ^ Accumulator -> (Ix1 -> ix -> a -> m a) -- ^ Action that accepts a linear index of the target, -- multi-dimensional index of the source and accumulator -> m a - iterTargetWithStrideAccST :: - Index ix + iterTargetWithStrideAccST + :: Index ix => it - -> Scheduler s a -- ^ Scheduler to use - -> Ix1 -- ^ Target linear start index - -> Sz ix -- ^ Target size - -> ix -- ^ Source start index - -> Stride ix -- ^ Source stride - -> a -- ^ Initial accumulator + -> Scheduler s a + -- ^ Scheduler to use + -> Ix1 + -- ^ Target linear start index + -> Sz ix + -- ^ Target size + -> ix + -- ^ Source start index + -> Stride ix + -- ^ Source stride + -> a + -- ^ Initial accumulator -> (a -> ST s (a, a)) -- ^ Splitting action that produces new accumulators for separate worker threads. - -> (Ix1 -> ix -> a -> ST s a) -- ^ Action + -> (Ix1 -> ix -> a -> ST s a) + -- ^ Action -> ST s a - iterTargetWithStrideAccST_ :: - Index ix + iterTargetWithStrideAccST_ + :: Index ix => it - -> Scheduler s () -- ^ Scheduler to use - -> Ix1 -- ^ Target linear start index - -> Sz ix -- ^ Target size - -> ix -- ^ Start - -> Stride ix -- ^ Stride - -> a -- ^ Initial accumulator + -> Scheduler s () + -- ^ Scheduler to use + -> Ix1 + -- ^ Target linear start index + -> Sz ix + -- ^ Target size + -> ix + -- ^ Start + -> Stride ix + -- ^ Stride + -> a + -- ^ Initial accumulator -> (a -> ST s (a, a)) -- ^ Splitting action that produces new accumulators for separate worker threads. - -> (Ix1 -> ix -> a -> ST s a) -- ^ Action + -> (Ix1 -> ix -> a -> ST s a) + -- ^ Action -> ST s () -- | Iterate over a region with a monadic action and accumulator. -- -- @since 1.0.2 - iterFullM :: - (Index ix, Monad m) + iterFullM + :: (Index ix, Monad m) => it - -> ix -- ^ Source start index - -> Sz ix -- ^ Source size - -> a -- ^ Accumulator + -> ix + -- ^ Source start index + -> Sz ix + -- ^ Source size + -> a + -- ^ Accumulator -> (ix -> a -> m a) -- ^ Action that accepts a linear index of the target, -- multi-dimensional index of the source and accumulator @@ -123,11 +149,13 @@ -- | Iterate over a region with an applicative action ignoring the result. -- -- @since 1.0.2 - iterFullA_ :: - (Index ix, Applicative f) + iterFullA_ + :: (Index ix, Applicative f) => it - -> ix -- ^ Source start index - -> Sz ix -- ^ Source size + -> ix + -- ^ Source start index + -> Sz ix + -- ^ Source size -> (ix -> f a) -- ^ Action that accepts a linear index of the target, -- multi-dimensional index of the source and accumulator @@ -137,55 +165,78 @@ {-# INLINE iterFullA_ #-} -- | Iterate over a region in a ST monad with access to `Scheduler`. - iterFullAccST :: - Index ix - => it -- ^ Scheduler multiplying factor. Must be positive - -> Scheduler s a -- ^ Scheduler to use - -> ix -- ^ Start index - -> Sz ix -- ^ Size - -> a -- ^ Initial accumulator - -> (a -> ST s (a, a)) -- ^ Function that splits accumulator for each scheduled job. - -> (ix -> a -> ST s a) -- ^ Action + iterFullAccST + :: Index ix + => it + -- ^ Scheduler multiplying factor. Must be positive + -> Scheduler s a + -- ^ Scheduler to use + -> ix + -- ^ Start index + -> Sz ix + -- ^ Size + -> a + -- ^ Initial accumulator + -> (a -> ST s (a, a)) + -- ^ Function that splits accumulator for each scheduled job. + -> (ix -> a -> ST s a) + -- ^ Action -> ST s a iterFullAccST it scheduler start sz acc splitAcc f = iterTargetAccST it scheduler 0 sz start acc splitAcc (const f) {-# INLINE iterFullAccST #-} - iterTargetFullAccST :: - Index ix + iterTargetFullAccST + :: Index ix => it - -> Scheduler s a -- ^ Scheduler to use - -> Ix1 -- ^ Target linear start index - -> Sz ix -- ^ Target size - -> a -- ^ Initial accumulator - -> (a -> ST s (a, a)) -- ^ Function that splits accumulator for each scheduled job. - -> (Ix1 -> ix -> a -> ST s a) -- ^ Action + -> Scheduler s a + -- ^ Scheduler to use + -> Ix1 + -- ^ Target linear start index + -> Sz ix + -- ^ Target size + -> a + -- ^ Initial accumulator + -> (a -> ST s (a, a)) + -- ^ Function that splits accumulator for each scheduled job. + -> (Ix1 -> ix -> a -> ST s a) + -- ^ Action -> ST s a iterTargetFullAccST it scheduler iStart sz = iterTargetFullWithStrideAccST it scheduler iStart sz oneStride {-# INLINE iterTargetFullAccST #-} - iterTargetFullAccST_ :: - Index ix + iterTargetFullAccST_ + :: Index ix => it - -> Scheduler s () -- ^ Scheduler to use - -> Ix1 -- ^ Target linear start index - -> Sz ix -- ^ Target size - -> a -- ^ Initial accumulator - -> (a -> ST s (a, a)) -- ^ Function that splits accumulator for each scheduled job. - -> (Ix1 -> ix -> a -> ST s a) -- ^ Action + -> Scheduler s () + -- ^ Scheduler to use + -> Ix1 + -- ^ Target linear start index + -> Sz ix + -- ^ Target size + -> a + -- ^ Initial accumulator + -> (a -> ST s (a, a)) + -- ^ Function that splits accumulator for each scheduled job. + -> (Ix1 -> ix -> a -> ST s a) + -- ^ Action -> ST s () iterTargetFullAccST_ it scheduler iStart sz = iterTargetFullWithStrideAccST_ it scheduler iStart sz oneStride {-# INLINE iterTargetFullAccST_ #-} - iterTargetFullST_ :: - Index ix + iterTargetFullST_ + :: Index ix => it - -> Scheduler s () -- ^ Scheduler to use - -> Ix1 -- ^ Target linear start index - -> Sz ix -- ^ Target size - -> (Ix1 -> ix -> ST s ()) -- ^ Action + -> Scheduler s () + -- ^ Scheduler to use + -> Ix1 + -- ^ Target linear start index + -> Sz ix + -- ^ Target size + -> (Ix1 -> ix -> ST s ()) + -- ^ Action -> ST s () iterTargetFullST_ it scheduler iStart sz = iterTargetST_ it scheduler iStart sz (pureIndex 0) @@ -193,23 +244,29 @@ -- NOTE: this function does not have to be part of the class, but for some -- reason it creates a severe regression when moved outside. + -- | Iterate over a target array with a stride without an accumulator - iterTargetWithStrideST_ :: - Index ix + iterTargetWithStrideST_ + :: Index ix => it - -> Scheduler s () -- ^ Scheduler to use - -> Ix1 -- ^ Target linear start index - -> Sz ix -- ^ Target size - -> ix -- ^ Start - -> Stride ix -- ^ Stride - -> (Ix1 -> ix -> ST s a) -- ^ Action + -> Scheduler s () + -- ^ Scheduler to use + -> Ix1 + -- ^ Target linear start index + -> Sz ix + -- ^ Target size + -> ix + -- ^ Start + -> Stride ix + -- ^ Stride + -> (Ix1 -> ix -> ST s a) + -- ^ Action -> ST s () iterTargetWithStrideST_ it scheduler i sz ix stride action = iterTargetWithStrideAccST_ it scheduler i sz ix stride () noSplit $ \j jx _ -> void $ action j jx {-# INLINE iterTargetWithStrideST_ #-} - -- | Default iterator that parallelizes work in linear chunks. Supplied factor -- will be used to schedule that many jobs per capability. -- @@ -220,11 +277,13 @@ defRowMajor :: RowMajor defRowMajor = RowMajorInternal 8 -pattern RowMajor :: Int - -- ^ Multiplier that will be used to scale number of jobs. - -> RowMajor +pattern RowMajor + :: Int + -- ^ Multiplier that will be used to scale number of jobs. + -> RowMajor pattern RowMajor f <- RowMajorInternal f - where RowMajor = RowMajorInternal . max 1 + where + RowMajor = RowMajorInternal . max 1 {-# COMPLETE RowMajor #-} instance Iterator RowMajor where @@ -248,7 +307,6 @@ iterTargetRowMajorAccST_ 0 fact scheduler i sz ix stride {-# INLINE iterTargetWithStrideAccST_ #-} - newtype RowMajorLinear = RowMajorLinear Int defRowMajorLinear :: RowMajorLinear @@ -265,18 +323,18 @@ {-# INLINE iterTargetA_ #-} iterTargetFullAccST it scheduler iStart sz acc splitAcc action = let !(RowMajorLinear fact) = it - in iterLinearAccST fact scheduler iStart 1 (totalElem sz) acc splitAcc $ \ !i -> - action i (fromLinearIndex sz i) + in iterLinearAccST fact scheduler iStart 1 (totalElem sz) acc splitAcc $ \ !i -> + action i (fromLinearIndex sz i) {-# INLINE iterTargetFullAccST #-} iterTargetFullAccST_ it scheduler iStart sz acc splitAcc action = let !(RowMajorLinear fact) = it - in iterLinearAccST_ fact scheduler iStart 1 (totalElem sz) acc splitAcc $ \ !i -> - action i (fromLinearIndex sz i) + in iterLinearAccST_ fact scheduler iStart 1 (totalElem sz) acc splitAcc $ \ !i -> + action i (fromLinearIndex sz i) {-# INLINE iterTargetFullAccST_ #-} iterTargetFullST_ it scheduler iStart sz action = let !(RowMajorLinear fact) = it - in iterLinearST_ fact scheduler iStart 1 (totalElem sz) $ \ !i -> - action i (fromLinearIndex sz i) + in iterLinearST_ fact scheduler iStart 1 (totalElem sz) $ \ !i -> + action i (fromLinearIndex sz i) {-# INLINE iterTargetFullST_ #-} iterTargetWithStrideAccST it scheduler iStart sz start (Stride stride) acc spliAcc action = let RowMajorLinear fact = it @@ -291,8 +349,6 @@ liftIndex2 (+) start (liftIndex2 (*) stride (fromLinearIndex sz i)) {-# INLINE iterTargetWithStrideAccST_ #-} - - -- | Parallelizing unbalanced computation (i.e. computing some elements of the -- array is much more expensive then the others) it can be benefitial to -- interleave iteration. Perfect example of this would be a ray tracer or the @@ -306,14 +362,15 @@ defRowMajorUnbalanced :: RowMajorUnbalanced defRowMajorUnbalanced = RowMajorUnbalancedInternal 8 -pattern RowMajorUnbalanced :: Int - -- ^ Multiplier that will be used to scale number of jobs. - -> RowMajorUnbalanced +pattern RowMajorUnbalanced + :: Int + -- ^ Multiplier that will be used to scale number of jobs. + -> RowMajorUnbalanced pattern RowMajorUnbalanced f <- RowMajorUnbalancedInternal f - where RowMajorUnbalanced = RowMajorUnbalancedInternal . max 1 + where + RowMajorUnbalanced = RowMajorUnbalancedInternal . max 1 {-# COMPLETE RowMajorUnbalanced #-} - instance Iterator RowMajorUnbalanced where iterFullM (RowMajorUnbalanced fact) = iterFullM (RowMajor fact) {-# INLINE iterFullM #-} @@ -327,15 +384,15 @@ {-# INLINE iterTargetWithStrideAccST #-} iterTargetWithStrideAccST_ it scheduler iStart sz start stride acc splitAcc' action = void $ - iterUnbalancedTargetWithStride innerLoop it scheduler iStart sz start stride acc splitAcc' action + iterUnbalancedTargetWithStride innerLoop it scheduler iStart sz start stride acc splitAcc' action where innerLoop initial condition increment initAcc f = void $ loopM initial condition increment initAcc f {-# INLINE innerLoop #-} {-# INLINE iterTargetWithStrideAccST_ #-} -iterUnbalancedTargetWithStride :: - Index ix +iterUnbalancedTargetWithStride + :: Index ix => (Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> t) -> ST s b) -> RowMajorUnbalanced -> Scheduler s b @@ -351,103 +408,129 @@ let RowMajorUnbalanced fact = it !n = totalElem sz !step = min (fact * numWorkers scheduler) n - in loopM 0 (< step) (+ 1) acc $ \ !istep !a -> do - (curAcc, nextAcc) <- splitAcc a - scheduleMassivWork scheduler $ - innerLoop istep (< n) (+ step) curAcc $ \i -> - action (iStart + i) $ - liftIndex2 (+) start (liftIndex2 (*) (unStride stride) (fromLinearIndex sz i)) - pure nextAcc + in loopM 0 (< step) (+ 1) acc $ \ !istep !a -> do + (curAcc, nextAcc) <- splitAcc a + scheduleMassivWork scheduler $ + innerLoop istep (< n) (+ step) curAcc $ \i -> + action (iStart + i) $ + liftIndex2 (+) start (liftIndex2 (*) (unStride stride) (fromLinearIndex sz i)) + pure nextAcc {-# INLINE iterUnbalancedTargetWithStride #-} - noSplit :: Applicative m => () -> m ((), ()) noSplit _ = pure ((), ()) - -iterTargetAccST :: - (Iterator it, Index ix) +iterTargetAccST + :: (Iterator it, Index ix) => it - -> Scheduler s a -- ^ Scheduler to use - -> Ix1 -- ^ Target linear start index - -> Sz ix -- ^ Target size - -> ix -- ^ Source start + -> Scheduler s a + -- ^ Scheduler to use + -> Ix1 + -- ^ Target linear start index + -> Sz ix + -- ^ Target size + -> ix + -- ^ Source start -> a -> (a -> ST s (a, a)) - -> (Ix1 -> ix -> a -> ST s a) -- ^ Action + -> (Ix1 -> ix -> a -> ST s a) + -- ^ Action -> ST s a iterTargetAccST it scheduler iStart sz ix = iterTargetWithStrideAccST it scheduler iStart sz ix oneStride {-# INLINE iterTargetAccST #-} -iterTargetAccST_ :: - (Iterator it, Index ix) +iterTargetAccST_ + :: (Iterator it, Index ix) => it - -> Scheduler s () -- ^ Scheduler to use - -> Ix1 -- ^ Target linear start index - -> Sz ix -- ^ Target size - -> ix -- ^ Source start + -> Scheduler s () + -- ^ Scheduler to use + -> Ix1 + -- ^ Target linear start index + -> Sz ix + -- ^ Target size + -> ix + -- ^ Source start -> a -> (a -> ST s (a, a)) - -> (Ix1 -> ix -> a -> ST s a) -- ^ Action + -> (Ix1 -> ix -> a -> ST s a) + -- ^ Action -> ST s () iterTargetAccST_ it scheduler iStart sz ix = iterTargetWithStrideAccST_ it scheduler iStart sz ix oneStride {-# INLINE iterTargetAccST_ #-} - -iterTargetFullWithStrideST_ :: - (Iterator it, Index ix) +iterTargetFullWithStrideST_ + :: (Iterator it, Index ix) => it - -> Scheduler s () -- ^ Scheduler to use - -> Ix1 -- ^ Target linear start index - -> Sz ix -- ^ Target size - -> Stride ix -- ^ Stride - -> (Ix1 -> ix -> ST s ()) -- ^ Action + -> Scheduler s () + -- ^ Scheduler to use + -> Ix1 + -- ^ Target linear start index + -> Sz ix + -- ^ Target size + -> Stride ix + -- ^ Stride + -> (Ix1 -> ix -> ST s ()) + -- ^ Action -> ST s () iterTargetFullWithStrideST_ it scheduler iStart sz = iterTargetWithStrideST_ it scheduler iStart sz (pureIndex 0) {-# INLINE iterTargetFullWithStrideST_ #-} -iterTargetST_ :: - (Iterator it, Index ix) +iterTargetST_ + :: (Iterator it, Index ix) => it - -> Scheduler s () -- ^ Scheduler to use - -> Ix1 -- ^ Target linear start index - -> Sz ix -- ^ Target size - -> ix -- ^ Start - -> (Ix1 -> ix -> ST s ()) -- ^ Action + -> Scheduler s () + -- ^ Scheduler to use + -> Ix1 + -- ^ Target linear start index + -> Sz ix + -- ^ Target size + -> ix + -- ^ Start + -> (Ix1 -> ix -> ST s ()) + -- ^ Action -> ST s () iterTargetST_ it scheduler iStart sz ix = iterTargetWithStrideST_ it scheduler iStart sz ix oneStride {-# INLINE iterTargetST_ #-} - -iterTargetFullWithStrideAccST :: - (Iterator it, Index ix) +iterTargetFullWithStrideAccST + :: (Iterator it, Index ix) => it - -> Scheduler s a -- ^ Scheduler to use - -> Ix1 -- ^ Target linear start index - -> Sz ix -- ^ Target size - -> Stride ix -- ^ Stride + -> Scheduler s a + -- ^ Scheduler to use + -> Ix1 + -- ^ Target linear start index + -> Sz ix + -- ^ Target size + -> Stride ix + -- ^ Stride -> a -> (a -> ST s (a, a)) - -> (Ix1 -> ix -> a -> ST s a) -- ^ Action + -> (Ix1 -> ix -> a -> ST s a) + -- ^ Action -> ST s a iterTargetFullWithStrideAccST it scheduler iStart sz = iterTargetWithStrideAccST it scheduler iStart sz (pureIndex 0) {-# INLINE iterTargetFullWithStrideAccST #-} -iterTargetFullWithStrideAccST_ :: - (Iterator it, Index ix) +iterTargetFullWithStrideAccST_ + :: (Iterator it, Index ix) => it - -> Scheduler s () -- ^ Scheduler to use - -> Ix1 -- ^ Target linear start index - -> Sz ix -- ^ Target size - -> Stride ix -- ^ Stride + -> Scheduler s () + -- ^ Scheduler to use + -> Ix1 + -- ^ Target linear start index + -> Sz ix + -- ^ Target size + -> Stride ix + -- ^ Stride -> a -> (a -> ST s (a, a)) - -> (Ix1 -> ix -> a -> ST s a) -- ^ Action + -> (Ix1 -> ix -> a -> ST s a) + -- ^ Action -> ST s () iterTargetFullWithStrideAccST_ it scheduler iStart sz = iterTargetWithStrideAccST_ it scheduler iStart sz (pureIndex 0) 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 @@ -10,6 +10,7 @@ {-# LANGUAGE TypeFamilyDependencies #-} {-# LANGUAGE TypeOperators #-} {-# LANGUAGE UndecidableInstances #-} + -- | -- Module : Data.Massiv.Core.Index.Ix -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -17,32 +18,32 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Core.Index.Ix - ( Ix - , IxN((:>)) - , type Sz - , pattern Sz - , type Ix1 - , pattern Ix1 - , pattern Sz1 - , type Ix2(Ix2, (:.)) - , pattern Sz2 - , type Ix3 - , pattern Ix3 - , pattern Sz3 - , type Ix4 - , pattern Ix4 - , pattern Sz4 - , type Ix5 - , pattern Ix5 - , pattern Sz5 - , HighIxN - ) where +module Data.Massiv.Core.Index.Ix ( + Ix, + IxN ((:>)), + type Sz, + pattern Sz, + type Ix1, + pattern Ix1, + pattern Sz1, + type Ix2 (Ix2, (:.)), + pattern Sz2, + type Ix3, + pattern Ix3, + pattern Sz3, + type Ix4, + pattern Ix4, + pattern Sz4, + type Ix5, + pattern Ix5, + pattern Sz5, + HighIxN, +) where import Control.DeepSeq -import Control.Monad.Catch (MonadThrow(..)) +import Control.Monad.Catch (MonadThrow (..)) import Data.Massiv.Core.Index.Internal +import Data.Massiv.Core.Loop import Data.Proxy import qualified Data.Vector.Generic as V import qualified Data.Vector.Generic.Mutable as VM @@ -50,15 +51,12 @@ import qualified GHC.Arr as I import GHC.TypeLits import System.Random.Stateful -import Data.Massiv.Core.Loop #if !MIN_VERSION_base(4,11,0) import Data.Semigroup #endif - infixr 5 :>, :. - -- | 2-dimensional index. This is also a base index for higher dimensions. -- -- @since 0.1.0 @@ -69,6 +67,7 @@ -- @since 0.1.0 pattern Ix2 :: Int -> Int -> Ix2 pattern Ix2 i2 i1 = i2 :. i1 + {-# COMPLETE Ix2 #-} -- | 2-dimensional size constructor. @(Sz2 i j) == Sz (i :. j)@ @@ -76,6 +75,7 @@ -- @since 0.3.0 pattern Sz2 :: Int -> Int -> Sz Ix2 pattern Sz2 i2 i1 = Sz (i2 :. i1) + {-# COMPLETE Sz2 #-} -- | 3-dimensional type synonym. Useful as a alternative to enabling @DataKinds@ and using type @@ -89,6 +89,7 @@ -- @since 0.1.0 pattern Ix3 :: Int -> Int -> Int -> Ix3 pattern Ix3 i3 i2 i1 = i3 :> i2 :. i1 + {-# COMPLETE Ix3 #-} -- | 3-dimensional size constructor. @(Sz3 i j k) == Sz (i :> j :. k)@ @@ -96,6 +97,7 @@ -- @since 0.3.0 pattern Sz3 :: Int -> Int -> Int -> Sz Ix3 pattern Sz3 i3 i2 i1 = Sz (i3 :> i2 :. i1) + {-# COMPLETE Sz3 #-} -- | 4-dimensional type synonym. @@ -108,6 +110,7 @@ -- @since 0.1.0 pattern Ix4 :: Int -> Int -> Int -> Int -> Ix4 pattern Ix4 i4 i3 i2 i1 = i4 :> i3 :> i2 :. i1 + {-# COMPLETE Ix4 #-} -- | 4-dimensional size constructor. @(Sz4 i j k l) == Sz (i :> j :> k :. l)@ @@ -115,6 +118,7 @@ -- @since 0.3.0 pattern Sz4 :: Int -> Int -> Int -> Int -> Sz Ix4 pattern Sz4 i4 i3 i2 i1 = Sz (i4 :> i3 :> i2 :. i1) + {-# COMPLETE Sz4 #-} -- | 5-dimensional type synonym. @@ -127,6 +131,7 @@ -- @since 0.1.0 pattern Ix5 :: Int -> Int -> Int -> Int -> Int -> Ix5 pattern Ix5 i5 i4 i3 i2 i1 = i5 :> i4 :> i3 :> i2 :. i1 + {-# COMPLETE Ix5 #-} -- | 5-dimensional size constructor. @(Sz5 i j k l m) == Sz (i :> j :> k :> l :. m)@ @@ -134,6 +139,7 @@ -- @since 0.3.0 pattern Sz5 :: Int -> Int -> Int -> Int -> Int -> Sz Ix5 pattern Sz5 i5 i4 i3 i2 i1 = Sz (i5 :> i4 :> i3 :> i2 :. i1) + {-# COMPLETE Sz5 #-} -- | n-dimensional index. Needs a base case, which is the `Ix2`. @@ -150,11 +156,9 @@ Ix 2 = Ix2 Ix n = IxN n - type instance Lower Ix2 = Ix1 type instance Lower (IxN n) = Ix (n - 1) - instance Show Ix2 where showsPrec n (i :. j) = showsPrecWrapped n (shows i . (" :. " ++) . shows j) @@ -211,7 +215,6 @@ 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] (+) #-} @@ -260,8 +263,6 @@ fromInteger = pureIndex . fromInteger {-# INLINE [1] fromInteger #-} - - instance Bounded Ix2 where minBound = pureIndex minBound {-# INLINE minBound #-} @@ -286,21 +287,18 @@ instance NFData (IxN n) where rnf ix = ix `seq` () - instance Eq Ix2 where - (i1 :. j1) == (i2 :. j2) = i1 == i2 && j1 == j2 + (i1 :. j1) == (i2 :. j2) = i1 == i2 && j1 == j2 instance Eq (Ix (n - 1)) => Eq (IxN n) where (i1 :> ix1) == (i2 :> ix2) = i1 == i2 && ix1 == ix2 - instance Ord Ix2 where compare (i1 :. j1) (i2 :. j2) = compare i1 i2 <> compare j1 j2 instance Ord (Ix (n - 1)) => Ord (IxN n) where compare (i1 :> ix1) (i2 :> ix2) = compare i1 i2 <> compare ix1 ix2 - instance Index Ix2 where type Dimensions Ix2 = 2 dimensions _ = 2 @@ -325,19 +323,19 @@ {-# INLINE [1] unsnocDim #-} getDimM (i2 :. _) 2 = pure i2 getDimM (_ :. i1) 1 = pure i1 - getDimM ix d = throwM $ IndexDimensionException ix d + getDimM ix d = throwM $ IndexDimensionException ix d {-# INLINE [1] getDimM #-} setDimM (_ :. i1) 2 i2 = pure (i2 :. i1) setDimM (i2 :. _) 1 i1 = pure (i2 :. i1) - setDimM ix d _ = throwM $ IndexDimensionException ix d + setDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] setDimM #-} pullOutDimM (i2 :. i1) 2 = pure (i2, i1) pullOutDimM (i2 :. i1) 1 = pure (i1, i2) - pullOutDimM ix d = throwM $ IndexDimensionException ix d + pullOutDimM ix d = throwM $ IndexDimensionException ix d {-# INLINE [1] pullOutDimM #-} insertDimM i1 2 i2 = pure (i2 :. i1) insertDimM i2 1 i1 = pure (i2 :. i1) - insertDimM ix d _ = throwM $ IndexDimensionException ix d + insertDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] insertDimM #-} pureIndex i = i :. i {-# INLINE [1] pureIndex #-} @@ -353,7 +351,6 @@ loopF sIxL (`cond` eIxL) (+ incIxL) g $ \ !j -> f (i :. j) {-# INLINE iterF #-} - instance {-# OVERLAPPING #-} Index (IxN 3) where type Dimensions Ix3 = 3 dimensions _ = 3 @@ -380,25 +377,25 @@ {-# INLINE [1] snocDim #-} unsnocDim (i3 :> i2 :. i1) = (i3 :. i2, i1) {-# INLINE [1] unsnocDim #-} - getDimM (i3 :> _ :. _) 3 = pure i3 - getDimM ( _ :> i2 :. _) 2 = pure i2 - getDimM ( _ :> _ :. i1) 1 = pure i1 - getDimM ix d = throwM $ IndexDimensionException ix d + getDimM (i3 :> _ :. _) 3 = pure i3 + getDimM (_ :> i2 :. _) 2 = pure i2 + getDimM (_ :> _ :. i1) 1 = pure i1 + getDimM ix d = throwM $ IndexDimensionException ix d {-# INLINE [1] getDimM #-} - setDimM ( _ :> i2 :. i1) 3 i3 = pure (i3 :> i2 :. i1) - setDimM (i3 :> _ :. i1) 2 i2 = pure (i3 :> i2 :. i1) - setDimM (i3 :> i2 :. _) 1 i1 = pure (i3 :> i2 :. i1) - setDimM ix d _ = throwM $ IndexDimensionException ix d + setDimM (_ :> i2 :. i1) 3 i3 = pure (i3 :> i2 :. i1) + setDimM (i3 :> _ :. i1) 2 i2 = pure (i3 :> i2 :. i1) + setDimM (i3 :> i2 :. _) 1 i1 = pure (i3 :> i2 :. i1) + setDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] setDimM #-} pullOutDimM (i3 :> i2 :. i1) 3 = pure (i3, i2 :. i1) pullOutDimM (i3 :> i2 :. i1) 2 = pure (i2, i3 :. i1) pullOutDimM (i3 :> i2 :. i1) 1 = pure (i1, i3 :. i2) - pullOutDimM ix d = throwM $ IndexDimensionException ix d + pullOutDimM ix d = throwM $ IndexDimensionException ix d {-# INLINE [1] pullOutDimM #-} insertDimM (i2 :. i1) 3 i3 = pure (i3 :> i2 :. i1) insertDimM (i3 :. i1) 2 i2 = pure (i3 :> i2 :. i1) insertDimM (i3 :. i2) 1 i1 = pure (i3 :> i2 :. i1) - insertDimM ix d _ = throwM $ IndexDimensionException ix d + insertDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] insertDimM #-} pureIndex i = i :> i :. i {-# INLINE [1] pureIndex #-} @@ -414,28 +411,28 @@ iShift = iStart + iAcc * n in loopM 0 (< n) (+ 1) initAcc $ \ !i !acc -> let (i3 :> i2 :. i1) = fromLinearIndex sz i - in action (iShift + i) ((b3 + s3 * i3) :> (b2 + s2 * i2) :. (b1 + s1 * i1)) acc + in action (iShift + i) ((b3 + s3 * i3) :> (b2 + s2 * i2) :. (b1 + s1 * i1)) acc {-# INLINE iterTargetRowMajorAccM #-} iterTargetRowMajorAccST_ iAcc fact scheduler iStart sz (b3 :> b2 :. b1) (s3 :> s2 :. s1) acc splitAcc action = let n = totalElem sz iShift = iStart + iAcc * n in iterLinearAccST_ fact scheduler 0 1 n acc splitAcc $ \ !i -> let (i3 :> i2 :. i1) = fromLinearIndex sz i - in action (iShift + i) ((b3 + s3 * i3) :> (b2 + s2 * i2) :. (b1 + s1 * i1)) + in action (iShift + i) ((b3 + s3 * i3) :> (b2 + s2 * i2) :. (b1 + s1 * i1)) {-# INLINE iterTargetRowMajorAccST_ #-} iterTargetRowMajorAccST iAcc fact scheduler iStart sz (b3 :> b2 :. b1) (s3 :> s2 :. s1) acc splitAcc action = let n = totalElem sz iShift = iStart + iAcc * n in iterLinearAccST fact scheduler 0 1 n acc splitAcc $ \ !i -> let (i3 :> i2 :. i1) = fromLinearIndex sz i - in action (iShift + i) ((b3 + s3 * i3) :> (b2 + s2 * i2) :. (b1 + s1 * i1)) + in action (iShift + i) ((b3 + s3 * i3) :> (b2 + s2 * i2) :. (b1 + s1 * i1)) {-# INLINE iterTargetRowMajorAccST #-} -- | Constraint synonym that encapsulates all constraints needed for dimension 4 and higher. -- -- @since 1.0.0 -type HighIxN n - = (4 <= n, KnownNat n, KnownNat (n - 1), Index (IxN (n - 1)), IxN (n - 1) ~ Ix (n - 1)) +type HighIxN n = + (4 <= n, KnownNat n, KnownNat (n - 1), Index (IxN (n - 1)), IxN (n - 1) ~ Ix (n - 1)) instance {-# OVERLAPPABLE #-} HighIxN n => Index (IxN n) where type Dimensions (IxN n) = n @@ -464,12 +461,12 @@ pullOutDimM ix@(i :> ixl) d | d == dimensions (Proxy :: Proxy (IxN n)) = pure (i, ixl) | otherwise = - maybe (throwM $ IndexDimensionException ix d) (pure . fmap (i :>)) (pullOutDimM ixl d) + maybe (throwM $ IndexDimensionException ix d) (pure . fmap (i :>)) (pullOutDimM ixl d) {-# INLINE [1] pullOutDimM #-} insertDimM ix@(i :> ixl) d di | d == dimensions (Proxy :: Proxy (IxN n)) = pure (di :> ix) | otherwise = - maybe (throwM $ IndexDimensionException ix d) (pure . (i :>)) (insertDimM ixl d di) + maybe (throwM $ IndexDimensionException ix d) (pure . (i :>)) (insertDimM ixl d di) {-# INLINE [1] insertDimM #-} pureIndex i = i :> (pureIndex i :: Ix (n - 1)) {-# INLINE [1] pureIndex #-} @@ -481,8 +478,6 @@ repairIndex (SafeSz n) i rBelow rOver :> repairIndex (SafeSz szL) ixL rBelow rOver {-# INLINE [1] repairIndex #-} - - ---- Unbox Ix -- | Unboxing of a `Ix2`. @@ -520,7 +515,6 @@ {-# INLINE basicInitialize #-} #endif - newtype instance VU.Vector Ix2 = V_Ix2 (VU.Vector (Int, Int)) instance V.Vector VU.Vector Ix2 where @@ -539,16 +533,12 @@ elemseq _ = seq {-# INLINE elemseq #-} - - ---- Unbox Ix - - -- | Unboxing of a `IxN`. instance (3 <= n, VU.Unbox (Ix (n - 1))) => VU.Unbox (IxN n) -newtype instance VU.MVector s (IxN n) = MV_IxN (VU.MVector s Int, VU.MVector s (Ix (n-1))) +newtype instance VU.MVector s (IxN n) = MV_IxN (VU.MVector s Int, VU.MVector s (Ix (n - 1))) instance (3 <= n, VU.Unbox (Ix (n - 1))) => VM.MVector VU.MVector (IxN n) where basicLength (MV_IxN (_, mvec)) = VM.basicLength mvec @@ -599,8 +589,7 @@ {-# INLINE basicInitialize #-} #endif - -newtype instance VU.Vector (IxN n) = V_IxN (VU.Vector Int, VU.Vector (Ix (n-1))) +newtype instance VU.Vector (IxN n) = V_IxN (VU.Vector Int, VU.Vector (Ix (n - 1))) instance (3 <= n, VU.Unbox (Ix (n - 1))) => V.Vector VU.Vector (IxN n) where basicUnsafeFreeze (MV_IxN (mvec1, mvec)) = do diff --git a/src/Data/Massiv/Core/Index/Stride.hs b/src/Data/Massiv/Core/Index/Stride.hs --- a/src/Data/Massiv/Core/Index/Stride.hs +++ b/src/Data/Massiv/Core/Index/Stride.hs @@ -1,6 +1,7 @@ {-# LANGUAGE CPP #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE PatternSynonyms #-} + -- | -- Module : Data.Massiv.Core.Index.Stride -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -8,20 +9,19 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Core.Index.Stride - ( Stride(SafeStride) - , pattern Stride - , unStride - , oneStride - , toLinearIndexStride - , strideStart - , strideSize - ) where +module Data.Massiv.Core.Index.Stride ( + Stride (SafeStride), + pattern Stride, + unStride, + oneStride, + toLinearIndexStride, + strideStart, + strideSize, +) where import Control.DeepSeq (NFData) import Data.Massiv.Core.Index.Internal -import System.Random.Stateful (Random, Uniform(..), UniformRange(..)) +import System.Random.Stateful (Random, Uniform (..), UniformRange (..)) -- | Stride provides a way to ignore elements of an array if an index is divisible by a -- corresponding value in a stride. So, for a @Stride (i :. j)@ only elements with indices will be @@ -51,21 +51,20 @@ -- @since 0.2.1 newtype Stride ix = SafeStride ix deriving (Eq, Ord, NFData) - -- | A safe bidirectional pattern synonym for `Stride` construction that will make sure stride -- elements are always positive. -- -- @since 0.2.1 pattern Stride :: Index ix => ix -> Stride ix -pattern Stride ix <- SafeStride ix where - Stride ix = SafeStride (liftIndex (max 1) ix) -{-# COMPLETE Stride #-} +pattern Stride ix <- SafeStride ix + where + Stride ix = SafeStride (liftIndex (max 1) ix) +{-# COMPLETE Stride #-} instance Index ix => Show (Stride ix) where showsPrec n (SafeStride ix) = showsPrecWrapped n (("Stride " ++) . showsPrec 1 ix) - instance (UniformRange ix, Index ix) => Uniform (Stride ix) where uniformM g = SafeStride <$> uniformRM (pureIndex 1, pureIndex maxBound) g {-# INLINE uniformM #-} @@ -76,7 +75,6 @@ instance (UniformRange ix, Index ix) => Random (Stride ix) - -- | Just a helper function for unwrapping `Stride`. -- -- @since 0.2.1 @@ -106,21 +104,21 @@ -- | Compute linear index with stride using the original size and index -- -- @since 0.2.1 -toLinearIndexStride :: - Index ix - => Stride ix -- ^ Stride - -> Sz ix -- ^ Size - -> ix -- ^ Index +toLinearIndexStride + :: Index ix + => Stride ix + -- ^ Stride + -> Sz ix + -- ^ Size + -> ix + -- ^ Index -> Int toLinearIndexStride (SafeStride stride) sz ix = toLinearIndex sz (liftIndex2 div ix stride) {-# INLINE toLinearIndexStride #-} - -- | A default stride of @1@, where all elements are kept -- -- @since 0.2.1 oneStride :: Index ix => Stride ix oneStride = SafeStride (pureIndex 1) {-# INLINE oneStride #-} - - diff --git a/src/Data/Massiv/Core/Index/Tuple.hs b/src/Data/Massiv/Core/Index/Tuple.hs --- a/src/Data/Massiv/Core/Index/Tuple.hs +++ b/src/Data/Massiv/Core/Index/Tuple.hs @@ -1,9 +1,10 @@ -{-# OPTIONS_GHC -fno-warn-orphans #-} {-# LANGUAGE BangPatterns #-} {-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE TypeFamilies #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} + -- | -- Module : Data.Massiv.Core.Index.Tuple -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -11,32 +12,40 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Core.Index.Tuple - ( -- * Tuple based indices - -- ** 1-dimensional - Ix1T - -- ** 2-dimensional - , Ix2T - , toIx2 - , fromIx2 - -- ** 3-dimensional - , Ix3T - , toIx3 - , fromIx3 - -- ** 4-dimensional - , Ix4T - , toIx4 - , fromIx4 - -- ** 5-dimensional - , Ix5T - , toIx5 - , fromIx5 - ) where +module Data.Massiv.Core.Index.Tuple ( + -- * Tuple based indices -import Control.Monad.Catch (MonadThrow(..)) -import Data.Massiv.Core.Index.Internal (Index(..), IndexException(..), Lower, - Sz(..)) + -- ** 1-dimensional + Ix1T, + + -- ** 2-dimensional + Ix2T, + toIx2, + fromIx2, + + -- ** 3-dimensional + Ix3T, + toIx3, + fromIx3, + + -- ** 4-dimensional + Ix4T, + toIx4, + fromIx4, + + -- ** 5-dimensional + Ix5T, + toIx5, + fromIx5, +) where + +import Control.Monad.Catch (MonadThrow (..)) +import Data.Massiv.Core.Index.Internal ( + Index (..), + IndexException (..), + Lower, + Sz (..), + ) import Data.Massiv.Core.Index.Ix -- | Another 1-dimensional index type synonym for `Int`, same as `Ix1` and is here just for @@ -60,8 +69,6 @@ type instance Lower Ix4T = Ix3T type instance Lower Ix5T = Ix4T - - -- | Convert an `Int` tuple to `Ix2` -- -- ==== __Example__ @@ -178,39 +185,38 @@ {-# INLINE [1] snocDim #-} unsnocDim = id {-# INLINE [1] unsnocDim #-} - getDimM (i2, _) 2 = pure i2 - getDimM ( _, i1) 1 = pure i1 - getDimM ix d = throwM $ IndexDimensionException ix d + getDimM (i2, _) 2 = pure i2 + getDimM (_, i1) 1 = pure i1 + getDimM ix d = throwM $ IndexDimensionException ix d {-# INLINE [1] getDimM #-} setDimM (_, i1) 2 i2 = pure (i2, i1) setDimM (i2, _) 1 i1 = pure (i2, i1) - setDimM ix d _ = throwM $ IndexDimensionException ix d + setDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] setDimM #-} - modifyDimM (i2, i1) 2 f = pure (i2, (f i2, i1)) - modifyDimM (i2, i1) 1 f = pure (i1, ( i2, f i1)) - modifyDimM ix d _ = throwM $ IndexDimensionException ix d + modifyDimM (i2, i1) 2 f = pure (i2, (f i2, i1)) + modifyDimM (i2, i1) 1 f = pure (i1, (i2, f i1)) + modifyDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] modifyDimM #-} pullOutDimM (i2, i1) 2 = pure (i2, i1) pullOutDimM (i2, i1) 1 = pure (i1, i2) - pullOutDimM ix d = throwM $ IndexDimensionException ix d + pullOutDimM ix d = throwM $ IndexDimensionException ix d {-# INLINE [1] pullOutDimM #-} insertDimM i1 2 i2 = pure (i2, i1) insertDimM i2 1 i1 = pure (i2, i1) - insertDimM ix d _ = throwM $ IndexDimensionException ix d + insertDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] insertDimM #-} pureIndex i = (i, i) {-# INLINE [1] pureIndex #-} liftIndex2 f (i2, i1) (i2', i1') = (f i2 i2', f i1 i1') {-# INLINE [1] liftIndex2 #-} - -- | -- @since 0.1.0 instance Index Ix3T where type Dimensions Ix3T = 3 dimensions _ = 3 {-# INLINE [1] dimensions #-} - totalElem (SafeSz (k3, k2, k1)) = k3 * k2 * k1 + totalElem (SafeSz (k3, k2, k1)) = k3 * k2 * k1 {-# INLINE [1] totalElem #-} consDim i3 (i2, i1) = (i3, i2, i1) {-# INLINE [1] consDim #-} @@ -220,36 +226,35 @@ {-# INLINE [1] snocDim #-} unsnocDim (i3, i2, i1) = ((i3, i2), i1) {-# INLINE [1] unsnocDim #-} - getDimM (i3, _, _) 3 = pure i3 - getDimM ( _, i2, _) 2 = pure i2 - getDimM ( _, _, i1) 1 = pure i1 - getDimM ix d = throwM $ IndexDimensionException ix d + getDimM (i3, _, _) 3 = pure i3 + getDimM (_, i2, _) 2 = pure i2 + getDimM (_, _, i1) 1 = pure i1 + getDimM ix d = throwM $ IndexDimensionException ix d {-# INLINE [1] getDimM #-} - setDimM ( _, i2, i1) 3 i3 = pure (i3, i2, i1) - setDimM (i3, _, i1) 2 i2 = pure (i3, i2, i1) - setDimM (i3, i2, _) 1 i1 = pure (i3, i2, i1) - setDimM ix d _ = throwM $ IndexDimensionException ix d + setDimM (_, i2, i1) 3 i3 = pure (i3, i2, i1) + setDimM (i3, _, i1) 2 i2 = pure (i3, i2, i1) + setDimM (i3, i2, _) 1 i1 = pure (i3, i2, i1) + setDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] setDimM #-} - modifyDimM (i3, i2, i1) 3 f = pure (i3, (f i3, i2, i1)) - modifyDimM (i3, i2, i1) 2 f = pure (i2, ( i3, f i2, i1)) - modifyDimM (i3, i2, i1) 1 f = pure (i1, ( i3, i2, f i1)) - modifyDimM ix d _ = throwM $ IndexDimensionException ix d + modifyDimM (i3, i2, i1) 3 f = pure (i3, (f i3, i2, i1)) + modifyDimM (i3, i2, i1) 2 f = pure (i2, (i3, f i2, i1)) + modifyDimM (i3, i2, i1) 1 f = pure (i1, (i3, i2, f i1)) + modifyDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] modifyDimM #-} pullOutDimM (i3, i2, i1) 3 = pure (i3, (i2, i1)) pullOutDimM (i3, i2, i1) 2 = pure (i2, (i3, i1)) pullOutDimM (i3, i2, i1) 1 = pure (i1, (i3, i2)) - pullOutDimM ix d = throwM $ IndexDimensionException ix d + pullOutDimM ix d = throwM $ IndexDimensionException ix d {-# INLINE [1] pullOutDimM #-} insertDimM (i2, i1) 3 i3 = pure (i3, i2, i1) insertDimM (i3, i1) 2 i2 = pure (i3, i2, i1) insertDimM (i3, i2) 1 i1 = pure (i3, i2, i1) - insertDimM ix d _ = throwM $ IndexDimensionException ix d + insertDimM ix d _ = throwM $ IndexDimensionException ix d pureIndex i = (i, i, i) {-# INLINE [1] pureIndex #-} liftIndex2 f (i3, i2, i1) (i3', i2', i1') = (f i3 i3', f i2 i2', f i1 i1') {-# INLINE [1] liftIndex2 #-} - instance Index Ix4T where type Dimensions Ix4T = 4 dimensions _ = 4 @@ -264,42 +269,41 @@ {-# INLINE [1] snocDim #-} unsnocDim (i4, i3, i2, i1) = ((i4, i3, i2), i1) {-# INLINE [1] unsnocDim #-} - getDimM (i4, _, _, _) 4 = pure i4 - getDimM ( _, i3, _, _) 3 = pure i3 - getDimM ( _, _, i2, _) 2 = pure i2 - getDimM ( _, _, _, i1) 1 = pure i1 - getDimM ix d = throwM $ IndexDimensionException ix d + getDimM (i4, _, _, _) 4 = pure i4 + getDimM (_, i3, _, _) 3 = pure i3 + getDimM (_, _, i2, _) 2 = pure i2 + getDimM (_, _, _, i1) 1 = pure i1 + getDimM ix d = throwM $ IndexDimensionException ix d {-# INLINE [1] getDimM #-} - setDimM ( _, i3, i2, i1) 4 i4 = pure (i4, i3, i2, i1) - setDimM (i4, _, i2, i1) 3 i3 = pure (i4, i3, i2, i1) - setDimM (i4, i3, _, i1) 2 i2 = pure (i4, i3, i2, i1) - setDimM (i4, i3, i2, _) 1 i1 = pure (i4, i3, i2, i1) - setDimM ix d _ = throwM $ IndexDimensionException ix d + setDimM (_, i3, i2, i1) 4 i4 = pure (i4, i3, i2, i1) + setDimM (i4, _, i2, i1) 3 i3 = pure (i4, i3, i2, i1) + setDimM (i4, i3, _, i1) 2 i2 = pure (i4, i3, i2, i1) + setDimM (i4, i3, i2, _) 1 i1 = pure (i4, i3, i2, i1) + setDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] setDimM #-} - modifyDimM (i4, i3, i2, i1) 4 f = pure (i4, (f i4, i3, i2, i1)) - modifyDimM (i4, i3, i2, i1) 3 f = pure (i3, ( i4, f i3, i2, i1)) - modifyDimM (i4, i3, i2, i1) 2 f = pure (i2, ( i4, i3, f i2, i1)) - modifyDimM (i4, i3, i2, i1) 1 f = pure (i1, ( i4, i3, i2, f i1)) - modifyDimM ix d _ = throwM $ IndexDimensionException ix d + modifyDimM (i4, i3, i2, i1) 4 f = pure (i4, (f i4, i3, i2, i1)) + modifyDimM (i4, i3, i2, i1) 3 f = pure (i3, (i4, f i3, i2, i1)) + modifyDimM (i4, i3, i2, i1) 2 f = pure (i2, (i4, i3, f i2, i1)) + modifyDimM (i4, i3, i2, i1) 1 f = pure (i1, (i4, i3, i2, f i1)) + modifyDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] modifyDimM #-} pullOutDimM (i4, i3, i2, i1) 4 = pure (i4, (i3, i2, i1)) pullOutDimM (i4, i3, i2, i1) 3 = pure (i3, (i4, i2, i1)) pullOutDimM (i4, i3, i2, i1) 2 = pure (i2, (i4, i3, i1)) pullOutDimM (i4, i3, i2, i1) 1 = pure (i1, (i4, i3, i2)) - pullOutDimM ix d = throwM $ IndexDimensionException ix d + pullOutDimM ix d = throwM $ IndexDimensionException ix d {-# INLINE [1] pullOutDimM #-} insertDimM (i3, i2, i1) 4 i4 = pure (i4, i3, i2, i1) insertDimM (i4, i2, i1) 3 i3 = pure (i4, i3, i2, i1) insertDimM (i4, i3, i1) 2 i2 = pure (i4, i3, i2, i1) insertDimM (i4, i3, i2) 1 i1 = pure (i4, i3, i2, i1) - insertDimM ix d _ = throwM $ IndexDimensionException ix d + insertDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] insertDimM #-} pureIndex i = (i, i, i, i) {-# INLINE [1] pureIndex #-} liftIndex2 f (i4, i3, i2, i1) (i4', i3', i2', i1') = (f i4 i4', f i3 i3', f i2 i2', f i1 i1') {-# INLINE [1] liftIndex2 #-} - instance Index Ix5T where type Dimensions Ix5T = 5 dimensions _ = 5 @@ -314,40 +318,40 @@ {-# INLINE [1] snocDim #-} unsnocDim (i5, i4, i3, i2, i1) = ((i5, i4, i3, i2), i1) {-# INLINE [1] unsnocDim #-} - getDimM (i5, _, _, _, _) 5 = pure i5 - getDimM ( _, i4, _, _, _) 4 = pure i4 - getDimM ( _, _, i3, _, _) 3 = pure i3 - getDimM ( _, _, _, i2, _) 2 = pure i2 - getDimM ( _, _, _, _, i1) 1 = pure i1 - getDimM ix d = throwM $ IndexDimensionException ix d + getDimM (i5, _, _, _, _) 5 = pure i5 + getDimM (_, i4, _, _, _) 4 = pure i4 + getDimM (_, _, i3, _, _) 3 = pure i3 + getDimM (_, _, _, i2, _) 2 = pure i2 + getDimM (_, _, _, _, i1) 1 = pure i1 + getDimM ix d = throwM $ IndexDimensionException ix d {-# INLINE [1] getDimM #-} - setDimM ( _, i4, i3, i2, i1) 5 i5 = pure (i5, i4, i3, i2, i1) - setDimM (i5, _, i3, i2, i1) 4 i4 = pure (i5, i4, i3, i2, i1) - setDimM (i5, i4, _, i2, i1) 3 i3 = pure (i5, i4, i3, i2, i1) - setDimM (i5, i4, i3, _, i1) 2 i2 = pure (i5, i4, i3, i2, i1) - setDimM (i5, i4, i3, i2, _) 1 i1 = pure (i5, i4, i3, i2, i1) - setDimM ix d _ = throwM $ IndexDimensionException ix d + setDimM (_, i4, i3, i2, i1) 5 i5 = pure (i5, i4, i3, i2, i1) + setDimM (i5, _, i3, i2, i1) 4 i4 = pure (i5, i4, i3, i2, i1) + setDimM (i5, i4, _, i2, i1) 3 i3 = pure (i5, i4, i3, i2, i1) + setDimM (i5, i4, i3, _, i1) 2 i2 = pure (i5, i4, i3, i2, i1) + setDimM (i5, i4, i3, i2, _) 1 i1 = pure (i5, i4, i3, i2, i1) + setDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] setDimM #-} - modifyDimM (i5, i4, i3, i2, i1) 5 f = pure (i5, (f i5, i4, i3, i2, i1)) - modifyDimM (i5, i4, i3, i2, i1) 4 f = pure (i4, ( i5, f i4, i3, i2, i1)) - modifyDimM (i5, i4, i3, i2, i1) 3 f = pure (i3, ( i5, i4, f i3, i2, i1)) - modifyDimM (i5, i4, i3, i2, i1) 2 f = pure (i2, ( i5, i4, i3, f i2, i1)) - modifyDimM (i5, i4, i3, i2, i1) 1 f = pure (i1, ( i5, i4, i3, i2, f i1)) - modifyDimM ix d _ = throwM $ IndexDimensionException ix d + modifyDimM (i5, i4, i3, i2, i1) 5 f = pure (i5, (f i5, i4, i3, i2, i1)) + modifyDimM (i5, i4, i3, i2, i1) 4 f = pure (i4, (i5, f i4, i3, i2, i1)) + modifyDimM (i5, i4, i3, i2, i1) 3 f = pure (i3, (i5, i4, f i3, i2, i1)) + modifyDimM (i5, i4, i3, i2, i1) 2 f = pure (i2, (i5, i4, i3, f i2, i1)) + modifyDimM (i5, i4, i3, i2, i1) 1 f = pure (i1, (i5, i4, i3, i2, f i1)) + modifyDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] modifyDimM #-} pullOutDimM (i5, i4, i3, i2, i1) 5 = pure (i5, (i4, i3, i2, i1)) pullOutDimM (i5, i4, i3, i2, i1) 4 = pure (i4, (i5, i3, i2, i1)) pullOutDimM (i5, i4, i3, i2, i1) 3 = pure (i3, (i5, i4, i2, i1)) pullOutDimM (i5, i4, i3, i2, i1) 2 = pure (i2, (i5, i4, i3, i1)) pullOutDimM (i5, i4, i3, i2, i1) 1 = pure (i1, (i5, i4, i3, i2)) - pullOutDimM ix d = throwM $ IndexDimensionException ix d + pullOutDimM ix d = throwM $ IndexDimensionException ix d {-# INLINE [1] pullOutDimM #-} insertDimM (i4, i3, i2, i1) 5 i5 = pure (i5, i4, i3, i2, i1) insertDimM (i5, i3, i2, i1) 4 i4 = pure (i5, i4, i3, i2, i1) insertDimM (i5, i4, i2, i1) 3 i3 = pure (i5, i4, i3, i2, i1) insertDimM (i5, i4, i3, i1) 2 i2 = pure (i5, i4, i3, i2, i1) insertDimM (i5, i4, i3, i2) 1 i1 = pure (i5, i4, i3, i2, i1) - insertDimM ix d _ = throwM $ IndexDimensionException ix d + insertDimM ix d _ = throwM $ IndexDimensionException ix d {-# INLINE [1] insertDimM #-} pureIndex i = (i, i, i, i, i) {-# INLINE [1] pureIndex #-} diff --git a/src/Data/Massiv/Core/List.hs b/src/Data/Massiv/Core/List.hs --- a/src/Data/Massiv/Core/List.hs +++ b/src/Data/Massiv/Core/List.hs @@ -9,6 +9,7 @@ {-# LANGUAGE TypeOperators #-} {-# LANGUAGE UndecidableInstances #-} {-# OPTIONS_GHC -fno-warn-orphans #-} + -- | -- Module : Data.Massiv.Core.List -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -16,16 +17,15 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Core.List - ( L(..) - , Array(..) - , List(..) - , toListArray - , showsArrayPrec - , showArrayList - , ListItem - ) where +module Data.Massiv.Core.List ( + L (..), + Array (..), + List (..), + toListArray, + showsArrayPrec, + showArrayList, + ListItem, +) where import Control.Monad (unless, when) import Control.Scheduler @@ -41,17 +41,15 @@ import GHC.TypeLits import System.IO.Unsafe (unsafePerformIO) - type family ListItem ix e :: Type where ListItem Ix1 e = e - ListItem ix e = [ListItem (Lower ix) e] + ListItem ix e = [ListItem (Lower ix) e] type family Elt ix e :: Type where Elt Ix1 e = e - Elt ix e = List (Lower ix) e - -newtype List ix e = List { unList :: [Elt ix e] } + Elt ix e = List (Lower ix) e +newtype List ix e = List {unList :: [Elt ix e]} instance Coercible (Elt ix e) (ListItem ix e) => IsList (List ix e) where type Item (List ix e) = ListItem ix e @@ -60,13 +58,12 @@ toList = coerce {-# INLINE toList #-} - data L = L -data instance Array L ix e = LArray { lComp :: Comp - , lData :: !(List ix e) - } - +data instance Array L ix e = LArray + { lComp :: Comp + , lData :: !(List ix e) + } instance Coercible (Elt ix e) (ListItem ix e) => IsList (Array L ix e) where type Item (Array L ix e) = ListItem ix e @@ -79,7 +76,7 @@ lengthHintList = \case [] -> LengthExact zeroSz - _ -> LengthUnknown + _ -> LengthUnknown {-# INLINE lengthHintList #-} instance Shape L Ix1 where @@ -101,8 +98,8 @@ {-# INLINE isNull #-} outerSize arr = case unList (lData arr) of - [] -> zeroSz - (x:xs) -> SafeSz ((1 + length xs) :. length (unList x)) + [] -> 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 @@ -114,16 +111,14 @@ {-# INLINE isNull #-} outerSize arr = case unList (lData arr) of - [] -> zeroSz - (x:xs) -> SafeSz ((1 + length xs) :> unSz (outerSize (LArray Seq x))) + [] -> zeroSz + (x : xs) -> SafeSz ((1 + length xs) :> unSz (outerSize (LArray Seq x))) {-# INLINE outerSize #-} - outerLength :: Array L ix e -> Sz Int outerLength = SafeSz . length . unList . lData {-# INLINE outerLength #-} - instance Ragged L Ix1 e where flattenRagged = id {-# INLINE flattenRagged #-} @@ -136,86 +131,88 @@ {-# INLINE generateRaggedM #-} loadRaggedST _scheduler xs uWrite start end sz = go (unList (lData xs)) start where - go (y:ys) i + 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 (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 + where + !sz = outerSize arr {-# INLINE iterArrayLinearST_ #-} instance Ragged L Ix2 e where generateRaggedM = unsafeGenerateParM {-# INLINE generateRaggedM #-} - flattenRagged arr = LArray {lComp = lComp arr, lData = coerce xs} + flattenRagged arr = LArray{lComp = lComp arr, lData = coerce xs} where xs = concatMap (unList . lData . flattenRagged . LArray (lComp arr)) (unList (lData arr)) {-# INLINE flattenRagged #-} 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) + let (k, szL) = unconsSz sz + step = totalElem szL + leftOver <- + loopM start (< end) (+ step) (coerce (lData xs)) $ \i zs -> + case zs of + [] -> throwM (DimTooShortException 2 k (outerLength xs)) + (y : ys) -> do + scheduleWork_ scheduler $ + let end' = i + step + go (a : as) j + | j < end' = uWrite j a >> go as (j + 1) + | otherwise = throwM $ DimTooLongException 1 szL (Sz (length y)) + go [] j = when (j /= end') $ throwM (DimTooShortException 1 szL (Sz (length y))) + in go y i + pure ys + unless (null leftOver) $ throwM $ DimTooLongException 2 k (outerLength xs) {-# INLINE loadRaggedST #-} raggedFormat f sep (LArray comp xs) = showN (\s y -> raggedFormat f s (LArray comp y :: Array L Ix1 e)) sep (coerce xs) -instance ( Shape L (IxN n) - , Ragged L (Ix (n - 1)) e - , Coercible (Elt (Ix (n - 1)) e) (ListItem (Ix (n - 1)) e) - ) => - Ragged L (IxN n) e where +instance + ( Shape L (IxN n) + , Ragged L (Ix (n - 1)) e + , Coercible (Elt (Ix (n - 1)) e) (ListItem (Ix (n - 1)) e) + ) + => Ragged L (IxN n) e + where generateRaggedM = unsafeGenerateParM {-# INLINE generateRaggedM #-} - flattenRagged arr = LArray {lComp = lComp arr, lData = coerce xs} + flattenRagged arr = LArray{lComp = lComp arr, lData = coerce xs} where xs = concatMap (unList . lData . flattenRagged . LArray (lComp arr)) (unList (lData arr)) {-# INLINE flattenRagged #-} loadRaggedST scheduler xs uWrite start end sz | isZeroSz sz = when (isNotNull (flattenRagged xs)) (throwM ShapeNonEmpty) | otherwise = do - let (k, szL) = unconsSz sz - step = totalElem szL - subScheduler - | end - start < numWorkers scheduler * step = scheduler - | otherwise = trivialScheduler_ - leftOver <- - loopM start (< end) (+ step) (unList (lData xs)) $ \i zs -> - case zs of - [] -> throwM (DimTooShortException (dimensions sz) k (outerLength xs)) - (y:ys) -> do - scheduleWork_ scheduler $ - loadRaggedST subScheduler (LArray Seq y) uWrite i (i + step) szL - pure ys - unless (null leftOver) $ throwM $ DimTooLongException (dimensions sz) k (outerLength xs) + let (k, szL) = unconsSz sz + step = totalElem szL + subScheduler + | end - start < numWorkers scheduler * step = scheduler + | otherwise = trivialScheduler_ + leftOver <- + loopM start (< end) (+ step) (unList (lData xs)) $ \i zs -> + case zs of + [] -> throwM (DimTooShortException (dimensions sz) k (outerLength xs)) + (y : ys) -> do + scheduleWork_ scheduler $ + loadRaggedST subScheduler (LArray Seq y) uWrite i (i + step) szL + pure ys + unless (null leftOver) $ throwM $ DimTooLongException (dimensions sz) k (outerLength xs) {-# INLINE loadRaggedST #-} raggedFormat f sep (LArray comp xs) = showN (\s y -> raggedFormat f s (LArray comp y :: Array L (Ix (n - 1)) e)) sep (coerce xs) -unsafeGenerateParM :: - (Elt ix e ~ List (Lower ix) e, Index ix, Monad m, Ragged L (Lower ix) e) +unsafeGenerateParM + :: (Elt ix e ~ List (Lower ix) e, Index ix, Monad m, Ragged L (Lower ix) e) => Comp -> Sz ix -> (ix -> m e) @@ -224,23 +221,23 @@ res <- sequence $ unsafePerformIO $ do let !(ksz, szL) = unconsSz sz !k = unSz ksz - withScheduler comp $ \ scheduler -> - splitLinearly (numWorkers scheduler) k $ \ chunkLength slackStart -> do + withScheduler comp $ \scheduler -> + splitLinearly (numWorkers scheduler) k $ \chunkLength slackStart -> do loopA_ 0 (< slackStart) (+ chunkLength) $ \ !start -> scheduleWork scheduler $ do res <- loopDeepM start (< (start + chunkLength)) (+ 1) [] $ \i acc -> - return (fmap lData (generateRaggedM Seq szL (\ !ixL -> f (consDim i ixL))):acc) + return (fmap lData (generateRaggedM Seq szL (\ !ixL -> f (consDim i ixL))) : acc) return $! sequence res when (slackStart < k) $ scheduleWork scheduler $ do res <- loopDeepM slackStart (< k) (+ 1) [] $ \i acc -> - return (fmap lData (generateRaggedM Seq szL (\ !ixL -> f (consDim i ixL))):acc) + return (fmap lData (generateRaggedM Seq szL (\ !ixL -> f (consDim i ixL))) : acc) return $! sequence res return $ LArray comp $ List $ concat res {-# INLINE unsafeGenerateParM #-} instance Strategy L where - setComp c arr = arr {lComp = c} + setComp c arr = arr{lComp = c} {-# INLINE setComp #-} getComp = lComp {-# INLINE getComp #-} @@ -263,7 +260,6 @@ -- return $! foldr' consR (emptyR comp) xs -- {-# INLINE unsafeGenerateN #-} - -- | Construct an array backed by linked lists from any source array -- -- @since 0.4.0 @@ -271,50 +267,57 @@ 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 n arr = showsArrayLAsPrec (Proxy :: Proxy L) (outerSize arr) n arr + showsPrec n arr = showsArrayLAsPrec (Proxy :: Proxy L) (outerSize arr) n arr instance (Ragged L ix e, Show e) => Show (List ix e) where show xs = " " ++ raggedFormat show "\n " arrL - where arrL = LArray Seq xs :: Array L ix e - + where + arrL = LArray Seq xs :: Array L ix e showN :: (String -> a -> String) -> String -> [a] -> String -showN _ _ [] = "[ ]" +showN _ _ [] = "[ ]" showN fShow lnPrefix ls = L.concat - (["[ "] ++ - L.intersperse (lnPrefix ++ ", ") (map (fShow (lnPrefix ++ " ")) ls) ++ [lnPrefix, "]"]) - + ( ["[ "] + ++ L.intersperse (lnPrefix ++ ", ") (map (fShow (lnPrefix ++ " ")) ls) + ++ [lnPrefix, "]"] + ) -showsArrayLAsPrec :: - forall r ix e. (Ragged L ix e, Typeable r, Show e) +showsArrayLAsPrec + :: forall r ix e + . (Ragged L ix e, Typeable r, Show e) => Proxy r -> Sz ix -> Int -> Array L ix e -- Array to show -> ShowS showsArrayLAsPrec pr sz n arr = - opp . - ("Array " ++) . - showsTypeRep (typeRep pr) . - (' ':) . - showsPrec 1 (getComp arr) . (" (" ++) . shows sz . (")\n" ++) . shows lnarr . clp + opp + . ("Array " ++) + . showsTypeRep (typeRep pr) + . (' ' :) + . showsPrec 1 (getComp arr) + . (" (" ++) + . shows sz + . (")\n" ++) + . shows lnarr + . clp where (opp, clp) = if n == 0 then (id, id) - else (('(':), ("\n)" ++)) + else (('(' :), ("\n)" ++)) lnarr = lData arr -- | Helper function for declaring `Show` instances for arrays -- -- @since 0.4.0 -showsArrayPrec :: - forall r r' ix e. (Ragged L ix e, Load r ix e, Load r' ix e, Source r' e, Show e) - => (Array r ix e -> Array r' ix e) -- ^ Modifier +showsArrayPrec + :: forall r r' ix e + . (Ragged L ix e, Load r ix e, Load r' ix e, Source r' e, Show e) + => (Array r ix e -> Array r' ix e) + -- ^ Modifier -> Int -> Array r ix e -- Array to show -> ShowS @@ -324,18 +327,16 @@ arr' = f arr larr = makeArray (getComp arr') sz (evaluate' arr') :: Array L ix e - -- | Helper function for declaring `Show` instances for arrays -- -- @since 0.4.0 showArrayList :: Show arr => [arr] -> String -> String -showArrayList arrs = ('[':) . go arrs . (']':) +showArrayList arrs = ('[' :) . go arrs . (']' :) where - go [] = id - go [x] = (' ':) . shows x . ('\n':) - go (x:xs) = (' ':) . shows x . ("\n," ++) . go xs - + go [] = id + go [x] = (' ' :) . shows x . ('\n' :) + go (x : xs) = (' ' :) . shows x . ("\n," ++) . go xs instance Stream L Ix1 e where toStream = S.fromList . unList . lData diff --git a/src/Data/Massiv/Core/Loop.hs b/src/Data/Massiv/Core/Loop.hs --- a/src/Data/Massiv/Core/Loop.hs +++ b/src/Data/Massiv/Core/Loop.hs @@ -2,6 +2,7 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE MonoLocalBinds #-} {-# LANGUAGE ScopedTypeVariables #-} + -- | -- Module : Data.Massiv.Core.Loop -- Copyright : (c) Alexey Kuleshevich 2018-2022 @@ -9,46 +10,54 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Core.Loop - ( loop - , loopF - , nextMaybeF - , loopA - , loopA_ - , loopM - , loopM_ - , iloopM - , iloopA_ - , loopNextM - , loopNextA_ - , loopDeepM - , splitLinearly - , splitLinearlyM - , splitLinearlyM_ - , splitLinearlyWith_ - , splitLinearlyWithM_ - , splitLinearlyWithStartAtM_ - , splitLinearlyWithStatefulM_ - , iterLinearST_ - , iterLinearAccST_ - , iterLinearAccST - , splitNumChunks - , stepStartAdjust +module Data.Massiv.Core.Loop ( + loop, + loopF, + nextMaybeF, + loopA, + loopA_, + loopM, + loopM_, + iloopM, + iloopA_, + loopNextM, + loopNextA_, + loopDeepM, + splitLinearly, + splitLinearlyM, + splitLinearlyM_, + splitLinearlyWith_, + splitLinearlyWithM_, + splitLinearlyWithStartAtM_, + splitLinearlyWithStatefulM_, + iterLinearST_, + iterLinearAccST_, + iterLinearAccST, + splitNumChunks, + stepStartAdjust, + -- * Experimental - , splitWorkWithFactorST - , scheduleMassivWork - , withMassivScheduler_ - ) where + splitWorkWithFactorST, + scheduleMassivWork, + withMassivScheduler_, +) where import Control.Monad (void, when) -import Control.Monad.IO.Unlift (MonadUnliftIO(..)) +import Control.Monad.IO.Unlift (MonadUnliftIO (..)) import Control.Monad.Primitive import Control.Monad.ST (ST) -import Control.Scheduler (Comp(..), Scheduler, SchedulerWS, - numWorkers, scheduleWork, scheduleWorkState_, - scheduleWork_, trivialScheduler_, unwrapSchedulerWS, - withScheduler_) +import Control.Scheduler ( + Comp (..), + Scheduler, + SchedulerWS, + numWorkers, + scheduleWork, + scheduleWorkState_, + scheduleWork_, + trivialScheduler_, + unwrapSchedulerWS, + withScheduler_, + ) import Control.Scheduler.Global (globalScheduler, withGlobalScheduler_) import Data.Coerce import Data.Functor.Identity @@ -76,15 +85,14 @@ | otherwise = pure acc {-# INLINE loopM #-} - -- | Efficient monadic loop with an accumulator and extra linear index incremented by 1. -- -- >>> iloopM 100 1 (< 20) (+ 2) [] (\i ix a -> Just ((i, ix) : a)) -- Just [(109,19),(108,17),(107,15),(106,13),(105,11),(104,9),(103,7),(102,5),(101,3),(100,1)] -- -- @since 1.0.2 -iloopM :: - Monad m => Int -> Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> Int -> a -> m a) -> m a +iloopM + :: Monad m => Int -> Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> Int -> a -> m a) -> m a iloopM !istart !initIx condition increment !initAcc f = go istart initIx initAcc where go !i !step !acc @@ -101,7 +109,7 @@ go !step | condition step = f step >> go (increment step) | otherwise = pure () - --loopF initial condition increment (pure ()) (\i ma -> f i >> ma) +-- loopF initial condition increment (pure ()) (\i ma -> f i >> ma) {-# INLINE loopM_ #-} {-# DEPRECATED loopM_ "In favor of `loopA_`" #-} @@ -115,8 +123,8 @@ -- (104,9) -- -- @since 1.0.2 -iloopA_ :: - Applicative f => Int -> Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> Int -> f a) -> f () +iloopA_ + :: Applicative f => Int -> Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> Int -> f a) -> f () iloopA_ !istart !initIx condition increment f = go istart initIx where go !i !step @@ -133,8 +141,8 @@ where go !step | condition step = - let !next = increment step - in f step next *> go next + let !next = increment step + in f step next *> go next | otherwise = pure () {-# INLINE loopNextA_ #-} @@ -147,8 +155,8 @@ where go !step !acc | condition step = - let !next = increment step - in f step next acc >>= go next + let !next = increment step + in f step next acc >>= go next | otherwise = pure acc {-# INLINE loopNextM #-} @@ -170,7 +178,6 @@ loopF initial condition increment lastAction (\i ma -> f i <*> ma) {-# INLINE loopA #-} - loopF :: Int -> (Int -> Bool) -> (Int -> Int) -> f a -> (Int -> f a -> f a) -> f a loopF !initial condition increment lastAction f = go initial where @@ -179,11 +186,10 @@ | otherwise = lastAction {-# INLINE loopF #-} - nextMaybeF :: Int -> (Int -> Bool) -> (Int -> Int) -> (Maybe Int -> f a) -> f a nextMaybeF !cur condition increment f = let !i = increment cur - in f $! if condition i then Just i else Nothing + in f $! if condition i then Just i else Nothing {-# INLINE nextMaybeF #-} -- | Similar to `loopM`, but way less efficient monadic loop with an accumulator that reverses @@ -203,14 +209,17 @@ loopF initial condition increment (pure initAcc) (\i ma -> ma >>= f i) {-# INLINE loopDeepM #-} - -- | Divide length in chunks and apply a function to the computed results -- -- @since 0.2.1 -splitLinearly :: Int -- ^ Number of chunks - -> Int -- ^ Total length - -> (Int -> Int -> a) -- ^ Function that accepts a chunk length and slack start index - -> a +splitLinearly + :: Int + -- ^ Number of chunks + -> Int + -- ^ Total length + -> (Int -> Int -> a) + -- ^ Function that accepts a chunk length and slack start index + -> a splitLinearly numChunks totalLength action = action chunkLength slackStart where !chunkLength = totalLength `quot` numChunks @@ -220,24 +229,25 @@ -- | Iterator that expects an action that accepts starting linear index as well as the ending -- -- @since 0.5.7 -splitLinearlyM_ :: - MonadPrimBase s m => Scheduler s () -> Int -> (Int -> Int -> m ()) -> m () +splitLinearlyM_ + :: MonadPrimBase s m => Scheduler s () -> Int -> (Int -> Int -> m ()) -> m () splitLinearlyM_ scheduler totalLength action = splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do - loopNextA_ 0 (< slackStart) (+ chunkLength) $ \ start next -> + loopNextA_ 0 (< slackStart) (+ chunkLength) $ \start next -> scheduleWork_ scheduler $ action start next when (slackStart < totalLength) $ - scheduleWork_ scheduler $ action slackStart totalLength + scheduleWork_ scheduler $ + action slackStart totalLength {-# INLINE splitLinearlyM_ #-} -- | Iterator that expects an action that accepts starting linear index as well as the ending -- -- @since 1.0.2 -splitLinearlyM :: - MonadPrimBase s m => Scheduler s a -> Int -> (Int -> Int -> m a) -> m () +splitLinearlyM + :: MonadPrimBase s m => Scheduler s a -> Int -> (Int -> Int -> m a) -> m () splitLinearlyM scheduler totalLength action = splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do - loopNextA_ 0 (< slackStart) (+ chunkLength) $ \ start next -> + loopNextA_ 0 (< slackStart) (+ chunkLength) $ \start next -> scheduleWork scheduler (action start next) when (slackStart < totalLength) $ scheduleWork scheduler (action slackStart totalLength) @@ -247,81 +257,87 @@ -- generator see `splitLinearlyWithM_`. -- -- @since 0.2.1 -splitLinearlyWith_ :: - MonadPrimBase s m => Scheduler s () -> Int -> (Int -> b) -> (Int -> b -> m ()) -> m () +splitLinearlyWith_ + :: MonadPrimBase s m => Scheduler s () -> Int -> (Int -> b) -> (Int -> b -> m ()) -> m () splitLinearlyWith_ scheduler totalLength index = splitLinearlyWithM_ scheduler totalLength (pure . index) {-# INLINE splitLinearlyWith_ #-} - -- | Iterator that can be used to split computation jobs -- -- @since 0.2.6 -splitLinearlyWithM_ :: - MonadPrimBase s m => Scheduler s () -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m () +splitLinearlyWithM_ + :: MonadPrimBase s m => Scheduler s () -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m () splitLinearlyWithM_ scheduler totalLength make write = splitLinearlyM_ scheduler totalLength go where - go start end = loopM_ start (< end) (+ 1) $ \ k -> make k >>= write k + go start end = loopM_ start (< end) (+ 1) $ \k -> make k >>= write k {-# INLINE go #-} {-# INLINE splitLinearlyWithM_ #-} - -- | Iterator that can be used to split computation jobs -- -- @since 0.3.0 -splitLinearlyWithStartAtM_ :: - MonadPrimBase s m => Scheduler s () -> Int -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m () +splitLinearlyWithStartAtM_ + :: MonadPrimBase s m => Scheduler s () -> Int -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m () splitLinearlyWithStartAtM_ scheduler startAt totalLength make write = splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do loopM_ startAt (< (slackStart + startAt)) (+ chunkLength) $ \ !start -> scheduleWork_ scheduler $ - loopM_ start (< (start + chunkLength)) (+ 1) $ \ !k -> make k >>= write k + loopM_ start (< (start + chunkLength)) (+ 1) $ + \ !k -> make k >>= write k when (slackStart < totalLength) $ scheduleWork_ scheduler $ - loopM_ (slackStart + startAt) (< (totalLength + startAt)) (+ 1) $ \ !k -> make k >>= write k + loopM_ (slackStart + startAt) (< (totalLength + startAt)) (+ 1) $ + \ !k -> make k >>= write k {-# INLINE splitLinearlyWithStartAtM_ #-} - - -- | Iterator that can be used to split computation jobs, while using a stateful scheduler. -- -- @since 0.3.4 -splitLinearlyWithStatefulM_ :: - MonadUnliftIO m +splitLinearlyWithStatefulM_ + :: MonadUnliftIO m => SchedulerWS ws () - -> Int -- ^ Total linear length - -> (Int -> ws -> m b) -- ^ Element producing action - -> (Int -> b -> m c) -- ^ Element storing action + -> Int + -- ^ Total linear length + -> (Int -> ws -> m b) + -- ^ Element producing action + -> (Int -> b -> m c) + -- ^ Element storing action -> m () splitLinearlyWithStatefulM_ schedulerWS totalLength make store = let nWorkers = numWorkers (unwrapSchedulerWS schedulerWS) in withRunInIO $ \run -> - splitLinearly nWorkers totalLength $ \chunkLength slackStart -> do - loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start -> + splitLinearly nWorkers totalLength $ \chunkLength slackStart -> do + loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start -> + scheduleWorkState_ schedulerWS $ \s -> + loopM_ start (< (start + chunkLength)) (+ 1) $ \ !k -> + run (make k s >>= store k) scheduleWorkState_ schedulerWS $ \s -> - loopM_ start (< (start + chunkLength)) (+ 1) $ \ !k -> + loopM_ slackStart (< totalLength) (+ 1) $ \ !k -> run (make k s >>= store k) - scheduleWorkState_ schedulerWS $ \s -> - loopM_ slackStart (< totalLength) (+ 1) $ \ !k -> - run (make k s >>= store k) {-# INLINE splitLinearlyWithStatefulM_ #-} - -- | This is a major helper function for fair splitting and parallelization of -- work with ability to use some arbitrary accumulator and splittable seed -- -- @since 1.0.2 -splitWorkWithFactorST :: - Int -- ^ Multiplying factor to be applied to number of workers for number - -- of jobs to schedule. Higher the factor, more jobs will be - -- scheduled. Only positive values are valid. +splitWorkWithFactorST + :: Int + -- ^ Multiplying factor to be applied to number of workers for number + -- of jobs to schedule. Higher the factor, more jobs will be + -- scheduled. Only positive values are valid. -> Scheduler s a - -> Int -- ^ Starting index - -> Int -- ^ Stepping value. Can be negative, but must not be zero. - -> Int -- ^ Total number of steps to be taken - -> b -- ^ Initial value for an accumulator - -> (b -> ST s (b, b)) -- ^ An action to split accumulator for multiple threads + -> Int + -- ^ Starting index + -> Int + -- ^ Stepping value. Can be negative, but must not be zero. + -> Int + -- ^ Total number of steps to be taken + -> b + -- ^ Initial value for an accumulator + -> (b -> ST s (b, b)) + -- ^ An action to split accumulator for multiple threads -> (Int -> Int -> Int -> Int -> b -> ST s a) -- ^ A job to be scheduled. Accepts: -- @@ -355,8 +371,8 @@ -- | Linear iterator that supports multiplying factor -- -- @since 1.0.2 -iterLinearST_ :: - Int +iterLinearST_ + :: Int -> Scheduler s () -> Int -> Int @@ -365,16 +381,16 @@ -> ST s () iterLinearST_ fact scheduler start step n action = do let totalLength = (n - start) `quot` step - splitWorkWithFactorST fact scheduler start step totalLength () (\_ -> pure ((), ())) - $ \ _ _ chunkStartAdj chunkStopAdj _ -> - loopA_ chunkStartAdj (< chunkStopAdj) (+ step) action + splitWorkWithFactorST fact scheduler start step totalLength () (\_ -> pure ((), ())) $ + \_ _ chunkStartAdj chunkStopAdj _ -> + loopA_ chunkStartAdj (< chunkStopAdj) (+ step) action {-# INLINE iterLinearST_ #-} -- | Linear iterator that supports multiplying factor and accumulator, but the results are discarded. -- -- @since 1.0.2 -iterLinearAccST_ :: - Int +iterLinearAccST_ + :: Int -> Scheduler s () -> Int -> Int @@ -385,33 +401,34 @@ -> ST s () iterLinearAccST_ fact scheduler start step n initAcc splitAcc action = do let totalLength = (n - start) `quot` step - void $ splitWorkWithFactorST fact scheduler start step totalLength initAcc splitAcc - $ \ _ _ chunkStartAdj chunkStopAdj accCur -> - void $ loopM chunkStartAdj (< chunkStopAdj) (+ step) accCur action + void $ + splitWorkWithFactorST fact scheduler start step totalLength initAcc splitAcc $ + \_ _ chunkStartAdj chunkStopAdj accCur -> + void $ loopM chunkStartAdj (< chunkStopAdj) (+ step) accCur action {-# INLINE iterLinearAccST_ #-} -- | Linear iterator that supports multiplying factor and accumulator. Results -- of actions are stored in the scheduler. -- -- @since 1.0.2 -iterLinearAccST :: - Int +iterLinearAccST + :: Int -> Scheduler s a -> Int - -> Int -- ^ Step. Must be non-zero -> Int + -- ^ Step. Must be non-zero + -> Int -> a -> (a -> ST s (a, a)) -> (Int -> a -> ST s a) -> ST s a iterLinearAccST fact scheduler start step n initAcc splitAcc action = do let totalLength = (n - start) `quot` step - splitWorkWithFactorST fact scheduler start step totalLength initAcc splitAcc - $ \ _ _ chunkStartAdj chunkStopAdj accCur -> - loopM chunkStartAdj (< chunkStopAdj) (+ step) accCur action + splitWorkWithFactorST fact scheduler start step totalLength initAcc splitAcc $ + \_ _ chunkStartAdj chunkStopAdj accCur -> + loopM chunkStartAdj (< chunkStopAdj) (+ step) accCur action {-# INLINE iterLinearAccST #-} - -- | Helper for figuring out the chunk length and slack start splitNumChunks :: Int -> Int -> Int -> (Int, Int) splitNumChunks fact nw totalLength = @@ -423,25 +440,23 @@ | otherwise = nw !chunkLength = totalLength `quot` numChunks !slackStart = chunkLength * numChunks - in (chunkLength, slackStart) - + in (chunkLength, slackStart) -- | Helper for adjusting stride of a chunk stepStartAdjust :: Int -> Int -> Int stepStartAdjust step ix = ix + ((step - (ix `mod` step)) `mod` step) {-# INLINE stepStartAdjust #-} - -- | Internal version of a `scheduleWork` that will be replaced by -- `scheduleWork_` by the compiler whenever action produces `()` scheduleMassivWork :: PrimBase m => Scheduler (PrimState m) a -> m a -> m () scheduleMassivWork = scheduleWork -{-# INLINE[0] scheduleMassivWork #-} +{-# INLINE [0] scheduleMassivWork #-} {-# RULES -"scheduleWork/scheduleWork_/ST" forall (scheduler :: Scheduler s ()) (action :: ST s ()) . scheduleMassivWork scheduler action = scheduleWork_ scheduler action -"scheduleWork/scheduleWork_/IO" forall (scheduler :: Scheduler RealWorld ()) (action :: IO ()) . scheduleMassivWork scheduler action = scheduleWork_ scheduler action - #-} +"scheduleWork/scheduleWork_/ST" forall (scheduler :: Scheduler s ()) (action :: ST s ()). scheduleMassivWork scheduler action = scheduleWork_ scheduler action +"scheduleWork/scheduleWork_/IO" forall (scheduler :: Scheduler RealWorld ()) (action :: IO ()). scheduleMassivWork scheduler action = scheduleWork_ scheduler action + #-} -- | Selects an optimal scheduler for the supplied strategy, but it works only in `IO` -- @@ -451,5 +466,5 @@ case comp of Par -> withGlobalScheduler_ globalScheduler f Seq -> f trivialScheduler_ - _ -> withScheduler_ comp f + _ -> withScheduler_ comp f {-# INLINE withMassivScheduler_ #-} diff --git a/src/Data/Massiv/Core/Operations.hs b/src/Data/Massiv/Core/Operations.hs --- a/src/Data/Massiv/Core/Operations.hs +++ b/src/Data/Massiv/Core/Operations.hs @@ -3,6 +3,7 @@ {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} + -- | -- Module : Data.Massiv.Core.Operations -- Copyright : (c) Alexey Kuleshevich 2019-2022 @@ -10,23 +11,20 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable -module Data.Massiv.Core.Operations - ( FoldNumeric(..) - , defaultPowerSumArray - , defaultUnsafeDotProduct - , defaultFoldArray - , Numeric(..) - , defaultUnsafeLiftArray - , defaultUnsafeLiftArray2 - , NumericFloat(..) - ) where +module Data.Massiv.Core.Operations ( + FoldNumeric (..), + defaultPowerSumArray, + defaultUnsafeDotProduct, + defaultFoldArray, + Numeric (..), + defaultUnsafeLiftArray, + defaultUnsafeLiftArray2, + NumericFloat (..), +) where import Data.Massiv.Core.Common - - class (Size r, Num e) => FoldNumeric r e where - {-# MINIMAL foldArray, powerSumArray, unsafeDotProduct #-} -- | Compute sum of all elements in the array @@ -58,9 +56,8 @@ -- @since 0.5.6 foldArray :: Index ix => (e -> e -> e) -> e -> Array r ix e -> e - -defaultUnsafeDotProduct :: - (Num e, Index ix, Source r e) => Array r ix e -> Array r ix e -> e +defaultUnsafeDotProduct + :: (Num e, Index ix, Source r e) => Array r ix e -> Array r ix e -> e defaultUnsafeDotProduct a1 a2 = go 0 0 where !len = totalElem (size a1) @@ -88,7 +85,6 @@ {-# INLINE defaultFoldArray #-} class FoldNumeric r e => Numeric r e where - {-# MINIMAL unsafeLiftArray, unsafeLiftArray2 #-} plusScalar :: Index ix => Array r ix e -> e -> Array r ix e @@ -126,6 +122,7 @@ -- TODO: -- - rename to powerScalar -- - add? powerPointwise :: Array r ix e -> Array r ix Int -> Array r ix e + -- | Raise each element of the array to the power powerPointwise :: Index ix => Array r ix e -> Int -> Array r ix e powerPointwise arr pow = unsafeLiftArray (^ pow) arr @@ -135,15 +132,13 @@ unsafeLiftArray2 :: Index ix => (e -> e -> e) -> Array r ix e -> Array r ix e -> Array r ix e - -defaultUnsafeLiftArray :: - (Load r ix e, Source r e) => (e -> e) -> Array r ix e -> Array r ix e +defaultUnsafeLiftArray + :: (Load r ix e, Source r e) => (e -> e) -> Array r ix e -> Array r ix e defaultUnsafeLiftArray f arr = makeArrayLinear (getComp arr) (size arr) (f . unsafeLinearIndex arr) {-# INLINE defaultUnsafeLiftArray #-} - -defaultUnsafeLiftArray2 :: - (Load r ix e, Source r e) +defaultUnsafeLiftArray2 + :: (Load r ix e, Source r e) => (e -> e -> e) -> Array r ix e -> Array r ix e @@ -153,9 +148,7 @@ f (unsafeLinearIndex a1 i) (unsafeLinearIndex a2 i) {-# INLINE defaultUnsafeLiftArray2 #-} - class (Numeric r e, Floating e) => NumericFloat r e where - divideScalar :: Index ix => Array r ix e -> e -> Array r ix e divideScalar arr e = unsafeLiftArray (/ e) arr {-# INLINE divideScalar #-} @@ -176,14 +169,13 @@ sqrtPointwise = unsafeLiftArray sqrt {-# INLINE sqrtPointwise #-} - -- floorPointwise :: (Index ix, Integral a) => Array r ix e -> Array r ix a - -- floorPointwise = unsafeLiftArray floor - -- {-# INLINE floorPointwise #-} - - -- ceilingPointwise :: (Index ix, Integral a) => Array r ix e -> Array r ix a - -- ceilingPointwise = unsafeLiftArray ceiling - -- {-# INLINE ceilingPointwise #-} +-- floorPointwise :: (Index ix, Integral a) => Array r ix e -> Array r ix a +-- floorPointwise = unsafeLiftArray floor +-- {-# INLINE floorPointwise #-} +-- ceilingPointwise :: (Index ix, Integral a) => Array r ix e -> Array r ix a +-- ceilingPointwise = unsafeLiftArray ceiling +-- {-# INLINE ceilingPointwise #-} -- class Equality r e where @@ -191,7 +183,6 @@ -- unsafeEqPointwise :: Index ix => Array r ix e -> Array r ix e -> Array r ix Bool - -- class Relation r e where -- unsafePointwiseLT :: Array r ix e -> Array r ix e -> Array r ix Bool @@ -206,5 +197,3 @@ -- unsafeMinimum :: Array r ix e -> e -- unsafeMaximum :: Array r ix e -> e - - diff --git a/src/Data/Massiv/Vector.hs b/src/Data/Massiv/Vector.hs --- a/src/Data/Massiv/Vector.hs +++ b/src/Data/Massiv/Vector.hs @@ -2,2752 +2,2839 @@ {-# LANGUAGE ExplicitForAll #-} {-# LANGUAGE FlexibleContexts #-} {-# OPTIONS_GHC -fno-warn-duplicate-exports #-} --- | --- Module : Data.Massiv.Vector --- Copyright : (c) Alexey Kuleshevich 2020-2022 --- License : BSD3 --- Maintainer : Alexey Kuleshevich --- Stability : experimental --- Portability : non-portable --- -module Data.Massiv.Vector - ( Vector - , MVector - -- * Accessors - -- *** Size - , slength - , maxLinearSize - , size - , isNull - , isNotNull - -- *** Indexing - , (!?) - , (!) - , index - , index' - , head' - , shead' - , last' - -- *** Monadic Indexing - , indexM - , headM - , sheadM - , lastM - , unconsM - , unsnocM - -- ** Slicing - , slice - , slice' - , sliceM - , sslice - , sliceAt - , sliceAt' - , sliceAtM - -- *** Init - , init - , init' - , initM - -- *** Tail - , tail - , tail' - , tailM - -- *** Take - , take - , take' - , takeM - , takeWhile - , stake - -- *** Drop - , drop - , dropWhile - , drop' - , dropM - , sdrop - -- * Construction - -- ** Initialization - , empty - , sempty - , singleton - , ssingleton - , cons - , snoc - , A.replicate - , sreplicate - , generate - , sgenerate - -- , iterateN - -- , iiterateN - , siterate - , siterateN - -- ** Monadic initialization - , sreplicateM - , sgenerateM - , siterateNM - -- , create - -- , createT - -- ** Unfolding - , sunfoldr - , sunfoldrM - , sunfoldrN - , sunfoldrNM - , sunfoldrExactN - , sunfoldrExactNM - -- , constructN - -- , constructrN - -- ** Enumeration - , (...) - , (..:) - , enumFromN - , senumFromN - , enumFromStepN - , senumFromStepN - -- ** Concatenation - -- , consS -- cons - -- , snocS -- snoc - , sappend -- (++) - , sconcat -- concat - -- -- ** Restricitng memory usage - -- , force - -- -- * Modifying - -- -- ** Bulk updates - -- , (//) - -- , update_ - -- -- ** Accumulations - -- , accum - -- , accumulate_ - -- -- ** Permutations - -- , reverse - -- , backpermute - -- -- ** Manifest updates - -- , modify - -- -- * Elementwise - -- -- ** Mapping - , smap - , simap - -- , sconcatMap - -- ** Monadic mapping - , straverse - , sitraverse - , smapM - , smapM_ - , simapM - , simapM_ - , sforM - , sforM_ - , siforM - , siforM_ - -- ** Zipping - , szip - , szip3 - , szip4 - , szip5 - , szip6 - , szipWith - , szipWith3 - , szipWith4 - , szipWith5 - , szipWith6 - , sizipWith - , sizipWith3 - , sizipWith4 - , sizipWith5 - , sizipWith6 - -- ** Monadic zipping - , szipWithM - , szipWith3M - , szipWith4M - , szipWith5M - , szipWith6M - , sizipWithM - , sizipWith3M - , sizipWith4M - , sizipWith5M - , sizipWith6M - - , szipWithM_ - , szipWith3M_ - , szipWith4M_ - , szipWith5M_ - , szipWith6M_ - , sizipWithM_ - , sizipWith3M_ - , sizipWith4M_ - , sizipWith5M_ - , sizipWith6M_ - -- * Predicates - -- ** Filtering - , sfilter - , sifilter - , sfilterM - , sifilterM - -- , uniq -- sunique? - , smapMaybe - , smapMaybeM - , scatMaybes - , simapMaybe - , simapMaybeM - -- , stakeWhile - -- , sdropWhile - -- -- ** Partitioning - -- , partition - -- , unstablePartition - -- , partitionWith - -- , span - -- , break - -- -- ** Searching - -- , elem - -- , notElem - -- , find - , findIndex - -- , findIndices - -- , elemIndex - -- , elemIndices - -- * Folding - , sfoldl - , sfoldlM - , sfoldlM_ - , sifoldl - , sifoldlM - , sifoldlM_ - , sfoldl1' - , sfoldl1M - , sfoldl1M_ - -- ** Specialized folds - , sor - , sand - , sall - , sany - , ssum - , sproduct - , smaximum' - , smaximumM - -- , maximumBy - , sminimum' - , sminimumM - -- , minimumBy - -- , minIndex - -- , minIndexBy - -- , maxIndex - -- , maxIndexBy - -- -- ** Prefix sums - -- , prescanl - -- , prescanl' - -- , postscanl - -- , postscanl' - -- , scanl - -- , scanl' - -- , scanl1 - -- , scanl1' - -- , prescanr - -- , prescanr' - -- , postscanr - -- , postscanr' - -- , scanr - -- , scanr' - -- , scanr1 - -- , scanr1' - -- * Conversions - -- ** Lists - , stoList - , fromList - , sfromList - , sfromListN - -- * Computation - , compute - , computeS - , computeIO - , computePrimM - , computeAs - , computeProxy - , computeSource - , computeWithStride - , computeWithStrideAs - , clone - , convert - , convertAs - , convertProxy - -- ** Re-exports - , module Data.Massiv.Core - , module Data.Massiv.Array.Delayed - , module Data.Massiv.Array.Manifest - , module Data.Massiv.Array.Mutable - ) where - -import Control.Monad hiding (filterM, replicateM) -import Data.Coerce -import Data.Massiv.Array.Delayed -import Data.Massiv.Array.Delayed.Pull -import Data.Massiv.Array.Delayed.Push -import Data.Massiv.Array.Delayed.Stream -import Data.Massiv.Array.Manifest -import Data.Massiv.Array.Manifest.Internal -import Data.Massiv.Array.Manifest.List (fromList) -import Data.Massiv.Array.Mutable -import Data.Massiv.Array.Ops.Construct -import qualified Data.Massiv.Array.Ops.Construct as A (replicate) -import Data.Massiv.Core -import Data.Massiv.Core.Common -import qualified Data.Massiv.Vector.Stream as S -import Data.Massiv.Vector.Unsafe -import Data.Maybe -import Prelude hiding (drop, dropWhile, init, length, null, replicate, splitAt, - tail, take, takeWhile) - --- ========= -- --- Accessors -- --- ========= -- - - ------------------------- --- Length information -- ------------------------- - --- | /O(1)/ - Get the length of a `Stream` array, but only if it is known exactly in --- constant time without looking at any of the elements in the array. --- --- /Related/: `maxLinearSize`, `size`, `elemsCount` and `totalElem` --- --- ==== __Examples__ --- --- >>> slength $ sfromList [] --- Nothing --- >>> slength $ sreplicate 5 () --- Just (Sz1 5) --- >>> slength $ makeArrayLinearR D Seq (Sz1 5) id --- Just (Sz1 5) --- >>> slength $ sunfoldr (\x -> Just (x, x)) (0 :: Int) --- Nothing --- >>> slength $ sunfoldrN 10 (\x -> Just (x, x)) (0 :: Int) --- Nothing --- >>> slength $ sunfoldrExactN 10 (\x -> (x, x)) (0 :: Int) --- Just (Sz1 10) --- --- /__Similar__/: --- --- [@Data.Foldable.`Data.Foldable.length`@] For some data structures, like a list for --- example, it is an /O(n)/ operation, because there is a need to evaluate the full spine --- and possibly even the elements in order to get the full length. With `Stream` vectors --- that is not always the case. --- --- [@Data.Vector.Generic.`Data.Vector.Generic.length`@] In the vector package this --- function will always break fusion, unless it is the only operation that is applied to --- the vector. --- --- @since 0.5.0 -slength :: - forall r ix e. Stream r ix e - => Array r ix e - -> Maybe Sz1 -slength v = - case stepsSize (toStream v) of - LengthExact sz -> Just sz - _ -> Nothing -{-# INLINE slength #-} - - --------------- --- Indexing -- --------------- - - --- | /O(1)/ - Get the first element of a `Source` vector. Throws an error on empty. --- --- /Related/: 'shead'', `headM`, `sheadM`, `unconsM`. --- --- ==== __Examples__ --- --- >>> head' (Ix1 10 ..: 10000000000000) --- 10 --- --- /__Similar__/: --- --- [@Data.List.`Data.List.head`@] Also constant time and partial. Fusion is broken if --- there other consumers of the list. --- --- [@Data.Vector.Generic.`Data.Vector.Generic.head`@] Also constant time and partial. Will --- cause materialization of the full vector if any other function is applied to the vector. --- --- @since 0.5.0 -head' :: - forall r e. (HasCallStack, Source r e) - => Vector r e - -> e -head' = throwEither . headM -{-# INLINE head' #-} - - --- | /O(1)/ - Get the first element of a `Source` vector. --- --- /Related/: 'head'', 'shead'', `sheadM`, `unconsM`. --- --- /__Throws Exceptions__/: `SizeEmptyException` when array is empty --- --- ==== __Examples__ --- --- >>> headM (Ix1 10 ..: 10000000000000) --- 10 --- >>> headM (Ix1 10 ..: 10000000000000) :: Maybe Int --- Just 10 --- >>> headM (empty :: Array D Ix1 Int) :: Maybe Int --- Nothing --- >>> either show (const "") $ headM (Ix1 10 ..: 10) --- "SizeEmptyException: (Sz1 0) corresponds to an empty array" --- --- /__Similar__/: --- --- [@Data.Maybe.`Data.Maybe.listToMaybe`@] It also a safe way to get the head of the list, --- except it is restricted to `Maybe` --- --- @since 0.5.0 -headM :: - forall r e m. (Source r e, MonadThrow m) - => Vector r e - -> m e -headM v - | elemsCount v == 0 = throwM $ SizeEmptyException (size v) - | otherwise = pure $ unsafeLinearIndex v 0 -{-# INLINE headM #-} - - --- | /O(1)/ - Get the first element of a `Stream` vector. Throws an error on empty. --- --- /Related/: 'head'', `headM`, `sheadM`, `unconsM`. --- --- ==== __Examples__ --- --- >>> shead' $ sunfoldr (\x -> Just (x, x)) (0 :: Int) --- 0 --- >>> shead' (Ix1 3 ... 5) --- 3 --- --- @since 0.5.0 -shead' :: - forall r e. (HasCallStack, Stream r Ix1 e) - => Vector r e - -> e -shead' = throwEither . sheadM -{-# INLINE shead' #-} - --- | /O(1)/ - Get the first element of a `Stream` vector. --- --- /Related/: 'head'', 'shead'', `headM`, `unconsM`. --- --- /__Throws Exceptions__/: `SizeEmptyException` --- --- ==== __Examples__ --- --- >>> maybe 101 id $ sheadM (empty :: Vector D Int) --- 101 --- >>> maybe 101 id $ sheadM (singleton 202 :: Vector D Int) --- 202 --- >>> sheadM $ sunfoldr (\x -> Just (x, x)) (0 :: Int) --- 0 --- >>> x <- sheadM $ sunfoldr (\_ -> Nothing) (0 :: Int) --- *** Exception: SizeEmptyException: (Sz1 0) corresponds to an empty array --- --- @since 0.5.0 -sheadM :: - forall r e m. (Stream r Ix1 e, MonadThrow m) - => Vector r e - -> m e -sheadM v = - case S.unId (S.headMaybe (toStream v)) of - Nothing -> throwM $ SizeEmptyException (zeroSz :: Sz1) - Just e -> pure e -{-# INLINE sheadM #-} - - --- | /O(1)/ - Take one element off of the `Source` vector from the left side, as well as --- the remaining part of the vector in delayed `D` representation. --- --- /Related/: 'head'', 'shead'', `headM`, `sheadM`, `cons` --- --- /__Throws Exceptions__/: `SizeEmptyException` --- --- ==== __Examples__ --- --- >>> unconsM (fromList Seq [1,2,3] :: Array P Ix1 Int) --- (1,Array P Seq (Sz1 2) --- [ 2, 3 ]) --- --- /__Similar__/: --- --- [@Data.List.`Data.List.uncons`@] Same concept, except it is restricted to `Maybe` instead of --- the more general `MonadThrow` --- --- @since 0.3.0 -unconsM :: - forall r e m. (MonadThrow m, Source r e) - => Vector r e - -> m (e, Vector r e) -unconsM arr - | 0 == totalElem sz = throwM $ SizeEmptyException sz - | otherwise = pure (unsafeLinearIndex arr 0, unsafeLinearSlice 1 (SafeSz (unSz sz - 1)) arr) - where - !sz = size arr -{-# INLINE unconsM #-} - --- | /O(1)/ - Take one element off of the vector from the right side, as well as the --- remaining part of the vector. --- --- /Related/: 'last'', `lastM`, `snoc` --- --- /__Throws Exceptions__/: `SizeEmptyException` --- --- ==== __Examples__ --- --- >>> unsnocM (fromList Seq [1,2,3] :: Array P Ix1 Int) --- (Array P Seq (Sz1 2) --- [ 1, 2 ],3) --- --- @since 0.3.0 -unsnocM :: - forall r e m. (MonadThrow m, Source r e) - => Vector r e - -> m (Vector r e, e) -unsnocM arr - | 0 == totalElem sz = throwM $ SizeEmptyException sz - | otherwise = pure (unsafeLinearSlice 0 (SafeSz k) arr, unsafeLinearIndex arr k) - where - !sz = size arr - !k = unSz sz - 1 -{-# INLINE unsnocM #-} - - --- | /O(1)/ - Get the last element of a `Source` vector. Throws an error on empty. --- --- /Related/: `lastM`, `unsnocM` --- --- ==== __Examples__ --- --- >>> last' (Ix1 10 ... 10000000000000) --- 10000000000000 --- --- /__Similar__/: --- --- [@Data.List.`Data.List.last`@] Also partial, but it has /O(n)/ complexity. Fusion is --- broken if there other consumers of the list. --- --- [@Data.Vector.Generic.`Data.Vector.Generic.last`@] Also constant time and partial. Will --- cause materialization of the full vector if any other function is applied to the vector. --- --- @since 0.5.0 -last' :: forall r e. (HasCallStack, Source r e) => Vector r e -> e -last' = throwEither . lastM -{-# INLINE last' #-} - - --- | /O(1)/ - Get the last element of a `Source` vector. --- --- /Related/: 'last'', `unsnocM` --- --- /__Throws Exceptions__/: `SizeEmptyException` --- --- ==== __Examples__ --- --- >>> lastM (Ix1 10 ... 10000000000000) --- 10000000000000 --- >>> lastM (Ix1 10 ... 10000000000000) :: Maybe Int --- Just 10000000000000 --- >>> either show (const "") $ lastM (fromList Seq [] :: Array P Ix1 Int) --- "SizeEmptyException: (Sz1 0) corresponds to an empty array" --- --- @since 0.5.0 -lastM :: forall r e m. (Source r e, MonadThrow m) => Vector r e -> m e -lastM v - | k == 0 = throwM $ SizeEmptyException (size v) - | otherwise = pure $ unsafeLinearIndex v (k - 1) - where k = unSz (size v) -{-# INLINE lastM #-} - - --- | /O(1)/ - Take a slice of a `Source` vector. Never fails, instead adjusts the indices. --- --- ==== __Examples__ --- --- >>> slice 10 5 (Ix1 0 ... 10000000000000) --- Array D Seq (Sz1 5) --- [ 10, 11, 12, 13, 14 ] --- >>> slice (-10) 5 (Ix1 0 ... 10000000000000) --- Array D Seq (Sz1 5) --- [ 0, 1, 2, 3, 4 ] --- >>> slice 9999999999998 50 (Ix1 0 ... 10000000000000) --- Array D Seq (Sz1 3) --- [ 9999999999998, 9999999999999, 10000000000000 ] --- --- @since 0.5.0 -slice :: forall r e. Source r e => Ix1 -> Sz1 -> Vector r e -> Vector r e -slice !i (Sz k) v = unsafeLinearSlice i' newSz v - where - !i' = min n (max 0 i) - !newSz = SafeSz (min (n - i') k) - Sz n = size v -{-# INLINE slice #-} - --- | /O(1)/ - Take a slice of a `Source` vector. Throws an error on incorrect indices. --- --- ==== __Examples__ --- --- >>> slice' 10 5 (Ix1 0 ... 100) --- Array D Seq (Sz1 5) --- [ 10, 11, 12, 13, 14 ] --- >>> slice' 9999999999998 3 (Ix1 0 ... 10000000000000) --- Array D Seq (Sz1 3) --- [ 9999999999998, 9999999999999, 10000000000000 ] --- --- @since 0.5.0 -slice' :: forall r e. (HasCallStack, Source r e) => Ix1 -> Sz1 -> Vector r e -> Vector r e -slice' i k = throwEither . sliceM i k -{-# INLINE slice' #-} - - --- | /O(1)/ - Take a slice of a `Source` vector. Throws an error on incorrect indices. --- --- /__Throws Exceptions__/: `SizeSubregionException` --- --- ==== __Examples__ --- --- >>> sliceM 10 5 (Ix1 0 ... 100) --- Array D Seq (Sz1 5) --- [ 10, 11, 12, 13, 14 ] --- >>> sliceM (-10) 5 (Ix1 0 ... 100) --- *** Exception: SizeSubregionException: (Sz1 101) is to small for -10 (Sz1 5) --- >>> sliceM 98 50 (Ix1 0 ... 100) --- *** Exception: SizeSubregionException: (Sz1 101) is to small for 98 (Sz1 50) --- >>> sliceM 9999999999998 3 (Ix1 0 ... 10000000000000) --- Array D Seq (Sz1 3) --- [ 9999999999998, 9999999999999, 10000000000000 ] --- --- @since 0.5.0 -sliceM :: - forall r e m. (Source r e, MonadThrow m) - => Ix1 - -- ^ Starting index - -> Sz1 - -- ^ Number of elements to take from the Source vector - -> Vector r e - -- ^ Source vector to take a slice from - -> m (Vector r e) -sliceM i newSz@(Sz k) v - | i >= 0 && k <= n - i = pure $ unsafeLinearSlice i newSz v - | otherwise = throwM $ SizeSubregionException sz i newSz - where - sz@(Sz n) = size v -{-# INLINE sliceM #-} - - --- | Take a slice of a `Stream` vector. Never fails, instead adjusts the indices. --- --- ==== __Examples__ --- --- >>> sslice 10 5 (Ix1 0 ... 10000000000000) --- Array DS Seq (Sz1 5) --- [ 10, 11, 12, 13, 14 ] --- >>> sslice 10 5 (sfromList [0 :: Int .. ]) --- Array DS Seq (Sz1 5) --- [ 10, 11, 12, 13, 14 ] --- >>> sslice (-10) 5 (Ix1 0 ... 10000000000000) --- Array DS Seq (Sz1 5) --- [ 0, 1, 2, 3, 4 ] --- --- Unlike `slice` it has to iterate through each element until the staring index is reached, --- therefore something like @sslice 9999999999998 50 (Ix1 0 ... 10000000000000)@ will not --- be feasable. --- --- >>> import System.Timeout (timeout) --- >>> let smallArr = sslice 9999999999998 50 (Ix1 0 ... 10000000000000) --- >>> timeout 500000 (computeIO smallArr :: IO (Array P Ix1 Int)) --- Nothing --- --- @since 0.5.0 -sslice :: - forall r e. Stream r Ix1 e - => Ix1 - -- ^ Starting index - -> Sz1 - -- ^ Number of elements to take from the stream vector - -> Vector r e - -- ^ Stream vector to take a slice from - -> Vector DS e -sslice !i !k = fromSteps . S.slice i k . S.toStream -{-# INLINE sslice #-} - - --- | /O(1)/ - Get a vector without the last element. Never fails. --- --- ==== __Examples__ --- --- >>> import Data.Massiv.Array as A --- >>> A.init (0 ..: 10) --- Array D Seq (Sz1 9) --- [ 0, 1, 2, 3, 4, 5, 6, 7, 8 ] --- >>> A.init (empty :: Array D Ix1 Int) --- Array D Seq (Sz1 0) --- [ ] --- --- @since 0.5.0 -init :: forall r e. Source r e => Vector r e -> Vector r e -init v = unsafeLinearSlice 0 (Sz (coerce (size v) - 1)) v -{-# INLINE init #-} - --- | /O(1)/ - Get a vector without the last element. Throws an error on empty --- --- ==== __Examples__ --- --- >>> init' (0 ..: 10) --- Array D Seq (Sz1 9) --- [ 0, 1, 2, 3, 4, 5, 6, 7, 8 ] --- --- @since 0.5.0 -init' :: forall r e. (HasCallStack, Source r e) => Vector r e -> Vector r e -init' = throwEither . initM -{-# INLINE init' #-} - --- | /O(1)/ - Get a vector without the last element. Throws an error on empty --- --- ==== __Examples__ --- --- >>> import Data.Massiv.Array as A --- >>> initM (0 ..: 10) --- Array D Seq (Sz1 9) --- [ 0, 1, 2, 3, 4, 5, 6, 7, 8 ] --- >>> maybe 0 A.sum $ initM (0 ..: 10) --- 36 --- >>> maybe 0 A.sum $ initM (empty :: Array D Ix1 Int) --- 0 --- --- @since 0.5.0 -initM :: forall r e m. (Source r e, MonadThrow m) => Vector r e -> m (Vector r e) -initM v = do - when (elemsCount v == 0) $ throwM $ SizeEmptyException $ size v - pure $ unsafeInit v -{-# INLINE initM #-} - - - --- | /O(1)/ - Get a vector without the first element. Never fails --- --- ==== __Examples__ --- --- >>> import Data.Massiv.Array as A --- >>> A.tail (0 ..: 10) --- Array D Seq (Sz1 9) --- [ 1, 2, 3, 4, 5, 6, 7, 8, 9 ] --- >>> A.tail (empty :: Array D Ix1 Int) --- Array D Seq (Sz1 0) --- [ ] --- --- @since 0.5.0 -tail :: forall r e. Source r e => Vector r e -> Vector r e -tail = drop oneSz -{-# INLINE tail #-} - - --- | /O(1)/ - Get a vector without the first element. Throws an error on empty --- --- ==== __Examples__ --- --- λ> tail' (0 ..: 10) --- Array D Seq (Sz1 9) --- [ 1, 2, 3, 4, 5, 6, 7, 8, 9 ] --- λ> tail' (empty :: Array D Ix1 Int) --- Array D *** Exception: SizeEmptyException: (Sz1 0) corresponds to an empty array --- --- @since 0.5.0 -tail' :: forall r e. (HasCallStack, Source r e) => Vector r e -> Vector r e -tail' = throwEither . tailM -{-# INLINE tail' #-} - - --- | /O(1)/ - Get the vector without the first element. Throws an error on empty --- --- ==== __Examples__ --- --- >>> import Data.Massiv.Array as A --- >>> tailM (0 ..: 10) --- Array D Seq (Sz1 9) --- [ 1, 2, 3, 4, 5, 6, 7, 8, 9 ] --- >>> maybe 0 A.sum $ tailM (0 ..: 10) --- 45 --- >>> maybe 0 A.sum $ tailM (empty :: Array D Ix1 Int) --- 0 --- --- @since 0.5.0 -tailM :: forall r e m. (Source r e, MonadThrow m) => Vector r e -> m (Vector r e) -tailM v = do - when (elemsCount v == 0) $ throwM $ SizeEmptyException $ size v - pure $ unsafeTail v -{-# INLINE tailM #-} - --- | /O(1)/ - Take first @n@ elements from a vector. This function never fails and has --- similar semantics as the `Data.List.take` for lists. --- --- ==== __Examples__ --- --- >>> import Data.Massiv.Array as A --- >>> A.take 5 (0 ..: 10) --- Array D Seq (Sz1 5) --- [ 0, 1, 2, 3, 4 ] --- >>> A.take 0 (0 ..: 10) --- Array D Seq (Sz1 0) --- [ ] --- >>> A.take 100 (0 ..: 10) --- Array D Seq (Sz1 10) --- [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9 ] --- --- @since 0.5.0 -take :: Source r e => Sz1 -> Vector r e -> Vector r e -take k = fst . sliceAt k -{-# INLINE take #-} - - --- | Slice a manifest vector in such a way that it will contain all initial elements that --- satisfy the supplied predicate. --- --- @since 0.5.5 -takeWhile :: Manifest r e => (e -> Bool) -> Vector r e -> Vector r e -takeWhile f v = take (go 0) v - where - !k = elemsCount v - go !i - | i < k && f (unsafeLinearIndex v i) = go (i + 1) - | otherwise = SafeSz i -{-# INLINE takeWhile #-} - - --- | /O(1)/ - Get the vector with the first @n@ elements. Throws an error size is less --- than @n@. --- --- ==== __Examples__ --- --- >>> take' 0 (0 ..: 0) --- Array D Seq (Sz1 0) --- [ ] --- >>> take' 5 (0 ..: 10) --- Array D Seq (Sz1 5) --- [ 0, 1, 2, 3, 4 ] --- --- @since 0.5.0 -take' :: forall r e. (HasCallStack, Source r e) => Sz1 -> Vector r e -> Vector r e -take' k = throwEither . takeM k -{-# INLINE take' #-} - --- | /O(1)/ - Get the vector with the first @n@ elements. Throws an error size is less than @n@ --- --- ==== __Examples__ --- --- >>> import Data.Massiv.Array as A --- >>> takeM 5 (0 ..: 10) --- Array D Seq (Sz1 5) --- [ 0, 1, 2, 3, 4 ] --- >>> maybe 0 A.sum $ takeM 5 (0 ..: 10) --- 10 --- >>> maybe (-1) A.sum $ takeM 15 (0 ..: 10) --- -1 --- >>> takeM 15 (0 ..: 10) --- *** Exception: SizeSubregionException: (Sz1 10) is to small for 0 (Sz1 15) --- --- @since 0.5.0 -takeM :: forall r e m. (Source r e, MonadThrow m) => Sz1 -> Vector r e -> m (Vector r e) -takeM k v = do - let sz = size v - when (k > sz) $ throwM $ SizeSubregionException sz 0 k - pure $ unsafeTake k v -{-# INLINE takeM #-} - --- | /O(1)/ - Create a `Stream` vector with the first @n@ elements. Never fails --- --- ==== __Examples__ --- --- @since 0.5.0 -stake :: forall r e. Stream r Ix1 e => Sz1 -> Vector r e -> Vector DS e -stake n = fromSteps . S.take n . S.toStream -{-# INLINE stake #-} - --- | /O(1)/ - Drop @n@ elements from a vector. This function never fails and has --- similar semantics as the `Data.List.drop` for lists. --- --- ==== __Examples__ --- --- >>> import Data.Massiv.Array as A --- >>> v = makeVectorR D Seq 10 id --- >>> v --- Array D Seq (Sz1 10) --- [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9 ] --- >>> A.drop 5 v --- Array D Seq (Sz1 5) --- [ 5, 6, 7, 8, 9 ] --- >>> A.drop 25 v --- Array D Seq (Sz1 0) --- [ ] --- --- @since 0.5.0 -drop :: forall r e. Source r e => Sz1 -> Vector r e -> Vector r e -drop k = snd . sliceAt k -{-# INLINE drop #-} - - --- | Slice a manifest vector in such a way that it will not contain all initial elements --- that satisfy the supplied predicate. --- --- @since 0.5.5 -dropWhile :: forall r e. Manifest r e => (e -> Bool) -> Vector r e -> Vector r e -dropWhile f v = drop (go 0) v - where - !k = elemsCount v - go !i - | i < k && f (unsafeLinearIndex v i) = go (i + 1) - | otherwise = SafeSz i -{-# INLINE dropWhile #-} - - --- | Keep all but the first @n@ elements from the delayed stream vector. --- --- ==== __Examples__ --- --- @since 0.5.0 -sdrop :: forall r e. Stream r Ix1 e => Sz1 -> Vector r e -> Vector DS e -sdrop n = fromSteps . S.drop n . S.toStream -{-# INLINE sdrop #-} - --- | /O(1)/ - Drop @n@ elements from a vector. Unlike `drop`, this function will --- produce an error when supplied number of elements to drop is larger than size --- of the supplied vector --- --- ==== __Examples__ --- --- @since 0.5.0 -drop' :: forall r e. (HasCallStack, Source r e) => Sz1 -> Vector r e -> Vector r e -drop' k = throwEither . dropM k -{-# INLINE drop' #-} - --- | --- --- ==== __Examples__ --- --- @since 0.5.0 -dropM :: forall r e m. (Source r e, MonadThrow m) => Sz1 -> Vector r e -> m (Vector r e) -dropM k@(Sz d) v = do - let sz@(Sz n) = size v - when (k > sz) $ throwM $ SizeSubregionException sz d (SafeSz (n - d)) - pure $ unsafeLinearSlice d (SafeSz (n - d)) v -{-# INLINE dropM #-} - - --- | Same as 'sliceAt'', except it never fails. --- --- ==== __Examples__ --- --- @since 0.5.0 -sliceAt :: forall r e. Source r e => Sz1 -> Vector r e -> (Vector r e, Vector r e) -sliceAt (Sz k) v = (unsafeTake d v, unsafeDrop d v) - where - !n = coerce (size v) - !d = SafeSz (min k n) -{-# INLINE sliceAt #-} - --- | Same as 'Data.Massiv.Array.splitAt'', except for a flat vector. --- --- ==== __Examples__ --- --- @since 0.5.0 -sliceAt' :: (HasCallStack, Source r e) => Sz1 -> Vector r e -> (Vector r e, Vector r e) -sliceAt' k = throwEither . sliceAtM k -{-# INLINE sliceAt' #-} - --- | Same as `Data.Massiv.Array.splitAtM`, except for a flat vector. --- --- ==== __Examples__ --- --- @since 0.5.0 -sliceAtM :: forall r e m. (Source r e, MonadThrow m) => Sz1 -> Vector r e -> m (Vector r e, Vector r e) -sliceAtM k v = do - l <- takeM k v - pure (l, unsafeDrop k v) -{-# INLINE sliceAtM #-} - - --- | Create an empty delayed stream vector --- --- ==== __Examples__ --- --- @since 0.5.0 -sempty :: Vector DS e -sempty = DSArray S.empty -{-# INLINE sempty #-} - --- | Create a delayed stream vector with a single element --- --- ==== __Examples__ --- --- @since 0.5.0 -ssingleton :: e -> Vector DS e -ssingleton = DSArray . S.singleton -{-# INLINE ssingleton #-} - --- | /O(1)/ - Add an element to the vector from the left side --- --- @since 0.3.0 -cons :: forall r e. (Size r, Load r Ix1 e) => e -> Vector r e -> Vector DL e -cons e v = - let dv = toLoadArray v - load scheduler startAt uWrite uSet = - uWrite startAt e >> dlLoad dv scheduler (startAt + 1) uWrite uSet - {-# INLINE load #-} - in dv {dlSize = SafeSz (1 + unSz (dlSize dv)), dlLoad = load} -{-# INLINE cons #-} - --- | /O(1)/ - Add an element to the vector from the right side --- --- @since 0.3.0 -snoc :: forall r e. (Size r, Load r Ix1 e) => Vector r e -> e -> Vector DL e -snoc v e = - let dv = toLoadArray v - !k = unSz (size dv) - load scheduler startAt uWrite uSet = - dlLoad dv scheduler startAt uWrite uSet >> uWrite (k + startAt) e - {-# INLINE load #-} - in dv {dlSize = SafeSz (1 + k), dlLoad = load} -{-# INLINE snoc #-} - - - --- | Replicate the same element @n@ times --- --- ==== __Examples__ --- --- @since 0.5.0 -sreplicate :: Sz1 -> e -> Vector DS e -sreplicate n = DSArray . S.replicate n -{-# INLINE sreplicate #-} - --- | Create a delayed vector of length @n@ with a function that maps an index to an --- element. Same as `makeLinearArray` --- --- ==== __Examples__ --- --- @since 0.5.0 -generate :: Comp -> Sz1 -> (Ix1 -> e) -> Vector D e -generate = makeArrayLinear -{-# INLINE generate #-} - --- | Create a delayed stream vector of length @n@ with a function that maps an index to an --- element. Same as `makeLinearArray` --- --- ==== __Examples__ --- --- @since 0.5.0 -sgenerate :: Sz1 -> (Ix1 -> e) -> Vector DS e -sgenerate n = DSArray . S.generate n -{-# INLINE sgenerate #-} - - --- | Create a delayed stream vector of infinite length by repeatedly applying a function to the --- initial value. --- --- ==== __Examples__ --- --- >>> stake 10 $ siterate succ 'a' --- Array DS Seq (Sz1 10) --- [ 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j' ] --- --- @since 0.5.2 -siterate :: (e -> e) -> e -> Vector DS e -siterate f = fromSteps . S.unfoldr (\a -> Just (a, f a)) -{-# INLINE siterate #-} - --- | Create a delayed stream vector of length @n@ by repeatedly applying a function to the --- initial value. --- --- ==== __Examples__ --- --- >>> siterateN 10 succ 'a' --- Array DS Seq (Sz1 10) --- [ 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j' ] --- --- @since 0.5.0 -siterateN :: Sz1 -> (e -> e) -> e -> Vector DS e -siterateN n f a = fromSteps $ S.iterateN n f a -{-# INLINE siterateN #-} - - --- | Create a vector by using the same monadic action @n@ times --- --- ==== __Examples__ --- --- @since 0.5.0 -sreplicateM :: forall e m. Monad m => Sz1 -> m e -> m (Vector DS e) -sreplicateM n f = fromStepsM $ S.replicateM n f -{-# INLINE sreplicateM #-} - - --- | Create a delayed stream vector of length @n@ with a monadic action that from an index --- generates an element. --- --- ==== __Examples__ --- --- @since 0.5.0 -sgenerateM :: forall e m. Monad m => Sz1 -> (Ix1 -> m e) -> m (Vector DS e) -sgenerateM n f = fromStepsM $ S.generateM n f -{-# INLINE sgenerateM #-} - - --- | Create a delayed stream vector of length @n@ by repeatedly apply a monadic action to --- the initial value. --- --- ==== __Examples__ --- --- @since 0.5.0 -siterateNM :: forall e m. Monad m => Sz1 -> (e -> m e) -> e -> m (Vector DS e) -siterateNM n f a = fromStepsM $ S.iterateNM n f a -{-# INLINE siterateNM #-} - - - - --- | Right unfolding function. Useful when it is unknown ahead of time how many --- elements a vector will have. --- --- ====__Example__ --- --- >>> import Data.Massiv.Array as A --- >>> sunfoldr (\i -> if i < 9 then Just (i * i, i + 1) else Nothing) (0 :: Int) --- Array DS Seq (Sz1 9) --- [ 0, 1, 4, 9, 16, 25, 36, 49, 64 ] --- --- @since 0.5.0 -sunfoldr :: forall e s. (s -> Maybe (e, s)) -> s -> Vector DS e -sunfoldr f = DSArray . S.unfoldr f -{-# INLINE sunfoldr #-} - - - --- | /O(n)/ - Right unfolding function with at most @n@ number of elements. --- --- ==== __Example__ --- --- >>> import Data.Massiv.Array as A --- >>> sunfoldrN 9 (\i -> Just (i*i, i + 1)) (0 :: Int) --- Array DS Seq (Sz1 9) --- [ 0, 1, 4, 9, 16, 25, 36, 49, 64 ] --- --- @since 0.5.0 -sunfoldrN :: - forall e s. - Sz1 - -- ^ @n@ - maximum number of elements that the vector will have - -> (s -> Maybe (e, s)) - -- ^ Unfolding function. Stops when `Nothing` is returned or maximum number of elements - -- is reached. - -> s -- ^ Inititial element. - -> Vector DS e -sunfoldrN n f = DSArray . S.unfoldrN n f -{-# INLINE sunfoldrN #-} - --- | /O(n)/ - Same as `sunfoldr`, but with monadic generating function. --- --- ==== __Examples__ --- --- >>> import Control.Monad (when, guard) --- >>> sunfoldrM (\i -> when (i == 0) (Left "Zero denominator") >> Right (guard (i < 5) >> Just (100 `div` i, i + 1))) (-10 :: Int) --- Left "Zero denominator" --- >>> sunfoldrM (\i -> when (i == 0) (Left "Zero denominator") >> Right (guard (i < -5) >> Just (100 `div` i, i + 1))) (-10 :: Int) --- Right (Array DS Seq (Sz1 5) --- [ -10, -12, -13, -15, -17 ] --- ) --- --- @since 0.5.0 -sunfoldrM :: forall e s m. Monad m => (s -> m (Maybe (e, s))) -> s -> m (Vector DS e) -sunfoldrM f = fromStepsM . S.unfoldrM f -{-# INLINE sunfoldrM #-} - --- | /O(n)/ - Same as `sunfoldrN`, but with monadic generating function. --- --- ==== __Examples__ --- --- >>> import Control.Monad (guard) --- >>> sunfoldrNM 6 (\i -> print i >> pure (guard (i < 5) >> Just (i * i, i + 1))) (10 :: Int) --- 10 --- Array DS Seq (Sz1 0) --- [ ] --- >>> sunfoldrNM 6 (\i -> print i >> pure (guard (i < 15) >> Just (i * i, i + 1))) (10 :: Int) --- 10 --- 11 --- 12 --- 13 --- 14 --- 15 --- Array DS Seq (Sz1 5) --- [ 100, 121, 144, 169, 196 ] --- --- --- @since 0.5.0 -sunfoldrNM :: forall e s m. Monad m => Sz1 -> (s -> m (Maybe (e, s))) -> s -> m (Vector DS e) -sunfoldrNM (Sz n) f = fromStepsM . S.unfoldrNM n f -{-# INLINE sunfoldrNM #-} - - --- | /O(n)/ - Similar to `sunfoldrN`, except the length of the resulting vector will be exactly @n@ --- --- ==== __Examples__ --- --- >>> sunfoldrExactN 10 (\i -> (i * i, i + 1)) (10 :: Int) --- Array DS Seq (Sz1 10) --- [ 100, 121, 144, 169, 196, 225, 256, 289, 324, 361 ] --- --- @since 0.5.0 -sunfoldrExactN :: forall e s. Sz1 -> (s -> (e, s)) -> s -> Vector DS e -sunfoldrExactN n f = fromSteps . S.unfoldrExactN n f -{-# INLINE sunfoldrExactN #-} - --- | /O(n)/ - Similar to `sunfoldrNM`, except the length of the resulting vector will be exactly @n@ --- --- ==== __Examples__ --- --- λ> sunfoldrExactNM 11 (\i -> pure (100 `div` i, i + 1)) (-10 :: Int) --- Array DS *** Exception: divide by zero --- λ> sunfoldrExactNM 11 (\i -> guard (i /= 0) >> Just (100 `div` i, i + 1)) (-10 :: Int) --- Nothing --- λ> sunfoldrExactNM 9 (\i -> guard (i /= 0) >> Just (100 `div` i, i + 1)) (-10 :: Int) --- Just (Array DS Seq (Sz1 9) --- [ -10, -12, -13, -15, -17, -20, -25, -34, -50 ] --- ) --- --- @since 0.5.0 -sunfoldrExactNM :: forall e s m. Monad m => Sz1 -> (s -> m (e, s)) -> s -> m (Vector DS e) -sunfoldrExactNM n f = fromStepsM . S.unfoldrExactNM n f -{-# INLINE sunfoldrExactNM #-} - - --- | /O(n)/ - Enumerate from a starting number @x@ exactly @n@ times with a step @1@. --- --- /Related/: `senumFromStepN`, `enumFromN`, `enumFromStepN`, `rangeSize`, --- `rangeStepSize`, `range`, 'rangeStep'' --- --- ==== __Examples__ --- --- >>> senumFromN (10 :: Int) 9 --- Array DS Seq (Sz1 9) --- [ 10, 11, 12, 13, 14, 15, 16, 17, 18 ] --- --- /__Similar__/: --- --- [@Prelude.`Prelude.enumFromTo`@] Very similar to @[x .. x + n - 1]@, except that --- `senumFromN` is faster and it only works for `Num` and not for `Enum` elements --- --- [@Data.Vector.Generic.`Data.Vector.Generic.enumFromN`@] Uses exactly the same --- implementation underneath. --- --- @since 0.5.0 -senumFromN :: - Num e - => e -- ^ @x@ - starting number - -> Sz1 -- ^ @n@ - length of resulting vector - -> Vector DS e -senumFromN x n = DSArray $ S.enumFromStepN x 1 n -{-# INLINE senumFromN #-} - --- | /O(n)/ - Enumerate from a starting number @x@ exactly @n@ times with a custom step value @dx@ --- --- ==== __Examples__ --- --- >>> senumFromStepN (5 :: Int) 2 10 --- Array DS Seq (Sz1 10) --- [ 5, 7, 9, 11, 13, 15, 17, 19, 21, 23 ] --- --- __/Similar/__: --- --- [@Prelude.`Prelude.enumFrom`@] Just like @take n [x, x + dx ..]@, except that --- `senumFromN` is faster and it only works for `Num` and not for `Enum` elements --- --- [@Data.Vector.Generic.`Data.Vector.Generic.enumFromStepN`@] Uses exactly the same --- implementation underneath. --- --- @since 0.5.0 -senumFromStepN :: - Num e - => e -- ^ @x@ - starting number - -> e -- ^ @dx@ - Step - -> Sz1 -- ^ @n@ - length of resulting vector - -> Vector DS e -senumFromStepN x step n = DSArray $ S.enumFromStepN x step n -{-# INLINE senumFromStepN #-} - - - --- | Append two vectors together --- --- /Related/: `appendM`, `appendOuterM`, --- --- ==== __Examples__ --- --- λ> sappend (1 ..: 6) (senumFromStepN 6 (-1) 6) --- Array DS Seq (Sz1 11) --- [ 1, 2, 3, 4, 5, 6, 5, 4, 3, 2, 1 ] --- --- __/Similar/__: --- --- [@Data.Semigroup.`Data.Semigroup.<>`@] `DS` and `DL` arrays have instances for --- `Semigroup`, so they will work in a similar fashion. `sappend` differs in that it accepts --- `Stream` arrays with possibly different representations. --- --- [@Data.List.`Data.List.++`@] Same operation, but for lists. --- --- [@Data.Vector.Generic.`Data.Vector.Generic.++`@] Uses exactly the same implementation --- underneath as `sappend`, except that it cannot append two vectors with different --- memory representations. --- --- @since 0.5.0 -sappend :: - forall r1 r2 e. (Stream r1 Ix1 e, Stream r2 Ix1 e) - => Vector r1 e - -> Vector r2 e - -> Vector DS e -sappend a1 a2 = fromSteps (toStream a1 `S.append` toStream a2) -{-# INLINE sappend #-} - - --- | Concat vectors together --- --- /Related/: `concatM`, `concatOuterM`, --- --- ==== __Examples__ --- --- >>> sconcat [2 ... 6, empty, singleton 1, generate Seq 5 id] --- Array DS Seq (Sz1 11) --- [ 2, 3, 4, 5, 6, 1, 0, 1, 2, 3, 4 ] --- >>> sconcat [senumFromN 2 5, sempty, ssingleton 1, sgenerate 5 id] --- Array DS Seq (Sz1 11) --- [ 2, 3, 4, 5, 6, 1, 0, 1, 2, 3, 4 ] --- --- __/Similar/__: --- --- [@Data.Monoid.`Data.Monoid.mconcat`@] `DS` and `DL` arrays have instances for `Monoid`, so --- they will work in a similar fashion. `sconcat` differs in that it accepts `Stream` --- arrays of other representations. --- --- [@Data.List.`Data.List.concat`@] Same operation, but for lists. --- --- [@Data.Vector.Generic.`Data.Vector.Generic.concat`@] Uses exactly the same --- implementation underneath as `sconcat`. --- --- @since 0.5.0 -sconcat :: forall r e. Stream r Ix1 e => [Vector r e] -> Vector DS e -sconcat = DSArray . foldMap toStream -{-# INLINE sconcat #-} - --- | Convert a list to a delayed stream vector --- --- /Related/: `fromList`, `fromListN`, `sfromListN` --- --- ==== __Examples__ --- --- >>> sfromList ([] :: [Int]) --- Array DS Seq (Sz1 0) --- [ ] --- >>> sfromList ([1,2,3] :: [Int]) --- Array DS Seq (Sz1 3) --- [ 1, 2, 3 ] --- --- @since 0.5.0 -sfromList :: [e] -> Vector DS e -sfromList = fromSteps . S.fromList -{-# INLINE sfromList #-} - --- | Convert a list to a delayed stream vector. Length of the resulting vector will be at --- most @n@. This version isn't really more efficient then `sfromList`, but there is --- `Data.Massiv.Array.Unsafe.unsafeFromListN` --- --- /Related/: `fromList`, `fromListN`, `sfromList` --- --- ==== __Examples__ --- --- >>> sfromListN 10 [1 :: Int ..] --- Array DS Seq (Sz1 10) --- [ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 ] --- >>> sfromListN 10 [1 :: Int .. 5] --- Array DS Seq (Sz1 5) --- [ 1, 2, 3, 4, 5 ] --- --- @since 0.5.1 -sfromListN :: Sz1 -> [e] -> Vector DS e -sfromListN (Sz n) = fromSteps . S.fromListN n -{-# INLINE sfromListN #-} - --- | Convert an array to a list by the means of a delayed stream vector. --- --- /Related/: `toList` --- --- ==== __Examples__ --- --- @since 0.5.0 -stoList :: forall r ix e. Stream r ix e => Array r ix e -> [e] -stoList = S.toList . toStream -{-# INLINE stoList #-} - - - --- | Sequentially filter out elements from the array according to the supplied predicate. --- --- ==== __Example__ --- --- >>> import Data.Massiv.Array as A --- >>> arr = makeArrayR D Seq (Sz2 3 4) fromIx2 --- >>> arr --- Array D Seq (Sz (3 :. 4)) --- [ [ (0,0), (0,1), (0,2), (0,3) ] --- , [ (1,0), (1,1), (1,2), (1,3) ] --- , [ (2,0), (2,1), (2,2), (2,3) ] --- ] --- >>> sfilter (even . fst) arr --- Array DS Seq (Sz1 8) --- [ (0,0), (0,1), (0,2), (0,3), (2,0), (2,1), (2,2), (2,3) ] --- --- @since 0.5.0 -sfilter :: forall r ix e. S.Stream r ix e => (e -> Bool) -> Array r ix e -> Vector DS e -sfilter f = DSArray . S.filter f . S.toStream -{-# INLINE sfilter #-} - - --- | Similar to `sfilter`, but filter with an index aware function. --- --- ==== __Examples__ --- --- @since 0.5.0 -sifilter :: forall r ix e. Stream r ix e => (ix -> e -> Bool) -> Array r ix e -> Vector DS e -sifilter f = - simapMaybe $ \ix e -> - if f ix e - then Just e - else Nothing -{-# INLINE sifilter #-} - - --- | Sequentially filter out elements from the array according to the supplied applicative predicate. --- --- ==== __Example__ --- --- >>> import Data.Massiv.Array as A --- >>> arr = makeArrayR D Seq (Sz2 3 4) fromIx2 --- >>> arr --- Array D Seq (Sz (3 :. 4)) --- [ [ (0,0), (0,1), (0,2), (0,3) ] --- , [ (1,0), (1,1), (1,2), (1,3) ] --- , [ (2,0), (2,1), (2,2), (2,3) ] --- ] --- >>> sfilterM (Just . odd . fst) arr --- Just (Array DS Seq (Sz1 4) --- [ (1,0), (1,1), (1,2), (1,3) ] --- ) --- >>> sfilterM (\ix@(_, j) -> print ix >> return (even j)) arr --- (0,0) --- (0,1) --- (0,2) --- (0,3) --- (1,0) --- (1,1) --- (1,2) --- (1,3) --- (2,0) --- (2,1) --- (2,2) --- (2,3) --- Array DS Seq (Sz1 6) --- [ (0,0), (0,2), (1,0), (1,2), (2,0), (2,2) ] --- --- @since 0.5.0 -sfilterM :: - forall r ix e f. (S.Stream r ix e, Applicative f) - => (e -> f Bool) - -> Array r ix e - -> f (Vector DS e) -sfilterM f arr = DSArray <$> S.filterA f (S.toStream arr) -{-# INLINE sfilterM #-} - - --- | Similar to `filterM`, but filter with an index aware function. --- --- Corresponds to: @`filterM` (uncurry f) . `simap` (,)@ --- --- @since 0.5.0 -sifilterM :: - forall r ix e f. (Stream r ix e, Applicative f) - => (ix -> e -> f Bool) - -> Array r ix e - -> f (Vector DS e) -sifilterM f = - simapMaybeM $ \ix e -> - (\p -> - if p - then Just e - else Nothing) <$> - f ix e -{-# INLINE sifilterM #-} - - --- | Apply a function to each element of the array, while discarding `Nothing` and --- keeping the `Maybe` result. --- --- ==== __Examples__ --- --- @since 0.5.0 -smapMaybe :: forall r ix a b. S.Stream r ix a => (a -> Maybe b) -> Array r ix a -> Vector DS b -smapMaybe f = DSArray . S.mapMaybe f . S.toStream -{-# INLINE smapMaybe #-} - - --- | Similar to `smapMaybe`, but map with an index aware function. --- --- ==== __Examples__ --- --- @since 0.5.0 -simapMaybe :: - forall r ix a b. Stream r ix a - => (ix -> a -> Maybe b) - -> Array r ix a - -> Vector DS b -simapMaybe f = DSArray . S.mapMaybe (uncurry f) . toStreamIx -{-# INLINE simapMaybe #-} - --- | Similar to `smapMaybeM`, but map with an index aware function. --- --- ==== __Examples__ --- --- @since 0.5.0 -simapMaybeM :: - forall r ix a b f. (Stream r ix a, Applicative f) - => (ix -> a -> f (Maybe b)) - -> Array r ix a - -> f (Vector DS b) -simapMaybeM f = fmap DSArray . S.mapMaybeA (uncurry f) . toStreamIx -{-# INLINE simapMaybeM #-} - - --- | Keep all `Maybe`s and discard the `Nothing`s. --- --- ==== __Examples__ --- --- @since 0.5.0 -scatMaybes :: forall r ix a. S.Stream r ix (Maybe a) => Array r ix (Maybe a) -> Vector DS a -scatMaybes = smapMaybe id -{-# INLINE scatMaybes #-} - - --- | Similar to `smapMaybe`, but with the `Applicative` function. --- --- Similar to @mapMaybe id <$> mapM f arr@ --- --- ==== __Examples__ --- --- @since 0.5.0 -smapMaybeM :: - forall r ix a b f. (S.Stream r ix a, Applicative f) - => (a -> f (Maybe b)) - -> Array r ix a - -> f (Vector DS b) -smapMaybeM f = fmap DSArray . S.mapMaybeA f . S.toStream -{-# INLINE smapMaybeM #-} - - - --- | Map a function over a stream vector --- --- ==== __Examples__ --- --- @since 0.5.0 -smap :: - forall r ix a b. S.Stream r ix a - => (a -> b) - -> Array r ix a - -> Vector DS b -smap f = fromSteps . S.map f . S.toStream -{-# INLINE smap #-} - --- | Map an index aware function over a stream vector --- --- ==== __Examples__ --- --- @since 0.5.0 -simap :: - forall r ix a b. S.Stream r ix a - => (ix -> a -> b) - -> Array r ix a - -> Vector DS b -simap f = fromSteps . S.map (uncurry f) . S.toStreamIx -{-# INLINE simap #-} - - --- | Traverse a stream vector with an applicative function. --- --- ==== __Examples__ --- --- @since 0.5.0 -straverse :: - forall r ix a b f. (S.Stream r ix a, Applicative f) - => (a -> f b) - -> Array r ix a - -> f (Vector DS b) -straverse f = fmap fromSteps . S.traverse f . S.toStream -{-# INLINE straverse #-} - - --- | Traverse a stream vector with an index aware applicative function. --- --- ==== __Examples__ --- --- @since 0.5.0 -sitraverse :: - forall r ix a b f. (S.Stream r ix a, Applicative f) - => (ix -> a -> f b) - -> Array r ix a - -> f (Vector DS b) -sitraverse f = fmap fromSteps . S.traverse (uncurry f) . S.toStreamIx -{-# INLINE sitraverse #-} - - --- | Traverse a stream vector with a monadic function. --- --- ==== __Examples__ --- --- @since 0.5.0 -smapM :: - forall r ix a b m. (S.Stream r ix a, Monad m) - => (a -> m b) - -> Array r ix a - -> m (Vector DS b) -smapM f = fromStepsM . S.mapM f . S.transStepsId . S.toStream -{-# INLINE smapM #-} - --- | Traverse a stream vector with a monadic index aware function. --- --- Corresponds to: @mapM (uncurry f) . imap (,) v@ --- --- ==== __Examples__ --- --- @since 0.5.0 -simapM :: - forall r ix a b m. (S.Stream r ix a, Monad m) - => (ix -> a -> m b) - -> Array r ix a - -> m (Vector DS b) -simapM f = fromStepsM . S.mapM (uncurry f) . S.transStepsId . S.toStreamIx -{-# INLINE simapM #-} - --- | Traverse a stream vector with a monadic function, while discarding the result --- --- ==== __Examples__ --- --- @since 0.5.0 -smapM_ :: - forall r ix a b m. (S.Stream r ix a, Monad m) - => (a -> m b) - -> Array r ix a - -> m () -smapM_ f = S.mapM_ f . S.transStepsId . S.toStream -{-# INLINE smapM_ #-} - --- | Traverse a stream vector with a monadic index aware function, while discarding the result --- --- ==== __Examples__ --- --- @since 0.5.0 -simapM_ :: - forall r ix a b m. (S.Stream r ix a, Monad m) - => (ix -> a -> m b) - -> Array r ix a - -> m () -simapM_ f = S.mapM_ (uncurry f) . S.transStepsId . S.toStreamIx -{-# INLINE simapM_ #-} - - --- | Same as `smapM`, but with arguments flipped. --- --- ==== __Examples__ --- --- @since 0.5.0 -sforM :: - forall r ix a b m. (S.Stream r ix a, Monad m) - => Array r ix a - -> (a -> m b) - -> m (Vector DS b) -sforM = flip smapM -{-# INLINE sforM #-} - --- | Same as `simapM`, but with arguments flipped. --- --- ==== __Examples__ --- --- @since 0.5.0 -siforM :: - forall r ix a b m. (S.Stream r ix a, Monad m) - => Array r ix a - -> (ix -> a -> m b) - -> m (Vector DS b) -siforM = flip simapM -{-# INLINE siforM #-} - --- | Same as `smapM_`, but with arguments flipped. --- --- ==== __Examples__ --- --- @since 0.5.0 -sforM_ :: (S.Stream r ix a, Monad m) => Array r ix a -> (a -> m b) -> m () -sforM_ = flip smapM_ -{-# INLINE sforM_ #-} - --- | Same as `simapM_`, but with arguments flipped. --- --- ==== __Examples__ --- --- @since 0.5.0 -siforM_ :: - forall r ix a b m. (S.Stream r ix a, Monad m) - => Array r ix a - -> (ix -> a -> m b) - -> m () -siforM_ = flip simapM_ -{-# INLINE siforM_ #-} - - - --- | Zip two vectors together into a vector. The length of a resulting vector will --- be the smallest length of the supplied vectors. --- --- ==== __Examples__ --- --- @since 0.5.0 -szip :: - forall ra rb a b. (S.Stream ra Ix1 a, S.Stream rb Ix1 b) - => Vector ra a - -> Vector rb b - -> Vector DS (a, b) -szip = szipWith (,) -{-# INLINE szip #-} - --- | Zip three vectors together into a vector. The length of a resulting vector will --- be the smallest length of the supplied vectors. --- --- @since 0.5.0 -szip3 :: - forall ra rb rc a b c. (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c) - => Vector ra a - -> Vector rb b - -> Vector rc c - -> Vector DS (a, b, c) -szip3 = szipWith3 (,,) -{-# INLINE szip3 #-} - --- | Zip four vectors together into a vector. The length of a resulting vector will --- be the smallest length of the supplied vectors. --- --- @since 0.5.0 -szip4 :: - forall ra rb rc rd a b c d. - (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, S.Stream rd Ix1 d) - => Vector ra a - -> Vector rb b - -> Vector rc c - -> Vector rd d - -> Vector DS (a, b, c, d) -szip4 = szipWith4 (,,,) -{-# INLINE szip4 #-} - --- | Zip five vectors together into a vector. The length of a resulting vector will --- be the smallest length of the supplied vectors. --- --- @since 0.5.0 -szip5 :: - forall ra rb rc rd re a b c d e. - (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, S.Stream rd Ix1 d, S.Stream re Ix1 e) - => Vector ra a - -> Vector rb b - -> Vector rc c - -> Vector rd d - -> Vector re e - -> Vector DS (a, b, c, d, e) -szip5 = szipWith5 (,,,,) -{-# INLINE szip5 #-} - --- | Zip six vectors together into a vector. The length of a resulting vector will --- be the smallest length of the supplied vectors. --- --- @since 0.5.0 -szip6 :: - forall ra rb rc rd re rf a b c d e f. - ( S.Stream ra Ix1 a - , S.Stream rb Ix1 b - , S.Stream rc Ix1 c - , S.Stream rd Ix1 d - , S.Stream re Ix1 e - , S.Stream rf Ix1 f - ) - => Vector ra a - -> Vector rb b - -> Vector rc c - -> Vector rd d - -> Vector re e - -> Vector rf f - -> Vector DS (a, b, c, d, e, f) -szip6 = szipWith6 (,,,,,) -{-# INLINE szip6 #-} - - - - --- | Zip two vectors together with a binary function into a vector. The length --- of a resulting vector will be the smallest length of the supplied vectors. --- --- ==== __Examples__ --- --- @since 0.5.0 -szipWith :: - forall ra rb a b c. - (S.Stream ra Ix1 a, S.Stream rb Ix1 b) - => (a -> b -> c) - -> Vector ra a - -> Vector rb b - -> Vector DS c -szipWith f v1 v2 = fromSteps $ S.zipWith f (S.toStream v1) (S.toStream v2) -{-# INLINE szipWith #-} - --- | Zip three vectors together with a ternary function into a vector. The length --- of a resulting vector will be the smallest length of the supplied vectors. --- --- @since 0.5.0 -szipWith3 :: - forall ra rb rc a b c d. - (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c) - => (a -> b -> c -> d) - -> Vector ra a - -> Vector rb b - -> Vector rc c - -> Vector DS d -szipWith3 f v1 v2 v3 = fromSteps $ S.zipWith3 f (S.toStream v1) (S.toStream v2) (S.toStream v3) -{-# INLINE szipWith3 #-} - --- | Zip four vectors together with a quaternary function into a vector. The length --- of a resulting vector will be the smallest length of the supplied vectors. --- --- @since 0.5.0 -szipWith4 :: - forall ra rb rc rd a b c d e. - (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, S.Stream rd Ix1 d) - => (a -> b -> c -> d -> e) - -> Vector ra a - -> Vector rb b - -> Vector rc c - -> Vector rd d - -> Vector DS e -szipWith4 f v1 v2 v3 v4 = - fromSteps $ S.zipWith4 f (S.toStream v1) (S.toStream v2) (S.toStream v3) (S.toStream v4) -{-# INLINE szipWith4 #-} - --- | Zip five vectors together with a quinary function into a vector. The length --- of a resulting vector will be the smallest length of the supplied vectors. --- --- @since 0.5.0 -szipWith5 :: - forall ra rb rc rd re a b c d e f. - (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, S.Stream rd Ix1 d, S.Stream re Ix1 e) - => (a -> b -> c -> d -> e -> f) - -> Vector ra a - -> Vector rb b - -> Vector rc c - -> Vector rd d - -> Vector re e - -> Vector DS f -szipWith5 f v1 v2 v3 v4 v5 = - fromSteps $ - S.zipWith5 f (S.toStream v1) (S.toStream v2) (S.toStream v3) (S.toStream v4) (S.toStream v5) -{-# INLINE szipWith5 #-} - --- | Zip six vectors together with a senary function into a vector. The length --- of a resulting vector will be the smallest length of the supplied vectors. --- --- @since 0.5.0 -szipWith6 :: - forall ra rb rc rd re rf a b c d e f g. - ( S.Stream ra Ix1 a - , S.Stream rb Ix1 b - , S.Stream rc Ix1 c - , S.Stream rd Ix1 d - , S.Stream re Ix1 e - , S.Stream rf Ix1 f - ) - => (a -> b -> c -> d -> e -> f -> g) - -> Vector ra a - -> Vector rb b - -> Vector rc c - -> Vector rd d - -> Vector re e - -> Vector rf f - -> Vector DS g -szipWith6 f v1 v2 v3 v4 v5 v6 = - fromSteps $ - S.zipWith6 - f - (S.toStream v1) - (S.toStream v2) - (S.toStream v3) - (S.toStream v4) - (S.toStream v5) - (S.toStream v6) -{-# INLINE szipWith6 #-} - --- | Just like `szipWith`, zip two vectors together, but with an index aware --- function. The length of a resulting vector will be the smallest length of the --- supplied vectors. --- --- ==== __Examples__ --- --- @since 0.5.0 -sizipWith :: - forall ra rb a b c. - (S.Stream ra Ix1 a, S.Stream rb Ix1 b) - => (Ix1 -> a -> b -> c) - -> Vector ra a - -> Vector rb b - -> Vector DS c -sizipWith f v1 v2 = fromSteps $ S.zipWith (uncurry f) (S.toStreamIx v1) (S.toStream v2) -{-# INLINE sizipWith #-} - --- | Just like `szipWith3`, zip three vectors together, but with an index aware --- function. The length of a resulting vector will be the smallest length of the --- supplied vectors. --- --- @since 0.5.0 -sizipWith3 :: - forall ra rb rc a b c d. - (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c) - => (Ix1 -> a -> b -> c -> d) - -> Vector ra a - -> Vector rb b - -> Vector rc c - -> Vector DS d -sizipWith3 f v1 v2 v3 = - fromSteps $ S.zipWith3 (uncurry f) (S.toStreamIx v1) (S.toStream v2) (S.toStream v3) -{-# INLINE sizipWith3 #-} - --- | Just like `szipWith4`, zip four vectors together, but with an index aware --- function. The length of a resulting vector will be the smallest --- length of the supplied vectors. --- --- @since 0.5.0 -sizipWith4 :: - forall ra rb rc rd a b c d e. - (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, S.Stream rd Ix1 d) - => (Ix1 -> a -> b -> c -> d -> e) - -> Vector ra a - -> Vector rb b - -> Vector rc c - -> Vector rd d - -> Vector DS e -sizipWith4 f v1 v2 v3 v4 = - fromSteps $ - S.zipWith4 (uncurry f) (S.toStreamIx v1) (S.toStream v2) (S.toStream v3) (S.toStream v4) -{-# INLINE sizipWith4 #-} - --- | Just like `szipWith5`, zip five vectors together, but with an index aware --- function. The length of a resulting vector will be the smallest length of the --- supplied vectors. --- --- @since 0.5.0 -sizipWith5 :: - forall ra rb rc rd re a b c d e f. - (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, S.Stream rd Ix1 d, S.Stream re Ix1 e) - => (Ix1 -> a -> b -> c -> d -> e -> f) - -> Vector ra a - -> Vector rb b - -> Vector rc c - -> Vector rd d - -> Vector re e - -> Vector DS f -sizipWith5 f v1 v2 v3 v4 v5 = - fromSteps $ - S.zipWith5 - (uncurry f) - (S.toStreamIx v1) - (S.toStream v2) - (S.toStream v3) - (S.toStream v4) - (S.toStream v5) -{-# INLINE sizipWith5 #-} - --- | Just like `szipWith6`, zip six vectors together, but with an index aware --- function. The length of a resulting vector will be the smallest length of the --- supplied vectors. --- --- @since 0.5.0 -sizipWith6 :: - forall ra rb rc rd re rf a b c d e f g. - ( S.Stream ra Ix1 a - , S.Stream rb Ix1 b - , S.Stream rc Ix1 c - , S.Stream rd Ix1 d - , S.Stream re Ix1 e - , S.Stream rf Ix1 f - ) - => (Ix1 -> a -> b -> c -> d -> e -> f -> g) - -> Vector ra a - -> Vector rb b - -> Vector rc c - -> Vector rd d - -> Vector re e - -> Vector rf f - -> Vector DS g -sizipWith6 f v1 v2 v3 v4 v5 v6 = - fromSteps $ - S.zipWith6 - (uncurry f) - (S.toStreamIx v1) - (S.toStream v2) - (S.toStream v3) - (S.toStream v4) - (S.toStream v5) - (S.toStream v6) -{-# INLINE sizipWith6 #-} - - --- | Zip two vectors together with a binary monadic action into a vector. The --- length of a resulting vector will be the smallest length of the supplied --- vectors. --- --- ==== __Examples__ --- --- @since 0.5.0 -szipWithM :: - forall ra rb a b c m. - (S.Stream ra Ix1 a, S.Stream rb Ix1 b, Monad m) - => (a -> b -> m c) - -> Vector ra a - -> Vector rb b - -> m (Vector DS c) -szipWithM f v1 v2 = fromStepsM $ S.zipWithM f (toStreamM v1) (toStreamM v2) -{-# INLINE szipWithM #-} - --- | Zip three vectors together with a ternary monadic action into a vector. The --- length of a resulting vector will be the smallest length of the supplied --- vectors. --- --- @since 0.5.0 -szipWith3M :: - forall ra rb rc a b c d m. - (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, Monad m) - => (a -> b -> c -> m d) - -> Vector ra a - -> Vector rb b - -> Vector rc c - -> m (Vector DS d) -szipWith3M f v1 v2 v3 = fromStepsM $ S.zipWith3M f (toStreamM v1) (toStreamM v2) (toStreamM v3) -{-# INLINE szipWith3M #-} - --- | Zip four vectors together with a quaternary monadic action into a vector. The --- length of a resulting vector will be the smallest length of the supplied --- vectors. --- --- @since 0.5.0 -szipWith4M :: - forall ra rb rc rd a b c d e m. - (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, S.Stream rd Ix1 d, Monad m) - => (a -> b -> c -> d -> m e) - -> Vector ra a - -> Vector rb b - -> Vector rc c - -> Vector rd d - -> m (Vector DS e) -szipWith4M f v1 v2 v3 v4 = - fromStepsM $ S.zipWith4M f (toStreamM v1) (toStreamM v2) (toStreamM v3) (toStreamM v4) -{-# INLINE szipWith4M #-} - --- | Zip five vectors together with a quinary monadic action into a vector. The --- length of a resulting vector will be the smallest length of the supplied --- vectors. --- --- @since 0.5.0 -szipWith5M :: - forall ra rb rc rd re a b c d e f m. - ( S.Stream ra Ix1 a - , S.Stream rb Ix1 b - , S.Stream rc Ix1 c - , S.Stream rd Ix1 d - , S.Stream re Ix1 e - , Monad m - ) - => (a -> b -> c -> d -> e -> m f) - -> Vector ra a - -> Vector rb b - -> Vector rc c - -> Vector rd d - -> Vector re e - -> m (Vector DS f) -szipWith5M f v1 v2 v3 v4 v5 = - fromStepsM $ - S.zipWith5M f (toStreamM v1) (toStreamM v2) (toStreamM v3) (toStreamM v4) (toStreamM v5) -{-# INLINE szipWith5M #-} - --- | Zip six vectors together with a senary monadic action into a vector. The --- length of a resulting vector will be the smallest length of the supplied --- vectors. --- --- @since 0.5.0 -szipWith6M :: - forall ra rb rc rd re rf a b c d e f g m. - ( S.Stream ra Ix1 a - , S.Stream rb Ix1 b - , S.Stream rc Ix1 c - , S.Stream rd Ix1 d - , S.Stream re Ix1 e - , S.Stream rf Ix1 f - , Monad m - ) - => (a -> b -> c -> d -> e -> f -> m g) - -> Vector ra a - -> Vector rb b - -> Vector rc c - -> Vector rd d - -> Vector re e - -> Vector rf f - -> m (Vector DS g) -szipWith6M f v1 v2 v3 v4 v5 v6 = - fromStepsM $ - S.zipWith6M - f - (toStreamM v1) - (toStreamM v2) - (toStreamM v3) - (toStreamM v4) - (toStreamM v5) - (toStreamM v6) -{-# INLINE szipWith6M #-} - - --- | Just like `szipWithM`, zip two vectors together, but with an index aware --- monadic action. The length of a resulting vector will be the smallest length of the --- supplied vectors. --- --- ==== __Examples__ --- --- @since 0.5.0 -sizipWithM :: - forall ra rb a b c m. - (S.Stream ra Ix1 a, S.Stream rb Ix1 b, Monad m) - => (Ix1 -> a -> b -> m c) - -> Vector ra a - -> Vector rb b - -> m (Vector DS c) -sizipWithM f v1 v2 = fromStepsM $ S.zipWithM (uncurry f) (toStreamIxM v1) (toStreamM v2) -{-# INLINE sizipWithM #-} - - --- | Just like `szipWith3M`, zip three vectors together, but with an index aware --- monadic action. The length of a resulting vector will be the smallest length of the --- supplied vectors. --- --- @since 0.5.0 -sizipWith3M :: - forall ra rb rc a b c d m. - (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, Monad m) - => (Ix1 -> a -> b -> c -> m d) - -> Vector ra a - -> Vector rb b - -> Vector rc c - -> m (Vector DS d) -sizipWith3M f v1 v2 v3 = - fromStepsM $ S.zipWith3M (uncurry f) (toStreamIxM v1) (toStreamM v2) (toStreamM v3) -{-# INLINE sizipWith3M #-} - --- | Just like `szipWith4M`, zip four vectors together, but with an index aware --- monadic action. The length of a resulting vector will be the smallest length of the --- supplied vectors. --- --- @since 0.5.0 -sizipWith4M :: - forall ra rb rc rd a b c d e m. - (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, S.Stream rd Ix1 d, Monad m) - => (Ix1 -> a -> b -> c -> d -> m e) - -> Vector ra a - -> Vector rb b - -> Vector rc c - -> Vector rd d - -> m (Vector DS e) -sizipWith4M f v1 v2 v3 v4 = - fromStepsM $ - S.zipWith4M (uncurry f) (toStreamIxM v1) (toStreamM v2) (toStreamM v3) (toStreamM v4) -{-# INLINE sizipWith4M #-} - --- | Just like `szipWith6M`, zip five vectors together, but with an index aware --- monadic action. The length of a resulting vector will be the smallest length of the --- supplied vectors. --- --- @since 0.5.0 -sizipWith5M :: - forall ra rb rc rd re a b c d e f m. - ( S.Stream ra Ix1 a - , S.Stream rb Ix1 b - , S.Stream rc Ix1 c - , S.Stream rd Ix1 d - , S.Stream re Ix1 e - , Monad m - ) - => (Ix1 -> a -> b -> c -> d -> e -> m f) - -> Vector ra a - -> Vector rb b - -> Vector rc c - -> Vector rd d - -> Vector re e - -> m (Vector DS f) -sizipWith5M f v1 v2 v3 v4 v5 = - fromStepsM $ - S.zipWith5M - (uncurry f) - (toStreamIxM v1) - (toStreamM v2) - (toStreamM v3) - (toStreamM v4) - (toStreamM v5) -{-# INLINE sizipWith5M #-} - --- | Just like `szipWith6M`, zip six vectors together, but with an index aware --- monadic action. The length of a resulting vector will be the smallest length of the --- supplied vectors. --- --- ==== __Examples__ --- --- @since 0.5.0 -sizipWith6M :: - forall ra rb rc rd re rf a b c d e f g m. - ( S.Stream ra Ix1 a - , S.Stream rb Ix1 b - , S.Stream rc Ix1 c - , S.Stream rd Ix1 d - , S.Stream re Ix1 e - , S.Stream rf Ix1 f - , Monad m - ) - => (Ix1 -> a -> b -> c -> d -> e -> f -> m g) - -> Vector ra a - -> Vector rb b - -> Vector rc c - -> Vector rd d - -> Vector re e - -> Vector rf f - -> m (Vector DS g) -sizipWith6M f v1 v2 v3 v4 v5 v6 = - fromStepsM $ - S.zipWith6M - (uncurry f) - (toStreamIxM v1) - (toStreamM v2) - (toStreamM v3) - (toStreamM v4) - (toStreamM v5) - (toStreamM v6) -{-# INLINE sizipWith6M #-} - - --- | Similar to `szipWithM`, zip two vectors together with a binary monadic --- action, while discarding its result. The action will be invoked as many times as --- the length of the smallest vector. --- --- ==== __Examples__ --- --- @since 0.5.0 -szipWithM_ :: - forall ra rb a b c m. (S.Stream ra Ix1 a, S.Stream rb Ix1 b, Monad m) - => (a -> b -> m c) - -> Vector ra a - -> Vector rb b - -> m () -szipWithM_ f v1 v2 = S.zipWithM_ f (toStreamM v1) (toStreamM v2) -{-# INLINE szipWithM_ #-} - --- | Similar to `szipWith3M`, zip three vectors together with a ternary monadic --- action, while discarding its result. The action will be invoked as many times as --- the length of the smallest vector. --- --- @since 0.5.0 -szipWith3M_ :: - forall ra rb rc a b c d m. - (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, Monad m) - => (a -> b -> c -> m d) - -> Vector ra a - -> Vector rb b - -> Vector rc c - -> m () -szipWith3M_ f v1 v2 v3 = S.zipWith3M_ f (toStreamM v1) (toStreamM v2) (toStreamM v3) -{-# INLINE szipWith3M_ #-} - --- | Similar to `szipWith4M`, zip four vectors together with a quaternary monadic --- action, while discarding its result. The action will be invoked as many times as --- the length of the smallest vector. --- --- @since 0.5.0 -szipWith4M_ :: - forall ra rb rc rd a b c d e m. - (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, S.Stream rd Ix1 d, Monad m) - => (a -> b -> c -> d -> m e) - -> Vector ra a - -> Vector rb b - -> Vector rc c - -> Vector rd d - -> m () -szipWith4M_ f v1 v2 v3 v4 = - S.zipWith4M_ f (toStreamM v1) (toStreamM v2) (toStreamM v3) (toStreamM v4) -{-# INLINE szipWith4M_ #-} - --- | Similar to `szipWith5M`, zip five vectors together with a quinary monadic --- action, while discarding its result. The action will be invoked as many times as --- the length of the smallest vector. --- --- @since 0.5.0 -szipWith5M_ :: - forall ra rb rc rd re a b c d e f m. - ( S.Stream ra Ix1 a - , S.Stream rb Ix1 b - , S.Stream rc Ix1 c - , S.Stream rd Ix1 d - , S.Stream re Ix1 e - , Monad m - ) - => (a -> b -> c -> d -> e -> m f) - -> Vector ra a - -> Vector rb b - -> Vector rc c - -> Vector rd d - -> Vector re e - -> m () -szipWith5M_ f v1 v2 v3 v4 v5 = - S.zipWith5M_ f (toStreamM v1) (toStreamM v2) (toStreamM v3) (toStreamM v4) (toStreamM v5) -{-# INLINE szipWith5M_ #-} - --- | Similar to `szipWith6M`, zip six vectors together with a senary monadic --- action, while discarding its result. The action will be invoked as many times as --- the length of the smallest vector. --- --- @since 0.5.0 -szipWith6M_ :: - forall ra rb rc rd re rf a b c d e f g m. - ( S.Stream ra Ix1 a - , S.Stream rb Ix1 b - , S.Stream rc Ix1 c - , S.Stream rd Ix1 d - , S.Stream re Ix1 e - , S.Stream rf Ix1 f - , Monad m - ) - => (a -> b -> c -> d -> e -> f -> m g) - -> Vector ra a - -> Vector rb b - -> Vector rc c - -> Vector rd d - -> Vector re e - -> Vector rf f - -> m () -szipWith6M_ f v1 v2 v3 v4 v5 v6 = - S.zipWith6M_ - f - (toStreamM v1) - (toStreamM v2) - (toStreamM v3) - (toStreamM v4) - (toStreamM v5) - (toStreamM v6) -{-# INLINE szipWith6M_ #-} - - - --- | Same as `szipWithM_`, zip two vectors together, but with an index aware --- monadic action. The action will be invoked as many times as the length of the --- smallest vector. --- --- ==== __Examples__ --- --- @since 0.5.0 -sizipWithM_ :: - forall ra rb a b c m. (S.Stream ra Ix1 a, S.Stream rb Ix1 b, Monad m) - => (Ix1 -> a -> b -> m c) - -> Vector ra a - -> Vector rb b - -> m () -sizipWithM_ f v1 v2 = S.zipWithM_ (uncurry f) (toStreamIxM v1) (toStreamM v2) -{-# INLINE sizipWithM_ #-} - - --- | Same as `szipWith3M_`, zip three vectors together, but with an index aware --- monadic action. The action will be invoked as many times as the length of the --- smallest vector. --- --- @since 0.5.0 -sizipWith3M_ :: - forall ra rb rc a b c d m. - (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, Monad m) - => (Ix1 -> a -> b -> c -> m d) - -> Vector ra a - -> Vector rb b - -> Vector rc c - -> m () -sizipWith3M_ f v1 v2 v3 = S.zipWith3M_ (uncurry f) (toStreamIxM v1) (toStreamM v2) (toStreamM v3) -{-# INLINE sizipWith3M_ #-} - --- | Same as `szipWith4M_`, zip four vectors together, but with an index aware --- monadic action. The action will be invoked as many times as the length of the --- smallest vector. --- --- @since 0.5.0 -sizipWith4M_ :: - forall ra rb rc rd a b c d e m. - (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, S.Stream rd Ix1 d, Monad m) - => (Ix1 -> a -> b -> c -> d -> m e) - -> Vector ra a - -> Vector rb b - -> Vector rc c - -> Vector rd d - -> m () -sizipWith4M_ f v1 v2 v3 v4 = - S.zipWith4M_ (uncurry f) (toStreamIxM v1) (toStreamM v2) (toStreamM v3) (toStreamM v4) -{-# INLINE sizipWith4M_ #-} - --- | Same as `szipWith5M_`, zip five vectors together, but with an index aware --- monadic action. The action will be invoked as many times as the length of the --- smallest vector. --- --- @since 0.5.0 -sizipWith5M_ :: - forall ra rb rc rd re a b c d e f m. - ( S.Stream ra Ix1 a - , S.Stream rb Ix1 b - , S.Stream rc Ix1 c - , S.Stream rd Ix1 d - , S.Stream re Ix1 e - , Monad m - ) - => (Ix1 -> a -> b -> c -> d -> e -> m f) - -> Vector ra a - -> Vector rb b - -> Vector rc c - -> Vector rd d - -> Vector re e - -> m () -sizipWith5M_ f v1 v2 v3 v4 v5 = - S.zipWith5M_ - (uncurry f) - (toStreamIxM v1) - (toStreamM v2) - (toStreamM v3) - (toStreamM v4) - (toStreamM v5) -{-# INLINE sizipWith5M_ #-} - --- | Same as `szipWith6M_`, zip six vectors together, but with an index aware --- monadic action. The action will be invoked as many times as the length of the --- smallest vector. --- --- @since 0.5.0 -sizipWith6M_ :: - forall ra rb rc rd re rf a b c d e f g m. - ( S.Stream ra Ix1 a - , S.Stream rb Ix1 b - , S.Stream rc Ix1 c - , S.Stream rd Ix1 d - , S.Stream re Ix1 e - , S.Stream rf Ix1 f - , Monad m - ) - => (Ix1 -> a -> b -> c -> d -> e -> f -> m g) - -> Vector ra a - -> Vector rb b - -> Vector rc c - -> Vector rd d - -> Vector re e - -> Vector rf f - -> m () -sizipWith6M_ f v1 v2 v3 v4 v5 v6 = - S.zipWith6M_ - (uncurry f) - (toStreamIxM v1) - (toStreamM v2) - (toStreamM v3) - (toStreamM v4) - (toStreamM v5) - (toStreamM v6) -{-# INLINE sizipWith6M_ #-} - - - - - --- | Streaming fold over an array in a row-major fashion with a left biased --- function and a strict accumulator. --- --- ==== __Examples__ --- --- @since 0.5.0 -sfoldl :: - forall r ix e a. Stream r ix e - => (a -> e -> a) - -> a - -> Array r ix e - -> a -sfoldl f acc = S.unId . S.foldl f acc . toStream -{-# INLINE sfoldl #-} - --- | --- --- ==== __Examples__ --- --- @since 0.5.0 -sfoldlM :: - forall r ix e a m. (Stream r ix e, Monad m) - => (a -> e -> m a) - -> a - -> Array r ix e - -> m a -sfoldlM f acc = S.foldlM f acc . S.transStepsId . toStream -{-# INLINE sfoldlM #-} - --- | --- --- ==== __Examples__ --- --- @since 0.5.0 -sfoldlM_ :: - forall r ix e a m. (Stream r ix e, Monad m) - => (a -> e -> m a) - -> a - -> Array r ix e - -> m () -sfoldlM_ f acc = void . sfoldlM f acc -{-# INLINE sfoldlM_ #-} - - --- | --- --- ==== __Examples__ --- --- @since 0.5.0 -sfoldl1' :: - forall r ix e. (HasCallStack, Stream r ix e) - => (e -> e -> e) - -> Array r ix e - -> e -sfoldl1' f = throwEither . sfoldl1M (\e -> pure . f e) -{-# INLINE sfoldl1' #-} - --- | --- --- ==== __Examples__ --- --- @since 0.5.0 -sfoldl1M :: - forall r ix e m. (Stream r ix e, MonadThrow m) - => (e -> e -> m e) - -> Array r ix e - -> m e -sfoldl1M f arr = do - let str = S.transStepsId $ toStream arr - isNullStream <- S.null str - when isNullStream $ throwM $ SizeEmptyException (outerSize arr) - S.foldl1M f str -{-# INLINE sfoldl1M #-} - --- | --- --- ==== __Examples__ --- --- @since 0.5.0 -sfoldl1M_ :: - forall r ix e m. (Stream r ix e, MonadThrow m) - => (e -> e -> m e) - -> Array r ix e - -> m () -sfoldl1M_ f = void . sfoldl1M f -{-# INLINE sfoldl1M_ #-} - - - --- | --- --- ==== __Examples__ --- --- @since 0.5.0 -sifoldl :: - forall r ix e a. Stream r ix e - => (a -> ix -> e -> a) - -> a - -> Array r ix e - -> a -sifoldl f acc = S.unId . S.foldl (\a (ix, e) -> f a ix e) acc . toStreamIx -{-# INLINE sifoldl #-} - --- | --- --- ==== __Examples__ --- --- @since 0.5.0 -sifoldlM :: - forall r ix e a m. (Stream r ix e, Monad m) - => (a -> ix -> e -> m a) - -> a - -> Array r ix e - -> m a -sifoldlM f acc = S.foldlM (\a (ix, e) -> f a ix e) acc . S.transStepsId . toStreamIx -{-# INLINE sifoldlM #-} - --- | --- --- ==== __Examples__ --- --- @since 0.5.0 -sifoldlM_ :: - forall r ix e a m. (Stream r ix e, Monad m) - => (a -> ix -> e -> m a) - -> a - -> Array r ix e - -> m () -sifoldlM_ f acc = void . sifoldlM f acc -{-# INLINE sifoldlM_ #-} - - --- | --- --- ==== __Examples__ --- --- @since 0.5.0 -sor :: - forall r ix. Stream r ix Bool - => Array r ix Bool - -> Bool -sor = S.unId . S.or . toStream -{-# INLINE sor #-} - - --- | --- --- ==== __Examples__ --- --- @since 0.5.0 -sand :: forall r ix. Stream r ix Bool => Array r ix Bool -> Bool -sand = S.unId . S.and . toStream -{-# INLINE sand #-} - - --- | --- --- ==== __Examples__ --- --- @since 0.5.0 -sany :: forall r ix e. Stream r ix e => (e -> Bool) -> Array r ix e -> Bool -sany f = S.unId . S.or . S.map f . toStream -{-# INLINE sany #-} - - --- | --- --- ==== __Examples__ --- --- @since 0.5.0 -sall :: forall r ix e. Stream r ix e => (e -> Bool) -> Array r ix e -> Bool -sall f = S.unId . S.and . S.map f . toStream -{-# INLINE sall #-} - - - --- | Add all elements of the array together --- --- /Related/: `sum`. --- --- ==== __Examples__ --- --- >>> import Data.Massiv.Vector as V --- >>> V.ssum $ V.sfromList [10, 3, 70, 5 :: Int] --- 88 --- --- @since 0.5.0 -ssum :: forall r ix e. (Num e, Stream r ix e) => Array r ix e -> e -ssum = sfoldl (+) 0 -{-# INLINE ssum #-} - --- | Multiply all elements of the array together --- --- /Related/: `product`. --- --- ==== __Examples__ --- --- >>> import Data.Massiv.Vector as V --- >>> V.sproduct $ V.sfromList [10, 3, 70, 5 :: Int] --- 10500 --- --- @since 0.5.0 -sproduct :: forall r ix e. (Num e, Stream r ix e) => Array r ix e -> e -sproduct = sfoldl (*) 1 -{-# INLINE sproduct #-} - - --- | /O(n)/ - Find the largest value in the array. Throws an error on empty. --- --- /Related/: `smaximumM`, `maximum`, `maximumM`. --- --- ==== __Examples__ --- --- >>> import Data.Massiv.Vector as V --- >>> V.smaximum' $ V.sfromList [10, 3, 70, 5 :: Int] --- 70 --- --- @since 0.5.0 -smaximum' :: forall r ix e. (HasCallStack, Ord e, Stream r ix e) => Array r ix e -> e -smaximum' = sfoldl1' max -{-# INLINE smaximum' #-} - --- | /O(n)/ - Find the largest value in the array. --- --- /Related/: `smaximum`, `maximum`, `maximumM`. --- --- /__Throws Exceptions__/: `SizeEmptyException` when array is empty --- --- ==== __Examples__ --- --- >>> import Data.Massiv.Vector as V --- >>> V.smaximumM $ V.sfromList [10, 3, 70, 5 :: Int] --- 70 --- >>> V.smaximumM (V.empty :: Vector D Int) :: Maybe Int --- Nothing --- --- @since 0.5.0 -smaximumM :: forall r ix e m. (Ord e, Stream r ix e, MonadThrow m) => Array r ix e -> m e -smaximumM = sfoldl1M (\e acc -> pure (max e acc)) -{-# INLINE smaximumM #-} - - --- | /O(n)/ - Find the smallest value in the array. Throws an error on empty. --- --- /Related/: `sminimumM`, `minimum`, `minimumM`. --- --- ==== __Examples__ --- --- >>> import Data.Massiv.Vector as V --- >>> V.sminimum' $ V.sfromList [10, 3, 70, 5 :: Int] --- 3 --- --- @since 0.5.0 -sminimum' :: forall r ix e. (HasCallStack, Ord e, Stream r ix e) => Array r ix e -> e -sminimum' = sfoldl1' min -{-# INLINE sminimum' #-} - --- | /O(n)/ - Find the smallest value in the array. --- --- /Related/: 'sminimum'', `minimum`, `minimumM`. --- --- /__Throws Exceptions__/: `SizeEmptyException` when array is empty --- --- ==== __Examples__ --- --- >>> import Data.Massiv.Vector as V --- >>> V.sminimumM $ V.sfromList [10, 3, 70, 5 :: Int] --- 3 --- >>> V.sminimumM (V.empty :: Array D Ix2 Int) :: Maybe Int --- Nothing --- --- @since 0.5.0 -sminimumM :: forall r ix e m. (Ord e, Stream r ix e, MonadThrow m) => Array r ix e -> m e -sminimumM = sfoldl1M (\e acc -> pure (min e acc)) -{-# INLINE sminimumM #-} + +-- | +-- Module : Data.Massiv.Vector +-- Copyright : (c) Alexey Kuleshevich 2020-2022 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +module Data.Massiv.Vector ( + Vector, + MVector, + + -- * Accessors + + -- *** Size + slength, + maxLinearSize, + size, + isNull, + isNotNull, + + -- *** Indexing + (!?), + (!), + index, + index', + head', + shead', + last', + + -- *** Monadic Indexing + indexM, + headM, + sheadM, + lastM, + unconsM, + unsnocM, + + -- ** Slicing + slice, + slice', + sliceM, + sslice, + sliceAt, + sliceAt', + sliceAtM, + + -- *** Init + init, + init', + initM, + + -- *** Tail + tail, + tail', + tailM, + + -- *** Take + take, + take', + takeM, + takeWhile, + stake, + + -- *** Drop + drop, + dropWhile, + drop', + dropM, + sdrop, + + -- * Construction + + -- ** Initialization + empty, + sempty, + singleton, + ssingleton, + cons, + snoc, + A.replicate, + sreplicate, + generate, + sgenerate, + -- , iterateN + -- , iiterateN + siterate, + siterateN, + + -- ** Monadic initialization + sreplicateM, + sgenerateM, + siterateNM, + -- , create + -- , createT + + -- ** Unfolding + sunfoldr, + sunfoldrM, + sunfoldrN, + sunfoldrNM, + sunfoldrExactN, + sunfoldrExactNM, + -- , constructN + -- , constructrN + + -- ** Enumeration + (...), + (..:), + enumFromN, + senumFromN, + enumFromStepN, + senumFromStepN, + + -- ** Concatenation + + -- , consS -- cons + -- , snocS -- snoc + sappend, -- (++) + sconcat, -- concat + -- -- ** Restricitng memory usage + -- , force + -- -- * Modifying + -- -- ** Bulk updates + -- , (//) + -- , update_ + -- -- ** Accumulations + -- , accum + -- , accumulate_ + -- -- ** Permutations + -- , reverse + -- , backpermute + -- -- ** Manifest updates + -- , modify + -- -- * Elementwise + -- -- ** Mapping + smap, + simap, + -- , sconcatMap + + -- ** Monadic mapping + straverse, + sitraverse, + smapM, + smapM_, + simapM, + simapM_, + sforM, + sforM_, + siforM, + siforM_, + + -- ** Zipping + szip, + szip3, + szip4, + szip5, + szip6, + szipWith, + szipWith3, + szipWith4, + szipWith5, + szipWith6, + sizipWith, + sizipWith3, + sizipWith4, + sizipWith5, + sizipWith6, + + -- ** Monadic zipping + szipWithM, + szipWith3M, + szipWith4M, + szipWith5M, + szipWith6M, + sizipWithM, + sizipWith3M, + sizipWith4M, + sizipWith5M, + sizipWith6M, + szipWithM_, + szipWith3M_, + szipWith4M_, + szipWith5M_, + szipWith6M_, + sizipWithM_, + sizipWith3M_, + sizipWith4M_, + sizipWith5M_, + sizipWith6M_, + + -- * Predicates + + -- ** Filtering + sfilter, + sifilter, + sfilterM, + sifilterM, + -- , uniq -- sunique? + smapMaybe, + smapMaybeM, + scatMaybes, + simapMaybe, + simapMaybeM, + -- , stakeWhile + -- , sdropWhile + -- -- ** Partitioning + -- , partition + -- , unstablePartition + -- , partitionWith + -- , span + -- , break + -- -- ** Searching + -- , elem + -- , notElem + -- , find + findIndex, + -- , findIndices + -- , elemIndex + -- , elemIndices + + -- * Folding + sfoldl, + sfoldlM, + sfoldlM_, + sifoldl, + sifoldlM, + sifoldlM_, + sfoldl1', + sfoldl1M, + sfoldl1M_, + + -- ** Specialized folds + sor, + sand, + sall, + sany, + ssum, + sproduct, + smaximum', + smaximumM, + -- , maximumBy + sminimum', + sminimumM, + -- , minimumBy + -- , minIndex + -- , minIndexBy + -- , maxIndex + -- , maxIndexBy + + -- ** Scanning + sprescanl, + spostscanl, + spostscanlAcc, + sscanl, + sscanl1, + -- sprescanr, + -- spostscanr, + -- sscanr, + -- sscanr1, + + -- * Conversions + + -- ** Lists + stoList, + fromList, + sfromList, + sfromListN, + + -- * Computation + compute, + computeS, + computeIO, + computePrimM, + computeAs, + computeProxy, + computeSource, + computeWithStride, + computeWithStrideAs, + clone, + convert, + convertAs, + convertProxy, + + -- ** Re-exports + module Data.Massiv.Core, + module Data.Massiv.Array.Delayed, + module Data.Massiv.Array.Manifest, + module Data.Massiv.Array.Mutable, +) where + +import Control.Monad hiding (filterM, replicateM) +import Data.Coerce +import Data.Massiv.Array.Delayed +import Data.Massiv.Array.Delayed.Pull +import Data.Massiv.Array.Delayed.Push +import Data.Massiv.Array.Delayed.Stream +import Data.Massiv.Array.Manifest +import Data.Massiv.Array.Manifest.Internal +import Data.Massiv.Array.Manifest.List (fromList) +import Data.Massiv.Array.Mutable +import Data.Massiv.Array.Ops.Construct +import qualified Data.Massiv.Array.Ops.Construct as A (replicate) +import Data.Massiv.Core +import Data.Massiv.Core.Common +import qualified Data.Massiv.Vector.Stream as S +import Data.Massiv.Vector.Unsafe +import Data.Maybe +import Prelude hiding ( + drop, + dropWhile, + init, + length, + null, + replicate, + splitAt, + tail, + take, + takeWhile, + ) + +-- ========= -- +-- Accessors -- +-- ========= -- + +------------------------ +-- Length information -- +------------------------ + +-- | /O(1)/ - Get the length of a `Stream` array, but only if it is known exactly in +-- constant time without looking at any of the elements in the array. +-- +-- /Related/: `maxLinearSize`, `size`, `elemsCount` and `totalElem` +-- +-- ==== __Examples__ +-- +-- >>> slength $ sfromList [] +-- Nothing +-- >>> slength $ sreplicate 5 () +-- Just (Sz1 5) +-- >>> slength $ makeArrayLinearR D Seq (Sz1 5) id +-- Just (Sz1 5) +-- >>> slength $ sunfoldr (\x -> Just (x, x)) (0 :: Int) +-- Nothing +-- >>> slength $ sunfoldrN 10 (\x -> Just (x, x)) (0 :: Int) +-- Nothing +-- >>> slength $ sunfoldrExactN 10 (\x -> (x, x)) (0 :: Int) +-- Just (Sz1 10) +-- +-- /__Similar__/: +-- +-- [@Data.Foldable.`Data.Foldable.length`@] For some data structures, like a list for +-- example, it is an /O(n)/ operation, because there is a need to evaluate the full spine +-- and possibly even the elements in order to get the full length. With `Stream` vectors +-- that is not always the case. +-- +-- [@Data.Vector.Generic.`Data.Vector.Generic.length`@] In the vector package this +-- function will always break fusion, unless it is the only operation that is applied to +-- the vector. +-- +-- @since 0.5.0 +slength + :: forall r ix e + . Stream r ix e + => Array r ix e + -> Maybe Sz1 +slength v = + case stepsSize (toStream v) of + LengthExact sz -> Just sz + _ -> Nothing +{-# INLINE slength #-} + +-------------- +-- Indexing -- +-------------- + +-- | /O(1)/ - Get the first element of a `Source` vector. Throws an error on empty. +-- +-- /Related/: 'shead'', `headM`, `sheadM`, `unconsM`. +-- +-- ==== __Examples__ +-- +-- >>> head' (Ix1 10 ..: 10000000000000) +-- 10 +-- +-- /__Similar__/: +-- +-- [@Data.List.`Data.List.head`@] Also constant time and partial. Fusion is broken if +-- there other consumers of the list. +-- +-- [@Data.Vector.Generic.`Data.Vector.Generic.head`@] Also constant time and partial. Will +-- cause materialization of the full vector if any other function is applied to the vector. +-- +-- @since 0.5.0 +head' + :: forall r e + . (HasCallStack, Source r e) + => Vector r e + -> e +head' = throwEither . headM +{-# INLINE head' #-} + +-- | /O(1)/ - Get the first element of a `Source` vector. +-- +-- /Related/: 'head'', 'shead'', `sheadM`, `unconsM`. +-- +-- /__Throws Exceptions__/: `SizeEmptyException` when array is empty +-- +-- ==== __Examples__ +-- +-- >>> headM (Ix1 10 ..: 10000000000000) +-- 10 +-- >>> headM (Ix1 10 ..: 10000000000000) :: Maybe Int +-- Just 10 +-- >>> headM (empty :: Array D Ix1 Int) :: Maybe Int +-- Nothing +-- >>> either show (const "") $ headM (Ix1 10 ..: 10) +-- "SizeEmptyException: (Sz1 0) corresponds to an empty array" +-- +-- /__Similar__/: +-- +-- [@Data.Maybe.`Data.Maybe.listToMaybe`@] It also a safe way to get the head of the list, +-- except it is restricted to `Maybe` +-- +-- @since 0.5.0 +headM + :: forall r e m + . (Source r e, MonadThrow m) + => Vector r e + -> m e +headM v + | elemsCount v == 0 = throwM $ SizeEmptyException (size v) + | otherwise = pure $ unsafeLinearIndex v 0 +{-# INLINE headM #-} + +-- | /O(1)/ - Get the first element of a `Stream` vector. Throws an error on empty. +-- +-- /Related/: 'head'', `headM`, `sheadM`, `unconsM`. +-- +-- ==== __Examples__ +-- +-- >>> shead' $ sunfoldr (\x -> Just (x, x)) (0 :: Int) +-- 0 +-- >>> shead' (Ix1 3 ... 5) +-- 3 +-- +-- @since 0.5.0 +shead' + :: forall r e + . (HasCallStack, Stream r Ix1 e) + => Vector r e + -> e +shead' = throwEither . sheadM +{-# INLINE shead' #-} + +-- | /O(1)/ - Get the first element of a `Stream` vector. +-- +-- /Related/: 'head'', 'shead'', `headM`, `unconsM`. +-- +-- /__Throws Exceptions__/: `SizeEmptyException` +-- +-- ==== __Examples__ +-- +-- >>> maybe 101 id $ sheadM (empty :: Vector D Int) +-- 101 +-- >>> maybe 101 id $ sheadM (singleton 202 :: Vector D Int) +-- 202 +-- >>> sheadM $ sunfoldr (\x -> Just (x, x)) (0 :: Int) +-- 0 +-- >>> x <- sheadM $ sunfoldr (\_ -> Nothing) (0 :: Int) +-- *** Exception: SizeEmptyException: (Sz1 0) corresponds to an empty array +-- +-- @since 0.5.0 +sheadM + :: forall r e m + . (Stream r Ix1 e, MonadThrow m) + => Vector r e + -> m e +sheadM v = + case S.unId (S.headMaybe (toStream v)) of + Nothing -> throwM $ SizeEmptyException (zeroSz :: Sz1) + Just e -> pure e +{-# INLINE sheadM #-} + +-- | /O(1)/ - Take one element off of the `Source` vector from the left side, as well as +-- the remaining part of the vector in delayed `D` representation. +-- +-- /Related/: 'head'', 'shead'', `headM`, `sheadM`, `cons` +-- +-- /__Throws Exceptions__/: `SizeEmptyException` +-- +-- ==== __Examples__ +-- +-- >>> unconsM (fromList Seq [1,2,3] :: Array P Ix1 Int) +-- (1,Array P Seq (Sz1 2) +-- [ 2, 3 ]) +-- +-- /__Similar__/: +-- +-- [@Data.List.`Data.List.uncons`@] Same concept, except it is restricted to `Maybe` instead of +-- the more general `MonadThrow` +-- +-- @since 0.3.0 +unconsM + :: forall r e m + . (MonadThrow m, Source r e) + => Vector r e + -> m (e, Vector r e) +unconsM arr + | 0 == totalElem sz = throwM $ SizeEmptyException sz + | otherwise = pure (unsafeLinearIndex arr 0, unsafeLinearSlice 1 (SafeSz (unSz sz - 1)) arr) + where + !sz = size arr +{-# INLINE unconsM #-} + +-- | /O(1)/ - Take one element off of the vector from the right side, as well as the +-- remaining part of the vector. +-- +-- /Related/: 'last'', `lastM`, `snoc` +-- +-- /__Throws Exceptions__/: `SizeEmptyException` +-- +-- ==== __Examples__ +-- +-- >>> unsnocM (fromList Seq [1,2,3] :: Array P Ix1 Int) +-- (Array P Seq (Sz1 2) +-- [ 1, 2 ],3) +-- +-- @since 0.3.0 +unsnocM + :: forall r e m + . (MonadThrow m, Source r e) + => Vector r e + -> m (Vector r e, e) +unsnocM arr + | 0 == totalElem sz = throwM $ SizeEmptyException sz + | otherwise = pure (unsafeLinearSlice 0 (SafeSz k) arr, unsafeLinearIndex arr k) + where + !sz = size arr + !k = unSz sz - 1 +{-# INLINE unsnocM #-} + +-- | /O(1)/ - Get the last element of a `Source` vector. Throws an error on empty. +-- +-- /Related/: `lastM`, `unsnocM` +-- +-- ==== __Examples__ +-- +-- >>> last' (Ix1 10 ... 10000000000000) +-- 10000000000000 +-- +-- /__Similar__/: +-- +-- [@Data.List.`Data.List.last`@] Also partial, but it has /O(n)/ complexity. Fusion is +-- broken if there other consumers of the list. +-- +-- [@Data.Vector.Generic.`Data.Vector.Generic.last`@] Also constant time and partial. Will +-- cause materialization of the full vector if any other function is applied to the vector. +-- +-- @since 0.5.0 +last' :: forall r e. (HasCallStack, Source r e) => Vector r e -> e +last' = throwEither . lastM +{-# INLINE last' #-} + +-- | /O(1)/ - Get the last element of a `Source` vector. +-- +-- /Related/: 'last'', `unsnocM` +-- +-- /__Throws Exceptions__/: `SizeEmptyException` +-- +-- ==== __Examples__ +-- +-- >>> lastM (Ix1 10 ... 10000000000000) +-- 10000000000000 +-- >>> lastM (Ix1 10 ... 10000000000000) :: Maybe Int +-- Just 10000000000000 +-- >>> either show (const "") $ lastM (fromList Seq [] :: Array P Ix1 Int) +-- "SizeEmptyException: (Sz1 0) corresponds to an empty array" +-- +-- @since 0.5.0 +lastM :: forall r e m. (Source r e, MonadThrow m) => Vector r e -> m e +lastM v + | k == 0 = throwM $ SizeEmptyException (size v) + | otherwise = pure $ unsafeLinearIndex v (k - 1) + where + k = unSz (size v) +{-# INLINE lastM #-} + +-- | /O(1)/ - Take a slice of a `Source` vector. Never fails, instead adjusts the indices. +-- +-- ==== __Examples__ +-- +-- >>> slice 10 5 (Ix1 0 ... 10000000000000) +-- Array D Seq (Sz1 5) +-- [ 10, 11, 12, 13, 14 ] +-- >>> slice (-10) 5 (Ix1 0 ... 10000000000000) +-- Array D Seq (Sz1 5) +-- [ 0, 1, 2, 3, 4 ] +-- >>> slice 9999999999998 50 (Ix1 0 ... 10000000000000) +-- Array D Seq (Sz1 3) +-- [ 9999999999998, 9999999999999, 10000000000000 ] +-- +-- @since 0.5.0 +slice :: forall r e. Source r e => Ix1 -> Sz1 -> Vector r e -> Vector r e +slice !i (Sz k) v = unsafeLinearSlice i' newSz v + where + !i' = min n (max 0 i) + !newSz = SafeSz (min (n - i') k) + Sz n = size v +{-# INLINE slice #-} + +-- | /O(1)/ - Take a slice of a `Source` vector. Throws an error on incorrect indices. +-- +-- ==== __Examples__ +-- +-- >>> slice' 10 5 (Ix1 0 ... 100) +-- Array D Seq (Sz1 5) +-- [ 10, 11, 12, 13, 14 ] +-- >>> slice' 9999999999998 3 (Ix1 0 ... 10000000000000) +-- Array D Seq (Sz1 3) +-- [ 9999999999998, 9999999999999, 10000000000000 ] +-- +-- @since 0.5.0 +slice' :: forall r e. (HasCallStack, Source r e) => Ix1 -> Sz1 -> Vector r e -> Vector r e +slice' i k = throwEither . sliceM i k +{-# INLINE slice' #-} + +-- | /O(1)/ - Take a slice of a `Source` vector. Throws an error on incorrect indices. +-- +-- /__Throws Exceptions__/: `SizeSubregionException` +-- +-- ==== __Examples__ +-- +-- >>> sliceM 10 5 (Ix1 0 ... 100) +-- Array D Seq (Sz1 5) +-- [ 10, 11, 12, 13, 14 ] +-- >>> sliceM (-10) 5 (Ix1 0 ... 100) +-- *** Exception: SizeSubregionException: (Sz1 101) is to small for -10 (Sz1 5) +-- >>> sliceM 98 50 (Ix1 0 ... 100) +-- *** Exception: SizeSubregionException: (Sz1 101) is to small for 98 (Sz1 50) +-- >>> sliceM 9999999999998 3 (Ix1 0 ... 10000000000000) +-- Array D Seq (Sz1 3) +-- [ 9999999999998, 9999999999999, 10000000000000 ] +-- +-- @since 0.5.0 +sliceM + :: forall r e m + . (Source r e, MonadThrow m) + => Ix1 + -- ^ Starting index + -> Sz1 + -- ^ Number of elements to take from the Source vector + -> Vector r e + -- ^ Source vector to take a slice from + -> m (Vector r e) +sliceM i newSz@(Sz k) v + | i >= 0 && k <= n - i = pure $ unsafeLinearSlice i newSz v + | otherwise = throwM $ SizeSubregionException sz i newSz + where + sz@(Sz n) = size v +{-# INLINE sliceM #-} + +-- | Take a slice of a `Stream` vector. Never fails, instead adjusts the indices. +-- +-- ==== __Examples__ +-- +-- >>> sslice 10 5 (Ix1 0 ... 10000000000000) +-- Array DS Seq (Sz1 5) +-- [ 10, 11, 12, 13, 14 ] +-- >>> sslice 10 5 (sfromList [0 :: Int .. ]) +-- Array DS Seq (Sz1 5) +-- [ 10, 11, 12, 13, 14 ] +-- >>> sslice (-10) 5 (Ix1 0 ... 10000000000000) +-- Array DS Seq (Sz1 5) +-- [ 0, 1, 2, 3, 4 ] +-- +-- Unlike `slice` it has to iterate through each element until the staring index is reached, +-- therefore something like @sslice 9999999999998 50 (Ix1 0 ... 10000000000000)@ will not +-- be feasable. +-- +-- >>> import System.Timeout (timeout) +-- >>> let smallArr = sslice 9999999999998 50 (Ix1 0 ... 10000000000000) +-- >>> timeout 500000 (computeIO smallArr :: IO (Array P Ix1 Int)) +-- Nothing +-- +-- @since 0.5.0 +sslice + :: forall r e + . Stream r Ix1 e + => Ix1 + -- ^ Starting index + -> Sz1 + -- ^ Number of elements to take from the stream vector + -> Vector r e + -- ^ Stream vector to take a slice from + -> Vector DS e +sslice !i !k = fromSteps . S.slice i k . S.toStream +{-# INLINE sslice #-} + +-- | /O(1)/ - Get a vector without the last element. Never fails. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> A.init (0 ..: 10) +-- Array D Seq (Sz1 9) +-- [ 0, 1, 2, 3, 4, 5, 6, 7, 8 ] +-- >>> A.init (empty :: Array D Ix1 Int) +-- Array D Seq (Sz1 0) +-- [ ] +-- +-- @since 0.5.0 +init :: forall r e. Source r e => Vector r e -> Vector r e +init v = unsafeLinearSlice 0 (Sz (coerce (size v) - 1)) v +{-# INLINE init #-} + +-- | /O(1)/ - Get a vector without the last element. Throws an error on empty +-- +-- ==== __Examples__ +-- +-- >>> init' (0 ..: 10) +-- Array D Seq (Sz1 9) +-- [ 0, 1, 2, 3, 4, 5, 6, 7, 8 ] +-- +-- @since 0.5.0 +init' :: forall r e. (HasCallStack, Source r e) => Vector r e -> Vector r e +init' = throwEither . initM +{-# INLINE init' #-} + +-- | /O(1)/ - Get a vector without the last element. Throws an error on empty +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> initM (0 ..: 10) +-- Array D Seq (Sz1 9) +-- [ 0, 1, 2, 3, 4, 5, 6, 7, 8 ] +-- >>> maybe 0 A.sum $ initM (0 ..: 10) +-- 36 +-- >>> maybe 0 A.sum $ initM (empty :: Array D Ix1 Int) +-- 0 +-- +-- @since 0.5.0 +initM :: forall r e m. (Source r e, MonadThrow m) => Vector r e -> m (Vector r e) +initM v = do + when (elemsCount v == 0) $ throwM $ SizeEmptyException $ size v + pure $ unsafeInit v +{-# INLINE initM #-} + +-- | /O(1)/ - Get a vector without the first element. Never fails +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> A.tail (0 ..: 10) +-- Array D Seq (Sz1 9) +-- [ 1, 2, 3, 4, 5, 6, 7, 8, 9 ] +-- >>> A.tail (empty :: Array D Ix1 Int) +-- Array D Seq (Sz1 0) +-- [ ] +-- +-- @since 0.5.0 +tail :: forall r e. Source r e => Vector r e -> Vector r e +tail = drop oneSz +{-# INLINE tail #-} + +-- | /O(1)/ - Get a vector without the first element. Throws an error on empty +-- +-- ==== __Examples__ +-- +-- λ> tail' (0 ..: 10) +-- Array D Seq (Sz1 9) +-- [ 1, 2, 3, 4, 5, 6, 7, 8, 9 ] +-- λ> tail' (empty :: Array D Ix1 Int) +-- Array D *** Exception: SizeEmptyException: (Sz1 0) corresponds to an empty array +-- +-- @since 0.5.0 +tail' :: forall r e. (HasCallStack, Source r e) => Vector r e -> Vector r e +tail' = throwEither . tailM +{-# INLINE tail' #-} + +-- | /O(1)/ - Get the vector without the first element. Throws an error on empty +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> tailM (0 ..: 10) +-- Array D Seq (Sz1 9) +-- [ 1, 2, 3, 4, 5, 6, 7, 8, 9 ] +-- >>> maybe 0 A.sum $ tailM (0 ..: 10) +-- 45 +-- >>> maybe 0 A.sum $ tailM (empty :: Array D Ix1 Int) +-- 0 +-- +-- @since 0.5.0 +tailM :: forall r e m. (Source r e, MonadThrow m) => Vector r e -> m (Vector r e) +tailM v = do + when (elemsCount v == 0) $ throwM $ SizeEmptyException $ size v + pure $ unsafeTail v +{-# INLINE tailM #-} + +-- | /O(1)/ - Take first @n@ elements from a vector. This function never fails and has +-- similar semantics as the `Data.List.take` for lists. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> A.take 5 (0 ..: 10) +-- Array D Seq (Sz1 5) +-- [ 0, 1, 2, 3, 4 ] +-- >>> A.take 0 (0 ..: 10) +-- Array D Seq (Sz1 0) +-- [ ] +-- >>> A.take 100 (0 ..: 10) +-- Array D Seq (Sz1 10) +-- [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9 ] +-- +-- @since 0.5.0 +take :: Source r e => Sz1 -> Vector r e -> Vector r e +take k = fst . sliceAt k +{-# INLINE take #-} + +-- | Slice a manifest vector in such a way that it will contain all initial elements that +-- satisfy the supplied predicate. +-- +-- @since 0.5.5 +takeWhile :: Manifest r e => (e -> Bool) -> Vector r e -> Vector r e +takeWhile f v = take (go 0) v + where + !k = elemsCount v + go !i + | i < k && f (unsafeLinearIndex v i) = go (i + 1) + | otherwise = SafeSz i +{-# INLINE takeWhile #-} + +-- | /O(1)/ - Get the vector with the first @n@ elements. Throws an error size is less +-- than @n@. +-- +-- ==== __Examples__ +-- +-- >>> take' 0 (0 ..: 0) +-- Array D Seq (Sz1 0) +-- [ ] +-- >>> take' 5 (0 ..: 10) +-- Array D Seq (Sz1 5) +-- [ 0, 1, 2, 3, 4 ] +-- +-- @since 0.5.0 +take' :: forall r e. (HasCallStack, Source r e) => Sz1 -> Vector r e -> Vector r e +take' k = throwEither . takeM k +{-# INLINE take' #-} + +-- | /O(1)/ - Get the vector with the first @n@ elements. Throws an error size is less than @n@ +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> takeM 5 (0 ..: 10) +-- Array D Seq (Sz1 5) +-- [ 0, 1, 2, 3, 4 ] +-- >>> maybe 0 A.sum $ takeM 5 (0 ..: 10) +-- 10 +-- >>> maybe (-1) A.sum $ takeM 15 (0 ..: 10) +-- -1 +-- >>> takeM 15 (0 ..: 10) +-- *** Exception: SizeSubregionException: (Sz1 10) is to small for 0 (Sz1 15) +-- +-- @since 0.5.0 +takeM :: forall r e m. (Source r e, MonadThrow m) => Sz1 -> Vector r e -> m (Vector r e) +takeM k v = do + let sz = size v + when (k > sz) $ throwM $ SizeSubregionException sz 0 k + pure $ unsafeTake k v +{-# INLINE takeM #-} + +-- | /O(1)/ - Create a `Stream` vector with the first @n@ elements. Never fails +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +stake :: forall r e. Stream r Ix1 e => Sz1 -> Vector r e -> Vector DS e +stake n = fromSteps . S.take n . S.toStream +{-# INLINE stake #-} + +-- | /O(1)/ - Drop @n@ elements from a vector. This function never fails and has +-- similar semantics as the `Data.List.drop` for lists. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> v = makeVectorR D Seq 10 id +-- >>> v +-- Array D Seq (Sz1 10) +-- [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9 ] +-- >>> A.drop 5 v +-- Array D Seq (Sz1 5) +-- [ 5, 6, 7, 8, 9 ] +-- >>> A.drop 25 v +-- Array D Seq (Sz1 0) +-- [ ] +-- +-- @since 0.5.0 +drop :: forall r e. Source r e => Sz1 -> Vector r e -> Vector r e +drop k = snd . sliceAt k +{-# INLINE drop #-} + +-- | Slice a manifest vector in such a way that it will not contain all initial elements +-- that satisfy the supplied predicate. +-- +-- @since 0.5.5 +dropWhile :: forall r e. Manifest r e => (e -> Bool) -> Vector r e -> Vector r e +dropWhile f v = drop (go 0) v + where + !k = elemsCount v + go !i + | i < k && f (unsafeLinearIndex v i) = go (i + 1) + | otherwise = SafeSz i +{-# INLINE dropWhile #-} + +-- | Keep all but the first @n@ elements from the delayed stream vector. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sdrop :: forall r e. Stream r Ix1 e => Sz1 -> Vector r e -> Vector DS e +sdrop n = fromSteps . S.drop n . S.toStream +{-# INLINE sdrop #-} + +-- | /O(1)/ - Drop @n@ elements from a vector. Unlike `drop`, this function will +-- produce an error when supplied number of elements to drop is larger than size +-- of the supplied vector +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +drop' :: forall r e. (HasCallStack, Source r e) => Sz1 -> Vector r e -> Vector r e +drop' k = throwEither . dropM k +{-# INLINE drop' #-} + +-- | +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +dropM :: forall r e m. (Source r e, MonadThrow m) => Sz1 -> Vector r e -> m (Vector r e) +dropM k@(Sz d) v = do + let sz@(Sz n) = size v + when (k > sz) $ throwM $ SizeSubregionException sz d (SafeSz (n - d)) + pure $ unsafeLinearSlice d (SafeSz (n - d)) v +{-# INLINE dropM #-} + +-- | Same as 'sliceAt'', except it never fails. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sliceAt :: forall r e. Source r e => Sz1 -> Vector r e -> (Vector r e, Vector r e) +sliceAt (Sz k) v = (unsafeTake d v, unsafeDrop d v) + where + !n = coerce (size v) + !d = SafeSz (min k n) +{-# INLINE sliceAt #-} + +-- | Same as 'Data.Massiv.Array.splitAt'', except for a flat vector. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sliceAt' :: (HasCallStack, Source r e) => Sz1 -> Vector r e -> (Vector r e, Vector r e) +sliceAt' k = throwEither . sliceAtM k +{-# INLINE sliceAt' #-} + +-- | Same as `Data.Massiv.Array.splitAtM`, except for a flat vector. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sliceAtM :: forall r e m. (Source r e, MonadThrow m) => Sz1 -> Vector r e -> m (Vector r e, Vector r e) +sliceAtM k v = do + l <- takeM k v + pure (l, unsafeDrop k v) +{-# INLINE sliceAtM #-} + +-- | Create an empty delayed stream vector +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sempty :: Vector DS e +sempty = DSArray S.empty +{-# INLINE sempty #-} + +-- | Create a delayed stream vector with a single element +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +ssingleton :: e -> Vector DS e +ssingleton = DSArray . S.singleton +{-# INLINE ssingleton #-} + +-- | /O(1)/ - Add an element to the vector from the left side +-- +-- @since 0.3.0 +cons :: forall r e. (Size r, Load r Ix1 e) => e -> Vector r e -> Vector DL e +cons e v = + let dv = toLoadArray v + load scheduler startAt uWrite uSet = + uWrite startAt e >> dlLoad dv scheduler (startAt + 1) uWrite uSet + {-# INLINE load #-} + in dv{dlSize = SafeSz (1 + unSz (dlSize dv)), dlLoad = load} +{-# INLINE cons #-} + +-- | /O(1)/ - Add an element to the vector from the right side +-- +-- @since 0.3.0 +snoc :: forall r e. (Size r, Load r Ix1 e) => Vector r e -> e -> Vector DL e +snoc v e = + let dv = toLoadArray v + !k = unSz (size dv) + load scheduler startAt uWrite uSet = + dlLoad dv scheduler startAt uWrite uSet >> uWrite (k + startAt) e + {-# INLINE load #-} + in dv{dlSize = SafeSz (1 + k), dlLoad = load} +{-# INLINE snoc #-} + +-- | Replicate the same element @n@ times +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sreplicate :: Sz1 -> e -> Vector DS e +sreplicate n = DSArray . S.replicate n +{-# INLINE sreplicate #-} + +-- | Create a delayed vector of length @n@ with a function that maps an index to an +-- element. Same as `makeLinearArray` +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +generate :: Comp -> Sz1 -> (Ix1 -> e) -> Vector D e +generate = makeArrayLinear +{-# INLINE generate #-} + +-- | Create a delayed stream vector of length @n@ with a function that maps an index to an +-- element. Same as `makeLinearArray` +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sgenerate :: Sz1 -> (Ix1 -> e) -> Vector DS e +sgenerate n = DSArray . S.generate n +{-# INLINE sgenerate #-} + +-- | Create a delayed stream vector of infinite length by repeatedly applying a function to the +-- initial value. +-- +-- ==== __Examples__ +-- +-- >>> stake 10 $ siterate succ 'a' +-- Array DS Seq (Sz1 10) +-- [ 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j' ] +-- +-- @since 0.5.2 +siterate :: (e -> e) -> e -> Vector DS e +siterate f = fromSteps . S.unfoldr (\a -> Just (a, f a)) +{-# INLINE siterate #-} + +-- | Create a delayed stream vector of length @n@ by repeatedly applying a function to the +-- initial value. +-- +-- ==== __Examples__ +-- +-- >>> siterateN 10 succ 'a' +-- Array DS Seq (Sz1 10) +-- [ 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j' ] +-- +-- @since 0.5.0 +siterateN :: Sz1 -> (e -> e) -> e -> Vector DS e +siterateN n f a = fromSteps $ S.iterateN n f a +{-# INLINE siterateN #-} + +-- | Create a vector by using the same monadic action @n@ times +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sreplicateM :: forall e m. Monad m => Sz1 -> m e -> m (Vector DS e) +sreplicateM n f = fromStepsM $ S.replicateM n f +{-# INLINE sreplicateM #-} + +-- | Create a delayed stream vector of length @n@ with a monadic action that from an index +-- generates an element. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sgenerateM :: forall e m. Monad m => Sz1 -> (Ix1 -> m e) -> m (Vector DS e) +sgenerateM n f = fromStepsM $ S.generateM n f +{-# INLINE sgenerateM #-} + +-- | Create a delayed stream vector of length @n@ by repeatedly apply a monadic action to +-- the initial value. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +siterateNM :: forall e m. Monad m => Sz1 -> (e -> m e) -> e -> m (Vector DS e) +siterateNM n f a = fromStepsM $ S.iterateNM n f a +{-# INLINE siterateNM #-} + +-- | Right unfolding function. Useful when it is unknown ahead of time how many +-- elements a vector will have. +-- +-- ====__Example__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> sunfoldr (\i -> if i < 9 then Just (i * i, i + 1) else Nothing) (0 :: Int) +-- Array DS Seq (Sz1 9) +-- [ 0, 1, 4, 9, 16, 25, 36, 49, 64 ] +-- +-- @since 0.5.0 +sunfoldr :: forall e s. (s -> Maybe (e, s)) -> s -> Vector DS e +sunfoldr f = DSArray . S.unfoldr f +{-# INLINE sunfoldr #-} + +-- | /O(n)/ - Right unfolding function with at most @n@ number of elements. +-- +-- ==== __Example__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> sunfoldrN 9 (\i -> Just (i*i, i + 1)) (0 :: Int) +-- Array DS Seq (Sz1 9) +-- [ 0, 1, 4, 9, 16, 25, 36, 49, 64 ] +-- +-- @since 0.5.0 +sunfoldrN + :: forall e s + . Sz1 + -- ^ @n@ - maximum number of elements that the vector will have + -> (s -> Maybe (e, s)) + -- ^ Unfolding function. Stops when `Nothing` is returned or maximum number of elements + -- is reached. + -> s + -- ^ Inititial element. + -> Vector DS e +sunfoldrN n f = DSArray . S.unfoldrN n f +{-# INLINE sunfoldrN #-} + +-- | /O(n)/ - Same as `sunfoldr`, but with monadic generating function. +-- +-- ==== __Examples__ +-- +-- >>> import Control.Monad (when, guard) +-- >>> sunfoldrM (\i -> when (i == 0) (Left "Zero denominator") >> Right (guard (i < 5) >> Just (100 `div` i, i + 1))) (-10 :: Int) +-- Left "Zero denominator" +-- >>> sunfoldrM (\i -> when (i == 0) (Left "Zero denominator") >> Right (guard (i < -5) >> Just (100 `div` i, i + 1))) (-10 :: Int) +-- Right (Array DS Seq (Sz1 5) +-- [ -10, -12, -13, -15, -17 ] +-- ) +-- +-- @since 0.5.0 +sunfoldrM :: forall e s m. Monad m => (s -> m (Maybe (e, s))) -> s -> m (Vector DS e) +sunfoldrM f = fromStepsM . S.unfoldrM f +{-# INLINE sunfoldrM #-} + +-- | /O(n)/ - Same as `sunfoldrN`, but with monadic generating function. +-- +-- ==== __Examples__ +-- +-- >>> import Control.Monad (guard) +-- >>> sunfoldrNM 6 (\i -> print i >> pure (guard (i < 5) >> Just (i * i, i + 1))) (10 :: Int) +-- 10 +-- Array DS Seq (Sz1 0) +-- [ ] +-- >>> sunfoldrNM 6 (\i -> print i >> pure (guard (i < 15) >> Just (i * i, i + 1))) (10 :: Int) +-- 10 +-- 11 +-- 12 +-- 13 +-- 14 +-- 15 +-- Array DS Seq (Sz1 5) +-- [ 100, 121, 144, 169, 196 ] +-- +-- +-- @since 0.5.0 +sunfoldrNM :: forall e s m. Monad m => Sz1 -> (s -> m (Maybe (e, s))) -> s -> m (Vector DS e) +sunfoldrNM (Sz n) f = fromStepsM . S.unfoldrNM n f +{-# INLINE sunfoldrNM #-} + +-- | /O(n)/ - Similar to `sunfoldrN`, except the length of the resulting vector will be exactly @n@ +-- +-- ==== __Examples__ +-- +-- >>> sunfoldrExactN 10 (\i -> (i * i, i + 1)) (10 :: Int) +-- Array DS Seq (Sz1 10) +-- [ 100, 121, 144, 169, 196, 225, 256, 289, 324, 361 ] +-- +-- @since 0.5.0 +sunfoldrExactN :: forall e s. Sz1 -> (s -> (e, s)) -> s -> Vector DS e +sunfoldrExactN n f = fromSteps . S.unfoldrExactN n f +{-# INLINE sunfoldrExactN #-} + +-- | /O(n)/ - Similar to `sunfoldrNM`, except the length of the resulting vector will be exactly @n@ +-- +-- ==== __Examples__ +-- +-- λ> sunfoldrExactNM 11 (\i -> pure (100 `div` i, i + 1)) (-10 :: Int) +-- Array DS *** Exception: divide by zero +-- λ> sunfoldrExactNM 11 (\i -> guard (i /= 0) >> Just (100 `div` i, i + 1)) (-10 :: Int) +-- Nothing +-- λ> sunfoldrExactNM 9 (\i -> guard (i /= 0) >> Just (100 `div` i, i + 1)) (-10 :: Int) +-- Just (Array DS Seq (Sz1 9) +-- [ -10, -12, -13, -15, -17, -20, -25, -34, -50 ] +-- ) +-- +-- @since 0.5.0 +sunfoldrExactNM :: forall e s m. Monad m => Sz1 -> (s -> m (e, s)) -> s -> m (Vector DS e) +sunfoldrExactNM n f = fromStepsM . S.unfoldrExactNM n f +{-# INLINE sunfoldrExactNM #-} + +-- | /O(n)/ - Enumerate from a starting number @x@ exactly @n@ times with a step @1@. +-- +-- /Related/: `senumFromStepN`, `enumFromN`, `enumFromStepN`, `rangeSize`, +-- `rangeStepSize`, `range`, 'rangeStep'' +-- +-- ==== __Examples__ +-- +-- >>> senumFromN (10 :: Int) 9 +-- Array DS Seq (Sz1 9) +-- [ 10, 11, 12, 13, 14, 15, 16, 17, 18 ] +-- +-- /__Similar__/: +-- +-- [@Prelude.`Prelude.enumFromTo`@] Very similar to @[x .. x + n - 1]@, except that +-- `senumFromN` is faster and it only works for `Num` and not for `Enum` elements +-- +-- [@Data.Vector.Generic.`Data.Vector.Generic.enumFromN`@] Uses exactly the same +-- implementation underneath. +-- +-- @since 0.5.0 +senumFromN + :: Num e + => e + -- ^ @x@ - starting number + -> Sz1 + -- ^ @n@ - length of resulting vector + -> Vector DS e +senumFromN x n = DSArray $ S.enumFromStepN x 1 n +{-# INLINE senumFromN #-} + +-- | /O(n)/ - Enumerate from a starting number @x@ exactly @n@ times with a custom step value @dx@ +-- +-- ==== __Examples__ +-- +-- >>> senumFromStepN (5 :: Int) 2 10 +-- Array DS Seq (Sz1 10) +-- [ 5, 7, 9, 11, 13, 15, 17, 19, 21, 23 ] +-- +-- __/Similar/__: +-- +-- [@Prelude.`Prelude.enumFrom`@] Just like @take n [x, x + dx ..]@, except that +-- `senumFromN` is faster and it only works for `Num` and not for `Enum` elements +-- +-- [@Data.Vector.Generic.`Data.Vector.Generic.enumFromStepN`@] Uses exactly the same +-- implementation underneath. +-- +-- @since 0.5.0 +senumFromStepN + :: Num e + => e + -- ^ @x@ - starting number + -> e + -- ^ @dx@ - Step + -> Sz1 + -- ^ @n@ - length of resulting vector + -> Vector DS e +senumFromStepN x step n = DSArray $ S.enumFromStepN x step n +{-# INLINE senumFromStepN #-} + +-- | Append two vectors together +-- +-- /Related/: `appendM`, `appendOuterM`, +-- +-- ==== __Examples__ +-- +-- λ> sappend (1 ..: 6) (senumFromStepN 6 (-1) 6) +-- Array DS Seq (Sz1 11) +-- [ 1, 2, 3, 4, 5, 6, 5, 4, 3, 2, 1 ] +-- +-- __/Similar/__: +-- +-- [@Data.Semigroup.`Data.Semigroup.<>`@] `DS` and `DL` arrays have instances for +-- `Semigroup`, so they will work in a similar fashion. `sappend` differs in that it accepts +-- `Stream` arrays with possibly different representations. +-- +-- [@Data.List.`Data.List.++`@] Same operation, but for lists. +-- +-- [@Data.Vector.Generic.`Data.Vector.Generic.++`@] Uses exactly the same implementation +-- underneath as `sappend`, except that it cannot append two vectors with different +-- memory representations. +-- +-- @since 0.5.0 +sappend + :: forall r1 r2 e + . (Stream r1 Ix1 e, Stream r2 Ix1 e) + => Vector r1 e + -> Vector r2 e + -> Vector DS e +sappend a1 a2 = fromSteps (toStream a1 `S.append` toStream a2) +{-# INLINE sappend #-} + +-- | Concat vectors together +-- +-- /Related/: `concatM`, `concatOuterM`, +-- +-- ==== __Examples__ +-- +-- >>> sconcat [2 ... 6, empty, singleton 1, generate Seq 5 id] +-- Array DS Seq (Sz1 11) +-- [ 2, 3, 4, 5, 6, 1, 0, 1, 2, 3, 4 ] +-- >>> sconcat [senumFromN 2 5, sempty, ssingleton 1, sgenerate 5 id] +-- Array DS Seq (Sz1 11) +-- [ 2, 3, 4, 5, 6, 1, 0, 1, 2, 3, 4 ] +-- +-- __/Similar/__: +-- +-- [@Data.Monoid.`Data.Monoid.mconcat`@] `DS` and `DL` arrays have instances for `Monoid`, so +-- they will work in a similar fashion. `sconcat` differs in that it accepts `Stream` +-- arrays of other representations. +-- +-- [@Data.List.`Data.List.concat`@] Same operation, but for lists. +-- +-- [@Data.Vector.Generic.`Data.Vector.Generic.concat`@] Uses exactly the same +-- implementation underneath as `sconcat`. +-- +-- @since 0.5.0 +sconcat :: forall r e. Stream r Ix1 e => [Vector r e] -> Vector DS e +sconcat = DSArray . foldMap toStream +{-# INLINE sconcat #-} + +-- | Convert a list to a delayed stream vector +-- +-- /Related/: `fromList`, `fromListN`, `sfromListN` +-- +-- ==== __Examples__ +-- +-- >>> sfromList ([] :: [Int]) +-- Array DS Seq (Sz1 0) +-- [ ] +-- >>> sfromList ([1,2,3] :: [Int]) +-- Array DS Seq (Sz1 3) +-- [ 1, 2, 3 ] +-- +-- @since 0.5.0 +sfromList :: [e] -> Vector DS e +sfromList = fromSteps . S.fromList +{-# INLINE sfromList #-} + +-- | Convert a list to a delayed stream vector. Length of the resulting vector will be at +-- most @n@. This version isn't really more efficient then `sfromList`, but there is +-- `Data.Massiv.Array.Unsafe.unsafeFromListN` +-- +-- /Related/: `fromList`, `fromListN`, `sfromList` +-- +-- ==== __Examples__ +-- +-- >>> sfromListN 10 [1 :: Int ..] +-- Array DS Seq (Sz1 10) +-- [ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 ] +-- >>> sfromListN 10 [1 :: Int .. 5] +-- Array DS Seq (Sz1 5) +-- [ 1, 2, 3, 4, 5 ] +-- +-- @since 0.5.1 +sfromListN :: Sz1 -> [e] -> Vector DS e +sfromListN (Sz n) = fromSteps . S.fromListN n +{-# INLINE sfromListN #-} + +-- | Convert an array to a list by the means of a delayed stream vector. +-- +-- /Related/: `toList` +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +stoList :: forall r ix e. Stream r ix e => Array r ix e -> [e] +stoList = S.toList . toStream +{-# INLINE stoList #-} + +-- | Sequentially filter out elements from the array according to the supplied predicate. +-- +-- ==== __Example__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> arr = makeArrayR D Seq (Sz2 3 4) fromIx2 +-- >>> arr +-- Array D Seq (Sz (3 :. 4)) +-- [ [ (0,0), (0,1), (0,2), (0,3) ] +-- , [ (1,0), (1,1), (1,2), (1,3) ] +-- , [ (2,0), (2,1), (2,2), (2,3) ] +-- ] +-- >>> sfilter (even . fst) arr +-- Array DS Seq (Sz1 8) +-- [ (0,0), (0,1), (0,2), (0,3), (2,0), (2,1), (2,2), (2,3) ] +-- +-- @since 0.5.0 +sfilter :: forall r ix e. S.Stream r ix e => (e -> Bool) -> Array r ix e -> Vector DS e +sfilter f = DSArray . S.filter f . S.toStream +{-# INLINE sfilter #-} + +-- | Similar to `sfilter`, but filter with an index aware function. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sifilter :: forall r ix e. Stream r ix e => (ix -> e -> Bool) -> Array r ix e -> Vector DS e +sifilter f = + simapMaybe $ \ix e -> + if f ix e + then Just e + else Nothing +{-# INLINE sifilter #-} + +-- | Sequentially filter out elements from the array according to the supplied applicative predicate. +-- +-- ==== __Example__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> arr = makeArrayR D Seq (Sz2 3 4) fromIx2 +-- >>> arr +-- Array D Seq (Sz (3 :. 4)) +-- [ [ (0,0), (0,1), (0,2), (0,3) ] +-- , [ (1,0), (1,1), (1,2), (1,3) ] +-- , [ (2,0), (2,1), (2,2), (2,3) ] +-- ] +-- >>> sfilterM (Just . odd . fst) arr +-- Just (Array DS Seq (Sz1 4) +-- [ (1,0), (1,1), (1,2), (1,3) ] +-- ) +-- >>> sfilterM (\ix@(_, j) -> print ix >> return (even j)) arr +-- (0,0) +-- (0,1) +-- (0,2) +-- (0,3) +-- (1,0) +-- (1,1) +-- (1,2) +-- (1,3) +-- (2,0) +-- (2,1) +-- (2,2) +-- (2,3) +-- Array DS Seq (Sz1 6) +-- [ (0,0), (0,2), (1,0), (1,2), (2,0), (2,2) ] +-- +-- @since 0.5.0 +sfilterM + :: forall r ix e f + . (S.Stream r ix e, Applicative f) + => (e -> f Bool) + -> Array r ix e + -> f (Vector DS e) +sfilterM f arr = DSArray <$> S.filterA f (S.toStream arr) +{-# INLINE sfilterM #-} + +-- | Similar to `filterM`, but filter with an index aware function. +-- +-- Corresponds to: @`filterM` (uncurry f) . `simap` (,)@ +-- +-- @since 0.5.0 +sifilterM + :: forall r ix e f + . (Stream r ix e, Applicative f) + => (ix -> e -> f Bool) + -> Array r ix e + -> f (Vector DS e) +sifilterM f = + simapMaybeM $ \ix e -> + ( \p -> + if p + then Just e + else Nothing + ) + <$> f ix e +{-# INLINE sifilterM #-} + +-- | Apply a function to each element of the array, while discarding `Nothing` and +-- keeping the `Maybe` result. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +smapMaybe :: forall r ix a b. S.Stream r ix a => (a -> Maybe b) -> Array r ix a -> Vector DS b +smapMaybe f = DSArray . S.mapMaybe f . S.toStream +{-# INLINE smapMaybe #-} + +-- | Similar to `smapMaybe`, but map with an index aware function. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +simapMaybe + :: forall r ix a b + . Stream r ix a + => (ix -> a -> Maybe b) + -> Array r ix a + -> Vector DS b +simapMaybe f = DSArray . S.mapMaybe (uncurry f) . toStreamIx +{-# INLINE simapMaybe #-} + +-- | Similar to `smapMaybeM`, but map with an index aware function. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +simapMaybeM + :: forall r ix a b f + . (Stream r ix a, Applicative f) + => (ix -> a -> f (Maybe b)) + -> Array r ix a + -> f (Vector DS b) +simapMaybeM f = fmap DSArray . S.mapMaybeA (uncurry f) . toStreamIx +{-# INLINE simapMaybeM #-} + +-- | Keep all `Maybe`s and discard the `Nothing`s. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +scatMaybes :: forall r ix a. S.Stream r ix (Maybe a) => Array r ix (Maybe a) -> Vector DS a +scatMaybes = smapMaybe id +{-# INLINE scatMaybes #-} + +-- | Similar to `smapMaybe`, but with the `Applicative` function. +-- +-- Similar to @mapMaybe id <$> mapM f arr@ +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +smapMaybeM + :: forall r ix a b f + . (S.Stream r ix a, Applicative f) + => (a -> f (Maybe b)) + -> Array r ix a + -> f (Vector DS b) +smapMaybeM f = fmap DSArray . S.mapMaybeA f . S.toStream +{-# INLINE smapMaybeM #-} + +-- | Map a function over a stream vector +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +smap + :: forall r ix a b + . S.Stream r ix a + => (a -> b) + -> Array r ix a + -> Vector DS b +smap f = fromSteps . S.map f . S.toStream +{-# INLINE smap #-} + +-- | Map an index aware function over a stream vector +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +simap + :: forall r ix a b + . S.Stream r ix a + => (ix -> a -> b) + -> Array r ix a + -> Vector DS b +simap f = fromSteps . S.map (uncurry f) . S.toStreamIx +{-# INLINE simap #-} + +-- | Traverse a stream vector with an applicative function. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +straverse + :: forall r ix a b f + . (S.Stream r ix a, Applicative f) + => (a -> f b) + -> Array r ix a + -> f (Vector DS b) +straverse f = fmap fromSteps . S.traverse f . S.toStream +{-# INLINE straverse #-} + +-- | Traverse a stream vector with an index aware applicative function. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sitraverse + :: forall r ix a b f + . (S.Stream r ix a, Applicative f) + => (ix -> a -> f b) + -> Array r ix a + -> f (Vector DS b) +sitraverse f = fmap fromSteps . S.traverse (uncurry f) . S.toStreamIx +{-# INLINE sitraverse #-} + +-- | Traverse a stream vector with a monadic function. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +smapM + :: forall r ix a b m + . (S.Stream r ix a, Monad m) + => (a -> m b) + -> Array r ix a + -> m (Vector DS b) +smapM f = fromStepsM . S.mapM f . S.transStepsId . S.toStream +{-# INLINE smapM #-} + +-- | Traverse a stream vector with a monadic index aware function. +-- +-- Corresponds to: @mapM (uncurry f) . imap (,) v@ +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +simapM + :: forall r ix a b m + . (S.Stream r ix a, Monad m) + => (ix -> a -> m b) + -> Array r ix a + -> m (Vector DS b) +simapM f = fromStepsM . S.mapM (uncurry f) . S.transStepsId . S.toStreamIx +{-# INLINE simapM #-} + +-- | Traverse a stream vector with a monadic function, while discarding the result +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +smapM_ + :: forall r ix a b m + . (S.Stream r ix a, Monad m) + => (a -> m b) + -> Array r ix a + -> m () +smapM_ f = S.mapM_ f . S.transStepsId . S.toStream +{-# INLINE smapM_ #-} + +-- | Traverse a stream vector with a monadic index aware function, while discarding the result +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +simapM_ + :: forall r ix a b m + . (S.Stream r ix a, Monad m) + => (ix -> a -> m b) + -> Array r ix a + -> m () +simapM_ f = S.mapM_ (uncurry f) . S.transStepsId . S.toStreamIx +{-# INLINE simapM_ #-} + +-- | Same as `smapM`, but with arguments flipped. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sforM + :: forall r ix a b m + . (S.Stream r ix a, Monad m) + => Array r ix a + -> (a -> m b) + -> m (Vector DS b) +sforM = flip smapM +{-# INLINE sforM #-} + +-- | Same as `simapM`, but with arguments flipped. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +siforM + :: forall r ix a b m + . (S.Stream r ix a, Monad m) + => Array r ix a + -> (ix -> a -> m b) + -> m (Vector DS b) +siforM = flip simapM +{-# INLINE siforM #-} + +-- | Same as `smapM_`, but with arguments flipped. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sforM_ :: (S.Stream r ix a, Monad m) => Array r ix a -> (a -> m b) -> m () +sforM_ = flip smapM_ +{-# INLINE sforM_ #-} + +-- | Same as `simapM_`, but with arguments flipped. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +siforM_ + :: forall r ix a b m + . (S.Stream r ix a, Monad m) + => Array r ix a + -> (ix -> a -> m b) + -> m () +siforM_ = flip simapM_ +{-# INLINE siforM_ #-} + +-- | Zip two vectors together into a vector. The length of a resulting vector will +-- be the smallest length of the supplied vectors. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +szip + :: forall ra rb a b + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b) + => Vector ra a + -> Vector rb b + -> Vector DS (a, b) +szip = szipWith (,) +{-# INLINE szip #-} + +-- | Zip three vectors together into a vector. The length of a resulting vector will +-- be the smallest length of the supplied vectors. +-- +-- @since 0.5.0 +szip3 + :: forall ra rb rc a b c + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c) + => Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector DS (a, b, c) +szip3 = szipWith3 (,,) +{-# INLINE szip3 #-} + +-- | Zip four vectors together into a vector. The length of a resulting vector will +-- be the smallest length of the supplied vectors. +-- +-- @since 0.5.0 +szip4 + :: forall ra rb rc rd a b c d + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, S.Stream rd Ix1 d) + => Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector DS (a, b, c, d) +szip4 = szipWith4 (,,,) +{-# INLINE szip4 #-} + +-- | Zip five vectors together into a vector. The length of a resulting vector will +-- be the smallest length of the supplied vectors. +-- +-- @since 0.5.0 +szip5 + :: forall ra rb rc rd re a b c d e + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, S.Stream rd Ix1 d, S.Stream re Ix1 e) + => Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector re e + -> Vector DS (a, b, c, d, e) +szip5 = szipWith5 (,,,,) +{-# INLINE szip5 #-} + +-- | Zip six vectors together into a vector. The length of a resulting vector will +-- be the smallest length of the supplied vectors. +-- +-- @since 0.5.0 +szip6 + :: forall ra rb rc rd re rf a b c d e f + . ( S.Stream ra Ix1 a + , S.Stream rb Ix1 b + , S.Stream rc Ix1 c + , S.Stream rd Ix1 d + , S.Stream re Ix1 e + , S.Stream rf Ix1 f + ) + => Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector re e + -> Vector rf f + -> Vector DS (a, b, c, d, e, f) +szip6 = szipWith6 (,,,,,) +{-# INLINE szip6 #-} + +-- | Zip two vectors together with a binary function into a vector. The length +-- of a resulting vector will be the smallest length of the supplied vectors. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +szipWith + :: forall ra rb a b c + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b) + => (a -> b -> c) + -> Vector ra a + -> Vector rb b + -> Vector DS c +szipWith f v1 v2 = fromSteps $ S.zipWith f (S.toStream v1) (S.toStream v2) +{-# INLINE szipWith #-} + +-- | Zip three vectors together with a ternary function into a vector. The length +-- of a resulting vector will be the smallest length of the supplied vectors. +-- +-- @since 0.5.0 +szipWith3 + :: forall ra rb rc a b c d + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c) + => (a -> b -> c -> d) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector DS d +szipWith3 f v1 v2 v3 = fromSteps $ S.zipWith3 f (S.toStream v1) (S.toStream v2) (S.toStream v3) +{-# INLINE szipWith3 #-} + +-- | Zip four vectors together with a quaternary function into a vector. The length +-- of a resulting vector will be the smallest length of the supplied vectors. +-- +-- @since 0.5.0 +szipWith4 + :: forall ra rb rc rd a b c d e + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, S.Stream rd Ix1 d) + => (a -> b -> c -> d -> e) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector DS e +szipWith4 f v1 v2 v3 v4 = + fromSteps $ S.zipWith4 f (S.toStream v1) (S.toStream v2) (S.toStream v3) (S.toStream v4) +{-# INLINE szipWith4 #-} + +-- | Zip five vectors together with a quinary function into a vector. The length +-- of a resulting vector will be the smallest length of the supplied vectors. +-- +-- @since 0.5.0 +szipWith5 + :: forall ra rb rc rd re a b c d e f + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, S.Stream rd Ix1 d, S.Stream re Ix1 e) + => (a -> b -> c -> d -> e -> f) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector re e + -> Vector DS f +szipWith5 f v1 v2 v3 v4 v5 = + fromSteps $ + S.zipWith5 f (S.toStream v1) (S.toStream v2) (S.toStream v3) (S.toStream v4) (S.toStream v5) +{-# INLINE szipWith5 #-} + +-- | Zip six vectors together with a senary function into a vector. The length +-- of a resulting vector will be the smallest length of the supplied vectors. +-- +-- @since 0.5.0 +szipWith6 + :: forall ra rb rc rd re rf a b c d e f g + . ( S.Stream ra Ix1 a + , S.Stream rb Ix1 b + , S.Stream rc Ix1 c + , S.Stream rd Ix1 d + , S.Stream re Ix1 e + , S.Stream rf Ix1 f + ) + => (a -> b -> c -> d -> e -> f -> g) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector re e + -> Vector rf f + -> Vector DS g +szipWith6 f v1 v2 v3 v4 v5 v6 = + fromSteps $ + S.zipWith6 + f + (S.toStream v1) + (S.toStream v2) + (S.toStream v3) + (S.toStream v4) + (S.toStream v5) + (S.toStream v6) +{-# INLINE szipWith6 #-} + +-- | Just like `szipWith`, zip two vectors together, but with an index aware +-- function. The length of a resulting vector will be the smallest length of the +-- supplied vectors. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sizipWith + :: forall ra rb a b c + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b) + => (Ix1 -> a -> b -> c) + -> Vector ra a + -> Vector rb b + -> Vector DS c +sizipWith f v1 v2 = fromSteps $ S.zipWith (uncurry f) (S.toStreamIx v1) (S.toStream v2) +{-# INLINE sizipWith #-} + +-- | Just like `szipWith3`, zip three vectors together, but with an index aware +-- function. The length of a resulting vector will be the smallest length of the +-- supplied vectors. +-- +-- @since 0.5.0 +sizipWith3 + :: forall ra rb rc a b c d + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c) + => (Ix1 -> a -> b -> c -> d) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector DS d +sizipWith3 f v1 v2 v3 = + fromSteps $ S.zipWith3 (uncurry f) (S.toStreamIx v1) (S.toStream v2) (S.toStream v3) +{-# INLINE sizipWith3 #-} + +-- | Just like `szipWith4`, zip four vectors together, but with an index aware +-- function. The length of a resulting vector will be the smallest +-- length of the supplied vectors. +-- +-- @since 0.5.0 +sizipWith4 + :: forall ra rb rc rd a b c d e + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, S.Stream rd Ix1 d) + => (Ix1 -> a -> b -> c -> d -> e) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector DS e +sizipWith4 f v1 v2 v3 v4 = + fromSteps $ + S.zipWith4 (uncurry f) (S.toStreamIx v1) (S.toStream v2) (S.toStream v3) (S.toStream v4) +{-# INLINE sizipWith4 #-} + +-- | Just like `szipWith5`, zip five vectors together, but with an index aware +-- function. The length of a resulting vector will be the smallest length of the +-- supplied vectors. +-- +-- @since 0.5.0 +sizipWith5 + :: forall ra rb rc rd re a b c d e f + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, S.Stream rd Ix1 d, S.Stream re Ix1 e) + => (Ix1 -> a -> b -> c -> d -> e -> f) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector re e + -> Vector DS f +sizipWith5 f v1 v2 v3 v4 v5 = + fromSteps $ + S.zipWith5 + (uncurry f) + (S.toStreamIx v1) + (S.toStream v2) + (S.toStream v3) + (S.toStream v4) + (S.toStream v5) +{-# INLINE sizipWith5 #-} + +-- | Just like `szipWith6`, zip six vectors together, but with an index aware +-- function. The length of a resulting vector will be the smallest length of the +-- supplied vectors. +-- +-- @since 0.5.0 +sizipWith6 + :: forall ra rb rc rd re rf a b c d e f g + . ( S.Stream ra Ix1 a + , S.Stream rb Ix1 b + , S.Stream rc Ix1 c + , S.Stream rd Ix1 d + , S.Stream re Ix1 e + , S.Stream rf Ix1 f + ) + => (Ix1 -> a -> b -> c -> d -> e -> f -> g) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector re e + -> Vector rf f + -> Vector DS g +sizipWith6 f v1 v2 v3 v4 v5 v6 = + fromSteps $ + S.zipWith6 + (uncurry f) + (S.toStreamIx v1) + (S.toStream v2) + (S.toStream v3) + (S.toStream v4) + (S.toStream v5) + (S.toStream v6) +{-# INLINE sizipWith6 #-} + +-- | Zip two vectors together with a binary monadic action into a vector. The +-- length of a resulting vector will be the smallest length of the supplied +-- vectors. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +szipWithM + :: forall ra rb a b c m + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, Monad m) + => (a -> b -> m c) + -> Vector ra a + -> Vector rb b + -> m (Vector DS c) +szipWithM f v1 v2 = fromStepsM $ S.zipWithM f (toStreamM v1) (toStreamM v2) +{-# INLINE szipWithM #-} + +-- | Zip three vectors together with a ternary monadic action into a vector. The +-- length of a resulting vector will be the smallest length of the supplied +-- vectors. +-- +-- @since 0.5.0 +szipWith3M + :: forall ra rb rc a b c d m + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, Monad m) + => (a -> b -> c -> m d) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> m (Vector DS d) +szipWith3M f v1 v2 v3 = fromStepsM $ S.zipWith3M f (toStreamM v1) (toStreamM v2) (toStreamM v3) +{-# INLINE szipWith3M #-} + +-- | Zip four vectors together with a quaternary monadic action into a vector. The +-- length of a resulting vector will be the smallest length of the supplied +-- vectors. +-- +-- @since 0.5.0 +szipWith4M + :: forall ra rb rc rd a b c d e m + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, S.Stream rd Ix1 d, Monad m) + => (a -> b -> c -> d -> m e) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> m (Vector DS e) +szipWith4M f v1 v2 v3 v4 = + fromStepsM $ S.zipWith4M f (toStreamM v1) (toStreamM v2) (toStreamM v3) (toStreamM v4) +{-# INLINE szipWith4M #-} + +-- | Zip five vectors together with a quinary monadic action into a vector. The +-- length of a resulting vector will be the smallest length of the supplied +-- vectors. +-- +-- @since 0.5.0 +szipWith5M + :: forall ra rb rc rd re a b c d e f m + . ( S.Stream ra Ix1 a + , S.Stream rb Ix1 b + , S.Stream rc Ix1 c + , S.Stream rd Ix1 d + , S.Stream re Ix1 e + , Monad m + ) + => (a -> b -> c -> d -> e -> m f) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector re e + -> m (Vector DS f) +szipWith5M f v1 v2 v3 v4 v5 = + fromStepsM $ + S.zipWith5M f (toStreamM v1) (toStreamM v2) (toStreamM v3) (toStreamM v4) (toStreamM v5) +{-# INLINE szipWith5M #-} + +-- | Zip six vectors together with a senary monadic action into a vector. The +-- length of a resulting vector will be the smallest length of the supplied +-- vectors. +-- +-- @since 0.5.0 +szipWith6M + :: forall ra rb rc rd re rf a b c d e f g m + . ( S.Stream ra Ix1 a + , S.Stream rb Ix1 b + , S.Stream rc Ix1 c + , S.Stream rd Ix1 d + , S.Stream re Ix1 e + , S.Stream rf Ix1 f + , Monad m + ) + => (a -> b -> c -> d -> e -> f -> m g) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector re e + -> Vector rf f + -> m (Vector DS g) +szipWith6M f v1 v2 v3 v4 v5 v6 = + fromStepsM $ + S.zipWith6M + f + (toStreamM v1) + (toStreamM v2) + (toStreamM v3) + (toStreamM v4) + (toStreamM v5) + (toStreamM v6) +{-# INLINE szipWith6M #-} + +-- | Just like `szipWithM`, zip two vectors together, but with an index aware +-- monadic action. The length of a resulting vector will be the smallest length of the +-- supplied vectors. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sizipWithM + :: forall ra rb a b c m + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, Monad m) + => (Ix1 -> a -> b -> m c) + -> Vector ra a + -> Vector rb b + -> m (Vector DS c) +sizipWithM f v1 v2 = fromStepsM $ S.zipWithM (uncurry f) (toStreamIxM v1) (toStreamM v2) +{-# INLINE sizipWithM #-} + +-- | Just like `szipWith3M`, zip three vectors together, but with an index aware +-- monadic action. The length of a resulting vector will be the smallest length of the +-- supplied vectors. +-- +-- @since 0.5.0 +sizipWith3M + :: forall ra rb rc a b c d m + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, Monad m) + => (Ix1 -> a -> b -> c -> m d) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> m (Vector DS d) +sizipWith3M f v1 v2 v3 = + fromStepsM $ S.zipWith3M (uncurry f) (toStreamIxM v1) (toStreamM v2) (toStreamM v3) +{-# INLINE sizipWith3M #-} + +-- | Just like `szipWith4M`, zip four vectors together, but with an index aware +-- monadic action. The length of a resulting vector will be the smallest length of the +-- supplied vectors. +-- +-- @since 0.5.0 +sizipWith4M + :: forall ra rb rc rd a b c d e m + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, S.Stream rd Ix1 d, Monad m) + => (Ix1 -> a -> b -> c -> d -> m e) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> m (Vector DS e) +sizipWith4M f v1 v2 v3 v4 = + fromStepsM $ + S.zipWith4M (uncurry f) (toStreamIxM v1) (toStreamM v2) (toStreamM v3) (toStreamM v4) +{-# INLINE sizipWith4M #-} + +-- | Just like `szipWith6M`, zip five vectors together, but with an index aware +-- monadic action. The length of a resulting vector will be the smallest length of the +-- supplied vectors. +-- +-- @since 0.5.0 +sizipWith5M + :: forall ra rb rc rd re a b c d e f m + . ( S.Stream ra Ix1 a + , S.Stream rb Ix1 b + , S.Stream rc Ix1 c + , S.Stream rd Ix1 d + , S.Stream re Ix1 e + , Monad m + ) + => (Ix1 -> a -> b -> c -> d -> e -> m f) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector re e + -> m (Vector DS f) +sizipWith5M f v1 v2 v3 v4 v5 = + fromStepsM $ + S.zipWith5M + (uncurry f) + (toStreamIxM v1) + (toStreamM v2) + (toStreamM v3) + (toStreamM v4) + (toStreamM v5) +{-# INLINE sizipWith5M #-} + +-- | Just like `szipWith6M`, zip six vectors together, but with an index aware +-- monadic action. The length of a resulting vector will be the smallest length of the +-- supplied vectors. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sizipWith6M + :: forall ra rb rc rd re rf a b c d e f g m + . ( S.Stream ra Ix1 a + , S.Stream rb Ix1 b + , S.Stream rc Ix1 c + , S.Stream rd Ix1 d + , S.Stream re Ix1 e + , S.Stream rf Ix1 f + , Monad m + ) + => (Ix1 -> a -> b -> c -> d -> e -> f -> m g) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector re e + -> Vector rf f + -> m (Vector DS g) +sizipWith6M f v1 v2 v3 v4 v5 v6 = + fromStepsM $ + S.zipWith6M + (uncurry f) + (toStreamIxM v1) + (toStreamM v2) + (toStreamM v3) + (toStreamM v4) + (toStreamM v5) + (toStreamM v6) +{-# INLINE sizipWith6M #-} + +-- | Similar to `szipWithM`, zip two vectors together with a binary monadic +-- action, while discarding its result. The action will be invoked as many times as +-- the length of the smallest vector. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +szipWithM_ + :: forall ra rb a b c m + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, Monad m) + => (a -> b -> m c) + -> Vector ra a + -> Vector rb b + -> m () +szipWithM_ f v1 v2 = S.zipWithM_ f (toStreamM v1) (toStreamM v2) +{-# INLINE szipWithM_ #-} + +-- | Similar to `szipWith3M`, zip three vectors together with a ternary monadic +-- action, while discarding its result. The action will be invoked as many times as +-- the length of the smallest vector. +-- +-- @since 0.5.0 +szipWith3M_ + :: forall ra rb rc a b c d m + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, Monad m) + => (a -> b -> c -> m d) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> m () +szipWith3M_ f v1 v2 v3 = S.zipWith3M_ f (toStreamM v1) (toStreamM v2) (toStreamM v3) +{-# INLINE szipWith3M_ #-} + +-- | Similar to `szipWith4M`, zip four vectors together with a quaternary monadic +-- action, while discarding its result. The action will be invoked as many times as +-- the length of the smallest vector. +-- +-- @since 0.5.0 +szipWith4M_ + :: forall ra rb rc rd a b c d e m + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, S.Stream rd Ix1 d, Monad m) + => (a -> b -> c -> d -> m e) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> m () +szipWith4M_ f v1 v2 v3 v4 = + S.zipWith4M_ f (toStreamM v1) (toStreamM v2) (toStreamM v3) (toStreamM v4) +{-# INLINE szipWith4M_ #-} + +-- | Similar to `szipWith5M`, zip five vectors together with a quinary monadic +-- action, while discarding its result. The action will be invoked as many times as +-- the length of the smallest vector. +-- +-- @since 0.5.0 +szipWith5M_ + :: forall ra rb rc rd re a b c d e f m + . ( S.Stream ra Ix1 a + , S.Stream rb Ix1 b + , S.Stream rc Ix1 c + , S.Stream rd Ix1 d + , S.Stream re Ix1 e + , Monad m + ) + => (a -> b -> c -> d -> e -> m f) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector re e + -> m () +szipWith5M_ f v1 v2 v3 v4 v5 = + S.zipWith5M_ f (toStreamM v1) (toStreamM v2) (toStreamM v3) (toStreamM v4) (toStreamM v5) +{-# INLINE szipWith5M_ #-} + +-- | Similar to `szipWith6M`, zip six vectors together with a senary monadic +-- action, while discarding its result. The action will be invoked as many times as +-- the length of the smallest vector. +-- +-- @since 0.5.0 +szipWith6M_ + :: forall ra rb rc rd re rf a b c d e f g m + . ( S.Stream ra Ix1 a + , S.Stream rb Ix1 b + , S.Stream rc Ix1 c + , S.Stream rd Ix1 d + , S.Stream re Ix1 e + , S.Stream rf Ix1 f + , Monad m + ) + => (a -> b -> c -> d -> e -> f -> m g) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector re e + -> Vector rf f + -> m () +szipWith6M_ f v1 v2 v3 v4 v5 v6 = + S.zipWith6M_ + f + (toStreamM v1) + (toStreamM v2) + (toStreamM v3) + (toStreamM v4) + (toStreamM v5) + (toStreamM v6) +{-# INLINE szipWith6M_ #-} + +-- | Same as `szipWithM_`, zip two vectors together, but with an index aware +-- monadic action. The action will be invoked as many times as the length of the +-- smallest vector. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sizipWithM_ + :: forall ra rb a b c m + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, Monad m) + => (Ix1 -> a -> b -> m c) + -> Vector ra a + -> Vector rb b + -> m () +sizipWithM_ f v1 v2 = S.zipWithM_ (uncurry f) (toStreamIxM v1) (toStreamM v2) +{-# INLINE sizipWithM_ #-} + +-- | Same as `szipWith3M_`, zip three vectors together, but with an index aware +-- monadic action. The action will be invoked as many times as the length of the +-- smallest vector. +-- +-- @since 0.5.0 +sizipWith3M_ + :: forall ra rb rc a b c d m + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, Monad m) + => (Ix1 -> a -> b -> c -> m d) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> m () +sizipWith3M_ f v1 v2 v3 = S.zipWith3M_ (uncurry f) (toStreamIxM v1) (toStreamM v2) (toStreamM v3) +{-# INLINE sizipWith3M_ #-} + +-- | Same as `szipWith4M_`, zip four vectors together, but with an index aware +-- monadic action. The action will be invoked as many times as the length of the +-- smallest vector. +-- +-- @since 0.5.0 +sizipWith4M_ + :: forall ra rb rc rd a b c d e m + . (S.Stream ra Ix1 a, S.Stream rb Ix1 b, S.Stream rc Ix1 c, S.Stream rd Ix1 d, Monad m) + => (Ix1 -> a -> b -> c -> d -> m e) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> m () +sizipWith4M_ f v1 v2 v3 v4 = + S.zipWith4M_ (uncurry f) (toStreamIxM v1) (toStreamM v2) (toStreamM v3) (toStreamM v4) +{-# INLINE sizipWith4M_ #-} + +-- | Same as `szipWith5M_`, zip five vectors together, but with an index aware +-- monadic action. The action will be invoked as many times as the length of the +-- smallest vector. +-- +-- @since 0.5.0 +sizipWith5M_ + :: forall ra rb rc rd re a b c d e f m + . ( S.Stream ra Ix1 a + , S.Stream rb Ix1 b + , S.Stream rc Ix1 c + , S.Stream rd Ix1 d + , S.Stream re Ix1 e + , Monad m + ) + => (Ix1 -> a -> b -> c -> d -> e -> m f) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector re e + -> m () +sizipWith5M_ f v1 v2 v3 v4 v5 = + S.zipWith5M_ + (uncurry f) + (toStreamIxM v1) + (toStreamM v2) + (toStreamM v3) + (toStreamM v4) + (toStreamM v5) +{-# INLINE sizipWith5M_ #-} + +-- | Same as `szipWith6M_`, zip six vectors together, but with an index aware +-- monadic action. The action will be invoked as many times as the length of the +-- smallest vector. +-- +-- @since 0.5.0 +sizipWith6M_ + :: forall ra rb rc rd re rf a b c d e f g m + . ( S.Stream ra Ix1 a + , S.Stream rb Ix1 b + , S.Stream rc Ix1 c + , S.Stream rd Ix1 d + , S.Stream re Ix1 e + , S.Stream rf Ix1 f + , Monad m + ) + => (Ix1 -> a -> b -> c -> d -> e -> f -> m g) + -> Vector ra a + -> Vector rb b + -> Vector rc c + -> Vector rd d + -> Vector re e + -> Vector rf f + -> m () +sizipWith6M_ f v1 v2 v3 v4 v5 v6 = + S.zipWith6M_ + (uncurry f) + (toStreamIxM v1) + (toStreamM v2) + (toStreamM v3) + (toStreamM v4) + (toStreamM v5) + (toStreamM v6) +{-# INLINE sizipWith6M_ #-} + +-- | Streaming fold over an array in a row-major fashion with a left biased +-- function and a strict accumulator. +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sfoldl + :: forall r ix e a + . Stream r ix e + => (a -> e -> a) + -> a + -> Array r ix e + -> a +sfoldl f acc = S.unId . S.foldl f acc . toStream +{-# INLINE sfoldl #-} + +-- | +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sfoldlM + :: forall r ix e a m + . (Stream r ix e, Monad m) + => (a -> e -> m a) + -> a + -> Array r ix e + -> m a +sfoldlM f acc = S.foldlM f acc . S.transStepsId . toStream +{-# INLINE sfoldlM #-} + +-- | +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sfoldlM_ + :: forall r ix e a m + . (Stream r ix e, Monad m) + => (a -> e -> m a) + -> a + -> Array r ix e + -> m () +sfoldlM_ f acc = void . sfoldlM f acc +{-# INLINE sfoldlM_ #-} + +-- | +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sfoldl1' + :: forall r ix e + . (HasCallStack, Stream r ix e) + => (e -> e -> e) + -> Array r ix e + -> e +sfoldl1' f = throwEither . sfoldl1M (\e -> pure . f e) +{-# INLINE sfoldl1' #-} + +-- | +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sfoldl1M + :: forall r ix e m + . (Stream r ix e, MonadThrow m) + => (e -> e -> m e) + -> Array r ix e + -> m e +sfoldl1M f arr = do + let str = S.transStepsId $ toStream arr + isNullStream <- S.null str + when isNullStream $ throwM $ SizeEmptyException (outerSize arr) + S.foldl1M f str +{-# INLINE sfoldl1M #-} + +-- | +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sfoldl1M_ + :: forall r ix e m + . (Stream r ix e, MonadThrow m) + => (e -> e -> m e) + -> Array r ix e + -> m () +sfoldl1M_ f = void . sfoldl1M f +{-# INLINE sfoldl1M_ #-} + +-- | +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sifoldl + :: forall r ix e a + . Stream r ix e + => (a -> ix -> e -> a) + -> a + -> Array r ix e + -> a +sifoldl f acc = S.unId . S.foldl (\a (ix, e) -> f a ix e) acc . toStreamIx +{-# INLINE sifoldl #-} + +-- | +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sifoldlM + :: forall r ix e a m + . (Stream r ix e, Monad m) + => (a -> ix -> e -> m a) + -> a + -> Array r ix e + -> m a +sifoldlM f acc = S.foldlM (\a (ix, e) -> f a ix e) acc . S.transStepsId . toStreamIx +{-# INLINE sifoldlM #-} + +-- | +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sifoldlM_ + :: forall r ix e a m + . (Stream r ix e, Monad m) + => (a -> ix -> e -> m a) + -> a + -> Array r ix e + -> m () +sifoldlM_ f acc = void . sifoldlM f acc +{-# INLINE sifoldlM_ #-} + +-- | +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sor + :: forall r ix + . Stream r ix Bool + => Array r ix Bool + -> Bool +sor = S.unId . S.or . toStream +{-# INLINE sor #-} + +-- | +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sand :: forall r ix. Stream r ix Bool => Array r ix Bool -> Bool +sand = S.unId . S.and . toStream +{-# INLINE sand #-} + +-- | +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sany :: forall r ix e. Stream r ix e => (e -> Bool) -> Array r ix e -> Bool +sany f = S.unId . S.or . S.map f . toStream +{-# INLINE sany #-} + +-- | +-- +-- ==== __Examples__ +-- +-- @since 0.5.0 +sall :: forall r ix e. Stream r ix e => (e -> Bool) -> Array r ix e -> Bool +sall f = S.unId . S.and . S.map f . toStream +{-# INLINE sall #-} + +-- | Add all elements of the array together +-- +-- /Related/: `sum`. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Vector as V +-- >>> V.ssum $ V.sfromList [10, 3, 70, 5 :: Int] +-- 88 +-- +-- @since 0.5.0 +ssum :: forall r ix e. (Num e, Stream r ix e) => Array r ix e -> e +ssum = sfoldl (+) 0 +{-# INLINE ssum #-} + +-- | Multiply all elements of the array together +-- +-- /Related/: `product`. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Vector as V +-- >>> V.sproduct $ V.sfromList [10, 3, 70, 5 :: Int] +-- 10500 +-- +-- @since 0.5.0 +sproduct :: forall r ix e. (Num e, Stream r ix e) => Array r ix e -> e +sproduct = sfoldl (*) 1 +{-# INLINE sproduct #-} + +-- | /O(n)/ - Find the largest value in the array. Throws an error on empty. +-- +-- /Related/: `smaximumM`, `maximum`, `maximumM`. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Vector as V +-- >>> V.smaximum' $ V.sfromList [10, 3, 70, 5 :: Int] +-- 70 +-- +-- @since 0.5.0 +smaximum' :: forall r ix e. (HasCallStack, Ord e, Stream r ix e) => Array r ix e -> e +smaximum' = sfoldl1' max +{-# INLINE smaximum' #-} + +-- | /O(n)/ - Find the largest value in the array. +-- +-- /Related/: `smaximum`, `maximum`, `maximumM`. +-- +-- /__Throws Exceptions__/: `SizeEmptyException` when array is empty +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Vector as V +-- >>> V.smaximumM $ V.sfromList [10, 3, 70, 5 :: Int] +-- 70 +-- >>> V.smaximumM (V.empty :: Vector D Int) :: Maybe Int +-- Nothing +-- +-- @since 0.5.0 +smaximumM :: forall r ix e m. (Ord e, Stream r ix e, MonadThrow m) => Array r ix e -> m e +smaximumM = sfoldl1M (\e acc -> pure (max e acc)) +{-# INLINE smaximumM #-} + +-- | /O(n)/ - Find the smallest value in the array. Throws an error on empty. +-- +-- /Related/: `sminimumM`, `minimum`, `minimumM`. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Vector as V +-- >>> V.sminimum' $ V.sfromList [10, 3, 70, 5 :: Int] +-- 3 +-- +-- @since 0.5.0 +sminimum' :: forall r ix e. (HasCallStack, Ord e, Stream r ix e) => Array r ix e -> e +sminimum' = sfoldl1' min +{-# INLINE sminimum' #-} + +-- | /O(n)/ - Find the smallest value in the array. +-- +-- /Related/: 'sminimum'', `minimum`, `minimumM`. +-- +-- /__Throws Exceptions__/: `SizeEmptyException` when array is empty +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Vector as V +-- >>> V.sminimumM $ V.sfromList [10, 3, 70, 5 :: Int] +-- 3 +-- >>> V.sminimumM (V.empty :: Array D Ix2 Int) :: Maybe Int +-- Nothing +-- +-- @since 0.5.0 +sminimumM :: forall r ix e m. (Ord e, Stream r ix e, MonadThrow m) => Array r ix e -> m e +sminimumM = sfoldl1M (\e acc -> pure (min e acc)) +{-# INLINE sminimumM #-} + +-- | /O(n)/ - left scan with strict accumulator. First element is the value of the +-- accumulator. Last element is not included. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Vector +-- >>> sprescanl min 6 $ sfromList [10, 5, 70, 3 :: Int] +-- Array DS Seq (Sz1 4) +-- [ 6, 6, 5, 5 ] +-- >>> sprescanl (+) 0 $ sfromList [10, 5, 70, 3 :: Int] +-- Array DS Seq (Sz1 4) +-- [ 0, 10, 15, 85 ] +-- +-- @since 1.0.3 +sprescanl :: Stream r ix e => (a -> e -> a) -> a -> Array r ix e -> Vector DS a +sprescanl f acc = DSArray . S.prescanlM (\a b -> pure (f a b)) acc . toStream +{-# INLINE sprescanl #-} + +-- | /O(n)/ - left scan with strict accumulator. First element is the result of applying +-- the supplied function. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Vector +-- >>> spostscanl min 6 $ sfromList [10, 5, 70, 3 :: Int] +-- Array DS Seq (Sz1 4) +-- [ 6, 5, 5, 3 ] +-- >>> spostscanl (+) 0 $ sfromList [10, 5, 70, 3 :: Int] +-- Array DS Seq (Sz1 4) +-- [ 10, 15, 85, 88 ] +-- +-- @since 1.0.3 +spostscanl :: Stream r ix e => (a -> e -> a) -> a -> Array r ix e -> Vector DS a +spostscanl f acc = DSArray . S.postscanlM (\a b -> pure (f a b)) acc . toStream +{-# INLINE spostscanl #-} + +-- | /O(n)/ - Just like `spostscanl` except it is possible to produce a vector with an +-- element type that differes from accumulator type. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Vector +-- >>> spostscanlAcc (\x y -> if x < y then (True, x) else (False, y)) 6 $ sfromList [10, 5, 70, 3 :: Int] +-- Array DS Seq (Sz1 4) +-- [ True, False, True, False ] +-- +-- @since 1.0.3 +spostscanlAcc :: Stream r ix e => (c -> e -> (a, c)) -> c -> Array r ix e -> Vector DS a +spostscanlAcc f acc = DSArray . S.postscanlAccM (\a b -> pure (f a b)) acc . toStream +{-# INLINE spostscanlAcc #-} + + +-- | /O(n)/ - left scan with strict accumulator. First element is the value of the accumulator. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Vector +-- >>> sscanl min 6 $ sfromList [10, 5, 70, 3 :: Int] +-- Array DS Seq (Sz1 5) +-- [ 6, 6, 5, 5, 3 ] +-- >>> sscanl (+) 0 $ sfromList [10, 5, 70, 3 :: Int] +-- Array DS Seq (Sz1 5) +-- [ 0, 10, 15, 85, 88 ] +-- +-- @since 1.0.3 +sscanl :: Stream r ix e => (a -> e -> a) -> a -> Array r ix e -> Vector DS a +sscanl f acc = DSArray . S.scanlM (\a b -> pure (f a b)) acc . toStream +{-# INLINE sscanl #-} + +-- | /O(n)/ - left scan with strict accumulator and no initial value for the accumulator. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Vector +-- >>> sscanl1 min $ sfromList [10, 5, 70, 3 :: Int] +-- Array DS Seq (Sz1 4) +-- [ 10, 5, 5, 3 ] +-- >>> sscanl1 (+) $ sfromList [10, 5, 70, 3 :: Int] +-- Array DS Seq (Sz1 4) +-- [ 10, 15, 85, 88 ] +-- >>> sscanl1 (+) $ sfromList ([] :: [Int]) +-- Array DS Seq (Sz1 0) +-- [ ] +-- +-- @since 1.0.3 +sscanl1 :: Stream r ix e => (e -> e -> e) -> Array r ix e -> Vector DS e +sscanl1 f = DSArray . S.scanl1M (\a b -> pure (f a b)) . toStream +{-# INLINE sscanl1 #-} diff --git a/src/Data/Massiv/Vector/Stream.hs b/src/Data/Massiv/Vector/Stream.hs --- a/src/Data/Massiv/Vector/Stream.hs +++ b/src/Data/Massiv/Vector/Stream.hs @@ -4,12 +4,14 @@ {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE LambdaCase #-} {-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE NamedFieldPuns #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TupleSections #-} {-# LANGUAGE TypeFamilies #-} {-# OPTIONS_GHC -fno-warn-orphans #-} {-# OPTIONS_HADDOCK hide, not-home #-} + -- | -- Module : Data.Massiv.Vector.Stream -- Copyright : (c) Alexey Kuleshevich 2019-2022 @@ -17,143 +19,170 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Vector.Stream - ( -- | This module has a similar purpose as the 'Data.Vector.Fusion.Bundle.Monadic', but - -- quite a bit simpler. - -- - -- __Important__ - This module is still experimental, as such it is considered - -- internal and exported for the curious users only. - Steps(..) - , Stream(..) +module Data.Massiv.Vector.Stream ( + -- | This module has a similar purpose as the 'Data.Vector.Fusion.Bundle.Monadic', but + -- quite a bit simpler. + -- + -- __Important__ - This module is still experimental, as such it is considered + -- internal and exported for the curious users only. + Steps (..), + Stream (..), + -- * Conversion - , steps - , isteps - , consume - , fromStream - , fromStreamM - , fromStreamExactM - , unstreamExact - , unstreamMax - , unstreamMaxM - , unstreamUnknown - , unstreamUnknownM - , unstreamIntoM - -- * Bundle - , toBundle - , fromBundle - , fromBundleM + steps, + isteps, + consume, + fromStream, + fromStreamM, + fromStreamExactM, + unstreamExact, + unstreamMax, + unstreamMaxM, + unstreamUnknown, + unstreamUnknownM, + unstreamIntoM, + -- * Operations on Steps - , length - , null - , empty - , singleton - , generate - , headMaybe - , last - , cons - , uncons - , snoc - , drop - , take - , slice - , iterateN - , iterateNM - , replicate - , replicateM - , generateM - , traverse - , map - , mapM - , mapM_ - , indexed - , concatMap - , append - , zipWith - , zipWith3 - , zipWith4 - , zipWith5 - , zipWith6 - , zipWithM - , zipWith3M - , zipWith4M - , zipWith5M - , zipWith6M - , zipWithM_ - , zipWith3M_ - , zipWith4M_ - , zipWith5M_ - , zipWith6M_ + length, + null, + empty, + singleton, + generate, + headMaybe, + last, + cons, + uncons, + snoc, + drop, + take, + slice, + iterateN, + iterateNM, + replicate, + replicateM, + generateM, + traverse, + map, + mapM, + mapM_, + indexed, + concatMap, + append, + zipWith, + zipWith3, + zipWith4, + zipWith5, + zipWith6, + zipWithM, + zipWith3M, + zipWith4M, + zipWith5M, + zipWith6M, + zipWithM_, + zipWith3M_, + zipWith4M_, + zipWith5M_, + zipWith6M_, + -- ** Folding - , foldl - , foldl1 - , foldlM - , foldl1M - , foldlLazy - , foldl1Lazy - , foldlLazyM - , foldl1LazyM - , foldrLazy - , foldr1Lazy - , foldrLazyM - , foldr1LazyM + foldl, + foldl1, + foldlM, + foldl1M, + foldlLazy, + foldl1Lazy, + foldlLazyM, + foldl1LazyM, + foldrLazy, + foldr1Lazy, + foldrLazyM, + foldr1LazyM, + or, + and, - , or - , and -- ** Unfolding - , unfoldr - , unfoldrN - , unsafeUnfoldrN - , unfoldrM - , unfoldrNM - , unsafeUnfoldrNM - , unfoldrExactN - , unfoldrExactNM + unfoldr, + unfoldrN, + unsafeUnfoldrN, + unfoldrM, + unfoldrNM, + unsafeUnfoldrNM, + unfoldrExactN, + unfoldrExactNM, + + -- ** Scanning + prescanlM, + postscanlM, + postscanlAccM, + scanlM, + scanl1M, + -- ** Enumeration - , enumFromStepN + enumFromStepN, + -- * Lists - , toList - , fromList - , fromListN - , unsafeFromListN + toList, + fromList, + fromListN, + unsafeFromListN, + -- ** Filter - , mapMaybe - , mapMaybeA - , mapMaybeM - , filter - , filterA - , filterM + mapMaybe, + mapMaybeA, + mapMaybeM, + filter, + filterA, + filterM, + -- * Transformations - , transSteps - , transStepsId + transSteps, + transStepsId, + -- * Useful re-exports - , module Data.Vector.Fusion.Util - , Id(..) - ) where + module Data.Vector.Fusion.Util, + Id (..), +) where import qualified Control.Monad as M import Control.Monad.ST -import qualified Data.Foldable as F -import Data.Massiv.Core.Common hiding (empty, singleton, replicate) import Data.Coerce +import qualified Data.Foldable as F +import Data.Massiv.Core.Common hiding (empty, replicate, singleton) import Data.Maybe (catMaybes) +import qualified Data.Stream.Monadic as S import qualified Data.Traversable as Traversable (traverse) -import qualified Data.Vector.Fusion.Bundle.Monadic as B 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(..)) +import qualified GHC.Exts (IsList (..)) +import Prelude hiding ( + and, + concatMap, + drop, + filter, + foldl, + foldl1, + foldr, + foldr1, + length, + map, + mapM, + mapM_, + null, + or, + replicate, + take, + traverse, + zipWith, + zipWith3, + ) instance Monad m => Functor (Steps m) where - fmap f str = str {stepsStream = S.map f (stepsStream str)} + fmap f str = str{stepsStream = S.map f (stepsStream str)} {-# INLINE fmap #-} (<$) e str = case stepsSize str of - LengthExact n -> str {stepsStream = S.replicate (coerce n) e} - _ -> fmap (const e) str + LengthExact n -> str{stepsStream = S.replicate (coerce n) e} + _ -> fmap (const e) str {-# INLINE (<$) #-} instance Monad m => Semigroup (Steps m e) where @@ -203,34 +232,31 @@ minimum = unId . foldl1 min {-# INLINE minimum #-} - -steps :: forall r ix e m . (Monad m, Index ix, Source r e) => Array r ix e -> Steps m e +steps :: forall r ix e m. (Monad m, Index ix, Source r e) => Array r ix e -> Steps m e steps !arr = case unsafePrefIndex arr of PrefIndex gix -> gix <$> ixRangeSteps (size arr) PrefIndexLinear gi -> Steps (S.Stream step 0) (LengthExact (coerce k)) where - k = totalElem $ size arr - step i - | i < k = - let e = gi i - in e `seq` pure $ S.Yield e (i + 1) + !k = totalElem $ size arr + step !i + | i < k = pure $ S.Yield (gi i) (i + 1) | otherwise = pure S.Done {-# INLINE [0] step #-} {-# INLINE [1] steps #-} -ixRangeSteps :: forall ix m . (Monad m, Index ix) => Sz ix -> Steps m ix +ixRangeSteps :: forall ix m. (Monad m, Index ix) => Sz ix -> Steps m ix ixRangeSteps sz = Steps (S.Stream step initStep) (LengthExact k) where !k = toLinearSz sz !initStep = if k == zeroSz then Nothing else Just zeroIndex - step (Just ix) = stepNextMF ix (unSz sz) oneIndex (<) $ \ mIx -> pure $ S.Yield ix mIx + step (Just ix) = stepNextMF ix (unSz sz) oneIndex (<) $ \mIx -> pure $ S.Yield ix mIx step Nothing = pure S.Done {-# INLINE [0] step #-} {-# INLINE [1] ixRangeSteps #-} -isteps :: forall r ix e m . (Monad m, Index ix, Source r e) => Array r ix e -> Steps m (ix, e) +isteps :: forall r ix e m. (Monad m, Index ix, Source r e) => Array r ix e -> Steps m (ix, e) isteps !arr = case unsafePrefIndex arr of PrefIndex gix -> (\ !ix -> let e = gix ix in e `seq` (ix, e)) <$> ixRangeSteps sz @@ -238,37 +264,20 @@ let k = totalElem sz step i | i < k = - let e = gi i - in e `seq` pure $ S.Yield (fromLinearIndex sz i, e) (i + 1) + let e = gi i + in e `seq` pure $ S.Yield (fromLinearIndex sz i, e) (i + 1) | otherwise = pure S.Done {-# INLINE [0] step #-} - in Steps (S.Stream step 0) (LengthExact (coerce k)) + in Steps (S.Stream step 0) (LengthExact (coerce k)) where !sz = size arr {-# INLINE isteps #-} - -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 (sizeHintToBundleSize k) -{-# INLINE toBundle #-} - -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, 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 . Manifest r e => B.Size -> S.Stream Id e -> Vector r e +fromStream :: forall r e. Manifest r e => B.Size -> S.Stream Id e -> Vector r e fromStream sz str = case B.upperBound sz of Nothing -> unstreamUnknown str - Just k -> unstreamMax k str + Just k -> unstreamMax k str {-# INLINE fromStream #-} fromStreamM :: forall r e m. (Monad m, Manifest r e) => B.Size -> S.Stream m e -> m (Vector r e) @@ -276,11 +285,12 @@ xs <- S.toList str case B.upperBound sz of Nothing -> pure $! unstreamUnknown (S.fromList xs) - Just k -> pure $! unstreamMax k (S.fromList xs) + Just k -> pure $! unstreamMax k (S.fromList xs) {-# INLINE fromStreamM #-} -fromStreamExactM :: - forall r ix e m. (Monad m, Manifest r e, Index ix) +fromStreamExactM + :: forall r ix e m + . (Monad m, Manifest r e, Index ix) => Sz ix -> S.Stream m e -> m (Array r ix e) @@ -289,9 +299,8 @@ pure $! unstreamExact sz (S.fromList xs) {-# INLINE fromStreamExactM #-} - -unstreamIntoM :: - (Manifest r a, PrimMonad m) +unstreamIntoM + :: (Manifest r a, PrimMonad m) => MVector (PrimState m) r a -> LengthHint -> S.Stream Id a @@ -299,14 +308,13 @@ unstreamIntoM marr sz str = case sz of LengthExact _ -> marr <$ unstreamMaxM marr str - LengthMax _ -> unsafeLinearShrink marr . SafeSz =<< unstreamMaxM marr str + LengthMax _ -> unsafeLinearShrink marr . SafeSz =<< unstreamMaxM marr str LengthUnknown -> unstreamUnknownM marr str {-# INLINE unstreamIntoM #-} - - -unstreamMax :: - forall r e. (Manifest r e) +unstreamMax + :: forall r e + . Manifest r e => Int -> S.Stream Id e -> Vector r e @@ -314,25 +322,21 @@ runST $ do marr <- unsafeNew (SafeSz kMax) k <- unstreamMaxM marr str - unsafeLinearShrink marr (SafeSz k) >>= unsafeFreeze Seq + marrShrunk <- + if k == kMax + then pure marr + else unsafeLinearShrink marr (SafeSz k) + unsafeFreeze Seq marrShrunk {-# INLINE unstreamMax #-} - -unstreamMaxM :: - (Manifest r a, Index ix, PrimMonad m) => MArray (PrimState m) r ix a -> S.Stream Id a -> m Int -unstreamMaxM marr (S.Stream step s) = stepLoad s 0 +unstreamMaxM + :: (Manifest r a, Index ix, PrimMonad m) => MArray (PrimState m) r ix a -> S.Stream Id a -> m Int +unstreamMaxM marr = S.foldlM' fillAtIndex 0 . S.trans (pure . unId) where - stepLoad t i = - case unId (step t) of - S.Yield e' t' -> do - unsafeLinearWrite marr i e' - stepLoad t' (i + 1) - S.Skip t' -> stepLoad t' i - S.Done -> pure i - {-# INLINE stepLoad #-} + fillAtIndex i x = (i + 1) <$ unsafeLinearWrite marr i x + {-# INLINE fillAtIndex #-} {-# INLINE unstreamMaxM #-} - unstreamUnknown :: Manifest r a => S.Stream Id a -> Vector r a unstreamUnknown str = runST $ do @@ -340,32 +344,30 @@ unstreamUnknownM marr str >>= unsafeFreeze Seq {-# INLINE unstreamUnknown #-} - -unstreamUnknownM :: - (Manifest r a, PrimMonad m) +unstreamUnknownM + :: (Manifest r a, PrimMonad m) => MVector (PrimState m) r a -> S.Stream Id a -> m (MVector (PrimState m) r a) -unstreamUnknownM marrInit (S.Stream step s) = stepLoad s 0 (unSz (sizeOfMArray marrInit)) marrInit +unstreamUnknownM marr str = do + (marr', k) <- S.foldlM' fillAtIndex (marr, 0) $ S.trans (pure . unId) str + if k < unSz (sizeOfMArray marr') + then unsafeLinearShrink marr' (SafeSz k) + else pure marr' where - stepLoad t i kMax marr - | i < kMax = - case unId (step t) of - S.Yield e' t' -> do - unsafeLinearWrite marr i e' - stepLoad t' (i + 1) kMax marr - S.Skip t' -> stepLoad t' i kMax marr - S.Done -> unsafeLinearShrink marr (SafeSz i) - | otherwise = do - let kMax' = max 1 (kMax * 2) - marr' <- unsafeLinearGrow marr (SafeSz kMax') - stepLoad t i kMax' marr' - {-# INLINE stepLoad #-} + fillAtIndex (!ma, !i) x = do + let k = unSz (sizeOfMArray ma) + ma' <- + if i < k + then pure ma + else unsafeLinearGrow ma (SafeSz (max 1 k * 2)) + (ma', i + 1) <$ unsafeLinearWrite ma' i x + {-# INLINE fillAtIndex #-} {-# INLINE unstreamUnknownM #-} - -unstreamExact :: - forall r ix e. (Manifest r e, Index ix) +unstreamExact + :: forall r ix e + . (Manifest r e, Index ix) => Sz ix -> S.Stream Id e -> Array r ix e @@ -380,15 +382,14 @@ length (Steps str sz) = case sz of LengthExact k -> pure $ coerce k - _ -> S.length str + _ -> S.length str {-# INLINE length #-} - null :: Monad m => Steps m a -> m Bool null (Steps str sz) = case sz of LengthExact k -> pure (k == zeroSz) - _ -> S.null str + _ -> S.null str {-# INLINE null #-} empty :: Monad m => Steps m e @@ -411,19 +412,18 @@ r <- step s case r of S.Yield x _ -> pure $ Just x - S.Skip s' -> headMaybeLoop S.SPEC s' - S.Done -> pure Nothing + S.Skip s' -> headMaybeLoop S.SPEC s' + S.Done -> pure Nothing {-# INLINE [0] headMaybeLoop #-} {-# INLINE headMaybe #-} - cons :: Monad m => e -> Steps m e -> Steps m e cons e (Steps str k) = Steps (S.cons e str) (k `addInt` 1) {-# INLINE cons #-} -- | First element of the `Steps` or `Nothing` if empty uncons :: Monad m => Steps m e -> m (Maybe (e, Steps m e)) -uncons sts = (\mx -> (, drop oneSz sts) <$> mx) <$> headMaybe sts +uncons sts = (\mx -> (,drop oneSz sts) <$> mx) <$> headMaybe sts {-# INLINE uncons #-} snoc :: Monad m => Steps m e -> e -> Steps m e @@ -463,14 +463,14 @@ 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 + :: Monad m => (a -> b -> c -> d -> e) -> Steps m a -> Steps m b -> Steps m c -> Steps m d -> Steps m e zipWith4 f (Steps sa ka) (Steps sb kb) (Steps sc kc) (Steps sd kd) = Steps (S.zipWith4 f sa sb sc sd) (minLengthHint ka (minLengthHint kb (minLengthHint kc kd))) {-# INLINE zipWith4 #-} -zipWith5 :: - Monad m +zipWith5 + :: Monad m => (a -> b -> c -> d -> e -> f) -> Steps m a -> Steps m b @@ -482,8 +482,8 @@ Steps (S.zipWith5 f sa sb sc sd se) (minLengthHint ka (minLengthHint kb (minLengthHint kc (minLengthHint kd ke)))) {-# INLINE zipWith5 #-} -zipWith6 :: - Monad m +zipWith6 + :: Monad m => (a -> b -> c -> d -> e -> f -> g) -> Steps m a -> Steps m b @@ -502,14 +502,13 @@ zipWithM f (Steps sa ka) (Steps sb kb) = Steps (S.zipWithM f sa sb) (minLengthHint ka kb) {-# INLINE zipWithM #-} - zipWith3M :: Monad m => (a -> b -> c -> m d) -> Steps m a -> Steps m b -> Steps m c -> Steps m d zipWith3M f (Steps sa ka) (Steps sb kb) (Steps sc kc) = Steps (S.zipWith3M f sa sb sc) (minLengthHint ka (minLengthHint kb kc)) {-# INLINE zipWith3M #-} -zipWith4M :: - Monad m +zipWith4M + :: Monad m => (a -> b -> c -> d -> m e) -> Steps m a -> Steps m b @@ -520,8 +519,8 @@ Steps (S.zipWith4M f sa sb sc sd) (minLengthHint ka (minLengthHint kb (minLengthHint kc kd))) {-# INLINE zipWith4M #-} -zipWith5M :: - Monad m +zipWith5M + :: Monad m => (a -> b -> c -> d -> e -> m f) -> Steps m a -> Steps m b @@ -533,8 +532,8 @@ Steps (S.zipWith5M f sa sb sc sd se) (minLengthHint ka (minLengthHint kb (minLengthHint kc (minLengthHint kd ke)))) {-# INLINE zipWith5M #-} -zipWith6M :: - Monad m +zipWith6M + :: Monad m => (a -> b -> c -> d -> e -> f -> m g) -> Steps m a -> Steps m b @@ -549,7 +548,6 @@ (minLengthHint ka (minLengthHint kb (minLengthHint kc (minLengthHint kd (minLengthHint ke kf))))) {-# INLINE zipWith6M #-} - zipWithM_ :: Monad m => (a -> b -> m c) -> Steps m a -> Steps m b -> m () zipWithM_ f (Steps str1 _) (Steps str2 _) = S.zipWithM_ f str1 str2 {-# INLINE zipWithM_ #-} @@ -558,9 +556,8 @@ zipWith3M_ f sa sb sc = consume $ zipWith3M f sa sb sc {-# INLINE zipWith3M_ #-} - -zipWith4M_ :: - Monad m +zipWith4M_ + :: Monad m => (a -> b -> c -> d -> m e) -> Steps m a -> Steps m b @@ -570,8 +567,8 @@ zipWith4M_ f sa sb sc sd = consume $ zipWith4M f sa sb sc sd {-# INLINE zipWith4M_ #-} -zipWith5M_ :: - Monad m +zipWith5M_ + :: Monad m => (a -> b -> c -> d -> e -> m f) -> Steps m a -> Steps m b @@ -582,8 +579,8 @@ zipWith5M_ f sa sb sc sd se = consume $ zipWith5M f sa sb sc sd se {-# INLINE zipWith5M_ #-} -zipWith6M_ :: - Monad m +zipWith6M_ + :: Monad m => (a -> b -> c -> d -> e -> f -> m g) -> Steps m a -> Steps m b @@ -595,8 +592,6 @@ zipWith6M_ f sa sb sc sd se sf = consume $ zipWith6M f sa sb sc sd se sf {-# INLINE zipWith6M_ #-} - - consume :: Monad m => Steps m a -> m () consume (Steps (S.Stream step t) _) = consumeLoop S.SPEC t where @@ -619,7 +614,6 @@ pure (Steps strN (LengthExact n)) {-# INLINE transSteps #-} - foldl :: Monad m => (b -> a -> b) -> b -> Steps m a -> m b foldl f acc = S.foldl' f acc . stepsStream {-# INLINE foldl #-} @@ -628,17 +622,14 @@ foldl1 f = S.foldl1' f . stepsStream {-# INLINE foldl1 #-} - foldlM :: Monad m => (a -> b -> m a) -> a -> Steps m b -> m a foldlM f acc = S.foldlM' f acc . stepsStream {-# INLINE foldlM #-} - foldl1M :: Monad m => (a -> a -> m a) -> Steps m a -> m a foldl1M f (Steps sts _) = S.foldl1M' f sts {-# INLINE foldl1M #-} - foldrLazy :: Monad m => (a -> b -> b) -> b -> Steps m a -> m b foldrLazy f acc = S.foldr f acc . stepsStream {-# INLINE foldrLazy #-} @@ -655,27 +646,22 @@ foldl1Lazy f = S.foldl1 f . stepsStream {-# INLINE foldl1Lazy #-} - foldlLazyM :: Monad m => (a -> b -> m a) -> a -> Steps m b -> m a foldlLazyM f acc = S.foldlM f acc . stepsStream {-# INLINE foldlLazyM #-} - foldl1LazyM :: Monad m => (a -> a -> m a) -> Steps m a -> m a foldl1LazyM f (Steps sts _) = S.foldl1M f sts {-# INLINE foldl1LazyM #-} - foldrLazyM :: Monad m => (b -> a -> m a) -> a -> Steps m b -> m a foldrLazyM f acc (Steps sts _) = S.foldrM f acc sts {-# INLINE foldrLazyM #-} - foldr1LazyM :: Monad m => (a -> a -> m a) -> Steps m a -> m a foldr1LazyM f = S.foldr1M f . stepsStream {-# INLINE foldr1LazyM #-} - or :: Monad m => Steps m Bool -> m Bool or = S.or . stepsStream {-# INLINE or #-} @@ -684,7 +670,6 @@ and = S.and . stepsStream {-# INLINE and #-} - mapMaybe :: Monad m => (a -> Maybe e) -> Steps m a -> Steps m e mapMaybe f (Steps str k) = Steps (S.mapMaybe f str) (toLengthMax k) {-# INLINE mapMaybe #-} @@ -693,7 +678,6 @@ concatMap f (Steps str _) = Steps (S.concatMap (stepsStream . f) str) LengthUnknown {-# INLINE concatMap #-} - mapMaybeA :: (Monad m, Applicative f) => (a -> f (Maybe e)) -> Steps Id a -> f (Steps m e) mapMaybeA f (Steps str k) = (`Steps` toLengthMax k) <$> liftListA (mapMaybeListA f) str {-# INLINE mapMaybeA #-} @@ -727,7 +711,6 @@ filter f (Steps str k) = Steps (S.filter f str) (toLengthMax k) {-# INLINE filter #-} - filterA :: (Monad m, Applicative f) => (e -> f Bool) -> Steps Id e -> f (Steps m e) filterA f (Steps str k) = (`Steps` toLengthMax k) <$> liftListA (M.filterM f) str {-# INLINE filterA #-} @@ -739,10 +722,10 @@ take :: Monad m => Sz1 -> Steps m a -> Steps m a take n (Steps str sz) = Steps (S.take (coerce n) str) $! - case sz of - LengthExact k -> LengthExact (inline0 min n k) - LengthMax k -> LengthMax (inline0 min n k) - LengthUnknown -> LengthUnknown + case sz of + LengthExact k -> LengthExact (inline0 min n k) + LengthMax k -> LengthMax (inline0 min n k) + LengthUnknown -> LengthUnknown {-# INLINE take #-} drop :: Monad m => Sz1 -> Steps m a -> Steps m a @@ -769,12 +752,10 @@ replicateM n f = Steps (S.replicateM (coerce n) f) (LengthExact n) {-# INLINE replicateM #-} - generateM :: Monad m => Sz1 -> (Int -> m a) -> Steps m a generateM n f = Steps (S.generateM (coerce n) f) (LengthExact n) {-# INLINE generateM #-} - unfoldr :: Monad m => (s -> Maybe (e, s)) -> s -> Steps m e unfoldr f e0 = Steps (S.unfoldr f e0) LengthUnknown {-# INLINE unfoldr #-} @@ -812,14 +793,10 @@ {-# INLINE [0] step #-} {-# INLINE unfoldrExactNM #-} - enumFromStepN :: (Num a, Monad m) => a -> a -> Sz1 -> Steps m a enumFromStepN x step k = Steps (S.enumFromStepN x step (coerce k)) (LengthExact k) {-# INLINE enumFromStepN #-} - - - toList :: Steps Id e -> [e] toList (Steps str _) = unId (S.toList str) {-# INLINE toList #-} @@ -829,18 +806,17 @@ {-# INLINE fromList #-} fromListN :: Monad m => Int -> [e] -> Steps m e -fromListN n = (`Steps` LengthUnknown) . S.fromListN n +fromListN n = (`Steps` LengthUnknown) . S.fromListN n {-# INLINE fromListN #-} unsafeFromListN :: Monad m => Sz1 -> [e] -> Steps m e -unsafeFromListN n = (`Steps` LengthMax n) . S.fromListN (coerce n) +unsafeFromListN n = (`Steps` LengthMax n) . S.fromListN (coerce n) {-# INLINE unsafeFromListN #-} liftListA :: (Monad m, Functor f) => ([a] -> f [b]) -> S.Stream Id a -> f (S.Stream m b) liftListA f str = S.fromList <$> f (unId (S.toList str)) {-# INLINE liftListA #-} - transListM :: (Monad m, Monad n) => S.Stream m a -> m (S.Stream n a) transListM str = do xs <- S.toList str @@ -853,17 +829,15 @@ 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, []) +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) + LengthMax k -> B.Max (coerce k) LengthUnknown -> B.Unknown {-# INLINE sizeHintToBundleSize #-} @@ -876,47 +850,89 @@ 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 +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 +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 +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 +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 +toLengthMax (LengthMax n) = LengthMax n +toLengthMax LengthUnknown = LengthUnknown {-# INLINE toLengthMax #-} +-- | Prefix scan with strict accumulator and a monadic operator +prescanlM :: Monad m => (a -> b -> m a) -> a -> Steps m b -> Steps m a +prescanlM f acc ss = ss{stepsStream = S.prescanlM' f acc (stepsStream ss)} +{-# INLINE prescanlM #-} +-- | Suffix scan with a monadic operator +postscanlM :: Monad m => (a -> b -> m a) -> a -> Steps m b -> Steps m a +postscanlM f acc ss = ss{stepsStream = S.postscanlM' f acc (stepsStream ss)} +{-# INLINE postscanlM #-} + +-- | Suffix scan with a monadic operator +postscanlAccM :: Monad m => (c -> b -> m (a, c)) -> c -> Steps m b -> Steps m a +postscanlAccM f acc ss = ss{stepsStream = postscanlAccStreamM f acc (stepsStream ss)} +{-# INLINE postscanlAccM #-} + +-- | Suffix scan with strict acccumulator and a monadic operator +postscanlAccStreamM :: Monad m => (c -> b -> m (a, c)) -> c -> S.Stream m b -> S.Stream m a +postscanlAccStreamM f w (S.Stream step t) = w `seq` S.Stream step' (t, w) + where + step' (s, x) = + x `seq` + do + r <- step s + case r of + S.Yield y s' -> do + (a, z) <- f x y + z `seq` return (S.Yield a (s', z)) + S.Skip s' -> return $ S.Skip (s', x) + S.Done -> return S.Done + {-# INLINE [0] step' #-} +{-# INLINE postscanlAccStreamM #-} + +-- | Haskell-style scan with a monadic operator +scanlM :: Monad m => (a -> b -> m a) -> a -> Steps m b -> Steps m a +scanlM f acc Steps{stepsStream, stepsSize} = + Steps + { stepsStream = S.scanlM' f acc stepsStream + , stepsSize = addLengthHint (LengthExact 1) stepsSize + } +{-# INLINE scanlM #-} + +-- | Initial-value free scan over a 'Stream' with a strict accumulator +-- and a monadic operator +scanl1M :: Monad m => (a -> a -> m a) -> Steps m a -> Steps m a +scanl1M f ss = ss{stepsStream = S.scanl1M' f (stepsStream ss)} +{-# INLINE scanl1M #-} diff --git a/src/Data/Massiv/Vector/Unsafe.hs b/src/Data/Massiv/Vector/Unsafe.hs --- a/src/Data/Massiv/Vector/Unsafe.hs +++ b/src/Data/Massiv/Vector/Unsafe.hs @@ -1,4 +1,5 @@ {-# LANGUAGE FlexibleContexts #-} + -- | -- Module : Data.Massiv.Vector.Unsafe -- Copyright : (c) Alexey Kuleshevich 2020-2022 @@ -6,23 +7,25 @@ -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable --- -module Data.Massiv.Vector.Unsafe - ( +module Data.Massiv.Vector.Unsafe ( -- * Vector + -- ** Accessors + -- *** Indexing - unsafeHead - , unsafeLast + unsafeHead, + unsafeLast, + -- *** Monadic Indexing - , unsafeIndexM - , unsafeHeadM - , unsafeLastM + unsafeIndexM, + unsafeHeadM, + unsafeLastM, + -- *** Slicing - , unsafeInit - , unsafeTail - , unsafeTake - , unsafeDrop + unsafeInit, + unsafeTail, + unsafeTake, + unsafeDrop, -- -- ** Modifying -- -- *** Bulk updates -- , unsafeUpdate @@ -33,15 +36,16 @@ -- , unsafeBackpermute -- -- ** Predicates -- , unsafePartition + -- ** Unbounded streams - , unsafeUnfoldrN - , unsafeUnfoldrNM - , unsafeFromListN - ) where + unsafeUnfoldrN, + unsafeUnfoldrNM, + unsafeFromListN, +) where import Data.Coerce -import Data.Massiv.Core.Common import Data.Massiv.Array.Delayed.Stream +import Data.Massiv.Core.Common import qualified Data.Massiv.Vector.Stream as S -- ========= -- @@ -52,7 +56,6 @@ -- Indexing -- -------------- - -- | -- -- @since 0.5.0 @@ -78,7 +81,6 @@ unsafeIndexM v i = pure $! unsafeLinearIndex v i {-# INLINE unsafeIndexM #-} - -- | -- -- @since 0.5.0 @@ -93,12 +95,10 @@ unsafeLastM v = pure $! unsafeLast v {-# INLINE unsafeLastM #-} - ------------- -- Slicing -- ------------- - -- | -- -- @since 0.5.0 @@ -106,7 +106,6 @@ unsafeInit v = unsafeLinearSlice 0 (SafeSz (coerce (size v) - 1)) v {-# INLINE unsafeInit #-} - -- | -- -- @since 0.5.0 @@ -114,7 +113,6 @@ unsafeTail = unsafeDrop oneSz {-# INLINE unsafeTail #-} - -- | -- -- @since 0.5.0 @@ -129,7 +127,6 @@ unsafeDrop (Sz d) v = unsafeLinearSlice d (SafeSz (coerce (size v) - d)) v {-# INLINE unsafeDrop #-} - -- | /O(n)/ - Convert a list of a known length to a delayed stream vector. -- -- /Unsafe/ - This function is unsafe because it will allocate enough space in memory for @@ -150,18 +147,17 @@ -- `Control.Exception.Base.HeapOverflow` exception. -- -- @since 0.5.1 -unsafeUnfoldrN :: - Sz1 +unsafeUnfoldrN + :: Sz1 -- ^ @n@ - maximum number of elements that the vector will have -> (s -> Maybe (e, s)) -- ^ Unfolding function. Stops when `Nothing` is returned or maximum number of elements -- is reached. - -> s -- ^ Inititial element. + -> s + -- ^ Inititial element. -> Vector DS e unsafeUnfoldrN n f = DSArray . S.unsafeUnfoldrN n f {-# INLINE unsafeUnfoldrN #-} - - -- | /O(n)/ - Same as `unsafeUnfoldrN`, but with monadic generating function. -- diff --git a/tests/doctests.hs b/tests/doctests.hs --- a/tests/doctests.hs +++ b/tests/doctests.hs @@ -1,4 +1,5 @@ {-# LANGUAGE CPP #-} + module Main where #if __GLASGOW_HASKELL__ >= 802 && __GLASGOW_HASKELL__ < 810