{-# LANGUAGE CPP, ScopedTypeVariables, FlexibleInstances, TypeSynonymInstances,
  MultiParamTypeClasses, FlexibleContexts #-}
{- |
   Module     : Data.ArrayBZ.Internals.Unboxed
   Copyright  : (c) The University of Glasgow 2001 & (c) 2006 Bulat Ziganshin
   License    : BSD3

   Maintainer : Bulat Ziganshin <Bulat.Ziganshin@gmail.com>
   Stability  : experimental
   Portability: GHC/Hugs

Unboxed arrays

Based on the idea of Oleg Kiselyov
  (see http://www.haskell.org/pipermail/haskell-cafe/2004-July/006400.html)
-}

module Data.ArrayBZ.Internals.Unboxed where

import Control.Monad.ST         (ST, runST)
import Control.Monad.ST.Lazy    ( strictToLazyST )
import qualified Control.Monad.ST.Lazy as Lazy (ST)
import Data.Ix
import Data.Typeable
#include "Typeable.h"

import Control.Monad.STorIO
import Data.ArrayBZ.Internals.IArray
import Data.ArrayBZ.Internals.MArray
import Data.HasDefaultValue
import Data.Unboxed

-- ---------------------------------------------------------------------------
-- | Unboxed mutable arrays

data UnboxedMutableArray s i e  =  UMA !i !i !(MUVec s e)

instance HasBounds (UnboxedMutableArray s) where
    {-# INLINE bounds #-}
    bounds (UMA l u _) = (l,u)

instance (STorIO m s) => HasMutableBounds (UnboxedMutableArray s) m where
    {-# INLINE getBounds #-}
    getBounds (UMA l u _) = return (l,u)

instance (STorIO m s, Unboxed e) => MArray (UnboxedMutableArray s) e m where
    newArray_ (l,u) = do arr <- allocUnboxed (rangeSize (l,u))
                         return (UMA l u arr)
    {-# INLINE unsafeRead #-}
    unsafeRead  (UMA _ _ arr) = readUnboxed arr
    {-# INLINE unsafeWrite #-}
    unsafeWrite (UMA _ _ arr) = writeUnboxed arr

-- ---------------------------------------------------------------------------
-- | Unboxed mutable arrays in ST monad

type STUArray = UnboxedMutableArray

stUArrayTc :: TyCon
INSTANCE_TYPEABLE3(STUArray,stUArrayTc,"STUArray")

-- ---------------------------------------------------------------------------
-- STUArray also works in Lazy ST monad

instance (Unboxed e) => MArray (STUArray s) e (Lazy.ST s) where
    {-# INLINE newArray_ #-}
    newArray_   (l,u)   = strictToLazyST (newArray_ (l,u))
    {-# INLINE unsafeRead #-}
    unsafeRead  arr i   = strictToLazyST (unsafeRead  arr i)
    {-# INLINE unsafeWrite #-}
    unsafeWrite arr i e = strictToLazyST (unsafeWrite arr i e)

-- ---------------------------------------------------------------------------
-- | Unboxed mutable arrays in IO monad

type IOUArray = IOSpecific3 UnboxedMutableArray

iOUArrayTc :: TyCon
INSTANCE_TYPEABLE2(IOUArray,iOUArrayTc,"IOUArray")

-- ---------------------------------------------------------------------------
-- | Unboxed arrays

data UArray i e  =  UA !i !i !(UVec e)

uArrayTc :: TyCon
INSTANCE_TYPEABLE2(UArray,uArrayTc,"UArray")

instance HasBounds UArray where
    {-# INLINE bounds #-}
    bounds (UA l u _) = (l,u)

instance (Unboxed e, HasDefaultValue e) => IArray UArray e where
    {-# INLINE unsafeArray #-}
    -- Create new array filled with (i,e) values
    unsafeArray lu ies = runST (withNewArray lu defaultValue (doReplace ies))
    {-# INLINE unsafeAt #-}
    unsafeAt (UA _ _ arr) = indexUnboxed arr
    {-# INLINE unsafeReplace #-}
    -- Make a copy of array and perform (i,e) replacements
    unsafeReplace arr ies = runST (withArrayCopy arr (doReplace ies))
    {-# INLINE unsafeAccum #-}
    -- Make a copy of array and perform (i,e) accumulation in new array
    unsafeAccum f arr ies = runST (withArrayCopy arr $ doAccum f ies)
    {-# INLINE unsafeAccumArray #-}
    -- Create new array accumulating (i,e) values
    unsafeAccumArray f int lu ies = runST (withNewArray lu int $ doAccum f ies)


-- Implementation helper functions -------------

-- Create new array and perform given action on it before freezing
withNewArray :: (STorIO t t2,
                       Ix i,
                          MArray (UnboxedMutableArray t2) e t) =>
               (i, i)
                   -> e
                   -> (UnboxedMutableArray t2 i e -> t t1)
                   -> t (UArray i e)
withNewArray lu int action = do
    marr <- newArray lu int
    action marr
    unsafeFreezeUA marr

-- Make a copy of array and perform given action on it before freezing
withArrayCopy :: (Ix t, Unboxed t1, STorIO t2 s) =>
                                             UArray t t1
                                             -> (UnboxedMutableArray s t t1 -> t2 t3)
                                             -> t2 (UArray t t1)
withArrayCopy arr action = do
    marr <- thawUA arr
    action marr
    unsafeFreezeUA marr

-- Perform (i,e) replaces in mutable array
doReplace :: (Ix i, MArray a e t) => [(Int, e)] -> a i e -> t ()
doReplace ies marr = do
    sequence_ [unsafeWrite marr i e | (i, e) <- ies]

-- Accumulate (i,e) values in mutable array
doAccum :: (Ix i, MArray a t1 t2) => (t1 -> t -> t1) -> [(Int, t)] -> a i t1 -> t2 ()
doAccum f ies marr = do
    sequence_ [do old <- unsafeRead marr i
                  unsafeWrite marr i (f old new)
              | (i, new) <- ies]

-- Mutable->immutable array conversion which takes a copy of contents
freezeUA :: (STorIO t3 t, Unboxed t2, Ix t1) => UnboxedMutableArray t t1 t2 -> t3 (UArray t1 t2)
freezeUA uma@(UMA l u marr) = do
    arr <- freezeUnboxed marr (sizeOfUMA uma)
    return (UA l u arr)

-- Immutable->mutable array conversion which takes a copy of contents
thawUA :: (STorIO t2 s, Unboxed t1, Ix t) => UArray t t1 -> t2 (UnboxedMutableArray s t t1)
thawUA ua@(UA l u arr) = do
    marr <- thawUnboxed arr (sizeOfUA ua)
    return (UMA l u marr)

-- On-the-place mutable->immutable array conversion
unsafeFreezeUA :: (STorIO t3 t1) => UnboxedMutableArray t1 t t2 -> t3 (UArray t t2)
unsafeFreezeUA (UMA l u marr) = do
    arr <- unsafeFreezeUnboxed marr
    return (UA l u arr)

-- On-the-place immutable->mutable array conversion
unsafeThawUA :: (STorIO t2 s) => UArray t t1 -> t2 (UnboxedMutableArray s t t1)
unsafeThawUA (UA l u arr) = do
    marr <- unsafeThawUnboxed arr
    return (UMA l u marr)

-- | Array size in bytes
sizeOfUA :: forall i e. (Ix i, Unboxed e) => UArray i e -> Int
sizeOfUA  arr =
     rangeSize (bounds arr) * sizeOfUnboxed (undefined :: e)

sizeOfUMA :: forall i e s. (Ix i, Unboxed e) => UnboxedMutableArray s i e -> Int
sizeOfUMA marr =
    rangeSize (bounds marr) * sizeOfUnboxed (undefined :: e)

-- ---------------------------------------------------------------------------
-- | Freeze/thaw rules for IOUArray

freezeIOUArray       :: (Unboxed e, HasDefaultValue e, Ix i) => IOUArray i e -> IO (UArray i e)
thawIOUArray         :: (Unboxed e, HasDefaultValue e, Ix i) => UArray i e -> IO (IOUArray i e)
unsafeFreezeIOUArray :: (Unboxed e, HasDefaultValue e, Ix i) => IOUArray i e -> IO (UArray i e)
unsafeThawIOUArray   :: (Unboxed e, HasDefaultValue e, Ix i) => UArray i e -> IO (IOUArray i e)

freezeIOUArray       = freezeUA
thawIOUArray         = thawUA
unsafeFreezeIOUArray = unsafeFreezeUA
unsafeThawIOUArray   = unsafeThawUA

{-# RULES
"freeze/IOUArray" forall (x :: (forall s e i . (Unboxed e, HasDefaultValue e) => IOUArray i e)) . freeze x = freezeIOUArray x
"thaw/IOUArray"   forall (x :: (forall   e i . (Unboxed e, HasDefaultValue e) =>   UArray i e)) . thaw   x = thawIOUArray   x
"unsafeFreeze/IOUArray" forall (x :: (forall s e i . (Unboxed e, HasDefaultValue e) => IOUArray i e)) . unsafeFreeze x = unsafeFreezeIOUArray x
"unsafeThaw/IOUArray"   forall (x :: (forall   e i . (Unboxed e, HasDefaultValue e) =>   UArray i e)) . unsafeThaw   x = unsafeThawIOUArray   x
    #-}

-- ---------------------------------------------------------------------------
-- | Freeze/thaw rules for STUArray

freezeSTUArray       :: (Unboxed e, HasDefaultValue e, Ix i) => STUArray s i e -> ST s (UArray i e)
thawSTUArray         :: (Unboxed e, HasDefaultValue e, Ix i) => UArray i e -> ST s (STUArray s i e)
unsafeFreezeSTUArray :: (Unboxed e, HasDefaultValue e, Ix i) => STUArray s i e -> ST s (UArray i e)
unsafeThawSTUArray   :: (Unboxed e, HasDefaultValue e, Ix i) => UArray i e -> ST s (STUArray s i e)

freezeSTUArray       = freezeUA
thawSTUArray         = thawUA
unsafeFreezeSTUArray = unsafeFreezeUA
unsafeThawSTUArray   = unsafeThawUA

{-# RULES
"freeze/STUArray" forall (x :: (forall s e i . (Unboxed e, HasDefaultValue e) => STUArray s i e)) . freeze x = freezeSTUArray x
"thaw/STUArray"   forall (x :: (forall   e i . (Unboxed e, HasDefaultValue e) =>   UArray   i e)) . thaw   x = thawSTUArray   x
"unsafeFreeze/STUArray" forall (x :: (forall s e i . (Unboxed e, HasDefaultValue e) => STUArray s i e)) . unsafeFreeze x = unsafeFreezeSTUArray x
"unsafeThaw/STUArray"   forall (x :: (forall   e i . (Unboxed e, HasDefaultValue e) =>   UArray   i e)) . unsafeThaw   x = unsafeThawSTUArray   x
    #-}

-- ---------------------------------------------------------------------------
-- | Casts to arrays with different element type

-- | Casts an 'UArray' with one element type into 'UArray' with a
-- different element type. All the elements of the resulting array
-- are undefined (unless you know what you\'re doing...).
-- Upper array bound is recalculated according to elements size,
-- for example UArray (1,2) Word32 -> UArray (1,8) Word8
castUArray :: forall i e e'. (Ix i, Enum i, Unboxed e, Unboxed e')
           => UArray i e -> UArray i e'
castUArray (UA l u vec) =
    (UA l u' (castUnboxed vec))
        where u' = toEnum (fromEnum l - 1 + newSize)
              newSize = rangeSize (l,u)   *   sizeOfUnboxed (undefined::e)
                                        `div` sizeOfUnboxed (undefined::e')

-- | Casts an 'IOUArray' with one element type into 'IOUArray' with a different
-- element type (upper bound is recalculated).
castIOUArray :: forall i e e'. (Ix i, Enum i, Unboxed e, Unboxed e')
             => IOUArray i e -> IOUArray i e'
castIOUArray (UMA l u mvec) =
    (UMA l u' (castMUnboxed mvec))
        where u' = toEnum (fromEnum l - 1 + newSize)
              newSize = rangeSize (l,u)   *   sizeOfUnboxed (undefined::e)
                                        `div` sizeOfUnboxed (undefined::e')

-- | Casts an 'STUArray' with one element type into 'STUArray' with a different
-- element type (upper bound is recalculated).
castSTUArray :: forall i e e' s. (Ix i, Enum i, Unboxed e, Unboxed e')
             => STUArray s i e -> STUArray s i e'
castSTUArray (UMA l u mvec) =
    (UMA l u' (castMUnboxed mvec))
        where u' = toEnum (fromEnum l - 1 + newSize)
              newSize = rangeSize (l,u)   *   sizeOfUnboxed (undefined::e)
                                        `div` sizeOfUnboxed (undefined::e')

-- ---------------------------------------------------------------------------
-- | Instances

instance (Ix i, Show i, Show e, Unboxed e, HasDefaultValue e) => Show (UArray i e) where
    showsPrec = showsIArray

instance (Ix i, Eq i, Eq e, Unboxed e, HasDefaultValue e) => Eq (UArray i e) where
    (==) = eqIArray

instance (Ix i, Ord i, Ord e, Unboxed e, HasDefaultValue e) => Ord (UArray i e) where
    compare = cmpIArray