{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE MultiParamTypeClasses #-}

-- | PrimitiveArray with 'Ix' keys.

module Data.PrimitiveArray.Ix where

import Data.Primitive
import Control.Monad.ST
import Control.Monad (liftM, forM_, zipWithM_, forM)
import Data.Primitive.Types
import qualified GHC.Arr as A

import Data.PrimitiveArray
import Data.PrimitiveArray.Internal



instance (Bounded a, A.Ix a, Prim b) => PrimArrayOps a b where
  data PrimArray a b = PaIxP {-# UNPACK #-} !a {-# UNPACK #-} !a {-# UNPACK #-} !ByteArray
  unsafeIndex (PaIxP lb ub arr) i = indexByteArray arr (A.unsafeIndex (lb,ub) i)
  assocs (pa@(PaIxP lb ub _)) = [(i, unsafeIndex pa i) | i<-A.range(lb,ub) ]
  fromAssocs l u z xs = runST $ do
    pam <- fromAssocsM l u z xs
    unsafeFreezeM pam
  bounds (PaIxP lb ub _) = (lb,ub)
  checkBounds (PaIxP lb ub _) i = A.inRange (lb,ub) i
  fromList l u xs = runST $ do
    pam <- fromListM l u xs
    unsafeFreezeM pam
  toList pa@(PaIxP lb ub _) = [unsafeIndex pa i | i <- A.range(lb,ub)]
  {-# INLINE unsafeIndex #-}
  {-# INLINE assocs #-}
  {-# INLINE fromAssocs #-}
  {-# INLINE bounds #-}
  {-# INLINE checkBounds #-}
  {-# INLINE fromList #-}
  {-# INLINE toList #-}

instance (Bounded a, A.Ix a, Prim b) => PrimArrayOpsM a b (ST s) where
  data PrimArrayM a b (ST s) = PaIxPM {-# UNPACK #-} !a {-# UNPACK #-} !a {-# UNPACK #-} !(MutableByteArray s)
  readM (PaIxPM lb ub marr) i = readByteArray marr (A.unsafeIndex (lb,ub) i)
  writeM (PaIxPM lb ub marr) i val = writeByteArray marr (A.unsafeIndex (lb,ub) i) val
  boundsM (PaIxPM lb ub _) = return (lb,ub)
  fromAssocsM lb ub z xs = do
    let l = A.rangeSize (lb,ub)
    pam <- PaIxPM lb ub `liftM` newWith l z
    forM_ xs $ uncurry (writeM pam)
    return pam
  unsafeFreezeM (PaIxPM lb ub marr) = do
    arr <- unsafeFreezeByteArray marr
    return $ PaIxP lb ub arr
  fromListM lb ub xs = do
    let l = A.rangeSize (lb,ub)
    pam <- PaIxPM lb ub `liftM` new l (undefined `asTypeOf` head xs)
    zipWithM_ (writeM pam) (A.range (lb,ub)) xs
    return pam
  toListM pam@(PaIxPM lb ub _) = forM (A.range (lb,ub)) (readM pam)
  {-# INLINE readM #-}
  {-# INLINE writeM #-}
  {-# INLINE boundsM #-}
  {-# INLINE fromAssocsM #-}
  {-# INLINE unsafeFreezeM #-}
  {-# INLINE fromListM #-}
  {-# INLINE toListM #-}