{-# OPTIONS_GHC -cpp -fglasgow-exts #-}
{- |
   Module     : GHC.Unboxed
   Copyright  : Copyright (C) 2006 Bulat Ziganshin
   License    : BSD3

   Maintainer : Bulat Ziganshin <bulatz@HotPOP.com>
   Stability  : experimental
   Portability: GHC

Unboxed values (simple datatypes that can be stored in ByteArrays,
  i.e. raw memory buffers allocated inside the Haskell heap)
Unboxed references
Unboxed arrays

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

module GHC.Unboxed where

import Data.Array.Base
import Data.Ix
import Foreign.Storable

import GHC.ST           ( ST(..), runST )
import GHC.IOBase       ( IO(..) )
import GHC.Base
import GHC.Word         ( Word(..) )
import GHC.Ptr          ( Ptr(..), FunPtr(..), nullPtr, nullFunPtr )
import GHC.Float        ( Float(..), Double(..) )
import GHC.Stable       ( StablePtr(..) )
import GHC.Int          ( Int8(..),  Int16(..),  Int32(..),  Int64(..) )
import GHC.Word         ( Word8(..), Word16(..), Word32(..), Word64(..) )

-- ---------------------------------------------------------------------------
-- | That's all we need to unify ST and IO operations!

class (Monad m) => STorIO m s | m->s where
    mLift :: (State# s -> (# State# s, a #)) -> m a

instance STorIO (ST s) s where
    {-# INLINE mLift #-}
    mLift = ST

instance STorIO IO RealWorld where
    {-# INLINE mLift #-}
    mLift = IO

-- | Type functions which converts universal ST/IO types to IO-specific ones
type IOSpecific  (a :: * -> *)           = a RealWorld
type IOSpecific2 (a :: * -> * -> *)      = a RealWorld
type IOSpecific3 (a :: * -> * -> * -> *) = a RealWorld

-- ---------------------------------------------------------------------------
-- | Immutable and mutable byte arrays

data BA    a  =  BA ByteArray#
data MBA s a  =  MBA (MutableByteArray# s)

-- | Alloc the mutable byte array
allocMBA :: (STorIO m s, Integral bytes, Unboxed a)
         => bytes -> m (MBA s a)
allocMBA bytes = mLift ( \s ->
  case newByteArray# (fromI# bytes) s of { (# s, arr #) ->
  (# s, MBA arr #) } )

-- | Alloc the mutable byte array having `elems` elements of required type
allocUnboxed :: forall m s elems a . (STorIO m s, Integral elems, Unboxed a)
               => elems -> m (MBA s a)
allocUnboxed elems = allocMBA (fromIntegral elems * sizeOfUnboxed (undefined::a))

-- | Mutable->immutable byte array on-place conversion
unsafeFreezeMBA :: (STorIO m s) => MBA s a -> m (BA a)
unsafeFreezeMBA (MBA mba) = mLift ( \s ->
  case unsafeFreezeByteArray# mba s of { (# s, ba #) ->
  (# s, BA ba #) } )

-- | Immutable->mutable byte array on-place conversion
unsafeThawBA :: (STorIO m s) => BA a -> m (MBA s a)
unsafeThawBA (BA ba) = mLift ( \s ->
  (# s, MBA (unsafeCoerce# ba) #) )

-- | Immutable->mutable byte array conversion which takes a copy of contents
thawBA :: (STorIO m s) => BA a -> m (MBA s a)
thawBA (BA arr#) = mLift ( \s1# ->
    case sizeofByteArray# arr#          of { n# ->
    case newByteArray# n# s1#           of { (# s2#, marr# #) ->
    case unsafeCoerce# memcpy marr# arr# n# s2# of { (# s3#, () #) ->
    (# s3#, MBA marr# #) }}} )

{-# INLINE fromI# #-}
-- Implementation helper function that converts any integral value to the Int#
fromI# :: (Integral n) => n -> Int#
fromI# n = n#   where I# n# = fromIntegral n

-- ---------------------------------------------------------------------------
-- | Unboxed is like Storable, but values are stored in byte arrays (i.e. inside the Haskell heap)

class Unboxed value where
    -- | Read the value from mutable byte array at given `index`
    readUnboxed    :: (STorIO m s, Integral index) => MBA s value -> index -> m value
    -- | Write the value to mutable byte array at given `index`
    writeUnboxed   :: (STorIO m s, Integral index) => MBA s value -> index -> value -> m ()
    -- | Read the value from immutable byte array at given `index`
    indexUnboxed   :: (Integral index) => BA value -> index -> value
    -- | Default value for this type
    default_value  :: value
    -- | How many bytes required to represent values of this type
    sizeOfUnboxed  :: value -> Int

-- Universal defition for Enum types having <= 256 variants
instance Unboxed Bool where
{
    {-# INLINE readUnboxed #-};
    readUnboxed (MBA arr) index  =  mLift ( \s ->
        case readInt8Array# arr (fromI# index) s of { (# s, value# #) ->
        (# s, tagToEnum# value# #) } );

    {-# INLINE writeUnboxed #-};
    writeUnboxed (MBA arr) index value  =  mLift ( \s ->
        case writeInt8Array# arr (fromI# index) (getTag value) s of { s ->
        (# s, () #) } );

    {-# INLINE indexUnboxed #-};
    indexUnboxed (BA arr) index  =  tagToEnum# (indexInt8Array# arr (fromI# index));

    {-# INLINE default_value #-};
    default_value = tagToEnum# 0#;

    {-# INLINE sizeOfUnboxed #-};
    sizeOfUnboxed _ = 1;
}

-- Universal defition for Storable types
#define InstanceUnboxed(type, cast, zero, read, write, at)            \
instance Unboxed type where                                           \
{                                                                     \
    {-# INLINE readUnboxed #-};                                       \
    readUnboxed (MBA arr) index  =  mLift ( \s ->                     \
        case read arr (fromI# index) s of { (# s, value# #) ->        \
        (# s, cast value# #) } );                                     \
                                                                      \
    {-# INLINE writeUnboxed #-};                                      \
    writeUnboxed (MBA arr) index (cast value#)  =  mLift ( \s ->      \
        case write arr (fromI# index) value# s of { s ->              \
        (# s, () #) } );                                              \
                                                                      \
    {-# INLINE indexUnboxed #-};                                      \
    indexUnboxed (BA arr) index  =  cast (at arr (fromI# index));     \
                                                                      \
    {-# INLINE default_value #-};                                     \
    default_value = zero;                                             \
                                                                      \
    {-# INLINE sizeOfUnboxed #-};                                     \
    sizeOfUnboxed = sizeOf;                                           \
}                                                                     \

InstanceUnboxed( Char,   C#, ' ', readWideCharArray#, writeWideCharArray#, indexWideCharArray#)
InstanceUnboxed( Int,    I#,   0, readIntArray#,      writeIntArray#,      indexIntArray#)
InstanceUnboxed( Int8,   I8#,  0, readInt8Array#,     writeInt8Array#,     indexInt8Array#)
InstanceUnboxed( Int16,  I16#, 0, readInt16Array#,    writeInt16Array#,    indexInt16Array#)
InstanceUnboxed( Int32,  I32#, 0, readInt32Array#,    writeInt32Array#,    indexInt32Array#)
InstanceUnboxed( Int64,  I64#, 0, readInt64Array#,    writeInt64Array#,    indexInt64Array#)
InstanceUnboxed( Word,   W#,   0, readWordArray#,     writeWordArray#,     indexWordArray#)
InstanceUnboxed( Word8,  W8#,  0, readWord8Array#,    writeWord8Array#,    indexWord8Array#)
InstanceUnboxed( Word16, W16#, 0, readWord16Array#,   writeWord16Array#,   indexWord16Array#)
InstanceUnboxed( Word32, W32#, 0, readWord32Array#,   writeWord32Array#,   indexWord32Array#)
InstanceUnboxed( Word64, W64#, 0, readWord64Array#,   writeWord64Array#,   indexWord64Array#)
InstanceUnboxed( Float,  F#,   0, readFloatArray#,    writeFloatArray#,    indexFloatArray#)
InstanceUnboxed( Double, D#,   0, readDoubleArray#,   writeDoubleArray#,   indexDoubleArray#)
InstanceUnboxed( (Ptr a),    Ptr,    nullPtr,    readAddrArray#, writeAddrArray#, indexAddrArray#)
InstanceUnboxed( (FunPtr a), FunPtr, nullFunPtr, readAddrArray#, writeAddrArray#, indexAddrArray#)
InstanceUnboxed( (StablePtr a), StablePtr, nullStablePtr, readStablePtrArray#, writeStablePtrArray#, indexStablePtrArray#)

-- ---------------------------------------------------------------------------
-- | Unboxed references (single-element arrays)

newtype URef s e  =  URef (MBA s e)

{-# INLINE newURef #-}
-- | Create new unboxed reference
newURef :: (STorIO m s, Unboxed e) => e -> m (URef s e)
newURef init = do var <- allocUnboxed 1
                  writeUnboxed var 0 init
                  return (URef var)

{-# INLINE readURef #-}
-- | Read current value of unboxed reference
readURef :: (STorIO m s, Unboxed e) => URef s e -> m e
readURef  (URef var)  =  readUnboxed  var 0

{-# INLINE writeURef #-}
-- | Change value of unboxed reference
writeURef :: (STorIO m s, Unboxed e) => URef s e -> e -> m ()
writeURef (URef var)  =  writeUnboxed var 0

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

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

instance HasBounds (UnboxedMutableArray s) where
    {-# INLINE bounds #-}
    bounds (UMA l u _) = (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) index  =  readUnboxed  arr index
    {-# INLINE unsafeWrite #-}
    unsafeWrite (UMA _ _ arr) index  =  writeUnboxed arr index

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

type UnboxedSTArray = UnboxedMutableArray

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

type UnboxedIOArray i e = IOSpecific3 UnboxedMutableArray i e

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

data UnboxedArray i e  =  UA !i !i !(BA e)

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

instance (Unboxed e) => IArray UnboxedArray e where
    {-# INLINE unsafeArray #-}
    unsafeArray lu ies = runST (unsafeArrayUA lu ies)
    {-# INLINE unsafeAt #-}
    unsafeAt (UA _ _ arr) index = indexUnboxed arr index
    {-# INLINE unsafeReplace #-}
    unsafeReplace arr ies = runST (unsafeReplaceUA arr ies)
    {-# INLINE unsafeAccum #-}
    unsafeAccum f arr ies = runST (unsafeAccumUA f arr ies)
    {-# INLINE unsafeAccumArray #-}
    unsafeAccumArray f init lu ies = runST (unsafeAccumArrayUA f init lu ies)

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

{-# INLINE unsafeArrayUA #-}
unsafeArrayUA :: (STorIO m s, Unboxed e, Ix i)
              => (i,i) -> [(Int, e)] -> m (UnboxedArray i e)
unsafeArrayUA (l,u) ies = do
    marr <- newArray (l,u) default_value
    sequence_ [unsafeWrite marr i e | (i, e) <- ies]
    unsafeFreezeUMA marr

unsafeFreezeUMA :: (STorIO m s)
                => UnboxedMutableArray s i e -> m (UnboxedArray i e)
unsafeFreezeUMA (UMA l u mba) = do
    ba <- unsafeFreezeMBA mba
    return (UA l u ba)

thawUA :: (STorIO m s)
       => UnboxedArray i e -> m (UnboxedMutableArray s i e)
thawUA (UA l u ba) = do
    mba <- thawBA ba
    return (UMA l u mba)

{-# RULES
"thaw/UnboxedSTArray" thaw = thawUA :: UnboxedArray i e -> ST s (UnboxedSTArray s i e)
"thaw/UnboxedIOArray" thaw = thawUA :: UnboxedArray i e -> IO   (UnboxedIOArray   i e)
    #-}

{-# INLINE unsafeReplaceUA #-}
unsafeReplaceUA :: (STorIO m s, Unboxed e, Ix i)
                => UnboxedArray i e -> [(Int, e)] -> m (UnboxedArray i e)
unsafeReplaceUA arr ies = do
    marr <- thawUA arr
    sequence_ [unsafeWrite marr i e | (i, e) <- ies]
    unsafeFreezeUMA marr

{-# INLINE unsafeAccumUA #-}
unsafeAccumUA :: (STorIO m s, Unboxed e, Ix i)
              => (e -> e' -> e) -> UnboxedArray i e -> [(Int, e')] -> m (UnboxedArray i e)
unsafeAccumUA f arr ies = do
    marr <- thawUA arr
    sequence_ [do
        old <- unsafeRead marr i
        unsafeWrite marr i (f old new)
        | (i, new) <- ies]
    unsafeFreezeUMA marr

{-# INLINE unsafeAccumArrayUA #-}
unsafeAccumArrayUA :: (STorIO m s, Unboxed e, Ix i)
                   => (e -> e' -> e) -> e -> (i,i) -> [(Int, e')] -> m (UnboxedArray i e)
unsafeAccumArrayUA f init (l,u) ies = do
    marr <- newArray (l,u) init
    sequence_ [do
        old <- unsafeRead marr i
        unsafeWrite marr i (f old new)
        | (i, new) <- ies]
    unsafeFreezeUMA marr

-- ---------------------------------------------------------------------------
-- Examples of using unboxed references and arrays
{-
main = do var <- newURef (31::Int)
          x <- readURef var
          print x

          var <- newURef True
          x <- readURef var
          print x

          arr <- newListArray (1,10) [11..20] :: IO (UnboxedIOArray Int Int)
          readArray arr 1 >>= print
          writeArray arr 1 333
          readArray arr 1 >>= print

          let arr = listArray (1,10) [11..20] :: UnboxedArray Int Int
          print (arr!1, arr!3)
-}