{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
-- |
-- Module      : Data.Vector.Generic.New
-- Copyright   : (c) Roman Leshchinskiy 2008-2010
--                   Alexey Kuleshevich 2020-2022
--                   Aleksey Khudyakov 2020-2022
--                   Andrew Lelechenko 2020-2022
-- License     : BSD-style
--
-- Maintainer  : Haskell Libraries Team <libraries@haskell.org>
-- Stability   : experimental
-- Portability : non-portable
--
-- Purely functional interface to initialisation of mutable vectors
--

module Data.Vector.Generic.New (
  -- * Array recycling primitives
  New(..), create, run, runPrim, apply, modify, modifyWithBundle,
  unstream, transform, unstreamR, transformR,
  slice, init, tail, take, drop,
  unsafeSlice, unsafeInit, unsafeTail
  -- * References
  -- $references
) where

import qualified Data.Vector.Generic.Mutable as MVector

import           Data.Vector.Generic.Base ( Vector, Mutable )

import           Data.Vector.Fusion.Bundle ( Bundle )
import qualified Data.Vector.Fusion.Bundle as Bundle
import           Data.Vector.Fusion.Stream.Monadic ( Stream )
import           Data.Vector.Fusion.Bundle.Size

import Control.Monad.Primitive
import Control.Monad.ST ( ST )
import Control.Monad  ( liftM )
import Prelude
  ( Monad, Int
  , return, seq
  , (.), (=<<) )

-- Data.Vector.Internal.Check is unused
#define NOT_VECTOR_MODULE
#include "vector.h"

-- | This data type is a wrapper around a monadic action which produces
-- a mutable vector. It's used by a number of rewrite rules in order to
-- facilitate the reuse of buffers allocated for vectors. See "Recycle
-- your arrays!" for a detailed explanation.
--
-- Note that this data type must be declared as @data@ and not @newtype@
-- since it's used for rewrite rules and rules won't fire with @newtype@.
data New v a = New (forall s. ST s (Mutable v s a))

create :: (forall s. ST s (Mutable v s a)) -> New v a
{-# INLINE create #-}
create :: forall (v :: * -> *) a. (forall s. ST s (Mutable v s a)) -> New v a
create forall s. ST s (Mutable v s a)
p = forall (v :: * -> *) a. (forall s. ST s (Mutable v s a)) -> New v a
New forall s. ST s (Mutable v s a)
p

run :: New v a -> ST s (Mutable v s a)
{-# INLINE run #-}
run :: forall (v :: * -> *) a s. New v a -> ST s (Mutable v s a)
run (New forall s. ST s (Mutable v s a)
p) = forall s. ST s (Mutable v s a)
p

runPrim :: PrimMonad m => New v a -> m (Mutable v (PrimState m) a)
{-# INLINE runPrim #-}
runPrim :: forall (m :: * -> *) (v :: * -> *) a.
PrimMonad m =>
New v a -> m (Mutable v (PrimState m) a)
runPrim (New forall s. ST s (Mutable v s a)
p) = forall (m1 :: * -> *) (m2 :: * -> *) a.
(PrimBase m1, PrimMonad m2, PrimState m1 ~ PrimState m2) =>
m1 a -> m2 a
primToPrim forall s. ST s (Mutable v s a)
p

apply :: (forall s. Mutable v s a -> Mutable v s a) -> New v a -> New v a
{-# INLINE apply #-}
apply :: forall (v :: * -> *) a.
(forall s. Mutable v s a -> Mutable v s a) -> New v a -> New v a
apply forall s. Mutable v s a -> Mutable v s a
f (New forall s. ST s (Mutable v s a)
p) = forall (v :: * -> *) a. (forall s. ST s (Mutable v s a)) -> New v a
New (forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM forall s. Mutable v s a -> Mutable v s a
f forall s. ST s (Mutable v s a)
p)

modify :: (forall s. Mutable v s a -> ST s ()) -> New v a -> New v a
{-# INLINE modify #-}
modify :: forall (v :: * -> *) a.
(forall s. Mutable v s a -> ST s ()) -> New v a -> New v a
modify forall s. Mutable v s a -> ST s ()
f (New forall s. ST s (Mutable v s a)
p) = forall (v :: * -> *) a. (forall s. ST s (Mutable v s a)) -> New v a
New (do { Mutable v s a
v <- forall s. ST s (Mutable v s a)
p; forall s. Mutable v s a -> ST s ()
f Mutable v s a
v; forall (m :: * -> *) a. Monad m => a -> m a
return Mutable v s a
v })

modifyWithBundle :: (forall s. Mutable v s a -> Bundle u b -> ST s ())
                 -> New v a -> Bundle u b -> New v a
{-# INLINE_FUSED modifyWithBundle #-}
modifyWithBundle :: forall (v :: * -> *) a (u :: * -> *) b.
(forall s. Mutable v s a -> Bundle u b -> ST s ())
-> New v a -> Bundle u b -> New v a
modifyWithBundle forall s. Mutable v s a -> Bundle u b -> ST s ()
f (New forall s. ST s (Mutable v s a)
p) Bundle u b
s = Bundle u b
s seq :: forall a b. a -> b -> b
`seq` forall (v :: * -> *) a. (forall s. ST s (Mutable v s a)) -> New v a
New (do { Mutable v s a
v <- forall s. ST s (Mutable v s a)
p; forall s. Mutable v s a -> Bundle u b -> ST s ()
f Mutable v s a
v Bundle u b
s; forall (m :: * -> *) a. Monad m => a -> m a
return Mutable v s a
v })

unstream :: Vector v a => Bundle v a -> New v a
{-# INLINE_FUSED unstream #-}
unstream :: forall (v :: * -> *) a. Vector v a => Bundle v a -> New v a
unstream Bundle v a
s = Bundle v a
s seq :: forall a b. a -> b -> b
`seq` forall (v :: * -> *) a. (forall s. ST s (Mutable v s a)) -> New v a
New (forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Bundle v a -> m (Mutable v (PrimState m) a)
MVector.vunstream Bundle v a
s)

transform
  :: Vector v a => (forall m. Monad m => Stream m a -> Stream m a)
                -> (Size -> Size) -> New v a -> New v a
{-# INLINE_FUSED transform #-}
transform :: forall (v :: * -> *) a.
Vector v a =>
(forall (m :: * -> *). Monad m => Stream m a -> Stream m a)
-> (Size -> Size) -> New v a -> New v a
transform forall (m :: * -> *). Monad m => Stream m a -> Stream m a
f Size -> Size
_ (New forall s. ST s (Mutable v s a)
p) = forall (v :: * -> *) a. (forall s. ST s (Mutable v s a)) -> New v a
New (forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
(Stream m a -> Stream m a)
-> v (PrimState m) a -> m (v (PrimState m) a)
MVector.transform forall (m :: * -> *). Monad m => Stream m a -> Stream m a
f forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall s. ST s (Mutable v s a)
p)

{-# RULES

"transform/transform [New]"
  forall (f1 :: forall m. Monad m => Stream m a -> Stream m a)
         (f2 :: forall m. Monad m => Stream m a -> Stream m a)
         g1 g2 p .
  transform f1 g1 (transform f2 g2 p) = transform (f1 . f2) (g1 . g2) p

"transform/unstream [New]"
  forall (f :: forall m. Monad m => Stream m a -> Stream m a)
         g s.
  transform f g (unstream s) = unstream (Bundle.inplace f g s)  #-}




unstreamR :: Vector v a => Bundle v a -> New v a
{-# INLINE_FUSED unstreamR #-}
unstreamR :: forall (v :: * -> *) a. Vector v a => Bundle v a -> New v a
unstreamR Bundle v a
s = Bundle v a
s seq :: forall a b. a -> b -> b
`seq` forall (v :: * -> *) a. (forall s. ST s (Mutable v s a)) -> New v a
New (forall (m :: * -> *) (v :: * -> * -> *) a (u :: * -> *).
(PrimMonad m, MVector v a) =>
Bundle u a -> m (v (PrimState m) a)
MVector.unstreamR Bundle v a
s)

transformR
  :: Vector v a => (forall m. Monad m => Stream m a -> Stream m a)
                -> (Size -> Size) -> New v a -> New v a
{-# INLINE_FUSED transformR #-}
transformR :: forall (v :: * -> *) a.
Vector v a =>
(forall (m :: * -> *). Monad m => Stream m a -> Stream m a)
-> (Size -> Size) -> New v a -> New v a
transformR forall (m :: * -> *). Monad m => Stream m a -> Stream m a
f Size -> Size
_ (New forall s. ST s (Mutable v s a)
p) = forall (v :: * -> *) a. (forall s. ST s (Mutable v s a)) -> New v a
New (forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
(Stream m a -> Stream m a)
-> v (PrimState m) a -> m (v (PrimState m) a)
MVector.transformR forall (m :: * -> *). Monad m => Stream m a -> Stream m a
f forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall s. ST s (Mutable v s a)
p)

{-# RULES

"transformR/transformR [New]"
  forall (f1 :: forall m. Monad m => Stream m a -> Stream m a)
         (f2 :: forall m. Monad m => Stream m a -> Stream m a)
         g1 g2
         p .
  transformR f1 g1 (transformR f2 g2 p) = transformR (f1 . f2) (g1 . g2) p

"transformR/unstreamR [New]"
  forall (f :: forall m. Monad m => Stream m a -> Stream m a)
         g s.
  transformR f g (unstreamR s) = unstreamR (Bundle.inplace f g s)  #-}



slice :: Vector v a => Int -> Int -> New v a -> New v a
{-# INLINE_FUSED slice #-}
slice :: forall (v :: * -> *) a.
Vector v a =>
Int -> Int -> New v a -> New v a
slice Int
i Int
n New v a
m = forall (v :: * -> *) a.
(forall s. Mutable v s a -> Mutable v s a) -> New v a -> New v a
apply (forall (v :: * -> * -> *) a s.
(HasCallStack, MVector v a) =>
Int -> Int -> v s a -> v s a
MVector.slice Int
i Int
n) New v a
m

init :: Vector v a => New v a -> New v a
{-# INLINE_FUSED init #-}
init :: forall (v :: * -> *) a. Vector v a => New v a -> New v a
init New v a
m = forall (v :: * -> *) a.
(forall s. Mutable v s a -> Mutable v s a) -> New v a -> New v a
apply forall (v :: * -> * -> *) a s. MVector v a => v s a -> v s a
MVector.init New v a
m

tail :: Vector v a => New v a -> New v a
{-# INLINE_FUSED tail #-}
tail :: forall (v :: * -> *) a. Vector v a => New v a -> New v a
tail New v a
m = forall (v :: * -> *) a.
(forall s. Mutable v s a -> Mutable v s a) -> New v a -> New v a
apply forall (v :: * -> * -> *) a s. MVector v a => v s a -> v s a
MVector.tail New v a
m

take :: Vector v a => Int -> New v a -> New v a
{-# INLINE_FUSED take #-}
take :: forall (v :: * -> *) a. Vector v a => Int -> New v a -> New v a
take Int
n New v a
m = forall (v :: * -> *) a.
(forall s. Mutable v s a -> Mutable v s a) -> New v a -> New v a
apply (forall (v :: * -> * -> *) a s. MVector v a => Int -> v s a -> v s a
MVector.take Int
n) New v a
m

drop :: Vector v a => Int -> New v a -> New v a
{-# INLINE_FUSED drop #-}
drop :: forall (v :: * -> *) a. Vector v a => Int -> New v a -> New v a
drop Int
n New v a
m = forall (v :: * -> *) a.
(forall s. Mutable v s a -> Mutable v s a) -> New v a -> New v a
apply (forall (v :: * -> * -> *) a s. MVector v a => Int -> v s a -> v s a
MVector.drop Int
n) New v a
m

unsafeSlice :: Vector v a => Int -> Int -> New v a -> New v a
{-# INLINE_FUSED unsafeSlice #-}
unsafeSlice :: forall (v :: * -> *) a.
Vector v a =>
Int -> Int -> New v a -> New v a
unsafeSlice Int
i Int
n New v a
m = forall (v :: * -> *) a.
(forall s. Mutable v s a -> Mutable v s a) -> New v a -> New v a
apply (forall (v :: * -> * -> *) a s.
MVector v a =>
Int -> Int -> v s a -> v s a
MVector.unsafeSlice Int
i Int
n) New v a
m

unsafeInit :: Vector v a => New v a -> New v a
{-# INLINE_FUSED unsafeInit #-}
unsafeInit :: forall (v :: * -> *) a. Vector v a => New v a -> New v a
unsafeInit New v a
m = forall (v :: * -> *) a.
(forall s. Mutable v s a -> Mutable v s a) -> New v a -> New v a
apply forall (v :: * -> * -> *) a s. MVector v a => v s a -> v s a
MVector.unsafeInit New v a
m

unsafeTail :: Vector v a => New v a -> New v a
{-# INLINE_FUSED unsafeTail #-}
unsafeTail :: forall (v :: * -> *) a. Vector v a => New v a -> New v a
unsafeTail New v a
m = forall (v :: * -> *) a.
(forall s. Mutable v s a -> Mutable v s a) -> New v a -> New v a
apply forall (v :: * -> * -> *) a s. MVector v a => v s a -> v s a
MVector.unsafeTail New v a
m

{-# RULES

"slice/unstream [New]" forall i n s.
  slice i n (unstream s) = unstream (Bundle.slice i n s)

"init/unstream [New]" forall s.
  init (unstream s) = unstream (Bundle.init s)

"tail/unstream [New]" forall s.
  tail (unstream s) = unstream (Bundle.tail s)

"take/unstream [New]" forall n s.
  take n (unstream s) = unstream (Bundle.take n s)

"drop/unstream [New]" forall n s.
  drop n (unstream s) = unstream (Bundle.drop n s)

"unsafeSlice/unstream [New]" forall i n s.
  unsafeSlice i n (unstream s) = unstream (Bundle.slice i n s)

"unsafeInit/unstream [New]" forall s.
  unsafeInit (unstream s) = unstream (Bundle.init s)

"unsafeTail/unstream [New]" forall s.
  unsafeTail (unstream s) = unstream (Bundle.tail s)   #-}


-- $references
--
-- * Leshchinskiy, Roman. "Recycle your arrays!". Practical Aspects of
--   Declarative Languages: 11th International Symposium, PADL 2009,
--   Savannah, GA, USA, January 19-20, 2009. Proceedings 11. Springer
--   Berlin Heidelberg, 2009.