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 = forall a. Show a => a -> String
show forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (v :: * -> *) a. Vector v a => v a -> [a]
V.toList forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 = forall (vec :: * -> *) k v. vec v -> (v -> k) -> CompactMap vec k v
CompactMap (forall (v :: * -> *) a. Vector v a => [a] -> v a
V.fromList forall a b. (a -> b) -> a -> b
$ 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) = 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 forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \(Int
_,v
r) -> if v -> k
f v
r forall a. Eq a => a -> a -> Bool
== k
k
                                                            then forall a. a -> Maybe a
Just v
r
                                                            else 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 forall a. Num a => a -> a -> a
- Int
1)
    where
        sz :: Int
sz = 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 forall a. Ord a => a -> a -> Bool
< Int
l     = forall a. Maybe a
Nothing
            | Int
m forall a. Ord a => a -> a -> Bool
> Int
h     = forall a. Maybe a
Nothing
            | k
k' forall a. Eq a => a -> a -> Bool
== k
k   = forall a. a -> Maybe a
Just (Int
m, v
x)
            | Int
l forall a. Ord a => a -> a -> Bool
>= Int
h    = Maybe (Int, v)
checkLower forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> forall a. a -> Maybe a
Just (Int
m, v
x)
            | k
k forall a. Ord a => a -> a -> Bool
< k
k'    = Maybe (Int, v)
checkLower forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Int -> Int -> Maybe (Int, v)
go Int
l (Int
m forall a. Num a => a -> a -> a
- Int
1)
            | Bool
otherwise = Int -> Int -> Maybe (Int, v)
go (Int
m forall a. Num a => a -> a -> a
+ Int
1) Int
h
            where
                checkLower :: Maybe (Int, v)
checkLower
                    | Int
m forall a. Ord a => a -> a -> Bool
< Int
1 = forall a. Maybe a
Nothing
                    | k
k forall a. Ord a => a -> a -> Bool
> k
k_1 = forall a. a -> Maybe a
Just (Int
m forall a. Num a => a -> a -> a
- Int
1, v
x_1)
                    | Bool
otherwise = forall a. Maybe a
Nothing
                m :: Int
m = (Int
l forall a. Num a => a -> a -> a
+ Int
h) forall a. Integral a => a -> a -> a
`div` Int
2
                x :: v
x = vec v
lst 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 forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
V.! (Int
m 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 forall (v :: * -> *) a. Vector v a => v a -> Int -> Maybe a
V.!? Int
i