module Data.CompactMap.Generic
    ( CompactMap
    , fromList
    , toVector
    , lookup
    , getLE
    , getIndex
    ) where

import qualified Data.Vector.Generic as V
import GHC.Exts (sortWith)
import Control.Applicative
import Prelude hiding (lookup)

data CompactMap vec k v = CompactMap { forall (vec :: * -> *) k v. CompactMap vec k v -> vec v
getMap :: vec v
                                     , forall (vec :: * -> *) k v. CompactMap vec k v -> v -> k
_gkf   :: v -> k
                                     }

instance (V.Vector vec v, Show v) => Show (CompactMap vec k v) where
    show :: CompactMap vec k v -> String
show = [v] -> String
forall a. Show a => a -> String
show ([v] -> String)
-> (CompactMap vec k v -> [v]) -> CompactMap vec k v -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. vec v -> [v]
forall (v :: * -> *) a. Vector v a => v a -> [a]
V.toList (vec v -> [v])
-> (CompactMap vec k v -> vec v) -> CompactMap vec k v -> [v]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CompactMap vec k v -> vec v
forall (vec :: * -> *) k v. CompactMap vec k v -> vec v
getMap

fromList :: (V.Vector vec v, Ord k) => [v] -> (v -> k) -> CompactMap vec k v
fromList :: forall (vec :: * -> *) v k.
(Vector vec v, Ord k) =>
[v] -> (v -> k) -> CompactMap vec k v
fromList [v]
lst v -> k
f = vec v -> (v -> k) -> CompactMap vec k v
forall (vec :: * -> *) k v. vec v -> (v -> k) -> CompactMap vec k v
CompactMap ([v] -> vec v
forall (v :: * -> *) a. Vector v a => [a] -> v a
V.fromList ([v] -> vec v) -> [v] -> vec v
forall a b. (a -> b) -> a -> b
$ (v -> k) -> [v] -> [v]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortWith v -> k
f [v]
lst) v -> k
f

toVector :: V.Vector vec v => CompactMap vec k v -> vec v
toVector :: forall (vec :: * -> *) v k.
Vector vec v =>
CompactMap vec k v -> vec v
toVector (CompactMap vec v
v v -> k
_) = vec v
v

lookup :: (V.Vector vec v, Ord k) => k -> CompactMap vec k v -> Maybe v
lookup :: forall (vec :: * -> *) v k.
(Vector vec v, Ord k) =>
k -> CompactMap vec k v -> Maybe v
lookup k
k cm :: CompactMap vec k v
cm@(CompactMap vec v
_ v -> k
f) = k -> CompactMap vec k v -> Maybe (Int, v)
forall (vec :: * -> *) v k.
(Vector vec v, Ord k) =>
k -> CompactMap vec k v -> Maybe (Int, v)
getLE k
k CompactMap vec k v
cm Maybe (Int, v) -> ((Int, v) -> Maybe v) -> Maybe v
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \(Int
_,v
r) -> if v -> k
f v
r k -> k -> Bool
forall a. Eq a => a -> a -> Bool
== k
k
                                                            then v -> Maybe v
forall a. a -> Maybe a
Just v
r
                                                            else Maybe v
forall a. Maybe a
Nothing

getLE :: (V.Vector vec v, Ord k) => k -> CompactMap vec k v -> Maybe (Int, v)
getLE :: forall (vec :: * -> *) v k.
(Vector vec v, Ord k) =>
k -> CompactMap vec k v -> Maybe (Int, v)
getLE k
k (CompactMap vec v
lst v -> k
f) = Int -> Int -> Maybe (Int, v)
go Int
0 (Int
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
    where
        sz :: Int
sz = vec v -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
V.length vec v
lst
        go :: Int -> Int -> Maybe (Int, v)
go Int
l Int
h
            | Int
m Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
l     = Maybe (Int, v)
forall a. Maybe a
Nothing
            | Int
m Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
h     = Maybe (Int, v)
forall a. Maybe a
Nothing
            | k
k' k -> k -> Bool
forall a. Eq a => a -> a -> Bool
== k
k   = (Int, v) -> Maybe (Int, v)
forall a. a -> Maybe a
Just (Int
m, v
x)
            | Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
h    = Maybe (Int, v)
checkLower Maybe (Int, v) -> Maybe (Int, v) -> Maybe (Int, v)
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> (Int, v) -> Maybe (Int, v)
forall a. a -> Maybe a
Just (Int
m, v
x)
            | k
k k -> k -> Bool
forall a. Ord a => a -> a -> Bool
< k
k'    = Maybe (Int, v)
checkLower Maybe (Int, v) -> Maybe (Int, v) -> Maybe (Int, v)
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Int -> Int -> Maybe (Int, v)
go Int
l (Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
            | Bool
otherwise = Int -> Int -> Maybe (Int, v)
go (Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
h
            where
                checkLower :: Maybe (Int, v)
checkLower
                    | Int
m Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 = Maybe (Int, v)
forall a. Maybe a
Nothing
                    | k
k k -> k -> Bool
forall a. Ord a => a -> a -> Bool
> k
k_1 = (Int, v) -> Maybe (Int, v)
forall a. a -> Maybe a
Just (Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1, v
x_1)
                    | Bool
otherwise = Maybe (Int, v)
forall a. Maybe a
Nothing
                m :: Int
m = (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
h) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
                x :: v
x = vec v
lst vec v -> Int -> v
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
V.! Int
m
                k' :: k
k' = v -> k
f v
x
                x_1 :: v
x_1 = vec v
lst vec v -> Int -> v
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
V.! (Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
                k_1 :: k
k_1 = v -> k
f v
x_1

getIndex :: V.Vector vec v => Int -> CompactMap vec k v -> Maybe v
getIndex :: forall (vec :: * -> *) v k.
Vector vec v =>
Int -> CompactMap vec k v -> Maybe v
getIndex Int
i (CompactMap vec v
lst v -> k
_) = vec v
lst vec v -> Int -> Maybe v
forall (v :: * -> *) a. Vector v a => v a -> Int -> Maybe a
V.!? Int
i