{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UnboxedTuples #-}

module Data.Map.Subset.Strict.Internal
  ( Map
  , lookup
  , empty
  , singleton
  , toList
  , fromList
  ) where

import Prelude hiding (lookup,concat)

import Data.Primitive.Contiguous (Contiguous,Element)
import Data.Set.Internal (Set(..))
import Data.Bifunctor (first)
import Data.Semigroup (Semigroup,(<>),First(..))

import qualified Data.Primitive.Contiguous as A
import qualified Data.Set.Internal as S
import qualified Data.Set.Lifted.Internal as SL
import qualified Data.Semigroup as SG
import qualified Prelude as P
import qualified Data.Foldable as F

-- There are two invariants for Map.
--
-- 1. The children of any Map may only contain keys that are
--    greater than the key in their parent.
-- 2. A parent's two children must not be equal.
--
-- This type cannot be a Functor since it needs to uses
-- an Eq instance for a kind of simple compression.
data Map k v
  = MapElement !k !(Map k v) !(Map k v)
  | MapValue !v
  | MapEmpty
  deriving (Map k v -> Map k v -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall k v. (Eq k, Eq v) => Map k v -> Map k v -> Bool
/= :: Map k v -> Map k v -> Bool
$c/= :: forall k v. (Eq k, Eq v) => Map k v -> Map k v -> Bool
== :: Map k v -> Map k v -> Bool
$c== :: forall k v. (Eq k, Eq v) => Map k v -> Map k v -> Bool
Eq,Map k v -> Map k v -> Bool
Map k v -> Map k v -> Ordering
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall {k} {v}. (Ord k, Ord v) => Eq (Map k v)
forall k v. (Ord k, Ord v) => Map k v -> Map k v -> Bool
forall k v. (Ord k, Ord v) => Map k v -> Map k v -> Ordering
forall k v. (Ord k, Ord v) => Map k v -> Map k v -> Map k v
min :: Map k v -> Map k v -> Map k v
$cmin :: forall k v. (Ord k, Ord v) => Map k v -> Map k v -> Map k v
max :: Map k v -> Map k v -> Map k v
$cmax :: forall k v. (Ord k, Ord v) => Map k v -> Map k v -> Map k v
>= :: Map k v -> Map k v -> Bool
$c>= :: forall k v. (Ord k, Ord v) => Map k v -> Map k v -> Bool
> :: Map k v -> Map k v -> Bool
$c> :: forall k v. (Ord k, Ord v) => Map k v -> Map k v -> Bool
<= :: Map k v -> Map k v -> Bool
$c<= :: forall k v. (Ord k, Ord v) => Map k v -> Map k v -> Bool
< :: Map k v -> Map k v -> Bool
$c< :: forall k v. (Ord k, Ord v) => Map k v -> Map k v -> Bool
compare :: Map k v -> Map k v -> Ordering
$ccompare :: forall k v. (Ord k, Ord v) => Map k v -> Map k v -> Ordering
Ord)

instance (Semigroup v, Eq v, Ord k) => Semigroup (Map k v) where
  <> :: Map k v -> Map k v -> Map k v
(<>) = forall k v.
(Semigroup v, Eq v, Ord k) =>
Map k v -> Map k v -> Map k v
append

instance (Semigroup v, Eq v, Ord k) => Monoid (Map k v) where
  mempty :: Map k v
mempty = forall k v. Map k v
empty
  mappend :: Map k v -> Map k v -> Map k v
mappend = forall a. Semigroup a => a -> a -> a
(SG.<>)
  -- mconcat = concat 

instance (Show k, Show v) => Show (Map k v) where
  showsPrec :: Int -> Map k v -> ShowS
showsPrec Int
p Map k v
xs = Bool -> ShowS -> ShowS
showParen (Int
p forall a. Ord a => a -> a -> Bool
> Int
10) forall a b. (a -> b) -> a -> b
$
    String -> ShowS
showString String
"fromList " forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> ShowS
shows (forall a b. (a -> b) -> [a] -> [b]
P.map (forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first forall a. Set Array a -> Set a
SL.Set) (forall (arr :: * -> *) k v.
(Contiguous arr, Element arr k) =>
Map k v -> [(Set arr k, v)]
toList Map k v
xs))

-- the functon f must satisfy the following:
-- a == b <=> f a == f b
unsafeMapValues :: (v -> w) -> Map k v -> Map k w
unsafeMapValues :: forall v w k. (v -> w) -> Map k v -> Map k w
unsafeMapValues v -> w
f = Map k v -> Map k w
go where
  go :: Map k v -> Map k w
go Map k v
MapEmpty = forall k v. Map k v
MapEmpty
  go (MapValue v
x) = forall k v. v -> Map k v
MapValue (v -> w
f v
x)
  go (MapElement k
k Map k v
present Map k v
absent) =
    forall k v. k -> Map k v -> Map k v -> Map k v
MapElement k
k (Map k v -> Map k w
go Map k v
present) (Map k v -> Map k w
go Map k v
absent)

toList :: (Contiguous arr, Element arr k)
  => Map k v
  -> [(Set arr k,v)]
toList :: forall (arr :: * -> *) k v.
(Contiguous arr, Element arr k) =>
Map k v -> [(Set arr k, v)]
toList = forall (arr :: * -> *) k v b.
(Contiguous arr, Element arr k) =>
(Set arr k -> v -> b -> b) -> b -> Map k v -> b
foldrWithKey (\Set arr k
k v
v [(Set arr k, v)]
xs -> (Set arr k
k,v
v) forall a. a -> [a] -> [a]
: [(Set arr k, v)]
xs) []

fromList :: (Contiguous arr, Element arr k, Ord k, Eq v)
  => [(Set arr k,v)]
  -> Map k v
fromList :: forall (arr :: * -> *) k v.
(Contiguous arr, Element arr k, Ord k, Eq v) =>
[(Set arr k, v)] -> Map k v
fromList = forall v w k. (v -> w) -> Map k v -> Map k w
unsafeMapValues forall a. First a -> a
getFirst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k v. (Ord k, Semigroup v, Eq v) => [Map k v] -> Map k v
concat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
P.map (\(Set arr k
s,v
v) -> forall v (arr :: * -> *) k.
(Eq v, Contiguous arr, Element arr k) =>
Set arr k -> v -> Map k v
singleton Set arr k
s (forall a. a -> First a
First v
v))

concat :: (Ord k,Semigroup v,Eq v)
  => [Map k v]
  -> Map k v
concat :: forall k v. (Ord k, Semigroup v, Eq v) => [Map k v] -> Map k v
concat = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
F.foldl' (\Map k v
r Map k v
x -> forall k v.
(Semigroup v, Eq v, Ord k) =>
Map k v -> Map k v -> Map k v
append Map k v
r Map k v
x) forall k v. Map k v
empty

