{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE BangPatterns #-}
{-# OPTIONS_GHC -funbox-strict-fields #-}
--
-- NOTE: Does not currently ensure that a 'Point' was indeed generated
-- by the specified 'PointSupply'.
--
module Data.UnionFind.IntMap 
    ( newPointSupply, fresh, repr, descriptor, union, equivalent,
      PointSupply, Point ) where

import qualified Data.IntMap as IM

data PointSupply a = PointSupply !Int (IM.IntMap (Link a))
  deriving Int -> PointSupply a -> ShowS
[PointSupply a] -> ShowS
PointSupply a -> String
(Int -> PointSupply a -> ShowS)
-> (PointSupply a -> String)
-> ([PointSupply a] -> ShowS)
-> Show (PointSupply a)
forall a. Show a => Int -> PointSupply a -> ShowS
forall a. Show a => [PointSupply a] -> ShowS
forall a. Show a => PointSupply a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [PointSupply a] -> ShowS
$cshowList :: forall a. Show a => [PointSupply a] -> ShowS
show :: PointSupply a -> String
$cshow :: forall a. Show a => PointSupply a -> String
showsPrec :: Int -> PointSupply a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> PointSupply a -> ShowS
Show

data Link a 
    = Info {-# UNPACK #-} !Int a
      -- ^ This is the descriptive element of the equivalence class
      -- and its rank.
    | Link {-# UNPACK #-} !Int
      -- ^ Pointer to some other element of the equivalence class.
     deriving Int -> Link a -> ShowS
[Link a] -> ShowS
Link a -> String
(Int -> Link a -> ShowS)
-> (Link a -> String) -> ([Link a] -> ShowS) -> Show (Link a)
forall a. Show a => Int -> Link a -> ShowS
forall a. Show a => [Link a] -> ShowS
forall a. Show a => Link a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Link a] -> ShowS
$cshowList :: forall a. Show a => [Link a] -> ShowS
show :: Link a -> String
$cshow :: forall a. Show a => Link a -> String
showsPrec :: Int -> Link a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Link a -> ShowS
Show

newtype Point a = Point Int

newPointSupply :: PointSupply a
newPointSupply :: PointSupply a
newPointSupply = Int -> IntMap (Link a) -> PointSupply a
forall a. Int -> IntMap (Link a) -> PointSupply a
PointSupply Int
0 IntMap (Link a)
forall a. IntMap a
IM.empty

fresh :: PointSupply a -> a -> (PointSupply a, Point a)
fresh :: PointSupply a -> a -> (PointSupply a, Point a)
fresh (PointSupply Int
next IntMap (Link a)
eqs) a
a =
  (Int -> IntMap (Link a) -> PointSupply a
forall a. Int -> IntMap (Link a) -> PointSupply a
PointSupply (Int
next Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> Link a -> IntMap (Link a) -> IntMap (Link a)
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert Int
next (Int -> a -> Link a
forall a. Int -> a -> Link a
Info Int
0 a
a) IntMap (Link a)
eqs), Int -> Point a
forall a. Int -> Point a
Point Int
next)

-- freshList :: PointSupply a -> [a] -> (PointSupply a, [Point a])
-- freshList 

repr :: PointSupply a -> Point a -> Point a
repr :: PointSupply a -> Point a -> Point a
repr PointSupply a
ps Point a
p = PointSupply a -> Point a -> (Int -> Int -> a -> Point a) -> Point a
forall a r. PointSupply a -> Point a -> (Int -> Int -> a -> r) -> r
reprInfo PointSupply a
ps Point a
p (\Int
n Int
_rank a
_a -> Int -> Point a
forall a. Int -> Point a
Point Int
n)

reprInfo :: PointSupply a -> Point a -> (Int -> Int -> a -> r) -> r
reprInfo :: PointSupply a -> Point a -> (Int -> Int -> a -> r) -> r
reprInfo (PointSupply Int
_next IntMap (Link a)
eqs) (Point Int
n) Int -> Int -> a -> r
k = Int -> r
go Int
n
  where
    go :: Int -> r
go !Int
i =
      case IntMap (Link a)
eqs IntMap (Link a) -> Int -> Link a
forall a. IntMap a -> Int -> a
IM.! Int
i of
        Link Int
i' -> Int -> r
go Int
i'
        Info Int
r a
a -> Int -> Int -> a -> r
k Int
i Int
r a
a
  
union :: PointSupply a -> Point a -> Point a -> PointSupply a
union :: PointSupply a -> Point a -> Point a -> PointSupply a
union ps :: PointSupply a
ps@(PointSupply Int
next IntMap (Link a)
eqs) Point a
p1 Point a
p2 =
  PointSupply a
-> Point a -> (Int -> Int -> a -> PointSupply a) -> PointSupply a
forall a r. PointSupply a -> Point a -> (Int -> Int -> a -> r) -> r
reprInfo PointSupply a
ps Point a
p1 ((Int -> Int -> a -> PointSupply a) -> PointSupply a)
-> (Int -> Int -> a -> PointSupply a) -> PointSupply a
forall a b. (a -> b) -> a -> b
$ \Int
i1 Int
r1 a
_a1 -> 
  PointSupply a
-> Point a -> (Int -> Int -> a -> PointSupply a) -> PointSupply a
forall a r. PointSupply a -> Point a -> (Int -> Int -> a -> r) -> r
reprInfo PointSupply a
ps Point a
p2 ((Int -> Int -> a -> PointSupply a) -> PointSupply a)
-> (Int -> Int -> a -> PointSupply a) -> PointSupply a
forall a b. (a -> b) -> a -> b
$ \Int
i2 Int
r2 a
a2 ->
  if Int
i1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
i2 then PointSupply a
ps else
    case Int
r1 Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` Int
r2 of
      Ordering
LT ->
        -- No rank or descriptor update necessary
        let !eqs1 :: IntMap (Link a)
eqs1 = Int -> Link a -> IntMap (Link a) -> IntMap (Link a)
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert Int
i1 (Int -> Link a
forall a. Int -> Link a
Link Int
i2) IntMap (Link a)
eqs in
        Int -> IntMap (Link a) -> PointSupply a
forall a. Int -> IntMap (Link a) -> PointSupply a
PointSupply Int
next IntMap (Link a)
eqs1
      Ordering
EQ ->
        let !eqs1 :: IntMap (Link a)
eqs1 = Int -> Link a -> IntMap (Link a) -> IntMap (Link a)
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert Int
i1 (Int -> Link a
forall a. Int -> Link a
Link Int
i2) IntMap (Link a)
eqs
            !eqs2 :: IntMap (Link a)
eqs2 = Int -> Link a -> IntMap (Link a) -> IntMap (Link a)
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert Int
i2 (Int -> a -> Link a
forall a. Int -> a -> Link a
Info (Int
r2 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) a
a2) IntMap (Link a)
eqs1 in
        Int -> IntMap (Link a) -> PointSupply a
forall a. Int -> IntMap (Link a) -> PointSupply a
PointSupply Int
next IntMap (Link a)
eqs2
      Ordering
GT ->
        let !eqs1 :: IntMap (Link a)
eqs1 = Int -> Link a -> IntMap (Link a) -> IntMap (Link a)
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert Int
i1 (Int -> a -> Link a
forall a. Int -> a -> Link a
Info Int
r2 a
a2) IntMap (Link a)
eqs
            !eqs2 :: IntMap (Link a)
eqs2 = Int -> Link a -> IntMap (Link a) -> IntMap (Link a)
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert Int
i2 (Int -> Link a
forall a. Int -> Link a
Link Int
i1) IntMap (Link a)
eqs1 in
        Int -> IntMap (Link a) -> PointSupply a
forall a. Int -> IntMap (Link a) -> PointSupply a
PointSupply Int
next IntMap (Link a)
eqs2

descriptor :: PointSupply a -> Point a -> a
descriptor :: PointSupply a -> Point a -> a
descriptor PointSupply a
ps Point a
p = PointSupply a -> Point a -> (Int -> Int -> a -> a) -> a
forall a r. PointSupply a -> Point a -> (Int -> Int -> a -> r) -> r
reprInfo PointSupply a
ps Point a
p (\Int
_ Int
_ a
a -> a
a)

equivalent :: PointSupply a -> Point a -> Point a -> Bool
equivalent :: PointSupply a -> Point a -> Point a -> Bool
equivalent PointSupply a
ps Point a
p1 Point a
p2 =
  PointSupply a -> Point a -> (Int -> Int -> a -> Bool) -> Bool
forall a r. PointSupply a -> Point a -> (Int -> Int -> a -> r) -> r
reprInfo PointSupply a
ps Point a
p1 ((Int -> Int -> a -> Bool) -> Bool)
-> (Int -> Int -> a -> Bool) -> Bool
forall a b. (a -> b) -> a -> b
$ \Int
i1 Int
_ a
_ ->
  PointSupply a -> Point a -> (Int -> Int -> a -> Bool) -> Bool
forall a r. PointSupply a -> Point a -> (Int -> Int -> a -> r) -> r
reprInfo PointSupply a
ps Point a
p2 ((Int -> Int -> a -> Bool) -> Bool)
-> (Int -> Int -> a -> Bool) -> Bool
forall a b. (a -> b) -> a -> b
$ \Int
i2 Int
_ a
_ ->
  Int
i1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
i2

{-
tst1 :: IO ()
tst1 = do
  let ps0 = newPointSupply
      (ps1, p1) = fresh ps0 "hello"
      (ps2, p2) = fresh ps1 "world"
      (ps3, p3) = fresh ps2 "you"
      (ps, p4) = fresh ps3 "there"
  let ps' = union ps p1 p2
  print (descr ps p1, descr ps p2, equivalent ps p1 p2)
  print (descr ps' p1, descr ps' p2, equivalent ps' p1 p2)
  let ps'' = union ps' p3 p1
  print (descr ps'' p1, descr ps'' p3, equivalent ps'' p1 p3)
  print ps''
-}

-- TODO: should fail
{-
tst2 :: IO ()
tst2 = do
  let as0 = newPointSupply
      (as, a1) = fresh as0 "foo"
      bs0 = newPointSupply
      (bs, b1) = fresh bs0 "bar"
  print $ union as a1 b1
  -}