-- SPDX-FileCopyrightText: 2022 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

module Morley.Util.Bimap
  ( Bimap(..)
  , empty
    -- * Optics
  , flipped
  ) where

import Prelude hiding (empty)

import Control.Lens (At(..), Index, Iso, IxValue, Ixed(..), iso)
import Data.Aeson (FromJSON(..), ToJSON(..))
import Data.Bimap qualified as Bimap
import Data.Coerce (coerce)
import GHC.Exts (IsList)

newtype Bimap a b = Bimap { forall a b. Bimap a b -> Bimap a b
unBimap :: Bimap.Bimap a b }
  deriving newtype (Int -> Bimap a b -> ShowS
[Bimap a b] -> ShowS
Bimap a b -> String
(Int -> Bimap a b -> ShowS)
-> (Bimap a b -> String)
-> ([Bimap a b] -> ShowS)
-> Show (Bimap a b)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall a b. (Show a, Show b) => Int -> Bimap a b -> ShowS
forall a b. (Show a, Show b) => [Bimap a b] -> ShowS
forall a b. (Show a, Show b) => Bimap a b -> String
showList :: [Bimap a b] -> ShowS
$cshowList :: forall a b. (Show a, Show b) => [Bimap a b] -> ShowS
show :: Bimap a b -> String
$cshow :: forall a b. (Show a, Show b) => Bimap a b -> String
showsPrec :: Int -> Bimap a b -> ShowS
$cshowsPrec :: forall a b. (Show a, Show b) => Int -> Bimap a b -> ShowS
Show, Bimap a b -> Bimap a b -> Bool
(Bimap a b -> Bimap a b -> Bool)
-> (Bimap a b -> Bimap a b -> Bool) -> Eq (Bimap a b)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall a b. (Eq a, Eq b) => Bimap a b -> Bimap a b -> Bool
/= :: Bimap a b -> Bimap a b -> Bool
$c/= :: forall a b. (Eq a, Eq b) => Bimap a b -> Bimap a b -> Bool
== :: Bimap a b -> Bimap a b -> Bool
$c== :: forall a b. (Eq a, Eq b) => Bimap a b -> Bimap a b -> Bool
Eq, Eq (Bimap a b)
Eq (Bimap a b)
-> (Bimap a b -> Bimap a b -> Ordering)
-> (Bimap a b -> Bimap a b -> Bool)
-> (Bimap a b -> Bimap a b -> Bool)
-> (Bimap a b -> Bimap a b -> Bool)
-> (Bimap a b -> Bimap a b -> Bool)
-> (Bimap a b -> Bimap a b -> Bimap a b)
-> (Bimap a b -> Bimap a b -> Bimap a b)
-> Ord (Bimap a b)
Bimap a b -> Bimap a b -> Bool
Bimap a b -> Bimap a b -> Ordering
Bimap a b -> Bimap a b -> Bimap a b
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} {b}. (Ord a, Ord b) => Eq (Bimap a b)
forall a b. (Ord a, Ord b) => Bimap a b -> Bimap a b -> Bool
forall a b. (Ord a, Ord b) => Bimap a b -> Bimap a b -> Ordering
forall a b. (Ord a, Ord b) => Bimap a b -> Bimap a b -> Bimap a b
min :: Bimap a b -> Bimap a b -> Bimap a b
$cmin :: forall a b. (Ord a, Ord b) => Bimap a b -> Bimap a b -> Bimap a b
max :: Bimap a b -> Bimap a b -> Bimap a b
$cmax :: forall a b. (Ord a, Ord b) => Bimap a b -> Bimap a b -> Bimap a b
>= :: Bimap a b -> Bimap a b -> Bool
$c>= :: forall a b. (Ord a, Ord b) => Bimap a b -> Bimap a b -> Bool
> :: Bimap a b -> Bimap a b -> Bool
$c> :: forall a b. (Ord a, Ord b) => Bimap a b -> Bimap a b -> Bool
<= :: Bimap a b -> Bimap a b -> Bool
$c<= :: forall a b. (Ord a, Ord b) => Bimap a b -> Bimap a b -> Bool
< :: Bimap a b -> Bimap a b -> Bool
$c< :: forall a b. (Ord a, Ord b) => Bimap a b -> Bimap a b -> Bool
compare :: Bimap a b -> Bimap a b -> Ordering
$ccompare :: forall a b. (Ord a, Ord b) => Bimap a b -> Bimap a b -> Ordering
Ord, Int -> [Item (Bimap a b)] -> Bimap a b
[Item (Bimap a b)] -> Bimap a b
Bimap a b -> [Item (Bimap a b)]
([Item (Bimap a b)] -> Bimap a b)
-> (Int -> [Item (Bimap a b)] -> Bimap a b)
-> (Bimap a b -> [Item (Bimap a b)])
-> IsList (Bimap a b)
forall l.
([Item l] -> l)
-> (Int -> [Item l] -> l) -> (l -> [Item l]) -> IsList l
forall a b.
(Ord a, Ord b) =>
Int -> [Item (Bimap a b)] -> Bimap a b
forall a b. (Ord a, Ord b) => [Item (Bimap a b)] -> Bimap a b
forall a b. (Ord a, Ord b) => Bimap a b -> [Item (Bimap a b)]
toList :: Bimap a b -> [Item (Bimap a b)]
$ctoList :: forall a b. (Ord a, Ord b) => Bimap a b -> [Item (Bimap a b)]
fromListN :: Int -> [Item (Bimap a b)] -> Bimap a b
$cfromListN :: forall a b.
(Ord a, Ord b) =>
Int -> [Item (Bimap a b)] -> Bimap a b
fromList :: [Item (Bimap a b)] -> Bimap a b
$cfromList :: forall a b. (Ord a, Ord b) => [Item (Bimap a b)] -> Bimap a b
IsList)

