module Unison.Util.BiMap
  ( BiMap (..),
    empty,
    singleton,
    fromList,
    fromMap,
    toList,
    toMapL,
    toMapR,
    lookupL,
    lookupR,
    union,
    difference,
    insert,
    deleteL,
    deleteR,
    keysSetL,
    keysSetR,
  )
where

import Data.Map qualified as Map
import Data.Set (Set)
import Data.Tuple (swap)

-- | A bidirectional map between keys of type k and values of type v.
data BiMap k v = BiMap
  { forall k v. BiMap k v -> Map k v
forward :: Map.Map k v,
    forall k v. BiMap k v -> Map v k
backward :: Map.Map v k
  }
  deriving (BiMap k v -> BiMap k v -> Bool
(BiMap k v -> BiMap k v -> Bool)
-> (BiMap k v -> BiMap k v -> Bool) -> Eq (BiMap k v)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall k v. (Eq k, Eq v) => BiMap k v -> BiMap k v -> Bool
$c== :: forall k v. (Eq k, Eq v) => BiMap k v -> BiMap k v -> Bool
== :: BiMap k v -> BiMap k v -> Bool
$c/= :: forall k v. (Eq k, Eq v) => BiMap k v -> BiMap k v -> Bool
/= :: BiMap k v -> BiMap k v -> Bool
Eq, Eq (BiMap k v)
Eq (BiMap k v) =>
(BiMap k v -> BiMap k v -> Ordering)
-> (BiMap k v -> BiMap k v -> Bool)
-> (BiMap k v -> BiMap k v -> Bool)
-> (BiMap k v -> BiMap k v -> Bool)
-> (BiMap k v -> BiMap k v -> Bool)
-> (BiMap k v -> BiMap k v -> BiMap k v)
-> (BiMap k v -> BiMap k v -> BiMap k v)
-> Ord (BiMap k v)
BiMap k v -> BiMap k v -> Bool
BiMap k v -> BiMap k v -> Ordering
BiMap k v -> BiMap k v -> BiMap 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 (BiMap k v)
forall k v. (Ord k, Ord v) => BiMap k v -> BiMap k v -> Bool
forall k v. (Ord k, Ord v) => BiMap k v -> BiMap k v -> Ordering
forall k v. (Ord k, Ord v) => BiMap k v -> BiMap k v -> BiMap k v
$ccompare :: forall k v. (Ord k, Ord v) => BiMap k v -> BiMap k v -> Ordering
compare :: BiMap k v -> BiMap k v -> Ordering
$c< :: forall k v. (Ord k, Ord v) => BiMap k v -> BiMap k v -> Bool
< :: BiMap k v -> BiMap k v -> Bool
$c<= :: forall k v. (Ord k, Ord v) => BiMap k v -> BiMap k v -> Bool
<= :: BiMap k v -> BiMap k v -> Bool
$c> :: forall k v. (Ord k, Ord v) => BiMap k v -> BiMap k v -> Bool
> :: BiMap k v -> BiMap k v -> Bool
$c>= :: forall k v. (Ord k, Ord v) => BiMap k v -> BiMap k v -> Bool
>= :: BiMap k v -> BiMap k v -> Bool
$cmax :: forall k v. (Ord k, Ord v) => BiMap k v -> BiMap k v -> BiMap k v
max :: BiMap k v -> BiMap k v -> BiMap k v
$cmin :: forall k v. (Ord k, Ord v) => BiMap k v -> BiMap k v -> BiMap k v
min :: BiMap k v -> BiMap k v -> BiMap k v
Ord, Int -> BiMap k v -> ShowS
[BiMap k v] -> ShowS
BiMap k v -> String
(Int -> BiMap k v -> ShowS)
-> (BiMap k v -> String)
-> ([BiMap k v] -> ShowS)
-> Show (BiMap k v)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall k v. (Show k, Show v) => Int -> BiMap k v -> ShowS
forall k v. (Show k, Show v) => [BiMap k v] -> ShowS
forall k v. (Show k, Show v) => BiMap k v -> String
$cshowsPrec :: forall k v. (Show k, Show v) => Int -> BiMap k v -> ShowS
showsPrec :: Int -> BiMap k v -> ShowS
$cshow :: forall k v. (Show k, Show v) => BiMap k v -> String
show :: BiMap k v -> String
$cshowList :: forall k v. (Show k, Show v) => [BiMap k v] -> ShowS
showList :: [BiMap k v] -> ShowS
Show)

