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

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

Unboxed arrays

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

module Data.Array_.Internals.Unboxed where

import Control.Monad.ST         (ST, runST)
import Data.Array.Base
import Data.Ix

import Data.HasDefaultValue
import Data.Unboxed

-- ---------------------------------------------------------------------------
-- | 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, HasDefaultValue e) => IArray UnboxedArray 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) index = indexUnboxed arr index
    {-# 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 init lu ies = runST (withNewArray lu init (doAccum f ies))


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

-- Create new array and perform given action on it before freezing
withNewArray lu init action = do
    marr <- newArray lu init
    action marr
    unsafeFreezeUMA marr

-- Make a copy of array and perform given action on it before freezing
withArrayCopy arr action = do
    marr <- thawUA arr
    action marr
    unsafeFreezeUMA marr

-- Perform (i,e) replaces in mutable array
doReplace ies marr = do
    sequence_ [unsafeWrite marr i e | (i, e) <- ies]

-- Accumulate (i,e) values in mutable array
doAccum f ies marr = do
    sequence_ [do
        old <- unsafeRead marr i
        unsafeWrite marr i (f old new)
        | (i, new) <- ies]

-- On-the-place mutable->immutable array conversion
unsafeFreezeUMA (UMA l u mba) = do
    ba <- unsafeFreezeMBA mba
    return (UA l u ba)

-- Immutable->mutable array conversion which takes a copy of contents
thawUA ua@(UA l u ba) = do
    mba <- thawBA ba (sizeOfUA ua)
    return (UMA l u mba)

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

-- | Array size in bytes
sizeOfUA (arr :: UnboxedArray i e) =
    rangeSize (bounds arr) * sizeOfUnboxed (undefined::e)