-- Copyright 2013 Evan Laforge
-- This program is distributed under the terms of the GNU General Public
-- License 3.0, see COPYING or http://www.gnu.org/licenses/gpl-3.0.txt

-- | Array utilities.
--
-- IArray is awkward and I don't like it.  @vector@ is nicer but it's a big
-- library so I don't want to depend on it unless I have more than one data
-- structure using arrays.
module Util.Array where
import Prelude hiding (null, length)
import qualified Data.Array.IArray as IArray
import Data.Array.IArray ((!))
import qualified Data.List as List


type Array = IArray.Array Int

empty :: Array a
empty :: forall a. Array a
empty = forall a. [a] -> Array a
from_list []

null :: Array a -> Bool
null :: forall a. Array a -> Bool
null = (forall a. Eq a => a -> a -> Bool
==Int
0) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Array a -> Int
length

length :: Array a -> Int
length :: forall a. Array a -> Int
length Array a
a = Int
high forall a. Num a => a -> a -> a
- Int
low forall a. Num a => a -> a -> a
+ Int
1
    where (Int
low, Int
high) = forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
IArray.bounds Array a
a

from_list :: [a] -> Array a
from_list :: forall a. [a] -> Array a
from_list [a]
xs = forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [e] -> a i e
IArray.listArray (Int
0, forall (t :: * -> *) a. Foldable t => t a -> Int
List.length [a]
xs forall a. Num a => a -> a -> a
- Int
1) [a]
xs

-- | Like 'IArray.!', except throw a more informative error, with @msg@
-- prepended.
at :: String -> Array a -> Int -> a
at :: forall a. String -> Array a -> Int -> a
at String
msg Array a
a Int
i = Array a
a forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
! forall a. String -> Int -> Array a -> Int
assert_in_bounds String
msg Int
i Array a
a

assert_in_bounds :: String -> Int -> Array a -> Int
assert_in_bounds :: forall a. String -> Int -> Array a -> Int
assert_in_bounds String
msg Int
i Array a
a
    | forall a. Int -> Array a -> Bool
in_bounds Int
i Array a
a = Int
i
    | Bool
otherwise = forall a. HasCallStack => String -> a
error forall a b. (a -> b) -> a -> b
$ String
msg forall a. [a] -> [a] -> [a]
++ String
": index " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
i
        forall a. [a] -> [a] -> [a]
++ String
" out of range " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show (forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
IArray.bounds Array a
a)

-- | Is the given index within the array's bounds?
in_bounds :: Int -> Array a -> Bool
in_bounds :: forall a. Int -> Array a -> Bool
in_bounds Int
i Array a
a = let (Int
low, Int
high) = forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
IArray.bounds Array a
a in Int
low forall a. Ord a => a -> a -> Bool
<= Int
i Bool -> Bool -> Bool
&& Int
i forall a. Ord a => a -> a -> Bool
<= Int
high

-- | Just the array if the index is in bounds.
check :: Int -> Array a -> Maybe (Array a)
check :: forall a. Int -> Array a -> Maybe (Array a)
check Int
i Array a
a
    | forall a. Int -> Array a -> Bool
in_bounds Int
i Array a
a = forall a. a -> Maybe a
Just Array a
a
    | Bool
otherwise = forall a. Maybe a
Nothing

-- ** searching

-- | Find the index of the first element >= the given element in the sorted
-- array.
bsearch :: Ord k => Array k -> k -> Int
bsearch :: forall k. Ord k => Array k -> k -> Int
bsearch = forall k a. (k -> a -> Bool) -> Array a -> k -> Int
bsearch_with forall a. Ord a => a -> a -> Bool
(<=)

bsearch_on :: Ord k => (a -> k) -> Array a -> k -> Int
bsearch_on :: forall k a. Ord k => (a -> k) -> Array a -> k -> Int
bsearch_on a -> k
key Array a
a k
elt = forall k a. (k -> a -> Bool) -> Array a -> k -> Int
bsearch_with (\k
elt a
e1 -> (k
elt forall a. Ord a => a -> a -> Bool
<= a -> k
key a
e1)) Array a
a k
elt

bsearch_with :: (k -> a -> Bool) -> Array a -> k -> Int
bsearch_with :: forall k a. (k -> a -> Bool) -> Array a -> k -> Int
bsearch_with k -> a -> Bool
lte Array a
a k
elt = forall a. (a -> Bool) -> Array a -> Int -> Int -> Int
_do_bsearch (k -> a -> Bool
lte k
elt) Array a
a Int
low (Int
highforall a. Num a => a -> a -> a
+Int
1)
    where (Int
low, Int
high) = forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
IArray.bounds Array a
a

_do_bsearch :: (a -> Bool) -> Array a -> Int -> Int -> Int
_do_bsearch :: forall a. (a -> Bool) -> Array a -> Int -> Int -> Int
_do_bsearch a -> Bool
lte Array a
a Int
low Int
high
    | Int
low forall a. Eq a => a -> a -> Bool
== Int
high = Int
low
    | a -> Bool
lte (Array a
aforall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!Int
mid) = forall a. (a -> Bool) -> Array a -> Int -> Int -> Int
_do_bsearch a -> Bool
lte Array a
a Int
low Int
mid
    | Bool
otherwise = forall a. (a -> Bool) -> Array a -> Int -> Int -> Int
_do_bsearch a -> Bool
lte Array a
a (Int
midforall a. Num a => a -> a -> a
+Int
1) Int
high
    where
    mid :: Int
mid = (Int
low forall a. Num a => a -> a -> a
+ Int
high) forall a. Integral a => a -> a -> a
`div` Int
2