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

module Data.Map.Subset.Lazy.Internal
  ( Map
  , lookup
  , empty
  , singleton
  , antisingleton
  , fromPolarities
  , toList
  , fromList
  ) where

import Prelude hiding (lookup,concat)

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

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

-- 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.
--
-- Unlike the strict variant, which imposes an Eq constraint on
-- values, the lazy variant is able to have a Functor instance.
data Map k v
  = MapElement k (Map k v) (Map k v)
  | MapValue v
  | MapEmpty
  deriving (forall a b. a -> Map k b -> Map k a
forall a b. (a -> b) -> Map k a -> Map k b
forall k a b. a -> Map k b -> Map k a
forall k a b. (a -> b) -> Map k a -> Map k b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> Map k b -> Map k a
$c<$ :: forall k a b. a -> Map k b -> Map k a
fmap :: forall a b. (a -> b) -> Map k a -> Map k b
$cfmap :: forall k a b. (a -> b) -> Map k a -> Map k b
Functor,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, Ord k) => Semigroup (Map k v) where
  <> :: Map k v -> Map k v -> Map k v
(<>) = forall k v. (Semigroup v, Ord k) => Map k v -> Map k v -> Map k v
append

instance (Semigroup 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.
(ContiguousU arr, Element arr k) =>
Map k v -> [(Set arr k, v)]
toList Map k v
xs))

toList :: (ContiguousU arr, Element arr k)
  => Map k v
  -> [(Set arr k,v)]
toList :: forall (arr :: * -> *) k v.
(ContiguousU arr, Element arr k) =>
Map k v -> [(Set arr k, v)]
toList = forall (arr :: * -> *) k v b.
(ContiguousU 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 :: (ContiguousU arr, Element arr k, Ord k)
  => [(Set arr k,v)]
  -> Map k v
fromList :: forall (arr :: * -> *) k v.
(ContiguousU arr, Element arr k, Ord k) =>
[(Set arr k, v)] -> Map k v
fromList = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. First a -> a
getFirst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k v. (Ord k, Semigroup 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 (arr :: * -> *) k v.
(ContiguousU 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)
  => [Map k v]
  -> Map k v
concat :: forall k v. (Ord k, Semigroup 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, 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 :: (ContiguousU arr, Element arr k)
  => (Set arr k -> v -> b -> b)
  -> b
  -> Map k v
  -> b
foldrWithKey :: forall (arr :: * -> *) k v b.
(ContiguousU 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 :: (ContiguousU arr, Element arr k)
  => Set arr k
  -> v
  -> Map k v
singleton :: forall (arr :: * -> *) k v.
(ContiguousU 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

antisingleton :: (ContiguousU arr, Element arr k)
  => Set arr k
  -> v
  -> Map k v
antisingleton :: forall (arr :: * -> *) k v.
(ContiguousU arr, Element arr k) =>
Set arr k -> v -> Map k v
antisingleton 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 forall k v. Map k v
empty Map k v
m) (forall k v. v -> Map k v
MapValue v
v) Set arr k
s

fromPolarities :: (ContiguousU karr, Element karr k)
  => M.Map karr Array k Bool
  -> v
  -> Map k v
fromPolarities :: forall (karr :: * -> *) k v.
(ContiguousU karr, Element karr k) =>
Map karr Array k Bool -> v -> Map k v
fromPolarities Map karr Array k Bool
s v
v = forall (karr :: * -> *) k (varr :: * -> *) v b.
(ContiguousU karr, Element karr k, ContiguousU varr,
 Element varr v) =>
(k -> v -> b -> b) -> b -> Map karr varr k v -> b
M.foldrWithKey
  (\k
k Bool
p Map k v
m -> forall k v. k -> Map k v -> Map k v -> Map k v
MapElement k
k (forall a. a -> a -> Bool -> a
bool forall k v. Map k v
empty Map k v
m Bool
p) (forall a. a -> a -> Bool -> a
bool Map k v
m forall k v. Map k v
empty Bool
p))
  (forall k v. v -> Map k v
MapValue v
v) Map karr Array k Bool
s

lookup :: forall arr k v. (Ord k, ContiguousU arr, Element arr k)
  => Set arr k
  -> Map k v
  -> Maybe v
{-# INLINABLE lookup #-}
lookup :: forall (arr :: * -> *) k v.
(Ord k, ContiguousU 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 => (v -> v) -> v -> Map k v -> Map k v
augment :: forall k v. Eq k => (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 => (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 => (v -> v) -> v -> Map k v -> Map k v
augment v -> v
f v
v Map k v
absent
   in 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, Ord k) => Map k v -> Map k v -> Map k v
append :: forall k v. (Semigroup 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 => (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 => (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 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 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 forall k v. k -> Map k v -> Map k v -> Map k v
MapElement k
elemY Map k v
present Map k v
absent