module Data.Multwomap (
    Multwomap()
    , empty
    , null
    , size
    , union
    , insert
    , arbitraryKey
    , extract
) where



import           Data.Map (Map)
import qualified Data.Map as M
import           Prelude  hiding (null)

import Util



data OneTwo a = One a | Two a a deriving (OneTwo a -> OneTwo a -> Bool
(OneTwo a -> OneTwo a -> Bool)
-> (OneTwo a -> OneTwo a -> Bool) -> Eq (OneTwo a)
forall a. Eq a => OneTwo a -> OneTwo a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => OneTwo a -> OneTwo a -> Bool
== :: OneTwo a -> OneTwo a -> Bool
$c/= :: forall a. Eq a => OneTwo a -> OneTwo a -> Bool
/= :: OneTwo a -> OneTwo a -> Bool
Eq, Eq (OneTwo a)
Eq (OneTwo a)
-> (OneTwo a -> OneTwo a -> Ordering)
-> (OneTwo a -> OneTwo a -> Bool)
-> (OneTwo a -> OneTwo a -> Bool)
-> (OneTwo a -> OneTwo a -> Bool)
-> (OneTwo a -> OneTwo a -> Bool)
-> (OneTwo a -> OneTwo a -> OneTwo a)
-> (OneTwo a -> OneTwo a -> OneTwo a)
-> Ord (OneTwo a)
OneTwo a -> OneTwo a -> Bool
OneTwo a -> OneTwo a -> Ordering
OneTwo a -> OneTwo a -> OneTwo a
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 {a}. Ord a => Eq (OneTwo a)
forall a. Ord a => OneTwo a -> OneTwo a -> Bool
forall a. Ord a => OneTwo a -> OneTwo a -> Ordering
forall a. Ord a => OneTwo a -> OneTwo a -> OneTwo a
$ccompare :: forall a. Ord a => OneTwo a -> OneTwo a -> Ordering
compare :: OneTwo a -> OneTwo a -> Ordering
$c< :: forall a. Ord a => OneTwo a -> OneTwo a -> Bool
< :: OneTwo a -> OneTwo a -> Bool
$c<= :: forall a. Ord a => OneTwo a -> OneTwo a -> Bool
<= :: OneTwo a -> OneTwo a -> Bool
$c> :: forall a. Ord a => OneTwo a -> OneTwo a -> Bool
> :: OneTwo a -> OneTwo a -> Bool
$c>= :: forall a. Ord a => OneTwo a -> OneTwo a -> Bool
>= :: OneTwo a -> OneTwo a -> Bool
$cmax :: forall a. Ord a => OneTwo a -> OneTwo a -> OneTwo a
max :: OneTwo a -> OneTwo a -> OneTwo a
$cmin :: forall a. Ord a => OneTwo a -> OneTwo a -> OneTwo a
min :: OneTwo a -> OneTwo a -> OneTwo a
Ord, Int -> OneTwo a -> ShowS
[OneTwo a] -> ShowS
OneTwo a -> String
(Int -> OneTwo a -> ShowS)
-> (OneTwo a -> String) -> ([OneTwo a] -> ShowS) -> Show (OneTwo a)
forall a. Show a => Int -> OneTwo a -> ShowS
forall a. Show a => [OneTwo a] -> ShowS
forall a. Show a => OneTwo a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> OneTwo a -> ShowS
showsPrec :: Int -> OneTwo a -> ShowS
$cshow :: forall a. Show a => OneTwo a -> String
show :: OneTwo a -> String
$cshowList :: forall a. Show a => [OneTwo a] -> ShowS
showList :: [OneTwo a] -> ShowS
Show)

