{-# LANGUAGE CPP            #-}
{-# LANGUAGE TypeFamilies   #-}
{-# LANGUAGE TypeOperators  #-}
{-# LANGUAGE Rank2Types     #-}


#include "fusion-phases.h"

-----------------------------------------------------------------------------
-- |
-- Module      : Data.Array.Vector.UArr
-- Copyright   : (c) [2001..2002] Manuel M T Chakravarty & Gabriele Keller
--               (c) 2006         Manuel M T Chakravarty & Roman Leshchinskiy
-- License     : see libraries/ndp/LICENSE
-- 
-- Maintainer  : Roman Leshchinskiy <rl@cse.unsw.edu.au>
-- Stability   : internal
-- Portability : non-portable (GADTS)
--
-- Description ---------------------------------------------------------------
-- This module defines unlifted arrays generically as a GADT.
-- 
-- Slicing is implemented by each `BUArr' having the slicing information.  A
-- possible alternative design would be to maintain this information in
-- `UArr', but not in the representations, but at the root.  This may seem
-- attractive at first, but seems to be more disruptive without any real
-- benefits _ this is essentially, because we then need the slicing
-- information at each level; ie, also at the leafs where it is sufficient
-- using the current implementation.
--
-- Todo ----------------------------------------------------------------------
--
module Data.Array.Vector.UArr (

  -- * Array types and classes containing the admissable elements types
  UA, UArr, MUArr, UPrim(..),

  -- * Basic operations on parallel arrays
  lengthU, indexU, sliceU, {-extractU,-} unitsU, zipU, unzipU, fstU, sndU,
  newU, newDynU, newDynResU,
  lengthMU, newMU, readMU, writeMU, copyMU, unsafeFreezeMU, unsafeFreezeAllMU,
  memcpyMU, memcpyOffMU, memmoveOffMU,
  unsafeZipMU, unsafeUnzipMU,

  -- * I\/O
  UIO(..)

) where

-- standard libraries
import Control.Monad (liftM, liftM2)

-- For instances:
import Data.Complex
import GHC.Real

-- friends
import Data.Array.Vector.Prim.BUArr (
  BUArr, MBUArr, UAE,
  lengthBU, indexBU, sliceBU, hGetBU, hPutBU,
  lengthMBU, newMBU, readMBU, writeMBU, copyMBU, unsafeFreezeMBU,
  memcpyMBU, memcpyOffMBU, memmoveOffMBU)

import System.IO
import GHC.ST
import Data.Word
import Data.Int
import Data.Array.Vector.Prim.Debug
import Data.Array.Vector.Prim.Hyperstrict

infixl 9 `indexU`, `readMU`


-- *Basic operations on representation types
-- -----------------------------------------

-- |This type class determines the types that can be elements immutable
-- unboxed arrays. The representation type of these arrays is defined by way
-- of an associated type.  All representation-dependent functions are methods
-- of this class.
--
class UA e where
  -- |The basic array datatype.
  data UArr  e
  data MUArr e :: * -> *

  -- |/O(1?)/. Yield the length of an unboxed array.
  lengthU        :: UArr e                     -> Int

  indexU         :: UArr e -> Int              -> e

  -- |/O(1)/. 'sliceU' restricts access to a subrange of the original array 
  -- (no copying).
  sliceU         :: UArr e -> Int -> Int       -> UArr e

  ------------------------------------------------------------------------

  -- |/O(1)/. 'lengthMU' yields the length of a mutable unboxed array.
  lengthMU       :: MUArr e s                  -> Int

  -- |/O(1)/. 'newMU' allocates a mutable unboxed array of the specified length.
  newMU          :: Int                        -> ST s (MUArr e s)

  -- |/O(1)/. 'readMU' reads the element at the specified index of a mutable 
  -- unboxed array.
  readMU         :: MUArr e s -> Int           -> ST s e

  -- |/O(1)/. 'writeMU' writes a new value to the specified index of a
  -- mutable unboxed array.
  writeMU        :: MUArr e s -> Int -> e      -> ST s ()

  ------------------------------------------------------------------------

  -- |/O(n)/. 'copyMU' copies the contents of an immutable unboxed array into 
  -- a mutable one starting from the specified index.
  --
  copyMU         :: MUArr e s -> Int -> UArr e -> ST s ()

  -- |/O(1)/. 'unsafeFreezeMU' converts a prefix of a mutable array into an 
  -- immutable unboxed array, without copying. The mutable array must not be 
  -- mutated after this.
  --
  unsafeFreezeMU :: MUArr e s -> Int           -> ST s (UArr e)

  -- |Copy a portion of one mutable array to a second.
  memcpyMU       :: MUArr e s -> MUArr e s -> Int -> ST s ()

  -- |Copy a portion of one mutable array to a second, beginning at the
  -- specified offsets for each.
  memcpyOffMU    :: MUArr e s -> MUArr e s -> Int -> Int -> Int -> ST s ()

  -- |Copy a portion of one mutable array to a second, beginning at the
  -- specified offsets for each. This operation is safe even if the source
  -- and destination are the same.
  memmoveOffMU   :: MUArr e s -> MUArr e s -> Int -> Int -> Int -> ST s ()

-- instance HS e => HS (UArr e)
-- instance HS e => HS (MUArr e s)

class UAE e => UPrim e where
  mkUAPrim :: BUArr e -> UArr  e
  unUAPrim :: UArr  e -> BUArr e

  mkMUAPrim :: MBUArr s e -> MUArr  e s
  unMUAPrim :: MUArr  e s -> MBUArr s e

-- |/O(1)/. 'unsafeFreezeAllMU' converts an entire mutable array into an 
-- immutable array, without copying. The mutable array must not be mutated 
-- after this.
--
unsafeFreezeAllMU :: UA e => MUArr e s -> ST s (UArr e)
unsafeFreezeAllMU marr = unsafeFreezeMU marr (lengthMU marr)

-- *Creating unboxed arrays
-- ------------------------

-- |/O(n)/. 'newU' constructs an immutable array of the given size by 
-- performing the provided initialization function on a mutable representation
-- of the output array.
--
newU :: UA e => Int -> (forall s. MUArr e s -> ST s ()) -> UArr e
{-# INLINE_U newU #-}
newU n init = newDynU n (\ma -> init ma >> return n)

newDynU :: UA e => Int -> (forall s. MUArr e s -> ST s Int) -> UArr e
{-# INLINE_U newDynU #-}
newDynU n init =
  runST (do
           ma <- newMU n
           n' <- init ma
           unsafeFreezeMU ma n'
  )

newDynResU :: UA e
           => Int -> (forall s. MUArr e s -> ST s (Int :*: r)) -> UArr e :*: r
{-# INLINE_U newDynResU #-}
newDynResU n init =
  runST (do
           ma <- newMU n
           n' :*: r <- init ma
           arr <- unsafeFreezeMU ma n'
           return (arr :*: r)
  )

-- *Basic operations on unboxed arrays
-- -----------------------------------

-- |/O(1)/. Yield an array of units.
--
unitsU :: Int -> UArr ()
{-# INLINE_STREAM unitsU #-}
unitsU = UAUnit

-- |/O(1)/. Elementwise pairing of array elements. 
--
-- /N.B/: The output will be as long as the first array (and will thus 
-- access past the end of the second array), unlike its List counterpart.
-- This will not occur at the time zipU is called, but only after the resulting
-- array is accessed.
--
zipU :: (UA a, UA b) => UArr a -> UArr b -> UArr (a :*: b)
{-# INLINE_STREAM zipU #-}
zipU = UAProd

-- |/O(1)/. Elementwise unpairing of array elements.
--
unzipU :: (UA a, UA b) => UArr (a :*: b) -> (UArr a :*: UArr b)
{-# INLINE_STREAM unzipU #-}
unzipU (UAProd l r) = (l :*: r)

-- |/O(1)/. Yield the first components of an array of pairs.
--
fstU :: (UA a, UA b) => UArr (a :*: b) -> UArr a
{-# INLINE_STREAM fstU #-}
fstU (UAProd l r) = l

-- |/O(1)/. Yield the second components of an array of pairs.
--
sndU :: (UA a, UA b) => UArr (a :*: b) -> UArr b
{-# INLINE_STREAM sndU #-}
sndU (UAProd l r) = r

-- |Elementwise pairing of mutable arrays. This is an unsafe
-- operation, as no copying is performed, so changes to the
-- pair array will affect the original arrays, and vice versa.
unsafeZipMU :: (UA a, UA b) => MUArr a s -> MUArr b s -> MUArr (a :*: b) s
{-# INLINE_U unsafeZipMU #-}
unsafeZipMU = MUAProd

-- |Elementwise unpairing of mutable arrays. This is an unsafe
-- operation, as no copying is performed, so changes to the
-- unpaired arrays will affect the original, and vice versa.
unsafeUnzipMU :: (UA a, UA b) => MUArr (a :*: b) s -> MUArr a s :*: MUArr b s
{-# INLINE_U unsafeUnzipMU #-}
unsafeUnzipMU (MUAProd mua mub) = mua :*: mub

-- *Family of representation types
-- -------------------------------

-- |Array operations on the unit representation.
--
instance UA () where
  newtype UArr  ()   = UAUnit  Int
  newtype MUArr () s = MUAUnit Int

  lengthU (UAUnit n)     = n
  indexU  (UAUnit _) _   = ()
  sliceU  (UAUnit len) i n = if i == len then UAUnit 0 else UAUnit (min n (len - i))

  lengthMU (MUAUnit n)            = n
  newMU   n                       = return $ MUAUnit n
  readMU (MUAUnit _) _            = return ()
  writeMU (MUAUnit _) _ _         = return ()
  copyMU (MUAUnit _) _ (UAUnit _) = return ()
  unsafeFreezeMU (MUAUnit _) n    = return $ UAUnit n

  memcpyMU     (MUAUnit _) (MUAUnit _) _     = return ()
  memcpyOffMU  (MUAUnit _) (MUAUnit _) _ _ _ = return ()
  memmoveOffMU (MUAUnit _) (MUAUnit _) _ _ _ = return ()

-- |Array operations on the pair representation.
--
instance (UA a, UA b) => UA (a :*: b) where
  data UArr  (a :*: b)   = UAProd  !(UArr a)    !(UArr b)
  data MUArr (a :*: b) s = MUAProd !(MUArr a s) !(MUArr b s)

  -- TODO: changed from (lengthU l), as this causes problems when the length is used to
  --       limit the index
  lengthU (UAProd l r)     = checkEq "lengthU" "lengths of zipped arrays differ" (lengthU l) (lengthU r)
     (lengthU l) 
  {-# INLINE_U indexU #-}
  indexU  (UAProd l r) i   = indexU l i :*: indexU r i
  {-# INLINE_U sliceU #-}
  sliceU  (UAProd l r) i n = UAProd (sliceU l i n) (sliceU r i n)

  {-# INLINE_U lengthMU #-}
  lengthMU (MUAProd l r)   = lengthMU l

  {-# INLINE_U newMU #-}
  newMU n = 
    do
      a <- newMU n
      b <- newMU n
      return $ MUAProd a b
  {-# INLINE_U readMU #-}
  readMU (MUAProd a b) i = liftM2 (:*:) (a `readMU` i) (b `readMU` i)
  {-# INLINE_U writeMU #-}
  writeMU (MUAProd a b) i (x :*: y) = 
    do
      writeMU a i x
      writeMU b i y
  {-# INLINE_U copyMU #-}
  copyMU (MUAProd ma mb) i (UAProd a b) =
    do
      copyMU ma i a
      copyMU mb i b
  {-# INLINE_U unsafeFreezeMU #-}
  unsafeFreezeMU (MUAProd a b) n = 
    do
      a' <- unsafeFreezeMU a n
      b' <- unsafeFreezeMU b n
      return $ UAProd a' b'

  {-# INLINE_U memcpyMU #-}
  memcpyMU (MUAProd ma mb) (MUAProd ma' mb') l =
    do memcpyMU ma ma' l
       memcpyMU mb mb' l

  {-# INLINE_U memcpyOffMU #-}
  memcpyOffMU (MUAProd ma mb) (MUAProd ma' mb') s d l =
    do memcpyOffMU ma ma' s d l
       memcpyOffMU mb mb' s d l

  {-# INLINE_U memmoveOffMU #-}
  memmoveOffMU (MUAProd ma mb) (MUAProd ma' mb') s d l = 
    do memmoveOffMU ma ma' s d l
       memmoveOffMU mb mb' s d l

{-
-- |Selector for immutable arrays of sums
--
data USel = USel {
              selUS  :: !(BUArr Bool),  -- selector (False => left)
              lidxUS :: !(BUArr Int),   -- left indices
              ridxUS :: !(BUArr Int)    -- right indices
            }

--instance HS USel

-- |Selector for mutable arrays of sums
--
data MUSel s = MUSel {
                 selMUS  :: !(MBUArr s Bool),  -- selector (False => left)
                 lidxMUS :: !(MBUArr s Int),   -- left indices
                 ridxMUS :: !(MBUArr s Int)    -- right indices
               }

--instance HS (MUSel s)

-- |Array operations on the sum representation
--
instance (UA a, UA b) => UA (a :+: b) where
  lengthU (UASum sel _ _)     = lengthBU (selUS sel)
  {-# INLINE_U indexU #-}
  indexU  (UASum sel l r) i   = if (selUS sel)`indexBU`i then Inr $ indexU r i 
                                                         else Inl $ indexU l i
  {-# INLINE_U sliceU #-}
  sliceU  (UASum sel l r) i n = 
    let
      sel'        = sliceBU (selUS sel) i n
      li          = lidxUS sel`indexBU`i
      ri          = ridxUS sel`indexBU`i
      lidx        = mapBU (subtract li) $ sliceBU (lidxUS sel) i n
      ridx        = mapBU (subtract ri) $ sliceBU (ridxUS sel) i n
      (ln :*: rn) = if n == 0 then (0 :*: 0) else (lidx`indexBU`(n - 1) :*: 
                                                   ridx`indexBU`(n - 1))
    in
    UASum (USel sel' lidx ridx) (sliceU l li ln) (sliceU r ri rn)
  {-# INLINE_U extractU #-}
  extractU  (UASum sel l r) i n = 
    let
      sel'        = extractBU (selUS sel) i n
      li          = lidxUS sel`indexBU`i
      ri          = ridxUS sel`indexBU`i
      lidx        = mapBU (subtract li) $ sliceBU (lidxUS sel) i n
      ridx        = mapBU (subtract ri) $ sliceBU (ridxUS sel) i n
      (ln :*: rn) = if n == 0 then (0 :*: 0) else (lidx`indexBU`(n - 1) :*: 
                                                   ridx`indexBU`(n - 1))
    in
    UASum (USel sel' lidx ridx) (extractU l li ln) (extractU r ri rn)

instance (MUA a, MUA b) => MUA (a :+: b) where
  {-# INLINE_U newMU #-}
  newMU n = do
              sel  <- newMBU n
              lidx <- newMBU n
              ridx <- newMBU n
              a    <- newMU n
              b    <- newMU n
              return $ MUASum (MUSel sel lidx ridx) a b
  {-# INLINE_U writeMU #-}
  writeMU (MUASum sel l r) i (Inl x) = 
    do
      let lidx = lidxMUS sel
          ridx = ridxMUS sel
      writeMBU (selMUS sel) i False
      li <- if i == 0 then return 0 else liftM (+ 1) $ lidx`readMBU`(i - 1)
      ri <- if i == 0 then return 0 else               ridx`readMBU`(i - 1)
      writeMBU lidx i li
      writeMBU ridx i ri
      writeMU l li x
  writeMU (MUASum sel l r) i (Inr x) = 
    do
      let lidx = lidxMUS sel
          ridx = ridxMUS sel
      writeMBU (selMUS sel) i True
      li <- if i == 0 then return 0 else               lidx`readMBU`(i - 1)
      ri <- if i == 0 then return 0 else liftM (+ 1) $ ridx`readMBU`(i - 1)
      writeMBU lidx i li
      writeMBU ridx i ri
      writeMU r ri x
    --FIXME: that works only when the array is constructed left to right, but
    --not for something like permutations
  {-# INLINE_U unsafeFreezeMU #-}
  unsafeFreezeMU (MUASum sel l r) n = 
    do
      sel' <- unsafeFreezeMBU (selMUS  sel) n
      lidx <- unsafeFreezeMBU (lidxMUS sel) n
      ridx <- unsafeFreezeMBU (ridxMUS sel) n
      let ln = if n == 0 then 0 else lidx`indexBU`(n - 1)
          rn = if n == 0 then 0 else ridx`indexBU`(n - 1)
      l' <- unsafeFreezeMU l ln
      r' <- unsafeFreezeMU r rn
      return $ UASum (USel sel' lidx ridx) l' r'
-}

-- |Array operations on unboxed arrays
-- -
--
-- NB: We use instances for all possible unboxed types instead of re-using the
--     overloading provided by UAE to avoid having to store the UAE dictionary
--     in `UAPrimU'.

primLengthU :: UPrim e => UArr e -> Int
{-# INLINE_U primLengthU #-}
primLengthU = lengthBU . unUAPrim

primIndexU :: UPrim e => UArr e -> Int -> e
{-# INLINE_U primIndexU #-}
primIndexU = indexBU . unUAPrim

primSliceU :: UPrim e => UArr e -> Int -> Int -> UArr e
{-# INLINE_U primSliceU #-}
primSliceU arr i = mkUAPrim . sliceBU (unUAPrim arr) i

primLengthMU :: UPrim e => MUArr e s -> Int
{-# INLINE_U primLengthMU #-}
primLengthMU = lengthMBU . unMUAPrim

primNewMU :: UPrim e => Int -> ST s (MUArr e s)
{-# INLINE_U primNewMU #-}
primNewMU = liftM mkMUAPrim . newMBU

primReadMU :: UPrim e => MUArr e s -> Int -> ST s e
{-# INLINE_U primReadMU #-}
primReadMU = readMBU . unMUAPrim

primWriteMU :: UPrim e => MUArr e s -> Int -> e -> ST s ()
{-# INLINE_U primWriteMU #-}
primWriteMU = writeMBU . unMUAPrim

primCopyMU :: UPrim e => MUArr e s -> Int -> UArr e -> ST s ()
{-# INLINE_U primCopyMU #-}
primCopyMU ma i = copyMBU (unMUAPrim ma) i . unUAPrim

primUnsafeFreezeMU :: UPrim e => MUArr e s -> Int -> ST s (UArr e)
{-# INLINE_U primUnsafeFreezeMU #-}
primUnsafeFreezeMU ma = liftM mkUAPrim . unsafeFreezeMBU (unMUAPrim ma)

primMemcpyMU :: UPrim e => MUArr e s -> MUArr e s -> Int -> ST s ()
{-# INLINE_U primMemcpyMU #-}
primMemcpyMU src dst l = memcpyMBU (unMUAPrim src) (unMUAPrim dst) l

primMemcpyOffMU :: UPrim e => MUArr e s -> MUArr e s -> Int -> Int -> Int -> ST s ()
{-# INLINE_U primMemcpyOffMU #-}
primMemcpyOffMU src dst s d l = memcpyOffMBU (unMUAPrim src) (unMUAPrim dst) s d l

primMemmoveOffMU :: UPrim e => MUArr e s -> MUArr e s -> Int -> Int -> Int -> ST s ()
{-# INLINE_U primMemmoveOffMU #-}
primMemmoveOffMU src dst s d l = memmoveOffMBU (unMUAPrim src) (unMUAPrim dst) s d l

instance UPrim Bool where
  mkUAPrim                = UABool
  unUAPrim  (UABool  arr) = arr

  mkMUAPrim               = MUABool
  unMUAPrim (MUABool arr) = arr

instance UA Bool where
  newtype UArr  Bool   = UABool  (BUArr Bool)
  newtype MUArr Bool s = MUABool (MBUArr s Bool)

  lengthU        = primLengthU
  indexU         = primIndexU
  sliceU         = primSliceU

  lengthMU       = primLengthMU
  newMU          = primNewMU
  readMU         = primReadMU
  writeMU        = primWriteMU
  copyMU         = primCopyMU
  unsafeFreezeMU = primUnsafeFreezeMU

  memcpyMU     = primMemcpyMU
  memcpyOffMU  = primMemcpyOffMU
  memmoveOffMU = primMemmoveOffMU

instance UPrim Char where
  mkUAPrim                 = UAChar
  unUAPrim  (UAChar   arr) = arr

  mkMUAPrim                = MUAChar
  unMUAPrim (MUAChar arr)  = arr

instance UA Char where
  newtype UArr  Char   = UAChar  (BUArr Char)
  newtype MUArr Char s = MUAChar (MBUArr s Char)

  lengthU        = primLengthU
  indexU         = primIndexU
  sliceU         = primSliceU

  lengthMU       = primLengthMU
  newMU          = primNewMU
  readMU         = primReadMU
  writeMU        = primWriteMU
  copyMU         = primCopyMU
  unsafeFreezeMU = primUnsafeFreezeMU

  memcpyMU     = primMemcpyMU
  memcpyOffMU  = primMemcpyOffMU
  memmoveOffMU = primMemmoveOffMU


instance UPrim Int where
  mkUAPrim               = UAInt
  unUAPrim  (UAInt  arr) = arr

  mkMUAPrim              = MUAInt
  unMUAPrim (MUAInt arr) = arr

instance UA Int where
  newtype UArr  Int   = UAInt  (BUArr Int)
  newtype MUArr Int s = MUAInt (MBUArr s Int)

  lengthU        = primLengthU
  indexU         = primIndexU
  sliceU         = primSliceU

  lengthMU       = primLengthMU
  newMU          = primNewMU
  readMU         = primReadMU
  writeMU        = primWriteMU
  copyMU         = primCopyMU
  unsafeFreezeMU = primUnsafeFreezeMU

  memcpyMU     = primMemcpyMU
  memcpyOffMU  = primMemcpyOffMU
  memmoveOffMU = primMemmoveOffMU

  -- FIXME: For now, we assume that Int writes are atomic but we should really
  --        configure this.

instance UPrim Word where
  mkUAPrim               = UAWord
  unUAPrim  (UAWord  arr) = arr
  mkMUAPrim              = MUAWord
  unMUAPrim (MUAWord arr) = arr

instance UA Word where
  newtype UArr  Word   = UAWord  (BUArr Word)
  newtype MUArr Word s = MUAWord (MBUArr s Word)

  lengthU        = primLengthU
  indexU         = primIndexU
  sliceU         = primSliceU
  lengthMU       = primLengthMU
  newMU          = primNewMU
  readMU         = primReadMU
  writeMU        = primWriteMU
  copyMU         = primCopyMU
  unsafeFreezeMU = primUnsafeFreezeMU

  memcpyMU     = primMemcpyMU
  memcpyOffMU  = primMemcpyOffMU
  memmoveOffMU = primMemmoveOffMU
  -- FIXME: For now, we assume that Word writes are atomic but we should really
  --        configure this.

instance UPrim Float where
  mkUAPrim                 = UAFloat
  unUAPrim  (UAFloat  arr) = arr

  mkMUAPrim                = MUAFloat
  unMUAPrim (MUAFloat arr) = arr

instance UA Float where
  newtype UArr  Float   = UAFloat  (BUArr Float)
  newtype MUArr Float s = MUAFloat (MBUArr s Float)

  lengthU        = primLengthU
  indexU         = primIndexU
  sliceU         = primSliceU

  lengthMU       = primLengthMU
  newMU          = primNewMU
  readMU         = primReadMU
  writeMU        = primWriteMU
  copyMU         = primCopyMU
  unsafeFreezeMU = primUnsafeFreezeMU

  memcpyMU     = primMemcpyMU
  memcpyOffMU  = primMemcpyOffMU
  memmoveOffMU = primMemmoveOffMU

instance UPrim Double where
  mkUAPrim                  = UADouble
  unUAPrim  (UADouble  arr) = arr

  mkMUAPrim                 = MUADouble
  unMUAPrim (MUADouble arr) = arr

instance UA Double where
  newtype UArr  Double   = UADouble  (BUArr Double)
  newtype MUArr Double s = MUADouble (MBUArr s Double)

  lengthU        = primLengthU
  indexU         = primIndexU
  sliceU         = primSliceU

  lengthMU       = primLengthMU
  newMU          = primNewMU
  readMU         = primReadMU
  writeMU        = primWriteMU
  copyMU         = primCopyMU
  unsafeFreezeMU = primUnsafeFreezeMU

  memcpyMU     = primMemcpyMU
  memcpyOffMU  = primMemcpyOffMU
  memmoveOffMU = primMemmoveOffMU

instance UPrim Word8 where
  mkUAPrim               = UAWord8
  unUAPrim  (UAWord8  arr) = arr
  mkMUAPrim              = MUAWord8
  unMUAPrim (MUAWord8 arr) = arr

instance UA Word8 where
  newtype UArr  Word8   = UAWord8  (BUArr Word8)
  newtype MUArr Word8 s = MUAWord8 (MBUArr s Word8)

  lengthU        = primLengthU
  indexU         = primIndexU
  sliceU         = primSliceU
  lengthMU       = primLengthMU
  newMU          = primNewMU
  readMU         = primReadMU
  writeMU        = primWriteMU
  copyMU         = primCopyMU
  unsafeFreezeMU = primUnsafeFreezeMU

  memcpyMU     = primMemcpyMU
  memcpyOffMU  = primMemcpyOffMU
  memmoveOffMU = primMemmoveOffMU
  -- FIXME: For now, we assume that Word8 writes are atomic but we should really
  --        configure this.

instance UPrim Word16 where
  mkUAPrim               = UAWord16
  unUAPrim  (UAWord16  arr) = arr
  mkMUAPrim              = MUAWord16
  unMUAPrim (MUAWord16 arr) = arr

instance UA Word16 where
  newtype UArr  Word16   = UAWord16  (BUArr Word16)
  newtype MUArr Word16 s = MUAWord16 (MBUArr s Word16)

  lengthU        = primLengthU
  indexU         = primIndexU
  sliceU         = primSliceU
  lengthMU       = primLengthMU
  newMU          = primNewMU
  readMU         = primReadMU
  writeMU        = primWriteMU
  copyMU         = primCopyMU
  unsafeFreezeMU = primUnsafeFreezeMU

  memcpyMU     = primMemcpyMU
  memcpyOffMU  = primMemcpyOffMU
  memmoveOffMU = primMemmoveOffMU

instance UPrim Word32 where
  mkUAPrim               = UAWord32
  unUAPrim  (UAWord32  arr) = arr
  mkMUAPrim              = MUAWord32
  unMUAPrim (MUAWord32 arr) = arr

instance UA Word32 where
  newtype UArr  Word32   = UAWord32  (BUArr Word32)
  newtype MUArr Word32 s = MUAWord32 (MBUArr s Word32)

  lengthU        = primLengthU
  indexU         = primIndexU
  sliceU         = primSliceU
  lengthMU       = primLengthMU
  newMU          = primNewMU
  readMU         = primReadMU
  writeMU        = primWriteMU
  copyMU         = primCopyMU
  unsafeFreezeMU = primUnsafeFreezeMU

  memcpyMU     = primMemcpyMU
  memcpyOffMU  = primMemcpyOffMU
  memmoveOffMU = primMemmoveOffMU

instance UPrim Word64 where
  mkUAPrim               = UAWord64
  unUAPrim  (UAWord64  arr) = arr
  mkMUAPrim              = MUAWord64
  unMUAPrim (MUAWord64 arr) = arr

instance UA Word64 where
  newtype UArr  Word64   = UAWord64  (BUArr Word64)
  newtype MUArr Word64 s = MUAWord64 (MBUArr s Word64)

  lengthU        = primLengthU
  indexU         = primIndexU
  sliceU         = primSliceU
  lengthMU       = primLengthMU
  newMU          = primNewMU
  readMU         = primReadMU
  writeMU        = primWriteMU
  copyMU         = primCopyMU
  unsafeFreezeMU = primUnsafeFreezeMU

  memcpyMU     = primMemcpyMU
  memcpyOffMU  = primMemcpyOffMU
  memmoveOffMU = primMemmoveOffMU

instance UPrim Int8 where
  mkUAPrim               = UAInt8
  unUAPrim  (UAInt8  arr) = arr
  mkMUAPrim              = MUAInt8
  unMUAPrim (MUAInt8 arr) = arr

instance UA Int8 where
  newtype UArr  Int8   = UAInt8  (BUArr Int8)
  newtype MUArr Int8 s = MUAInt8 (MBUArr s Int8)

  lengthU        = primLengthU
  indexU         = primIndexU
  sliceU         = primSliceU
  lengthMU       = primLengthMU
  newMU          = primNewMU
  readMU         = primReadMU
  writeMU        = primWriteMU
  copyMU         = primCopyMU
  unsafeFreezeMU = primUnsafeFreezeMU

  memcpyMU     = primMemcpyMU
  memcpyOffMU  = primMemcpyOffMU
  memmoveOffMU = primMemmoveOffMU
  -- FIXME: For now, we assume that Int8 writes are atomic but we should really
  --        configure this.

instance UPrim Int16 where
  mkUAPrim               = UAInt16
  unUAPrim  (UAInt16  arr) = arr
  mkMUAPrim              = MUAInt16
  unMUAPrim (MUAInt16 arr) = arr

instance UA Int16 where
  newtype UArr  Int16   = UAInt16  (BUArr Int16)
  newtype MUArr Int16 s = MUAInt16 (MBUArr s Int16)

  lengthU        = primLengthU
  indexU         = primIndexU
  sliceU         = primSliceU
  lengthMU       = primLengthMU
  newMU          = primNewMU
  readMU         = primReadMU
  writeMU        = primWriteMU
  copyMU         = primCopyMU
  unsafeFreezeMU = primUnsafeFreezeMU

  memcpyMU     = primMemcpyMU
  memcpyOffMU  = primMemcpyOffMU
  memmoveOffMU = primMemmoveOffMU

instance UPrim Int32 where
  mkUAPrim               = UAInt32
  unUAPrim  (UAInt32  arr) = arr
  mkMUAPrim              = MUAInt32
  unMUAPrim (MUAInt32 arr) = arr

instance UA Int32 where
  newtype UArr  Int32   = UAInt32  (BUArr Int32)
  newtype MUArr Int32 s = MUAInt32 (MBUArr s Int32)

  lengthU        = primLengthU
  indexU         = primIndexU
  sliceU         = primSliceU
  lengthMU       = primLengthMU
  newMU          = primNewMU
  readMU         = primReadMU
  writeMU        = primWriteMU
  copyMU         = primCopyMU
  unsafeFreezeMU = primUnsafeFreezeMU

  memcpyMU     = primMemcpyMU
  memcpyOffMU  = primMemcpyOffMU
  memmoveOffMU = primMemmoveOffMU

instance UPrim Int64 where
  mkUAPrim               = UAInt64
  unUAPrim  (UAInt64  arr) = arr
  mkMUAPrim              = MUAInt64
  unMUAPrim (MUAInt64 arr) = arr

instance UA Int64 where
  newtype UArr  Int64   = UAInt64  (BUArr Int64)
  newtype MUArr Int64 s = MUAInt64 (MBUArr s Int64)

  lengthU        = primLengthU
  indexU         = primIndexU
  sliceU         = primSliceU
  lengthMU       = primLengthMU
  newMU          = primNewMU
  readMU         = primReadMU
  writeMU        = primWriteMU
  copyMU         = primCopyMU
  unsafeFreezeMU = primUnsafeFreezeMU

  memcpyMU     = primMemcpyMU
  memcpyOffMU  = primMemcpyOffMU
  memmoveOffMU = primMemmoveOffMU

------------------------------------------------------------------------

-- TODO could use a single array of 'a', doubly packed

instance (RealFloat a, UA a) => UA (Complex a) where

  newtype UArr  (Complex a)   = UAComplex  (UArr (a :*: a))
  newtype MUArr (Complex a) s = MUAComplex (MUArr (a :*: a) s)

  lengthU (UAComplex arr)   = lengthU arr

  indexU  (UAComplex arr) i = case indexU arr i of (a :*: b) -> a :+ b

  sliceU (UAComplex arr) i n = UAComplex (sliceU arr i n)

  lengthMU (MUAComplex arr)   = lengthMU arr

  newMU n = return . MUAComplex =<< newMU n

  readMU (MUAComplex arr) n = do (a :*: b) <- readMU arr n; return (a :+ b)

  writeMU (MUAComplex arr) i  (x :+ y) = writeMU arr i (x :*: y)

  copyMU (MUAComplex mua) n (UAComplex ua) = copyMU mua n ua

  unsafeFreezeMU (MUAComplex arr) n = do arr' <- unsafeFreezeMU arr n; return (UAComplex arr')

  memcpyMU     (MUAComplex src) (MUAComplex dst)     l = memcpyMU     src dst     l
  memcpyOffMU  (MUAComplex src) (MUAComplex dst) s d l = memcpyOffMU  src dst s d l
  memmoveOffMU (MUAComplex src) (MUAComplex dst) s d l = memmoveOffMU src dst s d l

instance (Integral a, UA a) => UA (Ratio a) where
  newtype UArr  (Ratio a)   = UARatio  (UArr  (a :*: a))
  newtype MUArr (Ratio a) s = MUARatio (MUArr (a :*: a) s)

  lengthU (UARatio arr)   = lengthU arr
  indexU  (UARatio arr) i = case indexU arr i of (a :*: b) -> a % b
  sliceU (UARatio arr) i n = UARatio (sliceU arr i n)

  lengthMU (MUARatio arr)   = lengthMU arr
  newMU n = return . MUARatio =<< newMU n
  readMU (MUARatio arr) n = do (a :*: b) <- readMU arr n; return (a % b)
  writeMU (MUARatio arr) i (n :% d) = writeMU arr i (n :*: d)
  copyMU (MUARatio mua) n (UARatio ua) = copyMU mua n ua
  unsafeFreezeMU (MUARatio arr) n = do arr' <- unsafeFreezeMU arr n; return (UARatio arr')

  memcpyMU     (MUARatio src) (MUARatio dst)     l = memcpyMU     src dst     l
  memcpyOffMU  (MUARatio src) (MUARatio dst) s d l = memcpyOffMU  src dst s d l
  memmoveOffMU (MUARatio src) (MUARatio dst) s d l = memmoveOffMU src dst s d l


------------------------------------------------------------------------

-- * I\/O
-- -----

class UA a => UIO a where
  hPutU :: Handle -> UArr a -> IO ()
  hGetU :: Handle -> IO (UArr a)

primPutU :: UPrim a => Handle -> UArr a -> IO ()
primPutU h = hPutBU h . unUAPrim

primGetU :: UPrim a => Handle -> IO (UArr a)
primGetU = liftM mkUAPrim . hGetBU

------------------------------------------------------------------------

instance UIO Bool   where hPutU = primPutU; hGetU = primGetU
instance UIO Char   where hPutU = primPutU; hGetU = primGetU
instance UIO Int    where hPutU = primPutU; hGetU = primGetU
instance UIO Word   where hPutU = primPutU; hGetU = primGetU
instance UIO Float  where hPutU = primPutU; hGetU = primGetU
instance UIO Double where hPutU = primPutU; hGetU = primGetU

instance UIO Word8  where hPutU = primPutU; hGetU = primGetU
instance UIO Word16 where hPutU = primPutU; hGetU = primGetU
instance UIO Word32 where hPutU = primPutU; hGetU = primGetU
instance UIO Word64 where hPutU = primPutU; hGetU = primGetU

instance UIO Int8   where hPutU = primPutU; hGetU = primGetU
instance UIO Int16  where hPutU = primPutU; hGetU = primGetU
instance UIO Int32  where hPutU = primPutU; hGetU = primGetU
instance UIO Int64  where hPutU = primPutU; hGetU = primGetU

------------------------------------------------------------------------

instance (UIO a, UIO b) => UIO (a :*: b) where
  hPutU h (UAProd xs ys) = do hPutU h xs
                              hPutU h ys
  hGetU h                = do xs <- hGetU h
                              ys <- hGetU h
                              return (UAProd xs ys)

instance (RealFloat a, UIO a) => UIO (Complex a) where
  hPutU h (UAComplex arr) = hPutU h arr
  hGetU h                 = do arr <- hGetU h
                               return (UAComplex arr)

instance (Integral a, UIO a) => UIO (Ratio a) where
  hPutU h (UARatio arr) = hPutU h arr
  hGetU h                 = do arr <- hGetU h
                               return (UARatio arr)