{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
# LANGUAGE MultiParamTypeClasses #
# LANGUAGE PolyKinds              #
{-# LANGUAGE Rank2Types            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
# LANGUAGE TypeFamilies           #
-- |
-- Type classes for vectors which are implemented on top of the arrays
-- and support in-place mutation. API is similar to one used in the
-- @vector@ package.
module Data.Vector.Fixed.Mutable (
    -- * Mutable vectors
    Arity
  , arity
  , Mutable
  , DimM
  , MVector(..)
  , lengthM
  , read
  , write
  , clone
    -- * Creation
  , replicate
  , replicateM
  , generate
  , generateM
    -- * Loops
  , forI
    -- * Immutable vectors
  , IVector(..)
  , index
  , freeze
  , thaw
    -- * Vector API
  , constructVec
  , inspectVec
  ) where

import Control.Applicative  (Const(..))
import Control.Monad.ST
import Control.Monad.Primitive
import Data.Typeable  (Proxy(..))
import GHC.TypeLits
import Data.Vector.Fixed.Cont (Dim,PeanoNum(..),Peano,Arity,Fun(..),Vector(..),ContVec,arity,apply,accum,length)
import Prelude hiding (read,length,replicate)


----------------------------------------------------------------
-- Type classes
----------------------------------------------------------------

-- | Mutable counterpart of fixed-length vector.
type family Mutable (v :: * -> *) :: * -> * -> *

-- | Dimension for mutable vector.
type family DimM (v :: * -> * -> *) :: Nat

-- | Type class for mutable vectors.
class (Arity (DimM v)) => MVector v a where
  | Copy vector . The two vectors may not overlap . Since vectors '
  --   length is encoded in the type there is no need in runtime checks.
  copy :: PrimMonad m
       => v (PrimState m) a    -- ^ Target
       -> v (PrimState m) a    -- ^ Source
       -> m ()
  | Copy vector . The two vectors may overlap . Since vectors ' length
  --   is encoded in the type there is no need in runtime checks.
  move :: PrimMonad m
       => v (PrimState m) a    -- ^ Target
       -> v (PrimState m) a    -- ^ Source
       -> m ()
  -- | Allocate new vector
  new   :: PrimMonad m => m (v (PrimState m) a)
  -- | Read value at index without bound checks.
  unsafeRead  :: PrimMonad m => v (PrimState m) a -> Int -> m a
  -- | Write value at index without bound checks.
  unsafeWrite :: PrimMonad m => v (PrimState m) a -> Int -> a -> m ()


-- | Length of mutable vector. Function doesn't evaluate its argument.
lengthM :: forall v s a. (Arity (DimM v)) => v s a -> Int
lengthM _ = arity (Proxy :: Proxy (DimM v))

-- | Create copy of vector.
--
--   Examples:
--
    > > > import Control . Monad . ST ( runST )
--   >>> import Data.Vector.Fixed (mk3)
--   >>> import Data.Vector.Fixed.Boxed (Vec3)
--   >>> import qualified Data.Vector.Fixed.Mutable as M
    > > > let x = runST ( do { v < - M.replicate 100 ; v ' < - clone v ; M.write v ' 0 2 ; M.unsafeFreeze v ' } ) : :  
--   >>> x
    fromList [ 2,100,100 ]
clone :: (PrimMonad m, MVector v a) => v (PrimState m) a -> m (v (PrimState m) a)
# INLINE clone #
clone v = do
  u <- new
  move u v
  return u

-- | Read value at index with bound checks.
read  :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> m a
# INLINE read #
read v i
  | i < 0 || i >= lengthM v = error "Data.Vector.Fixed.Mutable.read: index out of range"
  | otherwise               = unsafeRead v i

-- | Write value at index with bound checks.
write :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> a -> m ()
# INLINE write #
write v i x
  | i < 0 || i >= lengthM v = error "Data.Vector.Fixed.Mutable.write: index out of range"
  | otherwise               = unsafeWrite v i x


-- | Create new vector with all elements set to given value.
replicate :: (PrimMonad m, MVector v a) => a -> m (v (PrimState m) a)
# INLINE replicate #
replicate a = do
  v <- new
  forI v $ \i -> unsafeWrite v i a
  pure v

-- | Create new vector with all elements are generated by provided
--   monadic action.
replicateM :: (PrimMonad m, MVector v a) => m a -> m (v (PrimState m) a)
# INLINE replicateM #
replicateM m = do
  v <- new
  forI v $ \i -> unsafeWrite v i =<< m
  pure v

-- | Create new vector with using function from index to value.
generate :: (PrimMonad m, MVector v a) => (Int -> a) -> m (v (PrimState m) a)
# INLINE generate #
generate f = do
  v <- new
  forI v $ \i -> unsafeWrite v i $ f i
  pure v

-- | Create new vector with using monadic function from index to value.
generateM :: (PrimMonad m, MVector v a) => (Int -> m a) -> m (v (PrimState m) a)
# INLINE generateM #
generateM f = do
  v <- new
  forI v $ \i -> unsafeWrite v i =<< f i
  pure v

-- | Loop which calls function for each index
forI :: (PrimMonad m, MVector v a) => v (PrimState m) a -> (Int -> m ()) -> m ()
# INLINE forI #
forI v f = go 0
  where
    go i | i >= n    = pure ()
         | otherwise = f i >> go (i+1)
    n = lengthM v


----------------------------------------------------------------
-- Immutable
----------------------------------------------------------------

-- | Type class for immutable vectors
class (Dim v ~ DimM (Mutable v), MVector (Mutable v) a) => IVector v a where
  -- | Convert vector to immutable state. Mutable vector must not be
  --   modified afterwards.
  unsafeFreeze :: PrimMonad m => Mutable v (PrimState m) a -> m (v a)
  -- | /O(1)/ Unsafely convert immutable vector to mutable without
  --   copying.  Note that this is a very dangerous function and
  --   generally it's only safe to read from the resulting vector. In
  --   this case, the immutable vector could be used safely as well.
  --
  Problems with mutation happen because GHC has a lot of freedom to
  -- introduce sharing. As a result mutable vectors produced by
  -- @unsafeThaw@ may or may not share the same underlying buffer. For
  -- example:
  --
  -- > foo = do
  >    let vec = F.generate 10 i d
  >        
  -- >   do_something mvec
  --
  Here GHC could lift @vec@ outside of foo which means that all calls to
  -- @do_something@ will use same buffer with possibly disastrous
  -- results. Whether such aliasing happens or not depends on the program in
  question , optimization levels , and GHC flags .
  --
  All in all , attempts to modify a vector produced by @unsafeThaw@
  -- fall out of domain of software engineering and into realm of
  -- black magic, dark rituals, and unspeakable horrors. The only
  -- advice that could be given is: "Don't attempt to mutate a vector
  produced by  unless you know how to prevent GHC from
  -- aliasing buffers accidentally. We don't."
  unsafeThaw   :: PrimMonad m => v a -> m (Mutable v (PrimState m) a)
  -- | Get element at specified index without bounds check.
  unsafeIndex :: v a -> Int -> a

index :: IVector v a => v a -> Int -> a
# INLINE index #
index v i | i < 0 || i >= length v = error "Data.Vector.Fixed.Mutable.!: index out of bounds"
          | otherwise              = unsafeIndex v i


-- | Safely convert mutable vector to immutable.
freeze :: (PrimMonad m, IVector v a) => Mutable v (PrimState m) a -> m (v a)
# INLINE freeze #
freeze v = unsafeFreeze =<< clone v

-- | Safely convert immutable vector to mutable.
thaw :: (PrimMonad m, IVector v a) => v a -> m (Mutable v (PrimState m) a)
# INLINE thaw #
thaw v = clone =<< unsafeThaw v



----------------------------------------------------------------
-- Vector API
----------------------------------------------------------------

-- | Generic inspect implementation for array-based vectors.
inspectVec :: forall v a b. (Arity (Dim v), IVector v a) => v a -> Fun (Peano (Dim v)) a b -> b
# INLINE inspectVec #
inspectVec v
  = inspect cv
  where
    cv :: ContVec (Dim v) a
    cv = apply (\(Const i) -> (unsafeIndex v i, Const (i+1)))
               (Const 0 :: Const Int (Peano (Dim v)))

-- | Generic construct implementation for array-based vectors.
constructVec :: forall v a. (Arity (Dim v), IVector v a) => Fun (Peano (Dim v)) a (v a)
# INLINE constructVec #
constructVec =
  accum step
        (\(T_new _ st) -> runST $ unsafeFreeze =<< st :: v a)
        (T_new 0 new :: T_new v a (Peano (Dim v)))

data T_new v a n = T_new Int (forall s. ST s (Mutable v s a))

step :: (IVector v a) => T_new v a ('S n) -> a -> T_new v a n
step (T_new i st) x = T_new (i+1) $ do
  mv <- st
  unsafeWrite mv i x
  return mv

# LANGUAGE DataKinds             #
# LANGUAGE FlexibleContexts      #
# LANGUAGE Rank2Types            #
# LANGUAGE ScopedTypeVariables   #
 |
 Type classes for vectors which are implemented on top of the arrays
 and support in-place mutation. API is similar to one used in the
 @vector@ package.
 * Mutable vectors
 * Creation
 * Loops
 * Immutable vectors
 * Vector API
--------------------------------------------------------------
 Type classes
--------------------------------------------------------------
 | Mutable counterpart of fixed-length vector.
 | Dimension for mutable vector.
 | Type class for mutable vectors.
   length is encoded in the type there is no need in runtime checks.
 ^ Target
 ^ Source
   is encoded in the type there is no need in runtime checks.
 ^ Target
 ^ Source
 | Allocate new vector
 | Read value at index without bound checks.
 | Write value at index without bound checks.
 | Length of mutable vector. Function doesn't evaluate its argument.
 | Create copy of vector.

   Examples:

   >>> import Data.Vector.Fixed (mk3)
   >>> import Data.Vector.Fixed.Boxed (Vec3)
   >>> import qualified Data.Vector.Fixed.Mutable as M
   >>> x
 | Read value at index with bound checks.
 | Write value at index with bound checks.
 | Create new vector with all elements set to given value.
 | Create new vector with all elements are generated by provided
   monadic action.
 | Create new vector with using function from index to value.
 | Create new vector with using monadic function from index to value.
 | Loop which calls function for each index
--------------------------------------------------------------
 Immutable
--------------------------------------------------------------
 | Type class for immutable vectors
 | Convert vector to immutable state. Mutable vector must not be
   modified afterwards.
 | /O(1)/ Unsafely convert immutable vector to mutable without
   copying.  Note that this is a very dangerous function and
   generally it's only safe to read from the resulting vector. In
   this case, the immutable vector could be used safely as well.

 introduce sharing. As a result mutable vectors produced by
 @unsafeThaw@ may or may not share the same underlying buffer. For
 example:

 > foo = do
 >   do_something mvec

 @do_something@ will use same buffer with possibly disastrous
 results. Whether such aliasing happens or not depends on the program in

 fall out of domain of software engineering and into realm of
 black magic, dark rituals, and unspeakable horrors. The only
 advice that could be given is: "Don't attempt to mutate a vector
 aliasing buffers accidentally. We don't."
 | Get element at specified index without bounds check.
 | Safely convert mutable vector to immutable.
 | Safely convert immutable vector to mutable.
--------------------------------------------------------------
 Vector API
--------------------------------------------------------------
 | Generic inspect implementation for array-based vectors.
 | Generic construct implementation for array-based vectors.
# LANGUAGE MultiParamTypeClasses #
# LANGUAGE PolyKinds              #
# LANGUAGE TypeFamilies           #
module Data.Vector.Fixed.Mutable (
    Arity
  , arity
  , Mutable
  , DimM
  , MVector(..)
  , lengthM
  , read
  , write
  , clone
  , replicate
  , replicateM
  , generate
  , generateM
  , forI
  , IVector(..)
  , index
  , freeze
  , thaw
  , constructVec
  , inspectVec
  ) where

import Control.Applicative  (Const(..))
import Control.Monad.ST
import Control.Monad.Primitive
import Data.Typeable  (Proxy(..))
import GHC.TypeLits
import Data.Vector.Fixed.Cont (Dim,PeanoNum(..),Peano,Arity,Fun(..),Vector(..),ContVec,arity,apply,accum,length)
import Prelude hiding (read,length,replicate)



type family Mutable (v :: * -> *) :: * -> * -> *

type family DimM (v :: * -> * -> *) :: Nat

class (Arity (DimM v)) => MVector v a where
  | Copy vector . The two vectors may not overlap . Since vectors '
  copy :: PrimMonad m
       -> m ()
  | Copy vector . The two vectors may overlap . Since vectors ' length
  move :: PrimMonad m
       -> m ()
  new   :: PrimMonad m => m (v (PrimState m) a)
  unsafeRead  :: PrimMonad m => v (PrimState m) a -> Int -> m a
  unsafeWrite :: PrimMonad m => v (PrimState m) a -> Int -> a -> m ()


lengthM :: forall v s a. (Arity (DimM v)) => v s a -> Int
lengthM _ = arity (Proxy :: Proxy (DimM v))

    > > > import Control . Monad . ST ( runST )
    > > > let x = runST ( do { v < - M.replicate 100 ; v ' < - clone v ; M.write v ' 0 2 ; M.unsafeFreeze v ' } ) : :  
    fromList [ 2,100,100 ]
clone :: (PrimMonad m, MVector v a) => v (PrimState m) a -> m (v (PrimState m) a)
# INLINE clone #
clone v = do
  u <- new
  move u v
  return u

read  :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> m a
# INLINE read #
read v i
  | i < 0 || i >= lengthM v = error "Data.Vector.Fixed.Mutable.read: index out of range"
  | otherwise               = unsafeRead v i

write :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> a -> m ()
# INLINE write #
write v i x
  | i < 0 || i >= lengthM v = error "Data.Vector.Fixed.Mutable.write: index out of range"
  | otherwise               = unsafeWrite v i x


replicate :: (PrimMonad m, MVector v a) => a -> m (v (PrimState m) a)
# INLINE replicate #
replicate a = do
  v <- new
  forI v $ \i -> unsafeWrite v i a
  pure v

replicateM :: (PrimMonad m, MVector v a) => m a -> m (v (PrimState m) a)
# INLINE replicateM #
replicateM m = do
  v <- new
  forI v $ \i -> unsafeWrite v i =<< m
  pure v

generate :: (PrimMonad m, MVector v a) => (Int -> a) -> m (v (PrimState m) a)
# INLINE generate #
generate f = do
  v <- new
  forI v $ \i -> unsafeWrite v i $ f i
  pure v

generateM :: (PrimMonad m, MVector v a) => (Int -> m a) -> m (v (PrimState m) a)
# INLINE generateM #
generateM f = do
  v <- new
  forI v $ \i -> unsafeWrite v i =<< f i
  pure v

forI :: (PrimMonad m, MVector v a) => v (PrimState m) a -> (Int -> m ()) -> m ()
# INLINE forI #
forI v f = go 0
  where
    go i | i >= n    = pure ()
         | otherwise = f i >> go (i+1)
    n = lengthM v



class (Dim v ~ DimM (Mutable v), MVector (Mutable v) a) => IVector v a where
  unsafeFreeze :: PrimMonad m => Mutable v (PrimState m) a -> m (v a)
  Problems with mutation happen because GHC has a lot of freedom to
  >    let vec = F.generate 10 i d
  >        
  Here GHC could lift @vec@ outside of foo which means that all calls to
  question , optimization levels , and GHC flags .
  All in all , attempts to modify a vector produced by @unsafeThaw@
  produced by  unless you know how to prevent GHC from
  unsafeThaw   :: PrimMonad m => v a -> m (Mutable v (PrimState m) a)
  unsafeIndex :: v a -> Int -> a

index :: IVector v a => v a -> Int -> a
# INLINE index #
index v i | i < 0 || i >= length v = error "Data.Vector.Fixed.Mutable.!: index out of bounds"
          | otherwise              = unsafeIndex v i


freeze :: (PrimMonad m, IVector v a) => Mutable v (PrimState m) a -> m (v a)
# INLINE freeze #
freeze v = unsafeFreeze =<< clone v

thaw :: (PrimMonad m, IVector v a) => v a -> m (Mutable v (PrimState m) a)
# INLINE thaw #
thaw v = clone =<< unsafeThaw v




inspectVec :: forall v a b. (Arity (Dim v), IVector v a) => v a -> Fun (Peano (Dim v)) a b -> b
# INLINE inspectVec #
inspectVec v
  = inspect cv
  where
    cv :: ContVec (Dim v) a
    cv = apply (\(Const i) -> (unsafeIndex v i, Const (i+1)))
               (Const 0 :: Const Int (Peano (Dim v)))

constructVec :: forall v a. (Arity (Dim v), IVector v a) => Fun (Peano (Dim v)) a (v a)
# INLINE constructVec #
constructVec =
  accum step
        (\(T_new _ st) -> runST $ unsafeFreeze =<< st :: v a)
        (T_new 0 new :: T_new v a (Peano (Dim v)))

data T_new v a n = T_new Int (forall s. ST s (Mutable v s a))

step :: (IVector v a) => T_new v a ('S n) -> a -> T_new v a n
step (T_new i st) x = T_new (i+1) $ do
  mv <- st
  unsafeWrite mv i x
  return mv
