    module Data.Array.Repa.Repr.Partitioned
            ( P, Array (..)
            , Range(..)
            , inRange)
    where
    import Data.Array.Repa.Base
    import Data.Array.Repa.Shape
    import Data.Array.Repa.Eval
    import Data.Array.Repa.Repr.Delayed

    -- | Partitioned arrays.
    --   The last partition takes priority
    --
    --   These are produced by Repa's support functions and allow arrays to be defined
    --   using a different element function for each partition.
    --
    --   The basic idea is described in ``Efficient Parallel Stencil Convolution'',
    --   Ben Lippmeier and Gabriele Keller, Haskell 2011 -- though the underlying
    --   array representation has changed since this paper was published.
    --
    data P r1 r2

    data Range sh
            = Range !sh !sh                      -- indices defining the range
                    (sh -> Bool)                 -- predicate to check whether were in range

    -- | Check whether an index is within the given range.
    inRange :: Range sh -> sh -> Bool
    inRange (Range _ _ p) ix
            = p ix
    {-# INLINE inRange #-}

    -- Repr -----------------------------------------------------------------------
    -- | Read elements from a partitioned array.
    instance (Source r1 e, Source r2 e) => Source (P r1 r2) e where
     data Array (P r1 r2) sh e
            = APart !sh                          -- size of the whole array
                    !(Range sh) !(Array r1 sh e) -- if in range use this array
                    !(Array r2 sh e)             -- otherwise use this array

     index (APart _ range arr1 arr2) ix
       | inRange range ix   = index arr1 ix
       | otherwise          = index arr2 ix
     {-# INLINE index #-}

     linearIndex arr@(APart sh _ _ _) ix
            = index arr $ fromIndex sh ix
     {-# INLINE linearIndex #-}

     extent (APart sh _ _ _)
            = sh
     {-# INLINE extent #-}

     deepSeqArray (APart sh range arr1 arr2) y
      = sh `deepSeq` range `deepSeqRange` arr1 `deepSeqArray` arr2 `deepSeqArray` y
     {-# INLINE deepSeqArray #-}

    deepSeqRange :: Shape sh => Range sh -> b -> b
    deepSeqRange (Range ix sz f) y
            = ix `deepSeq` sz `deepSeq` f `seq` y
    {-# INLINE deepSeqRange #-}

    -- Load -----------------------------------------------------------------------
    instance (LoadRange r1 sh e, Load r2 sh e)
            => Load (P r1 r2) sh e where
     loadP (APart _ (Range ix sz _) arr1 arr2) marr
      = do  loadRangeP arr1 marr ix sz
            loadP arr2 marr
     {-# INLINE loadP #-}

     loadS (APart _ (Range ix sz _) arr1 arr2) marr
      = do  loadRangeS arr1 marr ix sz
            loadS arr2 marr
     {-# INLINE loadS #-}