-- | Combine two BiMaps. In case of key or value collisions, the entries from
-- the second BiMap take precedence.
instance (Ord k, Ord v) => Semigroup (BiMap k v) where
  BiMap Map k v
f1 Map v k
b1 <> :: BiMap k v -> BiMap k v -> BiMap k v
<> BiMap Map k v
f2 Map v k
b2 =
    Map k v -> Map v k -> BiMap k v
forall k v. Map k v -> Map v k -> BiMap k v
BiMap (Map k v -> Map k v -> Map k v
forall k a. Ord k => Map k a -> Map k a -> Map k a
Map.union Map k v
f2 Map k v
f1) (Map v k -> Map v k -> Map v k
forall k a. Ord k => Map k a -> Map k a -> Map k a
Map.union Map v k
b2 Map v k
b1)

instance (Ord k, Ord v) => Monoid (BiMap k v) where
  mempty :: BiMap k v
mempty = Map k v -> Map v k -> BiMap k v
forall k v. Map k v -> Map v k -> BiMap k v
BiMap Map k v
forall k a. Map k a
Map.empty Map v k
forall k a. Map k a
Map.empty

empty :: (Ord k, Ord v) => BiMap k v
empty :: forall k v. (Ord k, Ord v) => BiMap k v
empty = BiMap k v
forall a. Monoid a => a
mempty

singleton :: k -> v -> BiMap k v
singleton :: forall k v. k -> v -> BiMap k v
singleton k
k v
v = Map k v -> Map v k -> BiMap k v
forall k v. Map k v -> Map v k -> BiMap k v
BiMap (k -> v -> Map k v
forall k a. k -> a -> Map k a
Map.singleton k
k v
v) (v -> k -> Map v k
forall k a. k -> a -> Map k a
Map.singleton v
v k
k)

fromList :: (Ord k, Ord v) => [(k, v)] -> BiMap k v
fromList :: forall k v. (Ord k, Ord v) => [(k, v)] -> BiMap k v
fromList [(k, v)]
kvs =
  let forward :: Map k v
forward = [(k, v)] -> Map k v
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(k, v)]
kvs
      backward :: Map v k
