{-# LANGUAGE CPP, TypeSynonymInstances, FlexibleInstances,
  MultiParamTypeClasses, FlexibleContexts #-}
{- |
   Module     : Data.ArrayBZ.Internals.Boxed
   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

Boxed arrays

-}

module Data.ArrayBZ.Internals.Boxed
#ifdef __HUGS__
         ( -- * Types
           Array,
           IOArray,
           STArray,

           -- * Freeze/thaw operations
           freezeIOArray,
           thawIOArray,
           unsafeFreezeIOArray,
           freezeSTArray,
           thawSTArray,
           unsafeFreezeSTArray,
         )
#endif
   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 Data.ArrayBZ.Internals.IArray
import Data.ArrayBZ.Internals.MArray

#ifdef __HUGS__
-- ---------------------------------------------------------------------------
-- Hugs primitives are higher-level than GHC/NHC's
-- ---------------------------------------------------------------------------

import Hugs.Array as Arr
import Hugs.IOArray
import Hugs.ST

-----------------------------------------------------------------------------
-- Normal polymorphic arrays

instance HasBounds Array where
    bounds           = Arr.bounds

instance IArray Array e where
    unsafeArray      = Arr.unsafeArray
    unsafeAt         = Arr.unsafeAt
    unsafeReplace    = Arr.unsafeReplace
    unsafeAccum      = Arr.unsafeAccum
    unsafeAccumArray = Arr.unsafeAccumArray

-----------------------------------------------------------------------------
-- | Instance declarations for 'IOArray's

instance HasBounds IOArray where
    bounds      = boundsIOArray

instance HasMutableBounds IOArray IO where
    getBounds   = return . boundsIOArray

instance MArray IOArray e IO where
    newArray    = newIOArray
    unsafeRead  = unsafeReadIOArray
    unsafeWrite = unsafeWriteIOArray

-----------------------------------------------------------------------------
-- Polymorphic non-strict mutable arrays (ST monad)

instance HasBounds (STArray s) where
    bounds      = boundsSTArray

instance HasMutableBounds (STArray s) (ST s) where
    getBounds   = return . boundsSTArray

instance MArray (STArray s) e (ST s) where
    newArray    = newSTArray
    unsafeRead  = unsafeReadSTArray
    unsafeWrite = unsafeWriteSTArray

#else
-- ---------------------------------------------------------------------------
-- Non-Hugs implementation
-- ---------------------------------------------------------------------------

import Control.Monad.STorIO
import GHC.ArrBZ

-- ---------------------------------------------------------------------------
-- | Boxed mutable arrays

data BoxedMutableArray s i e  =  BMA !i !i !(MVec s e)

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

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

instance (STorIO m s) => MArray (BoxedMutableArray s) e m where
    newArray (l,u) initial = do arr <- allocBoxed (rangeSize (l,u)) initial
                                return (BMA l u arr)
    {-# INLINE unsafeRead #-}
    unsafeRead  (BMA _ _ arr) =  readBoxed  arr
    {-# INLINE unsafeWrite #-}
    unsafeWrite (BMA _ _ arr) =  writeBoxed arr

-- ---------------------------------------------------------------------------
-- | Boxed mutable arrays in ST monad

type STArray = BoxedMutableArray

-- ---------------------------------------------------------------------------
-- STArray also works in Lazy ST monad

instance MArray (STArray s) e (Lazy.ST s) where
    {-# INLINE newArray #-}
    newArray (l,u) initial = strictToLazyST (newArray (l,u) initial)
    {-# INLINE unsafeRead #-}
    unsafeRead  arr i   = strictToLazyST (unsafeRead  arr i)
    {-# INLINE unsafeWrite #-}
    unsafeWrite arr i e = strictToLazyST (unsafeWrite arr i e)

-- ---------------------------------------------------------------------------
-- | Boxed mutable arrays in IO monad

type IOArray = IOSpecific3 BoxedMutableArray

-- ---------------------------------------------------------------------------
-- | Boxed immutable arrays

data Array i e  =  BA !i !i !(Vec e)

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

instance IArray Array e where
    {-# INLINE unsafeArray #-}
    -- Create new array filled with (i,e) values
    unsafeArray lu ies = runST (withNewArray lu arrEleBottom (doReplace ies))
    {-# INLINE unsafeAt #-}
    unsafeAt (BA _ _ arr) = indexBoxed 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 initial lu ies = runST (withNewArray lu initial (doAccum f ies))


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

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

-- Make a copy of array and perform given action on it before freezing
withArrayCopy ::                              (Ix t, STorIO t2 s) =>
                                             Array t t1
                                             -> (BoxedMutableArray s t t1 -> t2 t3)
                                             -> t2 (Array t t1)
withArrayCopy arr action = do
    marr <- thawBA arr
    action marr
    unsafeFreezeBA 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
freezeBA :: (STorIO t3 t, Ix t1) => BoxedMutableArray t t1 t2 -> t3 (Array t1 t2)
freezeBA ua@(BMA l u marr) = do
    arr <- freezeBoxed marr (sizeOfBA ua) arrEleBottom
    return (BA l u arr)

-- Immutable->mutable array conversion which takes a copy of contents
thawBA :: (STorIO t2 s, Ix t) => Array t t1 -> t2 (BoxedMutableArray s t t1)
thawBA ua@(BA l u arr) = do
    marr <- thawBoxed arr (sizeOfBA ua) arrEleBottom
    return (BMA l u marr)

-- On-the-place mutable->immutable array conversion
unsafeFreezeBA :: (STorIO t3 t1) => BoxedMutableArray t1 t t2 -> t3 (Array t t2)
unsafeFreezeBA (BMA l u marr) = do
    arr <- unsafeFreezeBoxed marr
    return (BA l u arr)

-- On-the-place immutable->mutable array conversion
unsafeThawBA :: (STorIO t2 s) => Array t t1 -> t2 (BoxedMutableArray s t t1)
unsafeThawBA (BA l u arr) = do
    marr <- unsafeThawBoxed arr
    return (BMA l u marr)

-- | Number of array elements
sizeOfBA :: (Ix a1, HasBounds a) => a a1 e -> Int
sizeOfBA arr = rangeSize (bounds arr)

-- ---------------------------------------------------------------------------
-- | Freeze/thaw rules for IOArray

freezeIOArray       :: (Ix i) => IOArray i e -> IO (Array i e)
thawIOArray         :: (Ix i) => Array i e -> IO (IOArray i e)
unsafeFreezeIOArray :: (Ix i) => IOArray i e -> IO (Array i e)
unsafeThawIOArray   :: (Ix i) => Array i e -> IO (IOArray i e)

freezeIOArray       = freezeBA
thawIOArray         = thawBA
unsafeFreezeIOArray = unsafeFreezeBA
unsafeThawIOArray   = unsafeThawBA

{-# RULES
"freeze/IOArray" freeze = freezeIOArray
"thaw/IOArray"   thaw   = thawIOArray
"unsafeFreeze/IOArray" unsafeFreeze = unsafeFreezeIOArray
"unsafeThaw/IOArray"   unsafeThaw   = unsafeThawIOArray
    #-}

-- ---------------------------------------------------------------------------
-- | Freeze/thaw rules for STArray

freezeSTArray       :: (Ix i) => STArray s i e -> ST s (Array i e)
thawSTArray         :: (Ix i) => Array i e -> ST s (STArray s i e)
unsafeFreezeSTArray :: (Ix i) => STArray s i e -> ST s (Array i e)
unsafeThawSTArray   :: (Ix i) => Array i e -> ST s (STArray s i e)

freezeSTArray       = freezeBA
thawSTArray         = thawBA
unsafeFreezeSTArray = unsafeFreezeBA
unsafeThawSTArray   = unsafeThawBA

{-# RULES
"freeze/STArray" freeze = freezeSTArray
"thaw/STArray"   thaw   = thawSTArray
"unsafeFreeze/STArray" unsafeFreeze = unsafeFreezeSTArray
"unsafeThaw/STArray"   unsafeThaw   = unsafeThawSTArray
    #-}

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

instance (Ix i, Show i, Show e) => Show (Array i e) where
    showsPrec = showsIArray

instance (Ix i, Eq i, Eq e) => Eq (Array i e) where
    (==) = eqIArray

instance (Ix i, Ord i, Ord e) => Ord (Array i e) where
    compare = cmpIArray

#endif

INSTANCE_TYPEABLE2(IOArray,iOArrayTc,"IOArray")

INSTANCE_TYPEABLE3(STArray,stArrayTc,"STArray")