{-|
An implementation of bidirectional maps between values of two
key types. A 'Bimap' is essentially a bijection between subsets of
its two argument types.

Most functions implicitly consider the left-hand type to be the
key, and the right-hand type to be the value.
Functions with an @R@ suffix reverse this convention, treating the
left-hand type as the key.
-}
module Data.Bimap (
    -- * Bimap type
    Bimap(),
    -- * Query
    null,
    size,
    member,
    memberR,
    notMember,
    notMemberR,
    pairMember,
    pairNotMember,
    lookup,
    lookupR,
    (!),
    (!>),
    -- * Construction
    empty,
    singleton,
    -- * Update
    insert,
    delete,
    deleteR,
    -- * Conversion\/traversal
    fromList,
    toList,
    toAscList,
    toAscListR,
    keys,
    keysR,
    elems,
    assocs,
    fold,
    -- * Miscellaneous
    valid,
    twist,
    twisted,
) where

import Control.Arrow ((>>>))
import Control.Monad.Error () -- Monad instance for Either e
import Data.List (foldl', sort)
import qualified Data.Map as M
import Prelude hiding (lookup, null)


infixr 9 .:
(.:) :: (c -> d) -> (a -> b -> c) -> a -> b -> d
(.:) = (.).(.)

{-|
A bidirectional map between values of types @a@ and @b@.
-}
data Bimap a b = MkBimap !(M.Map a b) !(M.Map b a)

instance (Show a, Show b) => Show (Bimap a b) where
    show x = "fromList " ++ (show . toList $ x)

instance (Eq a, Eq b) => Eq (Bimap a b) where
    (==) bx by = toAscList bx == toAscList by

{-| The empty bimap. -}
empty :: Bimap a b
empty = MkBimap M.empty M.empty

{-| A bimap with a single element. -}
singleton :: a -> b -> Bimap a b
singleton x y = MkBimap (M.singleton x y) (M.singleton y x)

{-| Is the bimap empty? -}
null :: Bimap a b -> Bool
null (MkBimap left _) = M.null left

{-| The number of elements in the bimap. -}
size :: Bimap a b -> Int
size (MkBimap left _) = M.size left

{-| Is the specified value a member of the bimap? -}
member :: (Ord a, Ord b) => a -> Bimap a b -> Bool
member x (MkBimap left _) = M.member x left
{-| A version of 'member' specialized to the right key. -}
memberR :: (Ord a, Ord b) => b -> Bimap a b -> Bool
memberR y (MkBimap _ right) = M.member y right

{-| Is the specified value not a member of the bimap? -}
notMember :: (Ord a, Ord b) => a -> Bimap a b -> Bool
notMember = not .: member
{-| A version of 'notMember' specialized to the right key. -}
notMemberR :: (Ord a, Ord b) => b -> Bimap a b -> Bool
notMemberR = not .: memberR

{-| Are the two values associated /with each other/ in the bimap?

This function is uncurried in its first two arguments, so that it
can be used infix.
-}
pairMember :: (Ord a, Ord b)
           => (a, b) -> Bimap a b -> Bool
pairMember (x, y) (MkBimap left _) =
    maybe False (== y) (M.lookup x left)

{-| Are the two values not in the bimap, or not associated with
each other? (Complement of 'pairMember'.) -}
pairNotMember :: (Ord a, Ord b)
              => (a, b) -> Bimap a b -> Bool
pairNotMember = not .: pairMember

{-| Insert a pair of values into the bimap, associating them.
If either of the values is already in the bimap, any overlapping
bindings are deleted.
-}
insert :: (Ord a, Ord b)
        => a -> b -> Bimap a b -> Bimap a b
insert x y = delete  x
         >>> deleteR y
         >>> unsafeInsert x y

{-| Insert a pair of values into the bimap, without checking for
overlapping bindings. If either value is already in the bimap, and
is not bound to the other value, the bimap will become inconsistent.
-}
unsafeInsert :: (Ord a, Ord b)
             => a -> b -> Bimap a b -> Bimap a b
unsafeInsert x y (MkBimap left right) =
    MkBimap (M.insert x y left) (M.insert y x right)

{-| Common implementation for 'delete' and 'deleteR'. -}
deleteE :: (Ord a, Ord b)
       => Either a b -> Bimap a b -> Bimap a b
deleteE e (MkBimap left right) =
    MkBimap
        (perhaps M.delete x $ left)
        (perhaps M.delete y $ right)
    where
    perhaps = maybe id
    x = either Just (flip M.lookup right) e
    y = either (flip M.lookup left) Just  e

{-| Delete a value and its twin from a bimap.
When the value is not a member of the bimap, the original bimap is
returned.
-}
delete :: (Ord a, Ord b) => a -> Bimap a b -> Bimap a b
delete = deleteE . Left

{-| A version of 'delete' specialized to the right key. -}
deleteR :: (Ord a, Ord b) => b -> Bimap a b -> Bimap a b
deleteR = deleteE . Right

{-| Lookup a left key in the bimap, returning the associated right
key.

This function will @return@ the result in the monad, or @fail@ if
the value isn't in the bimap.
-}
lookup :: (Ord a, Ord b, Monad m)
        => a -> Bimap a b -> m b
lookup x (MkBimap left _) =
    maybe (fail "Data.Bimap.lookup: Left key not found")
          (return)
          (M.lookup x left)

{-| A version of 'lookup' that is specialized to the right key,
and returns only the left key. -}
lookupR :: (Ord a, Ord b, Monad m)
        => b -> Bimap a b -> m a
lookupR y (MkBimap _ right) =
    maybe (fail "Data.Bimap.lookupR: Right key not found")
          (return)
          (M.lookup y right)

{-| Find the right key corresponding to a given left key.
Calls @'error'@ when the key is not in the bimap.
-}
(!) :: (Ord a, Ord b) => Bimap a b -> a -> b
(!) bi x = either error id (lookup x bi)

{-| A version of @(!)@ that is specialized to the right key,
and returns only the left key. -}
(!>) :: (Ord a, Ord b) => Bimap a b -> b -> a
(!>) bi y = either error id (lookupR y bi)

{-| Build a map from a list of pairs. If there are any overlapping
pairs in the list, the later ones will override the earlier ones.
-}
fromList :: (Ord a, Ord b)
         => [(a, b)] -> Bimap a b
fromList xs = foldl' (flip . uncurry $ insert) empty xs

{-| Convert to a list of associated pairs. -}
toList :: Bimap a b -> [(a, b)]
toList = toAscList

{-| Convert to a list of associated pairs, with the left-hand
values in ascending order.
Since pair ordering is lexical, the pairs will also be in
ascending order.
-}
toAscList :: Bimap a b -> [(a, b)]
toAscList (MkBimap left _) = M.toList left

{-| Convert to a list of associated pairs, with the right-hand
values first in the pair and in ascending order.
Since pair ordering is lexical, the pairs will also be in
ascending order.
-}
toAscListR :: Bimap a b -> [(b, a)]
toAscListR = toAscList . twist

{-| Return all associated pairs in the bimap, with the left-hand
values in ascending order. -}
assocs :: Bimap a b -> [(a, b)]
assocs = toList

{-| Return all left-hand keys in the bimap in ascending order. -}
keys :: Bimap a b -> [a]
keys (MkBimap left _) = M.keys left

{-| Return all right-hand keys in the bimap in ascending order. -}
keysR :: Bimap a b -> [b]
keysR (MkBimap _ right) = M.keys right

{-| An alias for 'keysR'. -}
elems :: Bimap a b -> [b]
elems = keysR

{-| Test if the internal bimap structure is valid. -}
valid :: (Ord a, Ord b)
      => Bimap a b -> Bool
valid (MkBimap left right) = and
    [ M.valid left, M.valid right
    , (==)
        (sort .                M.toList $ left )
        (sort . map flipPair . M.toList $ right)
    ]
    where
    flipPair (x, y) = (y, x)

{-| Reverse the positions of the two element types in the bimap. -}
twist ::  Bimap a b -> Bimap b a
twist (MkBimap left right) = MkBimap right left

{-| Reverse the positions of the two element types in a bimap
transformation. -}
twisted :: (Bimap a b -> Bimap a b) -> (Bimap b a -> Bimap b a)
twisted f = twist . f . twist

{-| Fold the association pairs in the map, such that
@'fold' f z == 'foldr' f z . 'assocs'@.
-}
fold :: (a -> b -> c -> c) -> c -> Bimap a b -> c
fold f z = foldr (uncurry f) z . assocs