foldrWithKey :: (Contiguous arr, Element arr k)
  => (Set arr k -> v -> b -> b)
  -> b
  -> Map k v
  -> b
foldrWithKey :: forall (arr :: * -> *) k v b.
(Contiguous arr, Element arr k) =>
(Set arr k -> v -> b -> b) -> b -> Map k v -> b
foldrWithKey Set arr k -> v -> b -> b
f b
b0 = Int -> [k] -> b -> Map k v -> b
go Int
0 [] b
b0 where
  go :: Int -> [k] -> b -> Map k v -> b
go !Int
_ ![k]
_ b
b Map k v
MapEmpty = b
b
  go !Int
n ![k]
xs b
b (MapValue v
v) = Set arr k -> v -> b -> b
f (forall (arr :: * -> *) a. arr a -> Set arr a
Set (forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
Int -> [a] -> arr a
A.unsafeFromListReverseN Int
n [k]
xs)) v
v b
b
  go !Int
n ![k]
xs b
b (MapElement k
k Map k v
present Map k v
absent) =
    Int -> [k] -> b -> Map k v -> b
go (Int
n forall a. Num a => a -> a -> a
+ Int
1) (k
k forall a. a -> [a] -> [a]
: [k]
xs) (Int -> [k] -> b -> Map k v -> b
go Int
n [k]
xs b
b Map k v
absent) Map k v
present

empty :: Map k v
empty :: forall k v. Map k v
empty = forall k v. Map k v
MapEmpty