empty :: Bimap a b
empty :: forall a b. Bimap a b
empty = Bimap a b -> Bimap a b
coerce Bimap a b
forall a b. Bimap a b
Bimap.empty

type instance Index (Bimap k _) = k
type instance IxValue (Bimap _ v) = v

-- | Left-biased 'Ixed' instance.
-- It assumes the left value @a@ is the key (just like the @Ix (Map k v)@ instance).
--
-- To flip this assumption, use the 'flipped' optic.
instance (Ord k, Ord v) => Ixed (Bimap k v) where
  ix :: k -> Traversal' (Bimap k v) v
  ix :: k -> Traversal' (Bimap k v) v
ix k
k v -> f v
handler (Bimap Bimap k v
bmap) =
    case k -> Bimap k v -> Maybe v
forall a b (m :: * -> *).
(Ord a, Ord b, MonadThrow m) =>
a -> Bimap a b -> m b
Bimap.lookup k
k Bimap k v
bmap of
     Just v
v -> v -> f v
handler v
v f v -> (v -> Bimap k v) -> f (Bimap k v)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \v
v' -> Bimap k v -> Bimap k v
forall a b. Bimap a b -> Bimap a b
Bimap (Bimap k v -> Bimap k v) -> Bimap k v -> Bimap k v
forall a b. (a -> b) -> a -> b
$ k -> v -> Bimap k v -> Bimap k v
forall a b. (Ord a, Ord b) => a -> b -> Bimap a b -> Bimap a b
Bimap.insert k
k v
v' Bimap k v
bmap
     Maybe v
Nothing -> Bimap k v -> f (Bimap k v)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bimap k v -> f (Bimap k v)) -> Bimap k v -> f (Bimap k v)
forall a b. (a -> b) -> a -> b
$ Bimap k v -> Bimap k v
forall a b. Bimap a b -> Bimap a b
Bimap Bimap k v
bmap

-- | Left-biased 'At' instance.
-- It assumes the left value @a@ is the key (just like the @At (Map k v)@ instance).
--
-- To flip this assumption, use the 'flipped' optic.
instance (Ord k, Ord v) => At (Bimap k v) where
  at :: k -> Lens' (Bimap k v) (Maybe v)
  at :: k -> Lens' (Bimap k v) (Maybe v)
at k
k Maybe v -> f (Maybe v)
handler (Bimap Bimap k v
bmap) =
    Maybe v -> f (Maybe v)
