    {-# OPTIONS -fno-warn-incomplete-patterns #-}
    {-# LANGUAGE CPP #-}
    -----------------------------------------------------------------------------
    -- |
    -- Module      : Data.Array.Parallel.Unlifted.Sequential.Flat.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 ----------------------------------------------------------------------
    --

    #include "fusion-phases.h"

    module Data.Array.Parallel.Unlifted.Sequential.Flat.UArr (

      -- * Array types and classes containing the admissable elements types
      UA, UArr, MUArr, {-USel(..), MUSel(..),-}

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

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

    ) where

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

    -- friends
    import Data.Array.Parallel.Base
    import Data.Array.Parallel.Arr (
      BUArr, MBUArr, UAE,
      lengthBU, indexBU, sliceBU, hGetBU, hPutBU,
      lengthMBU, newMBU, readMBU, writeMBU, copyMBU, unsafeFreezeMBU)

    import System.IO

    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 HS e => UA e where
      data UArr  e
      data MUArr e :: * -> *

      -- |Yield the length of an unboxed array
      lengthU        :: UArr e                     -> Int

      -- |Extract an element out of an immutable unboxed array
      indexU         :: UArr e -> Int              -> e

      -- |Restrict access to a subrange of the original array (no copying)
      sliceU         :: UArr e -> Int -> Int       -> UArr e

      -- |Yield the length of a mutable unboxed array
      lengthMU       :: MUArr e s                  -> Int

      -- |Allocate a mutable unboxed array
      newMU          :: Int                        -> ST s (MUArr e s)

      -- |Read an element from a mutable unboxed array
      readMU         :: MUArr e s -> Int           -> ST s e

      -- |Update an element in a mutable unboxed array
      writeMU        :: MUArr e s -> Int -> e      -> ST s ()

      -- |Indicate whether the type supports atomic updates
      hasAtomicWriteMU :: e -> Bool

      -- |Atomically update an element in a mutable unboxed array if supported
      atomicWriteMU  :: MUArr e s -> Int -> e      -> ST s ()

      -- |Copy the contents of an immutable unboxed array into a mutable one
      -- from the specified position on
      copyMU         :: MUArr e s -> Int -> UArr e -> ST s ()

      -- |Convert a mutable into an immutable unboxed array
      unsafeFreezeMU :: MUArr e s -> Int           -> ST s (UArr e)

      hasAtomicWriteMU _  = False
      atomicWriteMU _ _ _ = error "atomicWriteMU: not supported"

    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

    unsafeFreezeAllMU :: UA e => MUArr e s -> ST s (UArr e)
    unsafeFreezeAllMU marr = unsafeFreezeMU marr (lengthMU marr)

    -- |Creating unboxed arrays
    -- ------------------------

    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
    -- -----------------------------------

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

    -- |Elementwise pairing of array elements.
    --
    zipU :: (UA a, UA b) => UArr a -> UArr b -> UArr (a :*: b)
    {-# INLINE_STREAM zipU #-}
    zipU = UAProd

    -- |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)

    -- |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

    -- |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

    -- |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 _) _ n = UAUnit n

      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

      hasAtomicWriteMU _ = True
      atomicWriteMU      = writeMU

    -- |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'

    {-
    -- |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)

    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

    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

    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

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

      hasAtomicWriteMU _ = True
      atomicWriteMU      = primWriteMU

    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

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

      hasAtomicWriteMU _ = True
      atomicWriteMU      = primWriteMU

    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

    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

    -- * 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 Int where
      hPutU = primPutU
      hGetU = primGetU

    instance UIO Double 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)