massiv 0.5.7.0 → 0.5.8.0

raw patch · 12 files changed

+210/−147 lines, 12 filesPVP: major bump suggested

API removals or changes: PVP suggests a major version bump

API changes (from Hackage documentation)

Files

• CHANGELOG.md +4/−0
• massiv.cabal +1/−1
• src/Data/Massiv/Array/Delayed/Push.hs +102/−84
• src/Data/Massiv/Array/Mutable.hs +5/−5
• src/Data/Massiv/Array/Ops/Construct.hs +11/−9
• src/Data/Massiv/Array/Ops/Transform.hs +23/−24
• src/Data/Massiv/Array/Unsafe.hs +0/−1
• src/Data/Massiv/Core.hs +4/−2
• src/Data/Massiv/Core/Common.hs +35/−3
• src/Data/Massiv/Core/Index.hs +13/−4
• src/Data/Massiv/Core/Index/Internal.hs +8/−12
• src/Data/Massiv/Vector.hs +4/−2

@@ -1,3 +1,7 @@+# 0.5.8++* Improve loading of push arrays by adding `loadArrayWithSetM` and deprecating `defaultElement`.+ # 0.5.7  * Improve performance of `><.` and `><!` while making their constraints a bit more relaxed.

@@ -1,5 +1,5 @@ name:                massiv-version:             0.5.7.0+version:             0.5.8.0 synopsis:            Massiv (Массив) is an Array Library. description:         Multi-dimensional Arrays with fusion, stencils and parallel computation. homepage:            https://github.com/lehins/massiv

