-- | Finite bijections (implemented as a pair of tree maps).

module Agda.Utils.BiMap where

import Prelude hiding (lookup, unzip)



import Data.Function

import qualified Data.List as List
import Data.Map (Map)
import qualified Data.Map as Map


-- | Finite bijective map from @a@ to @b@.  There, and back again.
data BiMap a b = BiMap
  { BiMap a b -> Map a b
biMapThere :: Map a b
  , BiMap a b -> Map b a
biMapBack  :: Map b a
  }

-- | Lookup. O(log n).
lookup :: Ord a => a -> BiMap a b -> Maybe b
lookup :: a -> BiMap a b -> Maybe b
lookup a
a = a -> Map a b -> Maybe b
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup a
a (Map a b -> Maybe b)
-> (BiMap a b -> Map a b) -> BiMap a b -> Maybe b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BiMap a b -> Map a b
forall a b. BiMap a b -> Map a b
biMapThere

-- | Inverse lookup.  O(log n).
invLookup :: Ord b => b -> BiMap a b -> Maybe a
invLookup :: b -> BiMap a b -> Maybe a
invLookup b
b = b -> Map b a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup b
b (Map b a -> Maybe a)
-> (BiMap a b -> Map b a) -> BiMap a b -> Maybe a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BiMap a b -> Map b a
forall a b. BiMap a b -> Map b a
biMapBack

-- | Empty bimap. O(1).
empty :: BiMap a b
empty :: BiMap a b
empty = Map a b -> Map b a -> BiMap a b
forall a b. Map a b -> Map b a -> BiMap a b
BiMap Map a b
forall k a. Map k a
Map.empty Map b a
forall k a. Map k a
Map.empty

-- | Singleton bimap. O(1).
singleton :: a -> b -> BiMap a b
singleton :: a -> b -> BiMap a b
singleton a
a b
b = Map a b -> Map b a -> BiMap a b
forall a b. Map a b -> Map b a -> BiMap a b
BiMap (a -> b -> Map a b
forall k a. k -> a -> Map k a
Map.singleton a
a b
b) (b -> a -> Map b a
forall k a. k -> a -> Map k a
Map.singleton b
b a
a)

-- | Insert.  Overwrites existing value if present. O(Map.insert).
insert :: (Ord a, Ord b) => a -> b -> BiMap a b -> BiMap a b
insert :: a -> b -> BiMap a b -> BiMap a b
insert a
a b
b (BiMap Map a b
t Map b a
u) = Map a b -> Map b a -> BiMap a b
forall a b. Map a b -> Map b a -> BiMap a b
BiMap (a -> b -> Map a b -> Map a b
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert a
a b
b Map a b
t) (b -> a -> Map b a -> Map b a
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert b
b a
a Map b a
u)

-- | Left-biased Union. O(Map.union).
union :: (Ord a, Ord b) => BiMap a b -> BiMap a b -> BiMap a b
union :: BiMap a b -> BiMap a b -> BiMap a b
union (BiMap Map a b
t1 Map b a
b1) (BiMap Map a b
t2 Map b a
b2) = Map a b -> Map b a -> BiMap a b
forall a b. Map a b -> Map b a -> BiMap a b
BiMap (Map a b -> Map a b -> Map a b
forall k a. Ord k => Map k a -> Map k a -> Map k a
Map.union Map a b
t1 Map a b
t2) (Map b a -> Map b a -> Map b a
forall k a. Ord k => Map k a -> Map k a -> Map k a
Map.union Map b a
b1 Map b a
b2)

-- | Construct from a list of pairs.
--
--   Does not check for actual bijectivity of constructed finite map. O(n log n)
fromList :: (Ord a, Ord b) => [(a,b)] -> BiMap a b
fromList :: [(a, b)] -> BiMap a b
fromList = (BiMap a b -> (a, b) -> BiMap a b)
-> BiMap a b -> [(a, b)] -> BiMap a b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl' (((a, b) -> BiMap a b -> BiMap a b)
-> BiMap a b -> (a, b) -> BiMap a b
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((a -> b -> BiMap a b -> BiMap a b)
-> (a, b) -> BiMap a b -> BiMap a b
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> b -> BiMap a b -> BiMap a b
forall a b. (Ord a, Ord b) => a -> b -> BiMap a b -> BiMap a b
insert)) BiMap a b
forall a b. BiMap a b
empty

-- | Turn into list, sorted ascendingly by first value.  O(Map.toList)
toList :: BiMap a b -> [(a,b)]
toList :: BiMap a b -> [(a, b)]
toList = Map a b -> [(a, b)]
forall k a. Map k a -> [(k, a)]
Map.toAscList (Map a b -> [(a, b)])
-> (BiMap a b -> Map a b) -> BiMap a b -> [(a, b)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BiMap a b -> Map a b
forall a b. BiMap a b -> Map a b
biMapThere

------------------------------------------------------------------------
-- Instances
------------------------------------------------------------------------

instance (Ord a, Ord b) => Eq (BiMap a b) where
  == :: BiMap a b -> BiMap a b -> Bool
(==) = Map a b -> Map a b -> Bool
forall a. Eq a => a -> a -> Bool
(==) (Map a b -> Map a b -> Bool)
-> (BiMap a b -> Map a b) -> BiMap a b -> BiMap a b -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` BiMap a b -> Map a b
forall a b. BiMap a b -> Map a b
biMapThere

instance (Ord a, Ord b) => Ord (BiMap a b) where
  compare :: BiMap a b -> BiMap a b -> Ordering
compare = Map a b -> Map a b -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Map a b -> Map a b -> Ordering)
-> (BiMap a b -> Map a b) -> BiMap a b -> BiMap a b -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` BiMap a b -> Map a b
forall a b. BiMap a b -> Map a b
biMapThere

instance (Show a, Show b) => Show (BiMap a b) where
  show :: BiMap a b -> String
show BiMap a b
bimap = String
"Agda.Utils.BiMap.fromList " String -> ShowS
forall a. [a] -> [a] -> [a]
++ [(a, b)] -> String
forall a. Show a => a -> String
show (BiMap a b -> [(a, b)]
forall a b. BiMap a b -> [(a, b)]
toList BiMap a b
bimap)