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

For functions with an @L@ or @R@ suffix, the letter indicates whether
the /parameter/ type is specialized to the left or right type of
the bimap.
-}
module Data.Bimap (
    -- * Bimap type
    Bimap(),
    -- * Query
    null,
    size,
    member,
    memberL,
    memberR,
    notMember,
    notMemberL,
    notMemberR,
    pairMember,
    pairNotMember,
    lookup,
    lookupL,
    lookupR,
    (!),
    (!<),
    (!>),
    -- * Construction
    empty,
    singleton,
    -- * Update
    insert,
    delete,
    deleteL,
    deleteR,
    -- * Conversion\/traversal
    fromList,
    toList,
    assocs,
    fold,
    -- * Miscellaneous
    valid,
    twist,
) where

import Control.Arrow ((>>>))
import Control.Monad (liftM)
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, Ord a, Show b, Ord b) => Show (Bimap a b) where
    show x = "fromList " ++ (show . toList $ x)

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

{-| A bimap with a single element. -}
singleton :: (Ord a, Ord b)
          => (a, b) -> Bimap a b
singleton xy = unsafeInsert xy empty

{-| 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)
       => Either a b -> Bimap a b -> Bool
member (Left  x) (MkBimap left  _) = M.member x left
member (Right y) (MkBimap _ right) = M.member y right

{-| A version of 'member' specialized to the left key. -}
memberL :: (Ord a, Ord b) => a -> Bimap a b -> Bool
memberL = member . Left
{-| A version of 'member' specialized to the right key. -}
memberR :: (Ord a, Ord b) => b -> Bimap a b -> Bool
memberR = member . Right

{-| Is the specified value not a member of the bimap? -}
notMember :: (Ord a, Ord b)
          => Either a b -> Bimap a b -> Bool
notMember = not .: member

{-| A version of 'notMember' specialized to the left key. -}
notMemberL :: (Ord a, Ord b) => a -> Bimap a b -> Bool
notMemberL = notMember . Left
{-| A version of 'notMember' specialized to the right key. -}
notMemberR :: (Ord a, Ord b) => b -> Bimap a b -> Bool
notMemberR = notMember . Right

{-| Are the two values associated /with each other/ in the bimap? -}
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 (Left x)
            >>> delete (Right 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)

{-| 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)
       => Either a b -> Bimap a b -> Bimap a b
delete 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

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

{-| Lookup the twin of a value in the bimap, returning both
associated values as a pair.

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)
       => Either a b -> Bimap a b -> m (a, b)
lookup (Left x)  (MkBimap left _) =
    maybe (fail "Data.Bimap.lookup: Left key not found")
          (\y -> return (x, y))
          (M.lookup x left)
lookup (Right y) (MkBimap _ right) =
    maybe (fail "Data.Bimap.lookup: Right key not found")
          (\x -> return (x, y))
          (M.lookup y right)

{-| A version of 'lookup' that is specialized to the left key,
and returns only the right key. -}
lookupL :: (Ord a, Ord b, Monad m)
        => a -> Bimap a b -> m b
lookupL = (liftM snd) .: lookup . 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 = (liftM fst) .: lookup . Right

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

{-| A version of '(!)' that is specialized to the left key,
and returns only the right key. -}
(!<) :: (Ord a, Ord b) => Bimap a b -> a -> b
(!<) bi x = snd $ bi ! Left  x
{-| 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 = fst $ bi ! Right y

{-| 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 insert) empty xs

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

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

{-| 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 :: (Ord a, Ord b)
      => Bimap a b -> Bimap b a
twist (MkBimap left right) = MkBimap right left

{-| 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 f z . assocs