@@ -1,8 +1,8 @@-{-# LANGUAGE LambdaCase #-} {-# LANGUAGE BangPatterns #-} {-# LANGUAGE CPP #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE LambdaCase #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE NamedFieldPuns #-} {-# LANGUAGE RankNTypes #-}@@ -31,10 +31,10 @@   ) where  import Control.Monad-import Data.Massiv.Core.Common-import Prelude hiding (map, zipWith) import Control.Scheduler as S (traverse_) import Data.Foldable as F+import Data.Massiv.Core.Common+import Prelude hiding (map, zipWith)  #include "massiv.h" @@ -45,21 +45,22 @@ data instance Array DL ix e = DLArray   { dlComp    :: !Comp   , dlSize    :: !(Sz ix)-  , dlDefault :: !(Maybe e)   , dlLoad    :: forall m . Monad m               => Scheduler m ()-              -> Int -- start loading at this linear index-              -> (Int -> e -> m ()) -- linear element writing action+              -> Ix1 -- start loading at this linear index+              -> (Ix1 -> e -> m ()) -- linear element writing action+              -> (Ix1 -> Sz1 -> e -> m ()) -- linear region setting action               -> m ()   }  instance Index ix => Construct DL ix e where   setComp c arr = arr {dlComp = c}   {-# INLINE setComp #-}-  makeArrayLinear comp sz f = DLArray comp sz Nothing load+  makeArrayLinear comp sz f = DLArray comp sz load     where-      load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()-      load scheduler startAt dlWrite =+      load :: Monad m =>+        Scheduler m () -> Ix1 -> (Ix1 -> e -> m ()) -> (Ix1 -> Sz1 -> e -> m ()) -> m ()+      load scheduler startAt dlWrite _ =         splitLinearlyWithStartAtM_ scheduler startAt (totalElem sz) (pure . f) dlWrite       {-# INLINE load #-}   {-# INLINE makeArrayLinear #-}@@ -73,9 +74,7 @@   {-# INLINE (<>) #-}  instance Monoid (Array DL Ix1 e) where-  mempty =-    DLArray-      {dlComp = mempty, dlSize = Sz zeroIndex, dlDefault = Nothing, dlLoad = \_ _ _ -> pure ()}+  mempty = DLArray {dlComp = mempty, dlSize = zeroSz, dlLoad = \_ _ _ _ -> pure ()}   {-# INLINE mempty #-}   mappend = mappendDL   {-# INLINE mappend #-}@@ -87,18 +86,15 @@  mconcatDL :: forall e . [Array DL Ix1 e] -> Array DL Ix1 e mconcatDL !arrs =-  DLArray {dlComp = foldMap getComp arrs, dlSize = SafeSz k, dlDefault = Nothing, dlLoad = load}+  DLArray {dlComp = foldMap getComp arrs, dlSize = SafeSz k, dlLoad = load}   where     !k = F.foldl' (+) 0 (unSz . size <$> arrs)-    load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()-    load scheduler startAt dlWrite =-      let loadArr !startAtCur DLArray {dlSize = SafeSz kCur, dlDefault, dlLoad} = do+    load :: Monad m =>+      Scheduler m () -> Ix1 -> (Ix1 -> e -> m ()) -> (Ix1 -> Sz1 -> e -> m ()) -> m ()+    load scheduler startAt dlWrite dlSet =+      let loadArr !startAtCur DLArray {dlSize = SafeSz kCur, dlLoad} = do             let !endAtCur = startAtCur + kCur-            scheduleWork_ scheduler $ do-              S.traverse_-                (\def -> loopM_ startAtCur (< endAtCur) (+ 1) (`dlWrite` def))-                dlDefault-              dlLoad scheduler startAtCur dlWrite+            scheduleWork_ scheduler $ dlLoad scheduler startAtCur dlWrite dlSet             pure endAtCur           {-# INLINE loadArr #-}        in foldM_ loadArr startAt arrs@@ -107,20 +103,16 @@   mappendDL :: forall e . Array DL Ix1 e -> Array DL Ix1 e -> Array DL Ix1 e-mappendDL (DLArray c1 sz1 mDef1 load1) (DLArray c2 sz2 mDef2 load2) =-  DLArray {dlComp = c1 <> c2, dlSize = SafeSz (k1 + k2), dlDefault = Nothing, dlLoad = load}+mappendDL (DLArray c1 sz1 load1) (DLArray c2 sz2 load2) =+  DLArray {dlComp = c1 <> c2, dlSize = SafeSz (k1 + k2), dlLoad = load}   where     !k1 = unSz sz1     !k2 = unSz sz2-    load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()-    load scheduler startAt dlWrite = do-      scheduleWork_ scheduler $ do-        S.traverse_ (\def1 -> loopM_ startAt (< k1) (+ 1) (`dlWrite` def1)) mDef1-        load1 scheduler startAt dlWrite-      scheduleWork_ scheduler $ do-        let startAt2 = startAt + k1-        S.traverse_ (\def2 -> loopM_ startAt2 (< startAt2 + k2) (+ 1) (`dlWrite` def2)) mDef2-        load2 scheduler startAt2 dlWrite+    load :: Monad n =>+      Scheduler n () -> Ix1 -> (Ix1 -> e -> n ()) -> (Ix1 -> Sz1 -> e -> n ()) -> n ()+    load scheduler !startAt dlWrite dlSet = do+      scheduleWork_ scheduler $ load1 scheduler startAt dlWrite dlSet+      scheduleWork_ scheduler $ load2 scheduler (startAt + k1) dlWrite dlSet     {-# INLINE load #-} {-# INLINE mappendDL #-} @@ -133,24 +125,17 @@   => Array DL ix e   -> Array DL ix e   -> m (Array DL ix e)-appendOuterM (DLArray c1 sz1 mDef1 load1) (DLArray c2 sz2 mDef2 load2) = do+appendOuterM (DLArray c1 sz1 load1) (DLArray c2 sz2 load2) = do   let (!i1, !szl1) = unconsSz sz1       (!i2, !szl2) = unconsSz sz2   unless (szl1 == szl2) $ throwM $ SizeMismatchException sz1 sz2   pure $-    DLArray {dlComp = c1 <> c2, dlSize = consSz (i1 + i2) szl1, dlDefault = Nothing, dlLoad = load}+    DLArray {dlComp = c1 <> c2, dlSize = consSz (i1 + i2) szl1, dlLoad = load}   where-    !k1 = totalElem sz1-    !k2 = totalElem sz2-    load :: Monad n => Scheduler n () -> Int -> (Int -> e -> n ()) -> n ()-    load scheduler !startAt dlWrite = do-      scheduleWork_ scheduler $ do-        S.traverse_ (\def1 -> loopM_ startAt (< k1) (+ 1) (`dlWrite` def1)) mDef1-        load1 scheduler startAt dlWrite-      scheduleWork_ scheduler $ do-        let !startAt2 = startAt + k1-        S.traverse_ (\def2 -> loopM_ startAt2 (< startAt2 + k2) (+ 1) (`dlWrite` def2)) mDef2-        load2 scheduler startAt2 dlWrite+    load :: Monad n => Scheduler n () -> Ix1 -> (Ix1 -> e -> n ()) -> (Ix1 -> Sz1 -> e -> n ()) -> n ()+    load scheduler !startAt dlWrite dlSet = do+      scheduleWork_ scheduler $ load1 scheduler startAt dlWrite dlSet+      scheduleWork_ scheduler $ load2 scheduler (startAt + totalElem sz1) dlWrite dlSet     {-# INLINE load #-} {-# INLINE appendOuterM #-} @@ -164,7 +149,7 @@   -> m (Array DL ix e) concatOuterM =   \case-    [] -> pure empty+    []     -> pure empty     (x:xs) -> F.foldlM appendOuterM x xs {-# INLINE concatOuterM #-} @@ -179,19 +164,20 @@   -- ^ Size of the resulting array   -> e   -- ^ Default value to use for all cells that might have been ommitted by the writing function-  -> (forall m. Monad m =>-                  (ix -> e -> m Bool) -> m ())+  -> (forall m. Monad m => (ix -> e -> m Bool) -> m ())   -- ^ Writing function that described which elements to write into the target array.   -> Array DL ix e-makeLoadArrayS sz defVal writer = DLArray Seq sz (Just defVal) load+makeLoadArrayS sz defVal writer = DLArray Seq sz load   where-    load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()-    load _scheduler !startAt uWrite =+    load :: Monad m =>+      Scheduler m () -> Ix1 -> (Ix1 -> e -> m ()) -> (Ix1 -> Sz1 -> e -> m ()) -> m ()+    load _scheduler !startAt uWrite uSet = do+      uSet startAt (toLinearSz sz) defVal       let safeWrite !ix !e             | isSafeIndex sz ix = uWrite (startAt + toLinearIndex sz ix) e >> pure True             | otherwise = pure False           {-# INLINE safeWrite #-}-       in writer safeWrite+      writer safeWrite     {-# INLINE load #-} {-# INLINE makeLoadArrayS #-} @@ -209,21 +195,22 @@   -- ^ Size of the resulting array   -> e   -- ^ Default value to use for all cells that might have been ommitted by the writing function-  -> (forall m. Monad m =>-                  Scheduler m () -> (ix -> e -> m Bool) -> m ())+  -> (forall m. Monad m => Scheduler m () -> (ix -> e -> m Bool) -> m ())   -- ^ Writing function that described which elements to write into the target array. It   -- accepts a scheduler, that can be used for parallelization, as well as a safe element   -- writing function.   -> Array DL ix e-makeLoadArray comp sz defVal writer = DLArray comp sz (Just defVal) load+makeLoadArray comp sz defVal writer = DLArray comp sz load   where-    load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()-    load scheduler !startAt uWrite =+    load :: Monad m =>+      Scheduler m () -> Ix1 -> (Ix1 -> e -> m ()) -> (Ix1 -> Sz1 -> e -> m ()) -> m ()+    load scheduler !startAt uWrite uSet = do+      uSet startAt (toLinearSz sz) defVal       let safeWrite !ix !e-            | isSafeIndex sz ix = uWrite (startAt + toLinearIndex sz ix) e >> pure True+            | isSafeIndex sz ix = True <$ uWrite (startAt + toLinearIndex sz ix) e             | otherwise = pure False           {-# INLINE safeWrite #-}-       in writer scheduler safeWrite+      writer scheduler safeWrite     {-# INLINE load #-} {-# INLINE makeLoadArray #-} @@ -234,7 +221,8 @@ -- -- @since 0.3.1 unsafeMakeLoadArray ::-     Comp+     forall ix e. Index ix+  => Comp   -- ^ Computation strategy to use. Directly affects the scheduler that gets created for   -- the loading function.   -> Sz ix@@ -242,7 +230,7 @@   -> Maybe e   -- ^ An element to use for initialization of the mutable array that will be created in   -- the future-  -> (forall m. Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m ())+  -> (forall m. Monad m => Scheduler m () -> Ix1 -> (Ix1 -> e -> m ()) -> m ())   -- ^ This function accepts:   --   -- * A scheduler that can be used for parallelization of loading@@ -252,7 +240,14 @@   --   -- * Linear element writing function   -> Array DL ix e-unsafeMakeLoadArray = DLArray+unsafeMakeLoadArray comp sz mDefVal writer = DLArray comp sz load+  where+    load :: Monad m =>+      Scheduler m () -> Ix1 -> (Ix1 -> e -> m ()) -> (Ix1 -> Sz1 -> e -> m ()) -> m ()+    load scheduler startAt uWrite uSet = do+      S.traverse_ (uSet startAt (toLinearSz sz)) mDefVal+      writer scheduler startAt uWrite+    {-# INLINE load #-} {-# INLINE unsafeMakeLoadArray #-}  -- | Same as `unsafeMakeLoadArray`, except will ensure that starting index is correctly@@ -260,17 +255,19 @@ -- -- @since 0.5.2 unsafeMakeLoadArrayAdjusted ::-     forall ix e.-     Comp+     forall ix e. Index ix+  => Comp   -> Sz ix   -> Maybe e-  -> (forall m. Monad m =>-                  Scheduler m () -> (Int -> e -> m ()) -> m ())+  -> (forall m. Monad m => Scheduler m () -> (Ix1 -> e -> m ()) -> m ())   -> Array DL ix e-unsafeMakeLoadArrayAdjusted comp sz mDefVal writer = DLArray comp sz mDefVal load+unsafeMakeLoadArrayAdjusted comp sz mDefVal writer = DLArray comp sz load   where-    load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()-    load scheduler !startAt uWrite = writer scheduler (\i -> uWrite (startAt + i))+    load :: Monad m =>+      Scheduler m () -> Ix1 -> (Ix1 -> e -> m ()) -> (Ix1 -> Sz1 -> e -> m ()) -> m ()+    load scheduler !startAt uWrite dlSet = do+      S.traverse_ (dlSet startAt (toLinearSz sz)) mDefVal+      writer scheduler (\i -> uWrite (startAt + i))     {-# INLINE load #-} {-# INLINE unsafeMakeLoadArrayAdjusted #-} @@ -281,10 +278,13 @@      forall r ix e. Load r ix e   => Array r ix e   -> Array DL ix e-toLoadArray arr = DLArray (getComp arr) (size arr) (defaultElement arr) load+toLoadArray arr = DLArray (getComp arr) sz load   where-    load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()-    load scheduler !startAt dlWrite = loadArrayM scheduler arr (dlWrite . (+ startAt))+    !sz = size arr+    load :: Monad m =>+      Scheduler m () -> Ix1 -> (Ix1 -> e -> m ()) -> (Ix1 -> Sz1 -> e -> m ()) -> m ()+    load scheduler !startAt dlWrite dlSet =+      loadArrayWithSetM scheduler arr (dlWrite . (+ startAt)) (\offset -> dlSet (offset + startAt))     {-# INLINE load #-} {-# INLINE[1] toLoadArray #-} {-# RULES "toLoadArray/id" toLoadArray = id #-}@@ -298,11 +298,12 @@   -> Array r ix e   -> Array DL ix e fromStrideLoad stride arr =-  DLArray (getComp arr) newsz Nothing load+  DLArray (getComp arr) newsz load   where-    newsz = strideSize stride (size arr)-    load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()-    load scheduler !startAt dlWrite =+    !newsz = strideSize stride (size arr)+    load :: Monad m =>+      Scheduler m () -> Ix1 -> (Ix1 -> e -> m ()) -> (Ix1 -> Sz1 -> e -> m ()) -> m ()+    load scheduler !startAt dlWrite _ =       loadArrayWithStrideM scheduler stride newsz arr (\ !i -> dlWrite (i + startAt))     {-# INLINE load #-} {-# INLINE fromStrideLoad #-}@@ -312,18 +313,35 @@   {-# INLINE size #-}   getComp = dlComp   {-# INLINE getComp #-}-  loadArrayM scheduler DLArray {dlLoad} = dlLoad scheduler 0-  {-# INLINE loadArrayM #-}-  defaultElement = dlDefault+  loadArrayWithSetM scheduler DLArray {dlLoad} = dlLoad scheduler 0+  {-# INLINE loadArrayWithSetM #-}+  defaultElement _ = Nothing   {-# INLINE defaultElement #-} -instance Functor (Array DL ix) where-  fmap f arr = arr {dlLoad = loadFunctor arr f, dlDefault = f <$> dlDefault arr}+instance Index ix => Functor (Array DL ix) where+  fmap f arr = arr {dlLoad = loadFunctor arr f}   {-# INLINE fmap #-}-  (<$) e arr = arr {dlLoad = \_ _ _ -> pure (), dlDefault = Just e}+  (<$) = overwriteFunctor   {-# INLINE (<$) #-} +overwriteFunctor :: forall ix a b. Index ix => a -> Array DL ix b -> Array DL ix a+overwriteFunctor e arr = arr {dlLoad = load}+  where+    load :: Scheduler m () -> Ix1 -> (Ix1 -> a -> m ()) -> (Ix1 -> Sz1 -> a -> m ()) -> m ()+    load _ !startAt _ dlSet = dlSet startAt (linearSize arr) e+    {-# INLINE load #-}+{-# INLINE overwriteFunctor #-}++ loadFunctor ::-     Monad m => Array DL ix t1 -> (t1 -> t2) -> Scheduler m () -> Int -> (Int -> t2 -> m ()) -> m ()-loadFunctor arr f scheduler startAt uWrite = dlLoad arr scheduler startAt (\ !i e -> uWrite i (f e))+     Monad m+  => Array DL ix a+  -> (a -> b)+  -> Scheduler m ()+  -> Ix1+  -> (Ix1 -> b -> m ())+  -> (Ix1 -> Sz1 -> b -> m ())+  -> m ()+loadFunctor arr f scheduler startAt uWrite uSet =+  dlLoad arr scheduler startAt (\ !i e -> uWrite i (f e)) (\o sz e -> uSet o sz (f e)) {-# INLINE loadFunctor #-}

@@ -251,10 +251,10 @@ {-# INLINE freezeS #-}  -newMaybeInitialized ::+unsafeNewUninitialized ::      (Load r' ix e, Mutable r ix e, PrimMonad m) => Array r' ix e -> m (MArray (PrimState m) r ix e)-newMaybeInitialized !arr = initializeNew (defaultElement arr) (fromMaybe zeroSz (maxSize arr))-{-# INLINE newMaybeInitialized #-}+unsafeNewUninitialized !arr = unsafeNew (fromMaybe zeroSz (maxSize arr))+{-# INLINE unsafeNewUninitialized #-}   -- | Load sequentially a pure array into the newly created mutable array.@@ -265,7 +265,7 @@   => Array r' ix e   -> m (MArray (PrimState m) r ix e) loadArrayS arr = do-  marr <- newMaybeInitialized arr+  marr <- unsafeNewUninitialized arr   unsafeLoadIntoS marr arr {-# INLINE loadArrayS #-} @@ -279,7 +279,7 @@   -> m (MArray RealWorld r ix e) loadArray arr =   liftIO $ do-    marr <- newMaybeInitialized arr+    marr <- unsafeNewUninitialized arr     unsafeLoadIntoM marr arr {-# INLINE loadArray #-} 

@@ -231,12 +231,13 @@   -> (a -> ix -> (e, a))   -> a   -> Array DL ix e-iunfoldrS_ sz f acc0 = DLArray {dlComp = Seq, dlSize = sz, dlDefault = Nothing, dlLoad = load}+iunfoldrS_ sz f acc0 = DLArray {dlComp = Seq, dlSize = sz, dlLoad = load}   where-    load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()-    load _ startAt dlWrite =+    load :: Monad m =>+      Scheduler m () -> Ix1 -> (Ix1 -> e -> m ()) -> (Ix1 -> Sz1 -> e -> m ()) -> m ()+    load _ startAt dlWrite _ =       void $-      loopM startAt (< (totalElem sz + startAt)) (+ 1) acc0 $ \ !i !acc ->+      loopM startAt (< totalElem sz + startAt) (+ 1) acc0 $ \ !i !acc ->         let (e, acc') = f acc $ fromLinearIndex sz (i - startAt)          in acc' <$ dlWrite i e     {-# INLINE load #-}@@ -261,12 +262,13 @@   -> (ix -> a -> (a, e))   -> a   -> Array DL ix e-iunfoldlS_ sz f acc0 = DLArray {dlComp = Seq, dlSize = sz, dlDefault = Nothing, dlLoad = load}+iunfoldlS_ sz f acc0 = DLArray {dlComp = Seq, dlSize = sz, dlLoad = load}   where-    load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()-    load _ startAt dlWrite =+    load :: Monad m =>+      Scheduler m () -> Ix1 -> (Ix1 -> e -> m ()) -> (Ix1 -> Sz1 -> e -> m ()) -> m ()+    load _ startAt dlWrite _ =       void $-      loopDeepM startAt (< (totalElem sz + startAt)) (+ 1) acc0 $ \ !i !acc ->+      loopDeepM startAt (< totalElem sz + startAt) (+ 1) acc0 $ \ !i !acc ->         let (acc', e) = f (fromLinearIndex sz (i - startAt)) acc          in acc' <$ dlWrite i e     {-# INLINE load #-}@@ -330,7 +332,7 @@                     then (genI, genI)                     else splitGen genI             scheduleWork_ scheduler $-              void $ loopM start (< (start + chunkLength)) (+ 1) genI0 writeRandom+              void $ loopM start (< start + chunkLength) (+ 1) genI0 writeRandom             pure genI1         when (slackStartAt < totalLength + startAt) $           scheduleWork_ scheduler $

@@ -442,8 +442,9 @@   unless (szl1 == szl2) $ throwM $ SizeMismatchException sz1 sz2   let !k1' = unSz k1   newSz <- insertSzM szl1 n (SafeSz (k1' + unSz k2))-  let load :: Monad n => Scheduler n () -> Int -> (Int -> e -> n ()) -> n ()-      load scheduler !startAt dlWrite = do+  let load :: Monad n =>+        Scheduler n () -> Ix1 -> (Ix1 -> e -> n ()) -> (Ix1 -> Sz1 -> e -> n ()) -> n ()+      load scheduler !startAt dlWrite _dlSet = do         scheduleWork scheduler $           iterM_ zeroIndex (unSz sz1) (pureIndex 1) (<) $ \ix ->             dlWrite (startAt + toLinearIndex newSz ix) (unsafeIndex arr1 ix)@@ -455,7 +456,7 @@       {-# INLINE load #-}   return $     DLArray-      {dlComp = getComp arr1 <> getComp arr2, dlSize = newSz, dlDefault = Nothing, dlLoad = load}+      {dlComp = getComp arr1 <> getComp arr2, dlSize = newSz, dlLoad = load} {-# INLINE appendM #-}  @@ -501,8 +502,9 @@         (dropWhile ((== szl) . snd) $ P.zip szs szls)       let kTotal = SafeSz $ F.foldl' (+) k ks       newSz <- insertSzM (SafeSz szl) n kTotal-      let load :: Monad n => Scheduler n () -> Int -> (Int -> e -> n ()) -> n ()-          load scheduler startAt dlWrite =+      let load :: Monad n =>+            Scheduler n () -> Ix1 -> (Ix1 -> e -> n ()) -> (Ix1 -> Sz1 -> e -> n ()) -> n ()+          load scheduler startAt dlWrite _dlSet =             let arrayLoader !kAcc (kCur, arr) = do                   scheduleWork scheduler $                     iforM_ arr $ \ix e ->@@ -513,7 +515,7 @@              in M.foldM_ arrayLoader 0 $ (k, a) : P.zip ks arrs           {-# INLINE load #-}       return $-        DLArray {dlComp = foldMap getComp arrsF, dlSize = newSz, dlDefault = Nothing, dlLoad = load}+        DLArray {dlComp = foldMap getComp arrsF, dlSize = newSz, dlLoad = load} {-# INLINE concatM #-}  @@ -593,17 +595,18 @@           len = SafeSz (F.length arrsF)       -- / make sure all arrays have the same size       M.forM_ arrsF $ \arr ->-        let sz' = size arr-         in unless (sz == sz') $ throwM (SizeMismatchException sz sz')+         unless (sz == size arr) $ throwM (SizeMismatchException sz (size arr))       newSz <- insertSzM sz dim len-      let load :: Monad n => Scheduler n () -> Int -> (Int -> e -> n ()) -> n ()-          load scheduler startAt dlWrite =+      let load :: Monad n =>+            Scheduler n () -> Ix1 -> (Ix1 -> e -> n ()) -> (Ix1 -> Sz1 -> e -> n ()) -> n ()+          load scheduler startAt dlWrite _dlSet =             let loadIndex k ix = dlWrite (toLinearIndex newSz (insertDim' ix dim k) + startAt)                 arrayLoader !k arr = (k + 1) <$ scheduleWork scheduler (imapM_ (loadIndex k) arr)+                {-# INLINE arrayLoader #-}              in M.foldM_ arrayLoader 0 arrsF           {-# INLINE load #-}       return $-        DLArray {dlComp = foldMap getComp arrs, dlSize = newSz, dlDefault = Nothing, dlLoad = load}+        DLArray {dlComp = foldMap getComp arrs, dlSize = newSz, dlLoad = load} {-# INLINE stackSlicesM #-}  -- | Specialized `stackSlicesM` to handling stacking from the outside. It is the inverse of@@ -850,15 +853,16 @@   -> Array r ix e   -> Array DL ix e downsample stride arr =-  DLArray {dlComp = getComp arr, dlSize = resultSize, dlDefault = defaultElement arr, dlLoad = load}+  DLArray {dlComp = getComp arr, dlSize = resultSize, dlLoad = load}   where     resultSize = strideSize stride (size arr)     strideIx = unStride stride     unsafeLinearWriteWithStride =       unsafeIndex arr . liftIndex2 (*) strideIx . fromLinearIndex resultSize     {-# INLINE unsafeLinearWriteWithStride #-}-    load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()-    load scheduler startAt dlWrite =+    load :: Monad m =>+      Scheduler m () -> Ix1 -> (Ix1 -> e -> m ()) -> (Ix1 -> Sz1 -> e -> m ()) -> m ()+    load scheduler startAt dlWrite _ =       splitLinearlyWithStartAtM_         scheduler         startAt@@ -908,19 +912,14 @@   DLArray     { dlComp = getComp arr     , dlSize = newsz-    , dlDefault = Just fillWith     , dlLoad = load     }   where-    load :: Monad m => Scheduler m () -> Int -> (Int -> e -> m ()) -> m ()-    load scheduler !startAt dlWrite = do-      M.forM_ (defaultElement arr) $ \prevFillWith ->-        loopM_-          startAt-          (< totalElem sz)-          (+ 1)-          (\i -> dlWrite (adjustLinearStride (i + startAt)) prevFillWith)-      loadArrayM scheduler arr (\i -> dlWrite (adjustLinearStride (i + startAt)))+    load :: Monad m =>+      Scheduler m () -> Ix1 -> (Ix1 -> e -> m ()) -> (Ix1 -> Sz1 -> e -> m ()) -> m ()+    load scheduler startAt uWrite uSet = do+      uSet startAt (toLinearSz newsz) fillWith+      loadArrayM scheduler arr (\i -> uWrite (adjustLinearStride (i + startAt)))     {-# INLINE load #-}     adjustLinearStride = toLinearIndex newsz . timesStride . fromLinearIndex sz     {-# INLINE adjustLinearStride #-}

@@ -2,7 +2,6 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-}-{-# LANGUAGE PatternSynonyms #-} -- | -- Module      : Data.Massiv.Array.Unsafe -- Copyright   : (c) Alexey Kuleshevich 2018-2020

@@ -14,7 +14,7 @@   , MMatrix   , Elt   , Construct-  , Load(R, loadArrayM, defaultElement)+  , Load(R, loadArrayM, loadArrayWithSetM, defaultElement)   , Stream(..)   , Source   , Resize@@ -31,6 +31,8 @@   , L(..)   , LN   , ListItem+  , Scheduler+  , SchedulerWS   , Comp(Seq, Par, Par', ParOn, ParN)   , WorkerStates   , initWorkerStates@@ -48,7 +50,7 @@   , PrimMonad(PrimState)   ) where -import Control.Scheduler (initWorkerStates)+import Control.Scheduler (SchedulerWS, initWorkerStates) import Data.Massiv.Core.Common import Data.Massiv.Core.Index import Data.Massiv.Core.List

@@ -54,6 +54,7 @@   , singleton   -- * Size   , elemsCount+  , linearSize   , isNotEmpty   , Sz(SafeSz)   , Size(..)@@ -268,6 +269,7 @@ class (Typeable r, Index ix) => Load r ix e where   type family R r :: *   type instance R r = r+  {-# MINIMAL getComp, size, (loadArrayM | loadArrayWithSetM) #-}    -- | Get computation strategy of this array   --@@ -281,7 +283,6 @@   -- @since 0.1.0   size :: Array r ix e -> Sz ix -   -- | Load an array into memory.   --   -- @since 0.3.0@@ -291,7 +292,27 @@     -> Array r ix e -- ^ Array that is being loaded     -> (Int -> e -> m ()) -- ^ Function that writes an element into target array     -> m ()+  loadArrayM scheduler arr uWrite =+    loadArrayWithSetM scheduler arr uWrite $ \offset sz e ->+      loopM_ offset (< (offset + unSz sz)) (+1) (\i -> uWrite i e)+  {-# INLINE loadArrayM #-} +  -- | Load an array into memory, just like `loadArrayM`. Except it also accepts a+  -- function that is potentially optimized for setting many cells in a region to the same+  -- value+  --+  -- @since 0.5.8+  loadArrayWithSetM+    :: Monad m =>+       Scheduler m ()+    -> Array r ix e -- ^ Array that is being loaded+    -> (Ix1 -> e -> m ()) -- ^ Function that writes an element into target array+    -> (Ix1 -> Sz1 -> e -> m ()) -- ^ Function that efficiently sets a region of an array+                                 -- to the supplied value target array+    -> m ()+  loadArrayWithSetM scheduler arr uWrite _ = loadArrayM scheduler arr uWrite+  {-# INLINE loadArrayWithSetM #-}+   defaultElement :: Array r ix e -> Maybe e   defaultElement _ = Nothing   {-# INLINE defaultElement #-}@@ -332,7 +353,7 @@     -> Array r ix e     -> m (MArray (PrimState m) r' ix e)   unsafeLoadIntoS marr arr =-    marr <$ loadArrayM trivialScheduler_ arr (unsafeLinearWrite marr)+    marr <$ loadArrayWithSetM trivialScheduler_ arr (unsafeLinearWrite marr) (unsafeLinearSet marr)   {-# INLINE unsafeLoadIntoS #-}    -- | Same as `unsafeLoadIntoS`, but respecting computation strategy.@@ -345,10 +366,13 @@     -> m (MArray RealWorld r' ix e)   unsafeLoadIntoM marr arr = do     liftIO $ withMassivScheduler_ (getComp arr) $ \scheduler ->-      loadArrayM scheduler arr (unsafeLinearWrite marr)+      loadArrayWithSetM scheduler arr (unsafeLinearWrite marr) (unsafeLinearSet marr)     pure marr   {-# INLINE unsafeLoadIntoM #-} +{-# DEPRECATED defaultElement "This is no longer used by any of the loading functions and will be removed in massiv-0.6" #-}++ -- | Selects an optimal scheduler for the supplied strategy, but it works only in `IO` withMassivScheduler_ :: Comp -> (Scheduler IO () -> IO ()) -> IO () withMassivScheduler_ comp f =@@ -910,6 +934,14 @@ elemsCount :: Load r ix e => Array r ix e -> Int elemsCount = totalElem . size {-# INLINE elemsCount #-}+++-- | Get the number of elements in the array+--+-- @since 0.5.8+linearSize :: Load r ix e => Array r ix e -> Sz1+linearSize = toLinearSz . size+{-# INLINE linearSize #-}   -- | /O(1)/ - Check if array has elements.

@@ -1,8 +1,8 @@-{-# LANGUAGE BangPatterns    #-}-{-# LANGUAGE DataKinds       #-}-{-# LANGUAGE GADTs           #-}+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE GADTs #-} {-# LANGUAGE PatternSynonyms #-}-{-# LANGUAGE TypeOperators   #-}+{-# LANGUAGE TypeOperators #-} -- | -- Module      : Data.Massiv.Core.Index -- Copyright   : (c) Alexey Kuleshevich 2018-2019@@ -40,6 +40,7 @@   , setSzM   , insertSzM   , pullOutSzM+  , toLinearSz   -- ** Dimension   , Dim(..)   , Dimension(Dim1, Dim2, Dim3, Dim4, Dim5, DimN)@@ -97,6 +98,7 @@ import Control.DeepSeq import Control.Exception (throw) import Control.Monad.Catch (MonadThrow(..))+import Data.Coerce import Data.Functor.Identity (runIdentity) import Data.Massiv.Core.Exception (guardNumberOfElements) import Data.Massiv.Core.Index.Internal@@ -224,6 +226,13 @@ -- - foldlIndex (*) 1 (unSz sz) /= 0 -- - foldlIndex (\a x -> a && x /= 0) True (unSz sz) -- - totalElem sz == 0++-- | Convert a size to a linear size.+--+-- @since 0.5.8+toLinearSz :: Index ix => Sz ix -> Sz1+toLinearSz = coerce . totalElem+{-# INLINE [1] toLinearSz #-}  -- | Get the outmost dimension of the index. --

@@ -499,9 +499,8 @@   toLinearIndex ::        Sz ix -- ^ Size     -> ix -- ^ Index-    -> Int-  default toLinearIndex :: Index (Lower ix) =>-    Sz ix -> ix -> Int+    -> Ix1+  default toLinearIndex :: Index (Lower ix) => Sz ix -> ix -> Ix1   toLinearIndex (SafeSz sz) !ix = toLinearIndex (SafeSz szL) ixL * n + i     where       !(szL, n) = unsnocDim sz@@ -512,9 +511,8 @@   -- 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 :: Ix1 -> ix -> ix -> Ix1+  default toLinearIndexAcc :: Index (Lower ix) => Ix1 -> ix -> ix -> Ix1   toLinearIndexAcc !acc !sz !ix = toLinearIndexAcc (acc * n + i) szL ixL     where       !(n, szL) = unconsDim sz@@ -524,9 +522,8 @@   -- | 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 :: Sz ix -> Ix1 -> ix+  default fromLinearIndex :: Index (Lower ix) => Sz ix -> Ix1 -> ix   fromLinearIndex (SafeSz sz) k = consDim q ixL     where       !(q, ixL) = fromLinearIndexAcc (snd (unconsDim sz)) k@@ -536,9 +533,8 @@   -- 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 -> Ix1 -> (Int, ix)+  default fromLinearIndexAcc :: Index (Lower ix) => ix -> Ix1 -> (Ix1, ix)   fromLinearIndexAcc ix' !k = (q, consDim r ixL)     where       !(m, ix) = unconsDim ix'

@@ -994,7 +994,8 @@ cons :: Load r Ix1 e => e -> Vector r e -> Vector DL e cons e v =   let dv = toLoadArray v-      load scheduler startAt uWrite = uWrite startAt e >> dlLoad dv scheduler (startAt + 1) uWrite+      load scheduler startAt uWrite uSet =+        uWrite startAt e >> dlLoad dv scheduler (startAt + 1) uWrite uSet       {-# INLINE load #-}    in dv {dlSize = SafeSz (1 + unSz (dlSize dv)), dlLoad = load} {-# INLINE cons #-}@@ -1006,7 +1007,8 @@ snoc v e =   let dv = toLoadArray v       !k = unSz (size dv)-      load scheduler startAt uWrite = dlLoad dv scheduler startAt uWrite >> uWrite (k + startAt) e+      load scheduler startAt uWrite uSet =+        dlLoad dv scheduler startAt uWrite uSet >> uWrite (k + startAt) e       {-# INLINE load #-}    in dv {dlSize = SafeSz (1 + k), dlLoad = load} {-# INLINE snoc #-}