{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}

-- |
-- Module      : Data.Massiv.Array.Mutable.Atomic
-- Copyright   : (c) Alexey Kuleshevich 2018-2022
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- 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 :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix Int -> ix -> m (Maybe Int)
atomicReadIntArray MArray (PrimState m) P ix Int
marr ix
ix
  | Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex (MArray (PrimState m) P ix Int -> Sz ix
forall ix s. Index ix => MArray s P ix Int -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) P ix Int
marr) ix
ix = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> m Int -> m (Maybe Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MArray (PrimState m) P ix Int -> ix -> m Int
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix Int -> ix -> m Int
unsafeAtomicReadIntArray MArray (PrimState m) P ix Int
marr ix
ix
  | Bool
otherwise = Maybe Int -> m (Maybe Int)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Int
forall a. Maybe a
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 :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix Int -> ix -> Int -> m Bool
atomicWriteIntArray MArray (PrimState m) P ix Int
marr ix
ix Int
f
  | Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex (MArray (PrimState m) P ix Int -> Sz ix
forall ix s. Index ix => MArray s P ix Int -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) P ix Int
marr) ix
ix = MArray (PrimState m) P ix Int -> ix -> Int -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix Int -> ix -> Int -> m ()
unsafeAtomicWriteIntArray MArray (PrimState m) P ix Int
marr ix
ix Int
f m () -> m Bool -> m Bool
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Bool -> m Bool
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
  | Bool
otherwise = Bool -> m Bool
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
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 :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix Int -> ix -> Int -> Int -> m (Maybe Int)
casIntArray MArray (PrimState m) P ix Int
marr ix
ix Int
e Int
n
  | Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex (MArray (PrimState m) P ix Int -> Sz ix
forall ix s. Index ix => MArray s P ix Int -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) P ix Int
marr) ix
ix = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> m Int -> m (Maybe Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MArray (PrimState m) P ix Int -> ix -> Int -> Int -> m Int
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix Int -> ix -> Int -> Int -> m Int
unsafeCasIntArray MArray (PrimState m) P ix Int
marr ix
ix Int
e Int
n
  | Bool
otherwise = Maybe Int -> m (Maybe Int)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Int
forall a. Maybe a
Nothing
{-# INLINE casIntArray #-}

-- | Atomically modify an `Int` element of the array. Returns the old value, unless the
-- supplied index was out of bounds.
--
-- @since 0.3.0
atomicModifyIntArray
  :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> (Int -> Int) -> m (Maybe Int)
atomicModifyIntArray :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix Int
-> ix -> (Int -> Int) -> m (Maybe Int)
atomicModifyIntArray MArray (PrimState m) P ix Int
marr ix
ix Int -> Int
f
  | Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex (MArray (PrimState m) P ix Int -> Sz ix
forall ix s. Index ix => MArray s P ix Int -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) P ix Int
marr) ix
ix = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> m Int -> m (Maybe Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MArray (PrimState m) P ix Int -> ix -> (Int -> Int) -> m Int
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix Int -> ix -> (Int -> Int) -> m Int
unsafeAtomicModifyIntArray MArray (PrimState m) P ix Int
marr ix
ix Int -> Int
f
  | Bool
otherwise = Maybe Int -> m (Maybe Int)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Int
forall a. Maybe a
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 :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix Int -> ix -> Int -> m (Maybe Int)
atomicAddIntArray MArray (PrimState m) P ix Int
marr ix
ix Int
e
  | Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex (MArray (PrimState m) P ix Int -> Sz ix
forall ix s. Index ix => MArray s P ix Int -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) P ix Int
marr) ix
ix = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> m Int -> m (Maybe Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MArray (PrimState m) P ix Int -> ix -> Int -> m Int
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix Int -> ix -> Int -> m Int
unsafeAtomicAddIntArray MArray (PrimState m) P ix Int
marr ix
ix Int
e
  | Bool
otherwise = Maybe Int -> m (Maybe Int)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Int
forall a. Maybe a
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 :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix Int -> ix -> Int -> m (Maybe Int)
atomicSubIntArray MArray (PrimState m) P ix Int
marr ix
ix Int
e
  | Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex (MArray (PrimState m) P ix Int -> Sz ix
forall ix s. Index ix => MArray s P ix Int -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) P ix Int
marr) ix
ix = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> m Int -> m (Maybe Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MArray (PrimState m) P ix Int -> ix -> Int -> m Int
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix Int -> ix -> Int -> m Int
unsafeAtomicSubIntArray MArray (PrimState m) P ix Int
marr ix
ix Int
e
  | Bool
otherwise = Maybe Int -> m (Maybe Int)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Int
forall a. Maybe a
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 :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix Int -> ix -> Int -> m (Maybe Int)
atomicAndIntArray MArray (PrimState m) P ix Int
marr ix
ix Int
e
  | Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex (MArray (PrimState m) P ix Int -> Sz ix
forall ix s. Index ix => MArray s P ix Int -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) P ix Int
marr) ix
ix = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> m Int -> m (Maybe Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MArray (PrimState m) P ix Int -> ix -> Int -> m Int
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix Int -> ix -> Int -> m Int
unsafeAtomicAndIntArray MArray (PrimState m) P ix Int
marr ix
ix Int
e
  | Bool
otherwise = Maybe Int -> m (Maybe Int)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Int
forall a. Maybe a
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 :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix Int -> ix -> Int -> m (Maybe Int)
atomicNandIntArray MArray (PrimState m) P ix Int
marr ix
ix Int
e
  | Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex (MArray (PrimState m) P ix Int -> Sz ix
forall ix s. Index ix => MArray s P ix Int -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) P ix Int
marr) ix
ix = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> m Int -> m (Maybe Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MArray (PrimState m) P ix Int -> ix -> Int -> m Int
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix Int -> ix -> Int -> m Int
unsafeAtomicNandIntArray MArray (PrimState m) P ix Int
marr ix
ix Int
e
  | Bool
otherwise = Maybe Int -> m (Maybe Int)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Int
forall a. Maybe a
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 :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix Int -> ix -> Int -> m (Maybe Int)
atomicOrIntArray MArray (PrimState m) P ix Int
marr ix
ix Int
e
  | Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex (MArray (PrimState m) P ix Int -> Sz ix
forall ix s. Index ix => MArray s P ix Int -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) P ix Int
marr) ix
ix = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> m Int -> m (Maybe Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MArray (PrimState m) P ix Int -> ix -> Int -> m Int
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix Int -> ix -> Int -> m Int
unsafeAtomicOrIntArray MArray (PrimState m) P ix Int
marr ix
ix Int
e
  | Bool
otherwise = Maybe Int -> m (Maybe Int)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Int
forall a. Maybe a
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 :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix Int -> ix -> Int -> m (Maybe Int)
atomicXorIntArray MArray (PrimState m) P ix Int
marr ix
ix Int
e
  | Sz ix -> ix -> Bool
forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex (MArray (PrimState m) P ix Int -> Sz ix
forall ix s. Index ix => MArray s P ix Int -> Sz ix
forall r e ix s.
(Manifest r e, Index ix) =>
MArray s r ix e -> Sz ix
sizeOfMArray MArray (PrimState m) P ix Int
marr) ix
ix = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> m Int -> m (Maybe Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MArray (PrimState m) P ix Int -> ix -> Int -> m Int
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix Int -> ix -> Int -> m Int
unsafeAtomicXorIntArray MArray (PrimState m) P ix Int
marr ix
ix Int
e
  | Bool
otherwise = Maybe Int -> m (Maybe Int)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Int
forall a. Maybe a
Nothing
{-# INLINE atomicXorIntArray #-}