{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
module Data.Massiv.Array.Unsafe
(
unsafeMakeArray
, unsafeGenerateArray
, unsafeGenerateArrayP
, unsafeGenerateM
, unsafeIndex
, unsafeLinearIndex
, unsafeLinearIndexM
, unsafeBackpermute
, unsafeTraverse
, unsafeTraverse2
, unsafeResize
, unsafeExtract
, unsafeSlice
, unsafeOuterSlice
, unsafeInnerSlice
, unsafeThaw
, unsafeFreeze
, unsafeNew
, unsafeNewZero
, unsafeRead
, unsafeLinearRead
, unsafeWrite
, unsafeLinearWrite
) 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)
unsafeBackpermute :: (Source r' ix' e, Index ix) =>
ix -> (ix -> ix') -> Array r' ix' e -> Array D ix e
unsafeBackpermute !sz ixF !arr =
unsafeMakeArray (getComp arr) sz $ \ !ix -> unsafeIndex arr (ixF ix)
{-# INLINE unsafeBackpermute #-}
unsafeTraverse
:: (Source r1 ix1 e1, Index ix)
=> ix
-> ((ix1 -> e1) -> ix -> e)
-> Array r1 ix1 e1
-> Array D ix e
unsafeTraverse sz f arr1 =
unsafeMakeArray (getComp arr1) sz (f (unsafeIndex arr1))
{-# INLINE unsafeTraverse #-}
unsafeTraverse2
:: (Source r1 ix1 e1, Source r2 ix2 e2, Index ix)
=> ix
-> ((ix1 -> e1) -> (ix2 -> e2) -> ix -> e)
-> Array r1 ix1 e1
-> Array r2 ix2 e2
-> Array D ix e
unsafeTraverse2 sz f arr1 arr2 =
unsafeMakeArray (getComp arr1) sz (f (unsafeIndex arr1) (unsafeIndex arr2))
{-# INLINE unsafeTraverse2 #-}
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 #-}
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 #-}
unsafeGenerateArray :: Mutable r ix e => 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 #-}
unsafeGenerateArrayP :: Mutable r ix e => [Int] -> 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 #-}