diff --git a/LICENSE b/LICENSE --- a/LICENSE +++ b/LICENSE @@ -1,4 +1,4 @@ -Copyright Alexey Kuleshevich (c) 2017 +Copyright Alexey Kuleshevich (c) 2017-2019 All rights reserved. @@ -13,7 +13,7 @@ disclaimer in the documentation and/or other materials provided with the distribution. - * Neither the name of Author name here nor the names of other + * Neither the name of Alexey Kuleshevich nor the names of other contributors may be used to endorse or promote products derived from this software without specific prior written permission. diff --git a/Setup.hs b/Setup.hs --- a/Setup.hs +++ b/Setup.hs @@ -1,3 +1,33 @@ +{-# LANGUAGE CPP #-} +{-# OPTIONS_GHC -Wall #-} +module Main (main) where + +#ifndef MIN_VERSION_cabal_doctest +#define MIN_VERSION_cabal_doctest(x,y,z) 0 +#endif + +#if MIN_VERSION_cabal_doctest(1,0,0) + +import Distribution.Extra.Doctest ( defaultMainWithDoctests ) +main :: IO () +main = defaultMainWithDoctests "doctests" + +#else + +#ifdef MIN_VERSION_Cabal +-- If the macro is defined, we have new cabal-install, +-- but for some reason we don't have cabal-doctest in package-db +-- +-- Probably we are running cabal sdist, when otherwise using new-build +-- workflow +#warning You are configuring this package without cabal-doctest installed. \ + The doctests test-suite will not work as a result. \ + To fix this, install cabal-doctest before configuring. +#endif + import Distribution.Simple + main :: IO () main = defaultMain + +#endif diff --git a/massiv.cabal b/massiv.cabal --- a/massiv.cabal +++ b/massiv.cabal @@ -1,5 +1,5 @@ name: massiv -version: 0.2.8.1 +version: 0.3.0.0 synopsis: Massiv (Массив) is an Array Library. description: Multi-dimensional Arrays with fusion, stencils and parallel computation. homepage: https://github.com/lehins/massiv @@ -7,9 +7,9 @@ license-file: LICENSE author: Alexey Kuleshevich maintainer: alexey@kuleshevi.ch -copyright: 2018 Alexey Kuleshevich -category: Data, Data Structures -build-type: Simple +copyright: 2018-2019 Alexey Kuleshevich +category: Data, Data Structures, Parallelism +build-type: Custom extra-source-files: README.md cabal-version: >=1.10 @@ -19,6 +19,12 @@ default: False manual: True +custom-setup + setup-depends: + base + , Cabal + , cabal-doctest >=1.0.6 + library hs-source-dirs: src exposed-modules: Data.Massiv.Array @@ -26,17 +32,18 @@ , Data.Massiv.Array.Manifest , Data.Massiv.Array.Manifest.Vector , Data.Massiv.Array.Mutable + , Data.Massiv.Array.Mutable.Atomic , Data.Massiv.Array.Numeric , Data.Massiv.Array.Numeric.Integral , Data.Massiv.Array.Stencil , Data.Massiv.Array.Stencil.Unsafe , Data.Massiv.Array.Unsafe , Data.Massiv.Core - , Data.Massiv.Core.Scheduler , Data.Massiv.Core.Index - other-modules: Data.Massiv.Array.Delayed.Internal - , Data.Massiv.Array.Delayed.Interleaved + other-modules: Data.Massiv.Array.Delayed.Interleaved + , Data.Massiv.Array.Delayed.Pull + , Data.Massiv.Array.Delayed.Push , Data.Massiv.Array.Delayed.Windowed , Data.Massiv.Array.Manifest.Boxed , Data.Massiv.Array.Manifest.Internal @@ -53,18 +60,21 @@ , Data.Massiv.Array.Stencil.Convolution , Data.Massiv.Array.Stencil.Internal , Data.Massiv.Core.Common - , Data.Massiv.Core.Computation - , Data.Massiv.Core.Index.Class + , Data.Massiv.Core.Exception + , Data.Massiv.Core.Index.Internal , Data.Massiv.Core.Index.Ix , Data.Massiv.Core.Index.Stride + , Data.Massiv.Core.Index.Tuple , Data.Massiv.Core.Iterator , Data.Massiv.Core.List - build-depends: base >= 4.8 && < 5 + build-depends: base >= 4.9 && < 5 , bytestring , data-default-class , deepseq - , ghc-prim + , exceptions + , scheduler , primitive + , unliftio-core , vector include-dirs: include install-includes: massiv.h @@ -72,24 +82,24 @@ if flag(unsafe-checks) cpp-options: -DMASSIV_UNSAFE_CHECKS - if impl(ghc < 8.0) - build-depends: semigroups - if impl(ghc < 7.10) - build-depends: unsupported-ghc-version > 1 && < 1 - default-language: Haskell2010 - ghc-options: -Wall - + ghc-options: -Wall + -Wincomplete-record-updates + -Wincomplete-uni-patterns + -Wredundant-constraints test-suite tests type: exitcode-stdio-1.0 hs-source-dirs: tests - main-is: Spec.hs - other-modules: Data.Massiv.Array.DelayedSpec + main-is: Main.hs + other-modules: Data.Massiv.ArraySpec + , Data.Massiv.Array.DelayedSpec + , Data.Massiv.Array.Delayed.InterleavedSpec , Data.Massiv.Array.Delayed.WindowedSpec , Data.Massiv.Array.ManifestSpec , Data.Massiv.Array.Manifest.VectorSpec , Data.Massiv.Array.MutableSpec + , Data.Massiv.Array.Numeric.IntegralSpec , Data.Massiv.Array.Ops.ConstructSpec , Data.Massiv.Array.Ops.FoldSpec , Data.Massiv.Array.Ops.SliceSpec @@ -99,20 +109,39 @@ , Data.Massiv.CoreArbitrary , Data.Massiv.Core.IndexSpec , Data.Massiv.Core.SchedulerSpec - build-depends: base >= 4.8 && < 5 + , Spec + build-depends: base + , async , bytestring - , deepseq + , containers , data-default - , safe-exceptions + , deepseq , massiv + , scheduler , hspec , QuickCheck + , unliftio , vector - if impl(ghc < 8.0) - build-depends: semigroups default-language: Haskell2010 - ghc-options: -Wall -fno-warn-orphans -threaded -with-rtsopts=-N2 + ghc-options: -Wall + -Wincomplete-record-updates + -Wincomplete-uni-patterns + -Wredundant-constraints + -fno-warn-orphans + -threaded + -with-rtsopts=-N2 + +test-suite doctests + type: exitcode-stdio-1.0 + hs-source-dirs: tests + main-is: doctests.hs + build-depends: base + , doctest >=0.15 + , QuickCheck + , massiv + , template-haskell + default-language: Haskell2010 source-repository head type: git diff --git a/src/Data/Massiv/Array.hs b/src/Data/Massiv/Array.hs --- a/src/Data/Massiv/Array.hs +++ b/src/Data/Massiv/Array.hs @@ -1,6 +1,6 @@ -- | -- Module : Data.Massiv.Array --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental @@ -11,9 +11,9 @@ -- sequentially. Depending on the representation (@__r__@), an @__`Array` r ix e__@ will have -- certain properties that are unique to that particular representation, but all of them will share -- the same trait, that an array is simply a mapping from an index (@__ix__@) of an arbitrary --- dimension to an element (@__e__@) of some value. Which means that some of the array types are --- pretty classic and are represented by a contiguous chunk of memory reserved for the elements, --- namely arrays with `Manifest` representations: +-- dimension to an element (@__e__@) of some value. Which means that some representations describe +-- classic arrays and are backed by a contiguous chunk of memory reserved for the elements (or +-- pointers to elements), namely arrays with `Manifest` representations: -- -- * `B` - The most basic type of array that can hold any type of element in a boxed form, i.e. each -- element is a pointer to the actual value, therefore it is also the slowest @@ -36,24 +36,30 @@ -- * `M` - General manifest array type, that any of the above representations can be converted to in -- constant time using `toManifest`. -- --- While at the same time, there are arrays that only describe how values for it's elements can be --- computed, and have no memory overhead on their own. +-- There are also array represnetation that only describe how values for its elements can be +-- computed or loaded into memory, as such, they are represented by functions and do not impose the +-- memory overhead, that is normally associated with arrays. They are needed for proper fusion and +-- parallelization of computation. -- --- * `D` - delayed array that is a mere function from an index to an element. Crucial representation --- for fusing computation. Use `computeAs` in order to load array into `Manifest` --- representation. +-- * `D` - delayed array that is a mere function from an index to an element. Also known as /Pull/ +-- array. Crucial representation for fusing computation. Use `computeAs` in order to load +-- array into `Manifest` representation. -- --- * `DI` - delayed interleaved array. Same as `D`, but performced better with unbalanced --- computation, when evaluation one element takes much longer than it's neighbor. +-- * `DL` - delayed load array representation that describes how an array can be loaded. Also known +-- as /Push/ array. Useful for fusing various array combining functions. Use `computeAs` in +-- order to load array into `Manifest` representation. -- +-- * `DI` - delayed interleaved array. Same as `D`, but performes better with unbalanced +-- computation, when evaluation of one element takes much longer than of its neighbor. +-- -- * `DW` - delayed windowed array. This peculiar representation allows for very fast `Stencil` -- computation. -- -- Other Array types: -- -- * `L` and `LN` - those types aren't particularly useful on their own, but because of their unique --- ability to be converted to and from nested lists in constant time, provide an amazing --- intermediary for list/array conversion. +-- ability to be converted to and from nested lists in constant time, provide a perfect +-- intermediary for lists <-> array conversion. -- -- Most of the `Manifest` arrays are capable of in-place mutation. Check out -- "Data.Massiv.Array.Mutable" module for available functionality. @@ -73,6 +79,7 @@ , getComp , setComp , compute + , computeS , computeAs , computeProxy , computeSource @@ -92,10 +99,13 @@ , (!?) , (!) , (??) + , indexM , index , index' , defaultIndex , borderIndex + , evaluateM + , evaluate' , evaluateAt -- * Mapping , module Data.Massiv.Array.Ops.Map @@ -123,35 +133,23 @@ , module Data.Massiv.Array.Numeric ) where -import Data.Massiv.Array.Delayed -import Data.Massiv.Array.Manifest -import Data.Massiv.Array.Manifest.Internal -import Data.Massiv.Array.Manifest.List -import Data.Massiv.Array.Mutable -import Data.Massiv.Array.Numeric -import Data.Massiv.Array.Ops.Construct -import Data.Massiv.Array.Ops.Fold -import Data.Massiv.Array.Ops.Map -import Data.Massiv.Array.Ops.Slice -import Data.Massiv.Array.Ops.Transform -import Data.Massiv.Array.Stencil -import Data.Massiv.Core -import Data.Massiv.Core.Common -import Prelude as P hiding ( all - , and - , any - , foldl - , foldr - , mapM - , mapM_ - , maximum - , minimum - , or - , product - , splitAt - , sum - , zip - ) +import Data.Massiv.Array.Delayed +import Data.Massiv.Array.Manifest +import Data.Massiv.Array.Manifest.Internal +import Data.Massiv.Array.Manifest.List +import Data.Massiv.Array.Mutable +import Data.Massiv.Array.Numeric +import Data.Massiv.Array.Ops.Construct +import Data.Massiv.Array.Ops.Fold +import Data.Massiv.Array.Ops.Map +import Data.Massiv.Array.Ops.Slice +import Data.Massiv.Array.Ops.Transform +import Data.Massiv.Array.Stencil +import Data.Massiv.Core +import Data.Massiv.Core.Common +import Prelude as P hiding (all, and, any, enumFromTo, foldl, foldr, mapM, + mapM_, maximum, minimum, or, product, replicate, splitAt, + sum, zip) {- $folding diff --git a/src/Data/Massiv/Array/Delayed.hs b/src/Data/Massiv/Array/Delayed.hs --- a/src/Data/Massiv/Array/Delayed.hs +++ b/src/Data/Massiv/Array/Delayed.hs @@ -1,24 +1,35 @@ -- | -- Module : Data.Massiv.Array.Delayed --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable -- module Data.Massiv.Array.Delayed - ( D(..) + ( -- * Delayed + -- ** Delayed Pull Array + D(..) , delay + -- ** Delayed Push Array + , DL(..) + , toLoadArray + , makeLoadArray + , fromStrideLoad + -- ** Delayed Interleaved Array , DI(..) , toInterleaved , fromInterleaved + -- ** Delayed Windowed Array , DW(..) , Window(..) + , insertWindow , getWindow + , dropWindow , makeWindowedArray ) where -import Data.Massiv.Array.Delayed.Interleaved -import Data.Massiv.Array.Delayed.Internal -import Data.Massiv.Array.Delayed.Windowed - +import Data.Massiv.Array.Delayed.Interleaved +import Data.Massiv.Array.Delayed.Pull +import Data.Massiv.Array.Delayed.Push +import Data.Massiv.Array.Delayed.Windowed diff --git a/src/Data/Massiv/Array/Delayed/Interleaved.hs b/src/Data/Massiv/Array/Delayed/Interleaved.hs --- a/src/Data/Massiv/Array/Delayed/Interleaved.hs +++ b/src/Data/Massiv/Array/Delayed/Interleaved.hs @@ -1,12 +1,14 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} -- | -- Module : Data.Massiv.Array.Delayed.Interleaved --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental @@ -18,9 +20,9 @@ , fromInterleaved ) where -import Data.Massiv.Array.Delayed.Internal -import Data.Massiv.Core.Common -import Data.Massiv.Core.Scheduler +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 fasion during parallel @@ -29,59 +31,53 @@ type instance EltRepr DI ix = DI -newtype instance Array DI ix e = DIArray { diArray :: Array D ix e } +newtype instance Array DI ix e = DIArray + { diArray :: Array D ix e + } deriving (Eq, Ord, Functor, Applicative, Foldable, Num, Floating, Fractional) -instance Index ix => Construct DI ix e where - getComp = dComp . diArray - {-# INLINE getComp #-} +instance (Ragged L ix e, Show e) => Show (Array DI ix e) where + showsPrec = showsArrayPrec diArray + showList = showArrayList +instance Index ix => Construct DI ix e where setComp c arr = arr { diArray = (diArray arr) { dComp = c } } {-# INLINE setComp #-} - unsafeMakeArray c sz = DIArray . unsafeMakeArray c sz - {-# INLINE unsafeMakeArray #-} - -instance Functor (Array DI ix) where - fmap f (DIArray arr) = DIArray (fmap f arr) - -instance Index ix => Size DI ix e where - size (DIArray arr) = size arr - {-# INLINE size #-} + makeArray c sz = DIArray . makeArray c sz + {-# INLINE makeArray #-} +instance Index ix => Resize DI ix where unsafeResize sz = DIArray . unsafeResize sz . diArray {-# INLINE unsafeResize #-} +instance Index ix => Extract DI ix e where unsafeExtract sIx newSz = DIArray . unsafeExtract sIx newSz . diArray {-# INLINE unsafeExtract #-} instance Index ix => Load DI ix e where - loadS (DIArray arr) = loadS arr - {-# INLINE loadS #-} - loadP wIds (DIArray (DArray _ sz f)) _ unsafeWrite = - withScheduler_ wIds $ \scheduler -> do - let !totalLength = totalElem sz - loopM_ 0 (< numWorkers scheduler) (+ 1) $ \ !start -> - scheduleWork scheduler $ - iterLinearM_ sz start totalLength (numWorkers scheduler) (<) $ \ !k !ix -> - unsafeWrite k $ f ix - {-# INLINE loadP #-} - loadArray numWorkers' scheduleWork' (DIArray (DArray _ sz f)) _ unsafeWrite = - loopM_ 0 (< numWorkers') (+ 1) $ \ !start -> - scheduleWork' $ - iterLinearM_ sz start (totalElem sz) numWorkers' (<) $ \ !k -> unsafeWrite k . f - {-# INLINE loadArray #-} - loadArrayWithStride numWorkers' scheduleWork' stride resultSize arr _ unsafeWrite = + size (DIArray arr) = size arr + {-# INLINE size #-} + getComp = dComp . diArray + {-# INLINE getComp #-} + loadArrayM Scheduler {numWorkers, scheduleWork} (DIArray (DArray _ sz f)) uWrite = + loopM_ 0 (< numWorkers) (+ 1) $ \ !start -> + scheduleWork $ + iterLinearM_ sz start (totalElem sz) numWorkers (<) $ \ !k -> uWrite k . f + {-# INLINE loadArrayM #-} + +instance Index ix => StrideLoad DI ix e where + loadArrayWithStrideM Scheduler {numWorkers, scheduleWork} stride resultSize arr uWrite = let strideIx = unStride stride DIArray (DArray _ _ f) = arr - in loopM_ 0 (< numWorkers') (+ 1) $ \ !start -> - scheduleWork' $ - iterLinearM_ resultSize start (totalElem resultSize) numWorkers' (<) $ - \ !i ix -> unsafeWrite i (f (liftIndex2 (*) strideIx ix)) - {-# INLINE loadArrayWithStride #-} + in loopM_ 0 (< numWorkers) (+ 1) $ \ !start -> + scheduleWork $ + iterLinearM_ resultSize start (totalElem resultSize) numWorkers (<) $ + \ !i ix -> uWrite i (f (liftIndex2 (*) strideIx ix)) + {-# INLINE loadArrayWithStrideM #-} -- | Convert a source array into an array that, when computed, will have its elemets evaluated out --- of order (interleaved amoungs cores), hence making unbalanced computation better parallelizable. +-- of order (interleaved amongst cores), hence making unbalanced computation better parallelizable. toInterleaved :: Source r ix e => Array r ix e -> Array DI ix e toInterleaved = DIArray . delay {-# INLINE toInterleaved #-} diff --git a/src/Data/Massiv/Array/Delayed/Internal.hs b/src/Data/Massiv/Array/Delayed/Internal.hs deleted file mode 100644 --- a/src/Data/Massiv/Array/Delayed/Internal.hs +++ /dev/null @@ -1,267 +0,0 @@ -{-# LANGUAGE NamedFieldPuns #-} -{-# LANGUAGE NamedFieldPuns #-} -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE CPP #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE UndecidableInstances #-} --- | --- Module : Data.Massiv.Array.Delayed.Internal --- Copyright : (c) Alexey Kuleshevich 2018 --- License : BSD3 --- Maintainer : Alexey Kuleshevich --- Stability : experimental --- Portability : non-portable --- -module Data.Massiv.Array.Delayed.Internal - ( D(..) - , Array(..) - , delay - , eq - , ord - , liftArray - , liftArray2 - ) where - -import Data.Foldable (Foldable (..)) -import Data.Massiv.Array.Ops.Fold.Internal as A -import Data.Massiv.Core.Common -import Data.Massiv.Core.Scheduler -import Data.Monoid ((<>)) -import GHC.Base (build) -import Prelude hiding (zipWith) - -#include "massiv.h" - --- | Delayed representation. -data D = D deriving Show - - -data instance Array D ix e = DArray { dComp :: !Comp - , dSize :: !ix - , dIndex :: ix -> e } -type instance EltRepr D ix = D - -instance Index ix => Construct D ix e where - getComp = dComp - {-# INLINE getComp #-} - - setComp c arr = arr { dComp = c } - {-# INLINE setComp #-} - - unsafeMakeArray = DArray - {-# INLINE unsafeMakeArray #-} - - -instance Index ix => Source D ix e where - unsafeIndex = INDEX_CHECK("(Source D ix e).unsafeIndex", size, dIndex) - {-# INLINE unsafeIndex #-} - -instance Index ix => Size D ix e where - size = dSize - {-# INLINE size #-} - - unsafeResize !sz !arr = - DArray (getComp arr) sz $ \ !ix -> - unsafeIndex arr (fromLinearIndex (size arr) (toLinearIndex sz ix)) - {-# INLINE unsafeResize #-} - - unsafeExtract !sIx !newSz !arr = - DArray (getComp arr) newSz $ \ !ix -> - unsafeIndex arr (liftIndex2 (+) ix sIx) - {-# INLINE unsafeExtract #-} - -instance ( Index ix - , Index (Lower ix) - , Elt D ix e ~ Array D (Lower ix) e - ) => - Slice D ix e where - unsafeSlice arr start cutSz dim = do - newSz <- dropDim cutSz dim - return $ unsafeResize newSz (unsafeExtract start cutSz arr) - {-# INLINE unsafeSlice #-} - - -instance (Elt D ix e ~ Array D (Lower ix) e, Index ix) => OuterSlice D ix e where - - unsafeOuterSlice !arr !i = - DArray (getComp arr) (tailDim (size arr)) (\ !ix -> unsafeIndex arr (consDim i ix)) - {-# INLINE unsafeOuterSlice #-} - -instance (Elt D ix e ~ Array D (Lower ix) e, Index ix) => InnerSlice D ix e where - - unsafeInnerSlice !arr !(szL, _) !i = - DArray (getComp arr) szL (\ !ix -> unsafeIndex arr (snocDim ix i)) - {-# INLINE unsafeInnerSlice #-} - - -instance (Eq e, Index ix) => Eq (Array D ix e) where - (==) = eq (==) - {-# INLINE (==) #-} - -instance (Ord e, Index ix) => Ord (Array D ix e) where - compare = ord compare - {-# INLINE compare #-} - -instance Functor (Array D ix) where - fmap f (DArray c sz g) = DArray c sz (f . g) - {-# INLINE fmap #-} - - -instance Index ix => Applicative (Array D ix) where - pure a = DArray Seq (liftIndex (+ 1) zeroIndex) (const a) - {-# INLINE pure #-} - (<*>) (DArray c1 sz1 uIndex1) (DArray c2 sz2 uIndex2) = - DArray (c1 <> c2) (liftIndex2 min sz1 sz2) $ \ !ix -> - (uIndex1 ix) (uIndex2 ix) - {-# INLINE (<*>) #-} - - --- | Row-major sequential folding over a Delayed array. -instance Index ix => Foldable (Array D ix) where - foldl = lazyFoldlS - {-# INLINE foldl #-} - foldl' = foldlS - {-# INLINE foldl' #-} - foldr = foldrFB - {-# INLINE foldr #-} - foldr' = foldrS - {-# INLINE foldr' #-} - null (DArray _ sz _) = totalElem sz == 0 - {-# INLINE null #-} - sum = foldl' (+) 0 - {-# INLINE sum #-} - product = foldl' (*) 1 - {-# INLINE product #-} - length = totalElem . size - {-# INLINE length #-} - toList arr = build (\ c n -> foldrFB c n arr) - {-# INLINE toList #-} - - -instance Index ix => Load D ix e where - loadS (DArray _ sz f) _ unsafeWrite = - iterM_ zeroIndex sz (pureIndex 1) (<) $ \ !ix -> unsafeWrite (toLinearIndex sz ix) (f ix) - {-# INLINE loadS #-} - loadP wIds (DArray _ sz f) _ unsafeWrite = - divideWork_ wIds sz $ \ !scheduler !chunkLength !totalLength !slackStart -> do - loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start -> - scheduleWork scheduler $ - iterLinearM_ sz start (start + chunkLength) 1 (<) $ \ !k !ix -> unsafeWrite k (f ix) - scheduleWork scheduler $ - iterLinearM_ sz slackStart totalLength 1 (<) $ \ !k !ix -> unsafeWrite k (f ix) - {-# INLINE loadP #-} - - -instance (Index ix, Num e) => Num (Array D ix e) where - (+) = liftArray2 (+) - {-# INLINE (+) #-} - (-) = liftArray2 (-) - {-# INLINE (-) #-} - (*) = liftArray2 (*) - {-# INLINE (*) #-} - abs = liftArray abs - {-# INLINE abs #-} - signum = liftArray signum - {-# INLINE signum #-} - fromInteger = singleton Seq . fromInteger - {-# INLINE fromInteger #-} - -instance (Index ix, Fractional e) => Fractional (Array D ix e) where - (/) = liftArray2 (/) - {-# INLINE (/) #-} - fromRational = singleton Seq . fromRational - {-# INLINE fromRational #-} - - -instance (Index ix, Floating e) => Floating (Array D ix e) where - pi = singleton Seq pi - {-# INLINE pi #-} - exp = liftArray exp - {-# INLINE exp #-} - log = liftArray log - {-# INLINE log #-} - sin = liftArray sin - {-# INLINE sin #-} - cos = liftArray cos - {-# INLINE cos #-} - asin = liftArray asin - {-# INLINE asin #-} - atan = liftArray atan - {-# INLINE atan #-} - acos = liftArray acos - {-# INLINE acos #-} - sinh = liftArray sinh - {-# INLINE sinh #-} - cosh = liftArray cosh - {-# INLINE cosh #-} - asinh = liftArray asinh - {-# INLINE asinh #-} - atanh = liftArray atanh - {-# INLINE atanh #-} - acosh = liftArray acosh - {-# INLINE acosh #-} - - - --- | /O(1)/ Conversion from a source array to `D` representation. -delay :: Source r ix e => Array r ix e -> Array D ix e -delay arr = DArray (getComp arr) (size arr) (unsafeIndex arr) -{-# INLINE delay #-} - - --- | /O(n1 + n2)/ - Compute array equality by applying a comparing function to each element. -eq :: (Source r1 ix e1, Source r2 ix e2) => - (e1 -> e2 -> Bool) -> Array r1 ix e1 -> Array r2 ix e2 -> Bool -eq f arr1 arr2 = - (size arr1 == size arr2) && - A.fold - (&&) - True - (DArray (getComp arr1 <> getComp arr2) (size arr1) $ \ix -> - f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)) -{-# INLINE eq #-} - --- | /O(n1 + n2)/ - Compute array ordering by applying a comparing function to each element. --- The exact ordering is unspecified so this is only intended for use in maps and the like where --- you need an ordering but do not care about which one is used. -ord :: (Source r1 ix e1, Source r2 ix e2) => - (e1 -> e2 -> Ordering) -> Array r1 ix e1 -> Array r2 ix e2 -> Ordering -ord f arr1 arr2 = - (compare (size arr1) (size arr2)) <> - A.fold - (<>) - mempty - (DArray (getComp arr1 <> getComp arr2) (size arr1) $ \ix -> - f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)) -{-# INLINE ord #-} - --- | The usual map. -liftArray :: Source r ix b => (b -> e) -> Array r ix b -> Array D ix e -liftArray f !arr = DArray (getComp arr) (size arr) (f . unsafeIndex arr) -{-# INLINE liftArray #-} - --- | Similar to `Data.Massiv.Array.zipWith`, except dimensions of both arrays either have to be the --- same, or at least one of the two array must be a singleton array, in which case it will behave as --- a `Data.Massiv.Array.map`. --- --- @since 0.1.4 -liftArray2 - :: (Source r1 ix a, Source r2 ix b) - => (a -> b -> e) -> Array r1 ix a -> Array r2 ix b -> Array D ix e -liftArray2 f !arr1 !arr2 - | sz1 == oneIndex = liftArray (f (unsafeIndex arr1 zeroIndex)) arr2 - | sz2 == oneIndex = liftArray (`f` (unsafeIndex arr2 zeroIndex)) arr1 - | sz1 == sz2 = - DArray (getComp arr1) sz1 (\ !ix -> f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)) - | otherwise = errorSizeMismatch "liftArray2" (size arr1) (size arr2) - where - oneIndex = pureIndex 1 - sz1 = size arr1 - sz2 = size arr2 -{-# INLINE liftArray2 #-} - - diff --git a/src/Data/Massiv/Array/Delayed/Pull.hs b/src/Data/Massiv/Array/Delayed/Pull.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Array/Delayed/Pull.hs @@ -0,0 +1,262 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} +-- | +-- Module : Data.Massiv.Array.Delayed.Pull +-- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Array.Delayed.Pull + ( D(..) + , Array(..) + , delay + , eq + , ord + , liftArray + , liftArray2 + ) where + +import qualified Data.Foldable as F +import Data.Massiv.Array.Ops.Fold.Internal as A +import Data.Massiv.Core.Common +import Data.Massiv.Core.Index.Internal +import Data.Massiv.Core.List (L, showArrayList, showsArrayPrec) +import GHC.Base (build) +import Prelude hiding (zipWith) + +#include "massiv.h" + +-- | Delayed representation. +data D = D deriving Show + +data instance Array D ix e = DArray { dComp :: !Comp + , dSize :: !(Sz ix) + , dIndex :: ix -> e } +type instance EltRepr D ix = D + +instance (Ragged L ix e, Show e) => Show (Array D ix e) where + showsPrec = showsArrayPrec id + showList = showArrayList + +instance Index ix => Resize D ix where + unsafeResize !sz !arr = + DArray (dComp arr) sz $ \ !ix -> + unsafeIndex arr (fromLinearIndex (size arr) (toLinearIndex sz ix)) + {-# INLINE unsafeResize #-} + +instance Index ix => Extract D ix e where + unsafeExtract !sIx !newSz !arr = + DArray (dComp arr) newSz $ \ !ix -> + unsafeIndex arr (liftIndex2 (+) ix sIx) + {-# INLINE unsafeExtract #-} + + +instance Index ix => Construct D ix e where + setComp c arr = arr { dComp = c } + {-# INLINE setComp #-} + + makeArray = DArray + {-# INLINE makeArray #-} + + +instance Index ix => Source D ix e where + unsafeIndex = INDEX_CHECK("(Source D ix e).unsafeIndex", size, dIndex) + {-# INLINE unsafeIndex #-} + + +instance ( Index ix + , Index (Lower ix) + , Elt D ix e ~ Array D (Lower ix) e + ) => + Slice D ix e where + unsafeSlice arr start cut@(SafeSz cutSz) dim = do + newSz <- dropDimM cutSz dim + return $ unsafeResize (SafeSz newSz) (unsafeExtract start cut arr) + {-# INLINE unsafeSlice #-} + + +instance (Elt D ix e ~ Array D (Lower ix) e, Index ix) => OuterSlice D ix e where + + unsafeOuterSlice !arr !i = + DArray (dComp arr) (snd (unconsSz (size arr))) (\ !ix -> unsafeIndex arr (consDim i ix)) + {-# INLINE unsafeOuterSlice #-} + +instance (Elt D ix e ~ Array D (Lower ix) e, Index ix) => InnerSlice D ix e where + + unsafeInnerSlice !arr (szL, _) !i = + DArray (dComp arr) szL (\ !ix -> unsafeIndex arr (snocDim ix i)) + {-# INLINE unsafeInnerSlice #-} + + +instance (Eq e, Index ix) => Eq (Array D ix e) where + (==) = eq (==) + {-# INLINE (==) #-} + +instance (Ord e, Index ix) => Ord (Array D ix e) where + compare = ord compare + {-# INLINE compare #-} + +instance Functor (Array D ix) where + fmap f (DArray c sz g) = DArray c sz (f . g) + {-# INLINE fmap #-} + + +instance Index ix => Applicative (Array D ix) where + pure = singleton + {-# INLINE pure #-} + (<*>) (DArray c1 (SafeSz sz1) uIndex1) (DArray c2 (SafeSz sz2) uIndex2) = + DArray (c1 <> c2) (SafeSz (liftIndex2 min sz1 sz2)) $ \ !ix -> + (uIndex1 ix) (uIndex2 ix) + {-# INLINE (<*>) #-} + + +-- | Row-major sequential folding over a Delayed array. +instance Index ix => Foldable (Array D ix) where + fold = A.fold + {-# INLINE fold #-} + foldMap = A.foldMono + {-# INLINE foldMap #-} + foldl = lazyFoldlS + {-# INLINE foldl #-} + foldl' = foldlS + {-# INLINE foldl' #-} + foldr = foldrFB + {-# INLINE foldr #-} + foldr' = foldrS + {-# INLINE foldr' #-} + null (DArray _ sz _) = totalElem sz == 0 + {-# INLINE null #-} + length = totalElem . size + {-# INLINE length #-} + toList arr = build (\ c n -> foldrFB c n arr) + {-# INLINE toList #-} + + +instance Index ix => Load D ix e where + size = dSize + {-# INLINE size #-} + getComp = dComp + {-# INLINE getComp #-} + loadArrayM !scheduler !arr = + splitLinearlyWith_ scheduler (elemsCount arr) (unsafeLinearIndex arr) + {-# INLINE loadArrayM #-} + +instance Index ix => StrideLoad D ix e + +instance (Index ix, Num e) => Num (Array D ix e) where + (+) = liftArray2 (+) + {-# INLINE (+) #-} + (-) = liftArray2 (-) + {-# INLINE (-) #-} + (*) = liftArray2 (*) + {-# INLINE (*) #-} + abs = liftArray abs + {-# INLINE abs #-} + signum = liftArray signum + {-# INLINE signum #-} + fromInteger = singleton . fromInteger + {-# INLINE fromInteger #-} + +instance (Index ix, Fractional e) => Fractional (Array D ix e) where + (/) = liftArray2 (/) + {-# INLINE (/) #-} + fromRational = singleton . fromRational + {-# INLINE fromRational #-} + + +instance (Index ix, Floating e) => Floating (Array D ix e) where + pi = singleton pi + {-# INLINE pi #-} + exp = liftArray exp + {-# INLINE exp #-} + log = liftArray log + {-# INLINE log #-} + sin = liftArray sin + {-# INLINE sin #-} + cos = liftArray cos + {-# INLINE cos #-} + asin = liftArray asin + {-# INLINE asin #-} + atan = liftArray atan + {-# INLINE atan #-} + acos = liftArray acos + {-# INLINE acos #-} + sinh = liftArray sinh + {-# INLINE sinh #-} + cosh = liftArray cosh + {-# INLINE cosh #-} + asinh = liftArray asinh + {-# INLINE asinh #-} + atanh = liftArray atanh + {-# INLINE atanh #-} + acosh = liftArray acosh + {-# INLINE acosh #-} + + + +-- | /O(1)/ Conversion from a source array to `D` representation. +delay :: Source r ix e => Array r ix e -> Array D ix e +delay arr = DArray (getComp arr) (size arr) (unsafeIndex arr) +{-# INLINE [1] delay #-} + +{-# RULES +"delay" [~1] forall (arr :: Array D ix e) . delay arr = arr + #-} + +-- TODO: switch to zipWith +-- | /O(min (n1, n2))/ - Compute array equality by applying a comparing function to each element. +eq :: (Source r1 ix e1, Source r2 ix e2) => + (e1 -> e2 -> Bool) -> Array r1 ix e1 -> Array r2 ix e2 -> Bool +eq f arr1 arr2 = + (size arr1 == size arr2) && + F.and + (DArray (getComp arr1 <> getComp arr2) (size arr1) $ \ix -> + f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)) +{-# INLINE eq #-} + +-- | /O(min (n1, n2))/ - Compute array ordering by applying a comparing function to each element. +-- The exact ordering is unspecified so this is only intended for use in maps and the like where +-- you need an ordering but do not care about which one is used. +ord :: (Source r1 ix e1, Source r2 ix e2) => + (e1 -> e2 -> Ordering) -> Array r1 ix e1 -> Array r2 ix e2 -> Ordering +ord f arr1 arr2 = + compare (size arr1) (size arr2) <> + A.fold + (DArray (getComp arr1 <> getComp arr2) (size arr1) $ \ix -> + f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)) +{-# INLINE ord #-} + +-- | The usual map. +liftArray :: Source r ix b => (b -> e) -> Array r ix b -> Array D ix e +liftArray f !arr = DArray (getComp arr) (size arr) (f . unsafeIndex arr) +{-# INLINE liftArray #-} + +-- | Similar to `Data.Massiv.Array.zipWith`, except dimensions of both arrays either have to be the +-- same, or at least one of the two array must be a singleton array, in which case it will behave as +-- a `Data.Massiv.Array.map`. +-- +-- @since 0.1.4 +liftArray2 + :: (Source r1 ix a, Source r2 ix b) + => (a -> b -> e) -> Array r1 ix a -> Array r2 ix b -> Array D ix e +liftArray2 f !arr1 !arr2 + | sz1 == oneSz = liftArray (f (unsafeIndex arr1 zeroIndex)) arr2 + | sz2 == oneSz = liftArray (`f` unsafeIndex arr2 zeroIndex) arr1 + | sz1 == sz2 = + DArray (getComp arr1 <> getComp arr2) sz1 (\ !ix -> f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)) + | otherwise = throw $ SizeMismatchException (size arr1) (size arr2) + where + sz1 = size arr1 + sz2 = size arr2 +{-# INLINE liftArray2 #-} + + diff --git a/src/Data/Massiv/Array/Delayed/Push.hs b/src/Data/Massiv/Array/Delayed/Push.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Array/Delayed/Push.hs @@ -0,0 +1,129 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} +-- | +-- Module : Data.Massiv.Array.Delayed.Push +-- Copyright : (c) Alexey Kuleshevich 2019 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Array.Delayed.Push + ( DL(..) + , Array(..) + , toLoadArray + , makeLoadArray + , fromStrideLoad + ) where + +import Data.Massiv.Core.Common +import Data.Massiv.Core.Index.Internal (Sz(SafeSz)) +import qualified Data.Semigroup as Semigroup +import Prelude hiding (map, zipWith) + +#include "massiv.h" + +-- | Delayed load representation. Also known as Push array. +data DL = DL deriving Show + + +data instance Array DL ix e = DLArray + { dlComp :: !Comp + , dlSize :: !(Sz ix) + , dlLoad :: forall m . Monad m + => Scheduler m () + -> Int -- start loading at this linear index + -> (Int -> e -> m ()) -- linear element writing action + -> m () + } + +type instance EltRepr DL ix = DL + +instance Index ix => Construct DL ix e where + setComp c arr = arr {dlComp = c} + {-# INLINE setComp #-} + makeArrayLinear comp sz f = + DLArray comp sz $ \scheduler startAt dlWrite -> + splitLinearlyWithStartAtM_ scheduler startAt (totalElem sz) (pure . f) dlWrite + {-# INLINE makeArrayLinear #-} + +instance Index ix => Resize DL ix where + unsafeResize !sz arr = arr { dlSize = sz } + {-# INLINE unsafeResize #-} + +instance Semigroup (Array DL Ix1 e) where + (<>) (DLArray c1 sz1 load1) (DLArray c2 sz2 load2) = + DLArray {dlComp = c1 <> c2, dlSize = SafeSz (k + unSz sz2), dlLoad = load} + where + !k = unSz sz1 + load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m () + load scheduler startAt dlWrite = do + load1 scheduler startAt dlWrite + load2 scheduler (startAt + k) dlWrite + {-# INLINE load #-} + {-# INLINE (<>) #-} + + +instance Monoid (Array DL Ix1 e) where + mempty = makeArray Seq zeroSz (const (throwImpossible Uninitialized)) + {-# INLINE mempty #-} + + mappend = (Semigroup.<>) + {-# INLINE mappend #-} + + +-- | Specify how an array can be loaded/computed through creation of a `DL` array. +-- +-- @since 0.3.0 +makeLoadArray :: + Comp + -> Sz ix + -> (forall m. Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()) + -> Array DL ix e +makeLoadArray = DLArray +{-# INLINE makeLoadArray #-} + +-- | Convert any `Load`able array into `DL` representation. +-- +-- @since 0.3.0 +toLoadArray :: Load r ix e => Array r ix e -> Array DL ix e +toLoadArray arr = + DLArray (getComp arr) (size arr) $ \scheduler startAt dlWrite -> + loadArrayM scheduler arr (\ !i -> dlWrite (i + startAt)) +{-# INLINE toLoadArray #-} + +-- | Convert an array that can be loaded with stride into `DL` representation. +-- +-- @since 0.3.0 +fromStrideLoad + :: StrideLoad r ix e => Stride ix -> Array r ix e -> Array DL ix e +fromStrideLoad stride arr = + DLArray (getComp arr) newsz $ \scheduler startAt dlWrite -> + loadArrayWithStrideM scheduler stride newsz arr (\ !i -> dlWrite (i + startAt)) + where + newsz = strideSize stride (size arr) +{-# INLINE fromStrideLoad #-} + +instance Index ix => Load DL ix e where + size = dlSize + {-# INLINE size #-} + getComp = dlComp + {-# INLINE getComp #-} + loadArrayM scheduler DLArray {dlLoad} = dlLoad scheduler 0 + {-# INLINE loadArrayM #-} + +instance Functor (Array DL ix) where + fmap f arr = + arr + { dlLoad = + \scheduler startAt uWrite -> dlLoad arr scheduler startAt (\ !i e -> uWrite i (f e)) + } + {-# INLINE fmap #-} 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 @@ -1,16 +1,16 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE NamedFieldPuns #-} -{-# LANGUAGE RecordWildCards #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} -- | -- Module : Data.Massiv.Array.Delayed.Windowed --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental @@ -20,73 +20,66 @@ ( DW(..) , Array(..) , Window(..) + , insertWindow , getWindow + , dropWindow , makeWindowedArray ) where -import Control.Monad (when) -import Data.Massiv.Array.Delayed.Internal -import Data.Massiv.Array.Manifest.Boxed -import Data.Massiv.Array.Manifest.Internal -import Data.Massiv.Core -import Data.Massiv.Core.Common -import Data.Massiv.Core.List (showArray) -import Data.Massiv.Core.Scheduler -import Data.Maybe (fromMaybe) -import Data.Proxy (Proxy (..)) -import Data.Typeable (showsTypeRep, typeRep) -import GHC.TypeLits +import Control.Exception (Exception(..)) +import Control.Monad (when) +import Data.Massiv.Array.Delayed.Pull +import Data.Massiv.Array.Manifest.Boxed +import Data.Massiv.Array.Manifest.Internal +import Data.Massiv.Core +import Data.Massiv.Core.Common +import Data.Massiv.Core.Index.Internal (Sz(..)) +import Data.Massiv.Core.List (L, showArrayList, showsArrayPrec) +import Data.Maybe (fromMaybe) +import GHC.TypeLits -- | Delayed Windowed Array representation. data DW = DW type instance EltRepr DW ix = D -data Window ix e = Window { windowStart :: !ix +data Window ix e = Window { windowStart :: !ix -- ^ Index of where window will start at. - , windowSize :: !ix + , windowSize :: !(Sz ix) -- ^ Size of the window - , windowIndex :: ix -> e + , 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 } data instance Array DW ix e = DWArray { dwArray :: !(Array D ix e) - , dwStencilSize :: !(Maybe ix) - -- ^ Setting this value during stencil - -- application improves cache utilization - -- while computing an array , dwWindow :: !(Maybe (Window ix e)) } -instance {-# OVERLAPPING #-} (Show e, Ragged L ix e, Load DW ix e) => - Show (Array DW ix e) where - show arr = showArray (showsTypeRep (typeRep (Proxy :: Proxy DW)) " ") (computeAs B arr) +instance (Ragged L ix e, Load DW ix e, Show e) => Show (Array DW ix e) where + showsPrec = showsArrayPrec (computeAs B) + showList = showArrayList instance Index ix => Construct DW ix e where - getComp = dComp . dwArray - {-# INLINE getComp #-} setComp c arr = arr { dwArray = (dwArray arr) { dComp = c } } {-# INLINE setComp #-} - unsafeMakeArray c sz f = - DWArray (unsafeMakeArray c sz f) Nothing Nothing - {-# INLINE unsafeMakeArray #-} + makeArray c sz f = DWArray (makeArray c sz f) Nothing + {-# INLINE makeArray #-} --- | Any resize or extract on Windowed Array will loose the interior window and all other --- optimizations, thus hurting the performance a lot. -instance Index ix => Size DW ix e where - size = dSize . dwArray - {-# INLINE size #-} - unsafeResize sz arr = arr { dwArray = unsafeResize sz (dwArray arr) - , dwWindow = Nothing - , dwStencilSize = Nothing } - unsafeExtract sIx newSz = unsafeExtract sIx newSz . dwArray +-- TODO: adjust in response to Window +-- instance Index ix => Extract DW ix e where +-- unsafeExtract sIx newSz = unsafeExtract sIx newSz . dwArray +-- {-# INLINE unsafeExtract #-} instance Functor (Array DW ix) where @@ -98,57 +91,115 @@ {-# INLINE fmap #-} --- | Supply a separate generating function for interior of an array. This is --- very usful for stencil mapping, where interior function does not perform --- boundary checks, thus significantly speeding up computation process. -- +-- +-- @since 0.3.0 +-- _makeWindowedArrayM +-- :: Source r ix e +-- => Array r ix e -- ^ Source array that will have a window inserted into it +-- -> ix -- ^ Start index for the window +-- -> Sz ix -- ^ Size of the window +-- -> (ix -> e) -- ^ Inside window indexing function +-- -> Array DW ix e +-- _makeWindowedArrayM !arr !windowStart !windowSize windowIndex +-- | not (isSafeIndex sz windowStart) = +-- error $ +-- "makeWindowedArray: Incorrect window starting index: (" ++ +-- show windowStart ++ ") for array size: (" ++ show (size arr) ++ ")" +-- | totalElem windowSize == 0 = +-- error $ +-- "makeWindowedArray: Window can't hold any elements with this size: (" ++ show windowSize ++ ")" +-- | not +-- (isSafeIndex +-- (Sz (liftIndex (+ 1) (unSz sz))) +-- (liftIndex2 (+) windowStart (unSz windowSize))) = +-- error $ +-- "makeWindowedArray: Incorrect window size: (" ++ +-- show windowSize ++ +-- ") and/or starting index: (" ++ +-- show windowStart ++ ") for array size: (" ++ show (size arr) ++ ")" +-- | otherwise = +-- DWArray {dwArray = delay arr, dwWindow = Just $! Window {..}} +-- where +-- windowUnrollIx2 = Nothing +-- sz = size arr +-- {-# INLINE _makeWindowedArrayM #-} + +-- | Construct a delayed windowed array by supply a separate element producing function for the +-- interior of an array. This is very usful for stencil mapping, where interior function does not +-- perform boundary checks, thus significantly speeding up computation process. +-- -- @since 0.1.3 makeWindowedArray :: Source r ix e => Array r ix e -- ^ Source array that will have a window inserted into it -> ix -- ^ Start index for the window -> Sz ix -- ^ Size of the window - -> (ix -> e) -- ^ Inside window indexing function + -> (ix -> e) -- ^ Indexing function foto use inside window -> Array DW ix e -makeWindowedArray !arr !windowStart !windowSize windowIndex - | not (isSafeIndex sz windowStart) = - error $ - "makeWindowedArray: Incorrect window starting index: (" ++ - show windowStart ++ ") for array size: (" ++ show (size arr) ++ ")" - | totalElem windowSize == 0 = - error $ - "makeWindowedArray: Window can't hold any elements with this size: (" ++ show windowSize ++ ")" - | not (isSafeIndex (liftIndex (+ 1) sz) (liftIndex2 (+) windowStart windowSize)) = - error $ - "makeWindowedArray: Incorrect window size: (" ++ - show windowSize ++ - ") and/or starting index: (" ++ - show windowStart ++ ") for array size: (" ++ show (size arr) ++ ")" - | otherwise = - DWArray - { dwArray = delay arr - , dwStencilSize = Nothing - , dwWindow = Just $! Window {..} - } - where - sz = size arr +makeWindowedArray !arr wStart wSize wIndex = + insertWindow (delay arr) $ + 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 +-- that array. +-- +-- @since 0.3.0 +insertWindow + :: Source D ix e + => Array D ix e -- ^ Source array that will have a window inserted into it + -> Window ix e -- ^ Window to place inside the delayed array + -> Array DW ix e +insertWindow !arr !window = + DWArray + { dwArray = delay arr + , dwWindow = + Just $! + Window + { windowStart = unSz (Sz (liftIndex2 min wStart (liftIndex (subtract 1) sz))) + , windowSize = Sz (liftIndex2 min wSize (liftIndex2 (-) sz wStart)) + , windowIndex = wIndex + , windowUnrollIx2 = wUnrollIx2 + } + } + where + Sz sz = size arr + Window { windowStart = wStart + , windowSize = Sz wSize + , windowIndex = wIndex + , windowUnrollIx2 = wUnrollIx2 + } = window +{-# INLINE insertWindow #-} --- | Get the `Window` from the Windowed array. + +-- | Get the `Window` from a windowed array. -- -- @since 0.2.1 getWindow :: Array DW ix e -> Maybe (Window ix e) getWindow = dwWindow {-# INLINE getWindow #-} +-- | Drop the `Window` from a windowed array. +-- +-- @since 0.3.0 +dropWindow :: Array DW ix e -> Array D ix e +dropWindow = dwArray +{-# INLINE dropWindow #-} + zeroWindow :: Index ix => Window ix e -zeroWindow = Window zeroIndex zeroIndex windowError +zeroWindow = Window zeroIndex zeroSz windowError Nothing {-# INLINE zeroWindow #-} +data EmptyWindowException = EmptyWindowException deriving (Eq, Show) + +instance Exception EmptyWindowException where + + displayException _ = "Index of zero size Window" + windowError :: a -windowError = error "Impossible: index of zeroWindow" +windowError = throwImpossible EmptyWindowException {-# NOINLINE windowError #-} @@ -157,65 +208,65 @@ => (m () -> m ()) -> Array DW Ix1 e -> (Ix1 -> e -> m a) - -> m ((Ix1, Ix1) -> m (), (Ix1, Ix1)) -loadWithIx1 with (DWArray (DArray _ sz indexB) _ window) unsafeWrite = do - let Window it wk indexW = fromMaybe zeroWindow window - wEnd = it + wk - with $ iterM_ 0 it 1 (<) $ \ !i -> unsafeWrite i (indexB i) - with $ iterM_ wEnd sz 1 (<) $ \ !i -> unsafeWrite i (indexB i) - return (\(from, to) -> with $ iterM_ from to 1 (<) $ \ !i -> unsafeWrite i (indexW i), (it, wEnd)) + -> m (Ix1 -> Ix1 -> m (), Ix1, Ix1) +loadWithIx1 with (DWArray (DArray _ sz indexB) mWindow) uWrite = do + let Window it wk indexW _ = fromMaybe zeroWindow mWindow + wEnd = it + unSz wk + with $ iterM_ 0 it 1 (<) $ \ !i -> uWrite i (indexB i) + with $ iterM_ wEnd (unSz sz) 1 (<) $ \ !i -> uWrite i (indexB i) + return (\from to -> with $ iterM_ from to 1 (<) $ \ !i -> uWrite i (indexW i), it, wEnd) {-# INLINE loadWithIx1 #-} instance Load DW Ix1 e where - loadS arr _ unsafeWrite = loadWithIx1 id arr unsafeWrite >>= uncurry ($) - {-# INLINE loadS #-} - loadP wIds arr _ unsafeWrite = - withScheduler_ wIds $ \scheduler -> do - (loadWindow, (wStart, wEnd)) <- loadWithIx1 (scheduleWork scheduler) arr unsafeWrite - let (chunkHeight, slackHeight) = (wEnd - wStart) `quotRem` numWorkers scheduler - loopM_ 0 (< numWorkers scheduler) (+ 1) $ \ !wid -> - let !it' = wid * chunkHeight + wStart - in loadWindow (it', it' + chunkHeight) - when (slackHeight > 0) $ - let !itSlack = numWorkers scheduler * chunkHeight + wStart - in loadWindow (itSlack, itSlack + slackHeight) - {-# INLINE loadP #-} - loadArray numWorkers' scheduleWork' arr = - loadArrayWithStride numWorkers' scheduleWork' oneStride (size arr) arr - {-# INLINE loadArray #-} - loadArrayWithStride numWorkers' scheduleWork' stride sz arr _ unsafeWrite = do - (loadWindow, (wStart, wEnd)) <- loadArrayWithIx1 scheduleWork' arr stride sz unsafeWrite - let (chunkHeight, slackHeight) = (wEnd - wStart) `quotRem` numWorkers' - loopM_ 0 (< numWorkers') (+ 1) $ \ !wid -> - let !it' = wid * chunkHeight + wStart - in loadWindow (it', it' + chunkHeight) - when (slackHeight > 0) $ - let !itSlack = numWorkers' * chunkHeight + wStart - in loadWindow (itSlack, itSlack + slackHeight) - {-# INLINE loadArrayWithStride #-} + size = dSize . dwArray + {-# INLINE size #-} + getComp = dComp . dwArray + {-# INLINE getComp #-} + loadArrayM Scheduler {numWorkers, scheduleWork} arr uWrite = do + (loadWindow, wStart, wEnd) <- loadWithIx1 scheduleWork arr uWrite + let (chunkWidth, slackWidth) = (wEnd - wStart) `quotRem` numWorkers + loopM_ 0 (< numWorkers) (+ 1) $ \ !wid -> + let !it' = wid * chunkWidth + wStart + in loadWindow it' (it' + chunkWidth) + when (slackWidth > 0) $ + let !itSlack = numWorkers * chunkWidth + wStart + in loadWindow itSlack (itSlack + slackWidth) + {-# INLINE loadArrayM #-} +instance StrideLoad DW Ix1 e where + loadArrayWithStrideM Scheduler {numWorkers, scheduleWork} stride sz arr uWrite = do + (loadWindow, (wStart, wEnd)) <- loadArrayWithIx1 scheduleWork arr stride sz uWrite + let (chunkWidth, slackWidth) = (wEnd - wStart) `quotRem` numWorkers + loopM_ 0 (< numWorkers) (+ 1) $ \ !wid -> + let !it' = wid * chunkWidth + wStart + in loadWindow (it', it' + chunkWidth) + when (slackWidth > 0) $ + let !itSlack = numWorkers * chunkWidth + wStart + in loadWindow (itSlack, itSlack + slackWidth) + {-# INLINE loadArrayWithStrideM #-} + loadArrayWithIx1 :: (Monad m) => (m () -> m ()) -> Array DW Ix1 e -> Stride Ix1 - -> Ix1 + -> Sz1 -> (Ix1 -> e -> m a) -> m ((Ix1, Ix1) -> m (), (Ix1, Ix1)) -loadArrayWithIx1 with (DWArray (DArray _ arrSz indexB) _ window) stride _ unsafeWrite = do - let Window it wk indexW = fromMaybe zeroWindow window - wEnd = it + wk +loadArrayWithIx1 with (DWArray (DArray _ arrSz indexB) mWindow) stride _ uWrite = do + let Window it wk indexW _ = fromMaybe zeroWindow mWindow + wEnd = it + unSz wk strideIx = unStride stride - with $ iterM_ 0 it strideIx (<) $ \ !i -> unsafeWrite (i `div` strideIx) (indexB i) + with $ iterM_ 0 it strideIx (<) $ \ !i -> uWrite (i `div` strideIx) (indexB i) with $ - iterM_ (strideStart stride wEnd) arrSz strideIx (<) $ \ !i -> - unsafeWrite (i `div` strideIx) (indexB i) + iterM_ (strideStart stride wEnd) (unSz arrSz) strideIx (<) $ \ !i -> + uWrite (i `div` strideIx) (indexB i) return ( \(from, to) -> with $ iterM_ (strideStart stride from) to strideIx (<) $ \ !i -> - unsafeWrite (i `div` strideIx) (indexW i) + uWrite (i `div` strideIx) (indexW i) , (it, wEnd)) {-# INLINE loadArrayWithIx1 #-} @@ -227,19 +278,16 @@ -> Array DW Ix2 t1 -> (Int -> t1 -> m ()) -> m (Ix2 -> m (), Ix2) -loadWithIx2 with arr unsafeWrite = do - let DWArray (DArray _ (m :. n) indexB) mStencilSize window = arr - let Window (it :. jt) (wm :. wn) indexW = fromMaybe zeroWindow window +loadWithIx2 with arr uWrite = do + let DWArray (DArray _ (Sz (m :. n)) indexB) window = arr + let Window (it :. jt) (Sz (wm :. wn)) indexW mUnrollHeight = fromMaybe zeroWindow window let ib :. jb = (wm + it) :. (wn + jt) - !blockHeight = - case mStencilSize of - Just (i :. _) -> min (max 1 i) 7 - _ -> 1 + !blockHeight = maybe 1 (min 7 . max 1) mUnrollHeight stride = oneStride !sz = strideSize stride $ size arr - writeB !ix = unsafeWrite (toLinearIndex sz ix) (indexB ix) + writeB !ix = uWrite (toLinearIndex sz ix) (indexB ix) {-# INLINE writeB #-} - writeW !ix = unsafeWrite (toLinearIndex sz ix) (indexW ix) + writeW !ix = uWrite (toLinearIndex sz ix) (indexW ix) {-# INLINE writeW #-} with $ iterM_ (0 :. 0) (it :. n) (1 :. 1) (<) writeB with $ iterM_ (ib :. 0) (m :. n) (1 :. 1) (<) writeB @@ -250,63 +298,30 @@ return (f, it :. ib) {-# INLINE loadWithIx2 #-} - -instance Load DW Ix2 e where - loadS arr _ unsafeWrite = loadWithIx2 id arr unsafeWrite >>= uncurry ($) - {-# INLINE loadS #-} - -- - loadP wIds arr _ unsafeWrite = - withScheduler_ wIds $ \scheduler -> do - (loadWindow, it :. ib) <- loadWithIx2 (scheduleWork scheduler) arr unsafeWrite - let !(chunkHeight, slackHeight) = (ib - it) `quotRem` numWorkers scheduler - loopM_ 0 (< numWorkers scheduler) (+ 1) $ \ !wid -> - let !it' = wid * chunkHeight + it - in loadWindow (it' :. (it' + chunkHeight)) - when (slackHeight > 0) $ - let !itSlack = numWorkers scheduler * chunkHeight + it - in loadWindow (itSlack :. (itSlack + slackHeight)) - {-# INLINE loadP #-} - loadArray numWorkers' scheduleWork' arr = - loadArrayWithStride numWorkers' scheduleWork' oneStride (size arr) arr - {-# INLINE loadArray #-} - loadArrayWithStride numWorkers' scheduleWork' stride sz arr _ unsafeWrite = do - (loadWindow, it :. ib) <- loadArrayWithIx2 scheduleWork' arr stride sz unsafeWrite - let !(chunkHeight, slackHeight) = (ib - it) `quotRem` numWorkers' - loopM_ 0 (< numWorkers') (+ 1) $ \ !wid -> - let !it' = wid * chunkHeight + it - in loadWindow (it' :. (it' + chunkHeight)) - when (slackHeight > 0) $ - let !itSlack = numWorkers' * chunkHeight + it - in loadWindow (itSlack :. (itSlack + slackHeight)) - {-# INLINE loadArrayWithStride #-} - loadArrayWithIx2 :: Monad m => (m () -> m ()) -> Array DW Ix2 e -> Stride Ix2 - -> Ix2 + -> Sz2 -> (Int -> e -> m ()) -> m (Ix2 -> m (), Ix2) -loadArrayWithIx2 with arr stride sz unsafeWrite = do - let DWArray (DArray _ (m :. n) indexB) mStencilSize window = arr - let Window (it :. jt) (wm :. wn) indexW = fromMaybe zeroWindow window +loadArrayWithIx2 with arr stride sz uWrite = do + let DWArray (DArray _ (Sz (m :. n)) indexB) window = arr + let Window (it :. jt) (Sz (wm :. wn)) indexW mUnrollHeight = fromMaybe zeroWindow window let ib :. jb = (wm + it) :. (wn + jt) - !blockHeight = - case mStencilSize of - Just (i :. _) -> min (max 1 i) 7 - _ -> 1 + !blockHeight = maybe 1 (min 7 . max 1) mUnrollHeight strideIx@(is :. js) = unStride stride - writeB !ix = unsafeWrite (toLinearIndexStride stride sz ix) (indexB ix) + writeB !ix = uWrite (toLinearIndexStride stride sz ix) (indexB ix) {-# INLINE writeB #-} - writeW !ix = unsafeWrite (toLinearIndexStride stride sz ix) (indexW ix) + writeW !ix = uWrite (toLinearIndexStride stride sz ix) (indexW ix) {-# INLINE writeW #-} with $ iterM_ (0 :. 0) (it :. n) strideIx (<) writeB with $ iterM_ (strideStart stride (ib :. 0)) (m :. n) strideIx (<) writeB with $ iterM_ (strideStart stride (it :. 0)) (ib :. jt) strideIx (<) writeB with $ iterM_ (strideStart stride (it :. jb)) (ib :. n) strideIx (<) writeB f <- - if is > 1 -- Turn off unrolling for vertical strides + if is > 1 || blockHeight <= 1 -- Turn off unrolling for vertical strides then return $ \(it' :. ib') -> iterM_ (strideStart stride (it' :. jt)) (ib' :. jb) strideIx (<) writeW else return $ \(it' :. ib') -> @@ -315,67 +330,87 @@ {-# INLINE loadArrayWithIx2 #-} +loadWindowIx2 :: Monad m => Int -> (Ix2 -> m ()) -> Ix2 -> m () +loadWindowIx2 numWorkers loadWindow (it :. ib) = do + let !(chunkHeight, slackHeight) = (ib - it) `quotRem` numWorkers + loopM_ 0 (< numWorkers) (+ 1) $ \ !wid -> + let !it' = wid * chunkHeight + it + in loadWindow (it' :. (it' + chunkHeight)) + when (slackHeight > 0) $ + let !itSlack = numWorkers * chunkHeight + it + in loadWindow (itSlack :. (itSlack + slackHeight)) +{-# INLINE loadWindowIx2 #-} + +instance Load DW Ix2 e where + size = dSize . dwArray + {-# INLINE size #-} + getComp = dComp . dwArray + {-# INLINE getComp #-} + loadArrayM Scheduler {numWorkers, scheduleWork} arr uWrite = + loadWithIx2 scheduleWork arr uWrite >>= uncurry (loadWindowIx2 numWorkers) + {-# INLINE loadArrayM #-} + +instance StrideLoad DW Ix2 e where + loadArrayWithStrideM Scheduler {numWorkers, scheduleWork} stride sz arr uWrite = + loadArrayWithIx2 scheduleWork arr stride sz uWrite >>= uncurry (loadWindowIx2 numWorkers) + {-# INLINE loadArrayWithStrideM #-} + + instance (Index (IxN n), Load DW (Ix (n - 1)) e) => Load DW (IxN n) e where - loadS = loadWithIxN id - {-# INLINE loadS #-} - loadP wIds arr unsafeRead unsafeWrite = - withScheduler_ wIds $ \scheduler -> - loadWithIxN (scheduleWork scheduler) arr unsafeRead unsafeWrite - {-# INLINE loadP #-} - loadArray numWorkers' scheduleWork' arr = - loadArrayWithStride numWorkers' scheduleWork' oneStride (size arr) arr - {-# INLINE loadArray #-} - loadArrayWithStride = loadArrayWithIxN - {-# INLINE loadArrayWithStride #-} + size = dSize . dwArray + {-# INLINE size #-} + getComp = dComp . dwArray + {-# INLINE getComp #-} + loadArrayM Scheduler {scheduleWork} = loadWithIxN scheduleWork + {-# INLINE loadArrayM #-} +instance (Index (IxN n), StrideLoad DW (Ix (n - 1)) e) => StrideLoad DW (IxN n) e where + loadArrayWithStrideM = loadArrayWithIxN + {-# INLINE loadArrayWithStrideM #-} + loadArrayWithIxN :: - (Index ix, Monad m, Load DW (Lower ix) e) - => Int - -> (m () -> m ()) + (Index ix, Monad m, StrideLoad DW (Lower ix) e) + => Scheduler m () -> Stride ix - -> ix + -> Sz ix -> Array DW ix e - -> (Int -> m e) -> (Int -> e -> m ()) -> m () -loadArrayWithIxN numWorkers' scheduleWork' stride szResult arr unsafeRead unsafeWrite = do - let DWArray darr mStencilSize window = arr +loadArrayWithIxN scheduler stride szResult arr uWrite = do + let DWArray darr window = arr DArray {dSize = szSource, dIndex = indexBorder} = darr - Window {windowStart, windowSize, windowIndex = indexWindow} = fromMaybe zeroWindow window - !(headSourceSize, lowerSourceSize) = unconsDim szSource - !lowerSize = tailDim szResult + Window {windowStart, windowSize, windowIndex, windowUnrollIx2} = fromMaybe zeroWindow window + !(headSourceSize, lowerSourceSize) = unconsSz szSource + !lowerSize = snd $ unconsSz szResult !(s, lowerStrideIx) = unconsDim $ unStride stride !(curWindowStart, lowerWindowStart) = unconsDim windowStart - !curWindowEnd = curWindowStart + headDim windowSize + !(headWindowSz, tailWindowSz) = unconsSz windowSize + !curWindowEnd = curWindowStart + unSz headWindowSz !pageElements = totalElem lowerSize - -- can safely drop the dim, only last 2 matter anyways - !mLowerStencilSize = fmap tailDim mStencilSize loadLower !i = let !lowerWindow = Window { windowStart = lowerWindowStart - , windowSize = tailDim windowSize - , windowIndex = indexWindow . consDim i + , windowSize = tailWindowSz + , windowIndex = windowIndex . consDim i + , windowUnrollIx2 = windowUnrollIx2 } !lowerArr = DWArray { dwArray = DArray Seq lowerSourceSize (indexBorder . consDim i) - , dwStencilSize = mLowerStencilSize , dwWindow = Just lowerWindow } - in loadArrayWithStride - numWorkers' - scheduleWork' + in loadArrayWithStrideM + scheduler (Stride lowerStrideIx) lowerSize lowerArr - (\k -> unsafeRead (k + pageElements * (i `div` s))) - (\k -> unsafeWrite (k + pageElements * (i `div` s))) + (\k -> uWrite (k + pageElements * (i `div` s))) {-# NOINLINE loadLower #-} loopM_ 0 (< headDim windowStart) (+ s) loadLower loopM_ (strideStart (Stride s) curWindowStart) (< curWindowEnd) (+ s) loadLower - loopM_ (strideStart (Stride s) curWindowEnd) (< headSourceSize) (+ s) loadLower + loopM_ (strideStart (Stride s) curWindowEnd) (< unSz headSourceSize) (+ s) loadLower {-# INLINE loadArrayWithIxN #-} @@ -384,41 +419,34 @@ (Index ix, Monad m, Load DW (Lower ix) e) => (m () -> m ()) -> Array DW ix e - -> (Int -> m e) -> (Int -> e -> m ()) -> m () -loadWithIxN with arr unsafeRead unsafeWrite = do - let DWArray darr mStencilSize window = arr +loadWithIxN with arr uWrite = do + let DWArray darr window = arr DArray {dSize = sz, dIndex = indexBorder} = darr - Window {windowStart, windowSize, windowIndex = indexWindow} = fromMaybe zeroWindow window - !szL = tailDim sz - !windowEnd = liftIndex2 (+) windowStart windowSize + Window {windowStart, windowSize, windowIndex, windowUnrollIx2} = fromMaybe zeroWindow window + !(si, szL) = unconsSz sz + !windowEnd = liftIndex2 (+) windowStart (unSz windowSize) !(t, windowStartL) = unconsDim windowStart !pageElements = totalElem szL - -- can safely drop the dim, only last 2 matter anyways - !stencilSizeLower = fmap tailDim mStencilSize loadLower !i = let !lowerWindow = Window { windowStart = windowStartL - , windowSize = tailDim windowSize - , windowIndex = indexWindow . consDim i + , windowSize = snd $ unconsSz windowSize + , windowIndex = windowIndex . consDim i + , windowUnrollIx2 = windowUnrollIx2 } !lowerArr = DWArray { dwArray = DArray Seq szL (indexBorder . consDim i) - , dwStencilSize = stencilSizeLower , dwWindow = Just lowerWindow } - in with $ - loadS - lowerArr - (\k -> unsafeRead (k + pageElements * i)) - (\k -> unsafeWrite (k + pageElements * i)) + in with $ loadArrayM (Scheduler 1 id) lowerArr (\k -> uWrite (k + pageElements * i)) {-# NOINLINE loadLower #-} loopM_ 0 (< headDim windowStart) (+ 1) loadLower loopM_ t (< headDim windowEnd) (+ 1) loadLower - loopM_ (headDim windowEnd) (< headDim sz) (+ 1) loadLower + loopM_ (headDim windowEnd) (< unSz si) (+ 1) loadLower {-# INLINE loadWithIxN #-} @@ -462,75 +490,74 @@ toIx2Window Window {..} = Window { windowStart = toIx2 windowStart - , windowSize = toIx2 windowSize + , windowSize = SafeSz (toIx2 $ unSz windowSize) , windowIndex = windowIndex . fromIx2 + , windowUnrollIx2 = windowUnrollIx2 } {-# INLINE toIx2Window #-} toIx2ArrayDW :: Array DW Ix2T e -> Array DW Ix2 e -toIx2ArrayDW DWArray {dwArray, dwStencilSize, dwWindow} = +toIx2ArrayDW DWArray {dwArray, dwWindow} = DWArray - { dwArray = dwArray {dIndex = dIndex dwArray . fromIx2, dSize = toIx2 (dSize dwArray)} - , dwStencilSize = fmap toIx2 dwStencilSize + { dwArray = + dwArray {dIndex = dIndex dwArray . fromIx2, dSize = SafeSz (toIx2 (unSz (dSize dwArray)))} , dwWindow = fmap toIx2Window dwWindow } {-# INLINE toIx2ArrayDW #-} instance Load DW Ix2T e where - loadS arr = loadS (toIx2ArrayDW arr) - {-# INLINE loadS #-} - loadP wIds arr = loadP wIds (toIx2ArrayDW arr) - {-# INLINE loadP #-} - loadArray numWorkers' scheduleWork' arr = - loadArrayWithStride numWorkers' scheduleWork' oneStride (size arr) arr - {-# INLINE loadArray #-} - loadArrayWithStride numWorkers' scheduleWork' stride sz arr = - loadArrayWithStride - numWorkers' - scheduleWork' + size = dSize . dwArray + {-# INLINE size #-} + getComp = dComp . dwArray + {-# INLINE getComp #-} + loadArrayM scheduler arr = + loadArrayWithStrideM scheduler oneStride (size arr) arr + {-# INLINE loadArrayM #-} + +instance StrideLoad DW Ix2T e where + loadArrayWithStrideM scheduler stride sz arr = + loadArrayWithStrideM + scheduler (Stride $ toIx2 $ unStride stride) - (toIx2 sz) + (SafeSz (toIx2 (unSz sz))) (toIx2ArrayDW arr) - {-# INLINE loadArrayWithStride #-} + {-# INLINE loadArrayWithStrideM #-} instance Load DW Ix3T e where - loadS = loadWithIxN id - {-# INLINE loadS #-} - loadP wIds arr unsafeRead unsafeWrite = - withScheduler_ wIds $ \scheduler -> - loadWithIxN (scheduleWork scheduler) arr unsafeRead unsafeWrite - {-# INLINE loadP #-} - loadArray numWorkers' scheduleWork' arr = - loadArrayWithStride numWorkers' scheduleWork' oneStride (size arr) arr - {-# INLINE loadArray #-} - loadArrayWithStride = loadArrayWithIxN - {-# INLINE loadArrayWithStride #-} + size = dSize . dwArray + {-# INLINE size #-} + getComp = dComp . dwArray + {-# INLINE getComp #-} + loadArrayM scheduler arr = + loadArrayWithStrideM scheduler oneStride (size arr) arr + {-# INLINE loadArrayM #-} +instance StrideLoad DW Ix3T e where + loadArrayWithStrideM = loadArrayWithIxN + {-# INLINE loadArrayWithStrideM #-} + instance Load DW Ix4T e where - loadS = loadWithIxN id - {-# INLINE loadS #-} - loadP wIds arr unsafeRead unsafeWrite = - withScheduler_ wIds $ \scheduler -> - loadWithIxN (scheduleWork scheduler) arr unsafeRead unsafeWrite - {-# INLINE loadP #-} - loadArray numWorkers' scheduleWork' arr = - loadArrayWithStride numWorkers' scheduleWork' oneStride (size arr) arr - {-# INLINE loadArray #-} - loadArrayWithStride = loadArrayWithIxN - {-# INLINE loadArrayWithStride #-} + size = dSize . dwArray + {-# INLINE size #-} + getComp = dComp . dwArray + {-# INLINE getComp #-} + loadArrayM scheduler arr = loadArrayWithStrideM scheduler oneStride (size arr) arr + {-# INLINE loadArrayM #-} +instance StrideLoad DW Ix4T e where + loadArrayWithStrideM = loadArrayWithIxN + {-# INLINE loadArrayWithStrideM #-} + instance Load DW Ix5T e where - loadS = loadWithIxN id - {-# INLINE loadS #-} - loadP wIds arr unsafeRead unsafeWrite = - withScheduler_ wIds $ \scheduler -> - loadWithIxN (scheduleWork scheduler) arr unsafeRead unsafeWrite - {-# INLINE loadP #-} - loadArray numWorkers' scheduleWork' arr = - loadArrayWithStride numWorkers' scheduleWork' oneStride (size arr) arr - {-# INLINE loadArray #-} - loadArrayWithStride = loadArrayWithIxN - {-# INLINE loadArrayWithStride #-} + size = dSize . dwArray + {-# INLINE size #-} + getComp = dComp . dwArray + {-# INLINE getComp #-} + loadArrayM scheduler arr = loadArrayWithStrideM scheduler oneStride (size arr) arr + {-# INLINE loadArrayM #-} +instance StrideLoad DW Ix5T e where + loadArrayWithStrideM = loadArrayWithIxN + {-# INLINE loadArrayWithStrideM #-} diff --git a/src/Data/Massiv/Array/Manifest.hs b/src/Data/Massiv/Array/Manifest.hs --- a/src/Data/Massiv/Array/Manifest.hs +++ b/src/Data/Massiv/Array/Manifest.hs @@ -1,13 +1,14 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} -- | -- Module : Data.Massiv.Array.Manifest --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental @@ -39,8 +40,10 @@ -- ** Conversion , toByteArray , fromByteArray + , fromByteArrayM , toMutableByteArray , fromMutableByteArray + , fromMutableByteArrayM -- * Storable , S(..) , Storable @@ -49,7 +52,6 @@ , toStorableMVector -- ** Direct Pointer Access , withPtr - , unsafeWithPtr -- * Unboxed , U(..) , Unbox @@ -58,22 +60,25 @@ , toUnboxedMVector -- * ByteString Conversion , fromByteString + , castFromByteString , toByteString + , castToByteString , toBuilder ) where -import Data.ByteString as S -import Data.ByteString.Builder -import Data.ByteString.Unsafe as SU -import Data.Massiv.Array.Manifest.Boxed -import Data.Massiv.Array.Manifest.Internal -import Data.Massiv.Array.Manifest.Primitive -import Data.Massiv.Array.Manifest.Storable -import Data.Massiv.Array.Manifest.Unboxed -import Data.Massiv.Array.Ops.Fold -import Data.Massiv.Core.Common -import Data.Word (Word8) - +import Data.ByteString as S +import Data.ByteString.Builder +import Data.ByteString.Internal +import Data.ByteString.Unsafe as SU +import Data.Massiv.Array.Manifest.Boxed +import Data.Massiv.Array.Manifest.Internal +import Data.Massiv.Array.Manifest.Primitive +import Data.Massiv.Array.Manifest.Storable +import Data.Massiv.Array.Manifest.Unboxed +import Data.Massiv.Array.Ops.Fold +import Data.Massiv.Core.Common +import Data.Massiv.Core.Index.Internal (Sz(..)) +import Data.Word (Word8) -- | /O(1)/ - Convert a strict ByteString into a manifest array. Will return `Nothing` if length @@ -84,27 +89,47 @@ Comp -- ^ Computation strategy -> ByteString -- ^ Strict ByteString to use as a source. -> Array M Ix1 Word8 -fromByteString comp bs = MArray comp (S.length bs) (SU.unsafeIndex bs) +fromByteString comp bs = MArray comp (SafeSz (S.length bs)) (SU.unsafeIndex bs) {-# INLINE fromByteString #-} --- | /O(n)/ - For now only sequentially convert an array into a strict ByteString +-- | /O(n)/ - Convert any source array into a strict `ByteString`. In case when the source array is +-- actually storable, no memory copy will occur. -- -- @since 0.2.1 toByteString :: - Source r ix Word8 + Load r ix Word8 => Array r ix Word8 -- ^ Source array -> ByteString -toByteString arr = - fst $ unfoldrN (totalElem (size arr)) (\ !i -> Just (unsafeLinearIndex arr i, i + 1)) 0 +toByteString = castToByteString . +#if __GLASGOW_HASKELL__ >= 820 + convert + {- For ghc-8.0 `covert` results in "internal error: ARR_WORDS object entered!" -} +#else + compute +#endif + --fst $ unfoldrN (totalElem (size arr)) (\ !i -> Just (unsafeLinearIndex arr i, i + 1)) 0 {-# INLINE toByteString #-} --- | /O(n)/ - Conversion of array monoidally into a string Builder. +-- | /O(n)/ - Conversion of array monoidally into a ByteString Builder. -- -- @since 0.2.1 toBuilder :: Source r ix e => (e -> Builder) -> Array r ix e -> Builder -toBuilder f = foldMono f +toBuilder = foldMono {-# INLINE toBuilder #-} +-- | /O(1)/ - Cast a `S`torable array into a strict `ByteString` +-- +-- @since 0.3.0 +castToByteString :: Array S ix Word8 -> ByteString +castToByteString = (\(fp, len) -> PS fp 0 len) . unsafeArrayToForeignPtr +{-# INLINE castToByteString #-} + +-- | /O(1)/ - Cast a strict `ByteString` into a `S`torable array +-- +-- @since 0.3.0 +castFromByteString :: Comp -> ByteString -> Array S Ix1 Word8 +castFromByteString comp (PS fp offset len) = unsafeArrayFromForeignPtr comp fp offset (Sz len) +{-# INLINE castFromByteString #-} -- $boxed_conversion_note -- diff --git a/src/Data/Massiv/Array/Manifest/Boxed.hs b/src/Data/Massiv/Array/Manifest/Boxed.hs --- a/src/Data/Massiv/Array/Manifest/Boxed.hs +++ b/src/Data/Massiv/Array/Manifest/Boxed.hs @@ -1,15 +1,16 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE CPP #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE MagicHash #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE UndecidableInstances #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} -- | -- Module : Data.Massiv.Array.Manifest.Boxed --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental @@ -19,7 +20,6 @@ ( B(..) , N(..) , Array(..) - , Uninitialized(..) , unwrapArray , evalArray , unwrapMutableArray @@ -29,34 +29,34 @@ , unwrapNormalFormMutableArray , evalNormalFormMutableArray , castArrayToVector - , vectorToArray , castVectorToArray , seqArray , deepseqArray ) where -import Control.DeepSeq (NFData (..), deepseq) -import Control.Exception -import Control.Monad.Primitive -import Control.Monad.ST (runST) -import qualified Data.Foldable as F (Foldable (..)) -import Data.Massiv.Array.Delayed.Internal (eq, ord) -import Data.Massiv.Array.Manifest.Internal (M, toManifest) -import Data.Massiv.Array.Manifest.List as L -import Data.Massiv.Array.Mutable -import Data.Massiv.Array.Ops.Fold.Internal -import Data.Massiv.Array.Unsafe (unsafeGenerateArray, - unsafeGenerateArrayP) -import Data.Massiv.Core.Common -import Data.Massiv.Core.List -import qualified Data.Primitive.Array as A -import qualified Data.Vector as VB -import qualified Data.Vector.Mutable as VB -import GHC.Base (build) -import GHC.Exts as GHC (IsList (..)) -import GHC.Prim -import GHC.Types -import Prelude hiding (mapM) +import Control.DeepSeq (NFData(..), deepseq) +import Control.Exception +import Control.Monad ((>=>)) +import Control.Monad.Primitive +import Control.Monad.ST (runST) +import qualified Data.Foldable as F (Foldable(..)) +import Data.Massiv.Array.Delayed.Pull (eq, ord) +import Data.Massiv.Array.Delayed.Push (DL) +import Data.Massiv.Array.Manifest.Internal (M, computeAs, toManifest) +import Data.Massiv.Array.Manifest.List as L +import Data.Massiv.Array.Mutable +import Data.Massiv.Array.Ops.Fold +import Data.Massiv.Array.Ops.Fold.Internal +import Data.Massiv.Array.Ops.Map (traverseA) +import Data.Massiv.Core.Common +import Data.Massiv.Core.List +import qualified Data.Primitive.Array as A +import qualified Data.Vector as VB +import qualified Data.Vector.Mutable as VB +import GHC.Base (build) +import GHC.Exts as GHC +import Prelude hiding (mapM) +import System.IO.Unsafe (unsafePerformIO) #include "massiv.h" @@ -81,12 +81,22 @@ type instance EltRepr B ix = M data instance Array B ix e = BArray { bComp :: !Comp - , bSize :: !ix + , bSize :: !(Sz ix) , bData :: {-# UNPACK #-} !(A.Array e) } +instance (Ragged L ix e, Show e) => Show (Array B ix e) where + showsPrec = showsArrayPrec id + showList = showArrayList + +instance (Ragged L ix e, Show e) => Show (Array DL ix e) where + showsPrec = showsArrayPrec (computeAs B) + showList = showArrayList + + instance (Index ix, NFData e) => NFData (Array B ix e) where rnf = (`deepseqArray` ()) + {-# INLINE rnf #-} instance (Index ix, Eq e) => Eq (Array B ix e) where (==) = eq (==) @@ -97,35 +107,28 @@ {-# INLINE compare #-} instance Index ix => Construct B ix e where - getComp = bComp - {-# INLINE getComp #-} - setComp c arr = arr { bComp = c } {-# INLINE setComp #-} - unsafeMakeArray Seq !sz f = unsafeGenerateArray sz f - unsafeMakeArray (ParOn wIds) !sz f = unsafeGenerateArrayP wIds sz f - {-# INLINE unsafeMakeArray #-} + makeArray !comp !sz f = unsafePerformIO $ generateArray comp sz (\ !ix -> return $! f ix) + {-# INLINE makeArray #-} instance Index ix => Source B ix e where unsafeLinearIndex (BArray _ _ a) = - INDEX_CHECK("(Source B ix e).unsafeLinearIndex", sizeofArray, A.indexArray) a + INDEX_CHECK("(Source B ix e).unsafeLinearIndex", Sz . sizeofArray, A.indexArray) a {-# INLINE unsafeLinearIndex #-} -instance Index ix => Size B ix e where - size = bSize - {-# INLINE size #-} - +instance Index ix => Resize B ix where unsafeResize !sz !arr = arr { bSize = sz } {-# INLINE unsafeResize #-} +instance Index ix => Extract B ix e where unsafeExtract !sIx !newSz !arr = unsafeExtract sIx newSz (toManifest arr) {-# INLINE unsafeExtract #-} -instance ( NFData e - , Index ix +instance ( Index ix , Index (Lower ix) , Elt M ix e ~ Array M (Lower ix) e , Elt B ix e ~ Array M (Lower ix) e @@ -134,8 +137,7 @@ unsafeOuterSlice arr = unsafeOuterSlice (toManifest arr) {-# INLINE unsafeOuterSlice #-} -instance ( NFData e - , Index ix +instance ( Index ix , Index (Lower ix) , Elt M ix e ~ Array M (Lower ix) e , Elt B ix e ~ Array M (Lower ix) e @@ -144,16 +146,20 @@ unsafeInnerSlice arr = unsafeInnerSlice (toManifest arr) {-# INLINE unsafeInnerSlice #-} +instance {-# OVERLAPPING #-} Slice B Ix1 e where + unsafeSlice arr i _ _ = pure (unsafeLinearIndex arr i) + {-# INLINE unsafeSlice #-} + instance Index ix => Manifest B ix e where unsafeLinearIndexM (BArray _ _ a) = - INDEX_CHECK("(Manifest B ix e).unsafeLinearIndexM", sizeofArray, A.indexArray) a + INDEX_CHECK("(Manifest B ix e).unsafeLinearIndexM", Sz . sizeofArray, A.indexArray) a {-# INLINE unsafeLinearIndexM #-} instance Index ix => Mutable B ix e where - data MArray s B ix e = MBArray !ix {-# UNPACK #-} !(A.MutableArray s e) + data MArray s B ix e = MBArray !(Sz ix) {-# UNPACK #-} !(A.MutableArray s e) msize (MBArray sz _) = sz {-# INLINE msize #-} @@ -167,20 +173,34 @@ unsafeNew sz = MBArray sz <$> A.newArray (totalElem sz) uninitialized {-# INLINE unsafeNew #-} - unsafeNewZero = unsafeNew - {-# INLINE unsafeNewZero #-} + initialize _ = return () + {-# INLINE initialize #-} unsafeLinearRead (MBArray _ ma) = - INDEX_CHECK("(Mutable B ix e).unsafeLinearRead", sizeofMutableArray, A.readArray) ma + INDEX_CHECK("(Mutable B ix e).unsafeLinearRead", Sz . sizeofMutableArray, A.readArray) ma {-# INLINE unsafeLinearRead #-} unsafeLinearWrite (MBArray _ ma) i e = e `seq` - INDEX_CHECK("(Mutable B ix e).unsafeLinearWrite", sizeofMutableArray, A.writeArray) ma i e + INDEX_CHECK("(Mutable B ix e).unsafeLinearWrite", Sz . sizeofMutableArray, A.writeArray) ma i e {-# INLINE unsafeLinearWrite #-} +instance Index ix => Load B ix e where + size = bSize + {-# INLINE size #-} + getComp = bComp + {-# INLINE getComp #-} + loadArrayM !scheduler !arr = splitLinearlyWith_ scheduler (elemsCount arr) (unsafeLinearIndex arr) + {-# INLINE loadArrayM #-} +instance Index ix => StrideLoad B ix e + + -- | Row-major sequential folding over a Boxed array. instance Index ix => Foldable (Array B ix) where + fold = fold + {-# INLINE fold #-} + foldMap = foldMono + {-# INLINE foldMap #-} foldl = lazyFoldlS {-# INLINE foldl #-} foldl' = foldlS @@ -191,16 +211,20 @@ {-# INLINE foldr' #-} null (BArray _ sz _) = totalElem sz == 0 {-# INLINE null #-} - sum = F.foldl' (+) 0 - {-# INLINE sum #-} - product = F.foldl' (*) 1 - {-# INLINE product #-} length = totalElem . size {-# INLINE length #-} toList arr = build (\ c n -> foldrFB c n arr) {-# INLINE toList #-} +instance Index ix => Functor (Array B ix) where + fmap f arr = makeArrayLinear (bComp arr) (bSize arr) (f . unsafeLinearIndex arr) + {-# INLINE fmap #-} + +instance Index ix => Traversable (Array B ix) where + traverse = traverseA + {-# INLINE traverse #-} + instance ( IsList (Array L ix e) , Nested LN ix e , Nested L ix e @@ -226,9 +250,13 @@ newtype instance Array N ix e = NArray { bArray :: Array B ix e } +instance (Ragged L ix e, Show e, NFData e) => Show (Array N ix e) where + showsPrec = showsArrayPrec bArray + showList = showArrayList + instance (Index ix, NFData e) => NFData (Array N ix e) where rnf (NArray barr) = barr `deepseqArray` () - + {-# INLINE rnf #-} instance (Index ix, NFData e, Eq e) => Eq (Array N ix e) where (==) = eq (==) @@ -240,29 +268,29 @@ instance (Index ix, NFData e) => Construct N ix e where - getComp = bComp . bArray - {-# INLINE getComp #-} - - setComp c (NArray arr) = NArray (arr { bComp = c }) + setComp c (NArray arr) = NArray (arr {bComp = c}) {-# INLINE setComp #-} - - unsafeMakeArray Seq !sz f = NArray $ unsafeGenerateArray sz f - unsafeMakeArray (ParOn wIds) !sz f = NArray $ unsafeGenerateArrayP wIds sz f - {-# INLINE unsafeMakeArray #-} + makeArray !comp !sz f = + unsafePerformIO $ + generateArray + comp + sz + (\ !ix -> + let res = f ix + in res `deepseq` return res) + {-# INLINE makeArray #-} instance (Index ix, NFData e) => Source N ix e where unsafeLinearIndex (NArray arr) = - INDEX_CHECK("(Source N ix e).unsafeLinearIndex", totalElem . size, unsafeLinearIndex) arr + INDEX_CHECK("(Source N ix e).unsafeLinearIndex", Sz . totalElem . size, unsafeLinearIndex) arr {-# INLINE unsafeLinearIndex #-} -instance (Index ix, NFData e) => Size N ix e where - size = bSize . bArray - {-# INLINE size #-} - +instance Index ix => Resize N ix where unsafeResize !sz = NArray . unsafeResize sz . bArray {-# INLINE unsafeResize #-} +instance (Index ix, NFData e) => Extract N ix e where unsafeExtract !sIx !newSz !arr = unsafeExtract sIx newSz (toManifest arr) {-# INLINE unsafeExtract #-} @@ -287,11 +315,15 @@ unsafeInnerSlice = unsafeInnerSlice . toManifest {-# INLINE unsafeInnerSlice #-} +instance {-# OVERLAPPING #-} NFData e => Slice N Ix1 e where + unsafeSlice arr i _ _ = pure (unsafeLinearIndex arr i) + {-# INLINE unsafeSlice #-} + instance (Index ix, NFData e) => Manifest N ix e where unsafeLinearIndexM (NArray arr) = - INDEX_CHECK("(Manifest N ix e).unsafeLinearIndexM", totalElem . size, unsafeLinearIndexM) arr + INDEX_CHECK("(Manifest N ix e).unsafeLinearIndexM", Sz . totalElem . size, unsafeLinearIndexM) arr {-# INLINE unsafeLinearIndexM #-} @@ -310,18 +342,29 @@ unsafeNew sz = MNArray <$> unsafeNew sz {-# INLINE unsafeNew #-} - unsafeNewZero = unsafeNew - {-# INLINE unsafeNewZero #-} + initialize _ = return () + {-# INLINE initialize #-} unsafeLinearRead (MNArray ma) = - INDEX_CHECK("(Mutable N ix e).unsafeLinearRead", totalElem . msize, unsafeLinearRead) ma + INDEX_CHECK("(Mutable N ix e).unsafeLinearRead", Sz . totalElem . msize, unsafeLinearRead) ma {-# INLINE unsafeLinearRead #-} unsafeLinearWrite (MNArray ma) i e = e `deepseq` - INDEX_CHECK("(Mutable N ix e).unsafeLinearWrite", totalElem . msize, unsafeLinearWrite) ma i e + INDEX_CHECK("(Mutable N ix e).unsafeLinearWrite", Sz . totalElem . msize, unsafeLinearWrite) ma i e {-# INLINE unsafeLinearWrite #-} +instance (Index ix, NFData e) => Load N ix e where + size = bSize . bArray + {-# INLINE size #-} + getComp = bComp . bArray + {-# INLINE getComp #-} + loadArrayM !scheduler !arr = splitLinearlyWith_ scheduler (elemsCount arr) (unsafeLinearIndex arr) + {-# INLINE loadArrayM #-} +instance (Index ix, NFData e) => StrideLoad N ix e + + + instance ( NFData e , IsList (Array L ix e) , Nested LN ix e @@ -340,16 +383,6 @@ -- Helper functions -- ---------------------- --- | An error that gets thrown when an unitialized element of a boxed array gets accessed. Can only --- happen when array was constructed with `unsafeNew`. -data Uninitialized = Uninitialized - -instance Show Uninitialized where - show Uninitialized = "Array element is uninitialized" - -instance Exception Uninitialized - - uninitialized :: a uninitialized = throw Uninitialized @@ -460,8 +493,8 @@ -> m (MArray (PrimState m) B Ix1 e) fromMutableArraySeq with mbarr = do let !sz = sizeofMutableArray mbarr - loopM_ 0 (< sz) (+ 1) $ \i -> A.readArray mbarr i >>= (`with` return ()) - return $! MBArray sz mbarr + loopM_ 0 (< sz) (+ 1) (A.readArray mbarr >=> (`with` return ())) + return $! MBArray (Sz sz) mbarr {-# INLINE fromMutableArraySeq #-} fromArraySeq :: @@ -469,7 +502,7 @@ -> Comp -> A.Array e -> a -fromArraySeq with comp barr = with (BArray comp (sizeofArray barr) barr) +fromArraySeq with comp barr = with (BArray comp (Sz (sizeofArray barr)) barr) {-# INLINE fromArraySeq #-} @@ -490,18 +523,6 @@ marr <- A.unsafeThawArray arr VB.unsafeFreeze $ VB.MVector 0 (sizeofArray arr) marr {-# INLINE castArrayToVector #-} - --- | Covert boxed `VB.Vector` into an `A.Array`. Sliced vectors will indure copying. -vectorToArray :: VB.Vector a -> A.Array a -vectorToArray v = - runST $ do - VB.MVector start len marr <- VB.unsafeThaw v - marr' <- - if start == 0 - then return marr - else A.cloneMutableArray marr start len - A.unsafeFreezeArray marr' -{-# INLINE vectorToArray #-} -- | Cast a Boxed Vector into an Array, but only if it wasn't previously sliced. diff --git a/src/Data/Massiv/Array/Manifest/Internal.hs b/src/Data/Massiv/Array/Manifest/Internal.hs --- a/src/Data/Massiv/Array/Manifest/Internal.hs +++ b/src/Data/Massiv/Array/Manifest/Internal.hs @@ -1,16 +1,17 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE CPP #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE MagicHash #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MagicHash #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} -- | -- Module : Data.Massiv.Array.Manifest.Internal --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental @@ -20,13 +21,12 @@ ( M , Manifest(..) , Array(..) - , makeBoxedVector , toManifest , compute + , computeS , computeAs , computeProxy , computeSource - , computeInto , computeWithStride , computeWithStrideAs , clone @@ -34,41 +34,35 @@ , convertAs , convertProxy , gcastArr - , loadMutableS - , loadMutableOnP - , sequenceP - , sequenceOnP - , fromRaggedArray + , fromRaggedArrayM , fromRaggedArray' + , fromRaggedArray , sizeofArray , sizeofMutableArray ) where -import Control.Exception (try) -import Control.Monad (unless) -import Control.Monad.ST (runST) -import Data.Foldable (Foldable (..)) -import Data.Massiv.Array.Delayed.Internal -import Data.Massiv.Array.Ops.Fold.Internal as M -import Data.Massiv.Array.Unsafe -import Data.Massiv.Core.Common -import Data.Massiv.Core.List -import Data.Massiv.Core.Scheduler -import Data.Maybe (fromMaybe) -import Data.Typeable -import qualified Data.Vector as V -import GHC.Base hiding (ord) -import System.IO.Unsafe (unsafePerformIO) +import Control.Exception (try) +import Control.Monad.ST +import Control.Scheduler +import qualified Data.Foldable as F (Foldable(..)) +import Data.Massiv.Array.Delayed.Pull +import Data.Massiv.Array.Mutable +import Data.Massiv.Array.Ops.Fold.Internal +import Data.Massiv.Core.Common +import Data.Massiv.Core.List +import Data.Maybe (fromMaybe) +import Data.Typeable +import qualified Data.Vector as V +import qualified Data.Vector.Mutable as MV +import GHC.Base hiding (ord) +import System.IO.Unsafe (unsafePerformIO) #if MIN_VERSION_primitive(0,6,2) -import Data.Primitive.Array (sizeofArray, - sizeofMutableArray) +import Data.Primitive.Array (sizeofArray, sizeofMutableArray) #else -import qualified Data.Primitive.Array as A (Array (..), - MutableArray (..)) -import GHC.Prim (sizeofArray#, - sizeofMutableArray#) +import qualified Data.Primitive.Array as A (Array(..), MutableArray(..)) +import GHC.Exts (sizeofArray#, sizeofMutableArray#) sizeofArray :: A.Array a -> Int sizeofArray (A.Array a) = I# (sizeofArray# a) @@ -84,10 +78,15 @@ data M data instance Array M ix e = MArray { mComp :: !Comp - , mSize :: !ix + , mSize :: !(Sz ix) , mLinearIndex :: Int -> e } type instance EltRepr M ix = M +instance (Ragged L ix e, Show e) => Show (Array M ix e) where + showsPrec = showsArrayPrec id + showList = showArrayList + + instance (Eq e, Index ix) => Eq (Array M ix e) where (==) = eq (==) {-# INLINE (==) #-} @@ -97,19 +96,17 @@ {-# INLINE compare #-} instance Index ix => Construct M ix e where - getComp = mComp - {-# INLINE getComp #-} - - setComp c arr = arr { mComp = c } + setComp c arr = arr {mComp = c} {-# INLINE setComp #-} - - unsafeMakeArray !c !sz f = MArray c sz (V.unsafeIndex (makeBoxedVector sz f)) - {-# INLINE unsafeMakeArray #-} - --- | Create a boxed from usual size and index to element function -makeBoxedVector :: Index ix => Sz ix -> (ix -> a) -> V.Vector a -makeBoxedVector !sz f = V.generate (totalElem sz) (f . fromLinearIndex sz) -{-# INLINE makeBoxedVector #-} + makeArrayLinear !comp !sz f = + unsafePerformIO $ do + let !k = totalElem sz + mv <- MV.unsafeNew k + withScheduler_ comp $ \scheduler -> + splitLinearlyWithM_ scheduler k (pure . f) (MV.unsafeWrite mv) + v <- V.unsafeFreeze mv + pure $ MArray comp sz (V.unsafeIndex v) + {-# INLINE makeArrayLinear #-} -- | /O(1)/ - Conversion of `Manifest` arrays to `M` representation. @@ -118,8 +115,12 @@ {-# INLINE toManifest #-} --- | Row-major sequential folding over a Manifest array. +-- | Row-major sequentia folding over a Manifest array. instance Index ix => Foldable (Array M ix) where + fold = fold + {-# INLINE fold #-} + foldMap = foldMono + {-# INLINE foldMap #-} foldl = lazyFoldlS {-# INLINE foldl #-} foldl' = foldlS @@ -130,10 +131,6 @@ {-# INLINE foldr' #-} null (MArray _ sz _) = totalElem sz == 0 {-# INLINE null #-} - sum = foldl' (+) 0 - {-# INLINE sum #-} - product = foldl' (*) 1 - {-# INLINE product #-} length = totalElem . size {-# INLINE length #-} toList arr = build (\ c n -> foldrFB c n arr) @@ -152,13 +149,11 @@ {-# INLINE unsafeLinearIndexM #-} -instance Index ix => Size M ix e where - size = mSize - {-# INLINE size #-} - +instance Index ix => Resize M ix where unsafeResize !sz !arr = arr { mSize = sz } {-# INLINE unsafeResize #-} +instance Index ix => Extract M ix e where unsafeExtract !sIx !newSz !arr = MArray (getComp arr) newSz $ \ i -> unsafeIndex arr (liftIndex2 (+) (fromLinearIndex newSz i) sIx) @@ -167,7 +162,7 @@ instance {-# OVERLAPPING #-} Slice M Ix1 e where - unsafeSlice arr i _ _ = Just (unsafeLinearIndex arr i) + unsafeSlice arr i _ _ = pure (unsafeLinearIndex arr i) {-# INLINE unsafeSlice #-} instance ( Index ix @@ -176,7 +171,7 @@ ) => Slice M ix e where unsafeSlice arr start cutSz dim = do - newSz <- dropDim cutSz dim + (_, newSz) <- pullOutSzM cutSz dim return $ unsafeResize newSz (unsafeExtract start cutSz arr) {-# INLINE unsafeSlice #-} @@ -186,7 +181,7 @@ instance (Elt M ix e ~ Array M (Lower ix) e, Index ix, Index (Lower ix)) => OuterSlice M ix e where unsafeOuterSlice !arr !i = - MArray (getComp arr) (tailDim (size arr)) (unsafeLinearIndex arr . (+ kStart)) + MArray (getComp arr) (snd (unconsSz (size arr))) (unsafeLinearIndex arr . (+ kStart)) where !kStart = toLinearIndex (size arr) (consDim i (zeroIndex :: Lower ix)) {-# INLINE unsafeOuterSlice #-} @@ -197,69 +192,44 @@ instance (Elt M ix e ~ Array M (Lower ix) e, Index ix, Index (Lower ix)) => InnerSlice M ix e where unsafeInnerSlice !arr (szL, m) !i = - MArray (getComp arr) szL (\k -> unsafeLinearIndex arr (k * m + kStart)) + MArray (getComp arr) szL (\k -> unsafeLinearIndex arr (k * unSz m + kStart)) where !kStart = toLinearIndex (size arr) (snocDim (zeroIndex :: Lower ix) i) {-# INLINE unsafeInnerSlice #-} instance Index ix => Load M ix e where - loadS (MArray _ sz f) _ uWrite = - iterM_ 0 (totalElem sz) 1 (<) $ \ !i -> - uWrite i (f i) - {-# INLINE loadS #-} - loadP wIds (MArray _ sz f) _ uWrite = - divideWork_ wIds (totalElem sz) $ \ !scheduler !chunkLength !totalLength !slackStart -> do - loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start -> - scheduleWork scheduler $ - iterM_ start (start + chunkLength) 1 (<) $ \ !i -> - uWrite i (f i) - scheduleWork scheduler $ - iterM_ slackStart totalLength 1 (<) $ \ !i -> - uWrite i (f i) - {-# INLINE loadP #-} - loadArray numWorkers' scheduleWork' (MArray _ sz f) _ = - splitLinearlyWith_ numWorkers' scheduleWork' (totalElem sz) f - {-# INLINE loadArray #-} - - -loadMutableS :: (Load r' ix e, Mutable r ix e) => - Array r' ix e -> Array r ix e -loadMutableS !arr = - runST $ do - mArr <- unsafeNew (size arr) - loadS arr (unsafeLinearRead mArr) (unsafeLinearWrite mArr) - unsafeFreeze Seq mArr -{-# INLINE loadMutableS #-} + size = mSize + {-# INLINE size #-} + getComp = mComp + {-# INLINE getComp #-} + loadArrayM scheduler (MArray _ sz f) = splitLinearlyWith_ scheduler (totalElem sz) f + {-# INLINE loadArrayM #-} -loadMutableOnP :: (Load r' ix e, Mutable r ix e) => - [Int] -> Array r' ix e -> IO (Array r ix e) -loadMutableOnP wIds !arr = do - mArr <- unsafeNew (size arr) - loadP wIds arr (unsafeLinearRead mArr) (unsafeLinearWrite mArr) - unsafeFreeze (ParOn wIds) mArr -{-# INLINE loadMutableOnP #-} +instance Index ix => StrideLoad M ix e -- | 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@ -compute :: (Load r' ix e, Mutable r ix e) => Array r' ix e -> Array r ix e -compute !arr = - case getComp arr of - Seq -> loadMutableS arr - ParOn wIds -> unsafePerformIO $ loadMutableOnP wIds arr +compute :: forall r ix e r' . (Mutable r ix e, Load r' ix e) => Array r' ix e -> Array r ix e +compute !arr = unsafePerformIO $ loadArray arr >>= unsafeFreeze (getComp arr) {-# INLINE compute #-} +computeS :: forall r ix e r' . (Mutable r ix e, Load r' ix e) => Array r' ix e -> Array r ix e +computeS !arr = runST $ loadArrayS arr >>= unsafeFreeze (getComp arr) +{-# INLINE computeS #-} + -- | Just as `compute`, but let's you supply resulting representation type as an argument. -- -- ====__Examples__ -- --- >>> computeAs P $ range Seq 0 10 --- (Array P Seq (10) --- [ 0,1,2,3,4,5,6,7,8,9 ]) +-- >>> import Data.Massiv.Array +-- >>> computeAs P $ range Seq (Ix1 0) 10 +-- Array P Seq (Sz1 10) +-- [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9 ] -- -computeAs :: (Load r' ix e, Mutable r ix e) => r -> Array r' ix e -> Array r ix e +computeAs :: (Mutable r ix e, Load r' ix e) => r -> Array r' ix e -> Array r ix e computeAs _ = compute {-# INLINE computeAs #-} @@ -267,69 +237,60 @@ -- | Same as `compute` and `computeAs`, but let's you supply resulting representation type as a proxy -- argument. -- --- @since 0.1.1 --- --- ====__Examples__ +-- ==== __Examples__ -- --- Useful for cases when representation constructor isn't available for some reason: +-- Useful only really for cases when representation constructor or @TypeApplications@ extension +-- aren't desireable for some reason: -- --- >>> computeProxy (Nothing :: Maybe P) $ range Seq 0 10 --- (Array P Seq (10) --- [ 0,1,2,3,4,5,6,7,8,9 ]) +-- >>> import Data.Proxy +-- >>> import Data.Massiv.Array +-- >>> computeProxy (Proxy :: Proxy P) $ (^ (2 :: Int)) <$> range Seq (Ix1 0) 10 +-- Array P Seq (Sz1 10) +-- [ 0, 1, 4, 9, 16, 25, 36, 49, 64, 81 ] -- -computeProxy :: (Load r' ix e, Mutable r ix e) => proxy r -> Array r' ix e -> Array r ix e +-- @since 0.1.1 +computeProxy :: (Mutable r ix e, Load r' ix e) => proxy r -> Array r' ix e -> Array r ix e computeProxy _ = compute {-# INLINE computeProxy #-} --- | Compute an Array while loading the results into the supplied mutable target array. Sizes of --- both arrays must agree, otherwise error. --- --- @since 0.1.3 -computeInto :: - (Load r' ix e, Mutable r ix e) - => MArray RealWorld r ix e -- ^ Target Array - -> Array r' ix e -- ^ Array to load - -> IO () -computeInto !mArr !arr = do - unless (msize mArr == size arr) $ errorSizeMismatch "computeInto" (msize mArr) (size arr) - case getComp arr of - Seq -> loadS arr (unsafeLinearRead mArr) (unsafeLinearWrite mArr) - ParOn wIds -> loadP wIds arr (unsafeLinearRead mArr) (unsafeLinearWrite mArr) -{-# INLINE computeInto #-} - - --- | This is just like `compute`, but can be applied to `Source` arrays and will be a noop if +-- | This is just like `convert`, but restricted to `Source` arrays. Will be a noop if -- resulting type is the same as the input. -computeSource :: forall r' r ix e . (Source r' ix e, Mutable r ix e) +-- +-- @since 0.1.0 +computeSource :: forall r ix e r' . (Mutable r ix e, Source r' ix e) => Array r' ix e -> Array r ix e -computeSource arr = - maybe (compute $ delay arr) (\Refl -> arr) (eqT :: Maybe (r' :~: r)) +computeSource arr = maybe (compute arr) (\Refl -> arr) (eqT :: Maybe (r' :~: r)) {-# INLINE computeSource #-} -- | /O(n)/ - Make an exact immutable copy of an Array. +-- +-- @since 0.1.0 clone :: Mutable r ix e => Array r ix e -> Array r ix e -clone = compute . toManifest +clone arr = unsafePerformIO $ thaw arr >>= unsafeFreeze (getComp arr) {-# INLINE clone #-} -- | /O(1)/ - Cast over Array representation -gcastArr :: forall r' r ix e. (Typeable r, Typeable r') +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 manifest types, except when source and result arrays --- are of the same representation, in which case it is an /O(1)/ operation. -convert :: (Manifest r' ix e, Mutable r ix e) +-- | /O(n)/ - conversion between array types. A full copy will occur, unless when the source and +-- result arrays are of the same representation, in which case it is an /O(1)/ operation. +-- +-- @since 0.1.0 +convert :: forall r ix e r' . (Mutable r ix e, Load r' ix e) => Array r' ix e -> Array r ix e -convert arr = - fromMaybe (compute $ toManifest arr) (gcastArr arr) +convert arr = fromMaybe (compute arr) (gcastArr arr) {-# INLINE convert #-} -- | Same as `convert`, but let's you supply resulting representation type as an argument. -convertAs :: (Manifest r' ix e, Mutable r ix e) +-- +-- @since 0.1.0 +convertAs :: (Mutable r ix e, Load r' ix e) => r -> Array r' ix e -> Array r ix e convertAs _ = convert {-# INLINE convertAs #-} @@ -339,101 +300,79 @@ -- proxy argument. -- -- @since 0.1.1 --- -convertProxy :: (Manifest r' ix e, Mutable r ix e) +convertProxy :: (Mutable r ix e, Load r' ix e) => proxy r -> Array r' ix e -> Array r ix e convertProxy _ = convert {-# INLINE convertProxy #-} -sequenceOnP :: (Source r1 ix (IO e), Mutable r ix e) => - [Int] -> Array r1 ix (IO e) -> IO (Array r ix e) -sequenceOnP wIds !arr = do - resArrM <- unsafeNew (size arr) - withScheduler_ wIds $ \scheduler -> - flip imapM_ arr $ \ !ix action -> - scheduleWork scheduler $ action >>= unsafeWrite resArrM ix - unsafeFreeze (getComp arr) resArrM -{-# INLINE sequenceOnP #-} - -sequenceP :: (Source r1 ix (IO e), Mutable r ix e) => Array r1 ix (IO e) -> IO (Array r ix e) -sequenceP = sequenceOnP [] -{-# INLINE sequenceP #-} - - -- | Convert a ragged array into a usual rectangular shaped one. -fromRaggedArray :: (Ragged r' ix e, Mutable r ix e) => - Array r' ix e -> Either ShapeError (Array r ix e) +fromRaggedArray :: (Mutable r ix e, Ragged r' ix e, Load r' ix e) => + Array r' ix e -> Either ShapeException (Array r ix e) fromRaggedArray arr = unsafePerformIO $ do - let sz = edgeSize arr + let !sz = edgeSize arr + !comp = getComp arr mArr <- unsafeNew sz - let loadWith using = loadRagged using (unsafeLinearWrite mArr) 0 (totalElem sz) sz arr - try $ - case getComp arr of - c -> do - loadWith id - unsafeFreeze c mArr - -- Seq -> do - -- loadWith id - -- unsafeFreeze Seq mArr - -- pComp@(ParOn ss) -> do - -- withScheduler_ ss (loadWith . scheduleWork) - -- unsafeFreeze pComp mArr + try $ do + withScheduler_ comp $ \scheduler -> + loadRagged (scheduleWork scheduler) (unsafeLinearWrite mArr) 0 (totalElem sz) sz arr + unsafeFreeze comp mArr {-# INLINE fromRaggedArray #-} +{-# DEPRECATED fromRaggedArray "In favor of a more general `fromRaggedArrayM`" #-} +-- | 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.3.0 +fromRaggedArrayM :: + forall r ix e r' m . (Mutable r ix e, Ragged r' ix e, Load r' ix e, MonadThrow m) + => Array r' ix e + -> m (Array r ix e) +fromRaggedArrayM arr = + let sz = edgeSize arr + in either (\(e :: ShapeException) -> throwM e) pure $ + unsafePerformIO $ do + marr <- unsafeNew sz + traverse (\_ -> unsafeFreeze (getComp arr) marr) =<< + try (withScheduler_ (getComp arr) $ \Scheduler {scheduleWork} -> + loadRagged scheduleWork (unsafeLinearWrite marr) 0 (totalElem sz) sz arr) +{-# INLINE fromRaggedArrayM #-} + + -- | Same as `fromRaggedArray`, but will throw an error if its shape is not -- rectangular. -fromRaggedArray' :: (Ragged r' ix e, Mutable r ix e) => - Array r' ix e -> Array r ix e -fromRaggedArray' arr = - case fromRaggedArray arr of - Left RowTooShortError -> error "Not enough elements in a row" - Left RowTooLongError -> error "Too many elements in a row" - Right resArr -> resArr +fromRaggedArray' :: + forall r ix e r'. (Mutable r ix e, Load r' ix e, Ragged r' ix e) + => Array r' ix e + -> Array r ix e +fromRaggedArray' arr = either throw id $ fromRaggedArrayM arr {-# INLINE fromRaggedArray' #-} - - -- | Same as `compute`, but with `Stride`. -computeWithStride :: (Load r' ix e, Mutable r ix e) => Stride ix -> Array r' ix e -> Array r ix e +-- +-- /O(n div k)/ - Where @n@ is numer of elements in the source array and @k@ is number of elemts in +-- the stride. +-- +-- @since 0.3.0 +computeWithStride :: + forall r ix e r'. (Mutable r ix e, StrideLoad r' ix e) + => Stride ix + -> Array r' ix e + -> Array r ix e computeWithStride stride !arr = unsafePerformIO $ do - let sz = strideSize stride (size arr) - comp = getComp arr - mArr <- unsafeNew sz - case comp of - Seq -> loadArrayWithStride 1 id stride sz arr (unsafeLinearRead mArr) (unsafeLinearWrite mArr) - ParOn wIds -> - withScheduler_ wIds $ \scheduler -> - loadArrayWithStride - (numWorkers scheduler) - (scheduleWork scheduler) - stride - sz - arr - (unsafeLinearRead mArr) - (unsafeLinearWrite mArr) - -- -- Alternative way to run computation sequentially: decreaseas compile time - -- let wIds = case comp of - -- Seq -> [1] - -- ParOn caps -> caps - -- withScheduler_ wIds $ \scheduler -> - -- loadArrayWithStride - -- (numWorkers scheduler) - -- (scheduleWork scheduler) - -- stride - -- sz - -- arr - -- (unsafeLinearRead mArr) - -- (unsafeLinearWrite mArr) - unsafeFreeze comp mArr + let !sz = strideSize stride (size arr) + createArray_ (getComp arr) sz $ \scheduler marr -> + loadArrayWithStrideM scheduler stride sz arr (unsafeLinearWrite marr) {-# INLINE computeWithStride #-} -- | Same as `computeWithStride`, but with ability to specify resulting array representation. +-- +-- @since 0.3.0 computeWithStrideAs :: - (Load r' ix e, Mutable r ix e) => r -> Stride ix -> Array r' ix e -> Array r ix e + (Mutable r ix e, StrideLoad r' ix e) => r -> Stride ix -> Array r' ix e -> Array r ix e computeWithStrideAs _ = computeWithStride {-# INLINE computeWithStrideAs #-} 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 @@ -1,13 +1,13 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} -- | -- Module : Data.Massiv.Array.Manifest.List --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental @@ -17,8 +17,9 @@ ( -- ** List fromList - , fromLists + , fromListsM , fromLists' + , fromLists , toList , toLists , toLists2 @@ -26,92 +27,109 @@ , toLists4 ) where -import Data.Massiv.Array.Manifest.Internal -import Data.Massiv.Array.Ops.Fold (foldrInner) -import Data.Massiv.Array.Ops.Fold.Internal (foldrFB) -import Data.Massiv.Core.Common -import Data.Massiv.Core.List -import GHC.Base (build) +import Data.Massiv.Array.Manifest.Internal +import Data.Massiv.Array.Ops.Fold (foldrInner) +import Data.Massiv.Array.Ops.Fold.Internal (foldrFB) +import Data.Massiv.Core.Common +import Data.Massiv.Core.List +import GHC.Base (build) -- | Convert a flat list into a vector -fromList :: (Nested LN Ix1 e, Nested L Ix1 e, Ragged L Ix1 e, Mutable r Ix1 e) - => Comp -- ^ Computation startegy to use - -> [e] -- ^ Flat list - -> Array r Ix1 e +-- +-- @since 0.1.0 +fromList :: + forall r e. Mutable r Ix1 e + => Comp -- ^ Computation startegy to use + -> [e] -- ^ Flat list + -> Array r Ix1 e fromList = 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. -- -- ==== __Examples__ -- --- >>> fromLists Seq [[1,2],[3,4]] :: Maybe (Array U Ix2 Int) --- Just (Array U Seq (2 :. 2) --- [ [ 1,2 ] --- , [ 3,4 ] --- ]) +-- >>> import Data.Massiv.Array as A +-- >>> fromListsM Seq [[1,2,3],[4,5,6]] :: Maybe (Array U Ix2 Int) +-- Just (Array U Seq (Sz (2 :. 3)) +-- [ [ 1, 2, 3 ] +-- , [ 4, 5, 6 ] +-- ] +-- ) -- --- >>> fromLists Par [[[1,2,3]],[[4,5,6]]] :: Maybe (Array U Ix3 Int) --- Just (Array U Par (2 :> 1 :. 3) --- [ [ [ 1,2,3 ] +-- >>> fromListsM Par [[[1,2,3]],[[4,5,6]]] :: Maybe (Array U Ix3 Int) +-- Just (Array U Par (Sz (2 :> 1 :. 3)) +-- [ [ [ 1, 2, 3 ] -- ] --- , [ [ 4,5,6 ] +-- , [ [ 4, 5, 6 ] -- ] --- ]) +-- ] +-- ) -- -- Elements of a boxed array could be lists themselves if necessary, but cannot be ragged: -- --- >>> fromLists Seq [[[1,2,3]],[[4,5]]] :: Maybe (Array B Ix2 [Int]) --- Just (Array B Seq (2 :. 1) +-- >>> fromListsM Seq [[[1,2,3]],[[4,5]]] :: Maybe (Array B Ix2 [Int]) +-- Just (Array B Seq (Sz (2 :. 1)) -- [ [ [1,2,3] ] -- , [ [4,5] ] --- ]) --- >>> fromLists Seq [[[1,2,3]],[[4,5]]] :: Maybe (Array B Ix3 Int) --- Nothing +-- ] +-- ) +-- >>> fromListsM Seq [[[1,2,3]],[[4,5]]] :: IO (Array B Ix3 Int) +-- *** Exception: DimTooShortException: expected (Sz1 3), got (Sz1 2) -- -fromLists :: (Nested LN ix e, Nested L ix e, Ragged L ix e, Mutable r ix e) +-- @since 0.3.0 +fromListsM :: forall r ix e m . (Nested LN ix e, Ragged L ix e, Mutable r ix e, MonadThrow m) + => Comp -> [ListItem ix e] -> m (Array r ix e) +fromListsM comp = fromRaggedArrayM . setComp comp . throughNested +{-# INLINE fromListsM #-} + +-- | Similar to `fromListsM`, but less general. +fromLists :: (Nested LN ix e, Ragged L ix e, Mutable r ix e) => Comp -> [ListItem ix e] -> Maybe (Array r ix e) -fromLists comp = either (const Nothing) Just . fromRaggedArray . setComp comp . throughNested +fromLists comp = fromRaggedArrayM . setComp comp . throughNested {-# INLINE fromLists #-} - +{-# DEPRECATED fromLists "In favor of a more general `fromListsM`" #-} +-- TODO: Figure out QuickCheck properties. Best guess idea so far IMHO is to add it as dependency +-- and move Arbitrary instances int the library +-- +-- prop> fromLists' Seq xs == fromList xs +-- -- | Same as `fromLists`, but will throw an error on irregular shaped lists. -- -- __Note__: This function is the same as if you would turn on @{-\# LANGUAGE OverloadedLists #-}@ -- extension. For that reason you can also use `GHC.Exts.fromList`. -- --- prop> fromLists' Seq xs == fromList xs --- --- ===__Examples__ +-- ====__Examples__ -- -- Convert a list of lists into a 2D Array -- --- >>> fromLists' Seq [[1,2],[3,4]] :: Array U Ix2 Int --- (Array U Seq (2 :. 2) --- [ [ 1,2 ] --- , [ 3,4 ] --- ]) +-- >>> import Data.Massiv.Array as A +-- >>> fromLists' Seq [[1,2,3],[4,5,6]] :: Array U Ix2 Int +-- Array U Seq (Sz (2 :. 3)) +-- [ [ 1, 2, 3 ] +-- , [ 4, 5, 6 ] +-- ] -- -- Above example implemented using GHC's `OverloadedLists` extension: -- -- >>> :set -XOverloadedLists --- >>> [[1,2],[3,4]] :: Array U Ix2 Int --- (Array U Seq (2 :. 2) --- [ [ 1,2 ] --- , [ 3,4 ] --- ]) +-- >>> [[1,2,3],[4,5,6]] :: Array U Ix2 Int +-- Array U Seq (Sz (2 :. 3)) +-- [ [ 1, 2, 3 ] +-- , [ 4, 5, 6 ] +-- ] -- -- Example of failure on conversion of an irregular nested list. -- -- >>> fromLists' Seq [[1],[3,4]] :: Array U Ix2 Int --- (Array U *** Exception: Too many elements in a row +-- Array U *** Exception: DimTooLongException -- -fromLists' :: (Nested LN ix e, Nested L ix e, Ragged L ix e, Mutable r ix e) +-- @since 0.1.0 +fromLists' :: forall r ix e . (Nested LN ix e, Ragged L ix e, Mutable r ix e) => Comp -- ^ Computation startegy to use -> [ListItem ix e] -- ^ Nested list -> Array r ix e @@ -119,7 +137,7 @@ {-# INLINE fromLists' #-} -throughNested :: forall ix e . (Nested LN ix e, Nested L ix e) => [ListItem ix e] -> Array L ix e +throughNested :: forall ix e . Nested LN ix e => [ListItem ix e] -> Array L ix e throughNested xs = fromNested (fromNested xs :: Array LN ix e) {-# INLINE throughNested #-} @@ -129,33 +147,38 @@ -- -- ==== __Examples__ -- --- >>> toList $ makeArrayR U Seq (2 :. 3) fromIx2 +-- >>> import Data.Massiv.Array +-- >>> toList $ makeArrayR U Seq (Sz (2 :. 3)) fromIx2 -- [(0,0),(0,1),(0,2),(1,0),(1,1),(1,2)] -- +-- @since 0.1.0 toList :: Source r ix e => Array r ix e -> [e] toList !arr = build (\ c n -> foldrFB c n arr) {-# INLINE toList #-} --- | /O(n)/ - Convert an array into a nested list. Array rank and list nestedness will always match, --- but you can use `toList`, `toLists2`, etc. if flattening of inner dimensions is desired. +-- | /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. -- -- __Note__: This function is almost the same as `GHC.Exts.toList`. -- --- ====__Examples__ +-- ==== __Examples__ -- --- >>> let arr = makeArrayR U Seq (2 :> 1 :. 3) fromIx3 --- >>> print arr --- (Array U Seq (2 :> 1 :. 3) --- [ [ [ (0,0,0),(0,0,1),(0,0,2) ] +-- >>> import Data.Massiv.Array +-- >>> arr = makeArrayR U Seq (Sz (2 :> 1 :. 3)) id +-- >>> arr +-- Array U Seq (Sz (2 :> 1 :. 3)) +-- [ [ [ 0 :> 0 :. 0, 0 :> 0 :. 1, 0 :> 0 :. 2 ] -- ] --- , [ [ (1,0,0),(1,0,1),(1,0,2) ] +-- , [ [ 1 :> 0 :. 0, 1 :> 0 :. 1, 1 :> 0 :. 2 ] -- ] --- ]) +-- ] -- >>> toLists arr --- [[[(0,0,0),(0,0,1),(0,0,2)]],[[(1,0,0),(1,0,1),(1,0,2)]]] +-- [[[0 :> 0 :. 0,0 :> 0 :. 1,0 :> 0 :. 2]],[[1 :> 0 :. 0,1 :> 0 :. 1,1 :> 0 :. 2]]] -- -toLists :: (Nested LN ix e, Nested L ix e, Construct L ix e, Source r ix e) +-- @since 0.1.0 +toLists :: (Nested LN ix e, Construct L ix e, Source r ix e) => Array r ix e -> [ListItem ix e] toLists = toNested . toNested . toListArray @@ -168,11 +191,13 @@ -- -- ==== __Examples__ -- --- >>> toList2 $ makeArrayR U Seq (2 :. 3) fromIx2 +-- >>> import Data.Massiv.Array +-- >>> toLists2 $ makeArrayR U Seq (Sz2 2 3) fromIx2 -- [[(0,0),(0,1),(0,2)],[(1,0),(1,1),(1,2)]] --- >>> toList2 $ makeArrayR U Seq (2 :> 1 :. 3) fromIx3 +-- >>> toLists2 $ makeArrayR U Seq (Sz3 2 1 3) fromIx3 -- [[(0,0,0),(0,0,1),(0,0,2)],[(1,0,0),(1,0,1),(1,0,2)]] -- +-- @since 0.1.0 toLists2 :: (Source r ix e, Index (Lower ix)) => Array r ix e -> [[e]] toLists2 = toList . foldrInner (:) [] {-# INLINE toLists2 #-} @@ -180,12 +205,16 @@ -- | Convert an array with at least 3 dimensions into a 3 deep nested list. Inner dimensions will -- get flattened. +-- +-- @since 0.1.0 toLists3 :: (Index (Lower (Lower ix)), Index (Lower ix), Source r ix e) => Array r ix e -> [[[e]]] toLists3 = toList . foldrInner (:) [] . foldrInner (:) [] {-# INLINE toLists3 #-} -- | Convert an array with at least 4 dimensions into a 4 deep nested list. Inner dimensions will -- get flattened. +-- +-- @since 0.1.0 toLists4 :: ( Index (Lower (Lower (Lower ix))) , Index (Lower (Lower ix)) 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,16 +1,16 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE CPP #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE MagicHash #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MagicHash #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE UnboxedTuples #-} -{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UnboxedTuples #-} +{-# LANGUAGE UndecidableInstances #-} -- | -- Module : Data.Massiv.Array.Manifest.Primitive --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental @@ -20,31 +20,39 @@ ( P(..) , Array(..) , Prim - , vectorToByteArray , toByteArray + , fromByteArrayM , fromByteArray , toMutableByteArray + , fromMutableByteArrayM , fromMutableByteArray + , unsafeAtomicReadIntArray + , unsafeAtomicWriteIntArray + , unsafeCasIntArray + , unsafeAtomicModifyIntArray + , unsafeAtomicAddIntArray + , unsafeAtomicSubIntArray + , unsafeAtomicAndIntArray + , unsafeAtomicNandIntArray + , unsafeAtomicOrIntArray + , unsafeAtomicXorIntArray ) where -import Control.DeepSeq (NFData (..), deepseq) -import Control.Monad.ST (runST) -import Data.Massiv.Array.Delayed.Internal (eq, ord) -import Data.Massiv.Array.Manifest.Internal -import Data.Massiv.Array.Manifest.List as A -import Data.Massiv.Array.Mutable -import Data.Massiv.Array.Unsafe (unsafeGenerateArray, - unsafeGenerateArrayP) -import Data.Massiv.Core.Common -import Data.Massiv.Core.List -import Data.Primitive (sizeOf) -import Data.Primitive.ByteArray -import Data.Primitive.Types -import qualified Data.Vector.Primitive as VP -import GHC.Base (Int (..)) -import GHC.Exts as GHC (IsList (..)) -import GHC.Prim -import Prelude hiding (mapM) +import Control.DeepSeq (NFData(..), deepseq) +import Control.Monad.Primitive (PrimMonad(primitive), PrimState, 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.List +import Data.Primitive (sizeOf) +import Data.Primitive.ByteArray +import Data.Primitive.Types +import GHC.Base (Int(..)) +import GHC.Exts as GHC +import Prelude hiding (mapM) +import System.IO.Unsafe (unsafePerformIO) #include "massiv.h" @@ -54,10 +62,14 @@ type instance EltRepr P ix = M data instance Array P ix e = PArray { pComp :: !Comp - , pSize :: !ix + , pSize :: !(Sz ix) , pData :: {-# UNPACK #-} !ByteArray } +instance (Ragged L ix e, Show e, Prim e) => Show (Array P ix e) where + showsPrec = showsArrayPrec id + showList = showArrayList + instance Index ix => NFData (Array P ix e) where rnf (PArray c sz a) = c `deepseq` sz `deepseq` a `seq` () {-# INLINE rnf #-} @@ -71,40 +83,30 @@ {-# INLINE compare #-} instance (Prim e, Index ix) => Construct P ix e where - getComp = pComp - {-# INLINE getComp #-} - setComp c arr = arr { pComp = c } {-# INLINE setComp #-} - unsafeMakeArray Seq !sz f = unsafeGenerateArray sz f - unsafeMakeArray (ParOn wIds) !sz f = unsafeGenerateArrayP wIds sz f - {-# INLINE unsafeMakeArray #-} - -elemsByteArray :: Prim a => a -> ByteArray -> Int -elemsByteArray dummy a = sizeofByteArray a `div` sizeOf dummy -{-# INLINE elemsByteArray #-} + makeArray !comp !sz f = unsafePerformIO $ generateArray comp sz (return . f) + {-# INLINE makeArray #-} instance (Prim e, Index ix) => Source P ix e where - unsafeLinearIndex (PArray _ _ a) = + unsafeLinearIndex _pa@(PArray _ _ a) = INDEX_CHECK("(Source P ix e).unsafeLinearIndex", - elemsByteArray (undefined :: e), indexByteArray) a + Sz . elemsBA _pa, indexByteArray) a {-# INLINE unsafeLinearIndex #-} -instance (Prim e, Index ix) => Size P ix e where - size = pSize - {-# INLINE size #-} - +instance Index ix => Resize P ix where unsafeResize !sz !arr = arr { pSize = sz } {-# INLINE unsafeResize #-} +instance (Prim e, Index ix) => Extract P ix e where unsafeExtract !sIx !newSz !arr = unsafeExtract sIx newSz (toManifest arr) {-# INLINE unsafeExtract #-} instance {-# OVERLAPPING #-} Prim e => Slice P Ix1 e where - unsafeSlice arr i _ _ = Just (unsafeLinearIndex arr i) + unsafeSlice arr i _ _ = pure (unsafeLinearIndex arr i) {-# INLINE unsafeSlice #-} @@ -149,18 +151,14 @@ instance (Index ix, Prim e) => Manifest P ix e where - unsafeLinearIndexM (PArray _ _ a) = + unsafeLinearIndexM _pa@(PArray _ _ a) = INDEX_CHECK("(Manifest P ix e).unsafeLinearIndexM", - elemsByteArray (undefined :: e), indexByteArray) a + Sz . elemsBA _pa, indexByteArray) a {-# INLINE unsafeLinearIndexM #-} -elemsMutableByteArray :: Prim a => a -> MutableByteArray s -> Int -elemsMutableByteArray dummy a = sizeofMutableByteArray a `div` sizeOf dummy -{-# INLINE elemsMutableByteArray #-} - instance (Index ix, Prim e) => Mutable P ix e where - data MArray s P ix e = MPArray !ix !(MutableByteArray s) + data MArray s P ix e = MPArray !(Sz ix) {-# UNPACK #-} !(MutableByteArray s) msize (MPArray sz _) = sz {-# INLINE msize #-} @@ -171,55 +169,43 @@ unsafeFreeze comp (MPArray sz a) = PArray comp sz <$> unsafeFreezeByteArray a {-# INLINE unsafeFreeze #-} - unsafeNew sz = MPArray sz <$> newByteArray (I# (totalSize# sz (undefined :: e))) + unsafeNew sz + | n <= (maxBound :: Int) `div` eSize = MPArray sz <$> newByteArray (n * eSize) + | otherwise = error $ "Array size is too big: " ++ show sz + where !n = totalElem sz + !eSize = sizeOf (undefined :: e) {-# INLINE unsafeNew #-} - unsafeNewZero sz = do - let !szBytes = I# (totalSize# sz (undefined :: e)) - barr <- newByteArray szBytes - fillByteArray barr 0 szBytes 0 - return $ MPArray sz barr - {-# INLINE unsafeNewZero #-} + initialize (MPArray sz mba) = + fillByteArray mba 0 (totalElem sz * sizeOf (undefined :: e)) 0 + {-# INLINE initialize #-} - unsafeLinearRead (MPArray _ ma) = + unsafeLinearRead _mpa@(MPArray _ ma) = INDEX_CHECK("(Mutable P ix e).unsafeLinearRead", - elemsMutableByteArray (undefined :: e), readByteArray) ma + Sz . elemsMBA _mpa, readByteArray) ma {-# INLINE unsafeLinearRead #-} - unsafeLinearWrite (MPArray _ ma) = + unsafeLinearWrite _mpa@(MPArray _ ma) = INDEX_CHECK("(Mutable P ix e).unsafeLinearWrite", - elemsMutableByteArray (undefined :: e), writeByteArray) ma + Sz . elemsMBA _mpa, writeByteArray) ma {-# INLINE unsafeLinearWrite #-} - unsafeNewA sz (State s#) = - let kb# = totalSize# sz (undefined :: e) - !(# s'#, mba# #) = newByteArray# kb# s# in - pure (State s'#, MPArray sz (MutableByteArray mba#)) - {-# INLINE unsafeNewA #-} - - unsafeThawA (PArray _ sz (ByteArray ba#)) s = - pure (s, MPArray sz (MutableByteArray (unsafeCoerce# ba#))) - {-# INLINE unsafeThawA #-} - - unsafeFreezeA comp (MPArray sz (MutableByteArray mba#)) (State s#) = - let !(# s'#, ba# #) = unsafeFreezeByteArray# mba# s# in - pure (State s'#, PArray comp sz (ByteArray ba#)) - {-# INLINE unsafeFreezeA #-} - - unsafeLinearWriteA (MPArray _ (MutableByteArray mba#)) (I# i#) val (State s#) = - pure (State (writeByteArray# mba# i# val s#)) - {-# INLINE unsafeLinearWriteA #-} - + unsafeLinearSet (MPArray _ ma) = setByteArray ma + {-# INLINE unsafeLinearSet #-} -totalSize# :: (Index ix, Prim e) => ix -> e -> Int# -totalSize# sz dummy = k# *# sizeOf# dummy - where - !(I# k#) = totalElem sz -{-# INLINE totalSize# #-} +instance (Prim e, Index ix) => Load P ix e where + size = pSize + {-# INLINE size #-} + getComp = pComp + {-# INLINE getComp #-} + loadArrayM !scheduler !arr = + splitLinearlyWith_ scheduler (elemsCount arr) (unsafeLinearIndex arr) + {-# INLINE loadArrayM #-} +instance (Prim e, Index ix) => StrideLoad P ix e -instance ( VP.Prim e +instance ( Prim e , IsList (Array L ix e) , Nested LN ix e , Nested L ix e @@ -233,21 +219,14 @@ {-# INLINE toList #-} -vectorToByteArray :: forall e . VP.Prim e => VP.Vector e -> ByteArray -vectorToByteArray (VP.Vector start len arr) = - if start == 0 - then arr - else runST $ do - marr <- newByteArray len - let elSize = sizeOf (undefined :: e) - copyByteArray marr 0 arr (start * elSize) (len * elSize) - unsafeFreezeByteArray marr -{-# INLINE vectorToByteArray #-} +elemsBA :: forall proxy e . Prim e => proxy e -> ByteArray -> Int +elemsBA _ a = sizeofByteArray a `div` sizeOf (undefined :: e) +{-# INLINE elemsBA #-} -primArrayDummy :: arr P ix e -> e -primArrayDummy = undefined -{-# INLINE primArrayDummy #-} +elemsMBA :: forall proxy e s . Prim e => proxy e -> MutableByteArray s -> Int +elemsMBA _ a = sizeofMutableByteArray a `div` sizeOf (undefined :: e) +{-# INLINE elemsMBA #-} -- | /O(1)/ - Extract the internal `ByteArray`. @@ -261,15 +240,21 @@ -- | /O(1)/ - Construct a primitive array from the `ByteArray`. Will return `Nothing` if number of -- elements doesn't match. -- --- @since 0.2.1 -fromByteArray :: (Index ix, Prim e) => Comp -> ix -> ByteArray -> Maybe (Array P ix e) -fromByteArray comp sz ba - | totalElem sz /= elemsByteArray (primArrayDummy arr) ba = Nothing - | otherwise = Just arr +-- @since 0.3.0 +fromByteArrayM :: (MonadThrow m, Index ix, Prim e) => Comp -> Sz ix -> ByteArray -> m (Array P ix e) +fromByteArrayM comp sz ba = + guardNumberOfElements sz (Sz (elemsBA arr ba)) >> pure arr where arr = PArray comp sz ba -{-# INLINE fromByteArray #-} +{-# INLINE fromByteArrayM #-} +-- | See `fromByteArrayM`. +-- +-- @since 0.2.1 +fromByteArray :: (Index ix, Prim e) => Comp -> Sz ix -> ByteArray -> Maybe (Array P ix e) +fromByteArray = fromByteArrayM +{-# INLINE fromByteArray #-} +{-# DEPRECATED fromByteArray "In favor of more general `fromByteArrayM`" #-} -- | /O(1)/ - Extract the internal `MutableByteArray`. @@ -283,11 +268,193 @@ -- | /O(1)/ - Construct a primitive mutable array from the `MutableByteArray`. Will return `Nothing` -- if number of elements doesn't match. -- --- @since 0.2.1 -fromMutableByteArray :: (Index ix, Prim e) => ix -> MutableByteArray s -> Maybe (MArray s P ix e) -fromMutableByteArray sz ba - | totalElem sz /= elemsMutableByteArray (primArrayDummy marr) ba = Nothing - | otherwise = Just marr +-- @since 0.3.0 +fromMutableByteArrayM :: + (MonadThrow m, Index ix, Prim e) => Sz ix -> MutableByteArray s -> m (MArray s P ix e) +fromMutableByteArrayM sz mba = + guardNumberOfElements sz (Sz (elemsMBA marr mba)) >> pure marr where - marr = MPArray sz ba + marr = MPArray sz mba +{-# INLINE fromMutableByteArrayM #-} + +-- | See `fromMutableByteArray`. +-- +-- @since 0.2.1 +fromMutableByteArray :: (Index ix, Prim e) => Sz ix -> MutableByteArray s -> Maybe (MArray s P ix e) +fromMutableByteArray = fromMutableByteArrayM {-# INLINE fromMutableByteArray #-} +{-# DEPRECATED fromMutableByteArray "In favor of more general `fromMutableByteArrayM`" #-} + + +-- | Atomically read an `Int` element from the array +-- +-- @since 0.3.0 +unsafeAtomicReadIntArray :: + (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> m Int +unsafeAtomicReadIntArray _mpa@(MPArray sz mba) ix = + INDEX_CHECK( "unsafeAtomicReadIntArray" + , Sz . elemsMBA _mpa + , \(MutableByteArray mba#) (I# i#) -> + primitive $ \s# -> + case atomicReadIntArray# mba# i# s# of + (# s'#, e# #) -> (# s'#, I# e# #)) + mba + (toLinearIndex sz ix) +{-# INLINE unsafeAtomicReadIntArray #-} + +-- | Atomically write an `Int` element int the array +-- +-- @since 0.3.0 +unsafeAtomicWriteIntArray :: + (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m () +unsafeAtomicWriteIntArray _mpa@(MPArray sz mba) ix (I# e#) = + INDEX_CHECK( "unsafeAtomicWriteIntArray" + , Sz . elemsMBA _mpa + , \(MutableByteArray mba#) (I# i#) -> + primitive_ (atomicWriteIntArray# mba# i# e#)) + mba + (toLinearIndex sz ix) +{-# INLINE unsafeAtomicWriteIntArray #-} + +-- | Atomically CAS an `Int` in the array. Returns the old value. +-- +-- @since 0.3.0 +unsafeCasIntArray :: + (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> Int -> m Int +unsafeCasIntArray _mpa@(MPArray sz mba) ix (I# e#) (I# n#) = + INDEX_CHECK( "unsafeCasIntArray" + , Sz . elemsMBA _mpa + , \(MutableByteArray mba#) (I# i#) -> + primitive $ \s# -> + case casIntArray# mba# i# e# n# s# of + (# s'#, o# #) -> (# s'#, I# o# #)) + mba + (toLinearIndex sz ix) +{-# INLINE unsafeCasIntArray #-} + + +-- | Atomically modify an `Int` element of the array. Returns the old value. +-- +-- @since 0.3.0 +unsafeAtomicModifyIntArray :: + (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> (Int -> Int) -> m Int +unsafeAtomicModifyIntArray _mpa@(MPArray sz mba) ix f = + INDEX_CHECK("unsafeAtomicModifyIntArray", Sz . elemsMBA _mpa, atomicModify) + mba + (toLinearIndex sz ix) + where + atomicModify (MutableByteArray mba#) (I# i#) = + let go s# o# = + let !(I# n#) = f (I# o#) + in case casIntArray# mba# i# o# n# s# of + (# s'#, o'# #) -> + case o# ==# o'# of + 0# -> go s# o'# + _ -> (# s'#, I# o# #) + in primitive $ \s# -> + case atomicReadIntArray# mba# i# s# of + (# s'#, o# #) -> go s'# o# + {-# INLINE atomicModify #-} +{-# INLINE unsafeAtomicModifyIntArray #-} + + +-- | Atomically add to an `Int` element in the array. Returns the old value. +-- +-- @since 0.3.0 +unsafeAtomicAddIntArray :: + (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int +unsafeAtomicAddIntArray _mpa@(MPArray sz mba) ix (I# e#) = + INDEX_CHECK( "unsafeAtomicAddIntArray" + , Sz . elemsMBA _mpa + , \(MutableByteArray mba#) (I# i#) -> + primitive $ \s# -> + case fetchAddIntArray# mba# i# e# s# of + (# s'#, p# #) -> (# s'#, I# p# #)) + mba + (toLinearIndex sz ix) +{-# INLINE unsafeAtomicAddIntArray #-} + + +-- | Atomically subtract from an `Int` element in the array. Returns the old value. +-- +-- @since 0.3.0 +unsafeAtomicSubIntArray :: + (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int +unsafeAtomicSubIntArray _mpa@(MPArray sz mba) ix (I# e#) = + INDEX_CHECK( "unsafeAtomicSubIntArray" + , Sz . elemsMBA _mpa + , \(MutableByteArray mba#) (I# i#) -> + primitive $ \s# -> + case fetchSubIntArray# mba# i# e# s# of + (# s'#, p# #) -> (# s'#, I# p# #)) + mba + (toLinearIndex sz ix) +{-# INLINE unsafeAtomicSubIntArray #-} + + +-- | Atomically AND an `Int` element in the array. Returns the old value. +-- +-- @since 0.3.0 +unsafeAtomicAndIntArray :: + (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int +unsafeAtomicAndIntArray _mpa@(MPArray sz mba) ix (I# e#) = + INDEX_CHECK( "unsafeAtomicAndIntArray" + , Sz . elemsMBA _mpa + , \(MutableByteArray mba#) (I# i#) -> + primitive $ \s# -> + case fetchAndIntArray# mba# i# e# s# of + (# s'#, p# #) -> (# s'#, I# p# #)) + mba + (toLinearIndex sz ix) +{-# INLINE unsafeAtomicAndIntArray #-} + + +-- | Atomically NAND an `Int` element in the array. Returns the old value. +-- +-- @since 0.3.0 +unsafeAtomicNandIntArray :: + (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int +unsafeAtomicNandIntArray _mpa@(MPArray sz mba) ix (I# e#) = + INDEX_CHECK( "unsafeAtomicNandIntArray" + , Sz . elemsMBA _mpa + , \(MutableByteArray mba#) (I# i#) -> + primitive $ \s# -> + case fetchNandIntArray# mba# i# e# s# of + (# s'#, p# #) -> (# s'#, I# p# #)) + mba + (toLinearIndex sz ix) +{-# INLINE unsafeAtomicNandIntArray #-} + + +-- | Atomically OR an `Int` element in the array. Returns the old value. +-- +-- @since 0.3.0 +unsafeAtomicOrIntArray :: + (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int +unsafeAtomicOrIntArray _mpa@(MPArray sz mba) ix (I# e#) = + INDEX_CHECK( "unsafeAtomicOrIntArray" + , Sz . elemsMBA _mpa + , \(MutableByteArray mba#) (I# i#) -> + primitive $ \s# -> + case fetchOrIntArray# mba# i# e# s# of + (# s'#, p# #) -> (# s'#, I# p# #)) + mba + (toLinearIndex sz ix) +{-# INLINE unsafeAtomicOrIntArray #-} + + +-- | Atomically XOR an `Int` element in the array. Returns the old value. +-- +-- @since 0.3.0 +unsafeAtomicXorIntArray :: + (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int +unsafeAtomicXorIntArray _mpa@(MPArray sz mba) ix (I# e#) = + INDEX_CHECK( "unsafeAtomicXorIntArray" + , Sz . elemsMBA _mpa + , \(MutableByteArray mba#) (I# i#) -> + primitive $ \s# -> + case fetchXorIntArray# mba# i# e# s# of + (# s'#, p# #) -> (# s'#, I# p# #)) + mba + (toLinearIndex sz ix) +{-# 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 @@ -1,14 +1,14 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE CPP #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} -- | -- Module : Data.Massiv.Array.Manifest.Storable --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental @@ -22,22 +22,30 @@ , toStorableMVector , withPtr , unsafeWithPtr + , unsafeArrayToForeignPtr + , unsafeMArrayToForeignPtr + , unsafeArrayFromForeignPtr + , unsafeArrayFromForeignPtr0 + , unsafeMArrayFromForeignPtr + , unsafeMArrayFromForeignPtr0 ) where -import Control.DeepSeq (NFData (..), deepseq) -import Data.Massiv.Array.Delayed.Internal (eq, ord) -import Data.Massiv.Array.Manifest.Internal -import Data.Massiv.Array.Manifest.List as A -import Data.Massiv.Array.Mutable -import Data.Massiv.Array.Unsafe (unsafeGenerateArray, - unsafeGenerateArrayP) -import Data.Massiv.Core.Common -import Data.Massiv.Core.List -import qualified Data.Vector.Storable as VS -import qualified Data.Vector.Storable.Mutable as MVS -import Foreign.Ptr -import GHC.Exts as GHC (IsList (..)) -import Prelude hiding (mapM) +import Control.DeepSeq (NFData(..), deepseq) +import Control.Monad.IO.Unlift +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.List +import qualified Data.Vector.Generic.Mutable as VGM +import qualified Data.Vector.Storable as VS +import qualified Data.Vector.Storable.Mutable as MVS +import Foreign.ForeignPtr +import Foreign.Ptr +import GHC.Exts as GHC (IsList(..)) +import Prelude hiding (mapM) +import System.IO.Unsafe (unsafePerformIO) #include "massiv.h" @@ -47,12 +55,17 @@ type instance EltRepr S ix = M data instance Array S ix e = SArray { sComp :: !Comp - , sSize :: !ix + , sSize :: !(Sz ix) , sData :: !(VS.Vector e) } -instance Index ix => NFData (Array S ix e) where +instance (Ragged L ix e, Show e, VS.Storable e) => Show (Array S ix e) where + showsPrec = showsArrayPrec id + showList = showArrayList + +instance NFData ix => NFData (Array S ix e) where rnf (SArray c sz v) = c `deepseq` sz `deepseq` v `deepseq` () + {-# INLINE rnf #-} instance (VS.Storable e, Eq e, Index ix) => Eq (Array S ix e) where (==) = eq (==) @@ -63,30 +76,23 @@ {-# INLINE compare #-} instance (VS.Storable e, Index ix) => Construct S ix e where - getComp = sComp - {-# INLINE getComp #-} - setComp c arr = arr { sComp = c } {-# INLINE setComp #-} - unsafeMakeArray Seq !sz f = unsafeGenerateArray sz f - unsafeMakeArray (ParOn wIds) !sz f = unsafeGenerateArrayP wIds sz f - {-# INLINE unsafeMakeArray #-} + makeArray !comp !sz f = unsafePerformIO $ generateArray comp sz (return . f) + {-# INLINE makeArray #-} instance (VS.Storable e, Index ix) => Source S ix e where unsafeLinearIndex (SArray _ _ v) = - INDEX_CHECK("(Source S ix e).unsafeLinearIndex", VS.length, VS.unsafeIndex) v + INDEX_CHECK("(Source S ix e).unsafeLinearIndex", Sz . VS.length, VS.unsafeIndex) v {-# INLINE unsafeLinearIndex #-} - -instance (VS.Storable e, Index ix) => Size S ix e where - size = sSize - {-# INLINE size #-} - +instance Index ix => Resize S ix where unsafeResize !sz !arr = arr { sSize = sz } {-# INLINE unsafeResize #-} +instance (VS.Storable e, Index ix) => Extract S ix e where unsafeExtract !sIx !newSz !arr = unsafeExtract sIx newSz (toManifest arr) {-# INLINE unsafeExtract #-} @@ -112,16 +118,20 @@ unsafeInnerSlice arr = unsafeInnerSlice (toManifest arr) {-# INLINE unsafeInnerSlice #-} +instance {-# OVERLAPPING #-} VS.Storable e => Slice S Ix1 e where + unsafeSlice arr i _ _ = pure (unsafeLinearIndex arr i) + {-# INLINE unsafeSlice #-} + instance (Index ix, VS.Storable e) => Manifest S ix e where unsafeLinearIndexM (SArray _ _ v) = - INDEX_CHECK("(Manifest S ix e).unsafeLinearIndexM", VS.length, VS.unsafeIndex) v + INDEX_CHECK("(Manifest S ix e).unsafeLinearIndexM", Sz . VS.length, VS.unsafeIndex) v {-# INLINE unsafeLinearIndexM #-} instance (Index ix, VS.Storable e) => Mutable S ix e where - data MArray s S ix e = MSArray !ix !(VS.MVector s e) + data MArray s S ix e = MSArray !(Sz ix) !(VS.MVector s e) msize (MSArray sz _) = sz {-# INLINE msize #-} @@ -135,18 +145,29 @@ unsafeNew sz = MSArray sz <$> MVS.unsafeNew (totalElem sz) {-# INLINE unsafeNew #-} - unsafeNewZero sz = MSArray sz <$> MVS.new (totalElem sz) - {-# INLINE unsafeNewZero #-} + initialize (MSArray _ marr) = VGM.basicInitialize marr + {-# INLINE initialize #-} unsafeLinearRead (MSArray _ mv) = - INDEX_CHECK("(Mutable S ix e).unsafeLinearRead", MVS.length, MVS.unsafeRead) mv + INDEX_CHECK("(Mutable S ix e).unsafeLinearRead", Sz . MVS.length, MVS.unsafeRead) mv {-# INLINE unsafeLinearRead #-} unsafeLinearWrite (MSArray _ mv) = - INDEX_CHECK("(Mutable S ix e).unsafeLinearWrite", MVS.length, MVS.unsafeWrite) mv + INDEX_CHECK("(Mutable S ix e).unsafeLinearWrite", Sz . MVS.length, MVS.unsafeWrite) mv {-# INLINE unsafeLinearWrite #-} +instance (Index ix, VS.Storable e) => Load S ix e where + size = sSize + {-# INLINE size #-} + getComp = sComp + {-# INLINE getComp #-} + loadArrayM !scheduler !arr = splitLinearlyWith_ scheduler (elemsCount arr) (unsafeLinearIndex arr) + {-# INLINE loadArrayM #-} + +instance (Index ix, VS.Storable e) => StrideLoad S ix e + + instance ( VS.Storable e , IsList (Array L ix e) , Nested LN ix e @@ -160,20 +181,21 @@ toList = GHC.toList . toListArray {-# INLINE toList #-} --- | A pointer to the beginning of the storable array. It is unsafe since it can break referential --- transparency. +-- | A pointer to the beginning of the storable array. It is unsafe since, if mutated, it can break +-- referential transparency. -- -- @since 0.1.3 -unsafeWithPtr :: VS.Storable a => Array S ix a -> (Ptr a -> IO b) -> IO b -unsafeWithPtr arr = VS.unsafeWith (sData arr) - +unsafeWithPtr :: (MonadUnliftIO m, VS.Storable a) => Array S ix a -> (Ptr a -> m b) -> m b +unsafeWithPtr arr f = withRunInIO $ \run -> VS.unsafeWith (sData arr) (run . f) +{-# INLINE unsafeWithPtr #-} -- | A pointer to the beginning of the mutable array. -- -- @since 0.1.3 -withPtr :: (Index ix, VS.Storable a) => MArray RealWorld S ix a -> (Ptr a -> IO b) -> IO b -withPtr (MSArray _ mv) = MVS.unsafeWith mv +withPtr :: (MonadUnliftIO m, VS.Storable a) => MArray RealWorld S ix a -> (Ptr a -> m b) -> m b +withPtr (MSArray _ mv) f = withRunInIO $ \run -> MVS.unsafeWith mv (run . f) +{-# INLINE withPtr #-} -- | /O(1)/ - Unwrap storable array and pull out the underlying storable vector. @@ -190,3 +212,55 @@ toStorableMVector :: MArray s S ix e -> VS.MVector s e toStorableMVector (MSArray _ mv) = mv {-# INLINE toStorableMVector #-} + + +-- | /O(1)/ - Yield the underlying `ForeignPtr` together with its length. +-- +-- @since 0.3.0 +unsafeArrayToForeignPtr :: VS.Storable e => Array S ix e -> (ForeignPtr e, Int) +unsafeArrayToForeignPtr = VS.unsafeToForeignPtr0 . toStorableVector +{-# INLINE unsafeArrayToForeignPtr #-} + +-- | /O(1)/ - Yield the underlying `ForeignPtr` together with its length. +-- +-- @since 0.3.0 +unsafeMArrayToForeignPtr :: VS.Storable e => MArray s S ix e -> (ForeignPtr e, Int) +unsafeMArrayToForeignPtr = MVS.unsafeToForeignPtr0 . toStorableMVector +{-# INLINE unsafeMArrayToForeignPtr #-} + +-- | /O(1)/ - Wrap a `ForeignPtr` and it's size into a pure storable array. +-- +-- @since 0.3.0 +unsafeArrayFromForeignPtr0 :: VS.Storable e => Comp -> ForeignPtr e -> Sz1 -> Array S Ix1 e +unsafeArrayFromForeignPtr0 comp ptr sz = + SArray {sComp = comp, sSize = sz, sData = VS.unsafeFromForeignPtr0 ptr (unSz sz)} +{-# INLINE unsafeArrayFromForeignPtr0 #-} + +-- | /O(1)/ - Wrap a `ForeignPtr`, an offset and it's size into a pure storable array. +-- +-- @since 0.3.0 +unsafeArrayFromForeignPtr :: VS.Storable e => Comp -> ForeignPtr e -> Int -> Sz1 -> Array S Ix1 e +unsafeArrayFromForeignPtr comp ptr offset sz = + SArray {sComp = comp, sSize = sz, sData = VS.unsafeFromForeignPtr ptr offset (unSz sz)} +{-# INLINE unsafeArrayFromForeignPtr #-} + + +-- | /O(1)/ - Wrap a `ForeignPtr` and it's size into a mutable storable array. It is still safe to +-- modify the pointer, unless the array gets frozen prior to modification. +-- +-- @since 0.3.0 +unsafeMArrayFromForeignPtr0 :: VS.Storable e => ForeignPtr e -> Sz1 -> MArray s S Ix1 e +unsafeMArrayFromForeignPtr0 fp sz = + MSArray sz (MVS.unsafeFromForeignPtr0 fp (unSz sz)) +{-# INLINE unsafeMArrayFromForeignPtr0 #-} + + +-- | /O(1)/ - Wrap a `ForeignPtr`, an offset and it's size into a mutable storable array. It is +-- still safe to modify the pointer, unless the array gets frozen prior to modification. +-- +-- @since 0.3.0 +unsafeMArrayFromForeignPtr :: VS.Storable e => ForeignPtr e -> Int -> Sz1 -> MArray s S Ix1 e +unsafeMArrayFromForeignPtr fp offset sz = + MSArray sz (MVS.unsafeFromForeignPtr fp offset (unSz sz)) +{-# INLINE unsafeMArrayFromForeignPtr #-} + 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,14 +1,15 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE CPP #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} -- | -- Module : Data.Massiv.Array.Manifest.Unboxed --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental @@ -22,19 +23,19 @@ , toUnboxedMVector ) where -import Control.DeepSeq (NFData (..), deepseq) -import Data.Massiv.Array.Delayed.Internal (eq, ord) -import Data.Massiv.Array.Manifest.Internal (M, toManifest) -import Data.Massiv.Array.Manifest.List as A -import Data.Massiv.Array.Mutable -import Data.Massiv.Array.Unsafe (unsafeGenerateArray, - unsafeGenerateArrayP) -import Data.Massiv.Core.Common -import Data.Massiv.Core.List -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 Control.DeepSeq (NFData(..), deepseq) +import Data.Massiv.Array.Delayed.Pull (eq, ord) +import Data.Massiv.Array.Manifest.Internal (M, toManifest) +import Data.Massiv.Array.Manifest.List as A +import Data.Massiv.Array.Mutable +import Data.Massiv.Core.Common +import Data.Massiv.Core.List +import qualified Data.Vector.Generic.Mutable as VGM +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 System.IO.Unsafe (unsafePerformIO) #include "massiv.h" @@ -44,25 +45,25 @@ type instance EltRepr U ix = M data instance Array U ix e = UArray { uComp :: !Comp - , uSize :: !ix + , uSize :: !(Sz ix) , uData :: !(VU.Vector e) } +instance (Ragged L ix e, Show e, VU.Unbox e) => Show (Array U ix e) where + showsPrec = showsArrayPrec id + showList = showArrayList -instance (Index ix, NFData e) => NFData (Array U ix e) where +instance NFData ix => NFData (Array U ix e) where rnf (UArray c sz v) = c `deepseq` sz `deepseq` v `deepseq` () + {-# INLINE rnf #-} instance (VU.Unbox e, Index ix) => Construct U ix e where - getComp = uComp - {-# INLINE getComp #-} - setComp c arr = arr { uComp = c } {-# INLINE setComp #-} - unsafeMakeArray Seq !sz f = unsafeGenerateArray sz f - unsafeMakeArray (ParOn wIds) !sz f = unsafeGenerateArrayP wIds sz f - {-# INLINE unsafeMakeArray #-} + makeArray !comp !sz f = unsafePerformIO $ generateArray comp sz (return . f) + {-# INLINE makeArray #-} instance (VU.Unbox e, Eq e, Index ix) => Eq (Array U ix e) where @@ -76,23 +77,31 @@ instance (VU.Unbox e, Index ix) => Source U ix e where unsafeLinearIndex (UArray _ _ v) = - INDEX_CHECK("(Source U ix e).unsafeLinearIndex", VU.length, VU.unsafeIndex) v + INDEX_CHECK("(Source U ix e).unsafeLinearIndex", Sz . VU.length, VU.unsafeIndex) v {-# INLINE unsafeLinearIndex #-} -instance (VU.Unbox e, Index ix) => Size U ix e where - size = uSize - {-# INLINE size #-} - +instance Index ix => Resize U ix where unsafeResize !sz !arr = arr { uSize = sz } {-# INLINE unsafeResize #-} +instance (VU.Unbox e, Index ix) => Extract U ix e where unsafeExtract !sIx !newSz !arr = unsafeExtract sIx newSz (toManifest arr) {-# INLINE unsafeExtract #-} +instance (VU.Unbox e, Index ix) => Load U ix e where + size = uSize + {-# INLINE size #-} + getComp = uComp + {-# INLINE getComp #-} + loadArrayM !scheduler !arr = splitLinearlyWith_ scheduler (elemsCount arr) (unsafeLinearIndex arr) + {-# INLINE loadArrayM #-} +instance (VU.Unbox e, Index ix) => StrideLoad U ix e + + instance {-# OVERLAPPING #-} VU.Unbox e => Slice U Ix1 e where - unsafeSlice arr i _ _ = Just (unsafeLinearIndex arr i) + unsafeSlice arr i _ _ = pure (unsafeLinearIndex arr i) {-# INLINE unsafeSlice #-} @@ -138,11 +147,12 @@ instance (VU.Unbox e, Index ix) => Manifest U ix e where unsafeLinearIndexM (UArray _ _ v) = - INDEX_CHECK("(Manifest U ix e).unsafeLinearIndexM", VU.length, VU.unsafeIndex) v + INDEX_CHECK("(Manifest U ix e).unsafeLinearIndexM", Sz . VU.length, VU.unsafeIndex) v {-# INLINE unsafeLinearIndexM #-} + instance (VU.Unbox e, Index ix) => Mutable U ix e where - data MArray s U ix e = MUArray ix (VU.MVector s e) + data MArray s U ix e = MUArray !(Sz ix) !(VU.MVector s e) msize (MUArray sz _) = sz {-# INLINE msize #-} @@ -156,15 +166,15 @@ unsafeNew sz = MUArray sz <$> MVU.unsafeNew (totalElem sz) {-# INLINE unsafeNew #-} - unsafeNewZero sz = MUArray sz <$> MVU.new (totalElem sz) - {-# INLINE unsafeNewZero #-} + initialize (MUArray _ marr) = VGM.basicInitialize marr + {-# INLINE initialize #-} unsafeLinearRead (MUArray _ mv) = - INDEX_CHECK("(Mutable U ix e).unsafeLinearRead", MVU.length, MVU.unsafeRead) mv + INDEX_CHECK("(Mutable U ix e).unsafeLinearRead", Sz . MVU.length, MVU.unsafeRead) mv {-# INLINE unsafeLinearRead #-} unsafeLinearWrite (MUArray _ mv) = - INDEX_CHECK("(Mutable U ix e).unsafeLinearWrite", MVU.length, MVU.unsafeWrite) mv + INDEX_CHECK("(Mutable U ix e).unsafeLinearWrite", Sz . MVU.length, MVU.unsafeWrite) mv {-# INLINE unsafeLinearWrite #-} 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 @@ -1,19 +1,21 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} -- | -- Module : Data.Massiv.Array.Manifest.Vector --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable -- module Data.Massiv.Array.Manifest.Vector - ( fromVector + ( fromVectorM + , fromVector' + , fromVector , castFromVector , toVector , castToVector @@ -21,35 +23,36 @@ , VRepr ) where -import Control.Monad (guard, join, msum) -import Data.Massiv.Array.Manifest.Boxed -import Data.Massiv.Array.Manifest.Internal -import Data.Massiv.Array.Manifest.Primitive -import Data.Massiv.Array.Manifest.Storable -import Data.Massiv.Array.Manifest.Unboxed -import Data.Massiv.Array.Mutable -import Data.Massiv.Core.Common -import Data.Typeable -import qualified Data.Vector as VB -import qualified Data.Vector.Generic as VG -import qualified Data.Vector.Primitive as VP -import qualified Data.Vector.Storable as VS -import qualified Data.Vector.Unboxed as VU +import Control.Monad (guard, join, msum) +import Data.Massiv.Array.Manifest.Boxed +import Data.Massiv.Array.Manifest.Internal +import Data.Massiv.Array.Manifest.Primitive +import Data.Massiv.Array.Manifest.Storable +import Data.Massiv.Array.Manifest.Unboxed +import Data.Massiv.Array.Mutable +import Data.Massiv.Core.Common +import Data.Maybe (fromMaybe) +import Data.Typeable +import qualified Data.Vector as VB +import qualified Data.Vector.Generic as VG +import qualified Data.Vector.Primitive as VP +import qualified Data.Vector.Storable as VS +import qualified Data.Vector.Unboxed as VU -- | Match vector type to array representation -type family ARepr (v :: * -> *) :: * --- | Match array representation to a vector type -type family VRepr r :: * -> * +type family ARepr (v :: * -> *) :: * where + ARepr VU.Vector = U + ARepr VS.Vector = S + ARepr VP.Vector = P + ARepr VB.Vector = B -type instance ARepr VU.Vector = U -type instance ARepr VS.Vector = S -type instance ARepr VP.Vector = P -type instance ARepr VB.Vector = B -type instance VRepr U = VU.Vector -type instance VRepr S = VS.Vector -type instance VRepr P = VP.Vector -type instance VRepr B = VB.Vector -type instance VRepr N = VB.Vector +-- | Match array representation to a vector type +type family VRepr r :: * -> * where + VRepr U = VU.Vector + VRepr S = VS.Vector + VRepr P = VP.Vector + VRepr B = VB.Vector + VRepr N = VB.Vector -- | /O(1)/ - conversion from vector to an array with a corresponding @@ -85,28 +88,55 @@ -- 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. -fromVector :: - (Typeable v, VG.Vector v a, Mutable (ARepr v) ix a, Mutable r ix a) +-- +-- @since 0.3.0 +fromVectorM :: + ( MonadThrow m + , Typeable v + , VG.Vector v a + , Mutable (ARepr v) ix a + , Construct r ix a + , Mutable r ix a + ) => Comp -> Sz ix -- ^ Resulting size of the array -> v a -- ^ Source Vector - -> Array r ix a -fromVector comp sz v = + -> m (Array r ix a) +fromVectorM comp sz v = case castFromVector comp sz v of - Just arr -> convert arr - Nothing -> - if (totalElem sz /= VG.length v) - then error $ - "Data.Array.Massiv.Manifest.fromVector: Supplied size: " ++ - show sz ++ " doesn't match vector length: " ++ show (VG.length v) - else makeArray comp sz ((VG.unsafeIndex v) . toLinearIndex sz) -{-# NOINLINE fromVector #-} + Just arr -> pure $ convert arr + Nothing -> do + guardNumberOfElements sz (Sz (VG.length v)) + pure (makeArrayLinear comp sz (VG.unsafeIndex v)) +{-# NOINLINE fromVectorM #-} +-- | Just like `fromVectorM`, but will throw an exception on a mismatched size. +-- +-- @since 0.3.0 +fromVector' :: + (Typeable v, VG.Vector v a, Mutable (ARepr v) ix a, Construct r ix a, Mutable r ix a) + => Comp + -> Sz ix -- ^ Resulting size of the array + -> v a -- ^ Source Vector + -> Array r ix a +fromVector' comp sz = either throw id . fromVectorM comp sz +{-# INLINE fromVector' #-} + +fromVector :: + (Typeable v, VG.Vector v a, Mutable (ARepr v) ix a, Construct r ix a, Mutable r ix a) + => Comp + -> Sz ix -- ^ Resulting size of the array + -> v a -- ^ Source Vector + -> Array r ix a +fromVector comp sz = either throw id . fromVectorM comp sz +{-# INLINE fromVector #-} +{-# DEPRECATED fromVector "In favor of safer `fromVectorM`" #-} + -- | /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 . (VG.Vector v e, Mutable r ix e, VRepr r ~ v) +castToVector :: forall v r ix e . (Mutable r ix e, VRepr r ~ v) => Array r ix e -> Maybe (v e) castToVector arr = msum @@ -137,8 +167,9 @@ -- In this example a `S`torable Array is created and then casted into a Storable -- `VS.Vector` in costant time: -- +-- >>> import Data.Massiv.Array as A -- >>> import qualified Data.Vector.Storable as VS --- >>> toVector (makeArrayR S Par (5 :. 6) (\(i :. j) -> i + j)) :: VS.Vector Int +-- >>> toVector (makeArrayR S Par (Sz2 5 6) (\(i :. j) -> i + j)) :: VS.Vector Int -- [0,1,2,3,4,5,1,2,3,4,5,6,2,3,4,5,6,7,3,4,5,6,7,8,4,5,6,7,8,9] -- -- While in this example `S`torable Array will first be converted into `U`nboxed @@ -146,7 +177,7 @@ -- `VU.Vector` in constant time. -- -- >>> import qualified Data.Vector.Unboxed as VU --- >>> toVector (makeArrayR S Par (5 :. 6) (\(i :. j) -> i + j)) :: VU.Vector Int +-- >>> 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 :: @@ -159,8 +190,8 @@ => Array r ix e -> v e toVector arr = - case castToVector (convert arr :: Array (ARepr v) ix e) of - Just v -> v - Nothing -> VG.generate (totalElem (size arr)) (unsafeLinearIndex arr) + fromMaybe + (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 @@ -1,20 +1,21 @@ -{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} -- | -- Module : Data.Massiv.Array.Mutable --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable -- module Data.Massiv.Array.Mutable - ( -- * Element-wise mutation - read + ( -- ** Size + msize + -- ** Element-wise mutation + , read , read' , write , write' @@ -23,114 +24,397 @@ , swap , swap' -- ** Operate over `MArray` - , Mutable - , MArray - , msize - -- *** Convert + -- *** Immutable conversion , new , thaw + , thawS , freeze - -- *** Create + , freezeS + -- *** Create mutable + , makeMArray + , makeMArrayLinear + , makeMArrayS + , makeMArrayLinearS + -- *** Create pure , createArray_ , createArray + , createArrayS_ + , createArrayS , createArrayST_ , createArrayST -- *** Generate , generateArray - , generateArrayIO + , generateArrayLinear + , generateArrayS + , generateArrayLinearS -- *** Unfold - , unfoldlPrim_ - , unfoldlPrim + , unfoldrPrimM_ + , iunfoldrPrimM_ + , unfoldrPrimM + , iunfoldrPrimM + , unfoldlPrimM_ + , iunfoldlPrimM_ + , unfoldlPrimM + , iunfoldlPrimM + -- *** Mapping + , forPrimM_ + , iforPrimM_ + , iforLinearPrimM_ -- *** Modify , withMArray , withMArrayST + -- *** Initialize + , initialize + , initializeNew -- ** Computation + , Mutable + , MArray , RealWorld , computeInto + , loadArray + , loadArrayS ) where -import Prelude hiding (mapM, read) - -import Control.Monad (unless) -import Control.Monad.Primitive (PrimMonad (..)) -import Control.Monad.ST -import Data.Massiv.Array.Manifest.Internal -import Data.Massiv.Array.Unsafe -import Data.Massiv.Core.Common -import Data.Massiv.Core.Scheduler +import Control.Monad (unless) +import Control.Monad.ST +import Control.Scheduler +import Data.Massiv.Core.Common +import Prelude hiding (mapM, read) --- | Initialize a new mutable array. Negative size will result in an empty array. -new :: (Mutable r ix e, PrimMonad m) => Sz ix -> m (MArray (PrimState m) r ix e) -new sz = unsafeNewZero (liftIndex (max 0) sz) +-- | /O(n)/ - Initialize a new mutable array. All elements will be set to some default value. For +-- boxed arrays in will be a thunk with `Uninitialized` exception, while for others it will be +-- simply zeros. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array +-- >>> marr <- new (Sz2 2 6) :: IO (MArray RealWorld P Ix2 Int) +-- >>> freeze Seq marr +-- Array P Seq (Sz (2 :. 6)) +-- [ [ 0, 0, 0, 0, 0, 0 ] +-- , [ 0, 0, 0, 0, 0, 0 ] +-- ] +-- +-- Or using @TypeApplications@: +-- +-- >>> :set -XTypeApplications +-- >>> new @P @Ix2 @Int (Sz2 2 6) >>= freezeS +-- Array P Seq (Sz (2 :. 6)) +-- [ [ 0, 0, 0, 0, 0, 0 ] +-- , [ 0, 0, 0, 0, 0, 0 ] +-- ] +-- >>> new @B @_ @Int (Sz2 2 6) >>= (`read'` 1) +-- *** Exception: Uninitialized +-- +-- @since 0.1.0 +new :: + forall r ix e m. (Mutable r ix e, PrimMonad m) + => Sz ix + -> m (MArray (PrimState m) r ix e) +new = initializeNew Nothing {-# INLINE new #-} --- | /O(n)/ - Yield a mutable copy of the immutable array -thaw :: (Mutable r ix e, PrimMonad m) => Array r ix e -> m (MArray (PrimState m) r ix e) -thaw = unsafeThaw . clone +-- | /O(n)/ - Make a mutable copy of a pure array. Keep in mind that both `freeze` and `thaw` trigger a +-- copy of the full array. +-- +-- ==== __Example__ +-- +-- >>> import Data.Massiv.Array +-- >>> :set -XTypeApplications +-- >>> arr <- fromListsM @U @Ix2 @Double Par [[12,21],[13,31]] +-- >>> marr <- thaw arr +-- >>> modify marr (+ 10) (1 :. 0) +-- True +-- >>> freeze Par marr +-- Array U Par (Sz (2 :. 2)) +-- [ [ 12.0, 21.0 ] +-- , [ 23.0, 31.0 ] +-- ] +-- +-- @since 0.1.0 +thaw :: forall r ix e m. (Mutable r ix e, MonadIO m) => Array r ix e -> m (MArray RealWorld r ix e) +thaw arr = liftIO $ makeMArrayLinear (getComp arr) (size arr) (pure . unsafeLinearIndexM arr) +-- TODO: use faster memcpy {-# INLINE thaw #-} --- | /O(n)/ - Yield an immutable copy of the mutable array -freeze :: (Mutable r ix e, PrimMonad m) => Comp -> MArray (PrimState m) r ix e -> m (Array r ix e) -freeze comp marr = clone <$> unsafeFreeze comp marr +-- | Same as `thaw`, but restrict computation to sequential only. +-- +-- ==== __Example__ +-- +-- >>> import Data.Massiv.Array +-- >>> :set -XOverloadedLists +-- >>> thawS @P @Ix1 @Double [1..10] +-- >>> marr <- thawS @P @Ix1 @Double [1..10] +-- >>> write' marr 5 100 +-- >>> freezeS marr +-- Array P Seq (Sz1 10) +-- [ 1.0, 2.0, 3.0, 4.0, 5.0, 100.0, 7.0, 8.0, 9.0, 10.0 ] +-- +-- @since 0.3.0 +thawS :: + forall r ix e m. (Mutable r ix e, PrimMonad m) + => Array r ix e + -> m (MArray (PrimState m) r ix e) +thawS arr = makeMArrayLinearS (size arr) (pure . unsafeLinearIndexM arr) +-- TODO: use faster memcpy +{-# INLINE thawS #-} + + +-- TODO: implement and benchmark parallel `thawIO` and `freezeIO` with memcpy + +-- | /O(n)/ - Yield an immutable copy of the mutable array. Note that mutable representations +-- have to be the same. +-- +-- ==== __Example__ +-- +-- >>> import Data.Massiv.Array +-- >>> marr <- new @P @_ @Int (Sz2 2 6) +-- >>> forM_ (range Seq 0 (Ix2 1 4)) $ \ix -> write marr ix 9 +-- >>> freeze Seq marr +-- Array P Seq (Sz (2 :. 6)) +-- [ [ 9, 9, 9, 9, 0, 0 ] +-- , [ 0, 0, 0, 0, 0, 0 ] +-- ] +-- +-- @since 0.1.0 +freeze :: + forall r ix e m. (Mutable r ix e, MonadIO m) + => Comp + -> MArray RealWorld r ix e + -> m (Array r ix e) +freeze comp marr = liftIO $ generateArrayLinear comp (msize marr) (unsafeLinearRead marr) {-# INLINE freeze #-} +-- | Same as `freeze`, but do the copy of supplied muable array sequentially. Also, unlike `freeze` +-- that has to be done in `IO`, `freezeS` can be used with `ST`. +-- +-- @since 0.3.0 +freezeS :: + forall r ix e m. (Mutable r ix e, PrimMonad m) + => MArray (PrimState m) r ix e + -> m (Array r ix e) +freezeS marr = generateArrayLinearS Seq (msize marr) (unsafeLinearRead marr) +{-# INLINE freezeS #-} + + +-- | Load sequentially a pure array into the newly created mutable array. +-- +-- @since 0.3.0 +loadArrayS :: + forall r ix e r' m. (Load r' ix e, Mutable r ix e, PrimMonad m) + => Array r' ix e + -> m (MArray (PrimState m) r ix e) +loadArrayS arr = do + marr <- unsafeNew (size arr) + loadArrayM (Scheduler 1 id) arr (unsafeLinearWrite marr) + pure marr +{-# INLINE loadArrayS #-} + + +-- | Load a pure array into the newly created mutable array, while respecting computation startegy. +-- +-- @since 0.3.0 +loadArray :: + forall r ix e r' m. (Load r' ix e, Mutable r ix e, MonadIO m) + => Array r' ix e + -> m (MArray RealWorld r ix e) +loadArray arr = + liftIO $ do + marr <- unsafeNew (size arr) + withScheduler_ (getComp arr) $ \scheduler -> loadArrayM scheduler arr (unsafeLinearWrite marr) + pure marr +{-# INLINE loadArray #-} + + + +-- | Compute an Array while loading the results into the supplied mutable target array. Number of +-- elements for arrays must agree, otherwise `SizeElementsMismatchException` exception is thrown. +-- +-- @since 0.1.3 +computeInto :: + (Load r' ix' e, Mutable r ix e, MonadIO m) + => MArray RealWorld r ix e -- ^ Target Array + -> Array r' ix' e -- ^ Array to load + -> m () +computeInto !mArr !arr = + liftIO $ do + unless (totalElem (msize mArr) == totalElem (size arr)) $ + throwM $ SizeElementsMismatchException (msize mArr) (size arr) + withScheduler_ (getComp arr) $ \scheduler -> loadArrayM scheduler arr (unsafeLinearWrite mArr) +{-# INLINE computeInto #-} + + +-- | Create a mutable array using an index aware generating action. +-- +-- @since 0.3.0 +makeMArrayS :: + forall r ix e m. (Mutable r ix e, PrimMonad m) + => Sz ix -- ^ Size of the create array + -> (ix -> m e) -- ^ Element generating action + -> m (MArray (PrimState m) r ix e) +makeMArrayS sz f = makeMArrayLinearS sz (f . fromLinearIndex sz) +{-# INLINE makeMArrayS #-} + + +-- | Same as `makeMArrayS`, but index supplied to the action is row-major linear index. +-- +-- @since 0.3.0 +makeMArrayLinearS :: + forall r ix e m. (Mutable r ix e, PrimMonad m) + => Sz ix + -> (Int -> m e) + -> m (MArray (PrimState m) r ix e) +makeMArrayLinearS sz f = do + marr <- unsafeNew sz + loopM_ 0 (< totalElem (msize marr)) (+ 1) (\ !i -> f i >>= unsafeLinearWrite marr i) + return marr +{-# INLINE makeMArrayLinearS #-} + +-- | Just like `makeMArrayS`, but also accepts computation strategy and runs in `IO`. +-- +-- @since 0.3.0 +makeMArray :: + forall r ix e m. (PrimMonad m, MonadUnliftIO m, Mutable r ix e) + => Comp + -> Sz ix + -> (ix -> m e) + -> m (MArray (PrimState m) r ix e) +makeMArray comp sz f = makeMArrayLinear comp sz (f . fromLinearIndex sz) +{-# INLINE makeMArray #-} + + +-- | Just like `makeMArrayLinearS`, but also accepts computation strategy and runs in `IO`. +-- +-- @since 0.3.0 +makeMArrayLinear :: + forall r ix e m. (PrimMonad m, MonadUnliftIO m, Mutable r ix e) + => Comp + -> Sz ix + -> (Int -> m e) + -> m (MArray (PrimState m) r ix e) +makeMArrayLinear comp sz f = do + marr <- unsafeNew sz + withScheduler_ comp $ \scheduler -> + splitLinearlyWithM_ scheduler (totalElem sz) f (unsafeLinearWrite marr) + return marr +{-# INLINE makeMArrayLinear #-} + + + + -- | Create a new array by supplying an action that will fill the new blank mutable array. Use -- `createArray` if you'd like to keep the result of the filling function. -- -- ====__Examples__ -- --- >>> createArray_ Seq (Sz1 2) (\ marr -> write marr 0 10 >> write marr 1 11) :: IO (Array P Ix1 Int) --- (Array P Seq (2) --- [ 10,11 ]) +-- >>> :set -XTypeApplications +-- >>> import Data.Massiv.Array +-- >>> createArray_ @P @_ @Int Seq (Sz1 2) (\ s marr -> scheduleWork s (write' marr 0 10) >> scheduleWork s (write' marr 1 11)) +-- Array P Seq (Sz1 2) +-- [ 10, 11 ] -- --- @since 0.2.6 +-- @since 0.3.0 -- createArray_ :: - (Mutable r ix e, PrimMonad m) + forall r ix e a m. (Mutable r ix e, PrimMonad m, MonadUnliftIO m) => Comp -- ^ Computation strategy to use after `MArray` gets frozen and onward. -> Sz ix -- ^ Size of the newly created array - -> (MArray (PrimState m) r ix e -> m a) + -> (Scheduler m () -> MArray (PrimState m) r ix e -> m a) -- ^ An action that should fill all elements of the brand new mutable array -> m (Array r ix e) -createArray_ comp sz action = fmap snd $ createArray comp sz action +createArray_ comp sz action = do + marr <- new sz + withScheduler_ comp (`action` marr) + unsafeFreeze comp marr {-# INLINE createArray_ #-} --- | Just like `createArray_`, but together with `Array` it returns the result of the filling action. +-- | Just like `createArray_`, but together with `Array` it returns results of scheduled filling +-- actions. -- --- @since 0.2.6 +-- @since 0.3.0 -- createArray :: - (Mutable r ix e, PrimMonad m) + forall r ix e a m. (Mutable r ix e, PrimMonad m, MonadUnliftIO m) => Comp -- ^ Computation strategy to use after `MArray` gets frozen and onward. -> Sz ix -- ^ Size of the newly created array + -> (Scheduler m a -> 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) +createArray comp sz action = do + marr <- new sz + a <- withScheduler comp (`action` marr) + arr <- unsafeFreeze comp marr + return (a, arr) +{-# INLINE createArray #-} + + +-- | Create a new array by supplying an action that will fill the new blank mutable array. Use +-- `createArrayS` if you'd like to keep the result of the filling function. +-- +-- ====__Examples__ +-- +-- >>> :set -XTypeApplications +-- >>> import Data.Massiv.Array +-- >>> createArrayS_ @P @_ @Int Seq (Sz1 2) (\ marr -> write marr 0 10 >> write marr 1 12) +-- Array P Seq (Sz1 2) +-- [ 10, 12 ] +-- +-- @since 0.3.0 +-- +createArrayS_ :: + forall r ix e a m. (Mutable r ix e, PrimMonad m) + => Comp -- ^ Computation strategy to use after `MArray` gets frozen and onward. + -> Sz ix -- ^ Size of the newly created array -> (MArray (PrimState m) r ix e -> m a) -- ^ An action that should fill all elements of the brand new mutable array + -> m (Array r ix e) +createArrayS_ comp sz action = snd <$> createArrayS comp sz action +{-# INLINE createArrayS_ #-} + +-- | Just like `createArray_`, but together with `Array` it returns the result of the filling action. +-- +-- @since 0.3.0 +-- +createArrayS :: + forall r ix e a m. (Mutable r ix e, PrimMonad m) + => Comp -- ^ Computation strategy to use after `MArray` gets frozen and onward. + -> Sz ix -- ^ Size of the newly created array + -> (MArray (PrimState m) r ix e -> m a) + -- ^ An action that should fill all elements of the brand new mutable array -> m (a, Array r ix e) -createArray comp sz action = do +createArrayS comp sz action = do marr <- new sz a <- action marr arr <- unsafeFreeze comp marr return (a, arr) -{-# INLINE createArray #-} +{-# INLINE createArrayS #-} --- | Just like `createArray_`, but restricted to `ST`. +-- | Just like `createArrayS_`, but restricted to `ST`. -- --- @since 0.2.6 +-- @since 0.3.0 -- createArrayST_ :: - Mutable r ix e => Comp -> Sz ix -> (forall s. MArray s r ix e -> ST s a) -> Array r ix e -createArrayST_ comp sz action = runST $ createArray_ comp sz action + forall r ix e a. Mutable r ix e + => Comp + -> Sz ix + -> (forall s. MArray s r ix e -> ST s a) + -> Array r ix e +createArrayST_ comp sz action = runST $ createArrayS_ comp sz action {-# INLINE createArrayST_ #-} --- | Just like `createArray`, but restricted to `ST`. +-- | Just like `createArrayS`, but restricted to `ST`. -- -- @since 0.2.6 -- createArrayST :: - Mutable r ix e => Comp -> Sz ix -> (forall s. MArray s r ix e -> ST s a) -> (a, Array r ix e) -createArrayST comp sz action = runST $ createArray comp sz action + forall r ix e a. Mutable r ix e + => Comp + -> Sz ix + -> (forall s. MArray s r ix e -> ST s a) + -> (a, Array r ix e) +createArrayST comp sz action = runST $ createArrayS comp sz action {-# INLINE createArrayST #-} @@ -144,6 +428,7 @@ -- -- ====__Examples__ -- +-- >>> import Data.Massiv.Array -- >>> import Data.IORef -- >>> ref <- newIORef (0 :: Int) -- >>> generateArray Seq (Sz1 6) (\ i -> modifyIORef' ref (+i) >> print i >> pure i) :: IO (Array U Ix1 Int) @@ -153,61 +438,71 @@ -- 3 -- 4 -- 5 --- (Array U Seq (6) --- [ 0,1,2,3,4,5 ]) +-- Array U Seq (Sz1 6) +-- [ 0, 1, 2, 3, 4, 5 ] -- >>> readIORef ref -- 15 -- -generateArray :: - (Mutable r ix e, PrimMonad m) +generateArrayS :: + forall r ix e m. (Mutable r ix e, PrimMonad m) => Comp -- ^ Computation strategy (ingored during generation) -> Sz ix -- ^ Resulting size of the array -> (ix -> m e) -- ^ Element producing generator -> m (Array r ix e) -generateArray comp sz' gen = do - let sz = liftIndex (max 0) sz' +generateArrayS comp sz gen = generateArrayLinearS comp sz (gen . fromLinearIndex sz) +{-# INLINE generateArrayS #-} + +-- | Same as `generateArray` but with action takes row-major linear index. +-- +-- @since 0.3.0 +generateArrayLinearS :: + forall r ix e m. (Mutable r ix e, PrimMonad m) + => Comp -- ^ Computation strategy (ingored during generation) + -> Sz ix -- ^ Resulting size of the array + -> (Int -> m e) -- ^ Element producing generator + -> m (Array r ix e) +generateArrayLinearS comp sz gen = do marr <- unsafeNew sz - iterM_ zeroIndex (msize marr) (pureIndex 1) (<) $ \ix -> gen ix >>= write marr ix + loopM_ 0 (< totalElem (msize marr)) (+ 1) $ \i -> gen i >>= unsafeLinearWrite marr i unsafeFreeze comp marr -{-# INLINE generateArray #-} +{-# INLINE generateArrayLinearS #-} --- | Just like `generateArray`, except this generator __will__ respect the supplied computation +-- | 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 -generateArrayIO :: - (Mutable r ix e) +generateArray :: + forall r ix e m. (MonadUnliftIO m, PrimMonad m, Mutable r ix e) => Comp -> Sz ix - -> (ix -> IO e) - -> IO (Array r ix e) -generateArrayIO comp sz' gen = do - case comp of - Seq -> generateArray comp sz' gen - ParOn wids -> do - let sz = liftIndex (max 0) sz' - marr <- unsafeNew sz - withScheduler_ wids $ \scheduler -> - splitLinearlyWithM_ - (numWorkers scheduler) - (scheduleWork scheduler) - (totalElem sz) - (gen . fromLinearIndex sz) - (unsafeLinearWrite marr) - unsafeFreeze comp marr -{-# INLINE generateArrayIO #-} + -> (ix -> m e) + -> m (Array r ix e) +generateArray comp sz f = generateArrayLinear comp sz (f . fromLinearIndex sz) +{-# INLINE generateArray #-} --- | Sequentially unfold an array from the left. +-- | Just like `generateArrayIO`, but action supplied will receive a row-major linear index. -- --- @since 0.2.6 +-- @since 0.3.0 +generateArrayLinear :: + forall r ix e m. (MonadUnliftIO m, PrimMonad m, Mutable r ix e) + => Comp + -> Sz ix + -> (Int -> m e) + -> m (Array r ix e) +generateArrayLinear comp sz f = makeMArrayLinear comp sz f >>= unsafeFreeze comp +{-# INLINE generateArrayLinear #-} + + +-- | Sequentially unfold an array from the left. -- -- ====__Examples__ -- -- Create an array with Fibonacci numbers while performing and `IO` action on the accumulator for -- each element of the array. -- --- >>> unfoldlPrim_ Seq (Ix1 10) (\a@(f0, f1) _ -> let fn = f0 + f1 in print a >> return ((f1, fn), f0)) (0, 1) :: IO (Array P Ix1 Int) +-- >>> import Data.Massiv.Array +-- >>> unfoldrPrimM_ Seq (Sz1 10) (\a@(f0, f1) -> let fn = f0 + f1 in print a >> return (f0, (f1, fn))) (0, 1) :: IO (Array P Ix1 Int) -- (0,1) -- (1,1) -- (1,2) @@ -218,57 +513,234 @@ -- (13,21) -- (21,34) -- (34,55) --- (Array P Seq (10) --- [ 0,1,1,2,3,5,8,13,21,34 ]) +-- Array P Seq (Sz1 10) +-- [ 0, 1, 1, 2, 3, 5, 8, 13, 21, 34 ] -- -unfoldlPrim_ :: - (Mutable r ix e, PrimMonad m) +-- @since 0.3.0 +-- +unfoldrPrimM_ :: + forall r ix e a m. (Mutable r ix e, PrimMonad m) => Comp -- ^ Computation strategy (ignored during initial creation) -> Sz ix -- ^ Size of the desired array - -> (a -> ix -> m (a, e)) -- ^ Unfolding action + -> (a -> m (e, a)) -- ^ Unfolding action -> a -- ^ Initial accumulator -> m (Array r ix e) -unfoldlPrim_ comp sz gen acc0 = fmap snd $ unfoldlPrim comp sz gen acc0 -{-# INLINE unfoldlPrim_ #-} +unfoldrPrimM_ comp sz gen acc0 = snd <$> unfoldrPrimM comp sz gen acc0 +{-# INLINE unfoldrPrimM_ #-} +-- | Same as `unfoldrPrim_` but do the unfolding with index aware function. +-- +-- @since 0.3.0 +-- +iunfoldrPrimM_ :: + forall r ix e a m. (Mutable r ix e, PrimMonad m) + => Comp -- ^ Computation strategy (ignored during initial creation) + -> Sz ix -- ^ Size of the desired array + -> (a -> ix -> m (e, a)) -- ^ Unfolding action + -> a -- ^ Initial accumulator + -> m (Array r ix e) +iunfoldrPrimM_ comp sz gen acc0 = snd <$> iunfoldrPrimM comp sz gen acc0 +{-# INLINE iunfoldrPrimM_ #-} --- | Just like `unfoldlPrim_`, but also returns the final value of the accumulator. + +-- | Just like `iunfoldrPrim_`, but also returns the final value of the accumulator. -- --- @since 0.2.6 +-- @since 0.3.0 +iunfoldrPrimM :: + forall r ix e a m. (Mutable r ix e, PrimMonad m) + => Comp -- ^ Computation strategy (ignored during initial creation) + -> Sz ix -- ^ Size of the desired array + -> (a -> ix -> m (e, a)) -- ^ Unfolding action + -> a -- ^ Initial accumulator + -> m (a, Array r ix e) +iunfoldrPrimM comp sz gen acc0 = + createArrayS comp sz $ \marr -> + let sz' = msize marr + in iterLinearM sz' 0 (totalElem sz') 1 (<) acc0 $ \i ix acc -> do + (e, acc') <- gen acc ix + unsafeLinearWrite marr i e + return acc' +{-# INLINE iunfoldrPrimM #-} + +-- | Just like `iunfoldrPrimM`, but do the unfolding with index aware function. -- -unfoldlPrim :: - (Mutable r ix e, PrimMonad m) +-- @since 0.3.0 +unfoldrPrimM :: + forall r ix e a m. (Mutable r ix e, PrimMonad m) => Comp -- ^ Computation strategy (ignored during initial creation) -> Sz ix -- ^ Size of the desired array + -> (a -> m (e, a)) -- ^ Unfolding action + -> a -- ^ Initial accumulator + -> m (a, Array r ix e) +unfoldrPrimM comp sz gen acc0 = + createArrayS comp sz $ \marr -> + let sz' = msize marr + in loopM 0 (< totalElem sz') (+1) acc0 $ \i acc -> do + (e, acc') <- gen acc + unsafeLinearWrite marr i e + return acc' +{-# INLINE unfoldrPrimM #-} + +-- | Sequentially unfold an array from the left. +-- +-- ====__Examples__ +-- +-- Create an array with Fibonacci numbers starting at the end while performing and `IO` action on +-- the accumulator for each element of the array. +-- +-- >>> import Data.Massiv.Array +-- >>> unfoldlPrimM_ Seq (Sz1 10) (\a@(f0, f1) -> let fn = f0 + f1 in print a >> return ((f1, fn), f0)) (0, 1) :: IO (Array P Ix1 Int) +-- (0,1) +-- (1,1) +-- (1,2) +-- (2,3) +-- (3,5) +-- (5,8) +-- (8,13) +-- (13,21) +-- (21,34) +-- (34,55) +-- Array P Seq (Sz1 10) +-- [ 34, 21, 13, 8, 5, 3, 2, 1, 1, 0 ] +-- +-- @since 0.3.0 +-- +unfoldlPrimM_ :: + forall r ix e a m. (Mutable r ix e, PrimMonad m) + => Comp -- ^ Computation strategy (ignored during initial creation) + -> Sz ix -- ^ Size of the desired array + -> (a -> m (a, e)) -- ^ Unfolding action + -> a -- ^ Initial accumulator + -> m (Array r ix e) +unfoldlPrimM_ comp sz gen acc0 = snd <$> unfoldlPrimM comp sz gen acc0 +{-# INLINE unfoldlPrimM_ #-} + +-- | Same as `unfoldlPrimM_` but do the unfolding with index aware function. +-- +-- @since 0.3.0 +-- +iunfoldlPrimM_ :: + forall r ix e a m. (Mutable r ix e, PrimMonad m) + => Comp -- ^ Computation strategy (ignored during initial creation) + -> Sz ix -- ^ Size of the desired array -> (a -> ix -> m (a, e)) -- ^ Unfolding action -> a -- ^ Initial accumulator + -> m (Array r ix e) +iunfoldlPrimM_ comp sz gen acc0 = snd <$> iunfoldlPrimM comp sz gen acc0 +{-# INLINE iunfoldlPrimM_ #-} + + +-- | Just like `iunfoldlPrimM_`, but also returns the final value of the accumulator. +-- +-- @since 0.3.0 +iunfoldlPrimM :: + forall r ix e a m. (Mutable r ix e, PrimMonad m) + => Comp -- ^ Computation strategy (ignored during initial creation) + -> Sz ix -- ^ Size of the desired array + -> (a -> ix -> m (a, e)) -- ^ Unfolding action + -> a -- ^ Initial accumulator -> m (a, Array r ix e) -unfoldlPrim comp sz gen acc0 = - createArray comp sz $ \marr -> +iunfoldlPrimM comp sz gen acc0 = + createArrayS comp sz $ \marr -> let sz' = msize marr - in iterLinearM sz' 0 (totalElem sz') 1 (<) acc0 $ \i ix acc -> do + in iterLinearM sz' (totalElem sz' - 1) 0 (negate 1) (>=) acc0 $ \i ix acc -> do (acc', e) <- gen acc ix unsafeLinearWrite marr i e return acc' -{-# INLINE unfoldlPrim #-} +{-# INLINE iunfoldlPrimM #-} +-- | Just like `iunfoldlPrimM`, but do the unfolding with index aware function. +-- +-- @since 0.3.0 +unfoldlPrimM :: + forall r ix e a m. (Mutable r ix e, PrimMonad m) + => Comp -- ^ Computation strategy (ignored during initial creation) + -> Sz ix -- ^ Size of the desired array + -> (a -> m (a, e)) -- ^ Unfolding action + -> a -- ^ Initial accumulator + -> m (a, Array r ix e) +unfoldlPrimM comp sz gen acc0 = + createArrayS comp sz $ \marr -> + let sz' = msize marr + in loopDeepM 0 (< totalElem sz') (+1) acc0 $ \i acc -> do + (acc', e) <- gen acc + unsafeLinearWrite marr i e + return acc' +{-# INLINE unfoldlPrimM #-} --- | Create a copy of a pure array, mutate it in place and return its frozen version. + +-- | Sequentially loop over a mutable array while modifying each element with an action. -- --- @since 0.2.2 +-- @since 0.3.0 +forPrimM_ :: (Mutable r ix e, PrimMonad m) => MArray (PrimState m) r ix e -> (e -> m e) -> m () +forPrimM_ marr f = + loopM_ 0 (< totalElem (msize marr)) (+1) (unsafeLinearModify marr (const f)) +{-# INLINE forPrimM_ #-} + + +-- | Sequentially loop over a mutable array while modifying each element with an index aware action. +-- +-- @since 0.3.0 +iforPrimM_ :: + (Mutable r ix e, PrimMonad m) => MArray (PrimState m) r ix e -> (ix -> e -> m e) -> m () +iforPrimM_ marr f = iforLinearPrimM_ marr (f . fromLinearIndex (msize marr)) +{-# INLINE iforPrimM_ #-} + + +-- | Sequentially loop over a mutable array while modifying each element with an index aware action. +-- +-- @since 0.3.0 +iforLinearPrimM_ :: + (Mutable r ix e, PrimMonad m) => MArray (PrimState m) r ix e -> (Int -> e -> m e) -> m () +iforLinearPrimM_ marr f = loopM_ 0 (< totalElem (msize marr)) (+ 1) (unsafeLinearModify marr f) +{-# INLINE iforLinearPrimM_ #-} + + +-- | Create a copy of a pure array, mutate it in place and return its frozen version. The big +-- difference between `withMArrayS` is that it's not only gonna respect the computation strategy +-- supplied to it, but it will also pass extra argumens to the action that suppose to modify the +-- mutable copy of the source array. These two extra arguments are: +-- +-- * Number of capabilities derived from the `Comp`utation strategy of the array. +-- +-- * An action that can be used to schedule arbitrary number of jobs that will be executed in +-- parallel. +-- +-- * And, of course, the mutable array itself. +-- +-- @since 0.3.0 withMArray :: - (Mutable r ix e, PrimMonad m) + (Mutable r ix e, MonadUnliftIO m) => Array r ix e - -> (MArray (PrimState m) r ix e -> m a) + -> (Int -> (m () -> m ()) -> MArray RealWorld r ix e -> m a) -> m (Array r ix e) withMArray arr action = do marr <- thaw arr + withScheduler_ (getComp arr) $ \scheduler -> + action (numWorkers scheduler) (scheduleWork scheduler) marr + liftIO $ unsafeFreeze (getComp arr) marr +{-# INLINE withMArray #-} + + +-- | Create a copy of a pure array, mutate it in place and return its frozen version. The important +-- benefit over doing a manual `thawS` followed by a `freezeS` is that an array will be only copied +-- once. +-- +-- @since 0.3.0 +withMArrayS :: + (Mutable r ix e, PrimMonad m) + => Array r ix e + -> (MArray (PrimState m) r ix e -> m a) + -> m (Array r ix e) +withMArrayS arr action = do + marr <- thawS arr _ <- action marr unsafeFreeze (getComp arr) marr -{-# INLINE withMArray #-} +{-# INLINE withMArrayS #-} --- | Same as `withMArray` but in `ST`. +-- | Same as `withMArrayS` but in `ST`. This is not only pure, but also the safest way to do +-- mutation to the array. -- -- @since 0.2.2 withMArrayST :: @@ -276,34 +748,34 @@ => Array r ix e -> (forall s . MArray s r ix e -> ST s a) -> Array r ix e -withMArrayST arr f = runST $ withMArray arr f +withMArrayST arr f = runST $ withMArrayS arr f {-# INLINE withMArrayST #-} --- | /O(1)/ - Lookup an element in the mutable array. Return `Nothing` when index is out of bounds. + +-- | /O(1)/ - Lookup an element in the mutable array. Returns `Nothing` when index is out of bounds. read :: (Mutable r ix e, PrimMonad m) => MArray (PrimState m) r ix e -> ix -> m (Maybe e) read marr ix = if isSafeIndex (msize marr) ix - then Just <$> unsafeRead marr ix - else return Nothing + then Just <$> unsafeRead marr ix + else return Nothing {-# INLINE read #-} --- | /O(1)/ - Same as `read`, but throws an error if index is out of bounds. -read' :: (Mutable r ix e, PrimMonad m) => - MArray (PrimState m) r ix e -> ix -> m e +-- | /O(1)/ - Same as `read`, but lives in IO and throws `IndexOutOfBoundsException` on invalid +-- index. +read' :: (Mutable r ix e, MonadThrow m, PrimMonad m) => MArray (PrimState m) r ix e -> ix -> m e read' marr ix = do mval <- read marr ix case mval of - Just e -> return e - Nothing -> errorIx "Data.Massiv.Array.Mutable.read'" (msize marr) ix + Just e -> pure e + Nothing -> throwM $ IndexOutOfBoundsException (msize marr) ix {-# INLINE read' #-} -- | /O(1)/ - Write an element into the cell of a mutable array. Returns `False` when index is out -- of bounds. -write :: (Mutable r ix e, PrimMonad m) => - MArray (PrimState m) r ix e -> ix -> e -> m Bool +write :: (Mutable r ix e, PrimMonad m) => MArray (PrimState m) r ix e -> ix -> e -> m Bool write marr ix e = if isSafeIndex (msize marr) ix then unsafeWrite marr ix e >> return True @@ -311,18 +783,17 @@ {-# INLINE write #-} --- | /O(1)/ - Same as `write`, but throws an error if index is out of bounds. -write' :: (Mutable r ix e, PrimMonad m) => - MArray (PrimState m) r ix e -> ix -> e -> m () -write' marr ix e = - write marr ix e >>= (`unless` errorIx "Data.Massiv.Array.Mutable.write'" (msize marr) ix) +-- | /O(1)/ - Same as `write`, but lives in IO and throws `IndexOutOfBoundsException` on invalid +-- index. +write' :: + (Mutable r ix e, MonadThrow m, PrimMonad m) => MArray (PrimState m) r ix e -> ix -> e -> m () +write' marr ix e = write marr ix e >>= (`unless` throwM (IndexOutOfBoundsException (msize marr) ix)) {-# INLINE write' #-} -- | /O(1)/ - Modify an element in the cell of a mutable array with a supplied function. Returns -- `False` when index is out of bounds. -modify :: (Mutable r ix e, PrimMonad m) => - MArray (PrimState m) r ix e -> (e -> e) -> ix -> m Bool +modify :: (Mutable r ix e, PrimMonad m) => MArray (PrimState m) r ix e -> (e -> e) -> ix -> m Bool modify marr f ix = if isSafeIndex (msize marr) ix then do @@ -334,39 +805,35 @@ -- | /O(1)/ - Same as `modify`, but throws an error if index is out of bounds. -modify' :: (Mutable r ix e, PrimMonad m) => +modify' :: (Mutable r ix e, MonadThrow m, PrimMonad m) => MArray (PrimState m) r ix e -> (e -> e) -> ix -> m () modify' marr f ix = - modify marr f ix >>= (`unless` errorIx "Data.Massiv.Array.Mutable.modify'" (msize marr) ix) + modify marr f ix >>= (`unless` throwM (IndexOutOfBoundsException (msize marr) ix)) {-# INLINE modify' #-} -- | /O(1)/ - Swap two elements in a mutable array by supplying their indices. Returns `False` when -- either one of the indices is out of bounds. -swap :: (Mutable r ix e, PrimMonad m) => - MArray (PrimState m) r ix e -> ix -> ix -> m Bool +swap :: (Mutable r ix e, PrimMonad m) => MArray (PrimState m) r ix e -> ix -> ix -> m Bool swap marr ix1 ix2 = do let sz = msize marr if isSafeIndex sz ix1 && isSafeIndex sz ix2 then do - val1 <- unsafeRead marr ix1 - val2 <- unsafeRead marr ix2 - unsafeWrite marr ix1 val2 - unsafeWrite marr ix2 val1 + val <- unsafeRead marr ix1 + unsafeRead marr ix2 >>= unsafeWrite marr ix1 + unsafeWrite marr ix2 val return True else return False {-# INLINE swap #-} --- | /O(1)/ - Same as `swap`, but throws an error if index is out of bounds. -swap' :: (Mutable r ix e, PrimMonad m) => - MArray (PrimState m) r ix e -> ix -> ix -> m () -swap' marr ix1 ix2 = do - success <- swap marr ix1 ix2 - unless success $ - errorIx "Data.Massiv.Array.Mutable.swap'" (msize marr) $ - if isSafeIndex (msize marr) ix1 - then ix2 - else ix1 +-- | /O(1)/ - Same as `swap`, but throws an `IndexOutOfBoundsException` on invalid indices. +swap' :: + (Mutable r ix e, MonadThrow m, PrimMonad m) => MArray (PrimState m) r ix e -> ix -> ix -> m () +swap' marr ix1 ix2 = + swap marr ix1 ix2 >>= + (`unless` if isSafeIndex (msize marr) ix1 + then throwM $ IndexOutOfBoundsException (msize marr) ix2 + else throwM $ IndexOutOfBoundsException (msize marr) ix1) {-# INLINE swap' #-} diff --git a/src/Data/Massiv/Array/Mutable/Atomic.hs b/src/Data/Massiv/Array/Mutable/Atomic.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Array/Mutable/Atomic.hs @@ -0,0 +1,148 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +-- | +-- Module : Data.Massiv.Array.Mutable.Atomic +-- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Array.Mutable.Atomic + ( -- * Atomic element-wise mutation + atomicReadIntArray + , atomicWriteIntArray + , atomicModifyIntArray + , atomicAddIntArray + , atomicSubIntArray + , atomicAndIntArray + , atomicNandIntArray + , atomicOrIntArray + , atomicXorIntArray + , casIntArray + ) where + +import Control.Monad.Primitive +import Data.Massiv.Array.Manifest.Primitive +import Data.Massiv.Core.Common + +-- Atomic operations + +-- | 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 marr ix + | isSafeIndex (msize 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 marr ix f + | isSafeIndex (msize 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 + -> m (Maybe Int) +casIntArray marr ix e n + | isSafeIndex (msize 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. +-- +-- @since 0.3.0 +atomicModifyIntArray :: + (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> (Int -> Int) -> m (Maybe Int) +atomicModifyIntArray marr ix f + | isSafeIndex (msize marr) ix = Just <$> unsafeAtomicModifyIntArray marr ix f + | 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 marr ix e + | isSafeIndex (msize 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 marr ix e + | isSafeIndex (msize 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 marr ix e + | isSafeIndex (msize 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 marr ix e + | isSafeIndex (msize 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 marr ix e + | isSafeIndex (msize 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 marr ix e + | isSafeIndex (msize marr) ix = Just <$> unsafeAtomicXorIntArray marr ix e + | otherwise = pure Nothing +{-# INLINE atomicXorIntArray #-} 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 @@ -1,10 +1,11 @@ -{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} -- | -- Module : Data.Massiv.Array.Numeric --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental @@ -62,37 +63,53 @@ , atan2A ) where -import Data.Massiv.Array.Delayed.Internal -import Data.Massiv.Array.Manifest.Internal -import Data.Massiv.Array.Ops.Fold as A -import Data.Massiv.Array.Ops.Map as A -import Data.Massiv.Array.Ops.Transform as A -import Data.Massiv.Core -import Data.Massiv.Core.Common -import Data.Monoid ((<>)) -import Prelude as P +import Data.Massiv.Array.Delayed.Pull +import Data.Massiv.Array.Manifest.Internal +import Data.Massiv.Array.Ops.Fold as A +import Data.Massiv.Array.Ops.Map as A +import Data.Massiv.Array.Ops.Transform as A +import Data.Massiv.Core +import Data.Massiv.Core.Common +import Data.Massiv.Core.Index.Internal (Sz(SafeSz)) +import Prelude as P infixr 8 .^, .^^ infixl 7 .*, ./, `quotA`, `remA`, `divA`, `modA` infixl 6 .+, .- +liftArray2Matching + :: (Source r1 ix a, Source r2 ix b) + => (a -> b -> e) -> Array r1 ix a -> Array r2 ix b -> Array D ix e +liftArray2Matching f !arr1 !arr2 + | sz1 == sz2 = + makeArray + (getComp arr1 <> getComp arr2) + sz1 + (\ !ix -> f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)) + | otherwise = throw $ SizeMismatchException (size arr1) (size arr2) + where + sz1 = size arr1 + sz2 = size arr2 +{-# INLINE liftArray2Matching #-} + + (.+) :: (Source r1 ix e, Source r2 ix e, Num e) => Array r1 ix e -> Array r2 ix e -> Array D ix e -(.+) = liftArray2 (+) +(.+) = liftArray2Matching (+) {-# INLINE (.+) #-} (.-) :: (Source r1 ix e, Source r2 ix e, Num e) => Array r1 ix e -> Array r2 ix e -> Array D ix e -(.-) = liftArray2 (-) +(.-) = liftArray2Matching (-) {-# INLINE (.-) #-} (.*) :: (Source r1 ix e, Source r2 ix e, Num e) => Array r1 ix e -> Array r2 ix e -> Array D ix e -(.*) = liftArray2 (*) +(.*) = liftArray2Matching (*) {-# INLINE (.*) #-} (.^) @@ -101,7 +118,7 @@ (.^) arr n = liftArray (^ n) arr {-# INLINE (.^) #-} --- | Perform matrix multiplication. Inner dimensions must agree, otherwise error. +-- | Perform matrix multiplication. Inner dimensions must agree, otherwise `SizeMismatchException`. (|*|) :: (Mutable r Ix2 e, Source r' Ix2 e, OuterSlice r Ix2 e, Source (EltRepr r Ix2) Ix1 e, Num e) => Array r Ix2 e @@ -112,7 +129,7 @@ {-# RULES "multDoubleTranspose" [~1] forall arr1 arr2 . arr1 |*| transpose arr2 = - multiplyTransposedFused arr1 (computeSource arr2) + multiplyTransposedFused arr1 (convert arr2) #-} multiplyTransposedFused :: @@ -154,16 +171,13 @@ -> Array r Ix2 e -> Array D Ix2 e multiplyTransposed arr1 arr2 - | n1 /= m2 = - error $ - "(|*|): Inner array dimensions must agree, but received: " ++ - show (size arr1) ++ " and " ++ show (size arr2) + | n1 /= m2 = throw $ SizeMismatchException (size arr1) (size arr2) | otherwise = - DArray (getComp arr1 <> getComp arr2) (m1 :. n2) $ \(i :. j) -> + DArray (getComp arr1 <> getComp arr2) (SafeSz (m1 :. n2)) $ \(i :. j) -> A.foldlS (+) 0 (A.zipWith (*) (unsafeOuterSlice arr1 i) (unsafeOuterSlice arr2 j)) where - (m1 :. n1) = size arr1 - (n2 :. m2) = size arr2 + SafeSz (m1 :. n1) = size arr1 + SafeSz (n2 :. m2) = size arr2 {-# INLINE multiplyTransposed #-} @@ -188,13 +202,13 @@ fromIntegerA :: (Index ix, Num e) => Integer -> Array D ix e -fromIntegerA = singleton Seq . fromInteger +fromIntegerA = singleton . fromInteger {-# INLINE fromIntegerA #-} (./) :: (Source r1 ix e, Source r2 ix e, Fractional e) => Array r1 ix e -> Array r2 ix e -> Array D ix e -(./) = liftArray2 (/) +(./) = liftArray2Matching (/) {-# INLINE (./) #-} (.^^) @@ -213,13 +227,13 @@ fromRationalA :: (Index ix, Fractional e) => Rational -> Array D ix e -fromRationalA = singleton Seq . fromRational +fromRationalA = singleton . fromRational {-# INLINE fromRationalA #-} piA :: (Index ix, Floating e) => Array D ix e -piA = singleton Seq pi +piA = singleton pi {-# INLINE piA #-} expA @@ -231,7 +245,7 @@ sqrtA :: (Source r ix e, Floating e) => Array r ix e -> Array D ix e -sqrtA = liftArray sqrt +sqrtA = liftArray exp {-# INLINE sqrtA #-} logA @@ -243,13 +257,13 @@ logBaseA :: (Source r1 ix e, Source r2 ix e, Floating e) => Array r1 ix e -> Array r2 ix e -> Array D ix e -logBaseA = liftArray2 logBase +logBaseA = liftArray2Matching logBase {-# INLINE logBaseA #-} (.**) :: (Source r1 ix e, Source r2 ix e, Floating e) => Array r1 ix e -> Array r2 ix e -> Array D ix e -(.**) = liftArray2 (**) +(.**) = liftArray2Matching (**) {-# INLINE (.**) #-} @@ -330,26 +344,26 @@ quotA :: (Source r1 ix e, Source r2 ix e, Integral e) => Array r1 ix e -> Array r2 ix e -> Array D ix e -quotA = liftArray2 (quot) +quotA = liftArray2Matching (quot) {-# INLINE quotA #-} remA :: (Source r1 ix e, Source r2 ix e, Integral e) => Array r1 ix e -> Array r2 ix e -> Array D ix e -remA = liftArray2 (rem) +remA = liftArray2Matching (rem) {-# INLINE remA #-} divA :: (Source r1 ix e, Source r2 ix e, Integral e) => Array r1 ix e -> Array r2 ix e -> Array D ix e -divA = liftArray2 (div) +divA = liftArray2Matching (div) {-# INLINE divA #-} modA :: (Source r1 ix e, Source r2 ix e, Integral e) => Array r1 ix e -> Array r2 ix e -> Array D ix e -modA = liftArray2 (mod) +modA = liftArray2Matching (mod) {-# INLINE modA #-} @@ -357,14 +371,14 @@ quotRemA :: (Source r1 ix e, Source r2 ix e, Integral e) => Array r1 ix e -> Array r2 ix e -> (Array D ix e, Array D ix e) -quotRemA arr1 = A.unzip . liftArray2 (quotRem) arr1 +quotRemA arr1 = A.unzip . liftArray2Matching (quotRem) arr1 {-# INLINE quotRemA #-} divModA :: (Source r1 ix e, Source r2 ix e, Integral e) => Array r1 ix e -> Array r2 ix e -> (Array D ix e, Array D ix e) -divModA arr1 = A.unzip . liftArray2 (divMod) arr1 +divModA arr1 = A.unzip . liftArray2Matching (divMod) arr1 {-# INLINE divModA #-} @@ -399,6 +413,6 @@ atan2A :: (Source r ix e, RealFloat e) => Array r ix e -> Array r ix e -> Array D ix e -atan2A = liftArray2 atan2 +atan2A = liftArray2Matching atan2 {-# INLINE atan2A #-} 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,13 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE FlexibleContexts #-} +-- | +-- Module : Data.Massiv.Array.Numeric.Integral +-- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- module Data.Massiv.Array.Numeric.Integral ( -- $integral_intro @@ -22,11 +30,14 @@ , fromFunction -- ** Sampled at the midpoint , fromFunctionMidpoint + -- * Helper functions ) where import Data.Coerce -import Data.Massiv.Array.Delayed +import Data.Massiv.Array.Delayed.Pull (D) +import Data.Massiv.Array.Delayed.Windowed (DW) import Data.Massiv.Array.Manifest.Internal +import Data.Massiv.Array.Ops.Construct (rangeInclusive) import Data.Massiv.Array.Ops.Transform (extract') import Data.Massiv.Array.Stencil import Data.Massiv.Core.Common @@ -46,7 +57,7 @@ -> Int -- ^ @n@ - number of sample points. -> Stencil ix e e midpointStencil dx dim k = - makeStencilDef 0 (setDim' (pureIndex 1) dim k) zeroIndex $ \g -> + makeStencilDef 0 (Sz (setDim' (pureIndex 1) dim k)) zeroIndex $ \g -> pure dx * loop 0 (< k) (+ 1) 0 (\i -> (+ g (setDim' zeroIndex dim i))) {-# INLINE midpointStencil #-} @@ -65,7 +76,7 @@ -> Int -- ^ @n@ - number of sample points. -> Stencil ix e e trapezoidStencil dx dim n = - makeStencilDef 0 (setDim' (pureIndex 1) dim (n + 1)) zeroIndex $ \g -> + makeStencilDef 0 (Sz (setDim' (pureIndex 1) dim (n + 1))) zeroIndex $ \g -> pure dx / 2 * (loop 1 (< n) (+ 1) (g zeroIndex) (\i -> (+ 2 * g (setDim' zeroIndex dim i))) + g (setDim' zeroIndex dim n)) @@ -90,7 +101,7 @@ error $ "Number of sample points for Simpson's rule stencil should be even, but received: " ++ show n | otherwise = - makeStencilDef 0 (setDim' (pureIndex 1) dim (n + 1)) zeroIndex $ \g -> + makeStencilDef 0 (Sz (setDim' (pureIndex 1) dim (n + 1))) zeroIndex $ \g -> let simAcc i (prev, acc) = let !fx3 = g (setDim' zeroIndex dim (i + 2)) !newAcc = acc + prev + 4 * g (setDim' zeroIndex dim (i + 1)) + fx3 @@ -100,7 +111,7 @@ -- | Integrate with a stencil along a particular dimension. integrateWith :: - (Fractional e, Load DW ix e, Mutable r ix e) + (Fractional e, StrideLoad DW ix e, Mutable r ix e) => (Dim -> Int -> Stencil ix e e) -> Dim -- ^ Dimension along which integration should be estimated. -> Int -- ^ @n@ - Number of samples @@ -115,7 +126,7 @@ -- | Compute an approximation of integral using a supplied rule in a form of `Stencil`. integralApprox :: - (Fractional e, Load DW ix e, Mutable r ix e) + (Fractional e, StrideLoad DW ix e, Mutable r ix e) => (e -> Dim -> Int -> Stencil ix e e) -- ^ Integration Stencil -> e -- ^ @d@ - Length of interval per cell -> Sz ix -- ^ @sz@ - Result size of the matrix @@ -133,7 +144,7 @@ -- | Use midpoint rule to approximate an integral. midpointRule :: - (Fractional e, Load DW ix e, Mutable r ix e) + (Fractional e, StrideLoad DW ix e, Mutable r ix e) => Comp -- ^ Computation strategy. -> r -- ^ Intermediate array representation. -> ((Int -> e) -> ix -> e) -- ^ @f(x,y,...)@ - Function to integrate @@ -149,14 +160,14 @@ -- | Use trapezoid rule to approximate an integral. trapezoidRule :: - (Fractional e, Load DW ix e, Mutable r ix e) + (Fractional e, StrideLoad DW ix e, Mutable r ix e) => Comp -- ^ Computation strategy -> r -- ^ Intermediate array representation -> ((Int -> e) -> ix -> e) -- ^ @f(x,y,...)@ - function to integrate - -> e -- ^ @a@ - starting point - -> e -- ^ @d@ - distance per matrix cell - -> Sz ix -- ^ @sz@ - end matrix size - -> Int -- ^ @n@ - number of sample points per cell in each direction + -> e -- ^ @a@ - Starting value point. + -> e -- ^ @d@ - Distance per matrix cell. + -> Sz ix -- ^ @sz@ - Result matrix size. + -> Int -- ^ @n@ - Number of sample points per cell in each direction. -> Array M ix e trapezoidRule comp r f a d sz n = integralApprox trapezoidStencil d sz n $ computeAs r $ fromFunction comp f a d sz n @@ -164,15 +175,15 @@ -- | Use Simpson's rule to approximate an integral. simpsonsRule :: - (Fractional e, Load DW ix e, Mutable r ix e) + (Fractional e, StrideLoad DW ix e, Mutable r ix e) => Comp -- ^ Computation strategy -> r -- ^ Intermediate array representation - -> ((Int -> e) -> ix -> e) -- ^ @f(x,y,...)@ - function to integrate - -> e -- ^ @a@ - starting point - -> e -- ^ @d@ - distance per matrix cell - -> Sz ix -- ^ @sz@ - end matrix size - -> Int -- ^ @n@ - number of sample points per cell in each direction. This value must be even, - -- otherwise error.. + -> ((Int -> e) -> ix -> e) -- ^ @f(x,y,...)@ - Function to integrate + -> e -- ^ @a@ - Starting value point. + -> e -- ^ @d@ - Distance per matrix cell. + -> Sz ix -- ^ @sz@ - Result matrix size. + -> Int -- ^ @n@ - Number of sample points per cell in each direction. This value must be even, + -- otherwise error. -> Array M ix e simpsonsRule comp r f a d sz n = integralApprox simpsonsStencil d sz n $ computeAs r $ fromFunction comp f a d sz n @@ -181,18 +192,18 @@ -- | Create an array from a function with sample points at the edges -- --- >>> fromFunction Seq (\ scale (i :. j) -> scale i + scale j) (-2) 1 (4 :. 4) 2 --- (Array D Seq (9 :. 9) --- [ [ -4.0,-3.5,-3.0,-2.5,-2.0,-1.5,-1.0,-0.5,0.0 ] --- , [ -3.5,-3.0,-2.5,-2.0,-1.5,-1.0,-0.5,0.0,0.5 ] --- , [ -3.0,-2.5,-2.0,-1.5,-1.0,-0.5,0.0,0.5,1.0 ] --- , [ -2.5,-2.0,-1.5,-1.0,-0.5,0.0,0.5,1.0,1.5 ] --- , [ -2.0,-1.5,-1.0,-0.5,0.0,0.5,1.0,1.5,2.0 ] --- , [ -1.5,-1.0,-0.5,0.0,0.5,1.0,1.5,2.0,2.5 ] --- , [ -1.0,-0.5,0.0,0.5,1.0,1.5,2.0,2.5,3.0 ] --- , [ -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 Seq (\ scale (i :. j) -> scale i + scale j :: Double) (-2) 1 (Sz 4) 2 +-- Array D Seq (Sz (9 :. 9)) +-- [ [ -4.0, -3.5, -3.0, -2.5, -2.0, -1.5, -1.0, -0.5, 0.0 ] +-- , [ -3.5, -3.0, -2.5, -2.0, -1.5, -1.0, -0.5, 0.0, 0.5 ] +-- , [ -3.0, -2.5, -2.0, -1.5, -1.0, -0.5, 0.0, 0.5, 1.0 ] +-- , [ -2.5, -2.0, -1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5 ] +-- , [ -2.0, -1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, 2.0 ] +-- , [ -1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, 2.0, 2.5 ] +-- , [ -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0 ] +-- , [ -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) @@ -205,8 +216,8 @@ -> 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 n = - fmap (\ix -> f scale ix) $ range' comp zeroIndex (liftIndex (n *) sz) +fromFunction comp f a d (Sz sz) n = + (f scale) <$> rangeInclusive comp zeroIndex (liftIndex (n *) sz) where nFrac = fromIntegral n scale i = a + d * fromIntegral i / nFrac @@ -217,23 +228,23 @@ -- | Similar to `fromFunction`, but will create an array from a function with sample points in the -- middle of cells. -- --- >>> fromFunctionMidpoint Seq (\ scale (i :. j) -> scale i + scale j) (-2) 1 (4 :. 4) 2 --- (Array D Seq (8 :. 8) --- [ [ -3.5,-3.0,-2.5,-2.0,-1.5,-1.0,-0.5,0.0 ] --- , [ -3.0,-2.5,-2.0,-1.5,-1.0,-0.5,0.0,0.5 ] --- , [ -2.5,-2.0,-1.5,-1.0,-0.5,0.0,0.5,1.0 ] --- , [ -2.0,-1.5,-1.0,-0.5,0.0,0.5,1.0,1.5 ] --- , [ -1.5,-1.0,-0.5,0.0,0.5,1.0,1.5,2.0 ] --- , [ -1.0,-0.5,0.0,0.5,1.0,1.5,2.0,2.5 ] --- , [ -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 Seq (\ scale (i :. j) -> scale i + scale j :: Double) (-2) 1 (Sz 4) 2 +-- Array D Seq (Sz (8 :. 8)) +-- [ [ -3.5, -3.0, -2.5, -2.0, -1.5, -1.0, -0.5, 0.0 ] +-- , [ -3.0, -2.5, -2.0, -1.5, -1.0, -0.5, 0.0, 0.5 ] +-- , [ -2.5, -2.0, -1.5, -1.0, -0.5, 0.0, 0.5, 1.0 ] +-- , [ -2.0, -1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5 ] +-- , [ -1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, 2.0 ] +-- , [ -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, 2.0, 2.5 ] +-- , [ -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 -> ix -> Int -> Array D ix e -fromFunctionMidpoint comp f a d sz n = - fmap (\ix -> f scale ix) $ range' comp zeroIndex (liftIndex (\i -> n * i - 1) sz) + Comp -> ((Int -> a) -> ix -> e) -> a -> a -> Sz ix -> Int -> Array D ix e +fromFunctionMidpoint comp f a d (Sz sz) n = + (\ix -> f scale ix) <$> rangeInclusive comp zeroIndex (liftIndex (\i -> n * i - 1) sz) where nFrac = fromIntegral n dx2 = d / nFrac / 2 @@ -242,16 +253,6 @@ {-# INLINE fromFunctionMidpoint #-} --- TODO: make this function external --- https://github.com/lehins/massiv/issues/47 -range' :: Index ix => Comp -> ix -> ix -> Array D ix ix -range' comp ixFrom ixTo = - makeArray comp sz (\ix -> liftIndex2 (+) ixFrom ix) - where - sz = liftIndex2 (-) (liftIndex (+ 1) ixTo) ixFrom -{-# INLINE range' #-} - - -- $integral_intro -- -- Inspiration for the code in this module was taken from [Paul Dawkins Online @@ -262,7 +263,7 @@ -- Implementation-wise, integral approximation here relies heavily on stencils with stride, as such -- computation is fast and is automatically parallelizable. -- --- Here are some example of where this can be useful: +-- Here are some examples of where this can be useful: -- -- === Integral of a function on a region -- @@ -271,18 +272,20 @@ -- the function to that array, which will give us an array of @n + 1@ sample points, or looking from -- a different angle @n@ intervals. -- --- >>> f x = exp ( x ^ 2 ) --- >>> fromFunction Seq (\ scale x -> f (scale x)) 0 2 1 4 --- (Array D Seq (5) --- [ 1.0,1.2840254166877414,2.718281828459045,9.487735836358526,54.598150033144236 ]) +-- >>> import Data.Massiv.Array +-- >>> f x = exp ( x ^ (2 :: Int) ) :: Float +-- >>> fromFunction Seq (\ scale x -> f (scale x)) 0 2 (Sz1 1) 4 +-- Array D Seq (Sz1 5) +-- [ 1.0, 1.2840254, 2.7182817, 9.487736, 54.59815 ] -- --- Once we have that array of sample points ready we could use `integralApprox` and one of the +-- Once we have that array of sample points ready, we could use `integralApprox` and one of the -- stencils to compute an integral, but there are already functions that will do both steps for you: -- --- >>> simpsonsRule Seq U (\ scale x -> f (scale x)) 0 2 1 4 --- [ 17.353626450374566 ]) +-- >>> simpsonsRule Seq U (\ scale x -> f (scale x)) 0 2 (Sz1 1) 4 +-- Array M Seq (Sz1 1) +-- [ 17.353626 ] -- --- Scaling function @scale@ is what will change an array index into equally spaced and +-- @scale@ is the function that will change an array index into equally spaced and -- appropriately shifted values of @x, y, ...@ before they can get applied to @f(x, y, ...)@ -- -- === Accurate values of a function @@ -291,13 +294,13 @@ -- representation of a non-linear function is desired. Consider the same gaussian function applied -- to equally spaced values, with zero being in the middle of the vector: -- --- >>> xArr = makeArrayR D Seq (Ix1 4) $ \ i -> (fromIntegral i - 1.5 :: Float) +-- >>> xArr = makeArrayR D Seq (Sz1 4) $ \ i -> fromIntegral i - 1.5 :: Float -- >>> xArr --- (Array D Seq (4) --- [ -1.5,-0.5,0.5,1.5 ]) +-- Array D Seq (Sz1 4) +-- [ -1.5, -0.5, 0.5, 1.5 ] -- >>> fmap f xArr --- (Array D Seq (4) --- [ 9.487736,1.2840254,1.2840254,9.487736 ]) +-- Array D Seq (Sz1 4) +-- [ 9.487736, 1.2840254, 1.2840254, 9.487736 ] -- -- The problem with above example is that computed values do not accurately represent the total -- value contained within each vector cell. For that reason if your were to later use it for example @@ -307,16 +310,16 @@ -- -- >>> startValue = -2 :: Float -- >>> distPerCell = 1 :: Float --- >>> desiredSize = 4 :: Ix1 +-- >>> desiredSize = Sz1 4 :: Sz1 -- >>> numSamples = 4 :: Int -- >>> xArrX4 = fromFunction Seq ($) startValue distPerCell desiredSize numSamples -- >>> xArrX4 --- (Array D Seq (17) --- [ -2.0,-1.75,-1.5,-1.25,-1.0,-0.75,-0.5,-0.25,0.0,0.25,0.5,0.75,1.0,1.25,1.5,1.75,2.0 ]) +-- Array D Seq (Sz1 17) +-- [ -2.0, -1.75, -1.5, -1.25, -1.0, -0.75, -0.5, -0.25, 0.0, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0 ] -- >>> yArrX4 = computeAs U $ fmap f xArrX4 -- >>> integralApprox trapezoidStencil distPerCell desiredSize numSamples yArrX4 --- (Array M Seq (4) --- [ 16.074406,1.4906789,1.4906789,16.074408 ]) +-- Array M Seq (Sz1 4) +-- [ 16.074406, 1.4906789, 1.4906789, 16.074408 ] -- -- We can clearly see the difference is huge, but it doesn't mean it is much better than our -- previous estimate. In order to get more accurate results we can use a better Simpson's rule for @@ -324,5 +327,5 @@ -- and `yArr`, there are functions like `simpsonRule` that will take care it for you: -- -- >>> simpsonsRule Seq U (\ scale i -> f (scale i)) startValue distPerCell desiredSize 128 --- (Array M Seq (4) --- [ 14.989977,1.4626511,1.4626517,14.989977 ]) +-- Array M Seq (Sz1 4) +-- [ 14.989977, 1.4626511, 1.4626517, 14.989977 ] diff --git a/src/Data/Massiv/Array/Ops/Construct.hs b/src/Data/Massiv/Array/Ops/Construct.hs --- a/src/Data/Massiv/Array/Ops/Construct.hs +++ b/src/Data/Massiv/Array/Ops/Construct.hs @@ -1,31 +1,51 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE ExplicitForAll #-} -{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE ExplicitForAll #-} +{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} -- | -- Module : Data.Massiv.Array.Ops.Construct --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable -- module Data.Massiv.Array.Ops.Construct - ( -- ** From a function - makeArray + ( -- ** With constant value + empty + , singleton + , replicate + -- ** With a function + , makeArray + , makeArrayLinear , makeArrayR + , makeArrayLinearR , makeVectorR - , singleton + , iterateN + , iiterateN + , unfoldlS_ + , iunfoldlS_ + , unfoldrS_ + , iunfoldrS_ -- *** Applicative , makeArrayA , makeArrayAR -- ** Enumeration + , (...) + , (..:) , range + , rangeStepM + , rangeStep' , rangeStep + , rangeInclusive + , rangeStepInclusiveM + , rangeSize + , rangeStepSize , enumFromN , enumFromStepN -- ** Expansion @@ -35,98 +55,342 @@ , expandInner ) where -import Data.Massiv.Array.Delayed.Internal -import Data.Massiv.Core.Common -import Data.Massiv.Array.Ops.Map as A -import Prelude as P - +import Control.Applicative hiding (empty) +import Control.Monad (void) +import Control.Monad.ST +import Data.Massiv.Array.Delayed.Pull +import Data.Massiv.Array.Delayed.Push +import Data.Massiv.Core.Common +import Prelude as P 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. -- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array -- >>> makeArrayR U Par (Sz (2 :> 3 :. 4)) (\ (i :> j :. k) -> i * i + j * j == k * k) --- (Array U Par (2 :> 3 :. 4) --- [ [ [ True,False,False,False ] --- , [ False,True,False,False ] --- , [ False,False,True,False ] +-- Array U Par (Sz (2 :> 3 :. 4)) +-- [ [ [ True, False, False, False ] +-- , [ False, True, False, False ] +-- , [ False, False, True, False ] -- ] --- , [ [ False,True,False,False ] --- , [ False,False,False,False ] --- , [ False,False,False,False ] +-- , [ [ False, True, False, False ] +-- , [ False, False, False, False ] +-- , [ False, False, False, False ] -- ] --- ]) +-- ] -- +-- @since 0.1.0 makeArrayR :: Construct r ix e => r -> Comp -> Sz ix -> (ix -> e) -> Array r ix e makeArrayR _ = makeArray {-# INLINE makeArrayR #-} +-- | Same as `makeArrayLinear`, but with ability to supply resulting representation +-- +-- @since 0.3.0 +makeArrayLinearR :: Construct r ix e => r -> Comp -> Sz ix -> (Int -> e) -> Array r ix e +makeArrayLinearR _ = makeArrayLinear +{-# INLINE makeArrayLinearR #-} -- | Same as `makeArrayR`, but restricted to 1-dimensional arrays. +-- +-- @since 0.1.0 makeVectorR :: Construct r Ix1 e => r -> Comp -> Sz1 -> (Ix1 -> e) -> Array r Ix1 e makeVectorR _ = makeArray {-# INLINE makeVectorR #-} + +-- | Replicate the same element +-- +-- @since 0.3.0 +replicate :: forall r ix e . Construct r ix e => Comp -> Sz ix -> e -> Array r ix e +replicate comp sz e = makeArray comp sz (const e) +{-# INLINE replicate #-} + + +newtype STA r ix a = STA {_runSTA :: forall s. MArray s r ix a -> ST s (Array r ix a)} + +runSTA :: Mutable r ix e => Sz ix -> STA r ix e -> Array r ix e +runSTA !sz (STA m) = runST (unsafeNew sz >>= m) +{-# INLINE runSTA #-} + -- | Similar to `makeArray`, but construct the array sequentially using an `Applicative` interface -- disregarding the supplied `Comp`. -- +-- /Note/ - using `Data.Massiv.Array.Mutable.generateArray` or +-- `Data.Massiv.Array.Mutable.generateArrayS` will always be faster, althought not always possible. +-- +-- -- @since 0.2.6 -- -makeArrayA :: (Mutable r ix b, Applicative f) => Comp -> Sz ix -> (ix -> f b) -> f (Array r ix b) -makeArrayA comp sz f = traverseA f $ makeArrayR D comp sz id -{-# INLINE makeArrayA #-} +makeArrayA :: + forall r ix e f. (Mutable r ix e, Applicative f) + => Comp + -> Sz ix + -> (ix -> f e) + -> f (Array r ix e) +makeArrayA !comp !sz f = + let n = totalElem sz + go !i + | i < n = + liftA2 + (\e (STA st) -> STA (\ma -> unsafeLinearWrite ma i e >> st ma)) + (f (fromLinearIndex sz i)) + (go (i + 1)) + | otherwise = pure (STA (unsafeFreeze comp)) + in runSTA sz <$> go 0 +{-# INLINE makeArrayA #-} + -- | Same as `makeArrayA`, but with ability to supply result array representation. -- -- @since 0.2.6 -- -makeArrayAR :: (Mutable r ix b, Applicative f) => r -> Comp -> Sz ix -> (ix -> f b) -> f (Array r ix b) +makeArrayAR :: + forall r ix e f. (Mutable r ix e, Applicative f) + => r + -> Comp + -> Sz ix + -> (ix -> f e) + -> f (Array r ix e) makeArrayAR _ = makeArrayA {-# INLINE makeArrayAR #-} --- | Create a vector with a range of @Int@s incremented by 1. --- @range k0 k1 == rangeStep k0 k1 1@ + +-- | Sequentially iterate over each cell in the array in the row-major order while continuously +-- aplying the accumulator at each step. -- --- >>> range Seq 1 6 --- (Array D Seq (5) --- [ 1,2,3,4,5 ]) --- >>> range Seq (-2) 3 --- (Array D Seq (5) --- [ -2,-1,0,1,2 ]) +-- ==== __Example__ -- -range :: Comp -> Int -> Int -> Array D Ix1 Int -range comp !from !to = makeArray comp (max 0 (to - from)) (+ from) +-- >>> import Data.Massiv.Array +-- >>> iterateN Seq (Sz2 2 10) succ (10 :: Int) +-- Array DL Seq (Sz (2 :. 10)) +-- [ [ 11, 12, 13, 14, 15, 16, 17, 18, 19, 20 ] +-- , [ 21, 22, 23, 24, 25, 26, 27, 28, 29, 30 ] +-- ] +-- +-- @since 0.3.0 +iterateN :: forall ix e . Index ix => Comp -> Sz ix -> (e -> e) -> e -> Array DL ix e +iterateN comp sz f = unfoldrS_ comp 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 => Comp -> Sz ix -> (e -> ix -> e) -> e -> Array DL ix e +iiterateN comp sz f = iunfoldrS_ comp sz $ \a ix -> let !a' = f a ix in (a', a') +{-# INLINE iiterateN #-} + + +-- | +-- +-- @since 0.3.0 +unfoldrS_ :: forall ix e a . Construct DL ix e => Comp -> Sz ix -> (a -> (e, a)) -> a -> Array DL ix e +unfoldrS_ comp sz f = iunfoldrS_ comp sz (\a _ -> f a) +{-# INLINE unfoldrS_ #-} + +-- | +-- +-- @since 0.3.0 +iunfoldrS_ + :: Construct DL ix e => Comp -> Sz ix -> (a -> ix -> (e, a)) -> a -> Array DL ix e +iunfoldrS_ comp sz f acc0 = + DLArray + { dlComp = comp + , dlSize = sz + , dlLoad = + \_ startAt dlWrite -> + void $ + loopM startAt (< (totalElem sz + startAt)) (+ 1) acc0 $ \ !i !acc -> do + let (e, acc') = f acc $ fromLinearIndex sz (i - startAt) + dlWrite i e + pure acc' + } +{-# 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.unfoldlPrim` to achive such effect. +-- +-- @since 0.3.0 +unfoldlS_ :: Construct DL ix e => Comp -> Sz ix -> (a -> (a, e)) -> a -> Array DL ix e +unfoldlS_ comp sz f = iunfoldlS_ comp sz (const f) +{-# INLINE unfoldlS_ #-} + +-- | Unfold sequentially from the right with an index aware function. +-- +-- @since 0.3.0 +iunfoldlS_ + :: Construct DL ix e => Comp -> Sz ix -> (ix -> a -> (a, e)) -> a -> Array DL ix e +iunfoldlS_ comp sz f acc0 = + DLArray + { dlComp = comp + , dlSize = sz + , dlLoad = + \ _ startAt dlWrite -> + void $ loopDeepM startAt (< (totalElem sz + startAt)) (+ 1) acc0 $ \ !i !acc -> do + let (acc', e) = f (fromLinearIndex sz (i - startAt)) acc + dlWrite i e + pure acc' + } +{-# INLINE iunfoldlS_ #-} + + +-- | Handy synonym for `rangeInclusive` `Seq` +-- +-- >>> 4 ... 10 +-- Array D Seq (Sz1 7) +-- [ 4, 5, 6, 7, 8, 9, 10 ] +-- +-- @since 0.3.0 +(...) :: Int -> Int -> Array D Ix1 Int +(...) = rangeInclusive Seq +{-# INLINE (...) #-} + +-- | Handy synonym for `rangeInclusive` `Seq` +-- +-- >>> 4 ..: 10 +-- Array D Seq (Sz1 6) +-- [ 4, 5, 6, 7, 8, 9 ] +-- +-- @since 0.3.0 +(..:) :: Int -> Int -> Array D Ix1 Int +(..:) = 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 +-- incrimeted by one. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array +-- >>> range Seq (Ix1 1) 6 +-- Array D Seq (Sz1 5) +-- [ 1, 2, 3, 4, 5 ] +-- >>> fromIx2 <$> range Seq (-1) (2 :. 2) +-- Array D Seq (Sz (3 :. 3)) +-- [ [ (-1,-1), (-1,0), (-1,1) ] +-- , [ (0,-1), (0,0), (0,1) ] +-- , [ (1,-1), (1,0), (1,1) ] +-- ] +-- +-- @since 0.1.0 +range :: Index ix => Comp -> ix -> ix -> Array D ix ix +range comp !from !to = rangeSize comp from (Sz (liftIndex2 (-) to from)) {-# INLINE range #-} -- | Same as `range`, but with a custom step. -- --- >>> rangeStep Seq 1 2 6 --- (Array D Seq (3) --- [ 1,3,5 ]) +-- @since 0.1.0 +rangeStep :: Index ix => Comp -> ix -> ix -> ix -> Maybe (Array D ix ix) +rangeStep = rangeStepM +{-# INLINE rangeStep #-} +{-# DEPRECATED rangeStep "In favor of more general `rangeStepM`" #-} + +-- | Same as `range`, but with a custom step. -- -rangeStep :: Comp -- ^ Computation strategy - -> Int -- ^ Start - -> Int -- ^ Step (Can't be zero) - -> Int -- ^ End - -> Array D Ix1 Int -rangeStep comp !from !step !to - | step == 0 = error "rangeStep: Step can't be zero" +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array +-- >>> rangeStepM Seq (Ix1 1) 2 8 +-- Array D Seq (Sz1 4) +-- [ 1, 3, 5, 7 ] +-- >>> rangeStepM Seq (Ix1 1) 0 8 +-- *** Exception: IndexZeroException: 0 +-- +-- @since 0.3.0 +rangeStepM :: (Index ix, MonadThrow m) => + Comp -- ^ Computation strategy + -> ix -- ^ Start + -> ix -- ^ Step (Can't have zeros) + -> ix -- ^ End + -> m (Array D ix ix) +rangeStepM comp !from !step !to + | foldlIndex (\acc i -> acc || i == 0) False step = throwM $ IndexZeroException step | otherwise = - let (sz, r) = (to - from) `divMod` step - in makeArray comp (sz + signum r) (\i -> from + i * step) -{-# INLINE rangeStep #-} + 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 `rangeStep`, but will throw an error whenever @step@ contains zeros. +-- +-- ==== __Example__ +-- +-- >>> import Data.Massiv.Array +-- >>> rangeStep' Seq (Ix1 1) 2 6 +-- Array D Seq (Sz1 3) +-- [ 1, 3, 5 ] +-- +-- @since 0.3.0 +rangeStep' :: Index ix => Comp -> ix -> ix -> ix -> Array D ix ix +rangeStep' comp from step = either throw id . rangeStepM comp from step +{-# INLINE rangeStep' #-} +-- | Just like `range`, except the finish index is included. +-- +-- @since 0.3.0 +rangeInclusive :: Index ix => Comp -> ix -> ix -> Array D ix ix +rangeInclusive comp ixFrom ixTo = + rangeSize comp ixFrom (Sz (liftIndex2 (-) (liftIndex (+ 1) ixTo) ixFrom)) +{-# INLINE rangeInclusive #-} + + +-- | Just like `rangeStep`, except the finish index is included. +-- +-- @since 0.3.0 +rangeStepInclusiveM :: (MonadThrow m, Index ix) => Comp -> ix -> ix -> ix -> m (Array D ix ix) +rangeStepInclusiveM comp ixFrom step ixTo = rangeStepM comp ixFrom step (liftIndex (1 +) ixTo) +{-# INLINE rangeStepInclusiveM #-} + + +-- | Create an array of specified size with indices starting with some index at position @0@ and +-- incremented by @1@ until the end of the array is reached +-- +-- @since 0.3.0 +rangeSize :: Index ix => + Comp + -> ix -- ^ @x@ - start value + -> Sz ix -- ^ @sz@ - Size of resulting array + -> Array D ix ix +rangeSize comp !from !sz = makeArray comp sz (liftIndex2 (+) from) +{-# INLINE rangeSize #-} + +-- | Same as `rangeSize`, but with ability to specify the step. +-- +-- @since 0.3.0 +rangeStepSize :: Index ix => + Comp + -> ix -- ^ @x@ - start value + -> ix -- ^ @delta@ - step value + -> Sz ix -- ^ @sz@ - Size of resulting array + -> Array D ix ix +rangeStepSize comp !from !step !sz = + makeArray comp sz (liftIndex2 (+) from . liftIndex2 (*) step) +{-# INLINE rangeStepSize #-} + + -- | Same as `enumFromStepN` with step @delta = 1@. -- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array -- >>> enumFromN Seq (5 :: Double) 3 --- (Array D Seq (3) --- [ 5.0,6.0,7.0 ]) +-- Array D Seq (Sz1 3) +-- [ 5.0, 6.0, 7.0 ] -- +-- @since 0.1.0 enumFromN :: Num e => Comp -> e -- ^ @x@ - start value - -> Int -- ^ @n@ - length of resulting vector. + -> Sz1 -- ^ @n@ - length of resulting vector. -> Array D Ix1 e enumFromN comp !from !sz = makeArray comp sz $ \ i -> fromIntegral i + from {-# INLINE enumFromN #-} @@ -137,15 +401,19 @@ -- x + delta ..]@. Major difference is that `fromList` constructs an `Array` with manifest -- representation, while `enumFromStepN` is delayed. -- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array -- >>> enumFromStepN Seq 1 (0.1 :: Double) 5 --- (Array D Seq (5) --- [ 1.0,1.1,1.2,1.3,1.4 ]) +-- Array D Seq (Sz1 5) +-- [ 1.0, 1.1, 1.2, 1.3, 1.4 ] -- +-- @since 0.1.0 enumFromStepN :: Num e => Comp -> e -- ^ @x@ - start value -> e -- ^ @delta@ - step value - -> Int -- ^ @n@ - length of resulting vector + -> Sz1 -- ^ @n@ - length of resulting vector -> Array D Ix1 e enumFromStepN comp !from !step !sz = makeArray comp sz $ \ i -> from + fromIntegral i * step {-# INLINE enumFromStepN #-} @@ -166,46 +434,46 @@ -- you use 'expandWithin' and a creation function @a -> Int -> b@ to create an -- @`Array` `D` `Ix2` b@ or @`Array` `D` `Ix3` b@ respectfully. -- --- @since 0.2.6 --- -- ====__Examples__ -- --- >>> a = makeArrayR U Seq (Ix1 6) (+10) -- Imagine (+10) is some expensive function +-- >>> import Data.Massiv.Array +-- >>> a = makeArrayR U Seq (Sz1 6) (+10) -- Imagine (+10) is some expensive function -- >>> a --- (Array U Seq (6) --- [ 10,11,12,13,14,15 ]) +-- Array U Seq (Sz1 6) +-- [ 10, 11, 12, 13, 14, 15 ] -- >>> expandWithin Dim1 5 (\ e j -> (j + 1) * 100 + e) a :: Array D Ix2 Int --- (Array D Seq (6 :. 5) --- [ [ 110,210,310,410,510 ] --- , [ 111,211,311,411,511 ] --- , [ 112,212,312,412,512 ] --- , [ 113,213,313,413,513 ] --- , [ 114,214,314,414,514 ] --- , [ 115,215,315,415,515 ] --- ]) +-- Array D Seq (Sz (6 :. 5)) +-- [ [ 110, 210, 310, 410, 510 ] +-- , [ 111, 211, 311, 411, 511 ] +-- , [ 112, 212, 312, 412, 512 ] +-- , [ 113, 213, 313, 413, 513 ] +-- , [ 114, 214, 314, 414, 514 ] +-- , [ 115, 215, 315, 415, 515 ] +-- ] -- >>> expandWithin Dim2 5 (\ e j -> (j + 1) * 100 + e) a :: Array D Ix2 Int --- (Array D Seq (5 :. 6) --- [ [ 110,111,112,113,114,115 ] --- , [ 210,211,212,213,214,215 ] --- , [ 310,311,312,313,314,315 ] --- , [ 410,411,412,413,414,415 ] --- , [ 510,511,512,513,514,515 ] --- ]) +-- Array D Seq (Sz (5 :. 6)) +-- [ [ 110, 111, 112, 113, 114, 115 ] +-- , [ 210, 211, 212, 213, 214, 215 ] +-- , [ 310, 311, 312, 313, 314, 315 ] +-- , [ 410, 411, 412, 413, 414, 415 ] +-- , [ 510, 511, 512, 513, 514, 515 ] +-- ] -- -expandWithin - :: (IsIndexDimension ix n, Manifest r (Lower ix) a) +-- @since 0.2.6 +expandWithin :: + forall ix e r n a. (IsIndexDimension ix n, Manifest r (Lower ix) a) => Dimension n -> Int - -> (a -> Int -> b) + -> (a -> Int -> e) -> Array r (Lower ix) a - -> Array D ix b -expandWithin dim k f arr = do + -> Array D ix e +expandWithin dim k f arr = makeArray (getComp arr) sz $ \ix -> let (i, ixl) = pullOutDimension ix dim in f (unsafeIndex arr ixl) i where - szl = size arr - sz = insertDimension szl dim k + szl = unSz (size arr) + sz = Sz (insertDimension szl dim k) {-# INLINE expandWithin #-} -- | Similar to `expandWithin`, except that dimension is specified at a value level, which means it @@ -224,8 +492,8 @@ let (i, ixl) = pullOutDim' ix dim in f (unsafeIndex arr ixl) i where - szl = size arr - sz = insertDim' szl dim k + szl = unSz (size arr) + sz = Sz (insertDim' szl dim k) {-# INLINE expandWithin' #-} -- | Similar to `expandWithin`, except it uses the outermost dimension. @@ -243,7 +511,7 @@ in f (unsafeIndex arr ixl) i where szl = size arr - sz = consDim k szl + sz = consSz (Sz k) szl {-# INLINE expandOuter #-} -- | Similar to `expandWithin`, except it uses the innermost dimension. @@ -261,5 +529,5 @@ in f (unsafeIndex arr ixl) i where szl = size arr - sz = snocDim szl k + sz = snocSz szl (Sz k) {-# INLINE expandInner #-} 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 @@ -1,11 +1,11 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE GADTs #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE ScopedTypeVariables #-} -- | -- Module : Data.Massiv.Array.Ops.Fold --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental @@ -16,13 +16,16 @@ -- ** Unstructured folds -- $unstruct_folds - fold , ifoldMono , foldMono , ifoldSemi , foldSemi + , minimumM + , minimum' , minimum + , maximumM + , maximum' , maximum , sum , product @@ -84,24 +87,17 @@ , foldrP , ifoldlP , ifoldrP - , foldlOnP , ifoldlIO - , foldrOnP - , ifoldlOnP - , ifoldrOnP , ifoldrIO ) where -import Data.Massiv.Array.Delayed.Internal -import Data.Massiv.Array.Ops.Fold.Internal -import Data.Massiv.Array.Ops.Map (imap, map) -import Data.Massiv.Core -import Data.Massiv.Core.Common -import Data.Semigroup -import Prelude hiding (all, and, any, - foldl, foldr, map, - maximum, minimum, or, - product, sum) +import Data.Massiv.Array.Delayed.Pull +import Data.Massiv.Array.Ops.Fold.Internal +import Data.Massiv.Core +import Data.Massiv.Core.Common +import Data.Massiv.Core.Index.Internal (Sz(..)) +import Prelude hiding (all, and, any, foldl, foldr, map, maximum, minimum, or, + product, sum) -- | /O(n)/ - Monoidal fold over an array with an index aware function. Also known as reduce. -- @@ -111,22 +107,10 @@ => (ix -> e -> m) -- ^ Convert each element of an array to an appropriate `Monoid`. -> Array r ix e -- ^ Source array -> m -ifoldMono f = foldlInternal mappend mempty mappend mempty . imap f +ifoldMono f = ifoldlInternal (\a ix e -> a `mappend` f ix e) mempty mappend mempty {-# INLINE ifoldMono #-} --- | /O(n)/ - Monoidal fold over an array. Also known as reduce. --- --- @since 0.1.4 -foldMono :: - (Source r ix e, Monoid m) - => (e -> m) -- ^ Convert each element of an array to an appropriate `Monoid`. - -> Array r ix e -- ^ Source array - -> m -foldMono f = foldlInternal mappend mempty mappend mempty . map f -{-# INLINE foldMono #-} - - -- | /O(n)/ - Semigroup fold over an array with an index aware function. -- -- @since 0.2.4 @@ -136,7 +120,7 @@ -> m -- ^ Initial element that must be neutral to the (`<>`) function. -> Array r ix e -- ^ Source array -> m -ifoldSemi f m = foldlInternal (<>) m (<>) m . imap f +ifoldSemi f m = ifoldlInternal (\a ix e -> a <> f ix e) m (<>) m {-# INLINE ifoldSemi #-} @@ -149,7 +133,7 @@ -> m -- ^ Initial element that must be neutral to the (`<>`) function. -> Array r ix e -- ^ Source array -> m -foldSemi f m = foldlInternal (<>) m (<>) m . map f +foldSemi f m = foldlInternal (\a e -> a <> f e) m (<>) m {-# INLINE foldSemi #-} @@ -166,18 +150,20 @@ -- -- ====__Example__ -- --- >>> let arr = makeArrayR U Seq (2 :. 5) (toLinearIndex (2 :. 5)) +-- >>> import Data.Massiv.Array +-- >>> :set -XTypeApplications +-- >>> arr = makeArrayLinear @U Seq (Sz (2 :. 5)) id -- >>> arr --- (Array U Seq (2 :. 5) --- [ [ 0,1,2,3,4 ] --- , [ 5,6,7,8,9 ] --- ]) +-- Array U Seq (Sz (2 :. 5)) +-- [ [ 0, 1, 2, 3, 4 ] +-- , [ 5, 6, 7, 8, 9 ] +-- ] -- >>> foldlWithin Dim1 (flip (:)) [] arr --- (Array D Seq (2) --- [ [4,3,2,1,0],[9,8,7,6,5] ]) +-- Array D Seq (Sz1 2) +-- [ [4,3,2,1,0], [9,8,7,6,5] ] -- >>> foldlWithin Dim2 (flip (:)) [] arr --- (Array D Seq (5) --- [ [5,0],[6,1],[7,2],[8,3],[9,4] ]) +-- Array D Seq (Sz1 5) +-- [ [5,0], [6,1], [7,2], [8,3], [9,4] ] -- -- @since 0.2.4 foldlWithin :: (Index (Lower ix), IsIndexDimension ix n, Source r ix e) => @@ -211,7 +197,7 @@ ifoldlWithin' :: (Index (Lower ix), Source r ix e) => Dim -> (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a ifoldlWithin' dim f acc0 arr = - unsafeMakeArray (getComp arr) szl $ \ixl -> + makeArray (getComp arr) (SafeSz szl) $ \ixl -> iter (insertDim' ixl dim 0) (insertDim' ixl dim (k - 1)) @@ -220,7 +206,7 @@ acc0 (\ix acc' -> f ix acc' (unsafeIndex arr ix)) where - sz = size arr + SafeSz sz = size arr (k, szl) = pullOutDim' sz dim {-# INLINE ifoldlWithin' #-} @@ -243,7 +229,7 @@ ifoldrWithin' :: (Index (Lower ix), Source r ix e) => Dim -> (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a ifoldrWithin' dim f acc0 arr = - unsafeMakeArray (getComp arr) szl $ \ixl -> + makeArray (getComp arr) (SafeSz szl) $ \ixl -> iter (insertDim' ixl dim (k - 1)) (insertDim' ixl dim 0) @@ -252,7 +238,7 @@ acc0 (\ix acc' -> f ix (unsafeIndex arr ix) acc') where - sz = size arr + SafeSz sz = size arr (k, szl) = pullOutDim' sz dim {-# INLINE ifoldrWithin' #-} @@ -300,63 +286,103 @@ -- | /O(n)/ - Compute maximum of all elements. -maximum :: (Source r ix e, Ord e) => - Array r ix e -> e -maximum = \arr -> - if isEmpty arr - then error "Data.Massiv.Array.maximum - empty" - else fold max (evaluateAt arr zeroIndex) arr +-- +-- @since 0.3.0 +maximumM :: (MonadThrow m, Source r ix e, Ord e) => Array r ix e -> m e +maximumM arr = + if isEmpty arr + then throwM (SizeEmptyException (size arr)) + else let !e0 = unsafeIndex arr zeroIndex + in pure $ foldlInternal max e0 max e0 arr +{-# INLINE maximumM #-} + + +-- | /O(n)/ - Compute maximum of all elements. +-- +-- @since 0.1.0 +maximum :: (Source r ix e, Ord e) => Array r ix e -> e +maximum = maximum' {-# INLINE maximum #-} +{-# DEPRECATED maximum "In favor of a safer `maximumM` or an equivalent `maximum'`" #-} +-- | /O(n)/ - Compute maximum of all elements. +-- +-- @since 0.3.0 +maximum' :: (Source r ix e, Ord e) => Array r ix e -> e +maximum' = either throw id . maximumM +{-# INLINE maximum' #-} + -- | /O(n)/ - Compute minimum of all elements. -minimum :: (Source r ix e, Ord e) => - Array r ix e -> e -minimum = \arr -> - if isEmpty arr - then error "Data.Massiv.Array.minimum - empty" - else fold min (evaluateAt arr zeroIndex) arr -{-# INLINE minimum #-} +-- +-- @since 0.3.0 +minimumM :: (MonadThrow m, Source r ix e, Ord e) => Array r ix e -> m e +minimumM arr = + if isEmpty arr + then throwM (SizeEmptyException (size arr)) + else let !e0 = unsafeIndex arr zeroIndex + in pure $ foldlInternal min e0 min e0 arr +{-# INLINE minimumM #-} +-- | /O(n)/ - Compute minimum of all elements. +-- +-- @since 0.3.0 +minimum' :: (Source r ix e, Ord e) => Array r ix e -> e +minimum' = either throw id . minimumM +{-# INLINE minimum' #-} +-- | /O(n)/ - Compute minimum of all elements. +-- +-- @since 0.1.0 +minimum :: (Source r ix e, Ord e) => Array r ix e -> e +minimum = minimum' +{-# INLINE minimum #-} +{-# DEPRECATED minimum "In favor of a safer `minimumM` or an equivalent `minimum'`" #-} + -- | /O(n)/ - Compute sum of all elements. -sum :: (Source r ix e, Num e) => - Array r ix e -> e -sum = fold (+) 0 +-- +-- @since 0.1.0 +sum :: (Source r ix e, Num e) => Array r ix e -> e +sum = foldlInternal (+) 0 (+) 0 {-# INLINE sum #-} -- | /O(n)/ - Compute product of all elements. -product :: (Source r ix e, Num e) => - Array r ix e -> e -product = fold (*) 1 +-- +-- @since 0.1.0 +product :: (Source r ix e, Num e) => Array r ix e -> e +product = foldlInternal (*) 1 (*) 1 {-# INLINE product #-} -- | /O(n)/ - Compute conjunction of all elements. -and :: (Source r ix Bool) => - Array r ix Bool -> Bool -and = fold (&&) True +-- +-- @since 0.1.0 +and :: Source r ix Bool => Array r ix Bool -> Bool +and = foldlInternal (&&) True (&&) True {-# INLINE and #-} -- | /O(n)/ - Compute disjunction of all elements. -or :: Source r ix Bool => - Array r ix Bool -> Bool -or = fold (||) False +-- +-- @since 0.1.0 +or :: Source r ix Bool => Array r ix Bool -> Bool +or = foldlInternal (||) False (||) False {-# INLINE or #-} -- | /O(n)/ - Determines whether all element of the array satisfy the predicate. -all :: Source r ix e => - (e -> Bool) -> Array r ix e -> Bool -all f = foldlInternal (\acc el -> acc && f el) True (&&) True +-- +-- @since 0.1.0 +all :: Source r ix e => (e -> Bool) -> Array r ix e -> Bool +all f = foldlInternal (\acc e -> acc && f e) True (&&) True {-# INLINE all #-} -- | /O(n)/ - Determines whether any element of the array satisfies the predicate. -any :: Source r ix e => - (e -> Bool) -> Array r ix e -> Bool -any f = foldlInternal (\acc el -> acc || f el) False (||) False +-- +-- @since 0.1.0 +any :: Source r ix e => (e -> Bool) -> Array r ix e -> Bool +any f = foldlInternal (\acc e -> acc || f e) False (||) False {-# INLINE any #-} 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 @@ -1,11 +1,11 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE UndecidableInstances #-} -- | -- Module : Data.Massiv.Array.Ops.Fold.Internal --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental @@ -28,7 +28,9 @@ , ifoldrM_ --Special folds , fold + , foldMono , foldlInternal + , ifoldlInternal , foldrFB , lazyFoldlS , lazyFoldrS @@ -37,86 +39,106 @@ , foldrP , ifoldlP , ifoldrP - , foldlOnP , ifoldlIO - , foldrOnP - , ifoldlOnP - , ifoldrOnP , ifoldrIO ) where -import Control.Monad (void, when) -import qualified Data.Foldable as F -import Data.Functor.Identity (runIdentity) -import Data.Massiv.Core -import Data.Massiv.Core.Common -import Data.Massiv.Core.Scheduler -import Prelude hiding (all, and, any, foldl, foldr, - maximum, minimum, or, product, sum) -import System.IO.Unsafe (unsafePerformIO) - - +import Control.Scheduler +import Control.Monad (void, when) +import qualified Data.Foldable as F +import Data.Functor.Identity (runIdentity) +import Data.Massiv.Core.Common +import Prelude hiding (foldl, foldr) +import System.IO.Unsafe (unsafePerformIO) -- | /O(n)/ - Unstructured fold of an array. -fold :: Source r ix e => - (e -> e -> e) -- ^ Folding function (like with left fold, first argument - -- is an accumulator) - -> e -- ^ Initial element. Has to be neutral with respect to the folding - -- function. - -> Array r ix e -- ^ Source array - -> e -fold f initAcc = foldlInternal f initAcc f initAcc +-- +-- @since 0.3.0 +fold :: + (Monoid e, Source r ix e) + => Array r ix e -- ^ Source array + -> e +fold = foldlInternal mappend mempty mappend mempty {-# INLINE fold #-} +-- | /O(n)/ - This is exactly like `Data.Foldable.foldMap`, but for arrays. Fold over an array, +-- while converting each element into a `Monoid`. Also known as map-reduce. If elements of the array +-- are already a `Monoid` you can use `fold` instead. +-- +-- @since 0.1.4 +foldMono :: + (Source r ix e, Monoid m) + => (e -> m) -- ^ Convert each element of an array to an appropriate `Monoid`. + -> Array r ix e -- ^ Source array + -> m +foldMono f = foldlInternal (\a e -> a `mappend` f e) mempty mappend mempty +{-# INLINE foldMono #-} + -- | /O(n)/ - Monadic left fold. +-- +-- @since 0.1.0 foldlM :: (Source r ix e, Monad m) => (a -> e -> m a) -> a -> Array r ix e -> m a foldlM f = ifoldlM (\ a _ b -> f a b) {-# INLINE foldlM #-} -- | /O(n)/ - Monadic left fold, that discards the result. +-- +-- @since 0.1.0 foldlM_ :: (Source r ix e, Monad m) => (a -> e -> m a) -> a -> Array r ix e -> m () foldlM_ f = ifoldlM_ (\ a _ b -> f a b) {-# INLINE foldlM_ #-} -- | /O(n)/ - Monadic left fold with an index aware function. +-- +-- @since 0.1.0 ifoldlM :: (Source r ix e, Monad m) => (a -> ix -> e -> m a) -> a -> Array r ix e -> m a ifoldlM f !acc !arr = - iterM zeroIndex (size arr) (pureIndex 1) (<) acc $ \ !ix !a -> f a ix (unsafeIndex arr ix) + iterM zeroIndex (unSz (size arr)) (pureIndex 1) (<) acc $ \ !ix !a -> f a ix (unsafeIndex arr ix) {-# INLINE ifoldlM #-} -- | /O(n)/ - Monadic left fold with an index aware function, that discards the result. +-- +-- @since 0.1.0 ifoldlM_ :: (Source r ix e, Monad m) => (a -> ix -> e -> m a) -> a -> Array r ix e -> m () ifoldlM_ f acc = void . ifoldlM f acc {-# INLINE ifoldlM_ #-} -- | /O(n)/ - Monadic right fold. +-- +-- @since 0.1.0 foldrM :: (Source r ix e, Monad m) => (e -> a -> m a) -> a -> Array r ix e -> m a foldrM f = ifoldrM (\_ e a -> f e a) {-# INLINE foldrM #-} -- | /O(n)/ - Monadic right fold, that discards the result. +-- +-- @since 0.1.0 foldrM_ :: (Source r ix e, Monad m) => (e -> a -> m a) -> a -> Array r ix e -> m () foldrM_ f = ifoldrM_ (\_ e a -> f e a) {-# INLINE foldrM_ #-} -- | /O(n)/ - Monadic right fold with an index aware function. +-- +-- @since 0.1.0 ifoldrM :: (Source r ix e, Monad m) => (ix -> e -> a -> m a) -> a -> Array r ix e -> m a ifoldrM f !acc !arr = - iterM (liftIndex (subtract 1) (size arr)) zeroIndex (pureIndex (-1)) (>=) acc $ \ !ix !acc0 -> + iterM (liftIndex (subtract 1) (unSz (size arr))) zeroIndex (pureIndex (-1)) (>=) acc $ \ !ix !acc0 -> f ix (unsafeIndex arr ix) acc0 {-# INLINE ifoldrM #-} -- | /O(n)/ - Monadic right fold with an index aware function, that discards the result. +-- +-- @since 0.1.0 ifoldrM_ :: (Source r ix e, Monad m) => (ix -> e -> a -> m a) -> a -> Array r ix e -> m () ifoldrM_ f !acc !arr = void $ ifoldrM f acc arr {-# INLINE ifoldrM_ #-} @@ -124,6 +146,8 @@ -- | /O(n)/ - Left fold, computed sequentially with lazy accumulator. +-- +-- @since 0.1.0 lazyFoldlS :: Source r ix e => (a -> e -> a) -> a -> Array r ix e -> a lazyFoldlS f initAcc arr = go initAcc 0 where len = totalElem (size arr) @@ -133,18 +157,24 @@ -- | /O(n)/ - Right fold, computed sequentially with lazy accumulator. +-- +-- @since 0.1.0 lazyFoldrS :: Source r ix e => (e -> a -> a) -> a -> Array r ix e -> a lazyFoldrS = foldrFB {-# INLINE lazyFoldrS #-} -- | /O(n)/ - Left fold, computed sequentially. +-- +-- @since 0.1.0 foldlS :: Source r ix e => (a -> e -> a) -> a -> Array r ix e -> a foldlS f = ifoldlS (\ a _ e -> f a e) {-# INLINE foldlS #-} -- | /O(n)/ - Left fold with an index aware function, computed sequentially. +-- +-- @since 0.1.0 ifoldlS :: Source r ix e => (a -> ix -> e -> a) -> a -> Array r ix e -> a ifoldlS f acc = runIdentity . ifoldlM (\ a ix e -> return $ f a ix e) acc @@ -152,12 +182,24 @@ -- | /O(n)/ - Right fold, computed sequentially. +-- +-- @since 0.1.0 foldrS :: Source r ix e => (e -> a -> a) -> a -> Array r ix e -> a foldrS f = ifoldrS (\_ e a -> f e a) {-# INLINE foldrS #-} +-- | /O(n)/ - Right fold with an index aware function, computed sequentially. +-- +-- @since 0.1.0 +ifoldrS :: Source r ix e => (ix -> e -> a -> a) -> a -> Array r ix e -> a +ifoldrS f acc = runIdentity . ifoldrM (\ ix e a -> return $ f ix e a) acc +{-# INLINE ifoldrS #-} + + -- | Version of foldr that supports @foldr/build@ list fusion implemented by GHC. +-- +-- @since 0.1.0 foldrFB :: Source r ix e => (e -> b -> b) -> b -> Array r ix e -> b foldrFB c n arr = go 0 where @@ -169,184 +211,142 @@ --- | /O(n)/ - Right fold with an index aware function, computed sequentially. -ifoldrS :: Source r ix e => (ix -> e -> a -> a) -> a -> Array r ix e -> a -ifoldrS f acc = runIdentity . ifoldrM (\ ix e a -> return $ f ix e a) acc -{-# INLINE ifoldrS #-} - - - --- | /O(n)/ - Left fold, computed in parallel. Parallelization of folding is implemented in such a --- way that an array is split into a number of chunks of equal length, plus an extra one for the --- left over. Number of chunks is the same as number of available cores (capabilities) plus one, and --- each chunk is individually folded by a separate core with a function @g@. Results from folding --- each chunk are further folded with another function @f@, thus allowing us to use information --- about the structure of an array during folding. +-- | /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 +-- different, hence is the `MonadIO` constraint. -- -- ===__Examples__ -- --- >>> foldlP (flip (:)) [] (flip (:)) [] $ makeArrayR U Seq (Ix1 11) id --- [[10,9,8,7,6,5,4,3,2,1,0]] --- --- And this is how the result would look like if the above computation would be performed in a --- program executed with @+RTS -N3@, i.e. with 3 capabilities: --- --- >>> foldlOnP [1,2,3] (flip (:)) [] (flip (:)) [] $ makeArrayR U Seq (Ix1 11) id --- [[10,9],[8,7,6],[5,4,3],[2,1,0]] +-- >>> import Data.Massiv.Array +-- >>> foldlP (flip (:)) [] (flip (:)) [] $ makeArrayR D Seq (Sz1 6) id +-- [[5,4,3,2,1,0]] +-- >>> foldlP (flip (:)) [] (++) [] $ makeArrayR D Seq (Sz1 6) id +-- [5,4,3,2,1,0] +-- >>> foldlP (flip (:)) [] (flip (:)) [] $ makeArrayR D (ParN 3) (Sz1 6) id +-- [[5,4],[3,2],[1,0]] +-- >>> foldlP (flip (:)) [] (++) [] $ makeArrayR D (ParN 3) (Sz1 6) id +-- [1,0,3,2,5,4] -- -foldlP :: Source r ix e => +-- @since 0.1.0 +foldlP :: (MonadIO m, Source r ix e) => (a -> e -> a) -- ^ Folding function @g@. -> a -- ^ Accumulator. Will be applied to @g@ multiple times, thus must be neutral. -> (b -> a -> b) -- ^ Chunk results folding function @f@. -> b -- ^ Accumulator for results of chunks folding. - -> Array r ix e -> IO b -foldlP f = ifoldlP (\ x _ -> f x) + -> Array r ix e -> m b +foldlP f fAcc g gAcc = liftIO . ifoldlP (\ x _ -> f x) fAcc g gAcc {-# INLINE foldlP #-} - --- | Just like `foldlP`, but allows you to specify which cores (capabilities) to run computation --- on. The order in which chunked results will be supplied to function @f@ is guaranteed to be --- consecutive and aligned with the folding direction. -foldlOnP - :: Source r ix e - => [Int] -> (a -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> IO b -foldlOnP wIds f = ifoldlOnP wIds (\ x _ -> f x) -{-# INLINE foldlOnP #-} - - - --- | Parallel left fold. -ifoldlIO :: Source r ix e => - [Int] -- ^ List of capabilities - -> (a -> ix -> e -> IO a) -- ^ Index aware folding IO action - -> a -- ^ Accumulator - -> (b -> a -> IO b) -- ^ Folding action that is applied to results of parallel fold - -> b -- ^ Accumulator for chunks folding - -> Array r ix e -> IO b -ifoldlIO wIds f !initAcc g !tAcc !arr = do - let !sz = size arr - results <- - divideWork wIds sz $ \ !scheduler !chunkLength !totalLength !slackStart -> do - loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start -> - scheduleWork scheduler $ - iterLinearM sz start (start + chunkLength) 1 (<) initAcc $ \ !i ix !acc -> - f acc ix (unsafeLinearIndex arr i) - when (slackStart < totalLength) $ - scheduleWork scheduler $ - iterLinearM sz slackStart totalLength 1 (<) initAcc $ \ !i ix !acc -> - f acc ix (unsafeLinearIndex arr i) - F.foldlM g tAcc results -{-# INLINE ifoldlIO #-} - - --- | Just like `ifoldlP`, but allows you to specify which cores to run --- computation on. -ifoldlOnP :: Source r ix e => - [Int] -> (a -> ix -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> IO b -ifoldlOnP wIds f initAcc g = - ifoldlIO wIds (\acc ix -> return . f acc ix) initAcc (\acc -> return . g acc) -{-# INLINE ifoldlOnP #-} - - - -- | /O(n)/ - Left fold with an index aware function, computed in parallel. Just -- like `foldlP`, except that folding function will receive an index of an -- element it is being applied to. -ifoldlP :: Source r ix e => - (a -> ix -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> IO b -ifoldlP = ifoldlOnP [] +-- +-- @since 0.1.0 +ifoldlP :: (MonadIO m, Source r ix e) => + (a -> ix -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> m b +ifoldlP f fAcc g gAcc = + liftIO . ifoldlIO (\acc ix -> return . f acc ix) fAcc (\acc -> return . g acc) gAcc {-# INLINE ifoldlP #-} --- | /O(n)/ - Right fold, computed in parallel. Same as `foldlP`, except directed --- from the last element in the array towards beginning. +-- | /O(n)/ - Right fold, computed with respect to computation strategy. Same as `foldlP`, except +-- directed from the last element in the array towards beginning. -- -- ==== __Examples__ -- --- >>> foldrP (++) [] (:) [] $ makeArray2D (3,4) id --- [(0,0),(0,1),(0,2),(0,3),(1,0),(1,1),(1,2),(1,3),(2,0),(2,1),(2,2),(2,3)] +-- >>> import Data.Massiv.Array +-- >>> foldrP (:) [] (++) [] $ makeArrayR D (ParN 2) (Sz2 2 3) fromIx2 +-- [(0,0),(0,1),(0,2),(1,0),(1,1),(1,2)] +-- >>> foldrP (:) [] (:) [] $ makeArrayR D Seq (Sz1 6) id +-- [[0,1,2,3,4,5]] +-- >>> foldrP (:) [] (:) [] $ makeArrayR D (ParN 3) (Sz1 6) id +-- [[0,1],[2,3],[4,5]] -- -foldrP :: Source r ix e => - (e -> a -> a) -> a -> (a -> b -> b) -> b -> Array r ix e -> IO b -foldrP f = ifoldrP (const f) +-- @since 0.1.0 +foldrP :: (MonadIO m, Source r ix e) => + (e -> a -> a) -> a -> (a -> b -> b) -> b -> Array r ix e -> m b +foldrP f fAcc g gAcc = liftIO . ifoldrP (const f) fAcc g gAcc {-# INLINE foldrP #-} --- | Just like `foldrP`, but allows you to specify which cores to run --- computation on. --- --- ==== __Examples__ --- --- Number of wokers dictate the result structure: --- --- >>> foldrOnP [1,2,3] (:) [] (:) [] $ makeArray1D 9 id --- [[0,1,2],[3,4,5],[6,7,8]] --- >>> foldrOnP [1,2,3] (:) [] (:) [] $ makeArray1D 10 id --- [[0,1,2],[3,4,5],[6,7,8],[9]] --- >>> foldrOnP [1,2,3] (:) [] (:) [] $ makeArray1D 12 id --- [[0,1,2,3],[4,5,6,7],[8,9,10,11]] --- --- But most of the time that structure is of no importance: --- --- >>> foldrOnP [1,2,3] (++) [] (:) [] $ makeArray1D 10 id --- [0,1,2,3,4,5,6,7,8,9] --- --- Same as `foldlOnP`, order is guaranteed to be consecutive and in proper direction: --- --- >>> fmap snd $ foldrOnP [1,2,3] (\x (i, acc) -> (i + 1, (i, x):acc)) (1, []) (:) [] $ makeArray1D 11 id --- [(4,[0,1,2]),(3,[3,4,5]),(2,[6,7,8]),(1,[9,10])] --- >>> fmap (P.zip [4,3..]) <$> foldrOnP [1,2,3] (:) [] (:) [] $ makeArray1D 11 id --- [(4,[0,1,2]),(3,[3,4,5]),(2,[6,7,8]),(1,[9,10])] +-- | /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. -- -foldrOnP :: Source r ix e => - [Int] -> (e -> a -> a) -> a -> (a -> b -> b) -> b -> Array r ix e -> IO b -foldrOnP wIds f = ifoldrOnP wIds (const f) -{-# INLINE foldrOnP #-} +-- @since 0.1.0 +ifoldrP :: + (MonadIO m, Source r ix e) + => (ix -> e -> a -> a) + -> a + -> (a -> b -> b) + -> b + -> Array r ix e + -> m b +ifoldrP f fAcc g gAcc = liftIO . ifoldrIO (\ix e -> pure . f ix e) fAcc (\e -> pure . g e) gAcc +{-# INLINE ifoldrP #-} --- | Parallel right fold. Differs from `ifoldrP` in that it accepts `IO` actions instead of the --- usual pure functions as arguments. -ifoldrIO :: Source r ix e => - [Int] -> (ix -> e -> a -> IO a) -> a -> (a -> b -> IO b) -> b -> Array r ix e -> IO b -ifoldrIO wIds f !initAcc g !tAcc !arr = do - let !sz = size arr - results <- - divideWork wIds sz $ \ !scheduler !chunkLength !totalLength !slackStart -> do - when (slackStart < totalLength) $ - scheduleWork scheduler $ - iterLinearM sz (totalLength - 1) slackStart (-1) (>=) initAcc $ \ !i ix !acc -> - f ix (unsafeLinearIndex arr i) acc - loopM_ slackStart (> 0) (subtract chunkLength) $ \ !start -> - scheduleWork scheduler $ - iterLinearM sz (start - 1) (start - chunkLength) (-1) (>=) initAcc $ \ !i ix !acc -> - f ix (unsafeLinearIndex arr i) acc - F.foldlM (flip g) tAcc results -{-# INLINE ifoldrIO #-} +-- | This folding function breaks referential transparency on some functions +-- @f@, therefore it is kept here for internal use only. +foldlInternal :: Source r ix e => (a -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> b +foldlInternal g initAcc f resAcc = unsafePerformIO . foldlP g initAcc f resAcc +{-# INLINE foldlInternal #-} --- | /O(n)/ - Right fold with an index aware function, computed in parallel. --- Same as `ifoldlP`, except directed from the last element in the array towards --- beginning. -ifoldrOnP :: Source r ix e => - [Int] -> (ix -> e -> a -> a) -> a -> (a -> b -> b) -> b -> Array r ix e -> IO b -ifoldrOnP wIds f !initAcc g = - ifoldrIO wIds (\ix e -> return . f ix e) initAcc (\e -> return . g e) -{-# INLINE ifoldrOnP #-} +ifoldlInternal :: Source r ix 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 #-} --- | Just like `ifoldrOnP`, but allows you to specify which cores to run computation on. -ifoldrP :: Source r ix e => - (ix -> e -> a -> a) -> a -> (a -> b -> b) -> b -> Array r ix e -> IO b -ifoldrP = ifoldrOnP [] -{-# INLINE ifoldrP #-} +-- | Similar to `ifoldlP`, except that folding functions themselves do live in IO +-- +-- @since 0.1.0 +ifoldlIO :: + (MonadUnliftIO m, Source r ix e) + => (a -> ix -> e -> m a) -- ^ Index aware folding IO action + -> a -- ^ Accumulator + -> (b -> a -> m b) -- ^ Folding action that is applied to the results of a parallel fold + -> b -- ^ Accumulator for chunks folding + -> Array r ix e + -> m b +ifoldlIO f !initAcc g !tAcc !arr = do + let !sz = size arr + !totalLength = totalElem sz + results <- + withScheduler (getComp arr) $ \scheduler -> + splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do + loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start -> + scheduleWork scheduler $ + iterLinearM sz start (start + chunkLength) 1 (<) initAcc $ \ !i ix !acc -> + f acc ix (unsafeLinearIndex arr i) + when (slackStart < totalLength) $ + scheduleWork scheduler $ + iterLinearM sz slackStart totalLength 1 (<) initAcc $ \ !i ix !acc -> + f acc ix (unsafeLinearIndex arr i) + F.foldlM g tAcc results +{-# INLINE ifoldlIO #-} --- | This folding function breaks referential transparency on some functions --- @f@, therefore it is kept here for internal use only. -foldlInternal :: Source r ix e => - (a -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> b -foldlInternal g initAcc f resAcc = \ arr -> - case getComp arr of - Seq -> f resAcc (foldlS g initAcc arr) - ParOn wIds -> unsafePerformIO $ foldlOnP wIds g initAcc f resAcc arr -{-# INLINE foldlInternal #-} + +-- | Similar to `ifoldrP`, except that folding functions themselves do live in IO +-- +-- @since 0.1.0 +ifoldrIO :: (MonadUnliftIO m, Source r ix e) => + (ix -> e -> a -> m a) -> a -> (a -> b -> m b) -> b -> Array r ix e -> m b +ifoldrIO f !initAcc g !tAcc !arr = do + let !sz = size arr + !totalLength = totalElem sz + results <- + withScheduler (getComp arr) $ \ scheduler -> + splitLinearly (numWorkers scheduler) totalLength $ \ chunkLength slackStart -> do + when (slackStart < totalLength) $ + scheduleWork scheduler $ + iterLinearM sz (totalLength - 1) slackStart (-1) (>=) initAcc $ \ !i ix !acc -> + f ix (unsafeLinearIndex arr i) acc + loopM_ slackStart (> 0) (subtract chunkLength) $ \ !start -> + scheduleWork scheduler $ + iterLinearM sz (start - 1) (start - chunkLength) (-1) (>=) initAcc $ \ !i ix !acc -> + f ix (unsafeLinearIndex arr i) acc + F.foldlM (flip g) tAcc results +{-# INLINE ifoldrIO #-} 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,10 +1,10 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE ScopedTypeVariables #-} -- | -- Module : Data.Massiv.Array.Ops.Map --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental @@ -14,11 +14,21 @@ ( map , imap -- ** Traversing + -- *** Applicative , traverseA + , traverseA_ , itraverseA + , itraverseA_ , traverseAR , itraverseAR - -- ** Monadic + , sequenceA + , sequenceA_ + -- *** PrimMonad + , traversePrim + , itraversePrim + , traversePrimR + , itraversePrimR + -- ** Monadic mapping -- *** Sequential , mapM , mapMR @@ -41,8 +51,6 @@ , forIO_ , iforIO , iforIO_ - , mapP_ - , imapP_ -- ** Zipping , zip , zip3 @@ -53,24 +61,24 @@ , izipWith , izipWith3 , liftArray2 + -- *** Applicative + , zipWithA + , izipWithA + , zipWith3A + , izipWith3A ) where - -import Control.Monad (void, when) -import Control.Monad.ST (runST) -import Data.Foldable (foldlM) -import Data.Massiv.Array.Delayed.Internal -import Data.Massiv.Array.Mutable -import Data.Massiv.Array.Ops.Fold.Internal (foldrFB) -import Data.Massiv.Core.Common -import Data.Massiv.Core.Scheduler -import Data.Monoid ((<>)) -import GHC.Base (build) -import Prelude hiding (map, mapM, mapM_, - traverse, unzip, unzip3, - zip, zip3, zipWith, - zipWith3) -import qualified Prelude as Prelude (traverse) +import Control.Monad (void) +import Control.Monad.Primitive (PrimMonad) +import Control.Scheduler +import Data.Coerce +import Data.Massiv.Array.Delayed.Pull +import Data.Massiv.Array.Mutable +import Data.Massiv.Array.Ops.Construct (makeArrayA) +import Data.Massiv.Core.Common +import Data.Massiv.Core.Index.Internal (Sz(..)) +import Prelude hiding (map, mapM, mapM_, sequenceA, traverse, unzip, unzip3, + zip, zip3, zipWith, zipWith3) -------------------------------------------------------------------------------- -- map ------------------------------------------------------------------------- @@ -129,7 +137,9 @@ izipWith :: (Source r1 ix e1, Source r2 ix e2) => (ix -> e1 -> e2 -> e) -> Array r1 ix e1 -> Array r2 ix e2 -> Array D ix e izipWith f arr1 arr2 = - DArray (getComp arr1 <> getComp arr2) (liftIndex2 min (size arr1) (size arr2)) $ \ !ix -> + DArray + (getComp arr1 <> getComp arr2) + (SafeSz (liftIndex2 min (coerce (size arr1)) (coerce (size arr2)))) $ \ !ix -> f ix (unsafeIndex arr1 ix) (unsafeIndex arr2 ix) {-# INLINE izipWith #-} @@ -152,55 +162,148 @@ izipWith3 f arr1 arr2 arr3 = DArray (getComp arr1 <> getComp arr2 <> getComp arr3) - (liftIndex2 min (liftIndex2 min (size arr1) (size arr2)) (size arr3)) $ \ !ix -> + (SafeSz + (liftIndex2 + min + (liftIndex2 min (coerce (size arr1)) (coerce (size arr2))) + (coerce (size arr3)))) $ \ !ix -> f ix (unsafeIndex arr1 ix) (unsafeIndex arr2 ix) (unsafeIndex arr3 ix) {-# INLINE izipWith3 #-} + +-- | Similar to `zipWith`, except does it sequentiall and using the `Applicative`. Note that +-- resulting array has Mutable representation. +-- +-- @since 0.3.0 +zipWithA :: + (Source r1 ix e1, Source r2 ix e2, Applicative f, Mutable r ix e) + => (e1 -> e2 -> f e) + -> Array r1 ix e1 + -> Array r2 ix e2 + -> f (Array r ix e) +zipWithA f = izipWithA (const f) +{-# INLINE zipWithA #-} + +-- | Similar to `zipWith`, except does it sequentiall and using the `Applicative`. Note that +-- resulting array has Mutable representation. +-- +-- @since 0.3.0 +izipWithA :: + (Source r1 ix e1, Source r2 ix e2, Applicative f, Mutable r ix e) + => (ix -> e1 -> e2 -> f e) + -> Array r1 ix e1 + -> Array r2 ix e2 + -> f (Array r ix e) +izipWithA f arr1 arr2 = + makeArrayA + (getComp arr1 <> getComp arr2) + (SafeSz (liftIndex2 min (coerce (size arr1)) (coerce (size arr2)))) $ \ !ix -> + f ix (unsafeIndex arr1 ix) (unsafeIndex arr2 ix) +{-# INLINE izipWithA #-} + +-- | Same as `zipWithA`, but for three arrays. +-- +-- @since 0.3.0 +zipWith3A :: + (Source r1 ix e1, Source r2 ix e2, Source r3 ix e3, Applicative f, Mutable r ix e) + => (e1 -> e2 -> e3 -> f e) + -> Array r1 ix e1 + -> Array r2 ix e2 + -> Array r3 ix e3 + -> f (Array r ix e) +zipWith3A f = izipWith3A (const f) +{-# INLINE zipWith3A #-} + +-- | Same as `izipWithA`, but for three arrays. +-- +-- @since 0.3.0 +izipWith3A :: + (Source r1 ix e1, Source r2 ix e2, Source r3 ix e3, Applicative f, Mutable r ix e) + => (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 = + makeArrayA (getComp arr1 <> getComp arr2 <> getComp arr3) 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)) +{-# INLINE izipWith3A #-} + + -------------------------------------------------------------------------------- -- traverse -------------------------------------------------------------------- -------------------------------------------------------------------------------- -- | Traverse with an `Applicative` action over an array sequentially. -- +-- /Note/ - using `traversePrim` will always be faster, althought not always possible. +-- -- @since 0.2.6 -- traverseA :: - (Source r' ix a, Mutable r ix b, Applicative f) - => (a -> f b) + (Source r' ix a, Mutable r ix e, Applicative f) + => (a -> f e) -> Array r' ix a - -> f (Array r ix b) -traverseA f arr = loadList <$> Prelude.traverse f (build (\c n -> foldrFB c n arr)) - where - loadList xs = - runST $ do - marr <- unsafeNew (size arr) - _ <- foldlM (\i e -> unsafeLinearWrite marr i e >> return (i + 1)) 0 xs - unsafeFreeze (getComp arr) marr - {-# INLINE loadList #-} + -> f (Array r ix e) +traverseA f arr = makeArrayA (getComp arr) (size arr) (f . unsafeIndex arr) {-# INLINE traverseA #-} +-- | Traverse sequentially over a source array, while discarding the result. +-- +-- @since 0.3.0 +-- +traverseA_ :: (Source r ix a, Applicative f) => (a -> f e) -> Array r ix a -> f () +traverseA_ f arr = loopA_ 0 (< totalElem (size arr)) (+ 1) (f . unsafeLinearIndex arr) +{-# INLINE traverseA_ #-} + +-- | Sequence actions in a source array. +-- +-- @since 0.3.0 +-- +sequenceA :: + (Source r' ix (f e), Mutable r ix e, Applicative f) => Array r' ix (f e) -> f (Array r ix e) +sequenceA = traverseA id +{-# INLINE sequenceA #-} + +-- | Sequence actions in a source array, while discarding the result. +-- +-- @since 0.3.0 +-- +sequenceA_ :: (Source r ix (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 :: - (Source r' ix a, Mutable r ix b, Applicative f) - => (ix -> a -> f b) + (Source r' ix a, Mutable r ix e, Applicative f) + => (ix -> a -> f e) -> Array r' ix a - -> f (Array r ix b) -itraverseA f arr = - fmap loadList $ Prelude.traverse (uncurry f) $ build (\c n -> foldrFB c n (zipWithIndex arr)) - where - loadList xs = - runST $ do - marr <- unsafeNew (size arr) - _ <- foldlM (\i e -> unsafeLinearWrite marr i e >> return (i + 1)) 0 xs - unsafeFreeze (getComp arr) marr - {-# INLINE loadList #-} + -> f (Array r ix e) +itraverseA f arr = makeArrayA (getComp arr) (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_ :: (Source r ix a, Applicative f) => (ix -> a -> f e) -> Array r ix a -> f () +itraverseA_ f arr = + loopA_ 0 (< totalElem sz) (+ 1) (\ !i -> f (fromLinearIndex sz i) (unsafeLinearIndex arr i)) + where + sz = size arr +{-# INLINE itraverseA_ #-} + + -- | Same as `traverseA`, except with ability to specify representation. -- -- @since 0.2.6 @@ -227,10 +330,60 @@ itraverseAR _ = itraverseA {-# INLINE itraverseAR #-} -zipWithIndex :: forall r ix e . Source r ix e => Array r ix e -> Array D ix (ix, e) -zipWithIndex arr = zip (makeArray mempty (size arr) id :: Array D ix ix) arr -{-# INLINE zipWithIndex #-} + +-- | Traverse sequentially within `PrimMonad` over an array with an action. +-- +-- @since 0.3.0 +-- +traversePrim :: + (Source r' ix a, Mutable r ix b, PrimMonad m) + => (a -> m b) + -> Array r' ix a + -> m (Array r ix b) +traversePrim f arr = generateArrayS (getComp arr) (size arr) (f . unsafeIndex arr) +{-# INLINE traversePrim #-} + +-- | Same as `traversePrim`, but traverse with index aware action. +-- +-- @since 0.3.0 +-- +itraversePrim :: + (Source r' ix a, Mutable r ix b, PrimMonad m) + => (ix -> a -> m b) + -> Array r' ix a + -> m (Array r ix b) +itraversePrim f arr = generateArrayS (getComp arr) (size arr) (\ !ix -> f ix (unsafeIndex arr ix)) +{-# INLINE itraversePrim #-} + + +-- | Same as `traverseP`, but with ability to specify the desired representation. +-- +-- @since 0.3.0 +-- +traversePrimR :: + (Source r' ix a, Mutable r ix b, PrimMonad m) + => r + -> (a -> m b) + -> Array r' ix a + -> m (Array r ix b) +traversePrimR _ = traversePrim +{-# INLINE traversePrimR #-} + +-- | Same as `itraverseP`, but with ability to specify the desired representation. +-- +-- @since 0.3.0 +-- +itraversePrimR :: + (Source r' ix a, Mutable r ix b, PrimMonad m) + => r + -> (ix -> a -> m b) + -> Array r' ix a + -> m (Array r ix b) +itraversePrimR _ = itraversePrim +{-# INLINE itraversePrimR #-} + + -------------------------------------------------------------------------------- -- mapM ------------------------------------------------------------------------ -------------------------------------------------------------------------------- @@ -238,11 +391,10 @@ -- | Map a monadic action over an array sequentially. -- -- @since 0.2.6 --- mapM :: - (Source r' ix a, Mutable r ix b, Monad m) - => (a -> m b) - -> Array r' ix a + forall r ix b r' a m. (Source r' ix a, Mutable r ix b, Monad m) + => (a -> m b) -- ^ Mapping action + -> Array r' ix a -- ^ Source array -> m (Array r ix b) mapM = traverseA {-# INLINE mapM #-} @@ -251,9 +403,8 @@ -- | Same as `mapM`, except with ability to specify result representation. -- -- @since 0.2.6 --- mapMR :: - (Source r' ix a, Mutable r ix b, Monad m) + forall r ix b r' a m. (Source r' ix a, Mutable r ix b, Monad m) => r -> (a -> m b) -> Array r' ix a @@ -265,9 +416,8 @@ -- | Same as `mapM` except with arguments flipped. -- -- @since 0.2.6 --- forM :: - (Source r' ix a, Mutable r ix b, Monad m) + forall r ix b r' a m. (Source r' ix a, Mutable r ix b, Monad m) => Array r' ix a -> (a -> m b) -> m (Array r ix b) @@ -278,9 +428,8 @@ -- | Same as `forM`, except with ability to specify result representation. -- -- @since 0.2.6 --- forMR :: - (Source r' ix a, Mutable r ix b, Monad m) + forall r ix b r' a m. (Source r' ix a, Mutable r ix b, Monad m) => r -> Array r' ix a -> (a -> m b) @@ -293,9 +442,8 @@ -- | Map a monadic action over an array sequentially. -- -- @since 0.2.6 --- imapM :: - (Source r' ix a, Mutable r ix b, Monad m) + forall r ix b r' a m. (Source r' ix a, Mutable r ix b, Monad m) => (ix -> a -> m b) -> Array r' ix a -> m (Array r ix b) @@ -306,9 +454,8 @@ -- | Same as `imapM`, except with ability to specify result representation. -- -- @since 0.2.6 --- imapMR :: - (Source r' ix a, Mutable r ix b, Monad m) + forall r ix b r' a m. (Source r' ix a, Mutable r ix b, Monad m) => r -> (ix -> a -> m b) -> Array r' ix a @@ -321,9 +468,8 @@ -- | Same as `forM`, except map an index aware action. -- -- @since 0.2.6 --- iforM :: - (Source r' ix a, Mutable r ix b, Monad m) + forall r ix b r' a m. (Source r' ix a, Mutable r ix b, Monad m) => (ix -> a -> m b) -> Array r' ix a -> m (Array r ix b) @@ -336,7 +482,7 @@ -- @since 0.2.6 -- iforMR :: - (Source r' ix a, Mutable r ix b, Monad m) + forall r ix b r' a m. (Source r' ix a, Mutable r ix b, Monad m) => r -> (ix -> a -> m b) -> Array r' ix a @@ -349,15 +495,17 @@ -- -- ==== __Examples__ -- --- >>> mapM_ print $ rangeStep 10 12 60 +-- >>> import Data.Massiv.Array as A +-- >>> rangeStepM Par (Ix1 10) 12 60 >>= A.mapM_ print -- 10 -- 22 -- 34 -- 46 -- 58 -- +-- @since 0.1.0 mapM_ :: (Source r ix a, Monad m) => (a -> m b) -> Array r ix a -> m () -mapM_ f !arr = iterM_ zeroIndex (size arr) (pureIndex 1) (<) (f . unsafeIndex arr) +mapM_ f !arr = iterM_ zeroIndex (unSz (size arr)) (pureIndex 1) (<) (f . unsafeIndex arr) {-# INLINE mapM_ #-} @@ -368,10 +516,11 @@ -- Here is a common way of iterating N times using a for loop in an imperative -- language with mutation being an obvious side effect: -- --- >>> :m + Data.IORef --- >>> var <- newIORef 0 :: IO (IORef Int) --- >>> forM_ (range 0 1000) $ \ i -> modifyIORef' var (+i) --- >>> readIORef var +-- >>> import Data.Massiv.Array as A +-- >>> import Data.IORef +-- >>> ref <- newIORef 0 :: IO (IORef Int) +-- >>> A.forM_ (range Seq (Ix1 0) 1000) $ \ i -> modifyIORef' ref (+i) +-- >>> readIORef ref -- 499500 -- forM_ :: (Source r ix a, Monad m) => Array r ix a -> (a -> m b) -> m () @@ -391,7 +540,10 @@ -- -- @since 0.2.6 mapIO :: - (Source r' ix a, Mutable r ix b) => (a -> IO b) -> Array r' ix a -> IO (Array r ix b) + forall r ix b r' a m. (Source r' ix a, Mutable r ix b, MonadUnliftIO m, PrimMonad m) + => (a -> m b) + -> Array r' ix a + -> m (Array r ix b) mapIO action = imapIO (const action) {-# INLINE mapIO #-} @@ -399,26 +551,22 @@ -- array, therefore it is faster. Use this instead of `mapIO` when result is irrelevant. -- -- @since 0.2.6 -mapIO_ :: Source r b e => (e -> IO a) -> Array r b e -> IO () +mapIO_ :: (Source r b e, MonadUnliftIO m) => (e -> m a) -> Array r b e -> m () mapIO_ action = imapIO_ (const action) {-# INLINE mapIO_ #-} -- | Same as `mapIO_`, but map an index aware action instead. -- -- @since 0.2.6 -imapIO_ :: Source r ix e => (ix -> e -> IO a) -> Array r ix e -> IO () -imapIO_ action arr = - case getComp arr of - Seq -> imapM_ action arr - ParOn wids -> do - let sz = size arr - withScheduler_ wids $ \scheduler -> - splitLinearlyWith_ - (numWorkers scheduler) - (scheduleWork scheduler) - (totalElem sz) - (unsafeLinearIndex arr) - (\i -> void . action (fromLinearIndex sz i)) +imapIO_ :: (Source r ix e, MonadUnliftIO m) => (ix -> e -> m a) -> Array r ix e -> m () +imapIO_ action arr = do + let sz = size arr + withScheduler_ (getComp arr) $ \scheduler -> + splitLinearlyWith_ + scheduler + (totalElem sz) + (unsafeLinearIndex arr) + (\i -> void . action (fromLinearIndex sz i)) {-# INLINE imapIO_ #-} @@ -426,22 +574,40 @@ -- -- @since 0.2.6 imapIO :: - (Source r' ix a, Mutable r ix b) => (ix -> a -> IO b) -> Array r' ix a -> IO (Array r ix b) -imapIO action arr = generateArrayIO (getComp arr) (size arr) $ \ix -> action ix (unsafeIndex arr ix) + forall r ix b r' a m. (Source r' ix a, Mutable r ix b, MonadUnliftIO m, PrimMonad m) + => (ix -> a -> m b) + -> Array r' ix a + -> m (Array r ix b) +imapIO action arr = generateArray (getComp arr) (size arr) $ \ix -> action ix (unsafeIndex arr ix) {-# INLINE imapIO #-} -- | Same as `mapIO` but with arguments flipped. -- -- @since 0.2.6 forIO :: - (Source r' ix a, Mutable r ix b) => Array r' ix a -> (a -> IO b) -> IO (Array r ix b) + forall r ix b r' a m. (Source r' ix a, Mutable r ix b, MonadUnliftIO m, PrimMonad m) + => Array r' ix a + -> (a -> m b) + -> m (Array r ix b) forIO = flip mapIO {-# INLINE forIO #-} -- | Same as `mapIO_` but with arguments flipped. -- +-- ==== __Example__ +-- +-- This is the same example as in `forM_`, with important difference that accumulator `ref` will be +-- modified concurrently by as many threads as there are capabilities. +-- +-- >>> import Data.Massiv.Array +-- >>> import Data.IORef +-- >>> ref <- newIORef 0 :: IO (IORef Int) +-- >>> forIO_ (range Par (Ix1 0) 1000) $ \ i -> atomicModifyIORef' ref (\v -> (v+i, ())) +-- >>> readIORef ref +-- 499500 +-- -- @since 0.2.6 -forIO_ :: Source r ix e => Array r ix e -> (e -> IO a) -> IO () +forIO_ :: (Source r ix e, MonadUnliftIO m) => Array r ix e -> (e -> m a) -> m () forIO_ = flip mapIO_ {-# INLINE forIO_ #-} @@ -449,42 +615,16 @@ -- -- @since 0.2.6 iforIO :: - (Source r' ix a, Mutable r ix b) => Array r' ix a -> (ix -> a -> IO b) -> IO (Array r ix b) + forall r ix b r' a m. (Source r' ix a, Mutable r ix b, MonadUnliftIO m, PrimMonad m) + => Array r' ix a + -> (ix -> a -> m b) + -> m (Array r ix b) iforIO = flip imapIO {-# INLINE iforIO #-} -- | Same as `imapIO_` but with arguments flipped. -- -- @since 0.2.6 -iforIO_ :: Source r ix a => Array r ix a -> (ix -> a -> IO b) -> IO () +iforIO_ :: (Source r ix a, MonadUnliftIO m) => Array r ix a -> (ix -> a -> m b) -> m () iforIO_ = flip imapIO_ {-# INLINE iforIO_ #-} - - --- | Map an IO action, over an array in parallel, while discarding the result. -mapP_ :: Source r ix a => (a -> IO b) -> Array r ix a -> IO () -mapP_ f = imapP_ (const f) -{-# INLINE mapP_ #-} -{-# DEPRECATED mapP_ "In favor of 'mapIO_'" #-} - - --- | Map an index aware IO action, over an array in parallel, while --- discarding the result. -imapP_ :: Source r ix a => (ix -> a -> IO b) -> Array r ix a -> IO () -imapP_ f arr = do - let sz = size arr - wIds = - case getComp arr of - ParOn ids -> ids - _ -> [] - divideWork_ wIds sz $ \ !scheduler !chunkLength !totalLength !slackStart -> do - loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start -> - scheduleWork scheduler $ - iterLinearM_ sz start (start + chunkLength) 1 (<) $ \ !i ix -> do - void $ f ix (unsafeLinearIndex arr i) - when (slackStart < totalLength) $ - scheduleWork scheduler $ - iterLinearM_ sz slackStart totalLength 1 (<) $ \ !i ix -> do - void $ f ix (unsafeLinearIndex arr i) -{-# INLINE imapP_ #-} -{-# DEPRECATED imapP_ "In favor of 'imapIO_'" #-} diff --git a/src/Data/Massiv/Array/Ops/Slice.hs b/src/Data/Massiv/Array/Ops/Slice.hs --- a/src/Data/Massiv/Array/Ops/Slice.hs +++ b/src/Data/Massiv/Array/Ops/Slice.hs @@ -1,7 +1,7 @@ {-# LANGUAGE BangPatterns #-} -- | -- Module : Data.Massiv.Array.Ops.Slice --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental @@ -23,8 +23,8 @@ , () ) where -import Control.Monad (guard) -import Data.Massiv.Core.Common +import Control.Monad (unless) +import Data.Massiv.Core.Common infixl 4 !>, !?>, ??>, , , @@ -39,24 +39,25 @@ -- highermost dimension. For example with rank-3 arrays outer slice would be equivalent to getting a -- page: -- --- >>> let arr = makeArrayR U Seq (3 :> 2 :. 4) fromIx3 +-- >>> import Data.Massiv.Array +-- >>> arr = makeArrayR U Seq (Sz (3 :> 2 :. 4)) fromIx3 -- >>> arr --- (Array U Seq (3 :> 2 :. 4) --- [ [ [ (0,0,0),(0,0,1),(0,0,2),(0,0,3) ] --- , [ (0,1,0),(0,1,1),(0,1,2),(0,1,3) ] +-- Array U Seq (Sz (3 :> 2 :. 4)) +-- [ [ [ (0,0,0), (0,0,1), (0,0,2), (0,0,3) ] +-- , [ (0,1,0), (0,1,1), (0,1,2), (0,1,3) ] -- ] --- , [ [ (1,0,0),(1,0,1),(1,0,2),(1,0,3) ] --- , [ (1,1,0),(1,1,1),(1,1,2),(1,1,3) ] +-- , [ [ (1,0,0), (1,0,1), (1,0,2), (1,0,3) ] +-- , [ (1,1,0), (1,1,1), (1,1,2), (1,1,3) ] -- ] --- , [ [ (2,0,0),(2,0,1),(2,0,2),(2,0,3) ] --- , [ (2,1,0),(2,1,1),(2,1,2),(2,1,3) ] +-- , [ [ (2,0,0), (2,0,1), (2,0,2), (2,0,3) ] +-- , [ (2,1,0), (2,1,1), (2,1,2), (2,1,3) ] -- ] --- ]) +-- ] -- >>> arr !> 2 --- (Array M Seq (2 :. 4) --- [ [ (2,0,0),(2,0,1),(2,0,2),(2,0,3) ] --- , [ (2,1,0),(2,1,1),(2,1,2),(2,1,3) ] --- ]) +-- Array M Seq (Sz (2 :. 4)) +-- [ [ (2,0,0), (2,0,1), (2,0,2), (2,0,3) ] +-- , [ (2,1,0), (2,1,1), (2,1,2), (2,1,3) ] +-- ] -- -- There is nothing wrong with chaining, mixing and matching slicing operators, or even using them -- to index arrays: @@ -65,23 +66,26 @@ -- (2,0,3) -- >>> arr !> 2 >> arr !> 2 !> 0 !> 3 == arr ! 2 :> 0 :. 3 +-- >>> (arr !> 2 !> 0 !> 3) == (arr ! 2 :> 0 :. 3) -- True -- +-- +-- @since 0.1.0 (!>) :: OuterSlice r ix e => Array r ix e -> Int -> Elt r ix e -(!>) !arr !ix = - case arr !?> ix of - Just res -> res - Nothing -> errorIx "(!>)" (outerLength arr) ix +(!>) !arr !ix = either throw id (arr !?> ix) {-# INLINE (!>) #-} -- | /O(1)/ - Just like `!>` slices the array from the outside, but returns -- `Nothing` when index is out of bounds. -(!?>) :: OuterSlice r ix e => Array r ix e -> Int -> Maybe (Elt r ix e) +-- +-- @since 0.1.0 +(!?>) :: (MonadThrow m, OuterSlice r ix e) => Array r ix e -> Int -> m (Elt r ix e) (!?>) !arr !i - | isSafeIndex (outerLength arr) i = Just $ unsafeOuterSlice arr i - | otherwise = Nothing + | isSafeIndex sz i = pure $ unsafeOuterSlice arr i + | otherwise = throwM $ IndexOutOfBoundsException sz i + where + !sz = fst (unconsSz (size arr)) {-# INLINE (!?>) #-} @@ -90,79 +94,83 @@ -- -- ===__Examples__ -- --- >>> let arr = makeArrayR U Seq (3 :> 2 :. 4) fromIx3 --- >>> arr !?> 2 ??> 0 ??> 3 +-- >>> import Data.Massiv.Array +-- >>> arr = makeArrayR U Seq (Sz (3 :> 2 :. 4)) fromIx3 +-- >>> arr !?> 2 ??> 0 ??> 3 :: Maybe Ix3T -- Just (2,0,3) --- >>> arr !?> 2 ??> 0 ??> -1 --- Nothing --- >>> arr !?> -2 ??> 0 ?? 1 +-- >>> arr !?> 2 ??> 0 ??> -1 :: Maybe Ix3T -- Nothing +-- >>> arr !?> 2 ??> -10 ?? 1 +-- *** Exception: IndexOutOfBoundsException: -10 not safe for (Sz1 2) -- -(??>) :: OuterSlice r ix e => Maybe (Array r ix e) -> Int -> Maybe (Elt r ix e) -(??>) Nothing _ = Nothing -(??>) (Just arr) !ix = arr !?> ix +-- @since 0.1.0 +(??>) :: (MonadThrow m, OuterSlice r ix e) => m (Array r ix e) -> Int -> m (Elt r ix e) +(??>) marr !ix = marr >>= (!?> ix) {-# INLINE (??>) #-} -- | /O(1)/ - Safe slice from the inside -( Array r ix e -> Int -> Maybe (Elt r ix e) +-- +-- @since 0.1.0 +( Array r ix e -> Int -> m (Elt r ix e) (`) slice an array from an opposite direction. +-- +-- @since 0.1.0 ( Array r ix e -> Int -> Elt r ix e ( res - Nothing -> errorIx "( res + Left exc -> throw exc {-# INLINE ( Maybe (Array r ix e) -> Int -> Maybe (Elt r ix e) -( m (Array r ix e) -> Int -> m (Elt r ix e) +(>= (`), but fails gracefully with a `Nothing`, instead of an error -() :: Slice r ix e => Array r ix e -> (Dim, Int) -> Maybe (Elt r ix e) -() !arr !(dim, i) = do - m <- getDim (size arr) dim - guard $ isSafeIndex m i - start <- setDim zeroIndex dim i - cutSz <- setDim (size arr) dim 1 +-- +-- @since 0.1.0 +() :: (MonadThrow m, Slice r ix e) => Array r ix e -> (Dim, Int) -> m (Elt r ix e) +() !arr (dim, i) = do + (m, szl) <- pullOutSzM (size arr) dim + unless (isSafeIndex m i) $ throwM $ IndexOutOfBoundsException m i + start <- setDimM zeroIndex dim i + cutSz <- insertSzM szl dim oneSz unsafeSlice arr start cutSz dim {-# INLINE () #-} --- | /O(1)/ - Slices the array in any available dimension. Throws an error when --- index is out of bounds or dimensions is invalid. --- -- prop> arr !> i == arr (dimensions (size arr), i) -- prop> arr (1,i) -- +-- | /O(1)/ - Slices the array in any available dimension. Throws an error when +-- index is out of bounds or dimensions is invalid. +-- +-- @since 0.1.0 () :: Slice r ix e => Array r ix e -> (Dim, Int) -> Elt r ix e -() !arr !(dim, i) = - case arr (dim, i) of - Just res -> res - Nothing -> - let arrDims = dimensions (size arr) - in if dim < 1 || dim > arrDims - then error $ - "(): Invalid dimension: " ++ - show dim ++ " for Array of dimensions: " ++ show arrDims - else errorIx "()" (size arr) (dim, i) +() !arr !dix = + case arr dix of + Right res -> res + Left exc -> throw exc {-# INLINE () #-} -- | /O(1)/ - Safe slicing continuation from within. -() :: Slice r ix e => Maybe (Array r ix e) -> (Dim, Int) -> Maybe (Elt r ix e) -() Nothing _ = Nothing -() (Just arr) !ix = arr ix +-- +-- @since 0.1.0 +() :: (MonadThrow m, Slice r ix e) => m (Array r ix e) -> (Dim, Int) -> m (Elt r ix e) +() !marr !ix = marr >>= ( ix) {-# INLINE () #-} diff --git a/src/Data/Massiv/Array/Ops/Transform.hs b/src/Data/Massiv/Array/Ops/Transform.hs --- a/src/Data/Massiv/Array/Ops/Transform.hs +++ b/src/Data/Massiv/Array/Ops/Transform.hs @@ -1,10 +1,13 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE ExplicitForAll #-} +{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE TypeFamilies #-} +{-# OPTIONS_GHC -fno-warn-redundant-constraints #-} -- | -- Module : Data.Massiv.Array.Ops.Transform --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental @@ -16,36 +19,91 @@ , transposeInner , transposeOuter -- ** Backpermute + , backpermuteM + , backpermute' , backpermute -- ** Resize - , resize + , resizeM , resize' + , resize -- ** Extract + , extractM , extract , extract' + , extractFromToM , extractFromTo , extractFromTo' -- ** Append/Split + , cons + , unconsM + , snoc + , unsnocM + , appendM , append , append' + , concatM + , concat' + , splitAtM , splitAt , splitAt' - -- * Traverse + -- ** Upsample/Downsample + , upsample + , downsample + -- ** Transform + , transformM + , transform' + , transform2M + , transform2' + -- ** Traverse (deprecated) , traverse , traverse2 ) where -import Control.Monad (guard) -import Data.Massiv.Array.Delayed.Internal -import Data.Massiv.Array.Ops.Construct -import Data.Massiv.Core.Common -import Data.Maybe (fromMaybe) -import Prelude hiding (splitAt, traverse) +import Control.Scheduler (traverse_) +import Control.Monad as M (foldM_, unless) +import Data.Bifunctor (bimap) +import Data.Foldable as F (foldl', foldrM, toList) +import qualified Data.List as L (uncons) +import Data.Massiv.Array.Delayed.Pull +import Data.Massiv.Array.Delayed.Push +import Data.Massiv.Array.Mutable +import Data.Massiv.Array.Ops.Construct +import Data.Massiv.Core.Common +import Data.Massiv.Core.Index.Internal (Sz(SafeSz)) +import Prelude as P hiding (concat, splitAt, traverse) -- | Extract a sub-array from within a larger source array. Array that is being extracted must be +-- fully encapsulated in a source array, otherwise `SizeSubregionException` will be thrown. +extractM :: (MonadThrow m, Extract r ix e) + => ix -- ^ Starting index + -> Sz ix -- ^ Size of the resulting array + -> Array r ix e -- ^ Source array + -> m (Array (EltRepr r ix) ix e) +extractM !sIx !newSz !arr + | isSafeIndex sz1 sIx && isSafeIndex eIx1 sIx && isSafeIndex sz1 eIx = + pure $ unsafeExtract sIx newSz arr + | otherwise = throwM $ SizeSubregionException (size arr) sIx newSz + where + sz1 = Sz (liftIndex (+1) (unSz (size arr))) + eIx1 = Sz (liftIndex (+1) eIx) + eIx = liftIndex2 (+) sIx $ unSz newSz +{-# INLINE extractM #-} + +-- | Same as `extract`, but will throw an error if supplied dimensions are incorrect. +-- +-- @since 0.1.0 +extract' :: Extract r ix e + => ix -- ^ Starting index + -> Sz ix -- ^ Size of the resulting array + -> Array r ix e -- ^ Source array + -> Array (EltRepr r ix) ix e +extract' sIx newSz = either throw id . extractM sIx newSz +{-# INLINE extract' #-} + +-- | Extract a sub-array from within a larger source array. Array that is being extracted must be -- fully encapsulated in a source array, otherwise `Nothing` is returned, -extract :: Size r ix e +extract :: Extract r ix e => ix -- ^ Starting index -> Sz ix -- ^ Size of the resulting array -> Array r ix e -- ^ Source array @@ -55,86 +113,97 @@ Just $ unsafeExtract sIx newSz arr | otherwise = Nothing where - sz1 = liftIndex (+1) (size arr) - eIx1 = liftIndex (+1) eIx - eIx = liftIndex2 (+) sIx newSz + sz1 = Sz (liftIndex (+1) (unSz (size arr))) + eIx1 = Sz (liftIndex (+1) eIx) + eIx = liftIndex2 (+) sIx $ unSz newSz {-# INLINE extract #-} +{-# DEPRECATED extract "In favor of a more general `extractM`" #-} --- | Same as `extract`, but will throw an error if supplied dimensions are incorrect. -extract' :: Size r ix e - => ix -- ^ Starting index - -> Sz ix -- ^ Size of the resulting array - -> Array r ix e -- ^ Source array - -> Array (EltRepr r ix) ix e -extract' !sIx !newSz !arr = - case extract sIx newSz arr of - Just arr' -> arr' - Nothing -> - error $ - "Data.Massiv.Array.extract': Cannot extract an array of size " ++ - show newSz ++ - " starting at " ++ show sIx ++ " from within an array of size: " ++ show (size arr) -{-# INLINE extract' #-} +-- | Similar to `extractM`, except it takes starting and ending index. Result array will not include +-- the ending index. +-- +-- @since 0.3.0 +extractFromToM :: (MonadThrow m, Extract r ix e) => + ix -- ^ Starting index + -> ix -- ^ Index up to which elements should be extracted. + -> Array r ix e -- ^ Source array. + -> m (Array (EltRepr r ix) ix e) +extractFromToM sIx eIx = extractM sIx (Sz (liftIndex2 (-) eIx sIx)) +{-# INLINE extractFromToM #-} + -- | Similar to `extract`, except it takes starting and ending index. Result array will not include -- the ending index. -extractFromTo :: Size r ix e => +extractFromTo :: Extract r ix e => ix -- ^ Starting index -> ix -- ^ Index up to which elmenets should be extracted. -> Array r ix e -- ^ Source array. -> Maybe (Array (EltRepr r ix) ix e) -extractFromTo sIx eIx = extract sIx $ liftIndex2 (-) eIx sIx +extractFromTo sIx eIx = extract sIx $ Sz (liftIndex2 (-) eIx sIx) {-# INLINE extractFromTo #-} +{-# DEPRECATED extractFromTo "In favor of a more general `extractFromToM`" #-} -- | Same as `extractFromTo`, but throws an error on invalid indices. -- -- @since 0.2.4 -extractFromTo' :: Size r ix e => +extractFromTo' :: Extract r ix e => ix -- ^ Starting index -> ix -- ^ Index up to which elmenets should be extracted. -> Array r ix e -- ^ Source array. -> Array (EltRepr r ix) ix e -extractFromTo' sIx eIx = extract' sIx $ liftIndex2 (-) eIx sIx +extractFromTo' sIx eIx = extract' sIx $ Sz (liftIndex2 (-) eIx sIx) {-# INLINE extractFromTo' #-} -- | /O(1)/ - Changes the shape of an array. Returns `Nothing` if total -- number of elements does not match the source array. -resize :: (Index ix', Size r ix e) => Sz ix' -> Array r ix e -> Maybe (Array r ix' e) +resize :: + (Index ix', Load r ix e, Resize r ix) => Sz ix' -> Array r ix e -> Maybe (Array r ix' e) resize !sz !arr | totalElem sz == totalElem (size arr) = Just $ unsafeResize sz arr | otherwise = Nothing {-# INLINE resize #-} +{-# DEPRECATED resize "In favor of a more general `resizeM`" #-} --- | Same as `resize`, but will throw an error if supplied dimensions are incorrect. -resize' :: (Index ix', Size r ix e) => Sz ix' -> Array r ix e -> Array r ix' e -resize' !sz !arr = - maybe - (error $ - "Total number of elements do not match: " ++ - show sz ++ " vs " ++ show (size arr)) - id $ - resize sz arr +-- | /O(1)/ - Changes the shape of an array. Returns `Nothing` if total +-- number of elements does not match the source array. +-- +-- @since 0.3.0 +resizeM :: + (MonadThrow m, Index ix', Load r ix e, Resize r ix) + => Sz ix' + -> Array r ix e + -> m (Array r ix' e) +resizeM sz arr = guardNumberOfElements (size arr) sz >> pure (unsafeResize sz arr) +{-# INLINE resizeM #-} + +-- | Same as `resizeM`, but will throw an error if supplied dimensions are incorrect. +-- +-- @since 0.1.0 +resize' :: (Index ix', Load r ix e, Resize r ix) => Sz ix' -> Array r ix e -> Array r ix' e +resize' sz = either throw id . resizeM sz {-# INLINE resize' #-} -- | Transpose a 2-dimensional array -- --- ===__Examples__ +-- ==== __Examples__ -- --- >>> let arr = makeArrayR U Seq (2 :. 3) (toLinearIndex (2 :. 3)) +-- >>> import Data.Massiv.Array +-- >>> arr = makeArrayLinearR D Seq (Sz (2 :. 3)) id -- >>> arr --- (ArrayU Seq (2 :. 3) --- [ [ 0,1,2 ] --- , [ 3,4,5 ] --- ]) +-- Array D Seq (Sz (2 :. 3)) +-- [ [ 0, 1, 2 ] +-- , [ 3, 4, 5 ] +-- ] -- >>> transpose arr --- (Array D Seq (3 :. 2) --- [ [ 0,3 ] --- , [ 1,4 ] --- , [ 2,5 ] --- ]) +-- Array D Seq (Sz (3 :. 2)) +-- [ [ 0, 3 ] +-- , [ 1, 4 ] +-- , [ 2, 5 ] +-- ] -- +-- @since 0.1.0 transpose :: Source r Ix2 e => Array r Ix2 e -> Array D Ix2 e transpose = transposeInner {-# INLINE [1] transpose #-} @@ -150,226 +219,462 @@ -- -- ===__Examples__ -- --- >>> let arr = makeArrayR U Seq (2 :> 3 :. 4) fromIx3 +-- >>> import Data.Massiv.Array +-- >>> arr = makeArrayLinearR U Seq (Sz (2 :> 3 :. 4)) id -- >>> arr --- (Array U Seq (2 :> 3 :. 4) --- [ [ [ (0,0,0),(0,0,1),(0,0,2),(0,0,3) ] --- , [ (0,1,0),(0,1,1),(0,1,2),(0,1,3) ] --- , [ (0,2,0),(0,2,1),(0,2,2),(0,2,3) ] +-- Array U Seq (Sz (2 :> 3 :. 4)) +-- [ [ [ 0, 1, 2, 3 ] +-- , [ 4, 5, 6, 7 ] +-- , [ 8, 9, 10, 11 ] -- ] --- , [ [ (1,0,0),(1,0,1),(1,0,2),(1,0,3) ] --- , [ (1,1,0),(1,1,1),(1,1,2),(1,1,3) ] --- , [ (1,2,0),(1,2,1),(1,2,2),(1,2,3) ] +-- , [ [ 12, 13, 14, 15 ] +-- , [ 16, 17, 18, 19 ] +-- , [ 20, 21, 22, 23 ] -- ] --- ]) +-- ] -- >>> transposeInner arr --- (Array D Seq (3 :> 2 :. 4) --- [ [ [ (0,0,0),(0,0,1),(0,0,2),(0,0,3) ] --- , [ (1,0,0),(1,0,1),(1,0,2),(1,0,3) ] +-- Array D Seq (Sz (3 :> 2 :. 4)) +-- [ [ [ 0, 1, 2, 3 ] +-- , [ 12, 13, 14, 15 ] -- ] --- , [ [ (0,1,0),(0,1,1),(0,1,2),(0,1,3) ] --- , [ (1,1,0),(1,1,1),(1,1,2),(1,1,3) ] +-- , [ [ 4, 5, 6, 7 ] +-- , [ 16, 17, 18, 19 ] -- ] --- , [ [ (0,2,0),(0,2,1),(0,2,2),(0,2,3) ] --- , [ (1,2,0),(1,2,1),(1,2,2),(1,2,3) ] +-- , [ [ 8, 9, 10, 11 ] +-- , [ 20, 21, 22, 23 ] -- ] --- ]) +-- ] -- +-- @since 0.1.0 transposeInner :: (Index (Lower ix), Source r' ix e) => Array r' ix e -> Array D ix e -transposeInner !arr = unsafeMakeArray (getComp arr) (transInner (size arr)) newVal +transposeInner !arr = makeArray (getComp arr) newsz newVal where transInner !ix = - fromMaybe (errorImpossible "transposeInner" ix) $ do - n <- getDim ix (dimensions ix) - m <- getDim ix (dimensions ix - 1) - ix' <- setDim ix (dimensions ix) m - setDim ix' (dimensions ix - 1) n + either throwImpossible id $ do + n <- getDimM ix dix + m <- getDimM ix (dix - 1) + ix' <- setDimM ix dix m + setDimM ix' (dix - 1) n {-# INLINE transInner #-} newVal = unsafeIndex arr . transInner {-# INLINE newVal #-} + !newsz = Sz (transInner (unSz (size arr))) + !dix = dimensions newsz {-# INLINE [1] transposeInner #-} -- | Transpose outer two dimensions of at least rank-2 array. -- --- ===__Examples__ +-- ====__Examples__ -- --- >>> let arr = makeArrayR U Seq (2 :> 3 :. 4) fromIx3 +-- >>> import Data.Massiv.Array +-- >>> :set -XTypeApplications +-- >>> arr = makeArrayLinear @U Seq (Sz (2 :> 3 :. 4)) id -- >>> arr --- (Array U Seq (2 :> 3 :. 4) --- [ [ [ (0,0,0),(0,0,1),(0,0,2),(0,0,3) ] --- , [ (0,1,0),(0,1,1),(0,1,2),(0,1,3) ] --- , [ (0,2,0),(0,2,1),(0,2,2),(0,2,3) ] +-- Array U Seq (Sz (2 :> 3 :. 4)) +-- [ [ [ 0, 1, 2, 3 ] +-- , [ 4, 5, 6, 7 ] +-- , [ 8, 9, 10, 11 ] -- ] --- , [ [ (1,0,0),(1,0,1),(1,0,2),(1,0,3) ] --- , [ (1,1,0),(1,1,1),(1,1,2),(1,1,3) ] --- , [ (1,2,0),(1,2,1),(1,2,2),(1,2,3) ] +-- , [ [ 12, 13, 14, 15 ] +-- , [ 16, 17, 18, 19 ] +-- , [ 20, 21, 22, 23 ] -- ] --- ]) +-- ] -- >>> transposeOuter arr --- (Array D Seq (2 :> 4 :. 3) --- [ [ [ (0,0,0),(0,1,0),(0,2,0) ] --- , [ (0,0,1),(0,1,1),(0,2,1) ] --- , [ (0,0,2),(0,1,2),(0,2,2) ] --- , [ (0,0,3),(0,1,3),(0,2,3) ] +-- Array D Seq (Sz (2 :> 4 :. 3)) +-- [ [ [ 0, 4, 8 ] +-- , [ 1, 5, 9 ] +-- , [ 2, 6, 10 ] +-- , [ 3, 7, 11 ] -- ] --- , [ [ (1,0,0),(1,1,0),(1,2,0) ] --- , [ (1,0,1),(1,1,1),(1,2,1) ] --- , [ (1,0,2),(1,1,2),(1,2,2) ] --- , [ (1,0,3),(1,1,3),(1,2,3) ] +-- , [ [ 12, 16, 20 ] +-- , [ 13, 17, 21 ] +-- , [ 14, 18, 22 ] +-- , [ 15, 19, 23 ] -- ] --- ]) +-- ] -- +-- +-- @since 0.1.0 transposeOuter :: (Index (Lower ix), Source r' ix e) => Array r' ix e -> Array D ix e -transposeOuter !arr = unsafeMakeArray (getComp arr) (transOuter (size arr)) newVal +transposeOuter !arr = makeArray (getComp arr) newsz newVal where transOuter !ix = - fromMaybe (errorImpossible "transposeOuter" ix) $ do - n <- getDim ix 1 - m <- getDim ix 2 - ix' <- setDim ix 1 m - setDim ix' 2 n + either throwImpossible id $ do + n <- getDimM ix 1 + m <- getDimM ix 2 + ix' <- setDimM ix 1 m + setDimM ix' 2 n {-# INLINE transOuter #-} newVal = unsafeIndex arr . transOuter {-# INLINE newVal #-} + !newsz = Sz (transOuter (unSz (size arr))) {-# INLINE [1] transposeOuter #-} --- | Rearrange elements of an array into a new one. +-- | Rearrange elements of an array into a new one by using a function that maps indices of the +-- newly created one into the old one. This function can throw `IndexOutOfBoundsException`. -- -- ===__Examples__ -- --- >>> let arr = makeArrayR U Seq (2 :> 3 :. 4) fromIx3 +-- >>> import Data.Massiv.Array +-- >>> :set -XTypeApplications +-- >>> arr = makeArrayLinear @D Seq (Sz (2 :> 3 :. 4)) id -- >>> arr --- (Array U Seq (2 :> 3 :. 4) --- [ [ [ (0,0,0),(0,0,1),(0,0,2),(0,0,3) ] --- , [ (0,1,0),(0,1,1),(0,1,2),(0,1,3) ] --- , [ (0,2,0),(0,2,1),(0,2,2),(0,2,3) ] +-- Array D Seq (Sz (2 :> 3 :. 4)) +-- [ [ [ 0, 1, 2, 3 ] +-- , [ 4, 5, 6, 7 ] +-- , [ 8, 9, 10, 11 ] -- ] --- , [ [ (1,0,0),(1,0,1),(1,0,2),(1,0,3) ] --- , [ (1,1,0),(1,1,1),(1,1,2),(1,1,3) ] --- , [ (1,2,0),(1,2,1),(1,2,2),(1,2,3) ] +-- , [ [ 12, 13, 14, 15 ] +-- , [ 16, 17, 18, 19 ] +-- , [ 20, 21, 22, 23 ] -- ] --- ]) --- >>> backpermute (4 :. 3) (\(i :. j) -> 0 :> j :. i) arr --- (Array D Seq (4 :. 3) --- [ [ (0,0,0),(0,1,0),(0,2,0) ] --- , [ (0,0,1),(0,1,1),(0,2,1) ] --- , [ (0,0,2),(0,1,2),(0,2,2) ] --- , [ (0,0,3),(0,1,3),(0,2,3) ] --- ]) +-- ] +-- >>> backpermuteM @U (Sz (4 :. 2)) (\(i :. j) -> j :> j :. i) arr +-- Array U Seq (Sz (4 :. 2)) +-- [ [ 0, 16 ] +-- , [ 1, 17 ] +-- , [ 2, 18 ] +-- , [ 3, 19 ] +-- ] -- +-- @since 0.3.0 +backpermuteM :: + forall r ix e r' ix' m. + (Mutable r ix e, Source r' ix' e, MonadUnliftIO m, PrimMonad m, MonadThrow m) + => Sz ix -- ^ Size of the result array + -> (ix -> ix') -- ^ A function that maps indices of the new array into the source one. + -> Array r' ix' e -- ^ Source array. + -> m (Array r ix e) +backpermuteM sz ixF !arr = generateArray (getComp arr) sz (evaluateM arr . ixF) +{-# INLINE backpermuteM #-} + +-- | Similar to `backpermuteM`, with few notable differences: +-- +-- * Creates a delayed array, instead of manifest, therefore it can be fused +-- * Respects computation strategy, so it can be parallelized +-- * Throws a runtime `IndexOutOfBoundsException` from pure code. +-- +-- @since 0.3.0 +backpermute' :: (Source r' ix' e, Index ix) => + Sz ix -- ^ Size of the result array + -> (ix -> ix') -- ^ A function that maps indices of the new array into the source one. + -> Array r' ix' e -- ^ Source array. + -> Array D ix e +backpermute' sz ixF !arr = makeArray (getComp arr) sz (evaluate' arr . ixF) +{-# INLINE backpermute' #-} + +-- | See `backpermute'`. +-- +-- @since 0.1.0 backpermute :: (Source r' ix' e, Index ix) => Sz ix -- ^ Size of the result array -> (ix -> ix') -- ^ A function that maps indices of the new array into the source one. -> Array r' ix' e -- ^ Source array. -> Array D ix e -backpermute sz ixF !arr = makeArray (getComp arr) sz (evaluateAt arr . ixF) +backpermute = backpermute' {-# INLINE backpermute #-} +{-# DEPRECATED backpermute "In favor of a safe `backpermuteM` or an equivalent `backpermute'`" #-} +-- | /O(1)/ - Add an element to the vector from the left side +-- +-- @since 0.3.0 +cons :: e -> Array DL Ix1 e -> Array DL Ix1 e +cons e arr = + arr + { dlSize = SafeSz (1 + unSz (dlSize arr)) + , dlLoad = + \scheduler startAt uWrite -> + uWrite startAt e >> dlLoad arr scheduler (startAt + 1) uWrite + } +{-# INLINE cons #-} + +-- | /O(1)/ - Take one element off the vector from the left side. +-- +-- @since 0.3.0 +unconsM :: (MonadThrow m, Source r Ix1 e) => Array r Ix1 e -> m (e, Array D Ix1 e) +unconsM arr + | 0 == totalElem sz = throwM $ SizeEmptyException sz + | otherwise = + pure + ( unsafeLinearIndex arr 0 + , makeArray (getComp arr) (SafeSz (unSz sz - 1)) (\ !i -> unsafeLinearIndex arr (i + 1))) + where + !sz = size arr +{-# INLINE unconsM #-} + +-- | /O(1)/ - Add an element to the vector from the right side +-- +-- @since 0.3.0 +snoc :: Array DL Ix1 e -> e -> Array DL Ix1 e +snoc arr e = + arr + { dlSize = SafeSz (1 + k) + , dlLoad = + \scheduler startAt uWrite -> dlLoad arr scheduler startAt uWrite >> uWrite (k + startAt) e + } + where + !k = unSz (size arr) +{-# INLINE snoc #-} + + +-- | /O(1)/ - Take one element off the vector from the right side. +-- +-- @since 0.3.0 +unsnocM :: (MonadThrow m, Source r Ix1 e) => Array r Ix1 e -> m (Array D Ix1 e, e) +unsnocM arr + | 0 == totalElem sz = throwM $ SizeEmptyException sz + | otherwise = + pure (makeArray (getComp arr) (SafeSz k) (unsafeLinearIndex arr), unsafeLinearIndex arr k) + where + !sz = size arr + !k = unSz sz - 1 +{-# INLINE unsnocM #-} + + -- | Append two arrays together along a particular dimension. Sizes of both arrays must match, with -- an allowed exception of the dimension they are being appended along, otherwise `Nothing` is -- returned. -- -- ===__Examples__ -- --- Append two 2D arrays along both dimensions. Note that they have the same shape. +-- Append two 2D arrays along both dimensions. Note that they do agree on inner dimensions. -- --- >>> let arrA = makeArrayR U Seq (2 :. 3) (\(i :. j) -> ('A', i, j)) --- >>> let arrB = makeArrayR U Seq (2 :. 3) (\(i :. j) -> ('B', i, j)) --- >>> append 1 arrA arrB --- Just (Array D Seq (2 :. 6) --- [ [ ('A',0,0),('A',0,1),('A',0,2),('B',0,0),('B',0,1),('B',0,2) ] --- , [ ('A',1,0),('A',1,1),('A',1,2),('B',1,0),('B',1,1),('B',1,2) ] --- ]) --- >>> append 2 arrA arrB --- Just (Array D Seq (4 :. 3) --- [ [ ('A',0,0),('A',0,1),('A',0,2) ] --- , [ ('A',1,0),('A',1,1),('A',1,2) ] --- , [ ('B',0,0),('B',0,1),('B',0,2) ] --- , [ ('B',1,0),('B',1,1),('B',1,2) ] --- ]) +-- >>> import Data.Massiv.Array +-- >>> arrA = makeArrayR U Seq (Sz2 2 3) (\(i :. j) -> ('A', i, j)) +-- >>> arrB = makeArrayR U Seq (Sz2 2 3) (\(i :. j) -> ('B', i, j)) +-- >>> appendM 1 arrA arrB +-- Array DL Seq (Sz (2 :. 6)) +-- [ [ ('A',0,0), ('A',0,1), ('A',0,2), ('B',0,0), ('B',0,1), ('B',0,2) ] +-- , [ ('A',1,0), ('A',1,1), ('A',1,2), ('B',1,0), ('B',1,1), ('B',1,2) ] +-- ] +-- >>> appendM 2 arrA arrB +-- Array DL Seq (Sz (4 :. 3)) +-- [ [ ('A',0,0), ('A',0,1), ('A',0,2) ] +-- , [ ('A',1,0), ('A',1,1), ('A',1,2) ] +-- , [ ('B',0,0), ('B',0,1), ('B',0,2) ] +-- , [ ('B',1,0), ('B',1,1), ('B',1,2) ] +-- ] -- -- Now appending arrays with different sizes: -- --- >>> let arrC = makeArrayR U Seq (2 :. 4) (\(i :. j) -> ('C', i, j)) --- >>> append 1 arrA arrC --- Just (Array D Seq (2 :. 7) --- [ [ ('A',0,0),('A',0,1),('A',0,2),('C',0,0),('C',0,1),('C',0,2),('C',0,3) ] --- , [ ('A',1,0),('A',1,1),('A',1,2),('C',1,0),('C',1,1),('C',1,2),('C',1,3) ] --- ]) --- >>> append 2 arrA arrC --- Nothing +-- >>> arrC = makeArrayR U Seq (Sz (2 :. 4)) (\(i :. j) -> ('C', i, j)) +-- >>> appendM 1 arrA arrC +-- Array DL Seq (Sz (2 :. 7)) +-- [ [ ('A',0,0), ('A',0,1), ('A',0,2), ('C',0,0), ('C',0,1), ('C',0,2), ('C',0,3) ] +-- , [ ('A',1,0), ('A',1,1), ('A',1,2), ('C',1,0), ('C',1,1), ('C',1,2), ('C',1,3) ] +-- ] +-- >>> appendM 2 arrA arrC +-- *** Exception: SizeMismatchException: (Sz (2 :. 3)) vs (Sz (2 :. 4)) -- -append :: (Source r1 ix e, Source r2 ix e) => - Dim -> Array r1 ix e -> Array r2 ix e -> Maybe (Array D ix e) -append n !arr1 !arr2 = do - let sz1 = size arr1 - sz2 = size arr2 - k1 <- getDim sz1 n - k2 <- getDim sz2 n - sz1' <- setDim sz2 n k1 - guard $ sz1 == sz1' - newSz <- setDim sz1 n (k1 + k2) +-- @since 0.3.0 +appendM :: (MonadThrow m, Source r1 ix e, Source r2 ix e) => + Dim -> Array r1 ix e -> Array r2 ix e -> m (Array DL ix e) +appendM n !arr1 !arr2 = do + let !sz1 = size arr1 + !sz2 = size arr2 + (k1, szl1) <- pullOutSzM sz1 n + (k2, szl2) <- pullOutSzM sz2 n + unless (szl1 == szl2) $ throwM $ SizeMismatchException sz1 sz2 + let k1' = unSz k1 + newSz <- insertSzM szl1 n (SafeSz (k1' + unSz k2)) return $ - unsafeMakeArray (getComp arr1) newSz $ \ !ix -> - fromMaybe (errorImpossible "append" ix) $ do - k' <- getDim ix n - if k' < k1 - then Just (unsafeIndex arr1 ix) - else do - i <- getDim ix n - ix' <- setDim ix n (i - k1) - return $ unsafeIndex arr2 ix' + DLArray + { dlComp = getComp arr1 <> getComp arr2 + , dlSize = newSz + , dlLoad = + \Scheduler{scheduleWork} startAt dlWrite -> do + scheduleWork $ + iterM_ zeroIndex (unSz sz1) (pureIndex 1) (<) $ \ix -> + dlWrite (startAt + toLinearIndex newSz ix) (unsafeIndex arr1 ix) + scheduleWork $ + iterM_ zeroIndex (unSz sz2) (pureIndex 1) (<) $ \ix -> + let i = getDim' ix n + ix' = setDim' ix n (i + k1') + in dlWrite (startAt + toLinearIndex newSz ix') (unsafeIndex arr2 ix) + } +{-# INLINE appendM #-} + +-- | Append two arrays together along a specified dimension. +append :: (Source r1 ix e, Source r2 ix e) => + Dim -> Array r1 ix e -> Array r2 ix e -> Maybe (Array DL ix e) +append = appendM {-# INLINE append #-} +{-# DEPRECATED append "In favor of a more general `appendM`" #-} --- | Same as `append`, but will throw an error instead of returning `Nothing` on mismatched sizes. + +-- | Same as `appendM`, but will throw an exception in pure code on mismatched sizes. +-- +-- @since 0.3.0 append' :: (Source r1 ix e, Source r2 ix e) => - Dim -> Array r1 ix e -> Array r2 ix e -> Array D ix e -append' dim arr1 arr2 = - case append dim arr1 arr2 of - Just arr -> arr - Nothing -> - error $ - if 0 < dim && dim <= dimensions (size arr1) - then "append': Dimension mismatch: " ++ show (size arr1) ++ " and " ++ show (size arr2) - else "append': Invalid dimension: " ++ show dim + Dim -> Array r1 ix e -> Array r2 ix e -> Array DL ix e +append' dim arr1 arr2 = either throw id $ appendM dim arr1 arr2 {-# INLINE append' #-} +-- | Concat many arrays together along some dimension. +-- +-- @since 0.3.0 +concat' :: (Foldable f, Source r ix e) => Dim -> f (Array r ix e) -> Array DL ix e +concat' n arrs = either throw id $ concatM n arrs +{-# INLINE concat' #-} + +-- | Concatenate many arrays together along some dimension. It is important that all sizes are +-- equal, with an exception of the dimensions along which concatenation happens, otherwise it doues +-- result in a `SizeMismatchException` exception. +-- +-- @since 0.3.0 +concatM :: + (MonadThrow m, Foldable f, Source r ix e) => Dim -> f (Array r ix e) -> m (Array DL ix e) +concatM n !arrsF = + case L.uncons (F.toList arrsF) of + Nothing -> pure empty + Just (a, arrs) -> do + let sz = unSz (size a) + szs = P.map (unSz . size) arrs + (k, szl) <- pullOutDimM sz n + -- / remove the dimension out of all sizes along which concatenation will happen + (ks, szls) <- + F.foldrM (\ !csz (ks, szls) -> bimap (: ks) (: szls) <$> pullOutDimM csz n) ([], []) szs + -- / make sure to fail as soon as at least one of the arrays has a mismatching inner size + traverse_ + (\(sz', _) -> throwM (SizeMismatchException (SafeSz sz) (SafeSz sz'))) + (dropWhile ((== szl) . snd) $ zip szs szls) + let kTotal = SafeSz $ F.foldl' (+) k ks + newSz <- insertSzM (SafeSz szl) n kTotal + return $ + DLArray + { dlComp = mconcat $ P.map getComp arrs + , dlSize = newSz + , dlLoad = + \Scheduler{scheduleWork} startAt dlWrite -> + let arrayLoader !kAcc (kCur, arr) = do + scheduleWork $ + iterM_ zeroIndex (unSz (size arr)) (pureIndex 1) (<) $ \ix -> + let i = getDim' ix n + ix' = setDim' ix n (i + kAcc) + in dlWrite (startAt + toLinearIndex newSz ix') (unsafeIndex arr ix) + pure (kAcc + kCur) + in M.foldM_ arrayLoader 0 $ (k, a) : P.zip ks arrs + } +{-# INLINE concatM #-} + + -- | /O(1)/ - Split an array at an index along a specified dimension. +-- +-- @since 0.3.0 +splitAtM :: + (MonadThrow m, Extract r ix e, r' ~ EltRepr r ix) + => Dim -- ^ Dimension along which to split + -> Int -- ^ Index along the dimension to split at + -> Array r ix e -- ^ Source array + -> m (Array r' ix e, Array r' ix e) +splitAtM dim i arr = do + let Sz sz = size arr + eIx <- setDimM sz dim i + sIx <- setDimM zeroIndex dim i + arr1 <- extractFromToM zeroIndex eIx arr + arr2 <- extractFromToM sIx sz arr + return (arr1, arr2) +{-# INLINE splitAtM #-} + +-- | /O(1)/ - Split an array at an index along a specified dimension. splitAt :: - (Size r ix e, r' ~ EltRepr r ix) + (Extract r ix e, r' ~ EltRepr r ix) => Dim -- ^ Dimension along which to split -> Int -- ^ Index along the dimension to split at -> Array r ix e -- ^ Source array -> Maybe (Array r' ix e, Array r' ix e) splitAt dim i arr = do - let sz = size arr - eIx <- setDim sz dim i - sIx <- setDim zeroIndex dim i + let Sz sz = size arr + eIx <- setDimM sz dim i + sIx <- setDimM zeroIndex dim i arr1 <- extractFromTo zeroIndex eIx arr arr2 <- extractFromTo sIx sz arr return (arr1, arr2) {-# INLINE splitAt #-} +{-# DEPRECATED splitAt "In favor of a more general `splitAtM`" #-} -- | Same as `splitAt`, but will throw an error instead of returning `Nothing` on wrong dimension -- and index out of bounds. -splitAt' :: (Size r ix e, r' ~ EltRepr r ix) => +splitAt' :: (Extract r ix e, r' ~ EltRepr r ix) => Dim -> Int -> Array r ix e -> (Array r' ix e, Array r' ix e) -splitAt' dim i arr = - case splitAt dim i arr of - Just res -> res - Nothing -> - error $ - "Data.Massiv.Array.splitAt': " ++ - if 0 < dim && dim <= dimensions (size arr) - then "Index out of bounds: " ++ - show i ++ " for dimension: " ++ show dim ++ " and array with size: " ++ show (size arr) - else "Invalid dimension: " ++ show dim ++ " for array with size: " ++ show (size arr) +splitAt' dim i arr = either throw id $ splitAtM dim i arr {-# INLINE splitAt' #-} + +-- | Discard elements from the source array according to the stride. +-- +-- @since 0.3.0 +-- +downsample :: Source r ix e => Stride ix -> Array r ix e -> Array DL ix e +downsample !stride arr = + DLArray + { dlComp = getComp arr + , dlSize = resultSize + , dlLoad = + \scheduler startAt dlWrite -> + splitLinearlyWithStartAtM_ + scheduler + startAt + (totalElem resultSize) + (pure . unsafeLinearWriteWithStride) + dlWrite + } + where + resultSize = strideSize stride (size arr) + strideIx = unStride stride + unsafeLinearWriteWithStride = + unsafeIndex arr . liftIndex2 (*) strideIx . fromLinearIndex resultSize + {-# INLINE unsafeLinearWriteWithStride #-} +{-# INLINE downsample #-} + + +-- | Insert the same element into a `Load`able array according to the stride. +-- +-- @since 0.3.0 +upsample + :: Load r ix e => e -> Stride ix -> Array r ix e -> Array DL ix e +upsample !fillWith !safeStride arr = + DLArray + { dlComp = getComp arr + , dlSize = newsz + , dlLoad = + \scheduler startAt dlWrite -> do + unless (stride == pureIndex 1) $ + loopM_ startAt (< totalElem newsz) (+ 1) (`dlWrite` fillWith) + -- TODO: experiment a bit more. So far the fastest solution is to prefill the whole array + -- with default value and override non-stride elements afterwards. This approach seems a + -- bit wasteful, nevertheless it is fastest + -- + -- TODO: Is it possible to use fast fill operation that is available for MutableByteArray? + loadArrayM scheduler arr (\i -> dlWrite (adjustLinearStride (i + startAt))) + } + where + adjustLinearStride = toLinearIndex newsz . timesStride . fromLinearIndex sz + {-# INLINE adjustLinearStride #-} + timesStride !ix = liftIndex2 (*) stride ix + {-# INLINE timesStride #-} + !stride = unStride safeStride + !sz = size arr + !newsz = SafeSz (timesStride $ unSz sz) +{-# INLINE upsample #-} + + -- This was a sample optimization, that turned out to be significantly (~ x9) slower + -- makeLoadArray (getComp arr) newsz $ \numWorkers scheduleWith dlWrite -> do + -- iterM_ zeroIndex stride (pureIndex 1) (<) $ \ixs -> + -- if ixs == zeroIndex + -- then loadArray numWorkers scheduleWith arr $ \ !i -> dlWrite (adjustLinearStride i) + -- else let !is = toLinearIndex newsz ixs + -- in scheduleWith $ + -- loopM_ 0 (< totalElem sz) (+ 1) $ \ !i -> + -- dlWrite (is + adjustLinearStride i) fillWith + + + -- | Create an array by traversing a source array. traverse :: (Source r1 ix1 e1, Index ix) @@ -378,9 +683,9 @@ -- an index for an element it should return a value of. -> Array r1 ix1 e1 -- ^ Source array -> Array D ix e -traverse sz f arr1 = makeArray (getComp arr1) sz (f (evaluateAt arr1)) +traverse sz f arr1 = makeArray (getComp arr1) sz (f (evaluate' arr1)) {-# INLINE traverse #-} - +{-# DEPRECATED traverse "In favor of more general `transform'`" #-} -- | Create an array by traversing two source arrays. traverse2 @@ -390,12 +695,69 @@ -> Array r1 ix1 e1 -> Array r2 ix2 e2 -> Array D ix e -traverse2 sz f arr1 arr2 = makeArray (getComp arr1) sz (f (evaluateAt arr1) (evaluateAt arr2)) +traverse2 sz f arr1 arr2 = + makeArray (getComp arr1 <> getComp arr2) sz (f (evaluate' arr1) (evaluate' arr2)) {-# INLINE traverse2 #-} +{-# DEPRECATED traverse2 "In favor of more general `transform2'`" #-} +-- | General array transformation, that forces computation and produces a manifest array. +-- +-- @since 0.3.0 +transformM :: + forall r ix e r' ix' e' a m. + (Mutable r ix e, Source r' ix' e', MonadUnliftIO m, PrimMonad m, MonadThrow m) + => (Sz ix' -> m (Sz ix, a)) + -> (a -> (ix' -> m e') -> ix -> m e) + -> Array r' ix' e' + -> m (Array r ix e) +transformM getSzM getM arr = do + (sz, a) <- getSzM (size arr) + generateArray (getComp arr) sz (getM a (evaluateM arr)) +{-# INLINE transformM #-} --- | Throw an impossible error on a `Nothing` -errorImpossible :: Show c => String -> c -> a -errorImpossible fName cause = - error $ "Data.Massiv.Array." ++ fName ++ ": Impossible happened " ++ show cause -{-# NOINLINE errorImpossible #-} + +-- | General array transformation +-- +-- @since 0.3.0 +transform' :: + (Source r' ix' e', Index ix) + => (Sz ix' -> (Sz ix, a)) + -> (a -> (ix' -> e') -> ix -> e) + -> Array r' ix' e' + -> Array D ix e +transform' getSz get arr = makeArray (getComp arr) sz (get a (evaluate' arr)) + where + (sz, a) = getSz (size arr) +{-# INLINE transform' #-} + +-- | Same as `transformM`, but operates on two arrays +-- +-- @since 0.3.0 +transform2M :: + (Mutable r ix e, Source r1 ix1 e1, Source r2 ix2 e2, MonadUnliftIO m, PrimMonad m, MonadThrow m) + => (Sz ix1 -> Sz ix2 -> m (Sz ix, a)) + -> (a -> (ix1 -> m e1) -> (ix2 -> m e2) -> ix -> m e) + -> Array r1 ix1 e1 + -> Array r2 ix2 e2 + -> m (Array r ix e) +transform2M getSzM getM arr1 arr2 = do + (sz, a) <- getSzM (size arr1) (size arr2) + generateArray (getComp arr1 <> getComp arr2) sz (getM a (evaluateM arr1) (evaluateM arr2)) +{-# INLINE transform2M #-} + + +-- | Same as `transform'`, but operates on two arrays +-- +-- @since 0.3.0 +transform2' :: + (Source r1 ix1 e1, Source r2 ix2 e2, Index ix) + => (Sz ix1 -> Sz ix2 -> (Sz ix, a)) + -> (a -> (ix1 -> e1) -> (ix2 -> e2) -> ix -> e) + -> Array r1 ix1 e1 + -> Array r2 ix2 e2 + -> Array D ix e +transform2' getSz get arr1 arr2 = + makeArray (getComp arr1 <> getComp arr2) sz (get a (evaluate' arr1) (evaluate' arr2)) + where + (sz, a) = getSz (size arr1) (size arr2) +{-# INLINE transform2' #-} diff --git a/src/Data/Massiv/Array/Stencil.hs b/src/Data/Massiv/Array/Stencil.hs --- a/src/Data/Massiv/Array/Stencil.hs +++ b/src/Data/Massiv/Array/Stencil.hs @@ -1,9 +1,9 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE MultiParamTypeClasses #-} -- | -- Module : Data.Massiv.Array.Stencil --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental @@ -22,39 +22,41 @@ , rmapStencil -- * Convolution , module Data.Massiv.Array.Stencil.Convolution + -- * Re-export + , Default(def) ) where -import Data.Default.Class (Default (def)) -import Data.Massiv.Array.Delayed.Windowed -import Data.Massiv.Array.Manifest -import Data.Massiv.Array.Stencil.Convolution -import Data.Massiv.Array.Stencil.Internal -import Data.Massiv.Core.Common -import GHC.Exts (inline) +import Data.Default.Class (Default(def)) +import Data.Massiv.Array.Delayed.Windowed +import Data.Massiv.Array.Manifest +import Data.Massiv.Array.Stencil.Convolution +import Data.Massiv.Array.Stencil.Internal +import Data.Massiv.Core.Common +import GHC.Exts (inline) -- | Map a constructed stencil over an array. Resulting array must be `compute`d in order to be -- useful. +-- +-- @since 0.1.0 mapStencil :: (Source r ix e, Manifest r ix e) => Border e -- ^ Border resolution technique -> Stencil ix e a -- ^ Stencil to map over the array -> Array r ix e -- ^ Source array -> Array DW ix a -mapStencil b (Stencil sSz sCenter stencilF) !arr = - DWArray - (DArray (getComp arr) sz (unValue . stencilF (Value . borderIndex b arr))) - (Just sSz) - (Just window) +mapStencil b (Stencil sSz sCenter stencilF) !arr = insertWindow warr window where + !warr = DArray (getComp arr) sz (unValue . stencilF (Value . borderIndex b arr)) !window = Window - { windowStart = liftIndex2 min sCenter (liftIndex (max 0) (liftIndex (subtract 1) sz)) - , windowSize = liftIndex (max 0) (liftIndex2 min windowSz (liftIndex2 (-) sz sCenter)) + { windowStart = sCenter + , windowSize = windowSz , windowIndex = unValue . stencilF (Value . unsafeIndex arr) + , windowUnrollIx2 = unSz . fst <$> pullOutSzM sSz 2 } - !windowSz = liftIndex (max 0) (liftIndex2 (-) sz (liftIndex (subtract 1) sSz)) !sz = size arr + !windowSz = Sz (liftIndex2 (-) (unSz sz) (liftIndex (subtract 1) (unSz sSz))) {-# INLINE mapStencil #-} @@ -78,9 +80,10 @@ -- > get ( 1 :. -1) + get ( 1 :. 0) + get ( 1 :. 1) ) / 9 -- > {-# INLINE average3x3Stencil #-} -- +-- @since 0.1.0 makeStencil :: (Index ix, Default e) - => ix -- ^ Size of the stencil + => Sz ix -- ^ Size of the stencil -> ix -- ^ Center of the stencil -> ((ix -> Value e) -> Value a) -- ^ Stencil function that receives a "get" function as it's argument that can @@ -94,6 +97,8 @@ {-# INLINE makeStencil #-} -- | Same as `makeStencil`, but with ability to specify default value for stencil validation. +-- +-- @since 0.2.3 makeStencilDef :: Index ix => e 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,5 @@ +{-# LANGUAGE BangPatterns #-} {-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE BangPatterns #-} -- | -- Module : Data.Massiv.Array.Stencil.Convolution -- Copyright : (c) Alexey Kuleshevich 2018-2019 @@ -15,10 +15,10 @@ , makeCorrelationStencilFromKernel ) where -import Data.Massiv.Core.Common -import Data.Massiv.Array.Ops.Fold (ifoldlS) -import Data.Massiv.Array.Stencil.Internal -import GHC.Exts (inline) +import Data.Massiv.Array.Ops.Fold (ifoldlS) +import Data.Massiv.Array.Stencil.Internal +import Data.Massiv.Core.Common +import GHC.Exts (inline) -- | Create a convolution stencil by specifying border resolution technique and -- an accumulator function. @@ -44,7 +44,7 @@ makeConvolutionStencil !sz !sCenter relStencil = validateStencil 0 $ Stencil sz sInvertCenter stencil where - !sInvertCenter = liftIndex2 (-) (liftIndex (subtract 1) sz) sCenter + !sInvertCenter = liftIndex2 (-) (liftIndex (subtract 1) (unSz sz)) sCenter stencil getVal !ix = (inline relStencil $ \ !ixD !kVal !acc -> getVal (liftIndex2 (-) ix ixD) * kVal + acc) 0 {-# INLINE stencil #-} @@ -103,7 +103,7 @@ makeCorrelationStencilFromKernel kArr = Stencil sz sCenter stencil where !sz = size kArr - !sCenter = liftIndex (`div` 2) sz + !sCenter = liftIndex (`div` 2) $ unSz sz stencil getVal !ix = Value (ifoldlS accum 0 kArr) where !ixOff = liftIndex2 (-) ix sCenter accum !acc !kIx !kVal = diff --git a/src/Data/Massiv/Array/Stencil/Internal.hs b/src/Data/Massiv/Array/Stencil/Internal.hs --- a/src/Data/Massiv/Array/Stencil/Internal.hs +++ b/src/Data/Massiv/Array/Stencil/Internal.hs @@ -1,13 +1,12 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE CPP #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE RecordWildCards #-} -{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} -- | -- Module : Data.Massiv.Array.Stencil.Internal --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental @@ -22,19 +21,17 @@ , validateStencil ) where -import Control.Applicative -import Control.DeepSeq -import Data.Massiv.Array.Delayed.Internal -import Data.Massiv.Core.Common -#if !MIN_VERSION_base(4,11,0) -import Data.Semigroup -#endif +import Control.Applicative +import Control.DeepSeq +import Data.Massiv.Array.Delayed.Pull +import Data.Massiv.Core.Common +import Data.Massiv.Core.Index.Internal -- | Stencil is abstract description of how to handle elements in the neighborhood of every array -- cell in order to compute a value for the cells in the new array. Use `Data.Array.makeStencil` and -- `Data.Array.makeConvolutionStencil` in order to create a stencil. data Stencil ix e a = Stencil - { stencilSize :: !ix + { stencilSize :: !(Sz ix) , stencilCenter :: !ix , stencilFunc :: (ix -> Value e) -> ix -> Value a } @@ -139,29 +136,45 @@ -- Profunctor - +-- | A Profunctor dimap. Same caviat applies as in `lmapStencil` +-- +-- @since 0.2.3 dimapStencil :: (c -> d) -> (a -> b) -> Stencil ix d a -> Stencil ix c b -dimapStencil f g stencil@(Stencil {stencilFunc = sf}) = stencil {stencilFunc = sf'} +dimapStencil f g stencil@Stencil {stencilFunc = sf} = stencil {stencilFunc = sf'} where sf' s = Value . g . unValue . sf (Value . f . unValue . s) {-# INLINE sf' #-} {-# INLINE dimapStencil #-} - +-- | A contravariant map of a second type parameter. In other words map a function over each element +-- of the array, that the stencil will be applied to. +-- +-- __Note__: This map can be very inefficient, since for stencils larger than 1 element in size, the +-- supllied function will be repeatedly applied to the same element. It is better to simply map that +-- function over the source array instead. +-- +-- @since 0.2.3 lmapStencil :: (c -> d) -> Stencil ix d a -> Stencil ix c a -lmapStencil f stencil@(Stencil {stencilFunc = sf}) = stencil {stencilFunc = sf'} +lmapStencil f stencil@Stencil {stencilFunc = sf} = stencil {stencilFunc = sf'} where sf' s = sf (Value . f . unValue . s) {-# INLINE sf' #-} {-# INLINE lmapStencil #-} +-- | A covariant map over the right most type argument. In other words a usual Functor `fmap`: +-- +-- > fmap == rmapStencil +-- +-- @since 0.2.3 rmapStencil :: (a -> b) -> Stencil ix e a -> Stencil ix e b -rmapStencil f stencil@(Stencil {stencilFunc = sf}) = stencil {stencilFunc = sf'} +rmapStencil f stencil@Stencil {stencilFunc = sf} = stencil {stencilFunc = sf'} where sf' s = Value . f . unValue . sf s {-# INLINE sf' #-} {-# INLINE rmapStencil #-} + + -- TODO: Figure out interchange law (u <*> pure y = pure ($ y) <*> u) and issue -- with discarding size and center. Best idea so far is to increase stencil size to -- the maximum one and shift the center of the other stencil so that they both match @@ -169,17 +182,18 @@ -- Stencil - both stencils are trusted, increasing the size will not affect the -- safety. instance Index ix => Applicative (Stencil ix e) where - pure a = Stencil (pureIndex 1) zeroIndex (const (const (Value a))) + pure a = Stencil oneSz zeroIndex (const (const (Value a))) {-# INLINE pure #-} - (<*>) (Stencil sSz1 sC1 f1) (Stencil sSz2 sC2 f2) = Stencil newSz maxCenter stF + (<*>) (Stencil (SafeSz sSz1) sC1 f1) (Stencil (SafeSz sSz2) sC2 f2) = Stencil newSz maxCenter stF where - stF gV !ix = Value ((unValue (f1 gV ix)) (unValue (f2 gV ix))) + stF gV !ix = Value (unValue (f1 gV ix) (unValue (f2 gV ix))) {-# INLINE stF #-} !newSz = - liftIndex2 - (+) - maxCenter - (liftIndex2 max (liftIndex2 (-) sSz1 sC1) (liftIndex2 (-) sSz2 sC2)) + Sz + (liftIndex2 + (+) + maxCenter + (liftIndex2 max (liftIndex2 (-) sSz1 sC1) (liftIndex2 (-) sSz2 sC2))) !maxCenter = liftIndex2 max sC1 sC2 {-# INLINE (<*>) #-} @@ -249,9 +263,7 @@ safeStencilIndex :: Index ix => Array D ix e -> ix -> e safeStencilIndex DArray {..} ix | isSafeIndex dSize ix = dIndex ix - | otherwise = - error $ - "Index is out of bounds: " ++ show ix ++ " for stencil size: " ++ show dSize + | otherwise = throw $ IndexOutOfBoundsException dSize ix -- | Make sure constructed stencil doesn't index outside the allowed stencil size boundary. 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 @@ -1,10 +1,10 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} -- | -- Module : Data.Massiv.Array.Stencil.Unsafe --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental @@ -12,19 +12,20 @@ -- module Data.Massiv.Array.Stencil.Unsafe ( -- * Stencil - forStencilUnsafe + makeUnsafeStencil + , forStencilUnsafe ) where -import Data.Massiv.Core.Common -import Data.Massiv.Array.Delayed.Windowed (Window(..), DW, Array(..)) -import GHC.Exts (inline) +import Data.Massiv.Array.Delayed.Windowed (Array(..), DW, Window(..), + insertWindow) +import Data.Massiv.Array.Stencil.Internal +import Data.Massiv.Core.Common +import GHC.Exts (inline) --- | This is an unsafe version of the stencil computation. There are no bounds check further from --- the border, so if you make sure you don't go outside the size of the stencil, you will be safe, --- but this is not enforced. --- --- __/Note/__ - Still experimental and subject to change +-- | This is an unsafe version of the stencil computation. There are no bounds checking further from +-- the border, so if you do make sure you are not going outside the size of the stencil, you will be +-- safe, but this is not enforced. -- -- @since 0.1.7 forStencilUnsafe :: @@ -37,23 +38,41 @@ -- retrieve values of cells in the source array with respect to the center of -- the stencil. Stencil function must return a value that will be assigned to -- the cell in the result array. Offset supplied to the "get" function - -- cannot go outside the boundaries of the stencil, otherwise an error will be - -- raised during stencil creation. + -- cannot go outside the boundaries of the stencil. -> Array DW ix a forStencilUnsafe !arr !sSz !sCenter relStencil = - DWArray - (DArray (getComp arr) sz (stencil (index arr))) - (Just sSz) - (Just window) + insertWindow (DArray (getComp arr) sz (stencil (index arr))) window where !window = Window - { windowStart = liftIndex2 min sCenter (liftIndex (max 0) (liftIndex (subtract 1) sz)) - , windowSize = liftIndex (max 0) (liftIndex2 min windowSz (liftIndex2 (-) sz sCenter)) + { windowStart = sCenter + , windowSize = windowSz , windowIndex = stencil (Just . unsafeIndex arr) + , windowUnrollIx2 = unSz . fst <$> pullOutSzM windowSz 2 } - !windowSz = liftIndex (max 0) (liftIndex2 (-) sz (liftIndex (subtract 1) sSz)) + !sz = size arr + !windowSz = Sz (liftIndex2 (-) (unSz sz) (liftIndex (subtract 1) (unSz sSz))) stencil getVal !ix = inline relStencil $ \ !ixD -> getVal (liftIndex2 (+) ix ixD) {-# INLINE stencil #-} - !sz = size arr {-# INLINE forStencilUnsafe #-} + + + +-- | 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. +-- +-- @since 0.3.0 +makeUnsafeStencil + :: Index ix + => Sz ix -- ^ Size of the stencil + -> ix -- ^ Center of the stencil + -> (ix -> (ix -> e) -> a) + -- ^ Stencil function. + -> Stencil ix e a +makeUnsafeStencil !sSz !sCenter relStencil = Stencil sSz sCenter stencil + where + stencil getVal !ix = + Value $ inline $ relStencil ix (unValue . getVal . liftIndex2 (+) ix) + {-# INLINE stencil #-} +{-# INLINE makeUnsafeStencil #-} diff --git a/src/Data/Massiv/Array/Unsafe.hs b/src/Data/Massiv/Array/Unsafe.hs --- a/src/Data/Massiv/Array/Unsafe.hs +++ b/src/Data/Massiv/Array/Unsafe.hs @@ -1,10 +1,11 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE PatternSynonyms #-} -- | -- Module : Data.Massiv.Array.Unsafe --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental @@ -12,20 +13,23 @@ -- module Data.Massiv.Array.Unsafe ( -- * Creation - unsafeMakeArray - , unsafeGenerateArray - , unsafeGenerateArrayP - , unsafeGenerateM + -- , unsafeGenerateArray + -- , unsafeGenerateArrayP -- * Indexing + Sz(SafeSz) + , Stride(SafeStride) , unsafeIndex , unsafeLinearIndex , unsafeLinearIndexM -- * Manipulations , unsafeBackpermute - , unsafeTraverse - , unsafeTraverse2 , unsafeResize , unsafeExtract + , unsafeTransform + , unsafeTransform2 + -- ** Deprecated + , unsafeTraverse + , unsafeTraverse2 -- * Slicing , unsafeSlice , unsafeOuterSlice @@ -34,37 +38,56 @@ , unsafeThaw , unsafeFreeze , unsafeNew - , unsafeNewZero , unsafeRead , unsafeLinearRead , unsafeWrite , unsafeLinearWrite + , unsafeLinearSet + -- * Pointer access + , unsafeWithPtr + , unsafeArrayToForeignPtr + , unsafeMArrayToForeignPtr + , unsafeArrayFromForeignPtr + , unsafeArrayFromForeignPtr0 + , unsafeMArrayFromForeignPtr + , unsafeMArrayFromForeignPtr0 + -- ** Atomic Operations + , unsafeAtomicReadIntArray + , unsafeAtomicWriteIntArray + , unsafeAtomicModifyIntArray + , unsafeAtomicAddIntArray + , unsafeAtomicSubIntArray + , unsafeAtomicAndIntArray + , unsafeAtomicNandIntArray + , unsafeAtomicOrIntArray + , unsafeAtomicXorIntArray + , unsafeCasIntArray ) where -import Control.Monad.Primitive (PrimMonad (..)) -import Control.Monad.ST (runST) -import Data.Massiv.Array.Delayed.Internal (D) -import Data.Massiv.Core.Common -import Data.Massiv.Core.Scheduler -import System.IO.Unsafe (unsafePerformIO) - +import Data.Massiv.Array.Delayed.Pull (D) +import Data.Massiv.Array.Manifest.Primitive +import Data.Massiv.Array.Manifest.Storable +import Data.Massiv.Core.Common +import Data.Massiv.Core.Index.Internal (Sz(SafeSz)) +import Data.Massiv.Core.Index.Stride (Stride(SafeStride)) unsafeBackpermute :: (Source r' ix' e, Index ix) => Sz ix -> (ix -> ix') -> Array r' ix' e -> Array D ix e unsafeBackpermute !sz ixF !arr = - unsafeMakeArray (getComp arr) sz $ \ !ix -> unsafeIndex arr (ixF ix) + makeArray (getComp arr) sz $ \ !ix -> unsafeIndex arr (ixF ix) {-# INLINE unsafeBackpermute #-} unsafeTraverse - :: (Source r1 ix1 e1, Index ix) + :: (Source r ix' e', Index ix) => Sz ix - -> ((ix1 -> e1) -> ix -> e) - -> Array r1 ix1 e1 + -> ((ix' -> e') -> ix -> e) + -> Array r ix' e' -> Array D ix e unsafeTraverse sz f arr1 = - unsafeMakeArray (getComp arr1) sz (f (unsafeIndex arr1)) + makeArray (getComp arr1) sz (f (unsafeIndex arr1)) {-# INLINE unsafeTraverse #-} +{-# DEPRECATED unsafeTraverse "In favor of more general `unsafeTransform'`" #-} unsafeTraverse2 @@ -75,45 +98,38 @@ -> Array r2 ix2 e2 -> Array D ix e unsafeTraverse2 sz f arr1 arr2 = - unsafeMakeArray (getComp arr1) sz (f (unsafeIndex arr1) (unsafeIndex arr2)) + makeArray (getComp arr1 <> getComp arr2) sz (f (unsafeIndex arr1) (unsafeIndex arr2)) {-# INLINE unsafeTraverse2 #-} - - --- | Read an array element -unsafeRead :: (Mutable r ix e, PrimMonad m) => - MArray (PrimState m) r ix e -> ix -> m e -unsafeRead !marr !ix = unsafeLinearRead marr (toLinearIndex (msize marr) ix) -{-# INLINE unsafeRead #-} - --- | Write an element into array -unsafeWrite :: (Mutable r ix e, PrimMonad m) => - MArray (PrimState m) r ix e -> ix -> e -> m () -unsafeWrite !marr !ix = unsafeLinearWrite marr (toLinearIndex (msize marr) ix) -{-# INLINE unsafeWrite #-} - +{-# DEPRECATED unsafeTraverse2 "In favor of more general `unsafeTransform2'`" #-} --- | Create an array sequentially using mutable interface -unsafeGenerateArray :: Mutable r ix e => Sz ix -> (ix -> e) -> Array r ix e -unsafeGenerateArray !sz f = runST $ do - marr <- unsafeNew sz - iterLinearM_ sz 0 (totalElem sz) 1 (<) $ \ !k !ix -> - unsafeLinearWrite marr k (f ix) - unsafeFreeze Seq marr -{-# INLINE unsafeGenerateArray #-} +-- | Same `Data.Array.transform'`, except no bounds checking is performed, thus making it faster, +-- but unsafe. +-- +-- @since 0.3.0 +unsafeTransform :: + (Source r' ix' e', Index ix) + => (Sz ix' -> (Sz ix, a)) + -> (a -> (ix' -> e') -> ix -> e) + -> Array r' ix' e' + -> Array D ix e +unsafeTransform getSz get arr = makeArray (getComp arr) sz (get a (unsafeIndex arr)) + where + (sz, a) = getSz (size arr) +{-# INLINE unsafeTransform #-} --- | Create an array in parallel using mutable interface +-- | Same `Data.Array.transform2'`, except no bounds checking is performed, thus making it faster, +-- but unsafe. -- --- @since 0.1.5 -unsafeGenerateArrayP :: Mutable r ix e => [Int] -> Sz ix -> (ix -> e) -> Array r ix e -unsafeGenerateArrayP wIds !sz f = unsafePerformIO $ do - marr <- unsafeNew sz - divideWork_ wIds sz $ \ !scheduler !chunkLength !totalLength !slackStart -> do - loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start -> - scheduleWork scheduler $ - iterLinearM_ sz start (start + chunkLength) 1 (<) $ \ !k !ix -> - unsafeLinearWrite marr k (f ix) - scheduleWork scheduler $ - iterLinearM_ sz slackStart totalLength 1 (<) $ \ !k !ix -> - unsafeLinearWrite marr k (f ix) - unsafeFreeze (ParOn wIds) marr -{-# INLINE unsafeGenerateArrayP #-} +-- @since 0.3.0 +unsafeTransform2 :: + (Source r1 ix1 e1, Source r2 ix2 e2, Index ix) + => (Sz ix1 -> Sz ix2 -> (Sz ix, a)) + -> (a -> (ix1 -> e1) -> (ix2 -> e2) -> ix -> e) + -> Array r1 ix1 e1 + -> Array r2 ix2 e2 + -> Array D ix e +unsafeTransform2 getSz get arr1 arr2 = + makeArray (getComp arr1 <> getComp arr2) sz (get a (unsafeIndex arr1) (unsafeIndex arr2)) + where + (sz, a) = getSz (size arr1) (size arr2) +{-# INLINE unsafeTransform2 #-} diff --git a/src/Data/Massiv/Core.hs b/src/Data/Massiv/Core.hs --- a/src/Data/Massiv/Core.hs +++ b/src/Data/Massiv/Core.hs @@ -1,18 +1,6 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE CPP #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE UndecidableInstances #-} -#if __GLASGOW_HASKELL__ >= 800 - {-# OPTIONS_GHC -fno-warn-duplicate-exports #-} -#endif -- | -- Module : Data.Massiv.Core --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental @@ -23,31 +11,42 @@ , Elt , EltRepr , Construct + , Load(loadArrayM) , Source - , Load(..) - , Size + , Resize + , Extract + , StrideLoad(..) , Slice - , OuterSlice(outerLength) + , OuterSlice , InnerSlice , Manifest , Mutable - , Ragged(..) + , Ragged , Nested(..) , NestedStruct , L(..) , LN , ListItem -#if __GLASGOW_HASKELL__ >= 800 - , Comp(Seq, Par, ParOn) - , pattern Par -- already exported above and only needed for Haddock -#else - , Comp(..) - , pattern Par -#endif + , Comp(Seq, Par, ParOn, ParN) , module Data.Massiv.Core.Index + -- * Exceptions + , MonadThrow(..) + , throw + , Exception(..) + , SomeException + , IndexException(..) + , SizeException(..) + , ShapeException(..) + , module Data.Massiv.Core.Exception + -- * Stateful Monads + , MonadUnliftIO + , MonadIO(liftIO) + , PrimMonad(PrimState) ) where -import Data.Massiv.Core.Common hiding (indexWith, unsafeGenerateM) -import Data.Massiv.Core.List -import Data.Massiv.Core.Index +import Control.Exception (Exception(..), SomeException) +import Data.Massiv.Core.Common +import Data.Massiv.Core.Index +import Data.Massiv.Core.List +import Data.Massiv.Core.Exception diff --git a/src/Data/Massiv/Core/Common.hs b/src/Data/Massiv/Core/Common.hs --- a/src/Data/Massiv/Core/Common.hs +++ b/src/Data/Massiv/Core/Common.hs @@ -1,16 +1,14 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE CPP #-} -{-# LANGUAGE DefaultSignatures #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE MagicHash #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE DefaultSignatures #-} +{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE UnboxedTuples #-} -{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} -- | -- Module : Data.Massiv.Core.Common --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental @@ -22,18 +20,23 @@ , Construct(..) , Source(..) , Load(..) - , Size(..) + , StrideLoad(..) + , Resize(..) + , Extract(..) , Slice(..) , OuterSlice(..) , InnerSlice(..) , Manifest(..) , Mutable(..) - , State(..) - , WorldState + , Comp(..) + , Scheduler(..) + , unsafeRead + , unsafeWrite + , unsafeLinearModify , Ragged(..) , Nested(..) , NestedStruct - , makeArray + , empty , singleton -- * Size , elemsCount @@ -41,25 +44,44 @@ -- * Indexing , (!?) , index + , indexM , indexWith , (!) , index' , (??) , defaultIndex , borderIndex + , evaluateM + , evaluate' , evaluateAt , module Data.Massiv.Core.Index -- * Common Operations , imapM_ - , module Data.Massiv.Core.Computation + , Semigroup((<>)) + -- * Exceptions + , MonadThrow(..) + , throw + , IndexException(..) + , SizeException(..) + , ShapeException(..) + , module Data.Massiv.Core.Exception + -- * Stateful Monads + , MonadUnliftIO + , MonadIO(liftIO) + , PrimMonad(PrimState) ) where -import Control.Monad.Primitive -import Data.Massiv.Core.Computation -import Data.Massiv.Core.Index -import Data.Massiv.Core.Scheduler -import Data.Typeable -import GHC.Prim +#if !MIN_VERSION_base(4,11,0) +import Data.Semigroup +#endif +import Control.Exception (throw) +import Control.Monad.Catch (MonadThrow(..)) +import Control.Monad.IO.Unlift (MonadIO(liftIO), MonadUnliftIO) +import Control.Monad.Primitive +import Control.Scheduler (Comp(..), Scheduler(..)) +import Data.Massiv.Core.Exception +import Data.Massiv.Core.Index +import Data.Typeable #include "massiv.h" @@ -79,34 +101,86 @@ -- | Array types that can be constructed. class (Typeable r, Index ix) => Construct r ix e where - - -- | Get computation strategy of this array - getComp :: Array r ix e -> Comp + {-# MINIMAL setComp,(makeArray|makeArrayLinear) #-} -- | Set computation strategy for this array + -- + -- ==== __Example__ + -- + -- >>> :set -XTypeApplications + -- >>> import Data.Massiv.Array + -- >>> a = singleton @DL @Ix1 @Int 0 + -- >>> a + -- Array DL Seq (Sz1 1) + -- [ 0 ] + -- >>> setComp (ParN 6) a -- use 6 capabilities + -- Array DL (ParN 6) (Sz1 1) + -- [ 0 ] + -- setComp :: Comp -> Array r ix e -> Array r ix e - -- | Construct an array. No size validation is performed. - unsafeMakeArray :: Comp -> Sz ix -> (ix -> e) -> Array r ix e + -- | Construct an Array. Resulting type either has to be unambiguously inferred or restricted + -- manually, like in the example below. Use "Data.Massiv.Array.makeArrayR" if you'd like to + -- specify representation as an argument. + -- + -- >>> import Data.Massiv.Array + -- >>> makeArray Seq (Sz (3 :. 4)) (\ (i :. j) -> if i == j then i else 0) :: Array D Ix2 Int + -- Array D Seq (Sz (3 :. 4)) + -- [ [ 0, 0, 0, 0 ] + -- , [ 0, 1, 0, 0 ] + -- , [ 0, 0, 2, 0 ] + -- ] + -- + -- Instead of restricting the full type manually we can use `TypeApplications` as convenience: + -- + -- >>> :set -XTypeApplications + -- >>> makeArray @P @_ @Double Seq (Sz2 3 4) $ \(i :. j) -> logBase (fromIntegral i) (fromIntegral j) + -- Array P Seq (Sz (3 :. 4)) + -- [ [ NaN, -0.0, -0.0, -0.0 ] + -- , [ -Infinity, NaN, Infinity, Infinity ] + -- , [ -Infinity, 0.0, 1.0, 1.5849625007211563 ] + -- ] + -- + -- @since 0.1.0 + makeArray :: + Comp -- ^ Computation strategy. Useful constructors are `Seq` and `Par` + -> Sz ix -- ^ Size of the result array. + -> (ix -> e) -- ^ Function to generate elements at a particular index + -> Array r ix e + makeArray comp sz f = makeArrayLinear comp sz (f . fromLinearIndex sz) + {-# INLINE makeArray #-} + -- | Same as `makeArray`, but produce elements using linear row-major index. + -- + -- >>> import Data.Massiv.Array + -- >>> makeArrayLinear Seq (Sz (2 :. 4)) id :: Array D Ix2 Int + -- Array D Seq (Sz (2 :. 4)) + -- [ [ 0, 1, 2, 3 ] + -- , [ 4, 5, 6, 7 ] + -- ] + -- + -- @since 0.3.0 + makeArrayLinear :: Comp -> Sz ix -> (Int -> e) -> Array r ix e + makeArrayLinear comp sz f = makeArray comp sz (f . toLinearIndex sz) + {-# INLINE makeArrayLinear #-} --- | An array that contains size information. They can be resized and new arrays extracted from it --- in constant time. -class Construct r ix e => Size r ix e where - -- | /O(1)/ - Get the size of an array - size :: Array r ix e -> Sz ix - -- | /O(1)/ - Change the size of an array. New size is not validated. +class Index ix => Resize r ix where + -- | /O(1)/ - Change the size of an array. Total number of elements should be the same, but it is + -- not validated. unsafeResize :: Index ix' => Sz ix' -> Array r ix e -> Array r ix' e + +class Load r ix e => Extract r ix e where -- | /O(1)/ - Extract a portion of an array. Staring index and new size are -- not validated. unsafeExtract :: ix -> Sz ix -> Array r ix e -> Array (EltRepr r ix) ix e -- | Arrays that can be used as source to practically any manipulation function. -class Size r ix e => Source r ix e where +class Load r ix e => Source r ix e where + {-# MINIMAL (unsafeIndex|unsafeLinearIndex) #-} -- | Lookup element in the array. No bounds check is performed and access of -- arbitrary memory is possible when invalid index is supplied. @@ -122,177 +196,159 @@ unsafeLinearIndex !arr = unsafeIndex arr . fromLinearIndex (size arr) {-# INLINE unsafeLinearIndex #-} --- | Any array that can be computed -class Size r ix e => Load r ix e where +-- | Any array that can be computed and loaded into memory +class (Typeable r, Index ix) => Load r ix e where - -- | Load an array into memory sequentially - loadS + -- | Get computation strategy of this array + + getComp :: Array r ix e -> Comp + + -- | Get the size of an immutabe array + size :: Array r ix e -> Sz ix + + -- | Load an array into memory. + loadArrayM :: Monad m => - Array r ix e -- ^ Array that is being loaded - -> (Int -> m e) -- ^ Function that reads an element from target array + Scheduler m () + -> Array r ix e -- ^ Array that is being loaded -> (Int -> e -> m ()) -- ^ Function that writes an element into target array -> m () - loadS = loadArray 1 id - {-# INLINE loadS #-} - -- | Load an array into memory in parallel - loadP - :: [Int] -- ^ List of capabilities to run workers on, as described in - -- `Control.Concurrent.forkOn`. Empty list will imply all - -- capabilities, i.e. run on all cores available through @+RTS -N@. - -> Array r ix e -- ^ Array that is being loaded - -> (Int -> IO e) -- ^ Function that reads an element from target array - -> (Int -> e -> IO ()) -- ^ Function that writes an element into target array - -> IO () - loadP wIds arr unsafeRead unsafeWrite = - withScheduler_ wIds $ \scheduler -> - loadArray (numWorkers scheduler) (scheduleWork scheduler) arr unsafeRead unsafeWrite - {-# INLINE loadP #-} - - -- | Load an array into memory with stride. Default implementation can only handle the sequential - -- case and only if there is an instance of `Source`. - loadArrayWithStride +class Load r ix e => StrideLoad r ix e where + -- | Load an array into memory with stride. Default implementation requires an instance of + -- `Source`. + loadArrayWithStrideM :: Monad m => - Int -- ^ Total number of workers (for `Seq` it's always 1) - -> (m () -> m ()) -- ^ A monadic action that will schedule work for the workers (for `Seq` it's - -- always `id`) + Scheduler m () -> Stride ix -- ^ Stride to use -> Sz ix -- ^ Size of the target array affected by the stride. -> Array r ix e -- ^ Array that is being loaded - -> (Int -> m e) -- ^ Function that reads an element from target array -> (Int -> e -> m ()) -- ^ Function that writes an element into target array -> m () - default loadArrayWithStride + default loadArrayWithStrideM :: (Source r ix e, Monad m) => - Int - -> (m () -> m ()) + Scheduler m () -> Stride ix -> Sz ix -> Array r ix e - -> (Int -> m e) -> (Int -> e -> m ()) -> m () - loadArrayWithStride numWorkers' scheduleWork' stride resultSize arr _ = - splitLinearlyWith_ numWorkers' scheduleWork' (totalElem resultSize) unsafeLinearWriteWithStride + loadArrayWithStrideM scheduler stride resultSize arr = + splitLinearlyWith_ scheduler (totalElem resultSize) unsafeLinearWriteWithStride where - strideIx = unStride stride + !strideIx = unStride stride unsafeLinearWriteWithStride = unsafeIndex arr . liftIndex2 (*) strideIx . fromLinearIndex resultSize {-# INLINE unsafeLinearWriteWithStride #-} - {-# INLINE loadArrayWithStride #-} + {-# INLINE loadArrayWithStrideM #-} - -- TODO: this is the future replacement for loadS and loadP discussed in: - -- https://github.com/lehins/massiv/issues/41 - -- | Load an array into memory. Default implementation will respect the scheduler and use `Source` - -- instance to do loading in row-major fashion in parallel as well as sequentially. - loadArray - :: Monad m => - Int -- ^ Total number of workers (for `Seq` it's always 1) - -> (m () -> m ()) -- ^ A monadic action that will schedule work for the workers (for `Seq` it's - -- always `id`) - -> Array r ix e -- ^ Array that is being loaded - -> (Int -> m e) -- ^ Function that reads an element from target array - -> (Int -> e -> m ()) -- ^ Function that writes an element into target array - -> m () - default loadArray - :: (Source r ix e, Monad m) => - Int - -> (m () -> m ()) - -> Array r ix e - -> (Int -> m e) - -> (Int -> e -> m ()) - -> m () - loadArray numWorkers' scheduleWork' arr _ = - splitLinearlyWith_ numWorkers' scheduleWork' (totalElem (size arr)) (unsafeLinearIndex arr) - {-# INLINE loadArray #-} -class OuterSlice r ix e where +class Load r ix e => OuterSlice r ix e where -- | /O(1)/ - Take a slice out of an array from the outside unsafeOuterSlice :: Array r ix e -> Int -> Elt r ix e - outerLength :: Array r ix e -> Int - default outerLength :: Size r ix e => Array r ix e -> Int - outerLength = headDim . size - -class Size r ix e => InnerSlice r ix e where - unsafeInnerSlice :: Array r ix e -> (Sz (Lower ix), Sz1) -> Int -> Elt r ix e +class Load r ix e => InnerSlice r ix e where + unsafeInnerSlice :: Array r ix e -> (Sz (Lower ix), Sz Int) -> Int -> Elt r ix e -class Size r ix e => Slice r ix e where - unsafeSlice :: Array r ix e -> ix -> Sz ix -> Dim -> Maybe (Elt r ix e) +class Load r ix e => Slice r ix e where + unsafeSlice :: MonadThrow m => Array r ix e -> ix -> Sz ix -> Dim -> m (Elt r ix e) -- | Manifest arrays are backed by actual memory and values are looked up versus -- computed as it is with delayed arrays. Because of this fact indexing functions -- @(`!`)@, @(`!?`)@, etc. are constrained to manifest arrays only. -class Source r ix e => Manifest r ix e where +class (Load r ix e, Source r ix e) => Manifest r ix e where unsafeLinearIndexM :: Array r ix e -> Int -> e -data State s = State (State# s) - -type WorldState = State RealWorld - - class Manifest r ix e => Mutable r ix e where data MArray s r ix e :: * -- | Get the size of a mutable array. + -- + -- @since 0.1.0 msize :: MArray s r ix e -> Sz ix + -- | Convert immutable array into a mutable array without copy. + -- + -- @since 0.1.0 unsafeThaw :: PrimMonad m => Array r ix e -> m (MArray (PrimState m) r ix e) + -- | Convert mutable array into an immutable array without copy. + -- + -- @since 0.1.0 unsafeFreeze :: PrimMonad m => Comp -> MArray (PrimState m) r ix e -> m (Array r ix e) -- | Create new mutable array, leaving it's elements uninitialized. Size isn't validated -- either. + -- + -- @since 0.1.0 unsafeNew :: PrimMonad m => Sz ix -> m (MArray (PrimState m) r ix e) - -- | Create new mutable array, leaving it's elements uninitialized. Size isn't validated - -- either. - unsafeNewZero :: PrimMonad m => - Sz ix -> m (MArray (PrimState m) r ix e) - + -- | Read an element at linear row-major index + -- + -- @since 0.1.0 unsafeLinearRead :: PrimMonad m => MArray (PrimState m) r ix e -> Int -> m e + -- | Write an element into mutable array with linear row-major index + -- + -- @since 0.1.0 unsafeLinearWrite :: PrimMonad m => MArray (PrimState m) r ix e -> Int -> e -> m () - -- | Create new mutable array, leaving it's elements uninitialized. Size isn't validated - -- either. - unsafeNewA :: Applicative f => ix -> WorldState -> f (WorldState, MArray RealWorld r ix e) - unsafeNewA sz (State s#) = - case internal (unsafeNew sz :: IO (MArray RealWorld r ix e)) s# of - (# s'#, ma #) -> pure (State s'#, ma) - {-# INLINE unsafeNewA #-} + -- | Initialize mutable array to some default value. + -- + -- @since 0.3.0 + initialize :: PrimMonad m => MArray (PrimState m) r ix e -> m () - unsafeThawA :: Applicative m => - Array r ix e -> WorldState -> m (WorldState, MArray RealWorld r ix e) - unsafeThawA arr (State s#) = - case internal (unsafeThaw arr :: IO (MArray RealWorld r ix e)) s# of - (# s'#, ma #) -> pure (State s'#, ma) - {-# INLINE unsafeThawA #-} + -- | Create new mutable array while initializing all elements to some default value. + -- + -- @since 0.3.0 + initializeNew :: PrimMonad m => + Maybe e -> Sz ix -> m (MArray (PrimState m) r ix e) + initializeNew mdef sz = do + marr <- unsafeNew sz + case mdef of + Just val -> unsafeLinearSet marr 0 (totalElem sz) val + Nothing -> initialize marr + return marr + {-# INLINE initializeNew #-} - unsafeFreezeA :: Applicative m => - Comp -> MArray RealWorld r ix e -> WorldState -> m (WorldState, Array r ix e) - unsafeFreezeA comp marr (State s#) = - case internal (unsafeFreeze comp marr :: IO (Array r ix e)) s# of - (# s'#, a #) -> pure (State s'#, a) - {-# INLINE unsafeFreezeA #-} - unsafeLinearWriteA :: Applicative m => - MArray RealWorld r ix e -> Int -> e -> WorldState -> m WorldState - unsafeLinearWriteA marr i val (State s#) = - case internal (unsafeLinearWrite marr i val :: IO ()) s# of - (# s'#, _ #) -> pure (State s'#) - {-# INLINE unsafeLinearWriteA #-} + unsafeLinearSet :: PrimMonad m => + MArray (PrimState m) r ix e -> Int -> Int -> e -> m () + unsafeLinearSet marr offset len e = + loopM_ offset (< (offset + len)) (+1) (\i -> unsafeLinearWrite marr i e) + {-# INLINE unsafeLinearSet #-} +-- | Read an array element +unsafeRead :: (Mutable r ix e, PrimMonad m) => + MArray (PrimState m) r ix e -> ix -> m e +unsafeRead !marr !ix = unsafeLinearRead marr (toLinearIndex (msize marr) ix) +{-# INLINE unsafeRead #-} +-- | Write an element into array +unsafeWrite :: (Mutable r ix e, PrimMonad m) => + MArray (PrimState m) r ix e -> ix -> e -> m () +unsafeWrite !marr !ix = unsafeLinearWrite marr (toLinearIndex (msize marr) ix) +{-# INLINE unsafeWrite #-} + +-- | Modify an element in the array with an index aware action. +unsafeLinearModify :: (Mutable r ix e, PrimMonad m) => + MArray (PrimState m) r ix e -> (Int -> e -> m e) -> Int -> m () +unsafeLinearModify !marr f !i = do + v <- unsafeLinearRead marr i + v' <- f i v + unsafeLinearWrite marr i v' +{-# INLINE unsafeLinearModify #-} + + class Nested r ix e where fromNested :: NestedStruct r ix e -> Array r ix e @@ -301,26 +357,22 @@ class Construct r ix e => Ragged r ix e where - empty :: Comp -> Array r ix e + emptyR :: Comp -> Array r ix e isNull :: Array r ix e -> Bool - cons :: Elt r ix e -> Array r ix e -> Array r ix e - - uncons :: Array r ix e -> Maybe (Elt r ix e, Array r ix e) - - -- head :: Array r ix e -> Maybe (Elt r ix e, Array r ix e) + consR :: Elt r ix e -> Array r ix e -> Array r ix e - -- tail :: Array r ix e -> Maybe (Elt r ix e, Array r ix e) + unconsR :: Array r ix e -> Maybe (Elt r ix e, Array r ix e) - unsafeGenerateM :: Monad m => Comp -> Sz ix -> (ix -> m e) -> m (Array r ix e) + generateRaggedM :: Monad m => Comp -> Sz ix -> (ix -> m e) -> m (Array r ix e) edgeSize :: Array r ix e -> Sz ix flatten :: Array r ix e -> Array r Ix1 e loadRagged :: - (IO () -> IO ()) -> (Int -> e -> IO a) -> Int -> Int -> ix -> Array r ix e -> IO () + Monad m => (m () -> m ()) -> (Int -> e -> m a) -> Int -> Int -> Sz ix -> Array r ix e -> m () -- TODO: test property: -- (read $ raggedFormat show "\n" (ls :: Array L (IxN n) Int)) == ls @@ -328,151 +380,324 @@ --- | Create an Array. Resulting type either has to be unambiguously inferred or restricted manually, --- like in the example below. +-- | Create an Array with no elements. By itself it is not particularly useful, but it serves as a +-- nice base for constructing larger arrays. -- --- >>> makeArray Seq (Sz (3 :. 4)) (\ (i :. j) -> if i == j then i else 0) :: Array D Ix2 Int --- (Array D Seq (3 :. 4) --- [ [ 0,0,0,0 ] --- , [ 0,1,0,0 ] --- , [ 0,0,2,0 ] --- ]) +-- ==== __Examples__ -- -makeArray :: Construct r ix e => - Comp -- ^ Computation strategy. Useful constructors are `Seq` and `Par` - -> Sz ix -- ^ Size of the result array. Negative values will result in an empty array. - -> (ix -> e) -- ^ Function to generate elements at a particular index - -> Array r ix e -makeArray !c = unsafeMakeArray c . liftIndex (max 0) -{-# INLINE makeArray #-} - +-- >>> import Data.Massiv.Array as A +-- >>> :set -XTypeApplications +-- >>> xs = empty @DL @Ix1 @Double +-- >>> snoc (cons 4 (cons 5 xs)) 22 +-- Array DL Seq (Sz1 3) +-- [ 4.0, 5.0, 22.0 ] +-- +-- @since 0.3.0 +empty :: + forall r ix e. Construct r ix e + => Array r ix e +empty = makeArray Seq zeroSz (const (throwImpossible Uninitialized)) +{-# INLINE empty #-} -- | Create an Array with a single element. -singleton :: Construct r ix e => - Comp -- ^ Computation strategy - -> e -- ^ The element - -> Array r ix e -singleton !c = unsafeMakeArray c (pureIndex 1) . const +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> singleton 7 :: Array D Ix4 Double +-- Array D Seq (Sz (1 :> 1 :> 1 :. 1)) +-- [ [ [ [ 7.0 ] +-- ] +-- ] +-- ] +-- +-- Instead of specifying type signature we could use @TypeApplications@ +-- +-- >>> :set -XTypeApplications +-- >>> singleton @U @Ix4 @Double 7 +-- Array U Seq (Sz (1 :> 1 :> 1 :. 1)) +-- [ [ [ [ 7.0 ] +-- ] +-- ] +-- ] +-- +-- @since 0.1.0 +singleton :: + forall r ix e. Construct r ix e + => e -- ^ The only element + -> Array r ix e +singleton = makeArray Seq oneSz . const {-# INLINE singleton #-} infixl 4 !, !?, ?? -- | Infix version of `index'`. +-- +-- >>> import Data.Massiv.Array as A +-- >>> a = computeAs U $ iterateN Seq (Sz (2 :. 3)) succ (0 :: Int) +-- >>> a +-- Array U Seq (Sz (2 :. 3)) +-- [ [ 1, 2, 3 ] +-- , [ 4, 5, 6 ] +-- ] +-- >>> a ! 0 :. 2 +-- 3 +-- >>> a ! 0 :. 3 +-- *** Exception: IndexOutOfBoundsException: (0 :. 3) not safe for (Sz (2 :. 3)) +-- +-- @since 0.1.0 (!) :: Manifest r ix e => Array r ix e -> ix -> e (!) = index' {-# INLINE (!) #-} -- | Infix version of `index`. -(!?) :: Manifest r ix e => Array r ix e -> ix -> Maybe e -(!?) = index +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> :set -XTypeApplications +-- >>> a <- fromListsM @U @Ix2 @Int Seq [[1,2,3],[4,5,6]] +-- >>> a +-- Array U Seq (Sz (2 :. 3)) +-- [ [ 1, 2, 3 ] +-- , [ 4, 5, 6 ] +-- ] +-- >>> a !? 0 :. 2 +-- 3 +-- >>> a !? 0 :. 3 +-- *** Exception: IndexOutOfBoundsException: (0 :. 3) not safe for (Sz (2 :. 3)) +-- >>> a !? 0 :. 3 :: Maybe Int +-- Nothing +-- +-- @since 0.1.0 +(!?) :: (Manifest r ix e, MonadThrow m) => Array r ix e -> ix -> m e +(!?) = indexM {-# INLINE (!?) #-} --- | /O(1)/ - Lookup an element in the array, where array can itself be --- `Nothing`. This operator is useful when used together with slicing or other --- functions that return `Maybe` array: +-- | /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. -- --- >>> (fromList Seq [[[1,2,3]],[[4,5,6]]] :: Maybe (Array U Ix3 Int)) ??> 1 ?? (0 :. 2) +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array as A +-- >>> :set -XTypeApplications +-- >>> ma = fromListsM @U @Ix3 @Int @Maybe Seq [[[1,2,3]],[[4,5,6]]] +-- >>> ma +-- Just (Array U Seq (Sz (2 :> 1 :. 3)) +-- [ [ [ 1, 2, 3 ] +-- ] +-- , [ [ 4, 5, 6 ] +-- ] +-- ] +-- ) +-- >>> ma ??> 1 +-- Just (Array M Seq (Sz (1 :. 3)) +-- [ [ 4, 5, 6 ] +-- ] +-- ) +-- >>> ma ??> 1 ?? 0 :. 2 -- Just 6 +-- >>> ma ?? 1 :> 0 :. 2 +-- Just 6 -- -(??) :: Manifest r ix e => Maybe (Array r ix e) -> ix -> Maybe e -(??) Nothing = const Nothing -(??) (Just arr) = (arr !?) +-- @since 0.1.0 +(??) :: (Manifest r ix e, MonadThrow m) => m (Array r ix e) -> ix -> m e +(??) marr ix = marr >>= (!? ix) {-# INLINE (??) #-} --- | /O(1)/ - Lookup an element in the array. Returns `Nothing`, when index is out --- of bounds, `Just` element otherwise. +-- | /O(1)/ - Lookup an element in the array. Returns `Nothing`, when index is out of bounds and +-- returns the element at the supplied index otherwise. Use `indexM` instead, since it is more +-- generaland can just as well be used with `Maybe`. +-- +-- @since 0.1.0 index :: Manifest r ix e => Array r ix e -> ix -> Maybe e -index arr = handleBorderIndex (Fill Nothing) (size arr) (Just . unsafeIndex arr) +index = indexM {-# INLINE index #-} --- | /O(1)/ - Lookup an element in the array, while using default element when --- index is out of bounds. +-- | /O(1)/ - Lookup an element in the array. Throws `IndexOutOfBoundsException`, when index is out +-- of bounds and returns the element at the supplied index otherwise. +-- +-- @since 0.3.0 +indexM :: (Manifest r ix e, MonadThrow m) => Array r ix e -> ix -> m e +indexM = evaluateM +{-# INLINE indexM #-} + +-- | /O(1)/ - Lookup an element in the array, while using default element when index is out of +-- bounds. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array +-- >>> :set -XOverloadedLists +-- >>> xs = [0..100] :: Array P Ix1 Int +-- >>> defaultIndex 999 xs 100 +-- 100 +-- >>> defaultIndex 999 xs 101 +-- 999 +-- +-- @since 0.1.0 defaultIndex :: Manifest r ix e => e -> Array r ix e -> ix -> e defaultIndex defVal = borderIndex (Fill defVal) {-# INLINE defaultIndex #-} -- | /O(1)/ - Lookup an element in the array. Use a border resolution technique -- when index is out of bounds. +-- +-- >>> import Data.Massiv.Array as A +-- >>> :set -XOverloadedLists +-- >>> xs = [0..100] :: Array U Ix1 Int +-- >>> borderIndex Wrap xs <$> range Seq 99 104 +-- Array D Seq (Sz1 5) +-- [ 99, 100, 0, 1, 2 ] +-- +-- @since 0.1.0 borderIndex :: Manifest r ix e => Border e -> Array r ix e -> ix -> e borderIndex border arr = handleBorderIndex border (size arr) (unsafeIndex arr) {-# INLINE borderIndex #-} --- | /O(1)/ - Lookup an element in the array. Throw an error if index is out of bounds. +-- | /O(1)/ - Lookup an element in the array. This is a partial function and it can throw +-- `IndexOutOfBoundsException` inside pure code. It is safer to use `index` instead. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array +-- >>> :set -XOverloadedLists +-- >>> xs = [0..100] :: Array U Ix1 Int +-- >>> index' xs 50 +-- 50 +-- >>> index' xs 150 +-- *** Exception: IndexOutOfBoundsException: 150 not safe for (Sz1 101) +-- +-- @since 0.1.0 index' :: Manifest r ix e => Array r ix e -> ix -> e -index' arr ix = - borderIndex (Fill (errorIx "Data.Massiv.Array.index" (size arr) ix)) arr ix +index' = evaluate' {-# INLINE index' #-} - --- | This is just like `index'` function, but it allows getting values from --- delayed arrays as well as manifest. As the name suggests, indexing into a +-- | This is just like `indexM` function, but it allows getting values from +-- delayed arrays as well as `Manifest`. As the name suggests, indexing into a -- delayed array at the same index multiple times will cause evaluation of the -- value each time and can destroy the performace if used without care. -evaluateAt :: Source r ix e => Array r ix e -> ix -> e -evaluateAt !arr !ix = +-- +-- ==== __Examples__ +-- +-- >>> import Control.Exception +-- >>> import Data.Massiv.Array +-- >>> evaluateM (range Seq (Ix2 10 20) (100 :. 210)) 50 :: Either SomeException Ix2 +-- Right (60 :. 70) +-- >>> evaluateM (range Seq (Ix2 10 20) (100 :. 210)) 150 :: Either SomeException Ix2 +-- Left (IndexOutOfBoundsException: (150 :. 150) not safe for (Sz (90 :. 190))) +-- +-- @since 0.3.0 +evaluateM :: (Source r ix e, MonadThrow m) => Array r ix e -> ix -> m e +evaluateM arr ix = handleBorderIndex - (Fill (errorIx "Data.Massiv.Array.evaluateAt" (size arr) ix)) + (Fill (throwM (IndexOutOfBoundsException (size arr) ix))) (size arr) + (pure . unsafeIndex arr) + ix +{-# INLINE evaluateM #-} + +-- | Similar to `evaluateM`, but will throw an exception in pure code. +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array +-- >>> evaluate' (range Seq (Ix2 10 20) (100 :. 210)) 50 +-- 60 :. 70 +-- >>> evaluate' (range Seq (Ix2 10 20) (100 :. 210)) 150 +-- *** Exception: IndexOutOfBoundsException: (150 :. 150) not safe for (Sz (90 :. 190)) +-- +-- @since 0.3.0 +evaluate' :: Source r ix e => Array r ix e -> ix -> e +evaluate' arr ix = + handleBorderIndex + (Fill (throw (IndexOutOfBoundsException (size arr) ix))) + (size arr) (unsafeIndex arr) ix +{-# INLINE evaluate' #-} + +-- | See `evaluate'`. +-- +-- @since 0.1.0 +evaluateAt :: Source r ix e => Array r ix e -> ix -> e +evaluateAt = evaluate' {-# INLINE evaluateAt #-} +{-# DEPRECATED evaluateAt "In favor of a safe `evaluateM` or an equivalent `evaluate'`" #-} +-- | This is only used together with the @unsafe-checks@ cabal flag indexWith :: Index ix => String -- ^ Source file name, eg. __FILE__ -> Int -- ^ Line number in th source file, eg. __LINE__ -> String - -> (arr -> ix) -- ^ Get size of the array + -> (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 (getSize arr) ix = f arr ix - | otherwise = errorIx ("<" ++ fileName ++ ":" ++ show lineNo ++ "> " ++ funName) (getSize arr) ix -{-# INLINE indexWith #-} - - --- indexWith :: Size r ix e => String -> Int -> (Array r ix e -> ix -> e) -> Array r ix e -> ix -> e --- indexWith fileName lineNo f arr ix --- | isSafeIndex (size arr) ix = f arr ix --- | otherwise = errorIx ("<" ++ fileName ++ ":" ++ show lineNo ++ "> indexWith") (size arr) ix --- {-# INLINE indexWith #-} - --- errorImpossible :: String -> a --- errorImpossible loc = --- error $ "Please report this error. Impossible happend at: " ++ loc --- {-# NOINLINE errorImpossible #-} +indexWith fileName lineNo funName getSize' f arr ix + | isSafeIndex (getSize' arr) ix = f arr ix + | otherwise = errorIx ("<" ++ fileName ++ ":" ++ show lineNo ++ "> " ++ funName) (getSize' arr) ix +{-# NOINLINE indexWith #-} +-- | Helper function for throwing out of bounds error. Used by `indexWith` +errorIx :: (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 #-} -- | Map a monadic index aware function over an array sequentially, while discarding the result. -- -- ==== __Examples__ -- --- >>> imapM_ (curry print) $ range 10 15 +-- >>> import Data.Massiv.Array +-- >>> imapM_ (curry print) $ range Seq (Ix1 10) 15 -- (0,10) -- (1,11) -- (2,12) -- (3,13) -- (4,14) -- +-- @since 0.1.0 imapM_ :: (Source r ix a, Monad m) => (ix -> a -> m b) -> Array r ix a -> m () imapM_ f !arr = - iterM_ zeroIndex (size arr) (pureIndex 1) (<) $ \ !ix -> f ix (unsafeIndex arr ix) + iterM_ zeroIndex (unSz (size arr)) (pureIndex 1) (<) $ \ !ix -> f ix (unsafeIndex arr ix) {-# INLINE imapM_ #-} -- | /O(1)/ - Get the number of elements in the array -elemsCount :: Size r ix e => Array r ix e -> Int +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array +-- >>> elemsCount $ range Seq (Ix1 10) 15 +-- 5 +-- +-- @since 0.1.0 +elemsCount :: Load r ix e => Array r ix e -> Int elemsCount = totalElem . size {-# INLINE elemsCount #-} -- | /O(1)/ - Check if array has no elements. -isEmpty :: Size r ix e => Array r ix e -> Bool +-- +-- ==== __Examples__ +-- +-- >>> import Data.Massiv.Array +-- >>> isEmpty $ range Seq (Ix2 10 20) (11 :. 21) +-- False +-- >>> isEmpty $ range Seq (Ix2 10 20) (10 :. 21) +-- True +-- +-- @since 0.1.0 +isEmpty :: Load r ix e => Array r ix e -> Bool isEmpty !arr = 0 == elemsCount arr {-# INLINE isEmpty #-} diff --git a/src/Data/Massiv/Core/Computation.hs b/src/Data/Massiv/Core/Computation.hs deleted file mode 100644 --- a/src/Data/Massiv/Core/Computation.hs +++ /dev/null @@ -1,69 +0,0 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE ScopedTypeVariables #-} --- | --- Module : Data.Massiv.Core.Computation --- Copyright : (c) Alexey Kuleshevich 2018 --- License : BSD3 --- Maintainer : Alexey Kuleshevich --- Stability : experimental --- Portability : non-portable --- -module Data.Massiv.Core.Computation - ( Comp(..) - , pattern Par - ) where - -import Control.DeepSeq (NFData (..), deepseq) -#if !MIN_VERSION_base(4,11,0) -import Data.Semigroup -#endif - - --- | Computation type to use. -data Comp - = Seq -- ^ Sequential computation - | ParOn [Int] - -- ^ Use `Par` instead to use your CPU to the fullest. Also don't forget to compile - -- the program with @-threaded@ flag. - -- - -- Parallel computation with a list of capabilities to run computation - -- on. Specifying an empty list (@ParOn []@) or using `Par` will result in - -- utilization of all available capabilities, which are set at runtime by - -- @+RTS -Nx@ or at compile time by GHC flag @-with-rtsopts=-Nx@, - -- where @x@ is the number of capabilities. Ommiting @x@ in above flags - -- defaults to number available cores. - deriving (Show, Eq) - --- | Parallel computation using all available cores. -pattern Par :: Comp -pattern Par <- ParOn [] where - Par = ParOn [] - -instance NFData Comp where - rnf comp = - case comp of - Seq -> () - Par -> () - ParOn wIds -> wIds `deepseq` () - {-# INLINE rnf #-} - -instance Monoid Comp where - mempty = Seq - {-# INLINE mempty #-} - mappend = joinComp - {-# INLINE mappend #-} - -instance Semigroup Comp where - (<>) = joinComp - {-# INLINE (<>) #-} - - -joinComp :: Comp -> Comp -> Comp -joinComp Par _ = Par -joinComp _ Par = Par -joinComp (ParOn w1) (ParOn w2) = ParOn $ w1 ++ w2 -joinComp c@(ParOn _) _ = c -joinComp _ c@(ParOn _) = c -joinComp _ _ = Seq -{-# INLINE joinComp #-} diff --git a/src/Data/Massiv/Core/Exception.hs b/src/Data/Massiv/Core/Exception.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Core/Exception.hs @@ -0,0 +1,40 @@ +{-# LANGUAGE GADTs #-} +module Data.Massiv.Core.Exception + ( ImpossibleException(..) + , throwImpossible + , Uninitialized(..) + , guardNumberOfElements + ) where + +import Control.Exception +import Control.Monad +import Control.Monad.Catch +import Data.Massiv.Core.Index.Internal + +newtype ImpossibleException = + ImpossibleException SomeException + deriving (Show) + +throwImpossible :: Exception e => e -> a +throwImpossible = throw . ImpossibleException . toException +{-# NOINLINE throwImpossible #-} + +instance Exception ImpossibleException where + displayException (ImpossibleException exc) = + " ImpossibleException (" ++ + displayException exc ++ + "): Either one of the unsafe functions was used or it is a bug in the library. " ++ + "In latter case please report this error." + +-- | An error that gets thrown when an unitialized element of a boxed array gets accessed. Can only +-- happen when array was constructed with `unsafeNew`. +data Uninitialized = Uninitialized deriving Show + +instance Exception Uninitialized where + displayException Uninitialized = "Array element is uninitialized" + + +guardNumberOfElements :: (MonadThrow m, Index ix, Index ix') => Sz ix -> Sz ix' -> m () +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 @@ -5,16 +5,45 @@ {-# LANGUAGE TypeOperators #-} -- | -- Module : Data.Massiv.Core.Index --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 --- Maintainer : Alexey Kuleshevich +-- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable -- module Data.Massiv.Core.Index - ( module Data.Massiv.Core.Index.Ix - , Stride - , pattern Stride + ( Ix0(..) + , type Ix1 + , pattern Ix1 + , type Ix2(Ix2, (:.)) + , IxN((:>), Ix3, Ix4, Ix5) + , 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 + , consSz + , unconsSz + , snocSz + , unsnocSz + , setSzM + , insertSzM + , pullOutSzM + , Dim(..) + , Dimension(Dim1, Dim2, Dim3, Dim4, Dim5, DimN) + , IsIndexDimension + , Lower + , Index(..) + , Stride(Stride) , unStride , toLinearIndexStride , strideStart @@ -22,34 +51,48 @@ , oneStride , Border(..) , handleBorderIndex - , module Data.Massiv.Core.Index.Class , zeroIndex - , isSafeSize + , oneIndex , isNonEmpty , headDim , tailDim , lastDim , initDim - , getIndex' - , setIndex' , getDim' + , getDim , setDim' + , setDim + , dropDimM , dropDim' + , dropDim , pullOutDim' + , pullOutDim , insertDim' + , insertDim , fromDimension , getDimension , setDimension , dropDimension , pullOutDimension , insertDimension + -- * Iterators + , iter , iterLinearM , iterLinearM_ , module Data.Massiv.Core.Iterator + , module Data.Massiv.Core.Index.Tuple + -- * Exceptions + , IndexException(..) + , SizeException(..) + , ShapeException(..) ) where +import Control.Exception (throw) import Control.DeepSeq -import Data.Massiv.Core.Index.Class +import Control.Monad.Catch (MonadThrow(..)) +import Data.Functor.Identity (runIdentity) +import Data.Massiv.Core.Index.Internal +import Data.Massiv.Core.Index.Tuple import Data.Massiv.Core.Index.Ix import Data.Massiv.Core.Index.Stride import Data.Massiv.Core.Iterator @@ -107,6 +150,17 @@ -- | Apply a border resolution technique to an index +-- +-- ==== __Examples__ +-- +-- >>> handleBorderIndex (Fill 100) (Sz (2 :. 3)) id (2 :. 3) +-- 100 :. 100 +-- >>> handleBorderIndex Wrap (Sz (2 :. 3)) id (2 :. 3) +-- 0 :. 0 +-- >>> handleBorderIndex Edge (Sz (2 :. 3)) id (2 :. 3) +-- 1 :. 2 +-- +-- @since 0.1.0 handleBorderIndex :: Index ix => Border e -- ^ Broder resolution technique @@ -117,114 +171,247 @@ handleBorderIndex border !sz getVal !ix = case border of Fill val -> if isSafeIndex sz ix then getVal ix else val - Wrap -> getVal (repairIndex sz ix (flip mod) (flip mod)) - Edge -> getVal (repairIndex sz ix (const (const 0)) (\ !k _ -> k - 1)) - Reflect -> getVal (repairIndex sz ix (\ !k !i -> (abs i - 1) `mod` k) - (\ !k !i -> (-i - 1) `mod` k)) - Continue -> getVal (repairIndex sz ix (\ !k !i -> abs i `mod` k) - (\ !k !i -> (-i - 2) `mod` k)) + 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 +-- +-- ==== __Examples__ +-- +-- >>> zeroIndex :: Ix4 +-- 0 :> 0 :> 0 :. 0 +-- +-- @since 0.1.0 zeroIndex :: Index ix => ix zeroIndex = pureIndex 0 {-# INLINE [1] zeroIndex #-} --- | Checks whether the size is valid. +-- | Index with all ones -- --- __Note__ Will be removed in /massiv-0.3.0/ -isSafeSize :: Index ix => Sz ix -> Bool -isSafeSize = (zeroIndex >=) -{-# INLINE [1] isSafeSize #-} - +-- @since 0.3.0 +oneIndex :: Index ix => ix +oneIndex = pureIndex 1 +{-# INLINE [1] oneIndex #-} -- | Checks whether array with this size can hold at least one element. +-- +-- ==== __Examples__ +-- +-- >>> isNonEmpty (Sz3 1 0 2) +-- False +-- +-- @since 0.1.0 isNonEmpty :: Index ix => Sz ix -> Bool isNonEmpty !sz = isSafeIndex sz zeroIndex {-# INLINE [1] isNonEmpty #-} - +-- TODO: benchmark against (also adjust `isEmpty` with fastest): +-- - foldlIndex (*) 1 (unSz sz) /= 0 +-- - foldlIndex ((&&) . (==0)) True (unSz sz) +-- - totalElem sz == 0 +-- | Get the outmost dimension of the index. +-- +-- ==== __Examples__ +-- +-- >>> headDim (2 :> 3 :> 4 :. 5) +-- 2 +-- +-- @since 0.1.0 headDim :: Index ix => ix -> Int headDim = fst . unconsDim {-# INLINE [1] headDim #-} +-- | Drop the outmost dimension from the index +-- +-- ==== __Examples__ +-- +-- >>> tailDim (2 :> 3 :> 4 :. 5) +-- 3 :> 4 :. 5 +-- +-- @since 0.1.0 tailDim :: Index ix => ix -> Lower ix tailDim = snd . unconsDim {-# INLINE [1] tailDim #-} +-- | Get the innermost dimension from the index +-- +-- ==== __Examples__ +-- +-- >>> lastDim (2 :> 3 :> 4 :. 5) +-- 5 +-- +-- @since 0.1.0 lastDim :: Index ix => ix -> Int lastDim = snd . unsnocDim {-# INLINE [1] lastDim #-} +-- | Drop the innermost dimension from the index +-- +-- ==== __Examples__ +-- +-- >>> initDim (2 :> 3 :> 4 :. 5) +-- 2 :> 3 :. 4 +-- +-- @since 0.1.0 initDim :: Index ix => ix -> Lower ix initDim = fst . unsnocDim {-# INLINE [1] initDim #-} - - +-- | Change the value of a specific dimension within the index. Throws `IndexException`. See +-- `setDimM` for a safer version and `setDimension` for a type safe version. +-- +-- ==== __Examples__ +-- +-- >>> setDim' (2 :> 3 :> 4 :. 5) 3 10 +-- 2 :> 10 :> 4 :. 5 +-- +-- @since 0.2.4 setDim' :: Index ix => ix -> Dim -> Int -> ix -setDim' ix dim i = - case setDim ix dim i of - Just ix' -> ix' - Nothing -> errorDim "setDim'" dim +setDim' ix dim = either throw id . setDimM ix dim {-# INLINE [1] setDim' #-} +-- | See `setDimM` +-- +-- @since 0.2.4 +setDim :: Index ix => ix -> Dim -> Int -> Maybe ix +setDim = setDimM +{-# INLINE [1] setDim #-} +{-# DEPRECATED setDim "In favor of more general `setDimM`" #-} + +-- | Change the value from a specific dimension within the index. Throws `IndexException`. See +-- `getDimM` for a safer version and `getDimension` for a type safe version. +-- +-- ==== __Examples__ +-- +-- >>> getDim' (2 :> 3 :> 4 :. 5) 3 +-- 3 +-- >>> getDim' (2 :> 3 :> 4 :. 5) 0 +-- *** Exception: IndexDimensionException: (Dim 0) for 3 :> 4 :. 5 +-- +-- @since 0.2.4 getDim' :: Index ix => ix -> Dim -> Int -getDim' ix dim = - case getDim ix dim of - Just ix' -> ix' - Nothing -> errorDim "getDim'" dim +getDim' ix = either throw id . getDimM ix {-# INLINE [1] getDim' #-} --- | To be deprecated in favor of `setDim'`. -setIndex' :: Index ix => ix -> Dim -> Int -> ix -setIndex' ix dim i = - case setDim ix dim i of - Just ix' -> ix' - Nothing -> errorDim "setIndex'" dim -{-# INLINE [1] setIndex' #-} -{-# DEPRECATED setIndex' "In favor of `setDim'`" #-} +-- | See `getDimM` +-- +-- @since 0.2.4 +getDim :: Index ix => ix -> Dim -> Maybe Int +getDim = getDimM +{-# INLINE [1] getDim #-} +{-# DEPRECATED getDim "In favor of more general `getDimM`" #-} --- | To be deprecated in favor of `getDim'`. -getIndex' :: Index ix => ix -> Dim -> Int -getIndex' ix dim = - case getDim ix dim of - Just ix' -> ix' - Nothing -> errorDim "getIndex'" dim -{-# INLINE [1] getIndex' #-} -{-# DEPRECATED getIndex' "In favor of `getDim'`" #-} +-- | Remove a dimension from the index. +-- +-- ==== __Examples__ +-- +-- λ> dropDimM (2 :> 3 :> 4 :. 5) 3 :: Maybe Ix3 +-- Just (2 :> 4 :. 5) +-- λ> dropDimM (2 :> 3 :> 4 :. 5) 6 :: Maybe Ix3 +-- Nothing +-- +-- @since 0.3.0 +dropDimM :: (MonadThrow m, Index ix) => ix -> Dim -> m (Lower ix) +dropDimM ix = fmap snd . pullOutDimM ix +{-# INLINE [1] dropDimM #-} +-- | See `dropDimM` +-- +-- @since 0.1.0 +dropDim :: Index ix => ix -> Dim -> Maybe (Lower ix) +dropDim = dropDimM +{-# INLINE [1] dropDim #-} +{-# DEPRECATED dropDim "In favor of more general `dropDimM`" #-} + +-- | Remove a dimension from the index. +-- +-- ==== __Examples__ +-- +-- >>> dropDim' (2 :> 3 :> 4 :. 5) 3 +-- 2 :> 4 :. 5 +-- >>> dropDim' (2 :> 3 :> 4 :. 5) 6 +-- *** Exception: IndexDimensionException: (Dim 6) for 3 :> 4 :. 5 +-- +-- @since 0.2.4 dropDim' :: Index ix => ix -> Dim -> Lower ix -dropDim' ix dim = - case dropDim ix dim of - Just ixl -> ixl - Nothing -> errorDim "dropDim'" dim +dropDim' ix = either throw id . dropDimM ix {-# INLINE [1] dropDim' #-} +-- | Lower the dimension of the index by pulling the specified dimension. Throws `IndexException`. See +-- `pullOutDimM` for a safer version and `pullOutDimension` for a type safe version. +-- +-- ==== __Examples__ +-- +-- λ> pullOutDim' (2 :> 3 :> 4 :. 5) 3 +-- (3,2 :> 4 :. 5) +-- +-- @since 0.2.4 pullOutDim' :: Index ix => ix -> Dim -> (Int, Lower ix) -pullOutDim' ix dim = - case pullOutDim ix dim of - Just i_ixl -> i_ixl - Nothing -> errorDim "pullOutDim'" dim +pullOutDim' ix = either throw id . pullOutDimM ix {-# INLINE [1] pullOutDim' #-} +-- | See `pullOutDimM` +-- +-- @since 0.2.4 +pullOutDim :: Index ix => ix -> Dim -> Maybe (Int, Lower ix) +pullOutDim = pullOutDimM +{-# INLINE [1] pullOutDim #-} +{-# DEPRECATED pullOutDim "In favor of more general `pullOutDimM`" #-} + +-- | Raise the dimension of the index by inserting one in the specified dimension. Throws +-- `IndexException`. See `insertDimM` for a safer version and `insertDimension` for a type safe +-- version. +-- +-- ==== __Examples__ +-- +-- >>> insertDim' (2 :> 3 :> 4 :. 5) 3 10 :: Ix5 +-- 2 :> 3 :> 10 :> 4 :. 5 +-- >>> insertDim' (2 :> 3 :> 4 :. 5) 11 10 :: Ix5 +-- *** Exception: IndexDimensionException: (Dim 11) for 4 :. 5 +-- +-- @since 0.2.4 insertDim' :: Index ix => Lower ix -> Dim -> Int -> ix -insertDim' ix dim i = - case insertDim ix dim i of - Just ix' -> ix' - Nothing -> errorDim "insertDim'" dim +insertDim' ix dim = either throw id . insertDimM ix dim {-# INLINE [1] insertDim' #-} -errorDim :: String -> Dim -> a -errorDim funName dim = error $ funName ++ ": Dimension is out of reach: " ++ show dim -{-# NOINLINE errorDim #-} +-- | See `insertDimM` +-- +-- @since 0.2.4 +insertDim :: Index ix => Lower ix -> Dim -> Int -> Maybe ix +insertDim = insertDimM +{-# INLINE [1] insertDim #-} +{-# DEPRECATED insertDim "In favor of more general `insertDimM`" #-} +-- | Get the value level `Dim` from the type level equivalent. +-- +-- ==== __Examples__ +-- +-- >>> fromDimension Dim4 +-- (Dim 4) +-- >>> :set -XDataKinds +-- >>> fromDimension (DimN :: Dimension 10) +-- (Dim 10) +-- +-- @since 0.2.4 fromDimension :: KnownNat n => Dimension n -> Dim fromDimension = fromIntegral . natVal {-# INLINE [1] fromDimension #-} -- | Type safe way to set value of index at a particular dimension. -- +-- ==== __Examples__ +-- +-- >>> setDimension (2 :> 3 :> 4 :. 5) Dim4 10 +-- 10 :> 3 :> 4 :. 5 +-- -- @since 0.2.4 setDimension :: IsIndexDimension ix n => ix -> Dimension n -> Int -> ix setDimension ix d = setDim' ix (fromDimension d) @@ -232,6 +419,11 @@ -- | Type safe way to extract value of index at a particular dimension. -- +-- ==== __Examples__ +-- +-- >>> getDimension (2 :> 3 :> 4 :. 5) Dim2 +-- 4 +-- -- @since 0.2.4 getDimension :: IsIndexDimension ix n => ix -> Dimension n -> Int getDimension ix d = getDim' ix (fromDimension d) @@ -241,6 +433,11 @@ -- | Type safe way of dropping a particular dimension, thus lowering index -- dimensionality. -- +-- ==== __Examples__ +-- +-- >>> dropDimension (2 :> 3 :> 4 :. 5) Dim2 +-- 2 :> 3 :. 5 +-- -- @since 0.2.4 dropDimension :: IsIndexDimension ix n => ix -> Dimension n -> Lower ix dropDimension ix d = dropDim' ix (fromDimension d) @@ -249,6 +446,11 @@ -- | Type safe way of pulling out a particular dimension, thus lowering index -- dimensionality and returning the value at specified dimension. -- +-- ==== __Examples__ +-- +-- >>> pullOutDimension (2 :> 3 :> 4 :. 5) Dim2 +-- (4,2 :> 3 :. 5) +-- -- @since 0.2.4 pullOutDimension :: IsIndexDimension ix n => ix -> Dimension n -> (Int, Lower ix) pullOutDimension ix d = pullOutDim' ix (fromDimension d) @@ -256,14 +458,60 @@ -- | Type safe way of inserting a particular dimension, thus raising index dimensionality. -- +-- ==== __Examples__ +-- +-- >>> insertDimension (2 :> 3 :> 4 :. 5) Dim5 10 :: Ix5 +-- 10 :> 2 :> 3 :> 4 :. 5 +-- >>> insertDimension (2 :> 3 :> 4 :. 5) Dim4 10 :: Ix5 +-- 2 :> 10 :> 3 :> 4 :. 5 +-- >>> insertDimension (2 :> 3 :> 4 :. 5) Dim3 10 :: Ix5 +-- 2 :> 3 :> 10 :> 4 :. 5 +-- >>> insertDimension (2 :> 3 :> 4 :. 5) Dim2 10 :: Ix5 +-- 2 :> 3 :> 4 :> 10 :. 5 +-- >>> insertDimension (2 :> 3 :> 4 :. 5) Dim1 10 :: Ix5 +-- 2 :> 3 :> 4 :> 5 :. 10 +-- -- @since 0.2.5 insertDimension :: IsIndexDimension ix n => Lower ix -> Dimension n -> Int -> ix insertDimension ix d = insertDim' ix (fromDimension d) {-# INLINE [1] insertDimension #-} +-- | Row-major iterator for the index. Same as `iterM`, but pure. +-- +-- ==== __Examples__ +-- +-- >>> iter (Ix1 0) 1000 1 (<) 0 (+) +-- 499500 +-- >>> iter (0 :. 0) (2 :. 3) oneIndex (<) 100 $ \ (i :. j) acc -> (acc + i) * (j + 1) +-- 3615 +-- +-- @since 0.1.0 +iter :: Index ix + => ix -- ^ Start index + -> ix -- ^ End index + -> ix -- ^ Increment + -> (Int -> Int -> Bool) -- ^ Continuation confition + -> 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 lenarly from start to end in row-major fashion with an + +-- | Iterate over N-dimensional space linearly from start to end in row-major fashion with an -- accumulator +-- +-- ==== __Examples__ +-- +-- >>> sz = Sz2 3 4 +-- >>> iterLinearM sz 0 3 1 (<) 100 $ \ k ix acc -> print (fromLinearIndex sz k == ix) >> pure (acc + k) +-- True +-- True +-- True +-- 103 +-- +-- @since 0.1.0 iterLinearM :: (Index ix, Monad m) => Sz ix -- ^ Size -> Int -- ^ Linear start @@ -278,6 +526,16 @@ {-# INLINE iterLinearM #-} -- | Same as `iterLinearM`, except without an accumulator. +-- +-- ==== __Examples__ +-- +-- >>> sz = Sz2 3 4 +-- >>> iterLinearM_ sz 0 3 1 (<) $ \ k ix -> print (toLinearIndex sz ix == k) +-- True +-- True +-- True +-- +-- @since 0.1.0 iterLinearM_ :: (Index ix, Monad m) => Sz ix -- ^ Size -> Int -- ^ Start @@ -286,6 +544,6 @@ -> (Int -> Int -> Bool) -- ^ Continuation condition -> (Int -> ix -> m ()) -- ^ Monadic action that takes index in both forms -> m () -iterLinearM_ !sz !k0 !k1 !inc cond f = +iterLinearM_ sz !k0 !k1 !inc cond f = loopM_ k0 (`cond` k1) (+ inc) $ \ !i -> f i (fromLinearIndex sz i) {-# INLINE iterLinearM_ #-} diff --git a/src/Data/Massiv/Core/Index/Class.hs b/src/Data/Massiv/Core/Index/Class.hs deleted file mode 100644 --- a/src/Data/Massiv/Core/Index/Class.hs +++ /dev/null @@ -1,533 +0,0 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE DefaultSignatures #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE GeneralizedNewtypeDeriving #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} --- | --- Module : Data.Massiv.Core.Index.Class --- Copyright : (c) Alexey Kuleshevich 2018 --- License : BSD3 --- Maintainer : Alexey Kuleshevich --- Stability : experimental --- Portability : non-portable --- -module Data.Massiv.Core.Index.Class where - -import Control.DeepSeq (NFData (..)) -import Data.Functor.Identity (runIdentity) -import Data.Massiv.Core.Iterator -import GHC.TypeLits - - --- | Helpful type synonym for migrating to newer version @0.3@ -type Sz ix = ix - --- | Helpful pattern synonym for migrating to newer version @0.3@ -pattern Sz :: ix -> Sz ix -pattern Sz ix = ix - -type Sz1 = Sz Int -pattern Sz1 :: Int -> Int -pattern Sz1 ix = ix - - --- | A way to select Array dimension at a value level. -newtype Dim = Dim Int deriving (Show, Eq, Ord, Num, Real, Integral, Enum) - --- | A way to select Array dimension at a type level. -data Dimension (n :: Nat) where - Dim1 :: Dimension 1 - Dim2 :: Dimension 2 - Dim3 :: Dimension 3 - Dim4 :: Dimension 4 - Dim5 :: Dimension 5 - DimN :: (6 <= n, KnownNat n) => Dimension n - --- | A type level constraint that ensures index is indeed valid and that supplied dimension can be --- safely used with it. -type IsIndexDimension ix n = (1 <= n, n <= Dimensions ix, Index ix, KnownNat n) - --- | Zero-dimension, i.e. a scalar. Can't really be used directly as there are no instances of --- `Index` for it, and is included for completeness. -data Ix0 = Ix0 deriving (Eq, Ord, Show) - --- | 1-dimensional index. Synonym for `Int` and `Data.Massiv.Core.Index.Ix1`. -type Ix1T = Int - --- | 2-dimensional index as tuple of `Int`s. -type Ix2T = (Int, Int) - --- | 3-dimensional index as 3-tuple of `Int`s. -type Ix3T = (Int, Int, Int) - --- | 4-dimensional index as 4-tuple of `Int`s. -type Ix4T = (Int, Int, Int, Int) - --- | 5-dimensional index as 5-tuple of `Int`s. -type Ix5T = (Int, Int, Int, Int, Int) - --- | This type family will always point to a type for a dimension that is one lower than the type --- argument. -type family Lower ix :: * - -type instance Lower Ix1T = Ix0 -type instance Lower Ix2T = Ix1T -type instance Lower Ix3T = Ix2T -type instance Lower Ix4T = Ix3T -type instance Lower Ix5T = Ix4T - --- | 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) => Index ix where - type Dimensions ix :: Nat - - -- | Dimensions of an array that has this index type, i.e. what is the dimensionality. - dimensions :: ix -> Dim - - -- | Total number of elements in an array of this size. - totalElem :: Sz ix -> Int - - -- | Prepend a dimension to the index - consDim :: Int -> Lower ix -> ix - - -- | Take a dimension from the index from the outside - unconsDim :: ix -> (Int, Lower ix) - - -- | Apppend a dimension to the index - snocDim :: Lower ix -> Int -> ix - - -- | Take a dimension from the index from the inside - unsnocDim :: ix -> (Lower ix, Int) - - -- TODO: move out of the class - -- | Remove a dimension from the index - dropDim :: ix -> Dim -> Maybe (Lower ix) - dropDim ix = fmap snd . pullOutDim ix - {-# INLINE [1] dropDim #-} - - -- | Pull out value at specified dimension from the index, thus also lowering it dimensionality. - pullOutDim :: ix -> Dim -> Maybe (Int, Lower ix) - - -- | Insert a dimension into the index - insertDim :: Lower ix -> Dim -> Int -> Maybe ix - - -- | Extract the value index has at specified dimension. - getDim :: ix -> Dim -> Maybe Int - getDim = getIndex - {-# INLINE [1] getDim #-} - - -- | Set the value for an index at specified dimension. - setDim :: ix -> Dim -> Int -> Maybe ix - setDim = setIndex - {-# INLINE [1] setDim #-} - - -- | Extract the value index has at specified dimension. To be deprecated. - getIndex :: ix -> Dim -> Maybe Int - getIndex = getDim - {-# INLINE [1] getIndex #-} - - -- | Set the value for an index at specified dimension. To be deprecated. - setIndex :: ix -> Dim -> Int -> Maybe ix - setIndex = setDim - {-# INLINE [1] setIndex #-} - - -- | Lift an `Int` to any index by replicating the value as many times as there are dimensions. - pureIndex :: Int -> ix - - -- | Zip together two indices with a function - liftIndex2 :: (Int -> Int -> Int) -> ix -> ix -> ix - - -- | Map a function over an index - liftIndex :: (Int -> Int) -> ix -> ix - liftIndex f = liftIndex2 (\_ i -> f i) (pureIndex 0) - {-# INLINE [1] liftIndex #-} - - foldlIndex :: (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 #-} - - -- | Check whether index is within the size. - isSafeIndex :: Sz ix -- ^ Size - -> ix -- ^ Index - -> Bool - default isSafeIndex :: Index (Lower ix) => ix -> ix -> Bool - isSafeIndex !sz !ix = isSafeIndex n0 i0 && isSafeIndex szL ixL - where - !(n0, szL) = unconsDim sz - !(i0, ixL) = unconsDim ix - {-# INLINE [1] isSafeIndex #-} - - -- | Convert linear index from size and index - toLinearIndex :: Sz ix -- ^ Size - -> ix -- ^ Index - -> Int - default toLinearIndex :: Index (Lower ix) => ix -> ix -> Int - toLinearIndex !sz !ix = toLinearIndex szL ixL * n + i - where !(szL, n) = unsnocDim sz - !(ixL, i) = unsnocDim ix - {-# INLINE [1] toLinearIndex #-} - - -- | Convert linear index from size and index with an accumulator. Currently is useless and will - -- likley be removed in future versions. - toLinearIndexAcc :: Int -> ix -> ix -> Int - default toLinearIndexAcc :: Index (Lower ix) => Int -> ix -> ix -> Int - toLinearIndexAcc !acc !sz !ix = toLinearIndexAcc (acc * n + i) szL ixL - where !(n, szL) = unconsDim sz - !(i, ixL) = unconsDim ix - {-# INLINE [1] toLinearIndexAcc #-} - - -- | Compute an index from size and linear index - fromLinearIndex :: Sz ix -> Int -> ix - default fromLinearIndex :: Index (Lower ix) => Sz ix -> Int -> ix - fromLinearIndex sz k = consDim q ixL - where !(q, ixL) = fromLinearIndexAcc (snd (unconsDim sz)) k - {-# INLINE [1] fromLinearIndex #-} - - -- | Compute an index from size and linear index using an accumulator, thus trying to optimize for - -- tail recursion while getting the index computed. - fromLinearIndexAcc ::Sz ix -> Int -> (Int, ix) - default fromLinearIndexAcc :: Index (Lower ix) => Sz ix -> Int -> (Int, ix) - fromLinearIndexAcc ix' !k = (q, consDim r ixL) - where !(m, ix) = unconsDim ix' - !(kL, ixL) = fromLinearIndexAcc ix k - !(q, r) = quotRem kL m - {-# INLINE [1] fromLinearIndexAcc #-} - - -- | A way to make sure index is withing the bounds for the supplied size. Takes two functions - -- that will be invoked whenever index (2nd arg) is outsize the supplied size (1st arg) - repairIndex :: Sz ix -- ^ Size - -> ix -- ^ Index - -> (Int -> Int -> Int) -- ^ Repair when below zero - -> (Int -> Int -> Int) -- ^ Repair when higher than size - -> ix - default repairIndex :: Index (Lower ix) - => Sz ix -> ix -> (Int -> Int -> Int) -> (Int -> Int -> Int) -> ix - repairIndex !sz !ix rBelow rOver = - consDim (repairIndex n i rBelow rOver) (repairIndex szL ixL rBelow rOver) - where !(n, szL) = unconsDim sz - !(i, ixL) = unconsDim ix - {-# INLINE [1] repairIndex #-} - - -- | Iterator for the index. Same as `iterM`, but pure. - iter :: ix -> ix -> ix -> (Int -> Int -> Bool) -> a -> (ix -> a -> a) -> a - iter sIx eIx incIx cond acc f = - runIdentity $ iterM sIx eIx incIx cond acc (\ix -> return . f ix) - {-# INLINE iter #-} - - -- | This function is what makes it possible to iterate over an array of any dimension. - 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 - 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) - where - !(s, sIxL) = unconsDim sIx - !(e, eIxL) = unconsDim eIx - !(inc, incIxL) = unconsDim incIx - {-# INLINE iterM #-} - - -- TODO: Implement in terms of iterM, benchmark it and remove from `Index` - -- | Same as `iterM`, but don't bother with accumulator and return value. - iterM_ :: Monad m => ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> m a) -> m () - default iterM_ :: (Index (Lower ix), Monad m) - => ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> m a) -> m () - iterM_ !sIx !eIx !incIx cond f = - loopM_ s (`cond` e) (+ inc) $ \ !i -> - iterM_ sIxL eIxL incIxL cond $ \ !ix -> - f (consDim i ix) - where - !(s, sIxL) = unconsDim sIx - !(e, eIxL) = unconsDim eIx - !(inc, incIxL) = unconsDim incIx - {-# INLINE iterM_ #-} - -{-# DEPRECATED getIndex "In favor of 'getDim'" #-} -{-# DEPRECATED setIndex "In favor of 'setDim'" #-} - -instance Index Ix1T where - type Dimensions Ix1T = 1 - dimensions _ = 1 - {-# INLINE [1] dimensions #-} - totalElem = id - {-# INLINE [1] totalElem #-} - isSafeIndex !k !i = 0 <= i && i < k - {-# INLINE [1] isSafeIndex #-} - toLinearIndex _ = id - {-# INLINE [1] toLinearIndex #-} - toLinearIndexAcc !acc m i = acc * m + i - {-# INLINE [1] toLinearIndexAcc #-} - fromLinearIndex _ = id - {-# INLINE [1] fromLinearIndex #-} - fromLinearIndexAcc n k = k `quotRem` n - {-# INLINE [1] fromLinearIndexAcc #-} - repairIndex !k !i rBelow rOver - | i < 0 = rBelow k i - | i >= k = rOver k i - | otherwise = i - {-# INLINE [1] repairIndex #-} - consDim i _ = i - {-# INLINE [1] consDim #-} - unconsDim i = (i, Ix0) - {-# INLINE [1] unconsDim #-} - snocDim _ i = i - {-# INLINE [1] snocDim #-} - unsnocDim i = (Ix0, i) - {-# INLINE [1] unsnocDim #-} - getIndex i 1 = Just i - getIndex _ _ = Nothing - {-# INLINE [1] getIndex #-} - setIndex _ 1 i = Just i - setIndex _ _ _ = Nothing - {-# INLINE [1] setIndex #-} - dropDim _ 1 = Just Ix0 - dropDim _ _ = Nothing - {-# INLINE [1] dropDim #-} - pullOutDim i 1 = Just (i, Ix0) - pullOutDim _ _ = Nothing - {-# INLINE [1] pullOutDim #-} - insertDim Ix0 1 i = Just i - insertDim _ _ _ = Nothing - {-# INLINE [1] insertDim #-} - pureIndex i = i - {-# INLINE [1] pureIndex #-} - liftIndex f = f - {-# INLINE [1] liftIndex #-} - liftIndex2 f = f - {-# INLINE [1] liftIndex2 #-} - foldlIndex f = f - {-# INLINE [1] foldlIndex #-} - iter k0 k1 inc cond = loop k0 (`cond` k1) (+inc) - {-# INLINE iter #-} - iterM k0 k1 inc cond = loopM k0 (`cond` k1) (+inc) - {-# INLINE iterM #-} - iterM_ k0 k1 inc cond = loopM_ k0 (`cond` k1) (+inc) - {-# INLINE iterM_ #-} - - -instance Index Ix2T where - type Dimensions Ix2T = 2 - dimensions _ = 2 - {-# INLINE [1] dimensions #-} - totalElem (k2, k1) = k2 * k1 - {-# INLINE [1] totalElem #-} - toLinearIndex (_, k1) (i2, i1) = k1 * i2 + i1 - {-# INLINE [1] toLinearIndex #-} - fromLinearIndex (_, k1) !i = i `quotRem` k1 - {-# INLINE [1] fromLinearIndex #-} - consDim = (,) - {-# INLINE [1] consDim #-} - unconsDim = id - {-# INLINE [1] unconsDim #-} - snocDim = (,) - {-# INLINE [1] snocDim #-} - unsnocDim = id - {-# INLINE [1] unsnocDim #-} - getIndex (i2, _) 2 = Just i2 - getIndex ( _, i1) 1 = Just i1 - getIndex _ _ = Nothing - {-# INLINE [1] getIndex #-} - setIndex (_, i1) 2 i2 = Just (i2, i1) - setIndex (i2, _) 1 i1 = Just (i2, i1) - setIndex _ _ _ = Nothing - {-# INLINE [1] setIndex #-} - dropDim (_, i1) 2 = Just i1 - dropDim (i2, _) 1 = Just i2 - dropDim _ _ = Nothing - {-# INLINE [1] dropDim #-} - pullOutDim (i2, i1) 2 = Just (i2, i1) - pullOutDim (i2, i1) 1 = Just (i1, i2) - pullOutDim _ _ = Nothing - {-# INLINE [1] pullOutDim #-} - insertDim i1 2 i2 = Just (i2, i1) - insertDim i2 1 i1 = Just (i2, i1) - insertDim _ _ _ = Nothing - {-# INLINE [1] insertDim #-} - pureIndex i = (i, i) - {-# INLINE [1] pureIndex #-} - liftIndex2 f (i2, i1) (i2', i1') = (f i2 i2', f i1 i1') - {-# INLINE [1] liftIndex2 #-} - - -instance Index Ix3T where - type Dimensions Ix3T = 3 - dimensions _ = 3 - {-# INLINE [1] dimensions #-} - totalElem (k3, k2, k1) = k3 * k2 * k1 - {-# INLINE [1] totalElem #-} - consDim i3 (i2, i1) = (i3, i2, i1) - {-# INLINE [1] consDim #-} - unconsDim (i3, i2, i1) = (i3, (i2, i1)) - {-# INLINE [1] unconsDim #-} - snocDim (i3, i2) i1 = (i3, i2, i1) - {-# INLINE [1] snocDim #-} - unsnocDim (i3, i2, i1) = ((i3, i2), i1) - {-# INLINE [1] unsnocDim #-} - getIndex (i3, _, _) 3 = Just i3 - getIndex ( _, i2, _) 2 = Just i2 - getIndex ( _, _, i1) 1 = Just i1 - getIndex _ _ = Nothing - {-# INLINE [1] getIndex #-} - setIndex ( _, i2, i1) 3 i3 = Just (i3, i2, i1) - setIndex (i3, _, i1) 2 i2 = Just (i3, i2, i1) - setIndex (i3, i2, _) 1 i1 = Just (i3, i2, i1) - setIndex _ _ _ = Nothing - {-# INLINE [1] setIndex #-} - dropDim ( _, i2, i1) 3 = Just (i2, i1) - dropDim (i3, _, i1) 2 = Just (i3, i1) - dropDim (i3, i2, _) 1 = Just (i3, i2) - dropDim _ _ = Nothing - {-# INLINE [1] dropDim #-} - pullOutDim (i3, i2, i1) 3 = Just (i3, (i2, i1)) - pullOutDim (i3, i2, i1) 2 = Just (i2, (i3, i1)) - pullOutDim (i3, i2, i1) 1 = Just (i1, (i3, i2)) - pullOutDim _ _ = Nothing - {-# INLINE [1] pullOutDim #-} - insertDim (i2, i1) 3 i3 = Just (i3, i2, i1) - insertDim (i3, i1) 2 i2 = Just (i3, i2, i1) - insertDim (i3, i2) 1 i1 = Just (i3, i2, i1) - insertDim _ _ _ = Nothing - 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 - {-# INLINE [1] dimensions #-} - totalElem !(k4, k3, k2, k1) = k4 * k3 * k2 * k1 - {-# INLINE [1] totalElem #-} - consDim i4 (i3, i2, i1) = (i4, i3, i2, i1) - {-# INLINE [1] consDim #-} - unconsDim (i4, i3, i2, i1) = (i4, (i3, i2, i1)) - {-# INLINE [1] unconsDim #-} - snocDim (i4, i3, i2) i1 = (i4, i3, i2, i1) - {-# INLINE [1] snocDim #-} - unsnocDim (i4, i3, i2, i1) = ((i4, i3, i2), i1) - {-# INLINE [1] unsnocDim #-} - getIndex (i4, _, _, _) 4 = Just i4 - getIndex ( _, i3, _, _) 3 = Just i3 - getIndex ( _, _, i2, _) 2 = Just i2 - getIndex ( _, _, _, i1) 1 = Just i1 - getIndex _ _ = Nothing - {-# INLINE [1] getIndex #-} - setIndex ( _, i3, i2, i1) 4 i4 = Just (i4, i3, i2, i1) - setIndex (i4, _, i2, i1) 3 i3 = Just (i4, i3, i2, i1) - setIndex (i4, i3, _, i1) 2 i2 = Just (i4, i3, i2, i1) - setIndex (i4, i3, i2, _) 1 i1 = Just (i4, i3, i2, i1) - setIndex _ _ _ = Nothing - {-# INLINE [1] setIndex #-} - dropDim ( _, i3, i2, i1) 4 = Just (i3, i2, i1) - dropDim (i4, _, i2, i1) 3 = Just (i4, i2, i1) - dropDim (i4, i3, _, i1) 2 = Just (i4, i3, i1) - dropDim (i4, i3, i2, _) 1 = Just (i4, i3, i2) - dropDim _ _ = Nothing - {-# INLINE [1] dropDim #-} - pullOutDim (i4, i3, i2, i1) 4 = Just (i4, (i3, i2, i1)) - pullOutDim (i4, i3, i2, i1) 3 = Just (i3, (i4, i2, i1)) - pullOutDim (i4, i3, i2, i1) 2 = Just (i2, (i4, i3, i1)) - pullOutDim (i4, i3, i2, i1) 1 = Just (i1, (i4, i3, i2)) - pullOutDim _ _ = Nothing - {-# INLINE [1] pullOutDim #-} - insertDim (i3, i2, i1) 4 i4 = Just (i4, i3, i2, i1) - insertDim (i4, i2, i1) 3 i3 = Just (i4, i3, i2, i1) - insertDim (i4, i3, i1) 2 i2 = Just (i4, i3, i2, i1) - insertDim (i4, i3, i2) 1 i1 = Just (i4, i3, i2, i1) - insertDim _ _ _ = Nothing - {-# INLINE [1] insertDim #-} - 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 - {-# INLINE [1] dimensions #-} - totalElem !(n5, n4, n3, n2, n1) = n5 * n4 * n3 * n2 * n1 - {-# INLINE [1] totalElem #-} - consDim i5 (i4, i3, i2, i1) = (i5, i4, i3, i2, i1) - {-# INLINE [1] consDim #-} - unconsDim (i5, i4, i3, i2, i1) = (i5, (i4, i3, i2, i1)) - {-# INLINE [1] unconsDim #-} - snocDim (i5, i4, i3, i2) i1 = (i5, i4, i3, i2, i1) - {-# INLINE [1] snocDim #-} - unsnocDim (i5, i4, i3, i2, i1) = ((i5, i4, i3, i2), i1) - {-# INLINE [1] unsnocDim #-} - getIndex (i5, _, _, _, _) 5 = Just i5 - getIndex ( _, i4, _, _, _) 4 = Just i4 - getIndex ( _, _, i3, _, _) 3 = Just i3 - getIndex ( _, _, _, i2, _) 2 = Just i2 - getIndex ( _, _, _, _, i1) 1 = Just i1 - getIndex _ _ = Nothing - {-# INLINE [1] getIndex #-} - setIndex ( _, i4, i3, i2, i1) 5 i5 = Just (i5, i4, i3, i2, i1) - setIndex (i5, _, i3, i2, i1) 4 i4 = Just (i5, i4, i3, i2, i1) - setIndex (i5, i4, _, i2, i1) 3 i3 = Just (i5, i4, i3, i2, i1) - setIndex (i5, i4, i3, _, i1) 2 i2 = Just (i5, i4, i3, i2, i1) - setIndex (i5, i4, i3, i2, _) 1 i1 = Just (i5, i4, i3, i2, i1) - setIndex _ _ _ = Nothing - {-# INLINE [1] setIndex #-} - dropDim ( _, i4, i3, i2, i1) 5 = Just (i4, i3, i2, i1) - dropDim (i5, _, i3, i2, i1) 4 = Just (i5, i3, i2, i1) - dropDim (i5, i4, _, i2, i1) 3 = Just (i5, i4, i2, i1) - dropDim (i5, i4, i3, _, i1) 2 = Just (i5, i4, i3, i1) - dropDim (i5, i4, i3, i2, _) 1 = Just (i5, i4, i3, i2) - dropDim _ _ = Nothing - {-# INLINE [1] dropDim #-} - pullOutDim (i5, i4, i3, i2, i1) 5 = Just (i5, (i4, i3, i2, i1)) - pullOutDim (i5, i4, i3, i2, i1) 4 = Just (i4, (i5, i3, i2, i1)) - pullOutDim (i5, i4, i3, i2, i1) 3 = Just (i3, (i5, i4, i2, i1)) - pullOutDim (i5, i4, i3, i2, i1) 2 = Just (i2, (i5, i4, i3, i1)) - pullOutDim (i5, i4, i3, i2, i1) 1 = Just (i1, (i5, i4, i3, i2)) - pullOutDim _ _ = Nothing - {-# INLINE [1] pullOutDim #-} - insertDim (i4, i3, i2, i1) 5 i5 = Just (i5, i4, i3, i2, i1) - insertDim (i5, i3, i2, i1) 4 i4 = Just (i5, i4, i3, i2, i1) - insertDim (i5, i4, i2, i1) 3 i3 = Just (i5, i4, i3, i2, i1) - insertDim (i5, i4, i3, i1) 2 i2 = Just (i5, i4, i3, i2, i1) - insertDim (i5, i4, i3, i2) 1 i1 = Just (i5, i4, i3, i2, i1) - insertDim _ _ _ = Nothing - {-# INLINE [1] insertDim #-} - pureIndex i = (i, i, i, i, i) - {-# INLINE [1] pureIndex #-} - liftIndex2 f (i5, i4, i3, i2, i1) (i5', i4', i3', i2', i1') = - (f i5 i5', f i4 i4', f i3 i3', f i2 i2', f i1 i1') - {-# INLINE [1] liftIndex2 #-} - --- | Helper function for throwing out of bounds errors -errorIx :: (Show ix, Show ix') => String -> Sz ix -> ix' -> a -errorIx fName sz ix = - error $ - fName ++ - ": Index out of bounds: (" ++ show ix ++ ") for Array of size: (" ++ show sz ++ ")" -{-# NOINLINE errorIx #-} - - --- | Helper function for throwing error when sizes do not match -errorSizeMismatch :: (Show ix, Show ix') => String -> Sz ix -> ix' -> a -errorSizeMismatch fName sz sz' = - error $ fName ++ ": Mismatch in size of arrays " ++ show sz ++ " vs " ++ show sz' -{-# NOINLINE errorSizeMismatch #-} diff --git a/src/Data/Massiv/Core/Index/Internal.hs b/src/Data/Massiv/Core/Index/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Core/Index/Internal.hs @@ -0,0 +1,696 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DefaultSignatures #-} +{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE TypeSynonymInstances #-} + +#if __GLASGOW_HASKELL__ < 820 +{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-} +#endif +-- | +-- Module : Data.Massiv.Core.Index.Internal +-- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Core.Index.Internal + ( Sz(SafeSz) + , pattern Sz + , pattern Sz1 + , type Sz1 + , unSz + , zeroSz + , oneSz + , consSz + , unconsSz + , snocSz + , unsnocSz + , setSzM + , insertSzM + , pullOutSzM + , Dim(..) + , Dimension(DimN) + , pattern Dim1 + , pattern Dim2 + , pattern Dim3 + , pattern Dim4 + , pattern Dim5 + , IsIndexDimension + , Lower + , Index(..) + , Ix0(..) + , type Ix1 + , pattern Ix1 + , IndexException(..) + , SizeException(..) + , ShapeException(..) + ) where + +import Control.DeepSeq +import Control.Exception (Exception(..)) +import Control.Monad.Catch (MonadThrow(..)) +import Data.Coerce +import Data.Massiv.Core.Iterator +import Data.Typeable +import GHC.TypeLits + +-- | `Sz` provides type safety guarantees preventing mixup with index, which is used for looking into +-- array cells, from the size, that describes total number of elements along each dimension in the +-- array. Moreover the @Sz@ constructor will prevent creation of invalid sizes with negative numbers. +-- +-- @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 + deriving (Eq, Ord, NFData) + +-- | A safe bidirectional pattern synonym for `Sz` construction that will make sure that none of +-- the size elements are negative. +-- +-- @since 0.3.0 +pattern Sz :: Index ix => ix -> Sz ix +pattern Sz ix <- SafeSz ix where + Sz ix = SafeSz (liftIndex (max 0) ix) +{-# COMPLETE Sz #-} + +-- | 1-dimensional type synonym for size. +-- +-- @since 0.3.0 +type Sz1 = Sz Ix1 + +-- | 1-dimensional size constructor. Especially useful with literals: @(Sz1 5) == Sz (5 :: Int)@. +-- +-- @since 0.3.0 +pattern Sz1 :: Ix1 -> Sz1 +pattern Sz1 ix <- SafeSz ix where + Sz1 ix = SafeSz (max 0 ix) +{-# COMPLETE Sz1 #-} + + +instance Index ix => Show (Sz ix) where + showsPrec n sz@(SafeSz usz) s = + if n == 0 + then str ++ s + else '(' : str ++ ')' : s + where + str = + "Sz" ++ + case unDim (dimensions sz) of + 1 -> "1 " ++ show usz + _ -> " (" ++ show usz ++ ")" + +instance (Num ix, Index ix) => Num (Sz ix) where + (+) x y = SafeSz (coerce x + coerce y) + {-# INLINE (+) #-} + (-) x y = Sz (coerce x - coerce y) + {-# INLINE (-) #-} + (*) x y = SafeSz (coerce x * coerce y) + {-# INLINE (*) #-} + abs !x = x + {-# INLINE abs #-} + negate !_x = 0 + {-# INLINE negate #-} + signum x = SafeSz (signum (coerce x)) + {-# INLINE signum #-} + fromInteger = Sz . fromInteger + {-# INLINE fromInteger #-} + + +-- | Function for unwrapping `Sz`. +-- +-- ==== __Example__ +-- +-- >>> import Data.Massiv.Core.Index +-- >>> unSz $ Sz3 1 2 3 +-- 1 :> 2 :. 3 +-- +-- @since 0.3.0 +unSz :: Sz ix -> ix +unSz (SafeSz ix) = ix +{-# INLINE unSz #-} + +-- | An empty size with all elements in size equal to @0@. +-- +-- ==== __Example__ +-- +-- >>> import Data.Massiv.Core.Index +-- >>> zeroSz :: Sz5 +-- Sz (0 :> 0 :> 0 :> 0 :. 0) +-- +-- @since 0.3.0 +zeroSz :: Index ix => Sz ix +zeroSz = SafeSz (pureIndex 0) +{-# INLINE zeroSz #-} + +-- | A singleton size with all elements in size equal to @1@. +-- +-- ==== __Example__ +-- +-- >>> import Data.Massiv.Core.Index +-- >>> oneSz :: Sz3 +-- Sz (1 :> 1 :. 1) +-- +-- @since 0.3.0 +oneSz :: Index ix => Sz ix +oneSz = SafeSz (pureIndex 1) +{-# INLINE oneSz #-} + + +-- | Same as `consDim`, but for `Sz` +-- +-- ==== __Example__ +-- +-- >>> import Data.Massiv.Core.Index +-- >>> consSz (Sz1 1) (Sz2 2 3) :: Sz3 +-- Sz (1 :> 2 :. 3) +-- +-- @since 0.3.0 +consSz :: Index ix => Sz1 -> Sz (Lower ix) -> Sz ix +consSz (SafeSz i) (SafeSz ix) = SafeSz (consDim i ix) +{-# INLINE consSz #-} + + +-- | Same as `snocDim`, but for `Sz` +-- +-- ==== __Example__ +-- +-- >>> import Data.Massiv.Core.Index +-- >>> snocSz (Sz2 2 3) (Sz1 1) :: Sz3 +-- Sz (2 :> 3 :. 1) +-- +-- @since 0.3.0 +snocSz :: Index ix => Sz (Lower ix) -> Sz1 -> Sz ix +snocSz (SafeSz i) (SafeSz ix) = SafeSz (snocDim i ix) +{-# INLINE snocSz #-} + +-- | Same as `setDimM`, but for `Sz` +-- +-- ==== __Example__ +-- +-- >>> import Data.Massiv.Core.Index +-- >>> setSzM (Sz2 2 3) 2 (Sz1 1) :: IO Sz2 +-- Sz (1 :. 3) +-- >>> setSzM (Sz2 2 3) 3 (Sz1 1) :: IO Sz2 +-- *** Exception: IndexDimensionException: (Dim 3) for 2 :. 3 +-- +-- @since 0.3.0 +setSzM :: (MonadThrow m, Index ix) => Sz ix -> Dim -> Sz Int -> m (Sz ix) +setSzM (SafeSz sz) dim (SafeSz sz1) = SafeSz <$> setDimM sz dim sz1 +{-# INLINE setSzM #-} + +-- | Same as `insertDimM`, but for `Sz` +-- +-- ==== __Example__ +-- +-- >>> import Data.Massiv.Core.Index +-- >>> insertSzM (Sz2 2 3) 3 (Sz1 1) :: IO Sz3 +-- Sz (1 :> 2 :. 3) +-- >>> insertSzM (Sz2 2 3) 4 (Sz1 1) :: IO Sz3 +-- *** Exception: IndexDimensionException: (Dim 4) for 2 :. 3 +-- +-- @since 0.3.0 +insertSzM :: (MonadThrow m, Index ix) => Sz (Lower ix) -> Dim -> Sz Int -> m (Sz ix) +insertSzM (SafeSz sz) dim (SafeSz sz1) = SafeSz <$> insertDimM sz dim sz1 +{-# INLINE insertSzM #-} + +-- | Same as `unconsDim`, but for `Sz` +-- +-- ==== __Example__ +-- +-- >>> import Data.Massiv.Core.Index +-- >>> unconsSz $ Sz3 1 2 3 +-- (Sz1 1,Sz (2 :. 3)) +-- +-- @since 0.3.0 +unconsSz :: Index ix => Sz ix -> (Sz1, Sz (Lower ix)) +unconsSz (SafeSz sz) = coerce (unconsDim sz) +{-# INLINE unconsSz #-} + +-- | Same as `unsnocDim`, but for `Sz` +-- +-- ==== __Example__ +-- +-- >>> import Data.Massiv.Core.Index +-- >>> unsnocSz $ Sz3 1 2 3 +-- (Sz (1 :. 2),Sz1 3) +-- +-- @since 0.3.0 +unsnocSz :: Index ix => Sz ix -> (Sz (Lower ix), Sz1) +unsnocSz (SafeSz sz) = coerce (unsnocDim sz) +{-# INLINE unsnocSz #-} + +-- | Same as `pullOutDim`, but for `Sz` +-- +-- >>> import Data.Massiv.Core.Index +-- >>> pullOutSzM (Sz3 1 2 3) 3 +-- (Sz1 1,Sz (2 :. 3)) +-- >>> pullOutSzM (Sz3 1 2 3) 0 +-- *** Exception: IndexDimensionException: (Dim 0) for 1 :> 2 :. 3 +-- +-- @since 0.3.0 +pullOutSzM :: (MonadThrow m, Index ix) => Sz ix -> Dim -> m (Sz Ix1, Sz (Lower ix)) +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) + +instance Show Dim where + show (Dim d) = "(Dim " ++ show d ++ ")" + +-- | A way to select Array dimension at a type level. +-- +-- @since 0.2.4 +data Dimension (n :: Nat) where + DimN :: (1 <= n, KnownNat n) => Dimension n + +-- | Construct 1st dimension +-- +-- @since 0.2.4 +pattern Dim1 :: Dimension 1 +pattern Dim1 = DimN + +-- | Construct 2nd dimension +-- +-- @since 0.2.4 +pattern Dim2 :: Dimension 2 +pattern Dim2 = DimN + +-- | Construct 3rd dimension +-- +-- @since 0.2.4 +pattern Dim3 :: Dimension 3 +pattern Dim3 = DimN + +-- | Construct 4th dimension +-- +-- @since 0.2.4 +pattern Dim4 :: Dimension 4 +pattern Dim4 = DimN + +-- | Construct 5th dimension +-- +-- @since 0.2.4 +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 :: * + +-- | This is bread and butter of multi-dimensional array indexing. It is unlikely that any of the +-- functions in this class will be useful to a regular user, unless general algorithms are being +-- implemented that do span multiple dimensions. +class ( Eq ix + , Ord ix + , Show ix + , NFData ix + , Eq (Lower ix) + , Ord (Lower ix) + , Show (Lower ix) + , NFData (Lower ix) + ) => + Index ix + where + -- | Type level information on how many dimensions this index has. + -- + -- @since 0.2.0 + type Dimensions ix :: Nat + + -- | What is the dimensionality of this index. + -- + -- @since 0.2.0 + dimensions :: proxy ix -> Dim + + -- | Total number of elements in an array of this size. + -- + -- @since 0.1.0 + totalElem :: Sz ix -> Int + + -- | Prepend a dimension to the index + -- + -- @since 0.1.0 + consDim :: Int -> Lower ix -> ix + + -- | Take a dimension from the index from the outside + -- + -- @since 0.1.0 + unconsDim :: ix -> (Int, Lower ix) + + -- | Apppend a dimension to the index + -- + -- @since 0.1.0 + snocDim :: Lower ix -> Int -> ix + + -- | Take a dimension from the index from the inside + -- + -- @since 0.1.0 + unsnocDim :: ix -> (Lower ix, Int) + + -- | Pull out value at specified dimension from the index, thus also lowering it dimensionality. + -- + -- @since 0.2.5 + pullOutDimM :: MonadThrow m => ix -> Dim -> m (Int, Lower ix) + + -- | Insert a dimension into the index + insertDimM :: MonadThrow m => Lower ix -> Dim -> Int -> m ix + + -- | Extract the value index has at specified dimension. + getDimM :: MonadThrow m => ix -> Dim -> m Int + + -- | Set the value for an index at specified dimension. + setDimM :: MonadThrow m => ix -> Dim -> Int -> m ix + + -- | Lift an `Int` to any index by replicating the value as many times as there are dimensions. + -- + -- @since 0.1.0 + pureIndex :: Int -> ix + + -- | Zip together two indices with a function + -- + -- @since 0.1.0 + liftIndex2 :: (Int -> Int -> Int) -> ix -> ix -> ix + + -- | Map a function over an index + -- + -- @since 0.1.0 + liftIndex :: (Int -> Int) -> ix -> ix + liftIndex f = liftIndex2 (\_ i -> f i) (pureIndex 0) + {-# INLINE [1] liftIndex #-} + + -- | 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 + 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 + -> Bool + default isSafeIndex :: Index (Lower ix) => + Sz ix -> ix -> Bool + isSafeIndex sz !ix = isSafeIndex n0 i0 && isSafeIndex szL ixL + where + !(n0, szL) = unconsSz sz + !(i0, ixL) = unconsDim ix + {-# INLINE [1] isSafeIndex #-} + + -- | Convert linear index from size and index + -- + -- @since 0.1.0 + toLinearIndex :: + Sz ix -- ^ Size + -> ix -- ^ Index + -> Int + default toLinearIndex :: Index (Lower ix) => + Sz ix -> ix -> Int + toLinearIndex (SafeSz sz) !ix = toLinearIndex (SafeSz szL) ixL * n + i + where + !(szL, n) = unsnocDim sz + !(ixL, i) = unsnocDim ix + {-# INLINE [1] toLinearIndex #-} + + -- | Convert linear index from size and index with an accumulator. Currently is useless and will + -- likley be removed in future versions. + -- + -- @since 0.1.0 + toLinearIndexAcc :: Int -> ix -> ix -> Int + default toLinearIndexAcc :: Index (Lower ix) => + Int -> ix -> ix -> Int + toLinearIndexAcc !acc !sz !ix = toLinearIndexAcc (acc * n + i) szL ixL + where + !(n, szL) = unconsDim sz + !(i, ixL) = unconsDim ix + {-# INLINE [1] toLinearIndexAcc #-} + + -- | Compute an index from size and linear index + -- + -- @since 0.1.0 + fromLinearIndex :: Sz ix -> Int -> ix + default fromLinearIndex :: Index (Lower ix) => + Sz ix -> Int -> ix + fromLinearIndex (SafeSz sz) k = consDim q ixL + where + !(q, ixL) = fromLinearIndexAcc (snd (unconsDim sz)) k + {-# INLINE [1] fromLinearIndex #-} + + -- | Compute an index from size and linear index using an accumulator, thus trying to optimize for + -- tail recursion while getting the index computed. + -- + -- @since 0.1.0 + fromLinearIndexAcc :: ix -> Int -> (Int, ix) + default fromLinearIndexAcc :: Index (Lower ix) => + ix -> Int -> (Int, ix) + fromLinearIndexAcc ix' !k = (q, consDim r ixL) + where + !(m, ix) = unconsDim ix' + !(kL, ixL) = fromLinearIndexAcc ix k + !(q, r) = quotRem kL m + {-# INLINE [1] fromLinearIndexAcc #-} + + -- | A way to make sure index is withing the bounds for the supplied size. Takes two functions + -- that will be invoked whenever index (2nd arg) is outsize the supplied size (1st arg) + -- + -- @since 0.1.0 + repairIndex :: + Sz ix -- ^ Size + -> ix -- ^ Index + -> (Sz Int -> Int -> Int) -- ^ Repair when below zero + -> (Sz Int -> Int -> Int) -- ^ Repair when higher than size + -> 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 + !(n, szL) = unconsSz sz + !(i, ixL) = unconsDim ix + {-# INLINE [1] repairIndex #-} + + -- | 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 + -> 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) + where + !(s, sIxL) = unconsDim sIx + !(e, eIxL) = unconsDim eIx + !(inc, incIxL) = unconsDim incIx + {-# INLINE iterM #-} + + -- TODO: Implement in terms of iterM, benchmark it and remove from `Index` + -- | Same as `iterM`, but don't bother with accumulator and return value. + -- + -- @since 0.1.0 + iterM_ :: Monad m => ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> m a) -> m () + default iterM_ :: (Index (Lower ix), Monad m) => + ix -> ix -> ix -> (Int -> Int -> Bool) -> (ix -> m a) -> m () + iterM_ !sIx eIx !incIx cond f = + loopM_ s (`cond` e) (+ inc) $ \ !i -> iterM_ sIxL eIxL incIxL cond $ \ !ix -> f (consDim i ix) + where + !(s, sIxL) = unconsDim sIx + !(e, eIxL) = unconsDim eIx + !(inc, incIxL) = unconsDim incIx + {-# INLINE iterM_ #-} + +-- | Zero-dimension, i.e. a scalar. Can't really be used directly as there is no instance of +-- `Index` for it, and is included for completeness. +data Ix0 = Ix0 deriving (Eq, Ord, Show) + +instance NFData Ix0 where + rnf Ix0 = () + +-- | A type synonym for 1-dimensional index, i.e. `Int`. +-- +-- >>> 5 :: Ix1 +-- 5 +-- +-- @since 0.1.0 +type Ix1 = Int + +-- | This is a very handy pattern synonym to indicate that any arbitrary `Integral` literal is an +-- `Int`, e.g. a 1-dimensional index: @(Ix1 5) == (5 :: Int)@ +-- +-- >>> Ix1 5 +-- 5 +-- >>> :t Ix1 5 +-- Ix1 5 :: Ix1 +-- +-- @since 0.1.0 +pattern Ix1 :: Int -> Ix1 +pattern Ix1 i = i + +type instance Lower Int = Ix0 + + +instance Index Ix1 where + type Dimensions Ix1 = 1 + dimensions _ = 1 + {-# INLINE [1] dimensions #-} + totalElem = unSz + {-# INLINE [1] totalElem #-} + isSafeIndex (SafeSz k) !i = 0 <= i && i < k + {-# INLINE [1] isSafeIndex #-} + toLinearIndex _ = id + {-# INLINE [1] toLinearIndex #-} + toLinearIndexAcc !acc m i = acc * m + i + {-# INLINE [1] toLinearIndexAcc #-} + fromLinearIndex _ = id + {-# INLINE [1] fromLinearIndex #-} + fromLinearIndexAcc n k = k `quotRem` n + {-# INLINE [1] fromLinearIndexAcc #-} + repairIndex k@(SafeSz ksz) !i rBelow rOver + | i < 0 = rBelow k i + | i >= ksz = rOver k i + | otherwise = i + {-# INLINE [1] repairIndex #-} + consDim i _ = i + {-# INLINE [1] consDim #-} + unconsDim i = (i, Ix0) + {-# INLINE [1] unconsDim #-} + snocDim _ i = i + {-# INLINE [1] snocDim #-} + unsnocDim i = (Ix0, i) + {-# INLINE [1] unsnocDim #-} + getDimM i 1 = pure i + getDimM ix d = throwM $ IndexDimensionException ix d + {-# INLINE [1] getDimM #-} + setDimM _ 1 i = pure i + setDimM ix d _ = throwM $ IndexDimensionException ix d + {-# INLINE [1] setDimM #-} + pullOutDimM i 1 = pure (i, Ix0) + pullOutDimM ix d = throwM $ IndexDimensionException ix d + {-# INLINE [1] pullOutDimM #-} + insertDimM Ix0 1 i = pure i + insertDimM ix d _ = throwM $ IndexDimensionException ix d + {-# INLINE [1] insertDimM #-} + pureIndex i = i + {-# INLINE [1] pureIndex #-} + liftIndex f = f + {-# INLINE [1] liftIndex #-} + liftIndex2 f = f + {-# INLINE [1] liftIndex2 #-} + foldlIndex f = f + {-# INLINE [1] foldlIndex #-} + iterM k0 k1 inc cond = loopM k0 (`cond` k1) (+inc) + {-# INLINE iterM #-} + iterM_ k0 k1 inc cond = loopM_ k0 (`cond` k1) (+inc) + {-# INLINE iterM_ #-} + + +-- | Exceptions that get thrown when there is a problem with an index, size or dimension. +-- +-- @since 0.3.0 +data IndexException where + -- | Index contains a zero value along one of the dimensions. + IndexZeroException :: Index ix => !ix -> IndexException + -- | Dimension is out of reach. + IndexDimensionException :: (Show ix, Typeable ix) => !ix -> Dim -> IndexException + -- | Index is out of bounds. + IndexOutOfBoundsException :: Index ix => !(Sz ix) -> !ix -> IndexException + +instance Show IndexException where + show (IndexZeroException ix) = "IndexZeroException: " ++ show ix + show (IndexDimensionException ix dim) = + "IndexDimensionException: " ++ show dim ++ " for " ++ show ix + show (IndexOutOfBoundsException sz ix) = + "IndexOutOfBoundsException: " ++ showsPrec 1 ix " not safe for (" ++ show sz ++ ")" + showsPrec 0 arr s = show arr ++ s + showsPrec _ arr s = '(' : show arr ++ ")" ++ s + +instance Exception IndexException + +-- | Exception that indicates an issue with an array size. +-- +-- @since 0.3.0 +data SizeException where + -- | Two sizes are expected to be equal along some or all dimensions, but they are not. + SizeMismatchException :: Index ix => !(Sz ix) -> !(Sz ix) -> SizeException + -- | Total number of elements does not match between the two sizes. + SizeElementsMismatchException :: (Index ix, Index ix') => !(Sz ix) -> !(Sz ix') -> SizeException + -- | Described subregion is too big for the specified size. + SizeSubregionException :: Index ix => !(Sz ix) -> !ix -> !(Sz ix) -> SizeException + -- | An array with the size cannot contain any elements. + SizeEmptyException :: Index ix => !(Sz ix) -> SizeException + +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 ++ ")" + show (SizeEmptyException sz) = + "SizeEmptyException: (" ++ show sz ++ ") corresponds to an empty array" + showsPrec 0 arr s = show arr ++ s + showsPrec _ arr s = '(' : show arr ++ ")" ++ s + +-- | Exception that can happen upon conversion of a ragged type array into the rectangular kind. Which +-- means conversion from lists is susceptible to this exception. +-- +-- @since 0.3.0 +data ShapeException + = DimTooShortException !Sz1 !Sz1 + | DimTooLongException + deriving Eq + +instance Show ShapeException where + show (DimTooShortException sz sz') = + "DimTooShortException: expected (" ++ show sz ++ "), got (" ++ show sz' ++ ")" + show DimTooLongException = + "DimTooLongException" + showsPrec 0 arr s = show arr ++ s + showsPrec _ arr s = '(' : show arr ++ ")" ++ s + +instance Exception ShapeException diff --git a/src/Data/Massiv/Core/Index/Ix.hs b/src/Data/Massiv/Core/Index/Ix.hs --- a/src/Data/Massiv/Core/Index/Ix.hs +++ b/src/Data/Massiv/Core/Index/Ix.hs @@ -9,108 +9,184 @@ {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE TypeOperators #-} {-# LANGUAGE UndecidableInstances #-} - -#if __GLASGOW_HASKELL__ >= 800 {-# LANGUAGE TypeFamilyDependencies #-} -#else -{-# LANGUAGE GADTs #-} -#endif -- | -- Module : Data.Massiv.Core.Index.Ix --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable -- -module Data.Massiv.Core.Index.Ix where +module Data.Massiv.Core.Index.Ix + ( Ix + , IxN((:>)) + , type Sz + , pattern Sz + , type Ix1 + , pattern Ix1 + , type Sz1 + , pattern Sz1 + , type Ix2(Ix2, (:.)) + , type Sz2 + , pattern Sz2 + , type Ix3 + , pattern Ix3 + , type Sz3 + , pattern Sz3 + , type Ix4 + , pattern Ix4 + , type Sz4 + , pattern Sz4 + , type Ix5 + , pattern Ix5 + , type Sz5 + , pattern Sz5 + ) where +import Control.Monad.Catch (MonadThrow(..)) import Control.DeepSeq -import Control.Monad (liftM) -import Data.Massiv.Core.Index.Class -import Data.Monoid ((<>)) +import Control.Monad (liftM) +import Data.Massiv.Core.Index.Internal +import Data.Monoid ((<>)) import Data.Proxy -import qualified Data.Vector.Generic as V -import qualified Data.Vector.Generic.Mutable as VM -import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Generic as V +import qualified Data.Vector.Generic.Mutable as VM +import qualified Data.Vector.Unboxed as VU import GHC.TypeLits infixr 5 :>, :. --- | Another type synonym for 1-dimensional index, i.e. `Int` and `Ix1T`. Provided here purely for --- consistency. -type Ix1 = Int --- | This is a very handy pattern synonym to indicate that any arbitrary whole number is an `Int`, --- i.e. a 1-dimensional index: @(Ix1 i) == (i :: Int)@ -pattern Ix1 :: Int -> Ix1 -pattern Ix1 i = i - --- | 2-dimensional index. This also a base index for higher dimensions. -data Ix2 = (:.) {-# UNPACK #-} !Int {-# UNPACK #-} !Int +-- | 2-dimensional index. This is also a base index for higher dimensions. +-- +-- @since 0.1.0 +data Ix2 = {-# UNPACK #-} !Int :. {-# UNPACK #-} !Int -- | 2-dimensional index constructor. Useful when @TypeOperators@ extension isn't enabled, or simply -- infix notation is inconvenient. @(Ix2 i j) == (i :. j)@. +-- +-- @since 0.1.0 pattern Ix2 :: Int -> Int -> Ix2 -pattern Ix2 i j = i :. j +pattern Ix2 i2 i1 = i2 :. i1 +-- | 2-dimensional size type synonym. +-- +-- @since 0.3.0 +type Sz2 = Sz Ix2 + +-- | 2-dimensional size constructor. @(Sz2 i j) == Sz (i :. j)@. +-- +-- @since 0.3.0 +pattern Sz2 :: Int -> Int -> Sz2 +pattern Sz2 i2 i1 = Sz (i2 :. i1) +{-# COMPLETE Sz2 #-} + -- | 3-dimensional type synonym. Useful as a alternative to enabling @DataKinds@ and using type -- level Nats. +-- +-- @since 0.1.0 type Ix3 = IxN 3 -- | 3-dimensional index constructor. @(Ix3 i j k) == (i :> j :. k)@. +-- +-- @since 0.1.0 pattern Ix3 :: Int -> Int -> Int -> Ix3 -pattern Ix3 i j k = i :> j :. k +pattern Ix3 i3 i2 i1 = i3 :> i2 :. i1 +-- | 3-dimensional size type synonym. +-- +-- @since 0.3.0 +type Sz3 = Sz Ix3 + +-- | 3-dimensional size constructor. @(Sz3 i j k) == Sz (i :> j :. k)@. +-- +-- @since 0.3.0 +pattern Sz3 :: Int -> Int -> Int -> Sz3 +pattern Sz3 i3 i2 i1 = Sz (i3 :> i2 :. i1) +{-# COMPLETE Sz3 #-} + -- | 4-dimensional type synonym. +-- +-- @since 0.1.0 type Ix4 = IxN 4 + -- | 4-dimensional index constructor. @(Ix4 i j k l) == (i :> j :> k :. l)@. +-- +-- @since 0.1.0 pattern Ix4 :: Int -> Int -> Int -> Int -> Ix4 -pattern Ix4 i j k l = i :> j :> k :. l +pattern Ix4 i4 i3 i2 i1 = i4 :> i3 :> i2 :. i1 +-- | 4-dimensional size type synonym. +-- +-- @since 0.3.0 +type Sz4 = Sz Ix4 + +-- | 4-dimensional size constructor. @(Sz4 i j k l) == Sz (i :> j :> k :. l)@. +-- +-- @since 0.3.0 +pattern Sz4 :: Int -> Int -> Int -> Int -> Sz4 +pattern Sz4 i4 i3 i2 i1 = Sz (i4 :> i3 :> i2 :. i1) +{-# COMPLETE Sz4 #-} + -- | 5-dimensional type synonym. +-- +-- @since 0.1.0 type Ix5 = IxN 5 --- | 5-dimensional index constructor. @(Ix5 i j k l m) = (i :> j :> k :> l :. m)@. + +-- | 5-dimensional index constructor. @(Ix5 i j k l m) == (i :> j :> k :> l :. m)@. +-- +-- @since 0.1.0 pattern Ix5 :: Int -> Int -> Int -> Int -> Int -> Ix5 -pattern Ix5 i j k l m = i :> j :> k :> l :. m +pattern Ix5 i5 i4 i3 i2 i1 = i5 :> i4 :> i3 :> i2 :. i1 +-- | 5-dimensional size type synonym. +-- +-- @since 0.3.0 +type Sz5 = Sz Ix5 -#if __GLASGOW_HASKELL__ >= 800 +-- | 5-dimensional size constructor. @(Sz5 i j k l m) == Sz (i :> j :> k :> l :. m)@. +-- +-- @since 0.3.0 +pattern Sz5 :: Int -> Int -> Int -> Int -> Int -> Sz5 +pattern Sz5 i5 i4 i3 i2 i1 = Sz (i5 :> i4 :> i3 :> i2 :. i1) +{-# COMPLETE Sz5 #-} -- | n-dimensional index. Needs a base case, which is the `Ix2`. -data IxN (n :: Nat) = (:>) {-# UNPACK #-} !Int !(Ix (n - 1)) +-- +-- @since 0.1.0 +data IxN (n :: Nat) = {-# UNPACK #-} !Int :> !(Ix (n - 1)) -- | Defines n-dimensional index by relating a general `IxN` with few base cases. +-- +-- @since 0.1.0 type family Ix (n :: Nat) = r | r -> n where Ix 0 = Ix0 Ix 1 = Ix1 Ix 2 = Ix2 Ix n = IxN n -#else -data IxN (n :: Nat) where - (:>) :: Dimensions (Ix (n - 1)) ~ (n - 1) => {-# UNPACK #-} !Int -> !(Ix (n - 1)) -> IxN n - -type family Ix (n :: Nat) where - Ix 0 = Ix0 - Ix 1 = Ix1 - Ix 2 = Ix2 - Ix n = IxN n - -#endif - - type instance Lower Ix2 = Ix1 type instance Lower (IxN n) = Ix (n - 1) instance Show Ix2 where - show (i :. j) = show i ++ " :. " ++ show j + showsPrec n (i :. j) = + if n < 1 + then inner + else ("(" ++) . inner . (")" ++) + where + inner = shows i . (" :. " ++) . shows j instance Show (Ix (n - 1)) => Show (IxN n) where - show (i :> ix) = show i ++ " :> " ++ show ix + showsPrec n (i :> ix) = + if n < 1 + then inner + else ("(" ++) . inner . (")" ++) + where + inner = shows i . (" :> " ++) . shows ix instance Num Ix2 where @@ -149,9 +225,6 @@ instance {-# OVERLAPPABLE #-} (4 <= n, KnownNat n, Index (Ix (n - 1)), -#if __GLASGOW_HASKELL__ < 800 - Dimensions (Ix ((n - 1) - 1)) ~ ((n - 1) - 1), -#endif IxN (n - 1) ~ Ix (n - 1) ) => Num (IxN n) where (+) = liftIndex2 (+) @@ -186,9 +259,6 @@ instance {-# OVERLAPPABLE #-} (4 <= n, KnownNat n, Index (Ix (n - 1)), -#if __GLASGOW_HASKELL__ < 800 - Dimensions (Ix ((n - 1) - 1)) ~ ((n - 1) - 1), -#endif IxN (n - 1) ~ Ix (n - 1) ) => Bounded (IxN n) where minBound = pureIndex minBound @@ -216,59 +286,19 @@ instance Ord (Ix (n - 1)) => Ord (IxN n) where compare (i1 :> ix1) (i2 :> ix2) = compare i1 i2 <> compare ix1 ix2 --- | Convert a `Int` tuple to `Ix2` -toIx2 :: Ix2T -> Ix2 -toIx2 (i, j) = i :. j -{-# INLINE toIx2 #-} --- | Convert an `Ix2` to `Int` tuple -fromIx2 :: Ix2 -> Ix2T -fromIx2 (i :. j) = (i, j) -{-# INLINE fromIx2 #-} - --- | Convert a `Int` 3-tuple to `Ix3` -toIx3 :: Ix3T -> Ix3 -toIx3 (i, j, k) = i :> j :. k -{-# INLINE toIx3 #-} - --- | Convert an `Ix3` to `Int` 3-tuple -fromIx3 :: Ix3 -> Ix3T -fromIx3 (i :> j :. k) = (i, j, k) -{-# INLINE fromIx3 #-} - --- | Convert a `Int` 4-tuple to `Ix4` -toIx4 :: Ix4T -> Ix4 -toIx4 (i, j, k, l) = i :> j :> k :. l -{-# INLINE toIx4 #-} - --- | Convert an `Ix4` to `Int` 4-tuple -fromIx4 :: Ix4 -> Ix4T -fromIx4 (i :> j :> k :. l) = (i, j, k, l) -{-# INLINE fromIx4 #-} - --- | Convert a `Int` 5-tuple to `Ix5` -toIx5 :: Ix5T -> Ix5 -toIx5 (i, j, k, l, m) = i :> j :> k :> l :. m -{-# INLINE toIx5 #-} - --- | Convert an `Ix5` to `Int` 5-tuple -fromIx5 :: Ix5 -> Ix5T -fromIx5 (i :> j :> k :> l :. m) = (i, j, k, l, m) -{-# INLINE fromIx5 #-} - - -instance {-# OVERLAPPING #-} Index Ix2 where +instance Index Ix2 where type Dimensions Ix2 = 2 dimensions _ = 2 {-# INLINE [1] dimensions #-} - totalElem (k2 :. k1) = k2 * k1 + totalElem (SafeSz (k2 :. k1)) = k2 * k1 {-# INLINE [1] totalElem #-} - isSafeIndex (k2 :. k1) (i2 :. i1) = 0 <= i2 && 0 <= i1 && i2 < k2 && i1 < k1 + isSafeIndex (SafeSz (k2 :. k1)) (i2 :. i1) = 0 <= i2 && 0 <= i1 && i2 < k2 && i1 < k1 {-# INLINE [1] isSafeIndex #-} - toLinearIndex (_ :. k1) (i2 :. i1) = k1 * i2 + i1 + toLinearIndex (SafeSz (_ :. k1)) (i2 :. i1) = k1 * i2 + i1 {-# INLINE [1] toLinearIndex #-} - fromLinearIndex (_ :. k1) i = case i `quotRem` k1 of - (i2, i1) -> i2 :. i1 + fromLinearIndex (SafeSz (_ :. k1)) i = case i `quotRem` k1 of + (i2, i1) -> i2 :. i1 {-# INLINE [1] fromLinearIndex #-} consDim = (:.) {-# INLINE [1] consDim #-} @@ -278,33 +308,30 @@ {-# INLINE [1] snocDim #-} unsnocDim (i2 :. i1) = (i2, i1) {-# INLINE [1] unsnocDim #-} - getDim (i2 :. _) 2 = Just i2 - getDim ( _ :. i1) 1 = Just i1 - getDim _ _ = Nothing - {-# INLINE [1] getDim #-} - setDim ( _ :. i1) 2 i2 = Just (i2 :. i1) - setDim (i2 :. _) 1 i1 = Just (i2 :. i1) - setDim _ _ _ = Nothing - {-# INLINE [1] setDim #-} - dropDim ( _ :. i1) 2 = Just i1 - dropDim (i2 :. _) 1 = Just i2 - dropDim _ _ = Nothing - {-# INLINE [1] dropDim #-} - pullOutDim (i2 :. i1) 2 = Just (i2, i1) - pullOutDim (i2 :. i1) 1 = Just (i1, i2) - pullOutDim _ _ = Nothing - {-# INLINE [1] pullOutDim #-} - insertDim i1 2 i2 = Just (i2 :. i1) - insertDim i2 1 i1 = Just (i2 :. i1) - insertDim _ _ _ = Nothing + 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 + {-# INLINE [1] setDimM #-} + pullOutDimM (i2 :. i1) 2 = pure (i2, i1) + pullOutDimM (i2 :. i1) 1 = pure (i1, i2) + 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 + {-# INLINE [1] insertDimM #-} pureIndex i = i :. i {-# INLINE [1] pureIndex #-} liftIndex f (i2 :. i1) = f i2 :. f i1 {-# INLINE [1] liftIndex #-} liftIndex2 f (i2 :. i1) (i2' :. i1') = f i2 i2' :. f i1 i1' {-# INLINE [1] liftIndex2 #-} - repairIndex (k :. szL) (i :. ixL) rBelow rOver = - repairIndex k i rBelow rOver :. repairIndex szL ixL rBelow rOver + repairIndex (SafeSz (k :. szL)) (i :. ixL) rBelow rOver = + repairIndex (SafeSz k) i rBelow rOver :. repairIndex (SafeSz szL) ixL rBelow rOver {-# INLINE [1] repairIndex #-} @@ -312,14 +339,14 @@ type Dimensions Ix3 = 3 dimensions _ = 3 {-# INLINE [1] dimensions #-} - totalElem (k3 :> k2 :. k1) = k3 * k2 * k1 + totalElem (SafeSz (k3 :> k2 :. k1)) = k3 * k2 * k1 {-# INLINE [1] totalElem #-} - isSafeIndex (k3 :> k2 :. k1) (i3 :> i2 :. i1) = + isSafeIndex (SafeSz (k3 :> k2 :. k1)) (i3 :> i2 :. i1) = 0 <= i3 && 0 <= i2 && 0 <= i1 && i3 < k3 && i2 < k2 && i1 < k1 {-# INLINE [1] isSafeIndex #-} - toLinearIndex (_ :> k2 :. k1) (i3 :> i2 :. i1) = (k2 * i3 + i2) * k1 + i1 + toLinearIndex (SafeSz (_ :> k2 :. k1)) (i3 :> i2 :. i1) = (k2 * i3 + i2) * k1 + i1 {-# INLINE [1] toLinearIndex #-} - fromLinearIndex (_ :> ix) i = let !(q, ixL) = fromLinearIndexAcc ix i in q :> ixL + fromLinearIndex (SafeSz (_ :> ix)) i = let !(q, ixL) = fromLinearIndexAcc ix i in q :> ixL {-# INLINE [1] fromLinearIndex #-} consDim = (:>) {-# INLINE [1] consDim #-} @@ -329,53 +356,45 @@ {-# INLINE [1] snocDim #-} unsnocDim (i3 :> i2 :. i1) = (i3 :. i2, i1) {-# INLINE [1] unsnocDim #-} - getDim (i3 :> _ :. _) 3 = Just i3 - getDim ( _ :> i2 :. _) 2 = Just i2 - getDim ( _ :> _ :. i1) 1 = Just i1 - getDim _ _ = Nothing - {-# INLINE [1] getDim #-} - setDim ( _ :> i2 :. i1) 3 i3 = Just (i3 :> i2 :. i1) - setDim (i3 :> _ :. i1) 2 i2 = Just (i3 :> i2 :. i1) - setDim (i3 :> i2 :. _) 1 i1 = Just (i3 :> i2 :. i1) - setDim _ _ _ = Nothing - {-# INLINE [1] setDim #-} - dropDim ( _ :> i2 :. i1) 3 = Just (i2 :. i1) - dropDim (i3 :> _ :. i1) 2 = Just (i3 :. i1) - dropDim (i3 :> i2 :. _) 1 = Just (i3 :. i2) - dropDim _ _ = Nothing - {-# INLINE [1] dropDim #-} - pullOutDim (i3 :> i2 :. i1) 3 = Just (i3, i2 :. i1) - pullOutDim (i3 :> i2 :. i1) 2 = Just (i2, i3 :. i1) - pullOutDim (i3 :> i2 :. i1) 1 = Just (i1, i3 :. i2) - pullOutDim _ _ = Nothing - {-# INLINE [1] pullOutDim #-} - insertDim (i2 :. i1) 3 i3 = Just (i3 :> i2 :. i1) - insertDim (i3 :. i1) 2 i2 = Just (i3 :> i2 :. i1) - insertDim (i3 :. i2) 1 i1 = Just (i3 :> i2 :. i1) - insertDim _ _ _ = Nothing - {-# INLINE [1] insertDim #-} + 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 + {-# 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 + {-# 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 + {-# INLINE [1] insertDimM #-} pureIndex i = i :> i :. i {-# INLINE [1] pureIndex #-} liftIndex f (i3 :> i2 :. i1) = f i3 :> f i2 :. f i1 {-# INLINE [1] liftIndex #-} liftIndex2 f (i3 :> i2 :. i1) (i3' :> i2' :. i1') = f i3 i3' :> f i2 i2' :. f i1 i1' {-# INLINE [1] liftIndex2 #-} - repairIndex (n :> szL) (i :> ixL) rBelow rOver = - repairIndex n i rBelow rOver :> repairIndex szL ixL rBelow rOver + repairIndex (SafeSz (n :> szL)) (i :> ixL) rBelow rOver = + repairIndex (SafeSz n) i rBelow rOver :> repairIndex (SafeSz szL) ixL rBelow rOver {-# INLINE [1] repairIndex #-} instance {-# OVERLAPPABLE #-} (4 <= n, KnownNat n, Index (Ix (n - 1)), -#if __GLASGOW_HASKELL__ < 800 - Dimensions (Ix ((n - 1) - 1)) ~ ((n - 1) - 1), -#endif IxN (n - 1) ~ Ix (n - 1) ) => Index (IxN n) where type Dimensions (IxN n) = n dimensions _ = fromInteger $ natVal (Proxy :: Proxy n) {-# INLINE [1] dimensions #-} - totalElem (i :> ixl) = i * totalElem ixl + totalElem (SafeSz (i :> ixl)) = foldlIndex (*) i ixl {-# INLINE [1] totalElem #-} consDim = (:>) {-# INLINE [1] consDim #-} @@ -386,29 +405,26 @@ unsnocDim (i :> ixl) = case unsnocDim ixl of (ix, i1) -> (i :> ix, i1) {-# INLINE [1] unsnocDim #-} - getDim ix@(i :> ixl) d | d == dimensions ix = Just i - | otherwise = getDim ixl d - {-# INLINE [1] getDim #-} - setDim ix@(i :> ixl) d di | d == dimensions ix = Just (di :> ixl) - | otherwise = (i :>) <$> setDim ixl d di - {-# INLINE [1] setDim #-} - dropDim ix@(i :> ixl) d | d == dimensions ix = Just ixl - | otherwise = (i :>) <$> dropDim ixl d - {-# INLINE [1] dropDim #-} - pullOutDim ix@(i :> ixl) d | d == dimensions ix = Just (i, ixl) - | otherwise = fmap (i :>) <$> pullOutDim ixl d - {-# INLINE [1] pullOutDim #-} - insertDim ix@(i :> ixl) d di | d == dimensions ix + 1 = Just (di :> ix) - | otherwise = (i :>) <$> insertDim ixl d di - {-# INLINE [1] insertDim #-} + getDimM ix@(i :> ixl) d | d == dimensions (Just ix) = pure i + | otherwise = getDimM ixl d + {-# INLINE [1] getDimM #-} + setDimM ix@(i :> ixl) d di | d == dimensions (Just ix) = pure (di :> ixl) + | otherwise = (i :>) <$> setDimM ixl d di + {-# INLINE [1] setDimM #-} + pullOutDimM ix@(i :> ixl) d | d == dimensions (Just ix) = pure (i, ixl) + | otherwise = fmap (i :>) <$> pullOutDimM ixl d + {-# INLINE [1] pullOutDimM #-} + insertDimM ix@(i :> ixl) d di | d == dimensions (Just ix) + 1 = pure (di :> ix) + | otherwise = (i :>) <$> insertDimM ixl d di + {-# INLINE [1] insertDimM #-} pureIndex i = i :> (pureIndex i :: Ix (n - 1)) {-# INLINE [1] pureIndex #-} liftIndex f (i :> ix) = f i :> liftIndex f ix {-# INLINE [1] liftIndex #-} liftIndex2 f (i :> ix) (i' :> ix') = f i i' :> liftIndex2 f ix ix' {-# INLINE [1] liftIndex2 #-} - repairIndex (k :> szL) (i :> ixL) rBelow rOver = - repairIndex k i rBelow rOver :> repairIndex szL ixL rBelow rOver + repairIndex (SafeSz (n :> szL)) (i :> ixL) rBelow rOver = + repairIndex (SafeSz n) i rBelow rOver :> repairIndex (SafeSz szL) ixL rBelow rOver {-# INLINE [1] repairIndex #-} @@ -418,7 +434,7 @@ -- | Unboxing of a `Ix2`. instance VU.Unbox Ix2 -newtype instance VU.MVector s Ix2 = MV_Ix2 (VU.MVector s Ix2T) +newtype instance VU.MVector s Ix2 = MV_Ix2 (VU.MVector s (Int, Int)) instance VM.MVector VU.MVector Ix2 where basicLength (MV_Ix2 mvec) = VM.basicLength mvec @@ -429,15 +445,15 @@ {-# INLINE basicOverlaps #-} basicUnsafeNew len = MV_Ix2 `liftM` VM.basicUnsafeNew len {-# INLINE basicUnsafeNew #-} - basicUnsafeReplicate len val = MV_Ix2 `liftM` VM.basicUnsafeReplicate len (fromIx2 val) + basicUnsafeReplicate len (i :. j) = MV_Ix2 `liftM` VM.basicUnsafeReplicate len (i, j) {-# INLINE basicUnsafeReplicate #-} - basicUnsafeRead (MV_Ix2 mvec) idx = toIx2 `liftM` VM.basicUnsafeRead mvec idx + basicUnsafeRead (MV_Ix2 mvec) idx = uncurry (:.) `liftM` VM.basicUnsafeRead mvec idx {-# INLINE basicUnsafeRead #-} - basicUnsafeWrite (MV_Ix2 mvec) idx val = VM.basicUnsafeWrite mvec idx (fromIx2 val) + basicUnsafeWrite (MV_Ix2 mvec) idx (i :. j) = VM.basicUnsafeWrite mvec idx (i, j) {-# INLINE basicUnsafeWrite #-} basicClear (MV_Ix2 mvec) = VM.basicClear mvec {-# INLINE basicClear #-} - basicSet (MV_Ix2 mvec) val = VM.basicSet mvec (fromIx2 val) + basicSet (MV_Ix2 mvec) (i :. j) = VM.basicSet mvec (i, j) {-# INLINE basicSet #-} basicUnsafeCopy (MV_Ix2 mvec) (MV_Ix2 mvec') = VM.basicUnsafeCopy mvec mvec' {-# INLINE basicUnsafeCopy #-} @@ -451,7 +467,7 @@ #endif -newtype instance VU.Vector Ix2 = V_Ix2 (VU.Vector Ix2T) +newtype instance VU.Vector Ix2 = V_Ix2 (VU.Vector (Int, Int)) instance V.Vector VU.Vector Ix2 where basicUnsafeFreeze (MV_Ix2 mvec) = V_Ix2 `liftM` V.basicUnsafeFreeze mvec @@ -462,7 +478,7 @@ {-# INLINE basicLength #-} basicUnsafeSlice idx len (V_Ix2 vec) = V_Ix2 (V.basicUnsafeSlice idx len vec) {-# INLINE basicUnsafeSlice #-} - basicUnsafeIndexM (V_Ix2 vec) idx = toIx2 `liftM` V.basicUnsafeIndexM vec idx + basicUnsafeIndexM (V_Ix2 vec) idx = uncurry (:.) `liftM` V.basicUnsafeIndexM vec idx {-# INLINE basicUnsafeIndexM #-} basicUnsafeCopy (MV_Ix2 mvec) (V_Ix2 vec) = V.basicUnsafeCopy mvec vec {-# INLINE basicUnsafeCopy #-} @@ -476,20 +492,11 @@ -- | Unboxing of a `IxN`. -instance (3 <= n, -#if __GLASGOW_HASKELL__ < 800 - Dimensions (Ix (n - 1)) ~ (n - 1), -#endif - VU.Unbox (Ix (n-1))) => VU.Unbox (IxN n) +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))) -instance (3 <= n, -#if __GLASGOW_HASKELL__ < 800 - Dimensions (Ix (n - 1)) ~ (n - 1), -#endif - VU.Unbox (Ix (n - 1))) => - VM.MVector VU.MVector (IxN n) where +instance (3 <= n, VU.Unbox (Ix (n - 1))) => VM.MVector VU.MVector (IxN n) where basicLength (MV_IxN (_, mvec)) = VM.basicLength mvec {-# INLINE basicLength #-} basicUnsafeSlice idx len (MV_IxN (mvec1, mvec)) = @@ -541,11 +548,7 @@ newtype instance VU.Vector (IxN n) = V_IxN (VU.Vector Int, VU.Vector (Ix (n-1))) -instance (3 <= n, -#if __GLASGOW_HASKELL__ < 800 - Dimensions (Ix (n - 1)) ~ (n - 1), -#endif - VU.Unbox (Ix (n-1))) => V.Vector VU.Vector (IxN n) where +instance (3 <= n, VU.Unbox (Ix (n - 1))) => V.Vector VU.Vector (IxN n) where basicUnsafeFreeze (MV_IxN (mvec1, mvec)) = do iv <- V.basicUnsafeFreeze mvec1 ivs <- V.basicUnsafeFreeze mvec @@ -571,4 +574,3 @@ {-# INLINE basicUnsafeCopy #-} elemseq _ = seq {-# INLINE elemseq #-} - 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,16 +1,13 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE PatternSynonyms #-} -#if __GLASGOW_HASKELL__ >= 800 +#if __GLASGOW_HASKELL__ < 820 {-# OPTIONS_GHC -Wno-unrecognised-pragmas #-} -{-# LANGUAGE GeneralizedNewtypeDeriving #-} -#else -{-# LANGUAGE GADTs #-} -{-# LANGUAGE StandaloneDeriving #-} #endif -- | -- Module : Data.Massiv.Core.Index.Stride --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental @@ -26,8 +23,8 @@ , strideSize ) where -import Control.DeepSeq -import Data.Massiv.Core.Index.Class +import Control.DeepSeq +import Data.Massiv.Core.Index.Internal -- | 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 @@ -55,26 +52,15 @@ -- column intact then you'd use @Stride (5 :. 1)@. -- -#if __GLASGOW_HASKELL__ >= 800 newtype Stride ix = SafeStride ix deriving (Eq, Ord, NFData) -{-# COMPLETE Stride #-} -#else --- There is an issue in GHC 7.10 which prevents from placing `Index` constraint on a pattern. -data Stride ix where - SafeStride :: Index ix => ix -> Stride ix -deriving instance Eq ix => Eq (Stride ix) -deriving instance Ord ix => Ord (Stride ix) -instance NFData ix => NFData (Stride ix) where - rnf (SafeStride ix) = rnf ix -#endif - -- | A safe bidirectional pattern synonym for `Stride` construction that will make sure stride -- elements are always positive. pattern Stride :: Index ix => ix -> Stride ix pattern Stride ix <- SafeStride ix where Stride ix = SafeStride (liftIndex (max 1) ix) +{-# COMPLETE Stride #-} instance Index ix => Show (Stride ix) where @@ -86,7 +72,7 @@ unStride (SafeStride ix) = ix {-# INLINE unStride #-} --- | Adjust strating index according to the stride +-- | Adjust starting index according to the stride strideStart :: Index ix => Stride ix -> ix -> ix strideStart (SafeStride stride) ix = liftIndex2 @@ -96,14 +82,15 @@ {-# INLINE strideStart #-} -- | Adjust size according to the stride. -strideSize :: Index ix => Stride ix -> Sz ix -> ix -strideSize (SafeStride stride) sz = liftIndex (+ 1) $ liftIndex2 div (liftIndex (subtract 1) sz) stride +strideSize :: Index ix => Stride ix -> Sz ix -> Sz ix +strideSize (SafeStride stride) (SafeSz sz) = + SafeSz (liftIndex (+ 1) $ liftIndex2 div (liftIndex (subtract 1) sz) stride) {-# INLINE strideSize #-} -- | Compute an index with stride using the original size and index toLinearIndexStride :: Index ix => Stride ix -- ^ Stride - -> ix -- ^ Size + -> Sz ix -- ^ Size -> ix -- ^ Index -> Int toLinearIndexStride (SafeStride stride) sz ix = toLinearIndex sz (liftIndex2 div ix stride) diff --git a/src/Data/Massiv/Core/Index/Tuple.hs b/src/Data/Massiv/Core/Index/Tuple.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Core/Index/Tuple.hs @@ -0,0 +1,334 @@ +{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE TypeFamilies #-} +-- | +-- Module : Data.Massiv.Core.Index.Tuple +-- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Core.Index.Tuple + ( -- * Tuple based indices + -- ** 1-dimensional + Ix1T + -- ** 2-dimensional + , Ix2T + , toIx2 + , fromIx2 + -- ** 3-dimensional + , Ix3T + , toIx3 + , fromIx3 + -- ** 4-dimensional + , Ix4T + , toIx4 + , fromIx4 + -- ** 5-dimensional + , Ix5T + , toIx5 + , fromIx5 + ) where + +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 +-- consistency. +type Ix1T = Int + +-- | 2-dimensional index as tuple of `Int`s. +type Ix2T = (Int, Int) + +-- | 3-dimensional index as 3-tuple of `Int`s. +type Ix3T = (Int, Int, Int) + +-- | 4-dimensional index as 4-tuple of `Int`s. +type Ix4T = (Int, Int, Int, Int) + +-- | 5-dimensional index as 5-tuple of `Int`s. +type Ix5T = (Int, Int, Int, Int, Int) + +type instance Lower Ix2T = Ix1T +type instance Lower Ix3T = Ix2T +type instance Lower Ix4T = Ix3T +type instance Lower Ix5T = Ix4T + + + +-- | Convert an `Int` tuple to `Ix2` +-- +-- ==== __Example__ +-- +-- >>> toIx2 (2, 3) +-- 2 :. 3 +-- +-- @since 0.1.0 +toIx2 :: Ix2T -> Ix2 +toIx2 (i, j) = i :. j +{-# INLINE toIx2 #-} + +-- | Convert an `Ix2` to `Int` tuple +-- +-- ==== __Example__ +-- +-- >>> fromIx2 (2 :. 3) +-- (2,3) +-- +-- @since 0.1.0 +fromIx2 :: Ix2 -> Ix2T +fromIx2 (i :. j) = (i, j) +{-# INLINE fromIx2 #-} + +-- | Convert a `Int` 3-tuple to `Ix3` +-- +-- ==== __Example__ +-- +-- >>> toIx3 (1, 2, 3) +-- 1 :> 2 :. 3 +-- +-- @since 0.1.0 +toIx3 :: Ix3T -> Ix3 +toIx3 (i, j, k) = i :> j :. k +{-# INLINE toIx3 #-} + +-- | Convert an `Ix3` to `Int` 3-tuple +-- +-- ==== __Example__ +-- +-- >>> fromIx3 (1 :> 2 :. 3) +-- (1,2,3) +-- +-- @since 0.1.0 +fromIx3 :: Ix3 -> Ix3T +fromIx3 (i :> j :. k) = (i, j, k) +{-# INLINE fromIx3 #-} + +-- | Convert a `Int` 4-tuple to `Ix4` +-- +-- ==== __Example__ +-- +-- >>> toIx4 (1, 2, 3, 4) +-- 1 :> 2 :> 3 :. 4 +-- +-- @since 0.1.0 +toIx4 :: Ix4T -> Ix4 +toIx4 (i, j, k, l) = i :> j :> k :. l +{-# INLINE toIx4 #-} + +-- | Convert an `Ix4` to `Int` 4-tuple +-- +-- ==== __Example__ +-- +-- >>> fromIx4 (1 :> 2 :> 3 :. 4) +-- (1,2,3,4) +-- +-- @since 0.1.0 +fromIx4 :: Ix4 -> Ix4T +fromIx4 (i :> j :> k :. l) = (i, j, k, l) +{-# INLINE fromIx4 #-} + +-- | Convert a `Int` 5-tuple to `Ix5` +-- +-- ==== __Example__ +-- +-- >>> toIx5 (1, 2, 3, 4, 5) +-- 1 :> 2 :> 3 :> 4 :. 5 +-- +-- @since 0.1.0 +toIx5 :: Ix5T -> Ix5 +toIx5 (i, j, k, l, m) = i :> j :> k :> l :. m +{-# INLINE toIx5 #-} + +-- | Convert an `Ix5` to `Int` 5-tuple +-- +-- ==== __Example__ +-- +-- >>> fromIx5 (1 :> 2 :> 3 :> 4 :. 5) +-- (1,2,3,4,5) +-- +-- @since 0.1.0 +fromIx5 :: Ix5 -> Ix5T +fromIx5 (i :> j :> k :> l :. m) = (i, j, k, l, m) +{-# INLINE fromIx5 #-} + +-- | +-- @since 0.1.0 +instance Index Ix2T where + type Dimensions Ix2T = 2 + dimensions _ = 2 + {-# INLINE [1] dimensions #-} + totalElem (SafeSz (k2, k1)) = k2 * k1 + {-# INLINE [1] totalElem #-} + toLinearIndex (SafeSz (_, k1)) (i2, i1) = k1 * i2 + i1 + {-# INLINE [1] toLinearIndex #-} + fromLinearIndex (SafeSz (_, k1)) !i = i `quotRem` k1 + {-# INLINE [1] fromLinearIndex #-} + consDim = (,) + {-# INLINE [1] consDim #-} + unconsDim = id + {-# INLINE [1] unconsDim #-} + snocDim = (,) + {-# 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 + {-# INLINE [1] getDimM #-} + setDimM (_, i1) 2 i2 = pure (i2, i1) + setDimM (i2, _) 1 i1 = pure (i2, i1) + setDimM ix d _ = throwM $ IndexDimensionException ix d + {-# INLINE [1] setDimM #-} + pullOutDimM (i2, i1) 2 = pure (i2, i1) + pullOutDimM (i2, i1) 1 = pure (i1, i2) + pullOutDimM ix d = throwM $ IndexDimensionException ix d + {-# INLINE [1] pullOutDimM #-} + insertDimM i1 2 i2 = pure (i2, i1) + insertDimM i2 1 i1 = pure (i2, i1) + 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 + {-# INLINE [1] totalElem #-} + consDim i3 (i2, i1) = (i3, i2, i1) + {-# INLINE [1] consDim #-} + unconsDim (i3, i2, i1) = (i3, (i2, i1)) + {-# INLINE [1] unconsDim #-} + snocDim (i3, i2) i1 = (i3, i2, i1) + {-# 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 + {-# 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 + {-# 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 + {-# 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 + 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 + {-# INLINE [1] dimensions #-} + totalElem (SafeSz (k4, k3, k2, k1)) = k4 * k3 * k2 * k1 + {-# INLINE [1] totalElem #-} + consDim i4 (i3, i2, i1) = (i4, i3, i2, i1) + {-# INLINE [1] consDim #-} + unconsDim (i4, i3, i2, i1) = (i4, (i3, i2, i1)) + {-# INLINE [1] unconsDim #-} + snocDim (i4, i3, i2) i1 = (i4, i3, i2, i1) + {-# 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 + {-# 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 + {-# INLINE [1] setDimM #-} + 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 + {-# 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 + {-# 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 + {-# INLINE [1] dimensions #-} + totalElem (SafeSz (n5, n4, n3, n2, n1)) = n5 * n4 * n3 * n2 * n1 + {-# INLINE [1] totalElem #-} + consDim i5 (i4, i3, i2, i1) = (i5, i4, i3, i2, i1) + {-# INLINE [1] consDim #-} + unconsDim (i5, i4, i3, i2, i1) = (i5, (i4, i3, i2, i1)) + {-# INLINE [1] unconsDim #-} + snocDim (i5, i4, i3, i2) i1 = (i5, i4, i3, i2, i1) + {-# 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 + {-# 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 + {-# INLINE [1] setDimM #-} + 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 + {-# 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 + {-# INLINE [1] insertDimM #-} + pureIndex i = (i, i, i, i, i) + {-# INLINE [1] pureIndex #-} + liftIndex2 f (i5, i4, i3, i2, i1) (i5', i4', i3', i2', i1') = + (f i5 i5', f i4 i4', f i3 i3', f i2 i2', f i1 i1') + {-# INLINE [1] liftIndex2 #-} diff --git a/src/Data/Massiv/Core/Iterator.hs b/src/Data/Massiv/Core/Iterator.hs --- a/src/Data/Massiv/Core/Iterator.hs +++ b/src/Data/Massiv/Core/Iterator.hs @@ -1,7 +1,8 @@ -{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE BangPatterns #-} -- | -- Module : Data.Massiv.Core.Iterator --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental @@ -9,62 +10,98 @@ -- module Data.Massiv.Core.Iterator ( loop + , loopA_ , loopM , loopM_ , loopDeepM , splitLinearly , splitLinearlyWith_ , splitLinearlyWithM_ + , splitLinearlyWithStartAtM_ ) where +import Control.Scheduler (Scheduler(..)) -- | Efficient loop with an accumulator +-- +-- @since 0.1.0 loop :: Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> a) -> a loop !init' condition increment !initAcc f = go init' initAcc where - go !step !acc = - case condition step of - False -> acc - True -> go (increment step) (f step acc) + go !step !acc + | condition step = go (increment step) (f step acc) + | otherwise = acc {-# INLINE loop #-} --- | Very efficient monadic loop with an accumulator +-- | Efficient monadic loop with an accumulator +-- +-- >>> loopM 1 (< 20) (+ 2) [] (\i a -> Just (i:a)) +-- Just [19,17,15,13,11,9,7,5,3,1] +-- +-- @since 0.1.0 loopM :: Monad m => Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a loopM !init' condition increment !initAcc f = go init' initAcc where - go !step !acc = - case condition step of - False -> return acc - True -> f step acc >>= go (increment step) + go !step !acc + | condition step = f step acc >>= go (increment step) + | otherwise = return acc {-# INLINE loopM #-} -- | Efficient monadic loop. Result of each iteration is discarded. +-- +-- @since 0.1.0 loopM_ :: Monad m => Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m () loopM_ !init' condition increment f = go init' where - go !step = - case condition step of - False -> return () - True -> f step >> go (increment step) + go !step + | condition step = f step >> go (increment step) + | otherwise = pure () + {-# INLINE loopM_ #-} --- | Less efficient monadic loop with an accumulator that reverses the direction of action --- application +-- | Efficient Applicative loop. Result of each iteration is discarded. +-- +-- @since 0.3.0 +loopA_ :: Applicative f => Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f a) -> f () +loopA_ !init' condition increment f = go init' + where + go !step + | condition step = f step *> go (increment step) + | otherwise = pure () +{-# INLINE loopA_ #-} + + +-- | Similar to `loopM`, but slightly less efficient monadic loop with an accumulator that reverses +-- the direction of action application. eg: +-- +-- >>> loopDeepM 1 (< 20) (+ 2) [] (\i a -> Just (i:a)) +-- Just [1,3,5,7,9,11,13,15,17,19] +-- +-- Equivalent to: +-- +-- >>> loopM 19 (>= 1) (subtract 2) [] (\i a -> Just (i:a)) +-- Just [1,3,5,7,9,11,13,15,17,19] +-- +-- @since 0.1.0 loopDeepM :: Monad m => Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a loopDeepM !init' condition increment !initAcc f = go init' initAcc where - go !step !acc = - case condition step of - False -> return acc - True -> go (increment step) acc >>= f step + go !step !acc + | condition step = go (increment step) acc >>= f step + | otherwise = return acc {-# INLINE loopDeepM #-} - -splitLinearly :: Int -> Int -> (Int -> Int -> a) -> a +-- | Divide length in chunks and apply a function to the computed results +-- +-- @since 0.2.1 +splitLinearly :: Int -- ^ Number of chunks + -> Int -- ^ Total length + -> (Int -> Int -> a) -- ^ Function that accepts a chunk length and slack start index + -> a splitLinearly numChunks totalLength action = action chunkLength slackStart where !chunkLength = totalLength `quot` numChunks @@ -72,20 +109,43 @@ {-# INLINE splitLinearly #-} -splitLinearlyWith_ :: Monad m => Int -> (m () -> m a) -> Int -> (Int -> b) -> (Int -> b -> m ()) -> m a -splitLinearlyWith_ numChunks with totalLength index = - splitLinearlyWithM_ numChunks with totalLength (pure . index) +-- | Interator that can be used to split computation amongst different workers. For monadic +-- generator see `splitLinearlyWithM_`. +-- +-- @since 0.2.1 +splitLinearlyWith_ :: + Monad m => Scheduler m () -> Int -> (Int -> b) -> (Int -> b -> m ()) -> m () +splitLinearlyWith_ scheduler totalLength index = + splitLinearlyWithM_ scheduler totalLength (pure . index) {-# INLINE splitLinearlyWith_ #-} --- | Interator tha can be used to split computation jobs +-- | Interator that can be used to split computation jobs -- -- @since 0.2.6.0 splitLinearlyWithM_ :: - Monad m => Int -> (m () -> m a) -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m a -splitLinearlyWithM_ numChunks with totalLength make write = - splitLinearly numChunks totalLength $ \chunkLength slackStart -> do + Monad m => Scheduler m () -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m () +splitLinearlyWithM_ Scheduler {numWorkers, scheduleWork} totalLength make write = + splitLinearly numWorkers totalLength $ \chunkLength slackStart -> do loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start -> - with $ loopM_ start (< (start + chunkLength)) (+ 1) $ \ !k -> make k >>= write k - with $ loopM_ slackStart (< totalLength) (+ 1) $ \ !k -> make k >>= write k + scheduleWork $ loopM_ start (< (start + chunkLength)) (+ 1) $ \ !k -> make k >>= write k + scheduleWork $ loopM_ slackStart (< totalLength) (+ 1) $ \ !k -> make k >>= write k {-# INLINE splitLinearlyWithM_ #-} + + +-- | Interator that can be used to split computation jobs +-- +-- @since 0.2.6.0 +splitLinearlyWithStartAtM_ :: + Monad m => Scheduler m () -> Int -> Int -> (Int -> m b) -> (Int -> b -> m c) -> m () +splitLinearlyWithStartAtM_ Scheduler {numWorkers, scheduleWork} startAt totalLength make write = + splitLinearly numWorkers totalLength $ \chunkLength slackStart -> do + loopM_ startAt (< (slackStart + startAt)) (+ chunkLength) $ \ !start -> + scheduleWork $ loopM_ start (< (start + chunkLength)) (+ 1) $ \ !k -> make k >>= write k + scheduleWork $ + loopM_ (slackStart + startAt) (< (totalLength + startAt)) (+ 1) $ \ !k -> make k >>= write k +{-# INLINE splitLinearlyWithStartAtM_ #-} + + + + diff --git a/src/Data/Massiv/Core/List.hs b/src/Data/Massiv/Core/List.hs --- a/src/Data/Massiv/Core/List.hs +++ b/src/Data/Massiv/Core/List.hs @@ -1,16 +1,17 @@ +{-# LANGUAGE NamedFieldPuns #-} {-# OPTIONS_GHC -fno-warn-orphans #-} -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE RecordWildCards #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} -- | -- Module : Data.Massiv.Core.List --- Copyright : (c) Alexey Kuleshevich 2018 +-- Copyright : (c) Alexey Kuleshevich 2018-2019 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental @@ -21,23 +22,23 @@ , L(..) , Array(..) , toListArray - , showArray + , showsArrayPrec + , showArrayList , ListItem - , ShapeError(..) ) where -import Control.Exception -import Control.Monad (unless, when) -import Data.Coerce -import Data.Foldable (foldr') -import Data.Functor.Identity -import qualified Data.List as L -import Data.Massiv.Core.Common -import Data.Massiv.Core.Scheduler -import Data.Proxy -import Data.Typeable -import GHC.Exts -import System.IO.Unsafe (unsafePerformIO) +import Control.Exception +import Control.Scheduler +import Control.Monad (unless, when) +import Data.Coerce +import Data.Foldable (foldr') +import qualified Data.List as L +import Data.Massiv.Core.Common +import Data.Massiv.Core.Index.Internal +import Data.Proxy +import Data.Typeable +import GHC.Exts +import System.IO.Unsafe (unsafePerformIO) data LN @@ -86,14 +87,6 @@ , lData :: !(Array LN ix e) } - -data ShapeError = RowTooShortError - | RowTooLongError - deriving Show - -instance Exception ShapeError - - instance Nested L ix e where fromNested = LArray Seq {-# INLINE fromNested #-} @@ -108,42 +101,55 @@ toList = toNested . lData {-# INLINE toList #-} - - instance {-# OVERLAPPING #-} Ragged L Ix1 e where isNull = null . unList . lData {-# INLINE isNull #-} - empty comp = LArray comp (List []) - {-# INLINE empty #-} - edgeSize = length . unList . lData + emptyR comp = LArray comp (List []) + {-# INLINE emptyR #-} + edgeSize = SafeSz . length . unList . lData {-# INLINE edgeSize #-} - cons x arr = arr { lData = coerce (x : coerce (lData arr)) } - {-# INLINE cons #-} - uncons LArray {..} = + consR x arr = arr { lData = coerce (x : coerce (lData arr)) } + {-# INLINE consR #-} + unconsR LArray {..} = case L.uncons $ coerce lData of Nothing -> Nothing Just (x, xs) -> Just (x, LArray lComp (coerce xs)) - {-# INLINE uncons #-} + {-# INLINE unconsR #-} flatten = id {-# INLINE flatten #-} - unsafeGenerateM !comp !k f = do - xs <- loopDeepM 0 (< k) (+ 1) [] $ \i acc -> do + generateRaggedM !comp !k f = do + xs <- loopDeepM 0 (< coerce k) (+ 1) [] $ \i acc -> do e <- f i return (e:acc) return $ LArray comp $ coerce xs - {-# INLINE unsafeGenerateM #-} - loadRagged using uWrite start end _ xs = + {-# INLINE generateRaggedM #-} + loadRagged using uWrite start end sz xs = using $ do leftOver <- loopM start (< end) (+ 1) xs $ \i xs' -> - case uncons xs' of - Nothing -> throwIO RowTooShortError + case unconsR xs' of + Nothing -> return $! throw (DimTooShortException sz (outerLength xs)) Just (y, ys) -> uWrite i y >> return ys - unless (isNull leftOver) $ throwIO RowTooLongError + unless (isNull leftOver) (return $! throw DimTooLongException) {-# INLINE loadRagged #-} - raggedFormat f _ arr = L.concat $ "[ " : L.intersperse "," (map f (coerce (lData arr))) ++ [" ]"] + raggedFormat f _ arr = L.concat $ "[ " : L.intersperse ", " (map f (coerce (lData arr))) ++ [" ]"] +instance (Index ix, Ragged L ix e) => Load L ix e where + size = coerce . edgeSize + {-# INLINE size #-} + getComp = lComp + {-# INLINE getComp #-} + loadArrayM Scheduler {scheduleWork} arr uWrite = + loadRagged scheduleWork uWrite 0 (totalElem sz) sz arr + where !sz = edgeSize arr + {-# INLINE loadArrayM #-} + + + +outerLength :: Array L ix e -> Sz Int +outerLength = SafeSz . length . unList . lData + instance ( Index ix , Index (Lower ix) , Ragged L (Lower ix) e @@ -154,154 +160,89 @@ Ragged L ix e where isNull = null . unList . lData {-# INLINE isNull #-} - empty comp = LArray comp (List []) - {-# INLINE empty #-} + emptyR comp = LArray comp (List []) + {-# INLINE emptyR #-} edgeSize arr = - consDim (length (unList (lData arr))) $ - case uncons arr of - Nothing -> zeroIndex - Just (x, _) -> edgeSize x + SafeSz + (consDim (length (unList (lData arr))) $ + case unconsR arr of + Nothing -> zeroIndex + Just (x, _) -> coerce (edgeSize x)) {-# INLINE edgeSize #-} - cons (LArray _ x) arr = newArr + consR (LArray _ x) arr = newArr where - newArr = - arr {lData = coerce (x : coerce (lData arr))} - {-# INLINE cons #-} - uncons LArray {..} = + newArr = arr {lData = coerce (x : coerce (lData arr))} + {-# INLINE consR #-} + unconsR LArray {..} = case L.uncons (coerce lData) of Nothing -> Nothing Just (x, xs) -> let newArr = LArray lComp (coerce xs) newX = LArray lComp x - in Just (newX, newArr) - {-# INLINE uncons #-} - unsafeGenerateM Seq !sz f = do - let !(k, szL) = unconsDim sz - loopDeepM 0 (< k) (+ 1) (empty Seq) $ \i acc -> do - e <- unsafeGenerateM Seq szL (\ !ixL -> f (consDim i ixL)) - return (cons e acc) - unsafeGenerateM (ParOn wss) sz f = unsafeGenerateParM wss sz f - {-# INLINE unsafeGenerateM #-} + in Just (newX, newArr) + {-# INLINE unconsR #-} + -- generateRaggedM Seq !sz f = do + -- let !(k, szL) = unconsSz sz + -- loopDeepM 0 (< coerce k) (+ 1) (emptyR Seq) $ \i acc -> do + -- e <- generateRaggedM Seq szL (\ !ixL -> f (consDim i ixL)) + -- return (cons e acc) + generateRaggedM = unsafeGenerateParM + {-# INLINE generateRaggedM #-} flatten arr = LArray {lComp = lComp arr, lData = coerce xs} where - xs = - concatMap - (unList . lData . flatten . LArray (lComp arr)) - (unList (lData arr)) + xs = concatMap (unList . lData . flatten . LArray (lComp arr)) (unList (lData arr)) {-# INLINE flatten #-} loadRagged using uWrite start end sz xs = do - let szL = tailDim sz + let (k, szL) = unconsSz sz step = totalElem szL isZero = totalElem sz == 0 - when (isZero && not (isNull (flatten xs))) $ - throwIO RowTooLongError + when (isZero && not (isNull (flatten xs))) (return $! throw DimTooLongException) unless isZero $ do leftOver <- loopM start (< end) (+ step) xs $ \i zs -> - case uncons zs of - Nothing -> throwIO RowTooShortError + case unconsR zs of + Nothing -> return $! throw (DimTooShortException k (outerLength xs)) Just (y, ys) -> do _ <- loadRagged using uWrite i (i + step) szL y return ys - unless (isNull leftOver) $ throwIO RowTooLongError + unless (isNull leftOver) (return $! throw DimTooLongException) {-# INLINE loadRagged #-} raggedFormat f sep (LArray comp xs) = - showN - (\s y -> raggedFormat f s (LArray comp y :: Array L (Lower ix) e)) - sep - (coerce xs) - - --- unsafeGenerateParM :: --- (Elt LN ix e ~ Array LN (Lower ix) e, Index ix, Monad m, Ragged L (Lower ix) e) --- => [Int] --- -> ix --- -> (ix -> m e) --- -> m (Array L ix e) --- unsafeGenerateParM wws !sz f = do --- res <- sequence $ unsafePerformIO $ do --- let !(k, szL) = unconsDim sz --- resLs <- divideWork wws k $ \ !scheduler !chunkLength !totalLength !slackStart -> do --- loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start -> do --- scheduleWork scheduler $ do --- res <- loopM start (< (start + chunkLength)) (+ 1) [] $ \i acc -> do --- return (fmap lData (unsafeGenerateM Seq szL (\ !ixL -> f (consDim i ixL))):acc) --- return $! sequence res --- when (slackStart < totalLength) $ --- scheduleWork scheduler $ do --- res <- loopM (slackStart) (< totalLength) (+ 1) [] $ \i acc -> do --- return (fmap lData (unsafeGenerateM Seq szL (\ !ixL -> f (consDim i ixL))):acc) --- return $! sequence res --- return resLs --- return $ LArray (ParOn wws) $ List $ concat res --- {-# INLINE unsafeGenerateParM #-} + showN (\s y -> raggedFormat f s (LArray comp y :: Array L (Lower ix) e)) sep (coerce xs) unsafeGenerateParM :: (Elt LN ix e ~ Array LN (Lower ix) e, Index ix, Monad m, Ragged L (Lower ix) e) - => [Int] + => Comp -> Sz ix -> (ix -> m e) -> m (Array L ix e) -unsafeGenerateParM wws !sz f = do +unsafeGenerateParM comp !sz f = do res <- sequence $ unsafePerformIO $ do - let !(k, szL) = unconsDim sz - divideWork wws k $ \ !scheduler !chunkLength !totalLength !slackStart -> do - loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start -> - scheduleWork scheduler $ do - -- res <- loopM (start + chunkLength - 1) (>= start) (subtract 1) [] $ \i acc -> do - -- return (fmap lData (unsafeGenerateM Seq szL (\ !ixL -> f (consDim i ixL))):acc) - -- return $! sequence res - res <- loopDeepM start (< (start + chunkLength)) (+ 1) [] $ \i acc -> - return (fmap lData (unsafeGenerateM Seq szL (\ !ixL -> f (consDim i ixL))):acc) - return $! sequence res - when (slackStart < totalLength) $ - scheduleWork scheduler $ do - -- res <- loopM (totalLength - 1) (>= slackStart) (subtract 1) [] $ \i acc -> do - -- return (fmap lData (unsafeGenerateM Seq szL (\ !ixL -> f (consDim i ixL))):acc) - -- return $! sequence res - res <- loopDeepM slackStart (< totalLength) (+ 1) [] $ \i acc -> - return (fmap lData (unsafeGenerateM Seq szL (\ !ixL -> f (consDim i ixL))):acc) - return $! sequence res - return $ LArray (ParOn wws) $ List $ concat res + let !(ksz, szL) = unconsSz sz + !k = unSz ksz + withScheduler comp $ \ scheduler -> + splitLinearly (numWorkers scheduler) k $ \ chunkLength slackStart -> do + loopM_ 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 $! 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 $! sequence res + return $ LArray comp $ List $ concat res {-# INLINE unsafeGenerateParM #-} - --- unsafeGenerateParM :: --- (Elt LN ix e ~ Array LN (Lower ix) e, Index ix, Monad m, Ragged L (Lower ix) e) --- => [Int] --- -> ix --- -> (ix -> m e) --- -> m (Array L ix e) --- unsafeGenerateParM wws !sz f = do --- res <- sequence $ unsafePerformIO $ do --- let !(k, szL) = unconsDim sz --- resLs <- divideWork wws k $ \ !scheduler !chunkLength !totalLength !slackStart -> do --- when (slackStart < totalLength) $ --- scheduleWork scheduler $ do --- res <- loopM (totalLength - 1) (>= slackStart) (subtract 1) [] $ \i acc -> do --- return (fmap lData (unsafeGenerateM Seq szL (\ !ixL -> f (consDim i ixL))):acc) --- return $! sequence res --- loopM_ slackStart (> 0) (subtract chunkLength) $ \ !start -> do --- let !end = start - chunkLength --- scheduleWork scheduler $ do --- res <- loopM (start - 1) (>= end) (subtract 1) [] $ \i acc -> do --- return (fmap lData (unsafeGenerateM Seq szL (\ !ixL -> f (consDim i ixL))):acc) --- return $! sequence res --- return resLs --- return $ LArray (ParOn wws) $ List $ concat res --- {-# INLINE unsafeGenerateParM #-} - instance {-# OVERLAPPING #-} Construct L Ix1 e where - getComp = lComp - {-# INLINE getComp #-} setComp c arr = arr { lComp = c } {-# INLINE setComp #-} - unsafeMakeArray Seq sz f = runIdentity $ unsafeGenerateM Seq sz (return . f) - unsafeMakeArray (ParOn wss) sz f = LArray (ParOn wss) $ List $ unsafePerformIO $ - withScheduler' wss $ \scheduler -> - loopM_ 0 (< sz) (+ 1) (scheduleWork scheduler . return . f) - {-# INLINE unsafeMakeArray #-} + makeArray comp sz f = LArray comp $ List $ unsafePerformIO $ + withScheduler comp $ \scheduler -> + loopM_ 0 (< coerce sz) (+ 1) (scheduleWork scheduler . return . f) + {-# INLINE makeArray #-} instance ( Index ix @@ -310,12 +251,10 @@ , Elt L ix e ~ Array L (Lower ix) e ) => Construct L ix e where - getComp = lComp - {-# INLINE getComp #-} setComp c arr = arr {lComp = c} {-# INLINE setComp #-} - unsafeMakeArray = unsafeGenerateN - {-# INLINE unsafeMakeArray #-} + makeArray = unsafeGenerateN + {-# INLINE makeArray #-} -- TODO: benchmark against using unsafeGenerateM directly unsafeGenerateN :: @@ -326,31 +265,29 @@ -> Sz ix -> (ix -> e) -> Array r ix e -unsafeGenerateN Seq sz f = runIdentity $ unsafeGenerateM Seq sz (return . f) -unsafeGenerateN c@(ParOn wss) sz f = unsafePerformIO $ do - let !(m, szL) = unconsDim sz - xs <- withScheduler' wss $ \scheduler -> - loopM_ 0 (< m) (+ 1) $ \i -> scheduleWork scheduler $ - unsafeGenerateM c szL $ \ix -> return $ f (consDim i ix) - return $! foldr' cons (empty c) xs +unsafeGenerateN comp sz f = unsafePerformIO $ do + let !(m, szL) = unconsSz sz + xs <- withScheduler comp $ \scheduler -> + loopM_ 0 (< coerce m) (+ 1) $ \i -> scheduleWork scheduler $ + generateRaggedM comp szL $ \ix -> return $ f (consDim i ix) + return $! foldr' consR (emptyR comp) xs {-# INLINE unsafeGenerateN #-} toListArray :: (Construct L ix e, Source r ix e) => Array r ix e -> Array L ix e -toListArray !arr = - unsafeMakeArray (getComp arr) (size arr) (unsafeIndex arr) +toListArray !arr = makeArray (getComp arr) (size arr) (unsafeIndex arr) {-# INLINE toListArray #-} -instance {-# OVERLAPPING #-} (Ragged L ix e, Show e) => Show (Array L ix e) where - show arr = " " ++ raggedFormat show "\n " arr +instance (Ragged L ix e, Show e) => Show (Array L ix e) where + showsPrec = showsArrayLAsPrec (Proxy :: Proxy L) -instance {-# OVERLAPPING #-} (Ragged L ix e, Show e) => - Show (Array LN ix e) where - show arr = show (fromNested arr :: Array L ix e) +instance (Ragged L ix e, Show e) => Show (Array LN ix e) where + show arr = " " ++ raggedFormat show "\n " arrL + where arrL = fromNested arr :: Array L ix e showN :: (String -> a -> String) -> String -> [a] -> String @@ -361,46 +298,58 @@ L.intersperse (lnPrefix ++ ", ") (map (fShow (lnPrefix ++ " ")) ls) ++ [lnPrefix, "]"]) -instance ( Ragged L ix e - , Source r ix e - , Show e - ) => - Show (Array r ix e) where - show = showArray (showsTypeRep (typeRep (Proxy :: Proxy r)) " ") +showsArrayLAsPrec :: + forall r ix e. (Ragged L ix e, Typeable r, Show e) + => Proxy r + -> Int + -> Array L ix e -- Array to show + -> ShowS +showsArrayLAsPrec pr n arr = + opp . + ("Array " ++) . + showsTypeRep (typeRep pr) . + (' ':) . + showsPrec 1 (getComp arr) . (" (" ++) . shows (size arr) . (")\n" ++) . shows lnarr . clp + where + (opp, clp) = + if n == 0 + then (id, id) + else (('(':), ("\n)" ++)) + lnarr = toNested arr -showArray :: - forall r ix e. (Ragged L ix e, Source r ix e, Show e) - => String - -> Array r ix e - -> String -showArray tyStr arr = - "(Array " ++ tyStr ++ - showComp (getComp arr) ++ " (" ++ - (show (size arr)) ++ ")\n" ++ - show (makeArray (getComp arr) (size arr) (evaluateAt arr) :: Array L ix e) ++ ")" - where showComp Seq = "Seq" - showComp Par = "Par" - showComp c = "(" ++ show c ++ ")" +showsArrayPrec :: + forall r r' ix ix' e. (Ragged L ix' e, Load r ix e, Source r' ix' e, Show e) + => (Array r ix e -> Array r' ix' e) -- ^ Modifier + -> Int + -> Array r ix e -- Array to show + -> ShowS +showsArrayPrec f n arr = showsArrayLAsPrec (Proxy :: Proxy r) n larr + where + arr' = f arr + larr = makeArray (getComp arr') (size arr') (evaluate' arr') :: Array L ix' e +showArrayList + :: Show arr => [arr] -> String -> String +showArrayList arrs = ('[':) . go arrs . (']':) + where + go [] = id + go [x] = (' ':) . shows x . ('\n':) + go (x:xs) = (' ':) . shows x . ("\n," ++) . go xs instance {-# OVERLAPPING #-} OuterSlice L Ix1 e where unsafeOuterSlice (LArray _ xs) = (coerce xs !!) {-# INLINE unsafeOuterSlice #-} - outerLength = length . (coerce :: Array LN Ix1 e -> [e]). lData - {-# INLINE outerLength #-} instance Ragged L ix e => OuterSlice L ix e where unsafeOuterSlice arr' i = go 0 arr' where go n arr = - case uncons arr of - Nothing -> errorIx "Data.Massiv.Core.List.unsafeOuterSlice" (outerLength arr') i + case unconsR arr of + Nothing -> throw $ IndexOutOfBoundsException (Sz (headDim (unSz (size arr')))) i Just (x, _) | n == i -> x Just (_, xs) -> go (n + 1) xs {-# INLINE unsafeOuterSlice #-} - outerLength = length . (coerce :: Array LN ix e -> [Elt LN ix e]) . lData - {-# INLINE outerLength #-} diff --git a/src/Data/Massiv/Core/Scheduler.hs b/src/Data/Massiv/Core/Scheduler.hs deleted file mode 100644 --- a/src/Data/Massiv/Core/Scheduler.hs +++ /dev/null @@ -1,262 +0,0 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE RecordWildCards #-} -{-# LANGUAGE ScopedTypeVariables #-} --- | --- Module : Data.Massiv.Core.Scheduler --- Copyright : (c) Alexey Kuleshevich 2018 --- License : BSD3 --- Maintainer : Alexey Kuleshevich --- Stability : experimental --- Portability : non-portable --- -module Data.Massiv.Core.Scheduler - ( Scheduler - , numWorkers - , scheduleWork - , withScheduler - , withScheduler' - , withScheduler_ - , divideWork - , divideWork_ - ) where - -import Control.Concurrent (ThreadId, forkOnWithUnmask, - getNumCapabilities, killThread) -import Control.Concurrent.MVar -import Control.DeepSeq -import Control.Exception (SomeException, catch, mask, - mask_, throwIO, try, - uninterruptibleMask_) -import Control.Monad (void, forM) -import Control.Monad.Primitive (RealWorld) -import Data.IORef (IORef, atomicModifyIORef', - newIORef, readIORef) -import Data.Massiv.Core.Index.Class (Index (totalElem)) -import Data.Massiv.Core.Iterator (loop) -import Data.Primitive.Array (Array, MutableArray, indexArray, - newArray, unsafeFreezeArray, - writeArray) -import System.IO.Unsafe (unsafePerformIO) -import System.Mem.Weak - -data Job = Job (IO ()) - | Retire - -data Scheduler a = Scheduler - { jobsCountIORef :: !(IORef Int) - , jobQueueMVar :: !(MVar [Job]) - , resultsMVar :: !(MVar (MutableArray RealWorld a)) - , workers :: !Workers - , numCapabilities :: {-# UNPACK #-} !Int - } - -numWorkers :: Scheduler a -> Int -numWorkers = numCapabilities - -data Workers = Workers { workerThreadIds :: ![ThreadId] - , workerJobDone :: !(MVar (Maybe SomeException)) - , workerJobQueue :: !(MVar [Job]) - } - - --- | Helper function that allows scheduling work to be done in parallel. Use --- `withScheduler` to be able to get to a `Scheduler`. -scheduleWork :: Scheduler a -- ^ Scheduler to use - -> IO a -- ^ Action to hand of to a worker - -> IO () -scheduleWork Scheduler {..} jobAction = - modifyMVar_ jobQueueMVar $ \jobs -> do - jix <- atomicModifyIORef' jobsCountIORef $ \jc -> (jc + 1, jc) - let job = - Job $ do - jobResult <- jobAction - withMVar resultsMVar $ \resArray -> do - writeArray resArray jix jobResult - putMVar (workerJobDone workers) Nothing - return (job : jobs) - - -uninitialized :: a -uninitialized = error "Data.Array.Massiv.Scheduler: uncomputed job result" - - --- | Execute some action that needs a resource. Perform different cleanup actions depending if --- thataction resulted in an error or was successful. Sort of like `bracket` and `bracketOnError` --- with info about exception combined. -bracketWithException :: forall a b c d . - IO a -- ^ Acquire resource - -> (a -> IO b) -- ^ Run after successfull execution - -> (SomeException -> a -> IO c) -- ^ Run if execution resulted in exception. - -> (a -> IO d) -- ^ Execute an action that actually needs that resource. - -> IO d -bracketWithException before afterSuccess afterError thing = mask $ \restore -> do - x <- before - eRes <- try $ restore (thing x) - case eRes of - Left (exc :: SomeException) -> do - _ :: Either SomeException c <- try $ uninterruptibleMask_ $ afterError exc x - throwIO exc - Right y -> do - _ <- uninterruptibleMask_ $ afterSuccess x - return y - --- | Run arbitrary computations in parallel. A pool of workers is initialized, unless Worker --- Stations list is empty and a global worker pool is currently available. All of those workers will --- be stealing work that you can schedule using `scheduleWork`. The order in which work is scheduled --- will be the same as the order of the resuts of those computations, stored withing the resulting --- array. Size of the array, which is also the first element in the returned tuple, will match the --- number of times `scheduleWork` has been invoked. This function blocks until all of the submitted --- jobs has finished or one of them resulted in an exception, which will be re-thrown here. --- --- __Important__: In order to get work done truly in parallel, program needs to be compiled with --- @-threaded@ GHC flag and executed with @+RTS -N@. --- -withScheduler :: [Int] -- ^ Worker Stations, i.e. capabilities. Empty list will - -- result in utilization of all available capabilities. - -> (Scheduler a -> IO b) -- ^ Action that will be scheduling all - -- the work. - -> IO (Int, Array a) -withScheduler wss submitJobs = do - jobsCountIORef <- newIORef 0 - jobQueueMVar <- newMVar [] - resultsMVar <- newEmptyMVar - bracketWithException - (do mWeakWorkers <- - if null wss - then tryTakeMVar globalWorkersMVar - else return Nothing - mGlobalWorkers <- maybe (return Nothing) deRefWeak mWeakWorkers - let toWorkers w = return (mWeakWorkers, w) - maybe (hireWorkers wss >>= toWorkers) toWorkers mGlobalWorkers) - (\(mWeakWorkers, workers) -> - case mWeakWorkers of - Nothing -> - putMVar (workerJobQueue workers) $ - replicate (length (workerThreadIds workers)) Retire - Just weak -> putMVar globalWorkersMVar weak) - (\_ (mWeakWorkers, workers) -> - case mWeakWorkers of - Nothing -> mapM_ killThread (workerThreadIds workers) - Just weakWorkers -> do - finalize weakWorkers - newWeakWorkers <- hireWeakWorkers globalWorkersMVar - putMVar globalWorkersMVar newWeakWorkers) - (\(_, workers) -> do - let scheduler = - Scheduler {numCapabilities = length $ workerThreadIds workers, ..} - _ <- submitJobs scheduler - jobCount <- readIORef jobsCountIORef - marr <- newArray jobCount uninitialized - putMVar resultsMVar marr - jobQueue <- takeMVar jobQueueMVar - putMVar (workerJobQueue workers) $ reverse jobQueue - waitTillDone scheduler - arr <- unsafeFreezeArray marr - return (jobCount, arr)) - - --- | Just like `withScheduler`, but returns computed results in a list, instead --- of an array. -withScheduler' :: [Int] -> (Scheduler a -> IO b) -> IO [a] -withScheduler' wss submitJobs = do - (jc, arr) <- withScheduler wss submitJobs - return $ - loop (jc - 1) (>= 0) (subtract 1) [] $ \i acc -> indexArray arr i : acc - - --- | Just like `withScheduler`, but discards the results. -withScheduler_ :: [Int] -> (Scheduler a -> IO b) -> IO () -withScheduler_ wss submitJobs = void $ withScheduler wss submitJobs - - --- | Same as `divideWork`, but discard the result. -divideWork_ :: Index ix - => [Int] -> ix -> (Scheduler a -> Int -> Int -> Int -> IO b) -> IO () -divideWork_ wss sz submit = void $ divideWork wss sz submit - - --- | Linearly (row-major first) and equally divide work among available workers. Submit function --- will receive a `Scheduler`, length of each chunk, total number of elements, as well as where --- chunks end and slack begins. Slack work will get picked up by the first worker, that has finished --- working on his chunk. Returns list with results in the same order that work was submitted -divideWork :: Index ix - => [Int] -- ^ Worker Stations (capabilities) - -> ix -- ^ Size - -> (Scheduler a -> Int -> Int -> Int -> IO b) -- ^ Submit function - -> IO [a] -divideWork wss sz submit - | totalElem sz == 0 = return [] - | otherwise = - withScheduler' wss $ \scheduler -> do - let !totalLength = totalElem sz - !chunkLength = totalLength `quot` numWorkers scheduler - !slackStart = chunkLength * numWorkers scheduler - submit scheduler chunkLength totalLength slackStart - --- | Wait till workers finished with all submitted jobs, but raise an exception if either of them --- has died. Raised exception is the same one that was the cause of worker's death. -waitTillDone :: Scheduler a -> IO () -waitTillDone Scheduler {..} = readIORef jobsCountIORef >>= waitTill 0 - where - waitTill jobsDone jobsCount - | jobsDone == jobsCount = return () - | otherwise = do - mExc <- takeMVar (workerJobDone workers) - case mExc of - Just exc -> throwIO exc - Nothing -> waitTill (jobsDone + 1) jobsCount - - --- | Worker can either be doing work, waiting for a job, or going into retirement. Temp workers are --- rarely in waiting state, unless there is simply not enough work for all workers in the --- pool. Unlike temp workers, global workers do spend quite a bit of time waiting for work and they --- are never retired, but ruthlessly killed. -runWorker :: MVar [Job] -> IO () -runWorker jobsMVar = do - jobs <- takeMVar jobsMVar - case jobs of - (Job job:rest) -> putMVar jobsMVar rest >> job >> runWorker jobsMVar - (Retire:rest) -> putMVar jobsMVar rest - [] -> runWorker jobsMVar - - --- | Used whenever a pool of new workers is needed. If list is empty all capabilities are utilized, --- otherwise each element in the list will be an argument to `forkOn`. -hireWorkers :: [Int] -> IO Workers -hireWorkers wss = do - wss' <- - if null wss - then do - wNum <- getNumCapabilities - return [0 .. wNum - 1] - else return wss - workerJobQueue <- newEmptyMVar - workerJobDone <- newEmptyMVar - workerThreadIds <- - forM wss' $ \ws -> - mask_ $ - forkOnWithUnmask ws $ \unmask -> do - catch - (unmask $ runWorker workerJobQueue) - (unmask . putMVar workerJobDone . Just) - workerThreadIds `deepseq` return Workers {..} - --- | Global workers are the most utilized ones, therefore they are rarily restarted, in particular, --- only in case when one of them dies of an exception. Weak reference is used so workers don't --- continue running after MVar has been cleaned up by the GC. Each global worker has his own --- station, i.e. global workers always span all available capabilities. -globalWorkersMVar :: MVar (Weak Workers) -globalWorkersMVar = unsafePerformIO $ do - workersMVar <- newEmptyMVar - weakWorkers <- hireWeakWorkers workersMVar - putMVar workersMVar weakWorkers - return workersMVar -{-# NOINLINE globalWorkersMVar #-} - - --- | Hire workers under weak pointers. Finalizer will kill all the workers. These will be used as --- global workers -hireWeakWorkers :: key -> IO (Weak Workers) -hireWeakWorkers k = do - workers <- hireWorkers [] - mkWeak k workers (Just (mapM_ killThread (workerThreadIds workers))) diff --git a/tests/Data/Massiv/Array/Delayed/InterleavedSpec.hs b/tests/Data/Massiv/Array/Delayed/InterleavedSpec.hs new file mode 100644 --- /dev/null +++ b/tests/Data/Massiv/Array/Delayed/InterleavedSpec.hs @@ -0,0 +1,40 @@ +{-# LANGUAGE MonoLocalBinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeApplications #-} +module Data.Massiv.Array.Delayed.InterleavedSpec + ( spec + ) where + +import Data.Massiv.CoreArbitrary as A + + +prop_EqDelayed :: + (Ragged L ix Int, Load D ix Int, Load DI ix Int) + => Array D ix Int + -> Property +prop_EqDelayed arr = computeAs P arr === computeAs P (toInterleaved arr) + + +prop_Resize :: + (Ragged L ix Int, Load D ix Int, Load DI ix Int) + => Array DI ix Int + -> Property +prop_Resize arr = + computeAs P (resize' k arr) === computeAs P (resize' k arrD) + where + arrD = fromInterleaved arr + k = Sz (totalElem (size arr)) + +spec :: Spec +spec = + describe "Interleaved same as Delayed" $ do + it "EqDelayed Ix1" $ property $ prop_EqDelayed @Ix1 + it "EqDelayed Ix2" $ property $ prop_EqDelayed @Ix2 + it "EqDelayed Ix3" $ property $ prop_EqDelayed @Ix3 + it "EqDelayed Ix4" $ property $ prop_EqDelayed @Ix4 + it "EqDelayed Ix5" $ property $ prop_EqDelayed @Ix5 + it "Resize Ix1" $ property $ prop_Resize @Ix1 + it "Resize Ix2" $ property $ prop_Resize @Ix2 + it "Resize Ix3" $ property $ prop_Resize @Ix3 + it "Resize Ix4" $ property $ prop_Resize @Ix4 + it "Resize Ix5" $ property $ prop_Resize @Ix5 diff --git a/tests/Data/Massiv/Array/Delayed/WindowedSpec.hs b/tests/Data/Massiv/Array/Delayed/WindowedSpec.hs --- a/tests/Data/Massiv/Array/Delayed/WindowedSpec.hs +++ b/tests/Data/Massiv/Array/Delayed/WindowedSpec.hs @@ -1,19 +1,17 @@ -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE MonoLocalBinds #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MonoLocalBinds #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE NamedFieldPuns #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE UndecidableInstances #-} module Data.Massiv.Array.Delayed.WindowedSpec (spec) where -import Data.Massiv.Array.Delayed -import Data.Massiv.Array.Unsafe -import Data.Massiv.CoreArbitrary as A -import Data.Typeable -import Test.Hspec -import Test.QuickCheck +import Data.Massiv.Array.Delayed +import Data.Massiv.Array.Unsafe +import Data.Massiv.CoreArbitrary as A +import Data.Typeable data ArrDW ix e = ArrDW (Array D ix e) (Array DW ix e) @@ -41,11 +39,11 @@ let sz = size arr ArrDW arr <$> if totalElem sz == 0 - then return (unsafeMakeArray (getComp arr) sz (unsafeIndex arr)) + then return (makeArray (getComp arr) sz (unsafeIndex arr)) else do - wix <- flip (liftIndex2 mod) sz <$> arbitrary - wsz <- liftIndex (+1) . flip (liftIndex2 mod) (liftIndex2 (-) sz wix) <$> arbitrary - return $ makeWindowedArray arr wix wsz (unsafeIndex arr) + wix <- flip (liftIndex2 mod) (unSz sz) <$> arbitrary + wsz <- liftIndex (+1) . flip (liftIndex2 mod) (liftIndex2 (-) (unSz sz) wix) <$> arbitrary + return $ makeWindowedArray arr wix (Sz wsz) (unsafeIndex arr) prop_EqDelayed :: @@ -54,7 +52,7 @@ computeAs P arrD === computeAs P arrDW prop_EqDelayedStride :: - (Ragged L ix Int, Load DW ix Int) => Proxy ix -> Stride ix -> ArrDW ix Int -> Property + (Ragged L ix Int, StrideLoad DW ix Int) => Proxy ix -> Stride ix -> ArrDW ix Int -> Property prop_EqDelayedStride _ stride (ArrDW arrD arrDW) = computeWithStrideAs P stride arrD === computeWithStrideAs P stride arrDW diff --git a/tests/Data/Massiv/Array/DelayedSpec.hs b/tests/Data/Massiv/Array/DelayedSpec.hs --- a/tests/Data/Massiv/Array/DelayedSpec.hs +++ b/tests/Data/Massiv/Array/DelayedSpec.hs @@ -1,15 +1,12 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} module Data.Massiv.Array.DelayedSpec (spec) where -import Data.Massiv.Array.Delayed -import Data.Massiv.Array.Unsafe -import Data.Massiv.Array -import Data.Massiv.CoreArbitrary as A -import Data.Proxy -import Test.Hspec -import Test.QuickCheck +import Data.Massiv.Array +import Data.Massiv.Array.Unsafe +import Data.Massiv.CoreArbitrary as A +import Data.Proxy downsampleArr :: Source r ix e => Stride ix -> Array r ix e -> Array D ix e @@ -36,22 +33,24 @@ computeWithStride stride (toInterleaved arr) === computeAs U (downsampleArr stride arr) prop_computeWithStrideWindowedEqDownsample :: - (Ragged L ix Int, Load DW ix Int) + (Ragged L ix Int, StrideLoad DW ix Int) => proxy ix -> Stride ix -> ArrIx D ix Int -> Property prop_computeWithStrideWindowedEqDownsample _ stride (ArrIx arr _) = - computeWithStride stride (makeWindowedArray arr zeroIndex (size arr) (unsafeIndex arr)) === + computeWithStride stride (insertWindow arr (Window zeroIndex (size arr) (unsafeIndex arr) Nothing)) === + -- Below triggers a bug in ghc-8.0 which results in a deadlock. + -- computeWithStride stride (makeWindowedArray arr zeroIndex (size arr) (unsafeIndex arr)) === computeAs U (downsampleArr stride arr) delayedSpec :: - (Arbitrary ix, CoArbitrary ix, Load DW ix Int, Ragged L ix Int) + (Arbitrary ix, CoArbitrary ix, StrideLoad DW ix Int, Ragged L ix Int) => String -> proxy ix - -> SpecWith () -delayedSpec dimName proxy = do + -> Spec +delayedSpec dimName proxy = describe dimName $ do it "computeWithStrideEqDownsample" $ property $ prop_computeWithStrideEqDownsample proxy it "computeWithStrideInterleavedEqDownsample" $ diff --git a/tests/Data/Massiv/Array/Manifest/VectorSpec.hs b/tests/Data/Massiv/Array/Manifest/VectorSpec.hs --- a/tests/Data/Massiv/Array/Manifest/VectorSpec.hs +++ b/tests/Data/Massiv/Array/Manifest/VectorSpec.hs @@ -1,21 +1,19 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} module Data.Massiv.Array.Manifest.VectorSpec (spec) where -import Data.Massiv.CoreArbitrary -import Data.Massiv.Array.Manifest.Vector -import Data.Proxy -import Data.Typeable -import qualified Data.Vector as VB -import qualified Data.Vector.Generic as VG -import qualified Data.Vector.Primitive as VP -import qualified Data.Vector.Storable as VS -import qualified Data.Vector.Unboxed as VU -import Test.Hspec -import Test.QuickCheck +import Data.Massiv.Array.Manifest.Vector +import Data.Massiv.CoreArbitrary +import Data.Proxy +import Data.Typeable +import qualified Data.Vector as VB +import qualified Data.Vector.Generic as VG +import qualified Data.Vector.Primitive as VP +import qualified Data.Vector.Storable as VS +import qualified Data.Vector.Unboxed as VU prop_castToFromVector :: ( VG.Vector (VRepr r) Int @@ -32,7 +30,8 @@ prop_toFromVector :: forall r ix v. - ( Mutable r ix Int + ( Construct r ix Int + , Mutable r ix Int , Mutable (ARepr v) ix Int , VRepr (ARepr v) ~ v , Eq (Array r ix Int) @@ -46,11 +45,22 @@ -> Arr r ix Int -> Property prop_toFromVector _ _ _ (Arr arr) = - arr === fromVector (getComp arr) (size arr) (toVector arr :: v Int) + arr === fromVector' (getComp arr) (size arr) (toVector arr :: v Int) + toFromVectorSpec :: Spec -toFromVectorSpec = do - let it_prop name r = describe name $ do +toFromVectorSpec = do + it_prop "Unboxed" U + it_prop "Primitive" P + it_prop "Storable" S + it_prop "BoxedStrict" B + where + it_prop name r = + describe name $ do + describe "CastToFrom" $ do + it "Ix1" $ property $ prop_castToFromVector (Proxy :: Proxy Ix1) r + it "Ix2" $ property $ prop_castToFromVector (Proxy :: Proxy Ix2) r + it "Ix3" $ property $ prop_castToFromVector (Proxy :: Proxy Ix3) r describe "Through Boxed Vector" $ do it "Ix1" $ property $ prop_toFromVector (Proxy :: Proxy VB.Vector) (Proxy :: Proxy Ix1) r it "Ix2" $ property $ prop_toFromVector (Proxy :: Proxy VB.Vector) (Proxy :: Proxy Ix2) r @@ -63,25 +73,7 @@ describe "Through Storable Vector" $ do it "Ix1" $ property $ prop_toFromVector (Proxy :: Proxy VS.Vector) (Proxy :: Proxy Ix1) r it "Ix2" $ property $ prop_toFromVector (Proxy :: Proxy VS.Vector) (Proxy :: Proxy Ix2) r - it_prop "Unboxed" U - it_prop "Primitive" P - it_prop "Storable" S - it_prop "BoxedStrict" B -castToFromVectorSpec :: Spec -castToFromVectorSpec = do - let it_prop name r = describe name $ do - it "Ix1" $ property $ prop_castToFromVector (Proxy :: Proxy Ix1) r - it "Ix2" $ property $ prop_castToFromVector (Proxy :: Proxy Ix2) r - it "Ix3" $ property $ prop_castToFromVector (Proxy :: Proxy Ix3) r - it_prop "Unboxed" U - it_prop "Primitive" P - it_prop "Storable" S - it_prop "BoxedStrict" B - - spec :: Spec -spec = do - describe "toFromVector" toFromVectorSpec - describe "castToFromVector" castToFromVectorSpec +spec = describe "toFromVector" toFromVectorSpec diff --git a/tests/Data/Massiv/Array/ManifestSpec.hs b/tests/Data/Massiv/Array/ManifestSpec.hs --- a/tests/Data/Massiv/Array/ManifestSpec.hs +++ b/tests/Data/Massiv/Array/ManifestSpec.hs @@ -1,22 +1,25 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE MonoLocalBinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MonoLocalBinds #-} {-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeApplications #-} module Data.Massiv.Array.ManifestSpec (spec) where -import Data.ByteString as S -import Data.ByteString.Builder as S -import Data.ByteString.Lazy as SL -import Data.Massiv.CoreArbitrary as A -import Data.Word (Word8) -import Test.Hspec -import Test.QuickCheck +import Data.ByteString as S +import Data.ByteString.Builder as S +import Data.ByteString.Lazy as SL +import Data.Massiv.CoreArbitrary as A +import Data.Word (Word8) -- ByteString -prop_toFromByteString :: Array P Ix1 Word8 -> Property +prop_toFromByteString :: Manifest r Ix1 Word8 => Array r Ix1 Word8 -> Property prop_toFromByteString arr = toManifest arr === fromByteString (getComp arr) (toByteString arr) +prop_castToFromByteString :: Array S Ix1 Word8 -> Property +prop_castToFromByteString arr = arr === castFromByteString (getComp arr) (castToByteString arr) + + prop_fromToByteString :: Comp -> [Word8] -> Property prop_fromToByteString comp ls = bs === toByteString (fromByteString comp bs) where bs = S.pack ls @@ -26,9 +29,11 @@ where bs = toByteString arr conversionSpec :: Spec -conversionSpec = do +conversionSpec = describe "ByteString" $ do - it "to/from ByteString" $ property prop_toFromByteString + it "castTo/TromByteString" $ property prop_castToFromByteString + it "to/from ByteString P" $ property (prop_toFromByteString @P) + it "to/from ByteString S" $ property (prop_toFromByteString @S) it "from/to ByteString" $ property prop_fromToByteString it "toBuilder" $ property prop_toBuilder diff --git a/tests/Data/Massiv/Array/MutableSpec.hs b/tests/Data/Massiv/Array/MutableSpec.hs --- a/tests/Data/Massiv/Array/MutableSpec.hs +++ b/tests/Data/Massiv/Array/MutableSpec.hs @@ -1,18 +1,21 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE MonoLocalBinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MonoLocalBinds #-} {-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} module Data.Massiv.Array.MutableSpec (spec) where -import Control.Monad.ST -import Data.Functor.Identity -import Data.Massiv.CoreArbitrary as A -import Data.Proxy -import Test.Hspec -import Test.QuickCheck -import Test.QuickCheck.Monadic -import Test.QuickCheck.Function +import Control.Concurrent.Async +import Control.Monad.ST +import Data.Functor.Identity +import Data.List as L +import Data.Massiv.Array.Mutable.Atomic +import Data.Massiv.Array.Unsafe +import Data.Massiv.CoreArbitrary as A +import Data.Proxy + prop_MapMapM :: (Show (Array r ix Int), Eq (Array r ix Int), Mutable r ix Int) => r -> Proxy ix -> Fun Int Int -> ArrTiny D ix Int -> Property prop_MapMapM r _ f (ArrTiny arr) = @@ -25,17 +28,74 @@ runIdentity (A.imapMR r (\ix e -> return $ apply f (ix, e)) arr) -prop_generateMakeST :: (Show (Array r ix Int), Eq (Array r ix Int), Mutable r ix Int) => - r -> Proxy ix -> Arr r ix Int -> Property +prop_generateMakeST :: + (Show (Array r ix Int), Eq (Array r ix Int), Mutable r ix Int) + => r + -> Proxy ix + -> Arr r ix Int + -> Property prop_generateMakeST _ _ (Arr arr) = - arr === runST (generateArray (getComp arr) (size arr) (return . evaluateAt arr)) + arr === runST (generateArrayS (getComp arr) (size arr) (return . evaluate' arr)) prop_generateMakeIO :: (Show (Array r ix Int), Eq (Array r ix Int), Mutable r ix Int) => r -> Proxy ix -> Arr r ix Int -> Property prop_generateMakeIO _ _ (Arr arr) = monadicIO $ do - arr' <- run $ generateArray (getComp arr) (size arr) (return . evaluateAt arr) + arr' <- run $ generateArray (getComp arr) (size arr) (evaluateM arr) return (arr === arr') +prop_atomicModifyIntArrayMany :: ArrIx P Ix2 Int -> Array B Ix1 Int -> Property +prop_atomicModifyIntArrayMany (ArrIx arr ix) barr = + monadicIO $ do + xs <- + run $ do + marr <- thaw arr + mbarr' <- mapConcurrently (atomicModifyIntArray marr ix . const) barr + x <- A.read' marr ix + let xs = maybe (error "atomicModifyIntArray read'") toList (Prelude.sequenceA mbarr') + pure (x : xs) + return (L.sort (index' arr ix : toList barr) === L.sort xs) + + + +prop_atomicReadIntArrayMany :: Array P Ix2 Int -> Array B Ix1 Ix2 -> Property +prop_atomicReadIntArrayMany arr bix = monadicIO $ do + run $ do + marr <- thaw arr + as :: Array N Ix1 (Maybe Int) <- forM bix (A.read marr) + as' <- forM bix (atomicReadIntArray marr) + pure (as === as') + + +prop_atomicWriteIntArrayMany :: Array P Ix2 Int -> Array B Ix1 Ix2 -> (Fun Ix2 Int) -> Property +prop_atomicWriteIntArrayMany arr bix f = + monadicIO $ + run $ do + marr <- thaw arr + marr' <- unsafeThaw arr + bs :: Array N Ix1 Bool <- forM bix (\ix -> write marr ix (apply f ix)) + bs' <- forM bix (\ix -> atomicWriteIntArray marr' ix (apply f ix)) + arrRes <- unsafeFreeze (getComp arr) marr + arrRes' <- unsafeFreeze (getComp arr) marr' + pure (bs === bs' .&&. arrRes === arrRes') + + + +prop_unfoldrList :: Sz1 -> Fun Word (Int, Word) -> Word -> Property +prop_unfoldrList sz1 f i = + conjoin $ + L.zipWith + (===) + (A.toList (runST $ unfoldrPrimM_ @P Seq sz1 (pure . apply f) i)) + (L.unfoldr (Just . apply f) i) + +prop_unfoldrReverseUnfoldl :: Sz1 -> Fun Word (Int, Word) -> Word -> Property +prop_unfoldrReverseUnfoldl sz1 f i = + runST (unfoldrPrimM_ @P Seq sz1 (pure . apply f) i) === + rev (runST (unfoldlPrimM_ @P Seq sz1 (pure . swapTuple . apply f) i)) + where swapTuple (x, y) = (y, x) + rev a = computeAs P $ backpermute' sz1 (\ix1 -> unSz sz1 - ix1 - 1) a + + mutableSpec :: ( Show r , Show (Array r Ix3 Int) @@ -47,10 +107,13 @@ , Mutable r Ix3 Int , Mutable r Ix1 Int , Mutable r Ix2 Int + , Construct r Ix3 Int + , Construct r Ix1 Int + , Construct r Ix2 Int ) => r -> SpecWith () -mutableSpec r = do +mutableSpec r = describe (show r) $ do describe "map == mapM" $ do it "Ix1" $ property $ prop_MapMapM r (Proxy :: Proxy Ix1) @@ -80,4 +143,12 @@ spec :: Spec -spec = describe "GenerateM" generateSpec +spec = do + describe "GenerateM" generateSpec + describe "AtomicIntArraySpec" $ do + it "atomicReadIntArrayMany" $ property prop_atomicReadIntArrayMany + it "atomicWriteIntArrayMany" $ property prop_atomicWriteIntArrayMany + it "atomicModifyIntArrayMany" $ property prop_atomicModifyIntArrayMany + describe "Unfolding" $ do + it "unfoldrList" $ property prop_unfoldrList + it "unfoldrReverseUnfoldl" $ property prop_unfoldrReverseUnfoldl diff --git a/tests/Data/Massiv/Array/Numeric/IntegralSpec.hs b/tests/Data/Massiv/Array/Numeric/IntegralSpec.hs new file mode 100644 --- /dev/null +++ b/tests/Data/Massiv/Array/Numeric/IntegralSpec.hs @@ -0,0 +1,37 @@ +module Data.Massiv.Array.Numeric.IntegralSpec + ( spec + ) where + +import Data.Massiv.Array as A +import Data.Massiv.Array.Numeric.Integral +import Test.Hspec + +gaussian :: Float -> Float +gaussian x = exp (x ^ (2 :: Int)) + +spec :: Spec +spec = do + let (a, b) = (0, 2) + integrator rule = rule Seq N (\ scale -> gaussian . scale) a b (Sz1 1) + describe "Integral Approximation" $ do + it "Midpoint Rule" $ do + integrator midpointRule 4 ! 0 `shouldBe` 14.485613 + integrator midpointRule 8 ! 0 `shouldBe` 15.905677 + integrator midpointRule 16 ! 0 `shouldBe` 16.311854 + integrator midpointRule 32 ! 0 `shouldBe` 16.417171 + integrator midpointRule 64 ! 0 `shouldBe` 16.443748 + integrator midpointRule 128 ! 0 `shouldBe` 16.450407 + it "Trapezoid Rule" $ do + integrator trapezoidRule 4 ! 0 `shouldBe` 20.644558 + integrator trapezoidRule 8 ! 0 `shouldBe` 17.565086 + integrator trapezoidRule 16 ! 0 `shouldBe` 16.735381 + integrator trapezoidRule 32 ! 0 `shouldBe` 16.523618 + integrator trapezoidRule 64 ! 0 `shouldBe` 16.470394 + integrator trapezoidRule 128 ! 0 `shouldBe` 16.457073 + it "Simspon's Rule" $ do + integrator simpsonsRule 4 ! 0 `shouldBe` 17.353626 + integrator simpsonsRule 8 ! 0 `shouldBe` 16.538595 + integrator simpsonsRule 16 ! 0 `shouldBe` 16.458815 + integrator simpsonsRule 32 ! 0 `shouldBe` 16.453030 + integrator simpsonsRule 64 ! 0 `shouldBe` 16.452653 + integrator simpsonsRule 128 ! 0 `shouldBe` 16.452629 diff --git a/tests/Data/Massiv/Array/Ops/ConstructSpec.hs b/tests/Data/Massiv/Array/Ops/ConstructSpec.hs --- a/tests/Data/Massiv/Array/Ops/ConstructSpec.hs +++ b/tests/Data/Massiv/Array/Ops/ConstructSpec.hs @@ -1,32 +1,33 @@ {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE LambdaCase #-} {-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE GADTs #-} -module Data.Massiv.Array.Ops.ConstructSpec where - -import Data.Massiv.Array as A -import Data.Massiv.CoreArbitrary as A -import Data.Proxy -import qualified GHC.Exts as GHC (IsList (..)) -import Prelude as P -import Prelude hiding (map) -import Test.Hspec -import Test.QuickCheck +module Data.Massiv.Array.Ops.ConstructSpec (spec) where +import Data.List as L +import Data.Massiv.CoreArbitrary as A +import Data.Proxy +import qualified GHC.Exts as GHC (IsList(..)) +import Prelude as P prop_rangeEqRangeStep1 :: Int -> Int -> Property -prop_rangeEqRangeStep1 from to = range Seq from to === rangeStep Par from 1 to +prop_rangeEqRangeStep1 from to = range Seq from to === rangeStep' Par from 1 to prop_rangeEqEnumFromN :: Int -> Int -> Property -prop_rangeEqEnumFromN from to = range Seq from to === enumFromN Par from (to - from) +prop_rangeEqEnumFromN from to = range Seq from to === enumFromN Par from (Sz (to - from)) prop_rangeStepEqEnumFromStepN :: Int -> NonZero Int -> Int -> Property prop_rangeStepEqEnumFromStepN from (NonZero step) sz = - rangeStep Seq from step (from + step * sz) === enumFromStepN Par from step sz + rangeStep' Seq from step (from + step * sz) === enumFromStepN Par from step (Sz sz) prop_rangeStepExc :: Int -> Int -> Property prop_rangeStepExc from to = - assertSomeException (computeAs U (rangeStep Seq from 0 to)) + assertException + (\case + IndexZeroException _ -> True + _ -> False) + (computeAs U (rangeStep' Seq from 0 to)) prop_toFromListIsList :: (Show (Array U ix Int), GHC.IsList (Array U ix Int), Index ix) @@ -37,7 +38,7 @@ prop_toFromList :: - forall ix . (Show (Array U ix Int), Nested LN ix Int, Nested L ix Int, Ragged L ix Int) + forall ix . (Show (Array U ix Int), Nested LN ix Int, Ragged L ix Int) => Proxy ix -> Arr U ix Int -> Property @@ -61,7 +62,7 @@ prop_excFromToListIx3 comp ls3 | P.null (P.concat (P.concat ls3)) = classify True "Expected Success" $ counterexample (show arr) $ totalElem (size arr) === 0 - | P.all (head lsL ==) lsL && (P.and (P.map (P.all (head (head lsLL) ==)) lsLL)) = + | P.all (head lsL ==) lsL && P.all (P.all (head (head lsLL) ==)) lsLL = classify True "Expected Success" $ counterexample (show arr) $ resultLs === ls3 | otherwise = classify True "Expected Failure" $ assertSomeException resultLs where @@ -96,21 +97,29 @@ mkIntermediate t = A.fromList Seq [t + 50, t + 75] initArr :: Array N Ix1 (Array U Ix1 Int) -initArr = makeArray Seq 3 (\ x -> mkIntermediate x) +initArr = makeArray Seq (Sz1 3) mkIntermediate initArr2 :: Array N Ix2 (Array U Ix1 Int) -initArr2 = makeArray Seq (2 :. 2) (\ (x :. y) -> mkIntermediate (x+y)) +initArr2 = makeArray Seq (Sz 2) (\ (x :. y) -> mkIntermediate (x+y)) +prop_unfoldrList :: Sz1 -> Fun Word (Int, Word) -> Word -> Property +prop_unfoldrList sz1 f i = + conjoin $ + L.zipWith + (===) + (A.toList (computeAs P $ unfoldrS_ Seq sz1 (apply f) i)) + (L.unfoldr (Just . apply f) i) + specExpand :: Spec specExpand = do it "expandOuter" $ compute (expandOuter 2 A.index' initArr :: Array D Ix2 Int) `shouldBe` - resize' (2 :. 3) (fromList Seq [50, 51, 52, 75, 76, 77] :: Array U Ix1 Int) + resize' (Sz2 2 3) (fromList Seq [50, 51, 52, 75, 76, 77] :: Array U Ix1 Int) it "expandInner" $ compute (expandInner 2 A.index' initArr :: Array D Ix2 Int) `shouldBe` - resize' (3 :. 2) (fromList Seq [50, 75, 51, 76, 52, 77] :: Array U Ix1 Int) + resize' (Sz2 3 2) (fromList Seq [50, 75, 51, 76, 52, 77] :: Array U Ix1 Int) it "expandwithin" $ compute (expandWithin Dim1 2 A.index' initArr2 :: Array D Ix3 Int) `shouldBe` - resize' (2 :> 2 :. 2) (fromList Seq [50, 75, 51, 76, 51, 76, 52, 77] :: Array U Ix1 Int) + resize' (Sz 2) (fromList Seq [50, 75, 51, 76, 51, 76, 52, 77] :: Array U Ix1 Int) it "expandwithin'" $ compute (expandWithin' 1 2 A.index' initArr2 :: Array D Ix3 Int) `shouldBe` - resize' (2 :> 2 :. 2) (fromList Seq [50, 75, 51, 76, 51, 76, 52, 77] :: Array U Ix1 Int) + resize' (Sz 2) (fromList Seq [50, 75, 51, 76, 51, 76, 52, 77] :: Array U Ix1 Int) spec :: Spec spec = do @@ -118,3 +127,4 @@ describe "Ix2" specIx2 describe "Ix3" specIx3 describe "Expand" specExpand + describe "Unfolding" $ it "unfoldrS_" $ property prop_unfoldrList diff --git a/tests/Data/Massiv/Array/Ops/FoldSpec.hs b/tests/Data/Massiv/Array/Ops/FoldSpec.hs --- a/tests/Data/Massiv/Array/Ops/FoldSpec.hs +++ b/tests/Data/Massiv/Array/Ops/FoldSpec.hs @@ -1,20 +1,16 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE MonoLocalBinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MonoLocalBinds #-} {-# LANGUAGE MultiParamTypeClasses #-} module Data.Massiv.Array.Ops.FoldSpec (spec) where -import qualified Data.Foldable as F -import Data.Massiv.CoreArbitrary as A -import Data.Semigroup -import Prelude hiding (map, product, sum) -import qualified Prelude as P (length, sum) -import Test.Hspec -import Test.QuickCheck -import Test.QuickCheck.Function -import Test.QuickCheck.Monadic +import qualified Data.Foldable as F +import Data.Massiv.CoreArbitrary as A +import Data.Semigroup +import Prelude hiding (map, product, sum) + prop_SumSEqSumP :: Index ix => proxy ix -> Array D ix Int -> Bool prop_SumSEqSumP _ arr = sum arr == sum (setComp Par arr) @@ -25,43 +21,21 @@ prop_NestedFoldP :: Array D Ix1 (Array D Ix1 Int) -> Bool prop_NestedFoldP arr = sum (setComp Par (map sum $ setComp Par arr)) == sum (map sum arr) -prop_FoldrOnP :: Int -> [Int] -> ArrP D Ix1 Int -> Property -prop_FoldrOnP wId wIds (ArrP arr) = - P.length arr > P.length wIds ==> monadicIO $ do - res <- run $ ifoldrOnP wIdsNE (\_ -> (+)) 0 (:) [] arr - if P.length arr `mod` P.length wIdsNE == 0 - then assert (P.length res == P.length wIdsNE) - else assert (P.length res == P.length wIdsNE + 1) - assert (P.sum res == sum arr) - where - wIdsNE = wId : wIds -prop_FoldlOnP :: Int -> [Int] -> ArrP D Ix1 Int -> Property -prop_FoldlOnP wId wIds (ArrP arr) = - P.length arr > P.length wIds ==> monadicIO $ do - res <- run $ ifoldlOnP wIdsNE (\a _ x -> a + x) 0 (flip (:)) [] arr - if P.length arr `mod` P.length wIdsNE == 0 - then assert (P.length res == P.length wIdsNE) - else assert (P.length res == P.length wIdsNE + 1) - assert (P.sum res == sum arr) - where - wIdsNE = wId : wIds - - specFold :: (Arbitrary ix, CoArbitrary ix, Index ix, Show (Array D ix Int)) => proxy ix -> String -> Spec -specFold proxy dimStr = do +specFold proxy dimStr = describe dimStr $ do it "sumS Eq sumP" $ property $ prop_SumSEqSumP proxy it "prodS Eq prodP" $ property $ prop_ProdSEqProdP proxy foldOpsProp :: (Source P ix Int) => proxy ix -> Fun Int Bool -> ArrTiny1 P ix Int -> Property foldOpsProp _ f (ArrTiny1 arr) = - (A.maximum arr === getMax (foldMono Max arr)) .&&. - (A.minimum arr === getMin (foldSemi Min maxBound arr)) .&&. + (A.maximum' arr === getMax (foldMono Max arr)) .&&. + (A.minimum' arr === getMin (foldSemi Min maxBound arr)) .&&. (A.sum arr === F.sum ls) .&&. (A.product (A.map ((+ 0.1) . (fromIntegral :: Int -> Double)) arr) === getProduct (foldMono (Product . (+ 0.1) . fromIntegral) arr)) .&&. @@ -77,8 +51,6 @@ specFold (Nothing :: Maybe Ix1) "Ix1" specFold (Nothing :: Maybe Ix2) "Ix2" it "Nested Parallel Fold" $ property prop_NestedFoldP - it "FoldrOnP" $ property $ prop_FoldrOnP - it "FoldlOnP" $ property $ prop_FoldlOnP describe "Foldable Props" $ do it "Ix1" $ property $ foldOpsProp (Nothing :: Maybe Ix1) it "Ix2" $ property $ foldOpsProp (Nothing :: Maybe Ix2) diff --git a/tests/Data/Massiv/Array/Ops/MapSpec.hs b/tests/Data/Massiv/Array/Ops/MapSpec.hs --- a/tests/Data/Massiv/Array/Ops/MapSpec.hs +++ b/tests/Data/Massiv/Array/Ops/MapSpec.hs @@ -1,81 +1,109 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE GADTs #-} +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} {-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} module Data.Massiv.Array.Ops.MapSpec (spec) where -import Data.Massiv.CoreArbitrary as A -import Prelude as P -import Test.Hspec -import Test.QuickCheck +import Control.Monad.ST +import Data.Foldable as F +import Data.Massiv.Array.Unsafe +import Data.Massiv.CoreArbitrary as A +import Prelude as P + + prop_zipUnzip :: - (Arbitrary ix, CoArbitrary ix, Index ix, Show (Array D ix Int)) - => proxy ix - -> Array D ix Int + (Index ix, Show (Array D ix Int)) + => Array D ix Int -> Array D ix Int -> Property -prop_zipUnzip _ arr1 arr2 = +prop_zipUnzip arr1 arr2 = (extract' zeroIndex sz arr1, extract' zeroIndex sz arr2) === A.unzip (A.zip arr1 arr2) - where sz = liftIndex2 min (size arr1) (size arr2) + where sz = Sz (liftIndex2 min (unSz (size arr1)) (unSz (size arr2))) prop_zipFlip :: - (Arbitrary ix, CoArbitrary ix, Index ix, Show (Array D ix Int), Show (Array D ix (Int, Int))) - => proxy ix - -> Array D ix Int + (Index ix, Show (Array D ix (Int, Int))) + => Array D ix Int -> Array D ix Int -> Property -prop_zipFlip _ arr1 arr2 = +prop_zipFlip arr1 arr2 = A.zip arr1 arr2 === A.map (\(e2, e1) -> (e1, e2)) (A.zip arr2 arr1) prop_zipUnzip3 :: - (Arbitrary ix, CoArbitrary ix, Index ix, Show (Array D ix Int)) - => proxy ix - -> Array D ix Int + (Index ix, Show (Array D ix Int)) + => Array D ix Int -> Array D ix Int -> Array D ix Int -> Property -prop_zipUnzip3 _ arr1 arr2 arr3 = +prop_zipUnzip3 arr1 arr2 arr3 = (extract' zeroIndex sz arr1, extract' zeroIndex sz arr2, extract' zeroIndex sz arr3) === A.unzip3 (A.zip3 arr1 arr2 arr3) where - sz = liftIndex2 min (liftIndex2 min (size arr1) (size arr2)) (size arr3) + sz = + Sz (liftIndex2 min (liftIndex2 min (unSz (size arr1)) (unSz (size arr2))) (unSz (size arr3))) prop_zipFlip3 :: - ( Arbitrary ix - , CoArbitrary ix - , Index ix - , Show (Array D ix Int) - , Show (Array D ix (Int, Int, Int)) - ) - => proxy ix - -> Array D ix Int + (Index ix, Show (Array D ix (Int, Int, Int))) + => Array D ix Int -> Array D ix Int -> Array D ix Int -> Property -prop_zipFlip3 _ arr1 arr2 arr3 = +prop_zipFlip3 arr1 arr2 arr3 = A.zip3 arr1 arr2 arr3 === A.map (\(e3, e2, e1) -> (e1, e2, e3)) (A.zip3 arr3 arr2 arr1) + + +prop_itraverseA :: + (Index ix, Show (Array U ix Int)) => Array D ix Int -> Fun (ix, Int) Int -> Property +prop_itraverseA arr fun = + alt_imapM (\ix -> Just . applyFun2Compat fun ix) arr === + itraverseAR U (\ix -> Just . applyFun2Compat fun ix) arr + + mapSpec :: + forall ix. ( Arbitrary ix , CoArbitrary ix , Index ix + , Function ix + , Show (Array U ix Int) , Show (Array D ix Int) , Show (Array D ix (Int, Int)) , Show (Array D ix (Int, Int, Int)) ) - => proxy ix - -> Spec -mapSpec proxy = do + => Spec +mapSpec = do describe "Zipping" $ do - it "zipUnzip" $ property $ prop_zipUnzip proxy - it "zipFlip" $ property $ prop_zipFlip proxy - it "zipUnzip3" $ property $ prop_zipUnzip3 proxy - it "zipFlip3" $ property $ prop_zipFlip3 proxy + it "zipUnzip" $ property $ prop_zipUnzip @ix + it "zipFlip" $ property $ prop_zipFlip @ix + it "zipUnzip3" $ property $ prop_zipUnzip3 @ix + it "zipFlip3" $ property $ prop_zipFlip3 @ix + describe "Traversing" $ do + it "itraverseA" $ property $ prop_itraverseA @ix spec :: Spec spec = do - describe "Ix1" $ mapSpec (Nothing :: Maybe Ix1) - describe "Ix2" $ mapSpec (Nothing :: Maybe Ix2) - describe "Ix3" $ mapSpec (Nothing :: Maybe Ix3) - describe "Ix4" $ mapSpec (Nothing :: Maybe Ix4) + describe "Ix1" $ mapSpec @Ix1 + describe "Ix2" $ mapSpec @Ix2 + describe "Ix3" $ mapSpec @Ix3 + describe "Ix4" $ mapSpec @Ix4 + + + +alt_imapM + :: (Applicative f, Mutable r2 t1 b, Source r1 t1 t2) => + (t1 -> t2 -> f b) -> Array r1 t1 t2 -> f (Array r2 t1 b) +alt_imapM f arr = fmap loadList $ P.traverse (uncurry f) $ foldrS (:) [] (zipWithIndex arr) + where + loadList xs = + runST $ do + marr <- unsafeNew (size arr) + _ <- F.foldlM (\i e -> unsafeLinearWrite marr i e >> return (i + 1)) 0 xs + unsafeFreeze (getComp arr) marr + {-# INLINE loadList #-} + +zipWithIndex :: forall r ix e . Source r ix e => Array r ix e -> Array D ix (ix, e) +zipWithIndex arr = A.zip (range Seq zeroIndex (unSz (size arr))) arr +{-# INLINE zipWithIndex #-} diff --git a/tests/Data/Massiv/Array/Ops/SliceSpec.hs b/tests/Data/Massiv/Array/Ops/SliceSpec.hs --- a/tests/Data/Massiv/Array/Ops/SliceSpec.hs +++ b/tests/Data/Massiv/Array/Ops/SliceSpec.hs @@ -1,35 +1,43 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE MonoLocalBinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MonoLocalBinds #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeFamilies #-} module Data.Massiv.Array.Ops.SliceSpec (spec) where -import Data.Massiv.Array.Unsafe -import Data.Massiv.CoreArbitrary -import Test.Hspec -import Test.QuickCheck - +import Control.Applicative ((<|>)) +import Control.Exception +import Data.Massiv.Array.Unsafe +import Data.Massiv.CoreArbitrary ----------- -- Size -- ----------- --- extract - -prop_ExtractEqualsExtractFromTo - :: (Eq (Array (EltRepr r ix) ix e), Arbitrary (Array r ix e), Size r ix e) - => proxy (r, ix, e) -> SzIx ix -> Array r ix e -> Bool +prop_ExtractEqualsExtractFromTo :: + ( Eq (Array (EltRepr r ix) ix e) + , Show (Array (EltRepr r ix) ix e) + , Extract r ix e + ) + => proxy (r, ix, e) + -> SzIx ix + -> Array r ix e + -> Property prop_ExtractEqualsExtractFromTo _ (SzIx (Sz eIx) sIx) arr = - extractFromTo sIx eIx arr == extract sIx (liftIndex2 (-) eIx sIx) arr - - + (extractFromToM sIx eIx arr <|> Nothing) === extractM sIx (Sz (liftIndex2 (-) eIx sIx)) arr -specSizeN - :: (Eq (Array (EltRepr r ix) ix e), Arbitrary (Array r ix e), Show (Array r ix e), Arbitrary ix, Size r ix e) - => proxy (r, ix, e) -> Spec -specSizeN proxy = do - describe "extract" $ do +specSizeN :: + ( Eq (Array (EltRepr r ix) ix e) + , Show (Array (EltRepr r ix) ix e) + , Arbitrary (Array r ix e) + , Show (Array r ix e) + , Arbitrary ix + , Extract r ix e + ) + => proxy (r, ix, e) + -> Spec +specSizeN proxy = + describe "extract" $ it "ExtractEqualsExtractFromTo" $ property $ prop_ExtractEqualsExtractFromTo proxy @@ -38,137 +46,149 @@ ----------- -prop_SliceRight :: (Arbitrary (Array r ix e), Slice r ix e, OuterSlice r ix e, Eq (Elt r ix e)) - => proxy (r, ix, e) -> Int -> Array r ix e -> Bool -prop_SliceRight _ i arr = (arr !?> i) == (arr (dimensions (size arr), i)) +prop_SliceRight :: + (Slice r ix e, OuterSlice r ix e, Eq (Elt r ix e), Show (Elt r ix e)) + => proxy (r, ix, e) + -> Int + -> Array r ix e + -> Property +prop_SliceRight _ i arr = + either (Left . displayException) Right (arr !?> i) === + either (Left . displayException) Right (arr (dimensions (size arr), i)) -prop_SliceLeft :: (Arbitrary (Array r ix e), Slice r ix e, InnerSlice r ix e, Eq (Elt r ix e)) - => proxy (r, ix, e) -> Int -> Array r ix e -> Bool -prop_SliceLeft _ i arr = (arr (1, i)) +prop_SliceLeft :: + (Slice r ix e, InnerSlice r ix e, Eq (Elt r ix e), Show (Elt r ix e)) + => proxy (r, ix, e) + -> Int + -> Array r ix e + -> Property +prop_SliceLeft _ i arr = + either (Left . displayException) Right (arr (1, i)) -prop_SliceIndexDim2D :: ArrIx D Ix2 Int -> Bool +prop_SliceIndexDim2D :: ArrIx D Ix2 Int -> Property prop_SliceIndexDim2D (ArrIx arr ix@(i :. j)) = - val == evaluateAt (arr i) j + val === evaluate' (arr i) j where val = unsafeIndex arr ix -prop_SliceIndexDim2RankD :: ArrIx D Ix2 Int -> Bool +prop_SliceIndexDim2RankD :: ArrIx D Ix2 Int -> Property prop_SliceIndexDim2RankD (ArrIx arr ix@(i :. j)) = - val == evaluateAt (arr (2, i)) j && - val == evaluateAt (arr (1, j)) i + val === evaluate' (arr (2, i)) j .&&. + val === evaluate' (arr (1, j)) i where val = unsafeIndex arr ix -prop_SliceIndexDim3D :: ArrIx D Ix3 Int -> Bool +prop_SliceIndexDim3D :: ArrIx D Ix3 Int -> Property prop_SliceIndexDim3D (ArrIx arr ix@(i :> j :. k)) = - val == evaluateAt (arr i !> j) k && - val == evaluateAt (arr i) j && - val == evaluateAt (arr !> i i !> j) k .&&. + val === evaluate' (arr i) j .&&. + val === evaluate' (arr !> i Bool +prop_SliceIndexDim3RankD :: ArrIx D Ix3 Int -> Property prop_SliceIndexDim3RankD (ArrIx arr ix@(i :> j :. k)) = - val == evaluateAt (arr (3, i) (2, j)) k && - val == evaluateAt (arr (3, i) (1, k)) j && - val == evaluateAt (arr (2, j) (2, i)) k && - val == evaluateAt (arr (2, j) (1, k)) i && - val == evaluateAt (arr (1, k) (2, i)) j && - val == evaluateAt (arr (1, k) (1, j)) i + val === evaluate' (arr (3, i) (2, j)) k .&&. + val === evaluate' (arr (3, i) (1, k)) j .&&. + val === evaluate' (arr (2, j) (2, i)) k .&&. + val === evaluate' (arr (2, j) (1, k)) i .&&. + val === evaluate' (arr (1, k) (2, i)) j .&&. + val === evaluate' (arr (1, k) (1, j)) i where val = unsafeIndex arr ix -prop_SliceIndexDim2M :: ArrIx M Ix2 Int -> Bool +prop_SliceIndexDim2M :: ArrIx M Ix2 Int -> Property prop_SliceIndexDim2M (ArrIx arr ix@(i :. j)) = - val == (arr !> i ! j) && - val == (arr i ! j) .&&. + val === (arr Bool +prop_SliceIndexDim2RankM :: ArrIx M Ix2 Int -> Property prop_SliceIndexDim2RankM (ArrIx arr ix@(i :. j)) = - val == (arr (2, i) ! j) && - val == (arr (1, j) ! i) + val === (arr (2, i) ! j) .&&. + val === (arr (1, j) ! i) where val = unsafeIndex arr ix -prop_SliceIndexDim3M :: ArrIx M Ix3 Int -> Bool +prop_SliceIndexDim3M :: ArrIx M Ix3 Int -> Property prop_SliceIndexDim3M (ArrIx arr ix@(i :> j :. k)) = - val == (arr i !> j ! k) && - val == (arr i ! j) && - val == (arr !> i i !> j ! k) .&&. + val === (arr i ! j) .&&. + val === (arr !> i Bool +prop_SliceIndexDim3RankM :: ArrIx M Ix3 Int -> Property prop_SliceIndexDim3RankM (ArrIx arr ix@(i :> j :. k)) = - val == (arr (3, i) (2, j) ! k) && - val == (arr (3, i) (1, k) ! j) && - val == (arr (2, j) (2, i) ! k) && - val == (arr (2, j) (1, k) ! i) && - val == (arr (1, k) (2, i) ! j) && - val == (arr (1, k) (1, j) ! i) + val === (arr (3, i) (2, j) ! k) .&&. + val === (arr (3, i) (1, k) ! j) .&&. + val === (arr (2, j) (2, i) ! k) .&&. + val === (arr (2, j) (1, k) ! i) .&&. + val === (arr (1, k) (2, i) ! j) .&&. + val === (arr (1, k) (1, j) ! i) where val = unsafeIndex arr ix -prop_SliceIndexDim4D :: ArrIx D Ix4 Int -> Bool +prop_SliceIndexDim4D :: ArrIx D Ix4 Int -> Property prop_SliceIndexDim4D (ArrIx arr ix@(i1 :> i2 :> i3 :. i4)) = - val == evaluateAt (arr !> i1 !> i2 !> i3) i4 && - val == evaluateAt (arr !> i1 !> i2 i1 i1 i2) i3 && - val == evaluateAt (arr i1 !> i2) i3 && - val == evaluateAt (arr i1 i1) i2 + val === evaluate' (arr !> i1 !> i2 !> i3) i4 .&&. + val === evaluate' (arr !> i1 !> i2 i1 i1 i2) i3 .&&. + val === evaluate' (arr i1 !> i2) i3 .&&. + val === evaluate' (arr i1 i1) i2 where val = unsafeIndex arr ix -prop_SliceIndexDim4RankD :: ArrIx D Ix4 Int -> Bool +prop_SliceIndexDim4RankD :: ArrIx D Ix4 Int -> Property prop_SliceIndexDim4RankD (ArrIx arr ix@(i1 :> i2 :> i3 :. i4)) = - val == unsafeIndex (arr (4, i1) (3, i2) (2, i3)) i4 && - val == unsafeIndex (arr (4, i1) (2, i3) (3, i2) (3, i1)) (i3 :. i4) && - val == unsafeIndex (arr (2, i3) (2, i2)) (i1 :. i4) && - val == unsafeIndex (arr (2, i3) (1, i4) !> i1) i2 && - val == unsafeIndex (arr (1, i4) !> i1 !> i2) i3 + val === unsafeIndex (arr (4, i1) (3, i2) (2, i3)) i4 .&&. + val === unsafeIndex (arr (4, i1) (2, i3) (3, i2) (3, i1)) (i3 :. i4) .&&. + val === unsafeIndex (arr (2, i3) (2, i2)) (i1 :. i4) .&&. + val === unsafeIndex (arr (2, i3) (1, i4) !> i1) i2 .&&. + val === unsafeIndex (arr (1, i4) !> i1 !> i2) i3 where - val = evaluateAt arr ix + val = evaluate' arr ix -prop_SliceIndexDim4RankM :: ArrIx M Ix4 Int -> Bool +prop_SliceIndexDim4RankM :: ArrIx M Ix4 Int -> Property prop_SliceIndexDim4RankM (ArrIx arr ix@(i1 :> i2 :> i3 :. i4)) = - val == (arr (4, i1) (3, i2) (2, i3) ! i4) && - val == (arr (4, i1) (2, i3) (3, i2) (3, i1) ! (i3 :. i4)) && - val == (arr (2, i3) (2, i2) ! (i1 :. i4)) && - val == (arr (2, i3) (1, i4) !> i1 ! i2) && - val == (arr (1, i4) !> i1 !> i2 ! i3) + val === (arr (4, i1) (3, i2) (2, i3) ! i4) .&&. + val === (arr (4, i1) (2, i3) (3, i2) (3, i1) ! (i3 :. i4)) .&&. + val === (arr (2, i3) (2, i2) ! (i1 :. i4)) .&&. + val === (arr (2, i3) (1, i4) !> i1 ! i2) .&&. + val === (arr (1, i4) !> i1 !> i2 ! i3) where val = unsafeIndex arr ix -prop_SliceIndexDim4M :: ArrIx M Ix4 Int -> Bool +prop_SliceIndexDim4M :: ArrIx M Ix4 Int -> Property prop_SliceIndexDim4M (ArrIx arr ix@(i1 :> i2 :> i3 :. i4)) = - val == (arr !> i1 !> i2 !> i3 ! i4) && - val == (arr !> i1 !> i2 i1 i1 i2 ! i3) && - val == (arr i1 !> i2 ! i3) && - val == (arr i1 i1 ! i2) + val === (arr !> i1 !> i2 !> i3 ! i4) .&&. + val === (arr !> i1 !> i2 i1 i1 i2 ! i3) .&&. + val === (arr i1 !> i2 ! i3) .&&. + val === (arr i1 i1 ! i2) where val = unsafeIndex arr ix @@ -176,14 +196,14 @@ specSliceN :: ( Arbitrary (Array r ix e) , Show (Array r ix e) - , Arbitrary ix , Slice r ix e , OuterSlice r ix e , InnerSlice r ix e , Eq (Elt r ix e) + , Show (Elt r ix e) ) => proxy (r, ix, e) -> Spec -specSliceN proxy = do +specSliceN proxy = describe "Slice" $ do it "SliceRight" $ property $ prop_SliceRight proxy it "SliceLeft" $ property $ prop_SliceLeft proxy @@ -192,29 +212,29 @@ spec :: Spec spec = do - describe "Ix1" $ do + describe "Ix1" $ specSizeN (Nothing :: Maybe (D, Ix1, Int)) describe "Ix2" $ do specSizeN (Nothing :: Maybe (D, Ix2, Int)) specSliceN (Nothing :: Maybe (D, Ix2, Int)) describe "SliceIndex" $ do - it "Delayed" $ property $ prop_SliceIndexDim2D - it "Rank - Delayed" $ property $ prop_SliceIndexDim2RankD - it "Manifest" $ property $ prop_SliceIndexDim2M - it "Rank - Manifest" $ property $ prop_SliceIndexDim2RankM + it "Delayed" $ property prop_SliceIndexDim2D + it "Rank - Delayed" $ property prop_SliceIndexDim2RankD + it "Manifest" $ property prop_SliceIndexDim2M + it "Rank - Manifest" $ property prop_SliceIndexDim2RankM describe "Ix3" $ do specSizeN (Nothing :: Maybe (D, Ix3, Int)) specSliceN (Nothing :: Maybe (D, Ix3, Int)) describe "SliceIndex" $ do - it "Delayed" $ property $ prop_SliceIndexDim3D - it "Rank - Delayed" $ property $ prop_SliceIndexDim3RankD - it "Manifest" $ property $ prop_SliceIndexDim3M - it "Rank - Manifest" $ property $ prop_SliceIndexDim3RankM + it "Delayed" $ property prop_SliceIndexDim3D + it "Rank - Delayed" $ property prop_SliceIndexDim3RankD + it "Manifest" $ property prop_SliceIndexDim3M + it "Rank - Manifest" $ property prop_SliceIndexDim3RankM describe "Ix4" $ do specSizeN (Nothing :: Maybe (D, Ix4, Int)) specSliceN (Nothing :: Maybe (D, Ix4, Int)) describe "SliceIndex" $ do - it "Delayed" $ property $ prop_SliceIndexDim4D - it "Rank - Delayed" $ property $ prop_SliceIndexDim4RankD - it "Manifest" $ property $ prop_SliceIndexDim4M - it "Rank - Manifest" $ property $ prop_SliceIndexDim4RankM + it "Delayed" $ property prop_SliceIndexDim4D + it "Rank - Delayed" $ property prop_SliceIndexDim4RankD + it "Manifest" $ property prop_SliceIndexDim4M + it "Rank - Manifest" $ property prop_SliceIndexDim4RankM diff --git a/tests/Data/Massiv/Array/Ops/TransformSpec.hs b/tests/Data/Massiv/Array/Ops/TransformSpec.hs --- a/tests/Data/Massiv/Array/Ops/TransformSpec.hs +++ b/tests/Data/Massiv/Array/Ops/TransformSpec.hs @@ -1,51 +1,112 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE MonoLocalBinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE MonoLocalBinds #-} {-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeApplications #-} module Data.Massiv.Array.Ops.TransformSpec (spec) where -import Data.Massiv.CoreArbitrary as A -import Data.Typeable (Typeable) -import Test.Hspec -import Test.QuickCheck +import Data.Massiv.CoreArbitrary as A +import Data.Sequence as S +import Prelude as P +import Data.Foldable as F (foldl', toList) +import Data.Maybe +prop_transposeOuterInner :: Arr D Ix2 Int -> Property +prop_transposeOuterInner (Arr arr) = transposeOuter arr === transpose arr -prop_ExtractAppend - :: (Eq e, Size r ix e, Source r ix e, Source (EltRepr r ix) ix e, Arbitrary (ArrIx r ix e)) - => proxy (r, ix, e) -> DimIx ix -> ArrIx r ix e -> Bool -prop_ExtractAppend _ (DimIx dim) (ArrIx arr ix) = - maybe False ((delay arr ==) . uncurry (append' dim)) $ - A.splitAt dim (getDim' ix dim) arr +prop_upsampleDownsample :: + (Show (Array P ix Int), Index ix) => ArrTiny P ix Int -> Stride ix -> Int -> Property +prop_upsampleDownsample (ArrTiny arr) stride fill = + arr === compute (downsample stride (computeAs P (upsample fill stride arr))) +prop_ExtractAppend + :: (Show (Array P ix Int), Index ix) + => DimIx ix -> ArrIx P ix Int -> Property +prop_ExtractAppend (DimIx dim) (ArrIx arr ix) = + arr === compute (uncurry (append' dim) $ A.splitAt' dim (getDim' ix dim) arr) -prop_transposeOuterInner :: Arr D Ix2 Int -> Property -prop_transposeOuterInner (Arr arr) = transposeOuter arr === transpose arr +prop_ConcatAppend + :: (Show (Array P ix Int), Index ix) + => DimIx ix -> Comp -> Sz ix -> NonEmptyList (Fun ix Int) -> Property +prop_ConcatAppend (DimIx dim) comp sz (NonEmpty fns) = + foldl1 (\arr -> computeAs P . append' dim arr) arrs === + computeAs P (concat' dim arrs) + where + arrs = P.map (makeArrayR P comp sz . apply) fns +prop_AppendMappend + :: Array D Ix1 Int -> Array D Ix1 Int -> Property +prop_AppendMappend arr1 arr2 = + computeAs P (append' 1 arr1 arr2) === computeAs P (toLoadArray arr1 <> toLoadArray arr2) -specN :: - ( Eq e - , Size r ix e - , Source r ix e - , Source (EltRepr r ix) ix e - , Typeable e - , Show (Array r ix e) - , Arbitrary (ArrIx r ix e) - ) - => proxy (r, ix, e) - -> Spec -specN r = do - it "ExtractAppend" $ property $ prop_ExtractAppend r +prop_ConcatMconcat + :: [Array D Ix1 Int] -> Property +prop_ConcatMconcat arrs = + computeAs P (concat' 1 (A.empty : arrs)) === computeAs P (mconcat (fmap toLoadArray arrs)) spec :: Spec spec = do it "transposeOuterInner" $ property prop_transposeOuterInner - describe "Delayed" $ do - describe "Ix1" $ specN (Nothing :: Maybe (D, Ix1, Int)) - describe "Ix2" $ specN (Nothing :: Maybe (D, Ix2, Int)) - describe "Ix3" $ specN (Nothing :: Maybe (D, Ix3, Int)) - describe "Ix4" $ specN (Nothing :: Maybe (D, Ix4, Int)) - describe "Unboxed" $ do - describe "Ix1" $ specN (Nothing :: Maybe (U, Ix1, Int)) - describe "Ix2" $ specN (Nothing :: Maybe (U, Ix2, Int)) - describe "Ix3" $ specN (Nothing :: Maybe (U, Ix3, Int)) - describe "Ix4" $ specN (Nothing :: Maybe (U, Ix4, Int)) + describe "upsampleDownsample" $ do + it "Ix1" $ property (prop_upsampleDownsample @Ix1) + it "Ix2" $ property (prop_upsampleDownsample @Ix2) + it "Ix3" $ property (prop_upsampleDownsample @Ix3) + it "Ix4" $ property (prop_upsampleDownsample @Ix4) + describe "ExtractAppend" $ do + it "Ix1" $ property (prop_ExtractAppend @Ix1) -- modifyMaxSuccess (`div` 10) + it "Ix2" $ property (prop_ExtractAppend @Ix2) + it "Ix3" $ property (prop_ExtractAppend @Ix3) + it "Ix4" $ property (prop_ExtractAppend @Ix4) + describe "ConcatAppend" $ do + it "Ix1" $ property (prop_ConcatAppend @Ix1) + it "Ix2" $ property (prop_ConcatAppend @Ix2) + it "Ix3" $ property (prop_ConcatAppend @Ix3) + it "Ix4" $ property (prop_ConcatAppend @Ix4) + describe "Monoid" $ do + it "Ix1" $ property prop_AppendMappend + it "Ix1" $ property prop_ConcatMconcat + describe "Sequence" $ do + it "ConsSnoc" $ property prop_ConsSnoc + it "UnconsUnsnoc" $ property prop_UnconsUnsnoc + +prop_UnconsUnsnoc :: Array D Ix1 Int -> Bool -> Property +prop_UnconsUnsnoc arr unconsFirst = + preJust $ do + (arr', u, s) <- + if unconsFirst + then do + (u, au) <- unconsM arr + (as, s) <- unsnocM au + pure (as, u, s) + else do + (as, s) <- unsnocM arr + (u, au) <- unconsM as + pure (au, u, s) + pure (computeAs U (A.snoc (A.cons u (toLoadArray (computeAs U arr'))) s) === compute arr) + +preJust :: Testable prop => Maybe prop -> Property +preJust m = isJust m ==> fromJust m + +prop_ConsSnoc :: Array D Ix1 Int -> [SeqOp Int] -> Property +prop_ConsSnoc arr ops = + A.toList (computeAs U (foldl' applyArraySeqOp (toLoadArray arr) ops)) === + F.toList (foldl' applySequenceSeqOp (S.fromList (A.toList arr)) ops) + +data SeqOp e = Cons e | Snoc e deriving (Eq, Show) + +instance Arbitrary e => Arbitrary (SeqOp e) where + arbitrary = do + e <- arbitrary + elements [Cons e, Snoc e] + +applyArraySeqOp :: Array DL Ix1 e -> SeqOp e -> Array DL Ix1 e +applyArraySeqOp arr = \case + Cons x -> A.cons x arr + Snoc x -> A.snoc arr x + + +applySequenceSeqOp :: Seq a -> SeqOp a -> Seq a +applySequenceSeqOp arr = \case + Cons x -> x <| arr + Snoc x -> arr |> x diff --git a/tests/Data/Massiv/Array/StencilSpec.hs b/tests/Data/Massiv/Array/StencilSpec.hs --- a/tests/Data/Massiv/Array/StencilSpec.hs +++ b/tests/Data/Massiv/Array/StencilSpec.hs @@ -9,11 +9,7 @@ import Control.DeepSeq (deepseq) import Data.Default (Default(def)) import Data.Massiv.CoreArbitrary as A -import Data.Maybe (fromJust) import Data.Proxy -import Test.Hspec -import Test.QuickCheck -import Test.QuickCheck.Function -- sum3x3Stencil :: Fractional a => Stencil Ix2 a a -- sum3x3Stencil = makeConvolutionStencil (3 :. 3) (1 :. 1) $ \ get -> @@ -24,7 +20,8 @@ singletonStencil :: (Index ix) => (Int -> Int) -> Stencil ix Int Int -singletonStencil f = makeStencil (pureIndex 1) zeroIndex $ \ get -> fmap f (get zeroIndex) +singletonStencil f = + makeStencil oneSz zeroIndex $ \ get -> fmap f (get zeroIndex) {-# INLINE singletonStencil #-} @@ -33,7 +30,7 @@ prop_MapSingletonStencil _ f b (ArrP arr) = computeAs U (mapStencil b (singletonStencil (apply f)) arr) == computeAs U (A.map (apply f) arr) -prop_MapSingletonStencilWithStride :: (Load DW ix Int, Manifest U ix Int) => +prop_MapSingletonStencilWithStride :: (StrideLoad DW ix Int, Manifest U ix Int) => Proxy ix -> Fun Int Int -> Border Int -> ArrP U ix Int -> Bool prop_MapSingletonStencilWithStride _ f b (ArrP arr) = computeWithStride oneStride (mapStencil b (singletonStencil (apply f)) arr) == @@ -42,15 +39,12 @@ -- Tests out of bounds stencil indexing prop_DangerousStencil :: Index ix => Proxy ix -> NonZero Int -> DimIx ix -> SzIx ix -> Property -prop_DangerousStencil _ (NonZero s) (DimIx r) (SzIx (Sz sz) ix) = - ix' `deepseq` assertSomeException $ makeStencil sz ix $ \get -> (get ix' :: Value Int) +prop_DangerousStencil _ (NonZero s) (DimIx r) (SzIx sz ix) = + ix' `deepseq` assertSomeException $ makeStencil sz ix $ \get -> get ix' :: Value Int where - ix' = - liftIndex (* signum s) $ - fromJust $ do - i <- getDim sz r - setDim zeroIndex r i + ix' = liftIndex (* signum s) (setDim' zeroIndex r (getDim' (unSz sz) r)) + stencilSpec :: Spec stencilSpec = do describe "MapSingletonStencil" $ do @@ -71,13 +65,15 @@ stencilDirection :: (Default a, Unbox a, Manifest r Ix2 a) => Ix2 -> Array r Ix2 a -> Array U Ix2 a -stencilDirection ix = computeAs U . mapStencil (Fill def) (makeStencil (3 :. 3) (1 :. 1) $ \f -> f ix) +stencilDirection ix = + computeAs U . mapStencil (Fill def) (makeStencil (Sz 3) (1 :. 1) $ \f -> f ix) stencilCorners :: (Default a, Unbox a, Manifest r Ix2 a) => Ix2 -> Ix2 -> Array r Ix2 a -> Array U Ix2 a -stencilCorners ixC ix = computeAs U . mapStencil (Fill def) (makeStencil (3 :. 3) ixC $ \f -> f ix) +stencilCorners ixC ix = computeAs U . mapStencil (Fill def) (makeStencil (Sz 3) ixC $ \f -> f ix) + stencilConvolution :: Spec stencilConvolution = do let xs3 :: Array U Ix1 Int @@ -113,7 +109,7 @@ describe "makeConvolutionStencil == makeConvolutionStencilFromKernel" $ do it "Sobel Horizontal" $ property $ \(arr :: Array U Ix2 Int) -> - applyStencil (makeConvolutionStencil (3 :. 3) 1 sobelX) arr === + applyStencil (makeConvolutionStencil (Sz 3) 1 sobelX) arr === applyStencil (makeConvolutionStencilFromKernel sobelKernelX) arr it "1x3" $ property $ \(arr :: Array U Ix1 Int) -> @@ -126,7 +122,7 @@ describe "makeCorrelationStencil == makeCorrelationStencilFromKernel" $ do it "Sobel Horizontal" $ property $ \(arr :: Array U Ix2 Int) -> - applyStencil (makeCorrelationStencil (3 :. 3) 1 sobelX) arr === + applyStencil (makeCorrelationStencil (Sz 3) 1 sobelX) arr === applyStencil (makeCorrelationStencilFromKernel sobelKernelX) arr it "1x3" $ property $ \(arr :: Array U Ix1 Int) -> @@ -180,7 +176,7 @@ let strideArr = mapStencil (Fill 0) stencil arr in computeWithStrideAs U stride strideArr `shouldBe` [[-4, 8], [2, 14]] it "map stencil with stride on larger array" $ - let largeArr = makeArrayR U Seq (5 :. 5) (succ . toLinearIndex (5 :. 5)) + let largeArr = makeArrayR U Seq (Sz 5) (succ . toLinearIndex (Sz 5)) strideArr = mapStencil (Fill 0) stencil largeArr in computeWithStrideAs U stride strideArr `shouldBe` [[-6, 1, 14], [-13, 9, 43], [4, 21, 44]] @@ -197,5 +193,4 @@ , [-1, 0, 1] ] rotate180 :: (Num ix, Index ix) => Array U ix Int -> Array U ix Int -rotate180 arr = computeAs U $ backpermute sz (\ix -> sz - 1 - ix) arr - where sz = size arr +rotate180 = computeAs U . transform' (\sz -> (sz, sz)) (\(Sz sz) f ix -> f (sz - 1 - ix)) diff --git a/tests/Data/Massiv/ArraySpec.hs b/tests/Data/Massiv/ArraySpec.hs new file mode 100644 --- /dev/null +++ b/tests/Data/Massiv/ArraySpec.hs @@ -0,0 +1,141 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MonoLocalBinds #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +-- Here are contained tests for all instances for all main classes +module Data.Massiv.ArraySpec + ( spec + ) where + +import Data.Massiv.Core +import Data.Massiv.CoreArbitrary as A + + +prop_Construct_makeArray_Manifest :: + forall r ix. (Load D ix Int, Ragged L ix Int, Source r ix Int, Construct r ix Int) + => Comp + -> Sz ix + -> Fun Int Int + -> Property +prop_Construct_makeArray_Manifest comp sz f = + makeArrayLinearR D comp sz (apply f) === + delay (setComp Seq (makeArray comp sz (apply f . toLinearIndex sz) :: Array r ix Int)) + +prop_Construct_makeArray_Delayed :: + forall r ix. (Load D ix Int, Ragged L ix Int, Load r ix Int, Construct r ix Int) + => Comp + -> Sz ix + -> Fun Int Int + -> Property +prop_Construct_makeArray_Delayed comp sz f = + makeArrayLinearR P comp sz (apply f) === + compute (setComp Seq (makeArrayLinear comp sz (apply f)) :: Array r ix Int) + +prop_Functor :: + forall r ix. + (Load D ix Int, Ragged L ix Int, Load r ix Int, Construct r ix Int, Functor (Array r ix)) + => Comp + -> Sz ix + -> Fun Int Int + -> Fun Int Int + -> Property +prop_Functor comp sz f g = + makeArrayLinearR P comp sz (apply g . apply f) === + compute (fmap (apply g) (makeArrayLinear comp sz (apply f) :: Array r ix Int)) + +prop_Extract :: + forall r ix. + ( Load D ix Int + , Ragged L ix Int + , Load (EltRepr r ix) ix Int + , Construct r ix Int + , Extract r ix Int + ) + => Comp + -> Sz ix + -> Fun Int Int + -> ix + -> Sz ix + -> Property +prop_Extract comp sz f start newSize = + (computeAs P <$> toStringException (extractM start newSize arrD)) === + (compute <$> toStringException (extractM start newSize arr)) + where + arrD = makeArrayLinearR D comp sz (apply f) + arr = makeArrayLinear comp sz (apply f) :: Array r ix Int + +prop_IxUnbox :: + forall ix. + ( Load D ix ix + , Ragged L ix ix + , Construct U ix ix + , Source U ix ix + ) + => Comp + -> Sz ix + -> Fun Int ix + -> Property +prop_IxUnbox comp sz f = + makeArrayLinearR D comp sz (apply f) === + delay (makeArrayLinear comp sz (apply f) :: Array U ix ix) + +prop_computeWithStride :: + forall r ix. (Load D ix Int, Ragged L ix Int, StrideLoad r ix Int, Construct r ix Int) + => Comp + -> Sz ix + -> Fun Int Int + -> Stride ix + -> Property +prop_computeWithStride comp sz f stride = + arr === computeWithStride stride arrL .&&. + arr === compute (fromStrideLoad stride arrL) + where + arrL = makeArrayLinear comp sz (apply f) :: Array r ix Int + arr = computeWithStrideAs P stride (makeArrayLinearR D comp sz (apply f)) + + +specCommon :: + forall ix. + (Arbitrary ix, Load D ix Int, StrideLoad DW ix Int, Ragged L ix Int, Ragged L ix ix, Unbox ix) + => Spec +specCommon = + describe "Construct" $ do + it "Construct_makeArray B" $ property $ prop_Construct_makeArray_Manifest @B @ix + it "Construct_makeArray N" $ property $ prop_Construct_makeArray_Manifest @N @ix + it "Construct_makeArray S" $ property $ prop_Construct_makeArray_Manifest @S @ix + it "Construct_makeArray P" $ property $ prop_Construct_makeArray_Manifest @P @ix + it "Construct_makeArray U" $ property $ prop_Construct_makeArray_Manifest @U @ix + it "Construct_makeArray M" $ property $ prop_Construct_makeArray_Manifest @M @ix + it "Construct_makeArray_Delayed DI" $ property $ prop_Construct_makeArray_Delayed @DI @ix + it "Construct_makeArray_Delayed DL" $ property $ prop_Construct_makeArray_Delayed @DL @ix + it "Construct_makeArray_Delayed DW" $ property $ prop_Construct_makeArray_Delayed @DW @ix + it "Construct_makeArray_Delayed M" $ property $ prop_Construct_makeArray_Delayed @M @ix + it "Functor D" $ property $ prop_Functor @D @ix + it "Functor DI" $ property $ prop_Functor @DI @ix + it "Functor DL" $ property $ prop_Functor @DL @ix + it "Functor DW" $ property $ prop_Functor @DW @ix + it "Extract DI" $ property $ prop_Extract @DI @ix + it "Extract D" $ property $ prop_Extract @D @ix + it "Extract B" $ property $ prop_Extract @B @ix + it "Extract N" $ property $ prop_Extract @N @ix + it "Extract S" $ property $ prop_Extract @S @ix + it "Extract U" $ property $ prop_Extract @U @ix + it "Extract M" $ property $ prop_Extract @M @ix + it "computeWithStride DI" $ property $ prop_computeWithStride @DI @ix + it "computeWithStride DW" $ property $ prop_computeWithStride @DW @ix + it "computeWithStride B" $ property $ prop_computeWithStride @B @ix + it "computeWithStride N" $ property $ prop_computeWithStride @N @ix + it "computeWithStride S" $ property $ prop_computeWithStride @S @ix + it "computeWithStride U" $ property $ prop_computeWithStride @U @ix + it "computeWithStride M" $ property $ prop_computeWithStride @M @ix + it "IxUnbox" $ property $ prop_IxUnbox @ix + + +spec :: Spec +spec = do + specCommon @Ix1 + specCommon @Ix2 + specCommon @Ix3 + specCommon @Ix4 + specCommon @Ix5 diff --git a/tests/Data/Massiv/Core/IndexSpec.hs b/tests/Data/Massiv/Core/IndexSpec.hs --- a/tests/Data/Massiv/Core/IndexSpec.hs +++ b/tests/Data/Massiv/Core/IndexSpec.hs @@ -1,63 +1,53 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} {-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GADTs #-} {-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} +{-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE CPP #-} -module Data.Massiv.Core.IndexSpec (Sz(..), SzZ(..), SzIx(..), DimIx(..), spec) where +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeOperators #-} +module Data.Massiv.Core.IndexSpec (SzNE(..), SzIx(..), DimIx(..), spec) where -import Control.Monad -import Data.Massiv.Core.Index hiding (type Sz, pattern Sz) -import Data.Functor.Identity -import Test.Hspec -import Test.QuickCheck -#if !MIN_VERSION_QuickCheck(2,10,0) -import Test.QuickCheck.Function -#endif +import Control.Monad +import Data.Functor.Identity +import Data.Massiv.Core.Index +import Data.Proxy +import Test.Hspec +import Test.QuickCheck +import Test.QuickCheck.Function +import Data.Typeable -- | Size that will result in a non-empty array -newtype Sz ix = Sz ix deriving Show - --- | Size that can have zero elements -newtype SzZ ix = SzZ ix deriving Show +newtype SzNE ix = SzNE (Sz ix) deriving Show -- | Dimension that is always within bounds of an index newtype DimIx ix = DimIx Dim deriving Show -instance Functor Sz where - fmap f (Sz sz) = Sz (f sz) - -instance Functor SzZ where - fmap f (SzZ sz) = SzZ (f sz) - +-- | Non-empty size together with an index that is within bounds of that index. data SzIx ix = SzIx (Sz ix) ix deriving Show instance (Index ix, Arbitrary ix) => Arbitrary (Sz ix) where arbitrary = do - sz <- liftIndex ((+1) . abs) <$> arbitrary + sz <- Sz . liftIndex abs <$> arbitrary if totalElem sz > 200000 then arbitrary - else return $ Sz sz + else return sz -instance (Index ix, Arbitrary ix) => Arbitrary (SzZ ix) where - arbitrary = do - sz <- liftIndex abs <$> arbitrary - if totalElem sz > 200000 - then arbitrary - else return $ SzZ sz +instance (Index ix, Arbitrary ix) => Arbitrary (SzNE ix) where + arbitrary = SzNE . Sz . liftIndex (+1) . unSz <$> arbitrary instance (Index ix, Arbitrary ix) => Arbitrary (Stride ix) where arbitrary = do Positive (Small x) <- arbitrary - Stride <$> liftIndex ((+1) . (`mod` x)) <$> arbitrary + Stride . liftIndex ((+1) . (`mod` min 6 x)) <$> arbitrary instance (Index ix, Arbitrary ix) => Arbitrary (SzIx ix) where arbitrary = do - Sz sz <- arbitrary + SzNE sz <- arbitrary -- Make sure index is within bounds: - SzIx (Sz sz) <$> flip (liftIndex2 mod) sz <$> arbitrary + SzIx sz . flip (liftIndex2 mod) (unSz sz) <$> arbitrary instance Arbitrary e => Arbitrary (Border e) where @@ -74,7 +64,7 @@ instance Index ix => Arbitrary (DimIx ix) where arbitrary = do n <- arbitrary - return $ DimIx (1 + (Dim n `mod` (dimensions (undefined :: ix)))) + return $ DimIx (1 + (Dim n `mod` dimensions (Proxy :: Proxy ix))) arbitraryIntIx :: Gen Int arbitraryIntIx = sized (\s -> resize (floor $ (sqrt :: Double -> Double) $ fromIntegral s) arbitrary) @@ -118,26 +108,11 @@ function = functionMap fromIx5 toIx5 - --- instance Arbitrary Ix2 where --- arbitrary = (:.) <$> arbitrary <*> arbitrary - --- instance (Arbitrary (Ix (n - 1))) => Arbitrary (IxN n) where --- arbitrary = (:>) <$> arbitrary <*> arbitrary - --- instance CoArbitrary Ix2 where --- coarbitrary (i :. j) = coarbitrary i . coarbitrary j - --- instance CoArbitrary (Ix (n - 1)) => CoArbitrary (IxN n) where --- coarbitrary (i :> ix) = coarbitrary i . coarbitrary ix - - - prop_IsSafeIx :: Index ix => proxy ix -> SzIx ix -> Bool -prop_IsSafeIx _ (SzIx (Sz sz) ix) = isSafeIndex sz ix +prop_IsSafeIx _ (SzIx sz ix) = isSafeIndex sz ix prop_RepairSafeIx :: Index ix => proxy ix -> SzIx ix -> Bool -prop_RepairSafeIx _ (SzIx (Sz sz) ix) = +prop_RepairSafeIx _ (SzIx sz ix) = ix == repairIndex sz ix (error "Impossible") (error "Impossible") prop_UnconsCons :: Index ix => proxy ix -> ix -> Bool @@ -147,28 +122,30 @@ prop_UnsnocSnoc _ ix = ix == uncurry snocDim (unsnocDim ix) prop_ToFromLinearIndex :: Index ix => proxy ix -> SzIx ix -> Property -prop_ToFromLinearIndex _ (SzIx (Sz sz) ix) = +prop_ToFromLinearIndex _ (SzIx sz ix) = isSafeIndex sz ix ==> ix == fromLinearIndex sz (toLinearIndex sz ix) -prop_FromToLinearIndex :: Index ix => proxy ix -> Sz ix -> Int -> Property -prop_FromToLinearIndex _ (Sz sz) i = +prop_FromToLinearIndex :: Index ix => proxy ix -> SzNE ix -> NonNegative Int -> Property +prop_FromToLinearIndex _ (SzNE sz) (NonNegative i) = totalElem sz >= i ==> i == toLinearIndex sz (fromLinearIndex sz i) prop_CountElements :: Index ix => proxy ix -> Int -> Sz ix -> Property -prop_CountElements _ thresh (Sz sz) = +prop_CountElements _ thresh sz = totalElem sz < thresh ==> totalElem sz == - iter zeroIndex sz (pureIndex 1) (<) 0 (\_ acc -> (acc + 1)) + iter zeroIndex (unSz sz) (pureIndex 1) (<) 0 (const (+ 1)) prop_IterMonotonic :: Index ix => proxy ix -> Int -> Sz ix -> Property -prop_IterMonotonic _ thresh (Sz sz) = +prop_IterMonotonic _ thresh sz = totalElem sz < thresh ==> fst $ - iter (liftIndex succ zeroIndex) sz (pureIndex 1) (<) (True, zeroIndex) $ \curIx (prevMono, prevIx) -> - let isMono = prevMono && prevIx < curIx - in isMono `seq` (isMono, curIx) + iter (liftIndex succ zeroIndex) (unSz sz) (pureIndex 1) (<) (True, zeroIndex) mono + where + mono curIx (prevMono, prevIx) = + let isMono = prevMono && prevIx < curIx + in isMono `seq` (isMono, curIx) prop_IterMonotonic' :: Index ix => proxy ix -> Int -> Sz ix -> Property -prop_IterMonotonic' _ thresh (Sz sz) = +prop_IterMonotonic' _ thresh sz = totalElem sz < thresh ==> if isM @@ -176,110 +153,112 @@ else error (show a) where (isM, a, _) = - iter (liftIndex succ zeroIndex) sz (pureIndex 1) (<) (True, [], zeroIndex) $ - \ curIx (prevMono, acc, prevIx) -> - let nAcc = (prevIx, curIx, prevIx < curIx) : acc - isMono = prevMono && prevIx < curIx - in isMono `seq` (isMono, nAcc, curIx) + iter (liftIndex succ zeroIndex) (unSz sz) (pureIndex 1) (<) (True, [], zeroIndex) mono + mono curIx (prevMono, acc, prevIx) = + let nAcc = (prevIx, curIx, prevIx < curIx) : acc + isMono = prevMono && prevIx < curIx + in isMono `seq` (isMono, nAcc, curIx) prop_IterMonotonicBackwards' :: Index ix => proxy ix -> Int -> Sz ix -> Property -prop_IterMonotonicBackwards' _ thresh (Sz sz) = +prop_IterMonotonicBackwards' _ thresh sz@(Sz szix) = totalElem sz < thresh ==> if isM then isM else error (show a) where - (isM, a, _) = - iter (liftIndex pred sz) zeroIndex (pureIndex (-1)) (>=) (True, [], sz) $ - \ curIx (prevMono, acc, prevIx) -> + (isM, a, _) = iter (liftIndex pred szix) zeroIndex (pureIndex (-1)) (>=) (True, [], szix) mono + mono curIx (prevMono, acc, prevIx) = let isMono = prevMono && prevIx > curIx nAcc = (prevIx, curIx, prevIx > curIx) : acc - in isMono `seq` (isMono, nAcc, curIx) + in isMono `seq` (isMono, nAcc, curIx) prop_IterMonotonicM :: Index ix => proxy ix -> Int -> Sz ix -> Property -prop_IterMonotonicM _ thresh (Sz sz) = +prop_IterMonotonicM _ thresh sz = totalElem sz < thresh ==> fst $ - runIdentity $ - iterM (liftIndex succ zeroIndex) sz (pureIndex 1) (<) (True, zeroIndex) $ \curIx (prevMono, prevIx) -> - let isMono = prevMono && prevIx < curIx - in return $ isMono `seq` (isMono, curIx) + runIdentity $ iterM (liftIndex succ zeroIndex) (unSz sz) (pureIndex 1) (<) (True, zeroIndex) mono + where + mono curIx (prevMono, prevIx) = + let isMono = prevMono && prevIx < curIx + in return $ isMono `seq` (isMono, curIx) prop_IterMonotonicBackwards :: Index ix => proxy ix -> Int -> Sz ix -> Property -prop_IterMonotonicBackwards _ thresh (Sz sz) = +prop_IterMonotonicBackwards _ thresh sz@(Sz szix) = totalElem sz < thresh ==> fst $ - iter (liftIndex pred sz) zeroIndex (pureIndex (-1)) (>=) (True, sz) $ \ curIx (prevMono, prevIx) -> - let isMono = prevMono && prevIx > curIx in isMono `seq` (isMono, curIx) + iter (liftIndex pred szix) zeroIndex (pureIndex (-1)) (>=) (True, szix) mono + where + mono curIx (prevMono, prevIx) = + let isMono = prevMono && prevIx > curIx + in isMono `seq` (isMono, curIx) prop_IterMonotonicBackwardsM :: Index ix => proxy ix -> Int -> Sz ix -> Property -prop_IterMonotonicBackwardsM _ thresh (Sz sz) = - totalElem sz < thresh ==> fst $ runIdentity $ - iterM (liftIndex pred sz) zeroIndex (pureIndex (-1)) (>=) (True, sz) $ \ curIx (prevMono, prevIx) -> - let isMono = prevMono && prevIx > curIx in return $ isMono `seq` (isMono, curIx) +prop_IterMonotonicBackwardsM _ thresh sz@(Sz szix) = + totalElem sz < thresh ==> fst $ + runIdentity $ iterM (liftIndex pred szix) zeroIndex (pureIndex (-1)) (>=) (True, szix) mono + where + mono curIx (prevMono, prevIx) = + let isMono = prevMono && prevIx > curIx + in return $ isMono `seq` (isMono, curIx) prop_LiftLift2 :: Index ix => proxy ix -> ix -> Int -> Bool prop_LiftLift2 _ ix delta = liftIndex2 (+) ix (liftIndex (+delta) zeroIndex) == liftIndex (+delta) ix -instance Show (Ix1 -> Double) where - show _ = "Index Func: Ix1 -> Double" - - -prop_BorderRepairSafe :: Index ix => proxy ix -> Border ix -> Sz ix -> ix -> Property -prop_BorderRepairSafe _ border@(Fill defIx) (Sz sz) ix = +prop_BorderRepairSafe :: Index ix => proxy ix -> Border ix -> SzNE ix -> ix -> Property +prop_BorderRepairSafe _ border@(Fill defIx) (SzNE sz) ix = not (isSafeIndex sz ix) ==> handleBorderIndex border sz id ix == defIx -prop_BorderRepairSafe _ border (Sz sz) ix = +prop_BorderRepairSafe _ border (SzNE sz) ix = not (isSafeIndex sz ix) ==> isSafeIndex sz (handleBorderIndex border sz id ix) prop_GetDropInsert :: Index ix => proxy ix -> NonNegative Int -> ix -> Property prop_GetDropInsert _ (NonNegative d) ix = expected === do - i <- getDim ix dim - ixL <- dropDim ix dim - insertDim ixL dim i - where expected = if d >= 1 && dim <= dimensions ix then Just ix else Nothing + i <- getDimM ix dim + ixL <- dropDimM ix dim + insertDimM ixL dim i + where expected = if d >= 1 && dim <= dimensions (Just ix) then Just ix else Nothing dim = Dim d prop_PullOutInsert :: Index ix => proxy ix -> NonNegative Int -> ix -> Property prop_PullOutInsert _ (NonNegative d) ix = expected === do - (i, ixL) <- pullOutDim ix dim - insertDim ixL dim i - where expected = if d >= 1 && dim <= dimensions ix then Just ix else Nothing + (i, ixL) <- pullOutDimM ix dim + insertDimM ixL dim i + where expected = if d >= 1 && dim <= dimensions (Just ix) then Just ix else Nothing dim = Dim d prop_UnconsGetDrop :: (Index (Lower ix), Index ix) => proxy ix -> ix -> Property prop_UnconsGetDrop _ ix = Just (unconsDim ix) === do - i <- getDim ix (dimensions ix) - ixL <- dropDim ix (dimensions ix) + i <- getDimM ix (dimensions (Just ix)) + ixL <- dropDimM ix (dimensions (Just ix)) return (i, ixL) prop_UnsnocGetDrop :: (Index (Lower ix), Index ix) => proxy ix -> ix -> Property prop_UnsnocGetDrop _ ix = Just (unsnocDim ix) === do - i <- getDim ix 1 - ixL <- dropDim ix 1 + i <- getDimM ix 1 + ixL <- dropDimM ix 1 return (ixL, i) prop_SetAll :: Index ix => proxy ix -> ix -> Int -> Bool prop_SetAll _ ix i = - foldM (\cix d -> setDim cix d i) ix ([1 .. dimensions ix] :: [Dim]) == + foldM (\cix d -> setDimM cix d i) ix ([1 .. dimensions (Just ix)] :: [Dim]) == Just (pureIndex i) prop_SetGet :: Index ix => proxy ix -> ix -> DimIx ix -> Int -> Bool -prop_SetGet _ ix (DimIx dim) n = Just n == (setDim ix dim n >>= (`getDim` dim)) +prop_SetGet _ ix (DimIx dim) n = Just n == (setDimM ix dim n >>= (`getDimM` dim)) -prop_BorderIx1 :: Positive Int -> Border Double -> (Ix1 -> Double) -> Sz Ix1 -> Ix1 -> Bool -prop_BorderIx1 (Positive period) border getVal (Sz sz) ix = +prop_BorderIx1 :: Positive Int -> Border Char -> Fun Ix1 Char -> SzNE Ix1 -> Ix1 -> Bool +prop_BorderIx1 (Positive period) border getVal (SzNE sz) ix = if isSafeIndex sz ix - then getVal ix == val + then apply getVal ix == val else case border of Fill defVal -> defVal == val Wrap -> @@ -287,49 +266,49 @@ handleBorderIndex border sz - getVal - (liftIndex2 (+) (liftIndex (* period) sz) ix) + (apply getVal) + (liftIndex2 (+) (liftIndex (* period) (unSz sz)) ix) Edge -> if ix < 0 - then val == getVal (liftIndex (max 0) ix) + then val == apply getVal (liftIndex (max 0) ix) else val == - getVal (liftIndex2 min (liftIndex (subtract 1) sz) ix) + apply getVal (liftIndex2 min (liftIndex (subtract 1) (unSz sz)) ix) Reflect -> val == handleBorderIndex border sz - getVal - (liftIndex2 (+) (liftIndex (* (2 * signum ix * period)) sz) ix) + (apply getVal) + (liftIndex2 (+) (liftIndex (* (2 * signum ix * period)) (unSz sz)) ix) Continue -> val == handleBorderIndex Reflect sz - getVal + (apply getVal) (if ix < 0 then ix - 1 else ix + 1) where - val = handleBorderIndex border sz getVal ix + val = handleBorderIndex border sz (apply getVal) ix -specDimN :: (Index ix, Ord ix, CoArbitrary ix, Arbitrary ix) => proxy ix -> Spec +specDimN :: (Index ix, Arbitrary ix) => proxy ix -> Spec specDimN proxy = do describe "Safety" $ do it "isSafeIndex" $ property $ prop_IsSafeIx proxy it "RepairSafeIx" $ property $ prop_RepairSafeIx proxy - describe "Lifting" $ do + describe "Lifting" $ it "Lift/Lift2" $ property $ prop_LiftLift2 proxy describe "Linear" $ do it "ToFromLinearIndex" $ property $ prop_ToFromLinearIndex proxy it "FromToLinearIndex" $ property $ prop_FromToLinearIndex proxy describe "Iterator" $ do - it "CountElements" $ property $ prop_CountElements proxy (2000000) - it "Monotonic" $ property $ prop_IterMonotonic proxy (2000000) - it "MonotonicBackwards" $ property $ prop_IterMonotonicBackwards proxy (2000000) - it "MonotonicM" $ property $ prop_IterMonotonicM proxy (2000000) - it "MonotonicBackwardsM" $ property $ prop_IterMonotonicBackwardsM proxy (2000000) - describe "Border" $ do + it "CountElements" $ property $ prop_CountElements proxy 2000000 + it "Monotonic" $ property $ prop_IterMonotonic proxy 2000000 + it "MonotonicBackwards" $ property $ prop_IterMonotonicBackwards proxy 2000000 + it "MonotonicM" $ property $ prop_IterMonotonicM proxy 2000000 + it "MonotonicBackwardsM" $ property $ prop_IterMonotonicBackwardsM proxy 2000000 + describe "Border" $ it "BorderRepairSafe" $ property $ prop_BorderRepairSafe proxy describe "SetGetDrop" $ do it "SetAll" $ property $ prop_SetAll proxy @@ -338,22 +317,67 @@ it "PullOutInsert" $ property $ prop_PullOutInsert proxy specDim2AndUp - :: (Index ix, Index (Lower ix), Ord ix, CoArbitrary ix, Arbitrary ix) + :: (Index ix, Index (Lower ix), Arbitrary ix) => proxy ix -> Spec -specDim2AndUp proxy = do +specDim2AndUp proxy = describe "Higher/Lower" $ do it "UnconsCons" $ property $ prop_UnconsCons proxy it "UnsnocSnoc" $ property $ prop_UnsnocSnoc proxy it "UnconsGetDrop" $ property $ prop_UnconsGetDrop proxy it "UnsnocGetDrop" $ property $ prop_UnsnocGetDrop proxy +ixToList :: Index ix => ix -> [Int] +ixToList = reverse . foldlIndex (flip (:)) [] +prop_BinaryNumIx :: + (Num ix, Index ix) => (forall n . Num n => n -> n -> n) -> ix -> ix -> Property +prop_BinaryNumIx f ix1 ix2 = zipWith f (ixToList ix1) (ixToList ix2) === ixToList (f ix1 ix2) + +prop_UnaryNumIx :: + (Num ix, Index ix) => (forall n . Num n => n -> n) -> ix -> Property +prop_UnaryNumIx f ix = map f (ixToList ix) === ixToList (f ix) + +prop_BinaryNumSz :: + (Num ix, Index ix) => (forall n . Num n => n -> n -> n) -> Sz ix -> Sz ix -> Property +prop_BinaryNumSz f sz1 sz2 = + zipWith f' (ixToList (unSz sz1)) (ixToList (unSz sz2)) === ixToList (unSz (f sz1 sz2)) + where + f' x y = max 0 (f x y) + +prop_UnaryNumSz :: + (Num ix, Index ix) => (forall n . Num n => n -> n) -> Sz ix -> Property +prop_UnaryNumSz f sz = map f' (ixToList (unSz sz)) === ixToList (unSz (f sz)) + where + f' = max 0 . f + +specClasses :: forall ix . (Typeable ix, Num ix, Index ix, Arbitrary ix) => Spec +specClasses = do + describe ("Num (" ++ showsTypeRep (typeRep (Proxy :: Proxy ix)) ")") $ do + it "(+)" $ property $ prop_BinaryNumIx @ix (+) + it "(-)" $ property $ prop_BinaryNumIx @ix (-) + it "(*)" $ property $ prop_BinaryNumIx @ix (*) + it "negate" $ property $ prop_UnaryNumIx @ix negate + it "abs" $ property $ prop_UnaryNumIx @ix abs + it "signum" $ property $ prop_UnaryNumIx @ix signum + it "fromInteger" $ property $ \ (i :: Int) -> + (fromIntegral i :: ix) === liftIndex (const i) zeroIndex + describe ("Num (" ++ showsTypeRep (typeRep (Proxy :: Proxy (Sz ix))) ")") $ do + it "(+)" $ property $ prop_BinaryNumSz @ix (+) + it "(-)" $ property $ prop_BinaryNumSz @ix (-) + it "(*)" $ property $ prop_BinaryNumSz @ix (*) + it "negate" $ property $ prop_UnaryNumSz @ix negate + it "abs" $ property $ prop_UnaryNumSz @ix abs + it "signum" $ property $ prop_UnaryNumSz @ix signum + it "fromInteger" $ property $ \ (i :: Int) -> + (fromIntegral i :: ix) === liftIndex (const i) zeroIndex + it "fromIx" $ property $ \ (ix :: ix) -> unSz (Sz ix) == liftIndex (max 0) ix + spec :: Spec spec = do describe "Tuple based indices" $ do describe "Ix1T" $ do specDimN (Nothing :: Maybe Ix1T) - it "BorderIndex" $ property $ prop_BorderIx1 + it "prop_BorderIx1" $ property prop_BorderIx1 describe "Ix2T" $ do specDimN (Nothing :: Maybe Ix2T) specDim2AndUp (Nothing :: Maybe Ix2T) @@ -370,9 +394,9 @@ describe "Ix2" $ do -- These can be used to quickly debug monotonicity it "Monotonic'" $ - property $ prop_IterMonotonic' (Nothing :: Maybe Ix2) (20000) + property $ prop_IterMonotonic' (Nothing :: Maybe Ix2) 20000 it "MonotonicBackwards'" $ - property $ prop_IterMonotonicBackwards' (Nothing :: Maybe Ix2) (20000) + property $ prop_IterMonotonicBackwards' (Nothing :: Maybe Ix2) 20000 specDimN (Nothing :: Maybe Ix2) specDim2AndUp (Nothing :: Maybe Ix2) describe "Ix3" $ do @@ -384,3 +408,8 @@ describe "Ix5" $ do specDimN (Nothing :: Maybe Ix5) specDim2AndUp (Nothing :: Maybe Ix5) + specClasses @Ix1 + specClasses @Ix2 + specClasses @Ix3 + specClasses @Ix4 + specClasses @Ix5 diff --git a/tests/Data/Massiv/Core/SchedulerSpec.hs b/tests/Data/Massiv/Core/SchedulerSpec.hs --- a/tests/Data/Massiv/Core/SchedulerSpec.hs +++ b/tests/Data/Massiv/Core/SchedulerSpec.hs @@ -1,17 +1,11 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} module Data.Massiv.Core.SchedulerSpec (spec) where -import Control.Concurrent -import Control.Exception.Base (ArithException (DivideByZero), - AsyncException (ThreadKilled)) -import Data.Massiv.Core.Scheduler -import Data.Massiv.CoreArbitrary as A -import Prelude as P -import Test.Hspec -import Test.QuickCheck -import Test.QuickCheck.Monadic +import Control.Exception.Base (ArithException(DivideByZero)) +import Control.Scheduler +import Data.Massiv.CoreArbitrary as A +import Prelude as P -- | Ensure proper exception handling. @@ -28,8 +22,8 @@ (setComp (ParOn caps) arr)) -- | Ensure proper exception handling in nested parallel computation -_prop_CatchNested :: ArrIx D Ix1 (ArrIxP D Ix1 Int) -> [Int] -> Property -_prop_CatchNested (ArrIx arr ix) caps = +prop_CatchNested :: ArrIx D Ix1 (ArrIxP D Ix1 Int) -> [Int] -> Property +prop_CatchNested (ArrIx arr ix) caps = assertException (== DivideByZero) (computeAs U $ @@ -46,35 +40,18 @@ else iarr) (setComp (ParOn caps) arr)) --- | Make sure there is no deadlock if all workers get killed -prop_AllWorkersDied :: [Int] -> (Int, [Int]) -> Property -prop_AllWorkersDied wIds (hId, ids) = - assertExceptionIO - (== ThreadKilled) - (withScheduler_ [] $ \scheduler1 -> - scheduleWork - scheduler1 - (withScheduler_ wIds $ \scheduler -> - P.mapM_ - (\_ -> scheduleWork scheduler (myThreadId >>= killThread)) - (hId : ids))) - -- | Check weather all jobs have been completed and returned order is correct -prop_SchedulerAllJobsProcessed :: [Int] -> OrderedList Int -> Property -prop_SchedulerAllJobsProcessed wIds (Ordered jobs) = - monadicIO $ do - res <- (run $ withScheduler' wIds $ \scheduler -> - P.mapM_ (scheduleWork scheduler . return) jobs) - return (res === jobs) +prop_SchedulerAllJobsProcessed :: Comp -> OrderedList Int -> Property +prop_SchedulerAllJobsProcessed comp (Ordered jobs) = + monadicIO + ((=== jobs) <$> + run (withScheduler comp $ \scheduler -> P.mapM_ (scheduleWork scheduler . return) jobs)) spec :: Spec -spec = do +spec = describe "Exceptions" $ do it "CatchDivideByZero" $ property prop_CatchDivideByZero - it "CatchNested" $ do - pendingWith "Behaves weirdly with GHC 7.10 and whenever executed with --coverage" - --property prop_CatchNested - it "AllWorkersDied" $ property prop_AllWorkersDied + it "CatchNested" $ property prop_CatchNested it "SchedulerAllJobsProcessed" $ property prop_SchedulerAllJobsProcessed diff --git a/tests/Data/Massiv/CoreArbitrary.hs b/tests/Data/Massiv/CoreArbitrary.hs --- a/tests/Data/Massiv/CoreArbitrary.hs +++ b/tests/Data/Massiv/CoreArbitrary.hs @@ -1,11 +1,12 @@ -{-# LANGUAGE ExplicitNamespaces #-} -{-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} +{-# OPTIONS_GHC -Wno-redundant-constraints #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE StandaloneDeriving #-} -{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE UndecidableInstances #-} module Data.Massiv.CoreArbitrary ( Arr(..) , ArrTiny(..) @@ -14,24 +15,41 @@ , ArrP(..) , ArrS(..) , ArrIxP(..) - , Sz(..) , SzIx(..) - , SzZ(..) + , SzNE(..) , DimIx(..) + , toIx , assertException , assertSomeException , assertExceptionIO , assertSomeExceptionIO - , module Data.Massiv.Array + , toStringException + , Semigroup((<>)) + , applyFun2Compat + , module X ) where -import Control.DeepSeq (NFData, deepseq) -import Control.Exception (Exception, SomeException, catch) -import Data.Massiv.Array hiding (type Sz, pattern Sz) -import Data.Massiv.Core.IndexSpec hiding (spec) -import Data.Typeable -import Test.QuickCheck -import Test.QuickCheck.Monadic +import Control.DeepSeq (NFData, deepseq) +import UnliftIO.Exception (Exception, SomeException, catch, catchAny) +import Data.Foldable as F +import Data.Massiv.Array as X +import Data.Massiv.Core.IndexSpec hiding (spec) +import Data.Typeable +import Test.Hspec as X +import Test.QuickCheck as X hiding (resize) +import Test.QuickCheck.Function as X +import Test.QuickCheck.Monadic as X +#if !MIN_VERSION_base(4,11,0) +import Data.Semigroup +#endif +applyFun2Compat :: Fun (a, b) c -> (a -> b -> c) +#if MIN_VERSION_QuickCheck(2,10,0) +applyFun2Compat = applyFun2 +#else +applyFun2Compat (Fun _ f) a b = f (a, b) +instance Function Word where + function = functionMap fromIntegral fromInteger +#endif -- | Arbitrary non-empty array. Computation strategy can be either `Seq` or `Par`. newtype Arr r ix e = Arr {unArr :: Array r ix e} @@ -72,7 +90,7 @@ instance (CoArbitrary ix, Arbitrary ix, Typeable e, Construct r ix e, Arbitrary e) => Arbitrary (Array r ix e) where arbitrary = do - SzZ sz <- arbitrary + sz <- arbitrary func <- arbitrary comp <- oneof [pure Seq, pure Par] return $ makeArray comp sz func @@ -81,24 +99,24 @@ instance (CoArbitrary ix, Arbitrary ix, Typeable e, Construct r ix e, Arbitrary e) => Arbitrary (ArrTiny r ix e) where arbitrary = do - SzZ sz <- arbitrary + sz <- unSz <$> arbitrary func <- arbitrary comp <- oneof [pure Seq, pure Par] - return $ ArrTiny $ makeArray comp (liftIndex (`mod` 10) sz) func + return $ ArrTiny $ makeArray comp (Sz (liftIndex (`mod` 10) sz)) func -- | Arbitrary small and possibly empty array. Computation strategy can be either `Seq` or `Par`. instance (CoArbitrary ix, Arbitrary ix, Typeable e, Construct r ix e, Arbitrary e) => Arbitrary (ArrTiny1 r ix e) where arbitrary = do - SzZ sz <- arbitrary + sz <- unSz <$> arbitrary func <- arbitrary comp <- oneof [pure Seq, pure Par] - return $ ArrTiny1 $ makeArray comp (liftIndex (succ . (`mod` 10)) sz) func + return $ ArrTiny1 $ makeArray comp (Sz (liftIndex (succ . (`mod` 10)) sz)) func instance (CoArbitrary ix, Arbitrary ix, Typeable e, Construct r ix e, Arbitrary e) => Arbitrary (Arr r ix e) where arbitrary = do - Sz sz <- arbitrary + SzNE sz <- arbitrary func <- arbitrary comp <- oneof [pure Seq, pure Par] return $ Arr $ makeArray comp sz func @@ -108,9 +126,8 @@ instance (CoArbitrary ix, Arbitrary ix, Typeable e, Construct r ix e, Arbitrary e) => Arbitrary (ArrS r ix e) where arbitrary = do - Sz sz <- arbitrary - func <- arbitrary - return $ ArrS $ makeArray Seq sz func + SzNE sz <- arbitrary + ArrS . makeArray Seq sz <$> arbitrary instance (CoArbitrary ix, Arbitrary ix, Typeable e, Construct r ix e, Arbitrary e) => Arbitrary (ArrP r ix e) where @@ -121,7 +138,7 @@ instance (CoArbitrary ix, Arbitrary ix, Typeable e, Construct r ix e, Arbitrary e) => Arbitrary (ArrIx r ix e) where arbitrary = do - SzIx (Sz sz) ix <- arbitrary + SzIx sz ix <- arbitrary func <- arbitrary comp <- arbitrary return $ ArrIx (makeArray comp sz func) ix @@ -129,7 +146,7 @@ instance (CoArbitrary ix, Arbitrary ix, Typeable e, Construct r ix e, Arbitrary e) => Arbitrary (ArrIxS r ix e) where arbitrary = do - SzIx (Sz sz) ix <- arbitrary + SzIx sz ix <- arbitrary func <- arbitrary return $ ArrIxS (makeArray Seq sz func) ix @@ -142,15 +159,20 @@ return $ ArrIxP (setComp Par arrIx) ix +toIx :: (Dimensions ix' ~ Dimensions ix, Index ix', Index ix) => ix -> ix' +toIx ix = F.foldl' setEachIndex zeroIndex [1.. dimensions (Sz ix)] + where setEachIndex ix' d = setDim' ix' d (getDim' ix d) + + assertException :: (NFData a, Exception exc) => (exc -> Bool) -- ^ Return True if that is the exception that was expected -> a -- ^ Value that should throw an exception, when fully evaluated -> Property -assertException isExc action = assertExceptionIO isExc (return action) +assertException isExc = assertExceptionIO isExc . pure assertSomeException :: NFData a => a -> Property -assertSomeException = assertSomeExceptionIO . return +assertSomeException = assertSomeExceptionIO . pure assertExceptionIO :: (NFData a, Exception exc) => @@ -159,13 +181,23 @@ -> Property assertExceptionIO isExc action = monadicIO $ do - hasFailed <- + assert =<< run (catch (do res <- action - res `deepseq` return False) $ \exc -> - show exc `deepseq` return (isExc exc)) - assert hasFailed + res `deepseq` return False) + (\exc -> displayException exc `deepseq` return (isExc exc))) assertSomeExceptionIO :: NFData a => IO a -> Property -assertSomeExceptionIO = assertExceptionIO (\exc -> const True (exc :: SomeException)) +assertSomeExceptionIO action = + monadicIO $ do + assert =<< + run + (catchAny + (do res <- action + res `deepseq` return False) + (\exc -> displayException exc `deepseq` return True)) + + +toStringException :: Either SomeException a -> Either String a +toStringException = either (Left . displayException) Right diff --git a/tests/Main.hs b/tests/Main.hs new file mode 100644 --- /dev/null +++ b/tests/Main.hs @@ -0,0 +1,10 @@ +module Main where + +import System.IO (BufferMode(LineBuffering), hSetBuffering, stdout) +import Test.Hspec +import Spec + +main :: IO () +main = do + hSetBuffering stdout LineBuffering + hspec spec diff --git a/tests/Spec.hs b/tests/Spec.hs --- a/tests/Spec.hs +++ b/tests/Spec.hs @@ -1,40 +1,1 @@ -module Main where - -import Data.Massiv.Array.Delayed.WindowedSpec as Windowed -import Data.Massiv.Array.DelayedSpec as Delayed -import Data.Massiv.Array.Manifest.VectorSpec as Vector -import Data.Massiv.Array.ManifestSpec as Manifest -import Data.Massiv.Array.MutableSpec as Mutable -import Data.Massiv.Array.Ops.ConstructSpec as Construct -import Data.Massiv.Array.Ops.FoldSpec as Fold -import Data.Massiv.Array.Ops.MapSpec as Map -import Data.Massiv.Array.Ops.SliceSpec as Slice -import Data.Massiv.Array.Ops.TransformSpec as Transform -import Data.Massiv.Array.StencilSpec as Stencil -import Data.Massiv.Core.IndexSpec as Index -import Data.Massiv.Core.SchedulerSpec as Scheduler -import System.IO (BufferMode (LineBuffering), - hSetBuffering, stdout) -import Test.Hspec - - --- | Main entry point. Returns ExitFailure if a test fails. -main :: IO () -main = do - hSetBuffering stdout LineBuffering - hspec $ do - describe "Core" $ do - Scheduler.spec - Index.spec - describe "Ops" $ do - Construct.spec - Fold.spec - Slice.spec - Transform.spec - Map.spec - describe "Delayed" $ Delayed.spec - describe "Windowed" $ Windowed.spec - describe "Manifest" $ Manifest.spec - describe "Mutable" $ Mutable.spec - describe "Stencil" $ Stencil.spec - describe "Vector" $ Vector.spec +{-# OPTIONS_GHC -F -pgmF hspec-discover -optF --no-main #-} diff --git a/tests/doctests.hs b/tests/doctests.hs new file mode 100644 --- /dev/null +++ b/tests/doctests.hs @@ -0,0 +1,12 @@ +module Main where + +import Build_doctests (flags, pkgs, module_sources) +import Data.Foldable (traverse_) +import Test.DocTest (doctest) + +main :: IO () +main = do + traverse_ putStrLn args + doctest args + where + args = flags ++ pkgs ++ module_sources