-- | A multimap where each key can have at most two entries. Note that this is an
-- unsafe data structure: attempting to add a key twice will crash the program.
newtype Multwomap k v = Multwomap (Map k (OneTwo v)) deriving (Multwomap k v -> Multwomap k v -> Bool
(Multwomap k v -> Multwomap k v -> Bool)
-> (Multwomap k v -> Multwomap k v -> Bool) -> Eq (Multwomap k v)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall k v. (Eq k, Eq v) => Multwomap k v -> Multwomap k v -> Bool
$c== :: forall k v. (Eq k, Eq v) => Multwomap k v -> Multwomap k v -> Bool
== :: Multwomap k v -> Multwomap k v -> Bool
$c/= :: forall k v. (Eq k, Eq v) => Multwomap k v -> Multwomap k v -> Bool
/= :: Multwomap k v -> Multwomap k v -> Bool
Eq, Eq (Multwomap k v)
Eq (Multwomap k v)
-> (Multwomap k v -> Multwomap k v -> Ordering)
-> (Multwomap k v -> Multwomap k v -> Bool)
-> (Multwomap k v -> Multwomap k v -> Bool)
-> (Multwomap k v -> Multwomap k v -> Bool)
-> (Multwomap k v -> Multwomap k v -> Bool)
-> (Multwomap k v -> Multwomap k v -> Multwomap k v)
-> (Multwomap k v -> Multwomap k v -> Multwomap k v)
-> Ord (Multwomap k v)
Multwomap k v -> Multwomap k v -> Bool
Multwomap k v -> Multwomap k v -> Ordering
Multwomap k v -> Multwomap k v -> Multwomap k v
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 (Multwomap k v)
forall k v.
(Ord k, Ord v) =>
Multwomap k v -> Multwomap k v -> Bool
forall k v.
(Ord k, Ord v) =>
Multwomap k v -> Multwomap k v -> Ordering
forall k v.
(Ord k, Ord v) =>
Multwomap k v -> Multwomap k v -> Multwomap k v
$ccompare :: forall k v.
(Ord k, Ord v) =>
Multwomap k v -> Multwomap k v -> Ordering
compare :: Multwomap k v -> Multwomap k v -> Ordering
$c< :: forall k v.
(Ord k, Ord v) =>
Multwomap k v -> Multwomap k v -> Bool
< :: Multwomap k v -> Multwomap k v -> Bool
$c<= :: forall k v.
(Ord k, Ord v) =>
Multwomap k v -> Multwomap k v -> Bool
<= :: Multwomap k v -> Multwomap k v -> Bool
$c> :: forall k v.
(Ord k, Ord v) =>
Multwomap k v -> Multwomap k v -> Bool
> :: Multwomap k v -> Multwomap k v -> Bool
$c>= :: forall k v.
(Ord k, Ord v) =>
Multwomap k v -> Multwomap k v -> Bool
>= :: Multwomap k v -> Multwomap k v -> Bool
$cmax :: forall k v.
(Ord k, Ord v) =>
Multwomap k v -> Multwomap k v -> Multwomap k v
max :: Multwomap k v -> Multwomap k v -> Multwomap k v
$cmin :: forall k v.
(Ord k, Ord v) =>
Multwomap k v -> Multwomap k v -> Multwomap k v
min :: Multwomap k v -> Multwomap k v -> Multwomap k v
Ord, Int -> Multwomap k v -> ShowS
[Multwomap k v] -> ShowS
Multwomap k v -> String
(Int -> Multwomap k v -> ShowS)
-> (Multwomap k v -> String)
-> ([Multwomap k v] -> ShowS)
-> Show (Multwomap k v)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall k v. (Show k, Show v) => Int -> Multwomap k v -> ShowS
forall k v. (Show k, Show v) => [Multwomap k v] -> ShowS
forall k v. (Show k, Show v) => Multwomap k v -> String
$cshowsPrec :: forall k v. (Show k, Show v) => Int -> Multwomap k v -> ShowS
showsPrec :: Int -> Multwomap k v -> ShowS
$cshow :: forall k v. (Show k, Show v) => Multwomap k v -> String
show :: Multwomap k v -> String
$cshowList :: forall k v. (Show k, Show v) => [Multwomap k v] -> ShowS
showList :: [Multwomap k v] -> ShowS
Show)

empty :: Multwomap k v
empty :: forall k v. Multwomap k v
empty = Map k (OneTwo v) -> Multwomap k v
forall k v. Map k (OneTwo v) -> Multwomap k v
Multwomap Map k (OneTwo v)
forall k a. Map k a
M.empty

null :: Multwomap k v -> Bool
null :: forall k v. Multwomap k v -> Bool
null (Multwomap Map k (OneTwo v)
mmap) = Map k (OneTwo v) -> Bool
forall k a. Map k a -> Bool
M.null Map k (OneTwo v)
mmap

size :: Multwomap k v -> Int
size :: forall k v. Multwomap k v -> Int
size (Multwomap Map k (OneTwo v)
mmap) = Map k Int -> Int
forall a. Num a => Map k a -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((OneTwo v -> Int) -> Map k (OneTwo v) -> Map k Int
forall a b k. (a -> b) -> Map k a -> Map k b
M.map (\case One{} -> Int
1; Two{} -> Int
2) Map k (OneTwo v)
mmap)