backward = [(v, k)] -> Map v k
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ((k, v) -> (v, k)
forall a b. (a, b) -> (b, a)
swap ((k, v) -> (v, k)) -> [(k, v)] -> [(v, k)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map k v -> [(k, v)]
forall k a. Map k a -> [(k, a)]
Map.toList Map k v
forward)
   in BiMap {Map k v
forward :: Map k v
forward :: Map k v
forward, Map v k
backward :: Map v k
backward :: Map v k
backward}

fromMap :: (Ord k, Ord v) => Map.Map k v -> BiMap k v
fromMap :: forall k v. (Ord k, Ord v) => Map k v -> BiMap k v
fromMap Map k v
f =
  let b :: Map v k
b = [(v, k)] -> Map v k
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ((k, v) -> (v, k)
forall a b. (a, b) -> (b, a)
swap ((k, v) -> (v, k)) -> [(k, v)] -> [(v, k)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map k v -> [(k, v)]
forall k a. Map k a -> [(k, a)]
Map.toList Map k v
f)
   in Map k v -> Map v k -> BiMap k v
forall k v. Map k v -> Map v k -> BiMap k v
BiMap Map k v
f Map v k
b

toList :: BiMap k v -> [(k, v)]
toList :: forall k v. BiMap k v -> [(k, v)]
toList (BiMap Map k v
f Map v k
_) = Map k v -> [(k, v)]
forall k a. Map k a -> [(k, a)]
Map.toList Map k v
f

toMapL :: BiMap k v -> Map.Map k v
toMapL :: forall k v. BiMap k v -> Map k v
toMapL (BiMap Map k v
f Map v k
_) = Map k v
f

toMapR :: BiMap k v -> Map.Map v k
toMapR :: forall k v. BiMap k v -> Map v k
toMapR (BiMap Map k v
_ Map v k
b) = Map v k
b

lookupL :: (Ord k) => k -> BiMap k v -> Maybe v
lookupL :: forall k v. Ord k => k -> BiMap k v -> Maybe v
lookupL k
k (BiMap Map k v
f Map v k
_) = k -> Map k v -> Maybe v
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k Map k v
f

lookupR :: (Ord v) => v -> BiMap k v -> Maybe k
lookupR :: forall v k. Ord v => v -> BiMap k v -> Maybe k
lookupR v
v (BiMap Map k v
_ Map v k
b) = v -> Map v k -> Maybe k
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup v
v Map v k
b

union :: (Ord k, Ord v) => BiMap k v -> BiMap k v -> BiMap k v
union :: forall k v. (Ord k, Ord v) => BiMap k v -> BiMap k v -> BiMap k v
union = BiMap k v -> BiMap k v -> BiMap k v
forall a. Semigroup a => a -> a -> a
(<>)

difference :: (Ord k, Ord v) => BiMap k v -> BiMap k v -> BiMap k v
difference :: forall k v. (Ord k, Ord v) => BiMap k v -> BiMap k v -> BiMap k v
difference (BiMap Map k v
f1 Map v k
b1) (BiMap Map k v
f2 Map v k
b2) =
  let f' :: Map k v
f' = Map k v -> Map k v -> Map k v
forall k a b. Ord k => Map k a -> Map k b -> Map k a
Map.difference Map k v
f1 Map k v
f2
      b' :: Map v k
b' = Map v k -> Map v k -> Map v k
forall k a b. Ord k => Map k a -> Map k b -> Map k a
Map.difference Map v k
b1 Map v k
b2
   in Map k v -> Map v k -> BiMap k v
forall k v. Map k v -> Map v k -> BiMap k v
BiMap Map k v
f' Map v k
b'

insert :: (Ord k, Ord v) => k -> v -> BiMap k v -> BiMap k v
insert :: forall k v. (Ord k, Ord v) => k -> v -> BiMap k v -> BiMap k v
insert k
k v
v (BiMap Map k v
f Map v k
b) =
  Map k v -> Map v k -> BiMap k v
forall k v. Map k v -> Map v k -> BiMap k v
BiMap (k -> v -> Map k v -> Map k v
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert k
k v
v Map k v
f) (v -> k -> Map v k -> Map v k
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert v
v k
k Map v k
b)

deleteL :: (Ord k, Ord v) => k -> BiMap k v -> BiMap k v
deleteL :: forall k v. (Ord k, Ord v) => k -> BiMap k v -> BiMap k v
deleteL k
k BiMap k v
bm = case k -> BiMap k v -> Maybe v
forall k v. Ord k => k -> BiMap k v -> Maybe v
lookupL k
k BiMap k v
bm of
  Maybe v
Nothing -> BiMap k v
bm
  Just v
v ->
    Map k v -> Map v k -> BiMap k v
forall k v. Map k v -> Map v k -> BiMap k v
BiMap
      (k -> Map k v -> Map k v
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete k
k (BiMap k v -> Map k v
forall k v. BiMap k v -> Map k v
forward BiMap k v
bm))
      (v -> Map v k -> Map v k
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete v
v (BiMap k v -> Map v k
forall k v. BiMap k v -> Map v k
backward BiMap k v
bm))

deleteR :: (Ord v, Ord k) => v -> BiMap k v -> BiMap k v
deleteR :: forall v k. (Ord v, Ord k) => v -> BiMap k v -> BiMap k v
deleteR v
v BiMap k v
bm = BiMap v k -> BiMap k v
forall k v. BiMap k v -> BiMap v k
flipped (BiMap v k -> BiMap k v) -> BiMap v k -> BiMap k v
forall a b. (a -> b) -> a -> b
$ v -> BiMap v k -> BiMap v k
forall k v. (Ord k, Ord v) => k -> BiMap k v -> BiMap k v
deleteL v
v (BiMap k v -> BiMap v k
forall k v. BiMap k v -> BiMap v k
flipped BiMap k v
bm)

keysSetL :: BiMap k v -> Set k
keysSetL :: forall k v. BiMap k v -> Set k
keysSetL (BiMap Map k v
f Map v k
_) = Map k v -> Set k
forall k a. Map k a -> Set k
Map.keysSet Map k v
f

keysSetR :: BiMap k v -> Set v
keysSetR :: forall k v. BiMap k v -> Set v
keysSetR (BiMap Map k v
_ Map v k
b) = Map v k -> Set v
forall k a. Map k a -> Set k
Map.keysSet Map v k
b

flipped :: BiMap k v -> BiMap v k
flipped :: forall k v. BiMap k v -> BiMap v k
flipped (BiMap Map k v
f Map v k
b) = Map v k -> Map k v -> BiMap v k
forall k v. Map k v -> Map v k -> BiMap k v
BiMap Map v k
b Map k v
f