-- | Double-buffered storage
--
-- This module provides a safer alternative to the methods of the classes
-- 'Manifestable' and 'Manifestable2':
--
-- * 'store' instead of 'manifest'
-- * 'store2' instead of 'manifest2'
-- * 'setStore' instead of 'manifestStore'
-- * 'setStore2' instead of 'manifestStore2'
--
-- Consider the following example:
--
-- > bad = do
-- >   arr  <- newArr 20
-- >   vec1 <- manifest arr (1...20)
-- >   vec2 <- manifest arr $ map (*10) $ reverse vec1
-- >   printf "%d\n" $ sum vec2
--
-- First the vector @(1...20)@ is stored into @arr@. Then the result is used to
-- compute a new vector which is also stored into @arr@. So the storage is
-- updated while it is being read from, leading to unexpected results.
--
-- Using this module, we can make a small change to the program:
--
-- > good = do
-- >   st   <- newStore 20
-- >   vec1 <- store st (1...20)
-- >   vec2 <- store st $ map (*10) $ reverse vec1
-- >   printf "%d\n" $ sum vec2
--
-- Now the program works as expected; i.e. gives the same result as the normal
-- Haskell expression
--
-- > sum $ map (*10) $ reverse [1..20]
--
-- The price we have to pay for safe storage is that @`newStore` l@ allocates
-- twice as much memory as @`newArr` l@. However, none of the other functions in
-- this module allocate any memory.
--
-- Note that this module does not protect against improper use of
-- 'unsafeFreezeStore'. A vector from a frozen 'Store' is only valid as long as
-- the 'Store' is not updated.

module Feldspar.Data.Buffered
  ( Store
  , newStore
  , unsafeInplaceStore
  , unsafeFreezeStore
  , unsafeFreezeStore2
  , setStore
  , setStore2
  , store
  , store2
  , loopStore
  , loopStore2
  ) where

-- By only allowing `Store` to be created using `newStore`, we ensure that
-- `unsafeSwapArr` is only used in a safe way (on two arrays allocated in the
-- same scope).



import Prelude ()

import Control.Monad.State

import Feldspar.Representation
import Feldspar.Run
import Feldspar.Data.Vector



-- | Double-buffered storage
data Store a = Store
    { Store a -> Arr a
activeBuf :: Arr a
    , Store a -> Arr a
freeBuf   :: Arr a
    }