handler Maybe v
currentValueMaybe f (Maybe v) -> (Maybe v -> Bimap k v) -> f (Bimap k v)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \Maybe v
newValueMaybe ->
      Bimap k v -> Bimap k v
forall a b. Bimap a b -> Bimap a b
Bimap (Bimap k v -> Bimap k v) -> Bimap k v -> Bimap k v
forall a b. (a -> b) -> a -> b
$
        case (Maybe v
currentValueMaybe, Maybe v
newValueMaybe) of
          (Maybe v
Nothing, Maybe v
Nothing) -> Bimap k v
bmap
          (Just v
_, Maybe v
Nothing) -> k -> Bimap k v -> Bimap k v
forall a b. (Ord a, Ord b) => a -> Bimap a b -> Bimap a b
Bimap.delete k
k Bimap k v
bmap
          (Maybe v
_, Just v
newValue) -> k -> v -> Bimap k v -> Bimap k v
forall a b. (Ord a, Ord b) => a -> b -> Bimap a b -> Bimap a b
Bimap.insert k
k v
newValue Bimap k v
bmap
    where
      currentValueMaybe :: Maybe v
currentValueMaybe = k -> Bimap k v -> Maybe v
forall a b (m :: * -> *).
(Ord a, Ord b, MonadThrow m) =>
a -> Bimap a b -> m b
Bimap.lookup k
k Bimap k v
bmap

-- | Isomorphism between @Bimap a b@ and @Bimap b a@.
flipped :: Iso (Bimap a1 b1) (Bimap a2 b2) (Bimap b1 a1) (Bimap b2 a2)
flipped :: forall a1 b1 a2 b2.
Iso (Bimap a1 b1) (Bimap a2 b2) (Bimap b1 a1) (Bimap b2 a2)
flipped = (Bimap a1 b1 -> Bimap b1 a1)
-> (Bimap b2 a2 -> Bimap a2 b2)
-> Iso (Bimap a1 b1) (Bimap a2 b2) (Bimap b1 a1) (Bimap b2 a2)
forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso ((Bimap a1 b1 -> Bimap b1 a1) -> Bimap a1 b1 -> Bimap b1 a1
coerce Bimap a1 b1 -> Bimap b1 a1
forall a b. Bimap a b -> Bimap b a
Bimap.twist) ((Bimap b2 a2 -> Bimap a2 b2) -> Bimap b2 a2 -> Bimap a2 b2
coerce Bimap b2 a2 -> Bimap a2 b2
forall a b. Bimap a b -> Bimap b a
Bimap.twist)

instance (Ord a, Ord b, FromJSON a, FromJSON b) => FromJSON (Bimap a b) where
  parseJSON :: Value -> Parser (Bimap a b)
parseJSON = ([(a, b)] -> Bimap a b) -> Parser [(a, b)] -> Parser (Bimap a b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Bimap a b -> Bimap a b
forall a b. Bimap a b -> Bimap a b
Bimap (Bimap a b -> Bimap a b)
-> ([(a, b)] -> Bimap a b) -> [(a, b)] -> Bimap a b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(a, b)] -> Bimap a b
forall a b. (Ord a, Ord b) => [(a, b)] -> Bimap a b
Bimap.fromList) (Parser [(a, b)] -> Parser (Bimap a b))
-> (Value -> Parser [(a, b)]) -> Value -> Parser (Bimap a b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> Parser [(a, b)]
forall a. FromJSON a => Value -> Parser a
parseJSON

instance (ToJSON a, ToJSON b) => ToJSON (Bimap a b) where
  toJSON :: Bimap a b -> Value
toJSON = [(a, b)] -> Value
forall a. ToJSON a => a -> Value
toJSON ([(a, b)] -> Value)
-> (Bimap a b -> [(a, b)]) -> Bimap a b -> Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bimap a b -> [(a, b)]
forall a b. Bimap a b -> [(a, b)]
Bimap.toList (Bimap a b -> [(a, b)])
-> (Bimap a b -> Bimap a b) -> Bimap a b -> [(a, b)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bimap a b -> Bimap a b
forall a b. Bimap a b -> Bimap a b
unBimap