insert :: Ord k => k -> v -> Multwomap k v -> Multwomap k v
insert :: forall k v. Ord k => k -> v -> Multwomap k v -> Multwomap k v
insert k
k v
v (Multwomap Map k (OneTwo v)
mmap) = Map k (OneTwo v) -> Multwomap k v
forall k v. Map k (OneTwo v) -> Multwomap k v
Multwomap ((OneTwo v -> OneTwo v -> OneTwo v)
-> k -> OneTwo v -> Map k (OneTwo v) -> Map k (OneTwo v)
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith OneTwo v -> OneTwo v -> OneTwo v
forall a. OneTwo a -> OneTwo a -> OneTwo a
mergeOneTwo k
k (v -> OneTwo v
forall a. a -> OneTwo a
One v
v) Map k (OneTwo v)
mmap)

union :: Ord k => Multwomap k v -> Multwomap k v -> Multwomap k v
union :: forall k v.
Ord k =>
Multwomap k v -> Multwomap k v -> Multwomap k v
union (Multwomap Map k (OneTwo v)
mmap1) (Multwomap Map k (OneTwo v)
mmap2) = Map k (OneTwo v) -> Multwomap k v
forall k v. Map k (OneTwo v) -> Multwomap k v
Multwomap ((OneTwo v -> OneTwo v -> OneTwo v)
-> Map k (OneTwo v) -> Map k (OneTwo v) -> Map k (OneTwo v)
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWith OneTwo v -> OneTwo v -> OneTwo v
forall a. OneTwo a -> OneTwo a -> OneTwo a
mergeOneTwo Map k (OneTwo v)
mmap1 Map k (OneTwo v)
mmap2)

mergeOneTwo :: OneTwo a -> OneTwo a -> OneTwo a
mergeOneTwo :: forall a. OneTwo a -> OneTwo a -> OneTwo a
mergeOneTwo (One a
a) (One a
b) = a -> a -> OneTwo a
forall a. a -> a -> OneTwo a
Two a
a a
b
mergeOneTwo Two{} Two{} = String -> String -> OneTwo a
forall a. String -> String -> a
bugError String
"Multwomap" String
"Overflow: both args already have two targets"
mergeOneTwo Two{} One{} = String -> String -> OneTwo a
forall a. String -> String -> a
bugError String
"Multwomap" String
"Overflow: first arg already has two targets"
mergeOneTwo One{} Two{} = String -> String -> OneTwo a
forall a. String -> String -> a
bugError String
"Multwomap" String
"Overflow: second arg already has two targets"

-- | Get an arbitrary key contained in the Multwomap, or 'Nothing' if it’s empty.
arbitraryKey :: Multwomap k v -> Maybe k
arbitraryKey :: forall k v. Multwomap k v -> Maybe k
arbitraryKey (Multwomap Map k (OneTwo v)
mmap) = ((k, OneTwo v) -> k) -> Maybe (k, OneTwo v) -> Maybe k
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (k, OneTwo v) -> k
forall a b. (a, b) -> a
fst (Map k (OneTwo v) -> Maybe (k, OneTwo v)
forall k a. Map k a -> Maybe (k, a)
M.lookupMin Map k (OneTwo v)
mmap)

-- | Extract the value to the corresponding key, and return the Multwomap with the
-- key removed.
extract :: Ord k => k -> Multwomap k v -> Maybe (v, Multwomap k v)
extract :: forall k v. Ord k => k -> Multwomap k v -> Maybe (v, Multwomap k v)
extract k
k (Multwomap Map k (OneTwo v)
mmap) = case k -> Map k (OneTwo v) -> Maybe (OneTwo v)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
k Map k (OneTwo v)
mmap of
    Maybe (OneTwo v)
Nothing -> Maybe (v, Multwomap k v)
forall a. Maybe a
Nothing
    Just (One v
v) -> (v, Multwomap k v) -> Maybe (v, Multwomap k v)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (v
v, Map k (OneTwo v) -> Multwomap k v
forall k v. Map k (OneTwo v) -> Multwomap k v
Multwomap (k -> Map k (OneTwo v) -> Map k (OneTwo v)
forall k a. Ord k => k -> Map k a -> Map k a
M.delete k
k Map k (OneTwo v)
mmap))
    Just (Two v
v v
w) -> (v, Multwomap k v) -> Maybe (v, Multwomap k v)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (v
v, Map k (OneTwo v) -> Multwomap k v
forall k v. Map k (OneTwo v) -> Multwomap k v
Multwomap (k -> OneTwo v -> Map k (OneTwo v) -> Map k (OneTwo v)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
k (v -> OneTwo v
forall a. a -> OneTwo a
One v
w) Map k (OneTwo v)
mmap))