{-# LANGUAGE CPP #-}

module Data.Bijection where

#if MIN_VERSION_GLASGOW_HASKELL(9,6,0,0)
#else
import Control.Applicative (liftA2)
#endif
import Data.IntMap.Strict (IntMap)
import qualified Data.IntMap.Strict as IM
import qualified Data.IntMap.Merge.Strict as IM


compatibleInsert :: (Eq a) => Int -> a -> IntMap a -> Maybe (IntMap a)
compatibleInsert :: forall a. Eq a => Int -> a -> IntMap a -> Maybe (IntMap a)
compatibleInsert Int
x a
y = let
  f :: Maybe a -> Maybe (Maybe a)
f (Just a
z) = if a
y forall a. Eq a => a -> a -> Bool
== a
z then forall a. a -> Maybe a
Just (forall a. a -> Maybe a
Just a
y) else forall a. Maybe a
Nothing
  f Maybe a
Nothing  = forall a. a -> Maybe a
Just (forall a. a -> Maybe a
Just a
y)
  in forall (f :: * -> *) a.
Functor f =>
(Maybe a -> f (Maybe a)) -> Int -> IntMap a -> f (IntMap a)
IM.alterF Maybe a -> Maybe (Maybe a)
f Int
x

compatibleUnion :: (Eq a) => IntMap a -> IntMap a -> Maybe (IntMap a)
compatibleUnion :: forall a. Eq a => IntMap a -> IntMap a -> Maybe (IntMap a)
compatibleUnion = let
  f :: p -> a -> a -> Maybe a
f p
_ a
x a
y = if a
x forall a. Eq a => a -> a -> Bool
== a
y then forall a. a -> Maybe a
Just a
x else forall a. Maybe a
Nothing
  in forall (f :: * -> *) a c b.
Applicative f =>
WhenMissing f a c
-> WhenMissing f b c
-> WhenMatched f a b c
-> IntMap a
-> IntMap b
-> f (IntMap c)
IM.mergeA forall (f :: * -> *) x. Applicative f => WhenMissing f x x
IM.preserveMissing forall (f :: * -> *) x. Applicative f => WhenMissing f x x
IM.preserveMissing (forall (f :: * -> *) x y z.
Applicative f =>
(Int -> x -> y -> f z) -> WhenMatched f x y z
IM.zipWithAMatched forall {a} {p}. Eq a => p -> a -> a -> Maybe a
f)

data Bij = Bij {
  Bij -> IntMap Int
rightwards :: IntMap Int,
  Bij -> IntMap Int
leftwards :: IntMap Int
} deriving (Bij -> Bij -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Bij -> Bij -> Bool
$c/= :: Bij -> Bij -> Bool
== :: Bij -> Bij -> Bool
$c== :: Bij -> Bij -> Bool
Eq, Eq Bij
Bij -> Bij -> Bool
Bij -> Bij -> Ordering
Bij -> Bij -> Bij
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
min :: Bij -> Bij -> Bij
$cmin :: Bij -> Bij -> Bij
max :: Bij -> Bij -> Bij
$cmax :: Bij -> Bij -> Bij
>= :: Bij -> Bij -> Bool
$c>= :: Bij -> Bij -> Bool
> :: Bij -> Bij -> Bool
$c> :: Bij -> Bij -> Bool
<= :: Bij -> Bij -> Bool
$c<= :: Bij -> Bij -> Bool
< :: Bij -> Bij -> Bool
$c< :: Bij -> Bij -> Bool
compare :: Bij -> Bij -> Ordering
$ccompare :: Bij -> Bij -> Ordering
Ord)

empty :: Bij
empty :: Bij
empty = IntMap Int -> IntMap Int -> Bij
Bij forall a. IntMap a
IM.empty forall a. IntMap a
IM.empty

singleton :: Int -> Int -> Bij
singleton :: Int -> Int -> Bij
singleton Int
x Int
y = IntMap Int -> IntMap Int -> Bij
Bij (forall a. Int -> a -> IntMap a
IM.singleton Int
x Int
y) (forall a. Int -> a -> IntMap a
IM.singleton Int
y Int
x)

match :: Int -> Int -> Bij -> Maybe Bij
match :: Int -> Int -> Bij -> Maybe Bij
match Int
x Int
y (Bij IntMap Int
r IntMap Int
l) = forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 IntMap Int -> IntMap Int -> Bij
Bij (forall a. Eq a => Int -> a -> IntMap a -> Maybe (IntMap a)
compatibleInsert Int
x Int
y IntMap Int
r) (forall a. Eq a => Int -> a -> IntMap a -> Maybe (IntMap a)
compatibleInsert Int
y Int
x IntMap Int
l)

combine :: Bij -> Bij -> Maybe Bij
combine :: Bij -> Bij -> Maybe Bij
combine (Bij IntMap Int
r1 IntMap Int
l1) (Bij IntMap Int
r2 IntMap Int
l2) = forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 IntMap Int -> IntMap Int -> Bij
Bij (forall a. Eq a => IntMap a -> IntMap a -> Maybe (IntMap a)
compatibleUnion IntMap Int
r1 IntMap Int
r2) (forall a. Eq a => IntMap a -> IntMap a -> Maybe (IntMap a)
compatibleUnion IntMap Int
l1 IntMap Int
l2)

pop :: Bij -> Maybe ((Int, Int), Bij)
pop :: Bij -> Maybe ((Int, Int), Bij)
pop (Bij IntMap Int
r IntMap Int
l) = case forall a. IntMap a -> Maybe ((Int, a), IntMap a)
IM.minViewWithKey IntMap Int
r of
  Maybe ((Int, Int), IntMap Int)
Nothing -> forall a. Maybe a
Nothing
  Just ((Int
i, Int
j), IntMap Int
r') -> forall a. a -> Maybe a
Just ((Int
i, Int
j), IntMap Int -> IntMap Int -> Bij
Bij IntMap Int
r' (forall a. Int -> IntMap a -> IntMap a
IM.delete Int
j IntMap Int
l))

-- | Don't check consistency, just take a union
unsafeUnion :: Bij -> Bij -> Bij
unsafeUnion :: Bij -> Bij -> Bij
unsafeUnion (Bij IntMap Int
r1 IntMap Int
l1) (Bij IntMap Int
r2 IntMap Int
l2) = IntMap Int -> IntMap Int -> Bij
Bij (forall a. IntMap a -> IntMap a -> IntMap a
IM.union IntMap Int
r1 IntMap Int
r2) (forall a. IntMap a -> IntMap a -> IntMap a
IM.union IntMap Int
l1 IntMap Int
l2)

-- | Don't check consistency, just take a diff
unsafeDifference :: Bij -> Bij -> Bij
unsafeDifference :: Bij -> Bij -> Bij
unsafeDifference (Bij IntMap Int
r1 IntMap Int
l1) (Bij IntMap Int
r2 IntMap Int
l2) = IntMap Int -> IntMap Int -> Bij
Bij (forall a b. IntMap a -> IntMap b -> IntMap a
IM.difference IntMap Int
r1 IntMap Int
r2) (forall a b. IntMap a -> IntMap b -> IntMap a
IM.difference IntMap Int
l1 IntMap Int
l2)


-- | A newtype, just to get a partial monoidal structure representing consistent
-- unions.
newtype MaybeBij = MaybeBij {
  MaybeBij -> Maybe Bij
getMaybeBij :: Maybe Bij
} deriving (MaybeBij -> MaybeBij -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: MaybeBij -> MaybeBij -> Bool
$c/= :: MaybeBij -> MaybeBij -> Bool
== :: MaybeBij -> MaybeBij -> Bool
$c== :: MaybeBij -> MaybeBij -> Bool
Eq, Eq MaybeBij
MaybeBij -> MaybeBij -> Bool
MaybeBij -> MaybeBij -> Ordering
MaybeBij -> MaybeBij -> MaybeBij
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
min :: MaybeBij -> MaybeBij -> MaybeBij
$cmin :: MaybeBij -> MaybeBij -> MaybeBij
max :: MaybeBij -> MaybeBij -> MaybeBij
$cmax :: MaybeBij -> MaybeBij -> MaybeBij
>= :: MaybeBij -> MaybeBij -> Bool
$c>= :: MaybeBij -> MaybeBij -> Bool
> :: MaybeBij -> MaybeBij -> Bool
$c> :: MaybeBij -> MaybeBij -> Bool
<= :: MaybeBij -> MaybeBij -> Bool
$c<= :: MaybeBij -> MaybeBij -> Bool
< :: MaybeBij -> MaybeBij -> Bool
$c< :: MaybeBij -> MaybeBij -> Bool
compare :: MaybeBij -> MaybeBij -> Ordering
$ccompare :: MaybeBij -> MaybeBij -> Ordering
Ord)

instance Semigroup MaybeBij where
  MaybeBij Maybe Bij
Nothing <> :: MaybeBij -> MaybeBij -> MaybeBij
<> MaybeBij
_ = Maybe Bij -> MaybeBij
MaybeBij forall a. Maybe a
Nothing
  MaybeBij
_ <> MaybeBij Maybe Bij
Nothing = Maybe Bij -> MaybeBij
MaybeBij forall a. Maybe a
Nothing
  MaybeBij (Just Bij
a) <> MaybeBij (Just Bij
b) = Maybe Bij -> MaybeBij
MaybeBij (Bij -> Bij -> Maybe Bij
combine Bij
a Bij
b)

instance Monoid MaybeBij where
  mempty :: MaybeBij
mempty = Maybe Bij -> MaybeBij
MaybeBij (forall a. a -> Maybe a
Just Bij
empty)

msingleton :: Int -> Int -> MaybeBij
msingleton :: Int -> Int -> MaybeBij
msingleton Int
i Int
j = Maybe Bij -> MaybeBij
MaybeBij forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Int -> Int -> Bij
singleton Int
i Int
j


closeBijection :: (Int -> Int -> Maybe Bij) -> Bij -> Maybe Bij
closeBijection :: (Int -> Int -> Maybe Bij) -> Bij -> Maybe Bij
closeBijection Int -> Int -> Maybe Bij
f Bij
s = let
  inner :: Bij -> Bij -> Maybe Bij
inner Bij
a Bij
n = case Bij -> Maybe ((Int, Int), Bij)
pop Bij
n of
    Maybe ((Int, Int), Bij)
Nothing -> forall a. a -> Maybe a
Just Bij
a
    Just ((Int
i,Int
j), Bij
n') -> case Int -> Int -> Maybe Bij
f Int
i Int
j of
      Maybe Bij
Nothing -> forall a. Maybe a
Nothing
      Just Bij
b -> case Bij -> Bij -> Maybe Bij
combine Bij
a Bij
b of
        Maybe Bij
Nothing -> forall a. Maybe a
Nothing
        Just Bij
a' -> Bij -> Bij -> Maybe Bij
inner Bij
a' (Bij -> Bij -> Bij
unsafeUnion Bij
n' (Bij -> Bij -> Bij
unsafeDifference Bij
b Bij
a))
  in Bij -> Bij -> Maybe Bij
inner Bij
s Bij
s