{-# LANGUAGE BangPatterns #-}

-- | Functions specialised for arrays of dimension 2.
module Data.Array.Repa.Specialised.Dim2
        ( isInside2
        , isOutside2
        , clampToBorder2
        , makeBordered2)
where
import Data.Array.Repa.Index
import Data.Array.Repa.Base
import Data.Array.Repa.Repr.Partitioned
import Data.Array.Repa.Repr.Undefined


-- | Check if an index lies inside the given extent.
--   As opposed to `inRange` from "Data.Array.Repa.Index",
--   this is a short-circuited test that checks that lowest dimension first.
isInside2
        :: DIM2         -- ^ Extent of array.
        -> DIM2         -- ^ Index to check.
        -> Bool

{-# INLINE isInside2 #-}
isInside2 :: DIM2 -> DIM2 -> Bool
isInside2 DIM2
ex    = Bool -> Bool
not (Bool -> Bool) -> (DIM2 -> Bool) -> DIM2 -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DIM2 -> DIM2 -> Bool
isOutside2 DIM2
ex


-- | Check if an index lies outside the given extent.
--   As opposed to `inRange` from "Data.Array.Repa.Index",
--   this is a short-circuited test that checks the lowest dimension first.
isOutside2
        :: DIM2         -- ^ Extent of array.
        -> DIM2         -- ^ Index to check.
        -> Bool

{-# INLINE isOutside2 #-}
isOutside2 :: DIM2 -> DIM2 -> Bool
isOutside2 (DIM0
_ :. Int
yLen :. Int
xLen) (DIM0
_ :. Int
yy :. Int
xx)
        | Int
xx Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0        = Bool
True
        | Int
xx Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
xLen    = Bool
True
        | Int
yy Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0        = Bool
True
        | Int
yy Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
yLen    = Bool
True
        | Bool
otherwise     = Bool
False


-- | Given the extent of an array, clamp the components of an index so they
--   lie within the given array. Outlying indices are clamped to the index
--   of the nearest border element.
clampToBorder2
        :: DIM2         -- ^ Extent of array.
        -> DIM2         -- ^ Index to clamp.
        -> DIM2

{-# INLINE clampToBorder2 #-}
clampToBorder2 :: DIM2 -> DIM2 -> DIM2
clampToBorder2 (DIM0
_ :. Int
yLen :. Int
xLen) (DIM0
sh :. Int
j :. Int
i)
 = Int -> Int -> DIM2
clampX Int
j Int
i
 where  {-# INLINE clampX #-}
        clampX :: Int -> Int -> DIM2
clampX !Int
y !Int
x
          | Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0       = Int -> Int -> DIM2
clampY Int
y Int
0
          | Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
xLen   = Int -> Int -> DIM2
clampY Int
y (Int
xLen Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
          | Bool
otherwise   = Int -> Int -> DIM2
clampY Int
y Int
x

        {-# INLINE clampY #-}
        clampY :: Int -> Int -> DIM2
clampY !Int
y !Int
x
          | Int
y Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0       = DIM0
sh DIM0 -> Int -> DIM0 :. Int
forall tail head. tail -> head -> tail :. head
:. Int
0          (DIM0 :. Int) -> Int -> DIM2
forall tail head. tail -> head -> tail :. head
:. Int
x
          | Int
y Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
yLen   = DIM0
sh DIM0 -> Int -> DIM0 :. Int
forall tail head. tail -> head -> tail :. head
:. (Int
yLen Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (DIM0 :. Int) -> Int -> DIM2
forall tail head. tail -> head -> tail :. head
:. Int
x
          | Bool
otherwise   = DIM0
sh DIM0 -> Int -> DIM0 :. Int
forall tail head. tail -> head -> tail :. head
:. Int
y          (DIM0 :. Int) -> Int -> DIM2
forall tail head. tail -> head -> tail :. head
:. Int
x



-- | Make a 2D partitioned array from two others, one to produce the elements
--   in the internal region, and one to produce elements in the border region.
--   The two arrays must have the same extent.
--   The border must be the same width on all sides.
--
makeBordered2
        :: (Source r1 a, Source r2 a)
        => DIM2                 -- ^ Extent of array.
        -> Int                  -- ^ Width of border.
        -> Array r1 DIM2 a      -- ^ Array for internal elements.
        -> Array r2 DIM2 a      -- ^ Array for border elements.
        -> Array (P r1 (P r2 (P r2 (P r2 (P r2 X))))) DIM2 a

{-# INLINE makeBordered2 #-}
makeBordered2 :: DIM2
-> Int
-> Array r1 DIM2 a
-> Array r2 DIM2 a
-> Array (P r1 (P r2 (P r2 (P r2 (P r2 X))))) DIM2 a
makeBordered2 sh :: DIM2
sh@(DIM0
_ :. Int
aHeight :. Int
aWidth) Int
bWidth Array r1 DIM2 a
arrInternal Array r2 DIM2 a
arrBorder
 = ()
checkDims ()
-> Array (P r1 (P r2 (P r2 (P r2 (P r2 X))))) DIM2 a
-> Array (P r1 (P r2 (P r2 (P r2 (P r2 X))))) DIM2 a
`seq` 
   let
        -- minimum and maximum indicies of values in the inner part of the image.
        !inX :: Int
inX            = Int
bWidth
        !inY :: Int
inY            = Int
bWidth
        !inW :: Int
inW            = Int
aWidth  Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
bWidth 
        !inH :: Int
inH            = Int
aHeight Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
bWidth

        inInternal :: DIM2 -> Bool
inInternal (DIM0
Z :. Int
y :. Int
x)
                =  Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
inX Bool -> Bool -> Bool
&& Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< (Int
inX Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
inW)
                Bool -> Bool -> Bool
&& Int
y Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
inY Bool -> Bool -> Bool
&& Int
y Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< (Int
inY Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
inH)
        {-# INLINE inInternal #-}

        inBorder :: DIM2 -> Bool
inBorder        = Bool -> Bool
not (Bool -> Bool) -> (DIM2 -> Bool) -> DIM2 -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DIM2 -> Bool
inInternal
        {-# INLINE inBorder #-}

   in   
    --  internal region
        DIM2
-> Range DIM2
-> Array r1 DIM2 a
-> Array (P r2 (P r2 (P r2 (P r2 X)))) DIM2 a
-> Array (P r1 (P r2 (P r2 (P r2 (P r2 X))))) DIM2 a
forall r1 r2 sh e.
sh
-> Range sh
-> Array r1 sh e
-> Array r2 sh e
-> Array (P r1 r2) sh e
APart DIM2
sh (DIM2 -> DIM2 -> (DIM2 -> Bool) -> Range DIM2
forall sh. sh -> sh -> (sh -> Bool) -> Range sh
Range (DIM0
Z DIM0 -> Int -> DIM0 :. Int
forall tail head. tail -> head -> tail :. head
:. Int
inY     (DIM0 :. Int) -> Int -> DIM2
forall tail head. tail -> head -> tail :. head
:. Int
inX)       (DIM0
Z DIM0 -> Int -> DIM0 :. Int
forall tail head. tail -> head -> tail :. head
:. Int
inH (DIM0 :. Int) -> Int -> DIM2
forall tail head. tail -> head -> tail :. head
:. Int
inW )    DIM2 -> Bool
inInternal) Array r1 DIM2 a
arrInternal

    --  border regions
    (Array (P r2 (P r2 (P r2 (P r2 X)))) DIM2 a
 -> Array (P r1 (P r2 (P r2 (P r2 (P r2 X))))) DIM2 a)
-> Array (P r2 (P r2 (P r2 (P r2 X)))) DIM2 a
-> Array (P r1 (P r2 (P r2 (P r2 (P r2 X))))) DIM2 a
forall a b. (a -> b) -> a -> b
$   DIM2
-> Range DIM2
-> Array r2 DIM2 a
-> Array (P r2 (P r2 (P r2 X))) DIM2 a
-> Array (P r2 (P r2 (P r2 (P r2 X)))) DIM2 a
forall r1 r2 sh e.
sh
-> Range sh
-> Array r1 sh e
-> Array r2 sh e
-> Array (P r1 r2) sh e
APart DIM2
sh (DIM2 -> DIM2 -> (DIM2 -> Bool) -> Range DIM2
forall sh. sh -> sh -> (sh -> Bool) -> Range sh
Range (DIM0
Z DIM0 -> Int -> DIM0 :. Int
forall tail head. tail -> head -> tail :. head
:. Int
0         (DIM0 :. Int) -> Int -> DIM2
forall tail head. tail -> head -> tail :. head
:. Int
0)         (DIM0
Z DIM0 -> Int -> DIM0 :. Int
forall tail head. tail -> head -> tail :. head
:. Int
bWidth (DIM0 :. Int) -> Int -> DIM2
forall tail head. tail -> head -> tail :. head
:. Int
aWidth) DIM2 -> Bool
inBorder) Array r2 DIM2 a
arrBorder
    (Array (P r2 (P r2 (P r2 X))) DIM2 a
 -> Array (P r2 (P r2 (P r2 (P r2 X)))) DIM2 a)
-> Array (P r2 (P r2 (P r2 X))) DIM2 a
-> Array (P r2 (P r2 (P r2 (P r2 X)))) DIM2 a
forall a b. (a -> b) -> a -> b
$   DIM2
-> Range DIM2
-> Array r2 DIM2 a
-> Array (P r2 (P r2 X)) DIM2 a
-> Array (P r2 (P r2 (P r2 X))) DIM2 a
forall r1 r2 sh e.
sh
-> Range sh
-> Array r1 sh e
-> Array r2 sh e
-> Array (P r1 r2) sh e
APart DIM2
sh (DIM2 -> DIM2 -> (DIM2 -> Bool) -> Range DIM2
forall sh. sh -> sh -> (sh -> Bool) -> Range sh
Range (DIM0
Z DIM0 -> Int -> DIM0 :. Int
forall tail head. tail -> head -> tail :. head
:. Int
inY Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
inH (DIM0 :. Int) -> Int -> DIM2
forall tail head. tail -> head -> tail :. head
:. Int
0)         (DIM0
Z DIM0 -> Int -> DIM0 :. Int
forall tail head. tail -> head -> tail :. head
:. Int
bWidth (DIM0 :. Int) -> Int -> DIM2
forall tail head. tail -> head -> tail :. head
:. Int
aWidth) DIM2 -> Bool
inBorder) Array r2 DIM2 a
arrBorder
    (Array (P r2 (P r2 X)) DIM2 a
 -> Array (P r2 (P r2 (P r2 X))) DIM2 a)
-> Array (P r2 (P r2 X)) DIM2 a
-> Array (P r2 (P r2 (P r2 X))) DIM2 a
forall a b. (a -> b) -> a -> b
$   DIM2
-> Range DIM2
-> Array r2 DIM2 a
-> Array (P r2 X) DIM2 a
-> Array (P r2 (P r2 X)) DIM2 a
forall r1 r2 sh e.
sh
-> Range sh
-> Array r1 sh e
-> Array r2 sh e
-> Array (P r1 r2) sh e
APart DIM2
sh (DIM2 -> DIM2 -> (DIM2 -> Bool) -> Range DIM2
forall sh. sh -> sh -> (sh -> Bool) -> Range sh
Range (DIM0
Z DIM0 -> Int -> DIM0 :. Int
forall tail head. tail -> head -> tail :. head
:. Int
inY       (DIM0 :. Int) -> Int -> DIM2
forall tail head. tail -> head -> tail :. head
:. Int
0)         (DIM0
Z DIM0 -> Int -> DIM0 :. Int
forall tail head. tail -> head -> tail :. head
:. Int
inH    (DIM0 :. Int) -> Int -> DIM2
forall tail head. tail -> head -> tail :. head
:. Int
bWidth) DIM2 -> Bool
inBorder) Array r2 DIM2 a
arrBorder
    (Array (P r2 X) DIM2 a -> Array (P r2 (P r2 X)) DIM2 a)
-> Array (P r2 X) DIM2 a -> Array (P r2 (P r2 X)) DIM2 a
forall a b. (a -> b) -> a -> b
$   DIM2
-> Range DIM2
-> Array r2 DIM2 a
-> Array X DIM2 a
-> Array (P r2 X) DIM2 a
forall r1 r2 sh e.
sh
-> Range sh
-> Array r1 sh e
-> Array r2 sh e
-> Array (P r1 r2) sh e
APart DIM2
sh (DIM2 -> DIM2 -> (DIM2 -> Bool) -> Range DIM2
forall sh. sh -> sh -> (sh -> Bool) -> Range sh
Range (DIM0
Z DIM0 -> Int -> DIM0 :. Int
forall tail head. tail -> head -> tail :. head
:. Int
inY       (DIM0 :. Int) -> Int -> DIM2
forall tail head. tail -> head -> tail :. head
:. Int
inX Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
inW) (DIM0
Z DIM0 -> Int -> DIM0 :. Int
forall tail head. tail -> head -> tail :. head
:. Int
inH    (DIM0 :. Int) -> Int -> DIM2
forall tail head. tail -> head -> tail :. head
:. Int
bWidth) DIM2 -> Bool
inBorder) Array r2 DIM2 a
arrBorder
    (Array X DIM2 a -> Array (P r2 X) DIM2 a)
-> Array X DIM2 a -> Array (P r2 X) DIM2 a
forall a b. (a -> b) -> a -> b
$   DIM2 -> Array X DIM2 a
forall sh e. sh -> Array X sh e
AUndefined DIM2
sh

 where
        checkDims :: ()
checkDims
         = if (Array r1 DIM2 a -> DIM2
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r1 DIM2 a
arrInternal) DIM2 -> DIM2 -> Bool
forall a. Eq a => a -> a -> Bool
== (Array r2 DIM2 a -> DIM2
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r2 DIM2 a
arrBorder)
                then ()
                else [Char] -> ()
forall a. HasCallStack => [Char] -> a
error [Char]
"makeBordered2: internal and border arrays have different extents"
        {-# NOINLINE checkDims #-}
        --  NOINLINE because we don't want the branch in the core code.