singleton :: (Eq v, Contiguous arr, Element arr k)
  => Set arr k
  -> v
  -> Map k v
singleton :: forall v (arr :: * -> *) k.
(Eq v, Contiguous arr, Element arr k) =>
Set arr k -> v -> Map k v
singleton Set arr k
s v
v = forall (arr :: * -> *) a b.
(Contiguous arr, Element arr a) =>
(a -> b -> b) -> b -> Set arr a -> b
S.foldr (\k
k Map k v
m -> forall k v. k -> Map k v -> Map k v -> Map k v
MapElement k
k Map k v
m forall k v. Map k v
empty) (forall k v. v -> Map k v
MapValue v
v) Set arr k
s
  
lookup :: forall arr k v. (Ord k, Contiguous arr, Element arr k)
  => Set arr k
  -> Map k v
  -> Maybe v
{-# INLINABLE lookup #-}
lookup :: forall (arr :: * -> *) k v.
(Ord k, Contiguous arr, Element arr k) =>
Set arr k -> Map k v -> Maybe v
lookup (Set arr k
arr) = Int -> Map k v -> Maybe v
go Int
0 where
  !sz :: Int
sz = forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
A.size arr k
arr
  go :: Int -> Map k v -> Maybe v
  go :: Int -> Map k v -> Maybe v
go !Int
_ Map k v
MapEmpty = forall a. Maybe a
Nothing
  go !Int
_ (MapValue v
v) = forall a. a -> Maybe a
Just v
v
  go !Int
ix (MapElement k
element Map k v
present Map k v
absent) =
    Int -> k -> Map k v -> Map k v -> Maybe v
choose Int
ix k
element Map k v
present Map k v
absent
  choose :: Int -> k -> Map k v -> Map k v -> Maybe v
  choose :: Int -> k -> Map k v -> Map k v -> Maybe v
choose !Int
ix k
element Map k v
present Map k v
absent = if Int
ix forall a. Ord a => a -> a -> Bool
< Int
sz
    then 
      let (# k
k #) = forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> (# b #)
A.index# arr k
arr Int
ix
       in case forall a. Ord a => a -> a -> Ordering
compare k
k k
element of
            Ordering
EQ -> Int -> Map k v -> Maybe v
go (Int
ix forall a. Num a => a -> a -> a
+ Int
1) Map k v
present
            Ordering
LT -> Int -> k -> Map k v -> Map k v -> Maybe v
choose (Int
ix forall a. Num a => a -> a -> a
+ Int
1) k
element Map k v
present Map k v
absent
            Ordering
GT -> Int -> Map k v -> Maybe v
go Int
ix Map k v
absent
    else forall k v. Map k v -> Maybe v
followAbsent Map k v
absent

followAbsent :: Map k v -> Maybe v
followAbsent :: forall k v. Map k v -> Maybe v
followAbsent (MapElement k
_ Map k v
_ Map k v
x) = forall k v. Map k v -> Maybe v
followAbsent Map k v
x
followAbsent (MapValue v
v) = forall a. a -> Maybe a
Just v
v
followAbsent Map k v
MapEmpty = forall a. Maybe a
Nothing

augment :: (Eq k, Eq v) => (v -> v) -> v -> Map k v -> Map k v
augment :: forall k v. (Eq k, Eq v) => (v -> v) -> v -> Map k v -> Map k v
augment v -> v
_ v
v Map k v
MapEmpty = forall k v. v -> Map k v
MapValue v
v
augment v -> v
f v
_ (MapValue v
x) = forall k v. v -> Map k v
MapValue (v -> v
f v
x)
augment v -> v
f v
v (MapElement k
k Map k v
present Map k v
absent) =
  let present' :: Map k v
present' = forall k v. (Eq k, Eq v) => (v -> v) -> v -> Map k v -> Map k v
augment v -> v
f v
v Map k v
present
      absent' :: Map k v
absent' = forall k v. (Eq k, Eq v) => (v -> v) -> v -> Map k v -> Map k v
augment v -> v
f v
v Map k v
absent
   in if Map k v
present' forall a. Eq a => a -> a -> Bool
== Map k v
absent'
        then Map k v
present' 
        else forall k v. k -> Map k v -> Map k v -> Map k v
MapElement k
k Map k v
present' Map k v
absent'

append :: forall k v. (Semigroup v, Eq v, Ord k) => Map k v -> Map k v -> Map k v
append :: forall k v.
(Semigroup v, Eq v, Ord k) =>
Map k v -> Map k v -> Map k v
append = Map k v -> Map k v -> Map k v
go where
  go :: Map k v -> Map k v -> Map k v
  go :: Map k v -> Map k v -> Map k v
go Map k v
MapEmpty Map k v
m = Map k v
m
  go (MapValue v
x) (MapValue v
y) = forall k v. v -> Map k v
MapValue (v
x forall a. Semigroup a => a -> a -> a
<> v
y)
  go (MapValue v
x) Map k v
MapEmpty = forall k v. v -> Map k v
MapValue v
x
  go (MapValue v
x) (MapElement k
elemY Map k v
presentY Map k v
absentY) =
    forall k v. (Eq k, Eq v) => (v -> v) -> v -> Map k v -> Map k v
augment (v
x forall a. Semigroup a => a -> a -> a
SG.<>) v
x (forall k v. k -> Map k v -> Map k v -> Map k v
MapElement k
elemY Map k v
presentY Map k v
absentY)
  go (MapElement k
elemX Map k v
presentX Map k v
absentX) Map k v
MapEmpty =
    forall k v. k -> Map k v -> Map k v -> Map k v
MapElement k
elemX Map k v
presentX Map k v
absentX
  go (MapElement k
elemX Map k v
presentX Map k v
absentX) (MapValue v
y) =
    forall k v. (Eq k, Eq v) => (v -> v) -> v -> Map k v -> Map k v
augment (forall a. Semigroup a => a -> a -> a
SG.<> v
y) v
y (forall k v. k -> Map k v -> Map k v -> Map k v
MapElement k
elemX Map k v
presentX Map k v
absentX)
  go (MapElement k
elemX Map k v
presentX Map k v
absentX) (MapElement k
elemY Map k v
presentY Map k v
absentY) = case forall a. Ord a => a -> a -> Ordering
compare k
elemX k
elemY of
    Ordering
EQ -> 
      let present :: Map k v
present = Map k v -> Map k v -> Map k v
go Map k v
presentX Map k v
presentY
          absent :: Map k v
absent = Map k v -> Map k v -> Map k v
go Map k v
absentX Map k v
absentY
       in if Map k v
present forall a. Eq a => a -> a -> Bool
== Map k v
absent
            then Map k v
present
            else forall k v. k -> Map k v -> Map k v -> Map k v
MapElement k
elemX Map k v
present Map k v
absent
    Ordering
LT ->
      let present :: Map k v
present = Map k v -> Map k v -> Map k v
go Map k v
presentX (forall k v. k -> Map k v -> Map k v -> Map k v
MapElement k
elemY Map k v
presentY Map k v
absentY)
          absent :: Map k v
absent = Map k v -> Map k v -> Map k v
go Map k v
absentX (forall k v. k -> Map k v -> Map k v -> Map k v
MapElement k
elemY Map k v
presentY Map k v
absentY)
       in if Map k v
present forall a. Eq a => a -> a -> Bool
== Map k v
absent
            then Map k v
present
            else forall k v. k -> Map k v -> Map k v -> Map k v
MapElement k
elemX Map k v
present Map k v
absent
    Ordering
GT ->
      let present :: Map k v
present = Map k v -> Map k v -> Map k v
go (forall k v. k -> Map k v -> Map k v -> Map k v
MapElement k
elemX Map k v
presentX Map k v
absentX) Map k v
presentY
          absent :: Map k v
absent = Map k v -> Map k v -> Map k v
go (forall k v. k -> Map k v -> Map k v -> Map k v
MapElement k
elemX Map k v
presentX Map k v
absentX) Map k v
absentY
       in if Map k v
present forall a. Eq a => a -> a -> Bool
== Map k v
absent
            then Map k v
present
            else forall k v. k -> Map k v -> Map k v -> Map k v
MapElement k
elemY Map k v
present Map k v
absent