-- | Create a new double-buffered 'Store'
--
-- This operation allocates two arrays of the given length.
newStore :: (Syntax a, MonadComp m) => Data Length -> m (Store a)
newStore :: Data Length -> m (Store a)
newStore Data Length
l = Arr a -> Arr a -> Store a
forall a. Arr a -> Arr a -> Store a
Store (Arr a -> Arr a -> Store a) -> m (Arr a) -> m (Arr a -> Store a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> Data Length -> m (Arr a)
forall a (m :: * -> *).
(Type (Internal a), MonadComp m) =>
String -> Data Length -> m (Arr a)
newNamedArr String
"store" Data Length
l m (Arr a -> Store a) -> m (Arr a) -> m (Store a)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> String -> Data Length -> m (Arr a)
forall a (m :: * -> *).
(Type (Internal a), MonadComp m) =>
String -> Data Length -> m (Arr a)
newNamedArr String
"store" Data Length
l

-- | Create a new single-buffered 'Store'
--
-- Using 'unsafeInplaceStore' instead of 'newStore' allows double-buffered
-- algorithms to run inplace.
unsafeInplaceStore :: (Syntax a, MonadComp m) => Data Length -> m (Store a)
unsafeInplaceStore :: Data Length -> m (Store a)
unsafeInplaceStore Data Length
l = do
    Arr a
arr <- String -> Data Length -> m (Arr a)
forall a (m :: * -> *).
(Type (Internal a), MonadComp m) =>
String -> Data Length -> m (Arr a)
newNamedArr String
"store" Data Length
l
    Store a -> m (Store a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Store a -> m (Store a)) -> Store a -> m (Store a)
forall a b. (a -> b) -> a -> b
$ Arr a -> Arr a -> Store a
forall a. Arr a -> Arr a -> Store a
Store Arr a
arr Arr a
arr

-- | Read the contents of a 'Store' without making a copy. This is generally
-- only safe if the the 'Store' is not updated as long as the resulting vector
-- is alive.
unsafeFreezeStore :: (Syntax a, MonadComp m) =>
    Data Length -> Store a -> m (Manifest a)
unsafeFreezeStore :: Data Length -> Store a -> m (Manifest a)
unsafeFreezeStore Data Length
l = Data Length -> Arr a -> m (Manifest a)
forall (m :: * -> *) a.
MonadComp m =>
Data Length -> Arr a -> m (IArr a)
unsafeFreezeSlice Data Length
l (Arr a -> m (Manifest a))
-> (Store a -> Arr a) -> Store a -> m (Manifest a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Store a -> Arr a
forall a. Store a -> Arr a
activeBuf

-- | Read the contents of a 'Store' without making a copy (2-dimensional
-- version). This is generally only safe if the the 'Store' is not updated as
-- long as the resulting vector is alive.
unsafeFreezeStore2 :: (Syntax a, MonadComp m)
    => Data Length  -- ^ Number of rows
    -> Data Length  -- ^ Number of columns
    -> Store a
    -> m (Manifest2 a)
unsafeFreezeStore2 :: Data Length -> Data Length -> Store a -> m (Manifest2 a)
unsafeFreezeStore2 Data Length
r Data Length
c Store {Arr a
freeBuf :: Arr a
activeBuf :: Arr a
freeBuf :: forall a. Store a -> Arr a
activeBuf :: forall a. Store a -> Arr a
..} =
    Data Length -> Data Length -> IArr a -> Manifest2 a
forall a. Finite a => Data Length -> Data Length -> a -> Nest a
nest Data Length
r Data Length
c (IArr a -> Manifest2 a) -> m (IArr a) -> m (Manifest2 a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Data Length -> Arr a -> m (IArr a)
forall (m :: * -> *) a.
MonadComp m =>
Data Length -> Arr a -> m (IArr a)
unsafeFreezeSlice (Data Length
rData Length -> Data Length -> Data Length
forall a. Num a => a -> a -> a
*Data Length
c) Arr a
activeBuf

-- | Cheap swapping of the two buffers in a 'Store'
swapStore :: Syntax a => Store a -> Run ()
swapStore :: Store a -> Run ()
swapStore Store {Arr a
freeBuf :: Arr a
activeBuf :: Arr a
freeBuf :: forall a. Store a -> Arr a
activeBuf :: forall a. Store a -> Arr a
..} = Arr a -> Arr a -> Run ()
forall a. Arr a -> Arr a -> Run ()
unsafeSwapArr Arr a
activeBuf Arr a
freeBuf

-- | Write a 1-dimensional vector to a 'Store'. The operation may become a no-op
-- if the vector is already in the 'Store'.
setStore :: (Manifestable Run vec a, Finite vec, Syntax a) =>
    Store a -> vec -> Run ()
setStore :: Store a -> vec -> Run ()
setStore st :: Store a
st@Store {Arr a
freeBuf :: Arr a
activeBuf :: Arr a
freeBuf :: forall a. Store a -> Arr a
activeBuf :: forall a. Store a -> Arr a
..} vec
vec = case vec -> Maybe (Manifest a)
forall vec a. ViewManifest vec a => vec -> Maybe (Manifest a)
viewManifest vec
vec of
    Just Manifest a
iarr
      | Arr a -> Manifest a -> Bool
forall a. Arr a -> IArr a -> Bool
unsafeEqArrIArr Arr a
activeBuf Manifest a
iarr ->
          Data Bool -> Run () -> Run () -> Run ()
forall (m :: * -> *).
MonadComp m =>
Data Bool -> m () -> m () -> m ()
iff (Manifest a -> Data Length
forall a. IArr a -> Data Length
iarrOffset Manifest a
iarr Data Length -> Data Length -> Data Bool
forall a. PrimType a => Data a -> Data a -> Data Bool
== Arr a -> Data Length
forall a. Arr a -> Data Length
arrOffset Arr a
activeBuf)
            (() -> Run ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
            Run ()
saveAndSwap
          -- We don't check if `iarr` is equal to the free buffer, because that
          -- would mean that we're trying to overwrite a frozen buffer while
          -- reading it, which should lead to undefined behavior.
    Maybe (Manifest a)
_ -> Run ()
saveAndSwap
  where
    saveAndSwap :: Run ()
saveAndSwap = Arr a -> vec -> Run ()
forall (m :: * -> *) vec a.
(Manifestable m vec a, Syntax a) =>
Arr a -> vec -> m ()
manifestStore Arr a
freeBuf vec
vec Run () -> Run () -> Run ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Store a -> Run ()
forall a. Syntax a => Store a -> Run ()
swapStore Store a
st

-- | Write a 2-dimensional vector to a 'Store'. The operation may become a no-op
-- if the vector is already in the 'Store'.
setStore2 :: (Manifestable2 Run vec a, Finite2 vec, Syntax a) =>
    Store a -> vec -> Run ()
setStore2 :: Store a -> vec -> Run ()
setStore2 st :: Store a
st@Store {Arr a
freeBuf :: Arr a
activeBuf :: Arr a
freeBuf :: forall a. Store a -> Arr a
activeBuf :: forall a. Store a -> Arr a
..} vec
vec = case vec -> Maybe (Manifest2 a)
forall vec a. ViewManifest2 vec a => vec -> Maybe (Manifest2 a)
viewManifest2 vec
vec of
    Just Manifest2 a
arr
      | let iarr :: Manifest a
iarr = Manifest2 a -> Manifest a
forall a. Slicable a => Nest a -> a
unnest Manifest2 a
arr
      , Arr a -> Manifest a -> Bool
forall a. Arr a -> IArr a -> Bool
unsafeEqArrIArr Arr a
activeBuf Manifest a
iarr ->
          Data Bool -> Run () -> Run () -> Run ()
forall (m :: * -> *).
MonadComp m =>
Data Bool -> m () -> m () -> m ()
iff (Manifest a -> Data Length
forall a. IArr a -> Data Length
iarrOffset Manifest a
iarr Data Length -> Data Length -> Data Bool
forall a. PrimType a => Data a -> Data a -> Data Bool
== Arr a -> Data Length
forall a. Arr a -> Data Length
arrOffset Arr a
activeBuf)
            (() -> Run ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
            Run ()
saveAndSwap
          -- See comment to `setStore`
    Maybe (Manifest2 a)
_ -> Run ()
saveAndSwap
  where
    saveAndSwap :: Run ()
saveAndSwap = Arr a -> vec -> Run ()
forall (m :: * -> *) vec a.
(Manifestable2 m vec a, Syntax a) =>
Arr a -> vec -> m ()
manifestStore2 Arr a
freeBuf vec
vec Run () -> Run () -> Run ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Store a -> Run ()
forall a. Syntax a => Store a -> Run ()
swapStore Store a
st

-- | Write the contents of a vector to a 'Store' and get it back as a
-- 'Manifest' vector
store :: (Manifestable Run vec a, Finite vec, Syntax a) =>
    Store a -> vec -> Run (Manifest a)
store :: Store a -> vec -> Run (Manifest a)
store Store a
st vec
vec = Store a -> vec -> Run ()
forall vec a.
(Manifestable Run vec a, Finite vec, Syntax a) =>
Store a -> vec -> Run ()
setStore Store a
st vec
vec Run () -> Run (Manifest a) -> Run (Manifest a)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Data Length -> Store a -> Run (Manifest a)
forall a (m :: * -> *).
(Syntax a, MonadComp m) =>
Data Length -> Store a -> m (Manifest a)
unsafeFreezeStore (vec -> Data Length
forall a. Finite a => a -> Data Length
length vec
vec) Store a
st

-- | Write the contents of a vector to a 'Store' and get it back as a
-- 'Manifest2' vector
store2 :: (Manifestable2 Run vec a, Finite2 vec, Syntax a) =>
    Store a -> vec -> Run (Manifest2 a)
store2 :: Store a -> vec -> Run (Manifest2 a)
store2 Store a
st vec
vec = Store a -> vec -> Run ()
forall vec a.
(Manifestable2 Run vec a, Finite2 vec, Syntax a) =>
Store a -> vec -> Run ()
setStore2 Store a
st vec
vec Run () -> Run (Manifest2 a) -> Run (Manifest2 a)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Data Length -> Data Length -> Store a -> Run (Manifest2 a)
forall a (m :: * -> *).
(Syntax a, MonadComp m) =>
Data Length -> Data Length -> Store a -> m (Manifest2 a)
unsafeFreezeStore2 Data Length
r Data Length
c Store a
st
  where
    (Data Length
r,Data Length
c) = vec -> (Data Length, Data Length)
forall a. Finite2 a => a -> (Data Length, Data Length)
extent2 vec
vec

loopStore
    :: ( Integral i
       , PrimType i
       , Syntax a
       , Manifestable Run vec1 a
       , Finite vec1
       , Manifestable Run vec2 a
       , Finite vec2
       )
    => Store a
    -> IxRange (Data i)
    -> (Data i -> Manifest a -> Run vec1)
    -> vec2
    -> Run (Manifest a)
loopStore :: Store a
-> IxRange (Data i)
-> (Data i -> Manifest a -> Run vec1)
-> vec2
-> Run (Manifest a)
loopStore Store a
st IxRange (Data i)
rng Data i -> Manifest a -> Run vec1
body vec2
init = do
    Store a -> vec2 -> Run ()
forall vec a.
(Manifestable Run vec a, Finite vec, Syntax a) =>
Store a -> vec -> Run ()
setStore Store a
st vec2
init
    Ref (Data Length)
lr <- Data Length -> Run (Ref (Data Length))
forall a (m :: * -> *). (Syntax a, MonadComp m) => a -> m (Ref a)
initRef (Data Length -> Run (Ref (Data Length)))
-> Data Length -> Run (Ref (Data Length))
forall a b. (a -> b) -> a -> b
$ vec2 -> Data Length
forall a. Finite a => a -> Data Length
length vec2
init
    IxRange (Data i) -> (Data i -> Run ()) -> Run ()
forall (m :: * -> *) n.
(MonadComp m, Integral n, PrimType n) =>
IxRange (Data n) -> (Data n -> m ()) -> m ()
for IxRange (Data i)
rng ((Data i -> Run ()) -> Run ()) -> (Data i -> Run ()) -> Run ()
forall a b. (a -> b) -> a -> b
$ \Data i
i -> do
      Data Length
l    <- Ref (Data Length) -> Run (Data Length)
forall a (m :: * -> *). (Syntax a, MonadComp m) => Ref a -> m a
unsafeFreezeRef Ref (Data Length)
lr
      vec1
next <- Data i -> Manifest a -> Run vec1
body Data i
i (Manifest a -> Run vec1) -> Run (Manifest a) -> Run vec1
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Data Length -> Store a -> Run (Manifest a)
forall a (m :: * -> *).
(Syntax a, MonadComp m) =>
Data Length -> Store a -> m (Manifest a)
unsafeFreezeStore Data Length
l Store a
st
      Store a -> vec1 -> Run ()
forall vec a.
(Manifestable Run vec a, Finite vec, Syntax a) =>
Store a -> vec -> Run ()
setStore Store a
st vec1
next
      Ref (Data Length) -> Data Length -> Run ()
forall a (m :: * -> *).
(Syntax a, MonadComp m) =>
Ref a -> a -> m ()
setRef Ref (Data Length)
lr (Data Length -> Run ()) -> Data Length -> Run ()
forall a b. (a -> b) -> a -> b
$ vec1 -> Data Length
forall a. Finite a => a -> Data Length
length vec1
next
    Data Length
l <- Ref (Data Length) -> Run (Data Length)
forall a (m :: * -> *). (Syntax a, MonadComp m) => Ref a -> m a
unsafeFreezeRef Ref (Data Length)
lr
    Data Length -> Store a -> Run (Manifest a)
forall a (m :: * -> *).
(Syntax a, MonadComp m) =>
Data Length -> Store a -> m (Manifest a)
unsafeFreezeStore Data Length
l Store a
st

loopStore2
    :: ( Integral i
       , PrimType i
       , Syntax a
       , Manifestable2 Run vec1 a
       , Finite2 vec1
       , Manifestable2 Run vec2 a
       , Finite2 vec2
       )
    => Store a
    -> IxRange (Data i)
    -> (Data i -> Manifest2 a -> Run vec1)
    -> vec2
    -> Run (Manifest2 a)
loopStore2 :: Store a
-> IxRange (Data i)
-> (Data i -> Manifest2 a -> Run vec1)
-> vec2
-> Run (Manifest2 a)
loopStore2 Store a
st IxRange (Data i)
rng Data i -> Manifest2 a -> Run vec1
body vec2
init = do
    Store a -> vec2 -> Run ()
forall vec a.
(Manifestable2 Run vec a, Finite2 vec, Syntax a) =>
Store a -> vec -> Run ()
setStore2 Store a
st vec2
init
    Ref (Data Length)
rr <- Data Length -> Run (Ref (Data Length))
forall a (m :: * -> *). (Syntax a, MonadComp m) => a -> m (Ref a)
initRef (Data Length -> Run (Ref (Data Length)))
-> Data Length -> Run (Ref (Data Length))
forall a b. (a -> b) -> a -> b
$ vec2 -> Data Length
forall a. Finite2 a => a -> Data Length
numRows vec2
init
    Ref (Data Length)
cr <- Data Length -> Run (Ref (Data Length))
forall a (m :: * -> *). (Syntax a, MonadComp m) => a -> m (Ref a)
initRef (Data Length -> Run (Ref (Data Length)))
-> Data Length -> Run (Ref (Data Length))
forall a b. (a -> b) -> a -> b
$ vec2 -> Data Length
forall a. Finite2 a => a -> Data Length
numCols vec2
init
    IxRange (Data i) -> (Data i -> Run ()) -> Run ()
forall (m :: * -> *) n.
(MonadComp m, Integral n, PrimType n) =>
IxRange (Data n) -> (Data n -> m ()) -> m ()
for IxRange (Data i)
rng ((Data i -> Run ()) -> Run ()) -> (Data i -> Run ()) -> Run ()
forall a b. (a -> b) -> a -> b
$ \Data i
i -> do
      Data Length
r    <- Ref (Data Length) -> Run (Data Length)
forall a (m :: * -> *). (Syntax a, MonadComp m) => Ref a -> m a
unsafeFreezeRef Ref (Data Length)
rr
      Data Length
c    <- Ref (Data Length) -> Run (Data Length)
forall a (m :: * -> *). (Syntax a, MonadComp m) => Ref a -> m a
unsafeFreezeRef Ref (Data Length)
cr
      vec1
next <- Data i -> Manifest2 a -> Run vec1
body Data i
i (Manifest2 a -> Run vec1) -> Run (Manifest2 a) -> Run vec1
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Data Length -> Data Length -> Store a -> Run (Manifest2 a)
forall a (m :: * -> *).
(Syntax a, MonadComp m) =>
Data Length -> Data Length -> Store a -> m (Manifest2 a)
unsafeFreezeStore2 Data Length
r Data Length
c Store a
st
      Store a -> vec1 -> Run ()
forall vec a.
(Manifestable2 Run vec a, Finite2 vec, Syntax a) =>
Store a -> vec -> Run ()
setStore2 Store a
st vec1
next
      Ref (Data Length) -> Data Length -> Run ()
forall a (m :: * -> *).
(Syntax a, MonadComp m) =>
Ref a -> a -> m ()
setRef Ref (Data Length)
rr (Data Length -> Run ()) -> Data Length -> Run ()
forall a b. (a -> b) -> a -> b
$ vec1 -> Data Length
forall a. Finite2 a => a -> Data Length
numRows vec1
next
      Ref (Data Length) -> Data Length -> Run ()
forall a (m :: * -> *).
(Syntax a, MonadComp m) =>
Ref a -> a -> m ()
setRef Ref (Data Length)
cr (Data Length -> Run ()) -> Data Length -> Run ()
forall a b. (a -> b) -> a -> b
$ vec1 -> Data Length
forall a. Finite2 a => a -> Data Length
numCols vec1
next
    Data Length
r <- Ref (Data Length) -> Run (Data Length)
forall a (m :: * -> *). (Syntax a, MonadComp m) => Ref a -> m a
unsafeFreezeRef Ref (Data Length)
rr
    Data Length
c <- Ref (Data Length) -> Run (Data Length)
forall a (m :: * -> *). (Syntax a, MonadComp m) => Ref a -> m a
unsafeFreezeRef Ref (Data Length)
cr
    Data Length -> Data Length -> Store a -> Run (Manifest2 a)
forall a (m :: * -> *).
(Syntax a, MonadComp m) =>
Data Length -> Data Length -> Store a -> m (Manifest2 a)
unsafeFreezeStore2 Data Length
r Data Length
c Store a
st