module Bench.Vector.Algo.HybCC (hybcc) where

import Data.Vector.Unboxed as V

hybcc :: (Int, Vector Int, Vector Int) -> Vector Int
{-# NOINLINE hybcc #-}
hybcc :: (Int, Vector Int, Vector Int) -> Vector Int
hybcc (Int
n, Vector Int
e1, Vector Int
e2) = Vector (Int, Int) -> Int -> Vector Int
concomp (Vector Int -> Vector Int -> Vector (Int, Int)
forall a b.
(Unbox a, Unbox b) =>
Vector a -> Vector b -> Vector (a, b)
V.zip Vector Int
e1 Vector Int
e2) Int
n
    where
      concomp :: Vector (Int, Int) -> Int -> Vector Int
concomp Vector (Int, Int)
es Int
n
        | Vector (Int, Int) -> Bool
forall a. Unbox a => Vector a -> Bool
V.null Vector (Int, Int)
es = Int -> Int -> Vector Int
forall a. (Unbox a, Enum a) => a -> a -> Vector a
V.enumFromTo Int
0 (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
        | Bool
otherwise = Vector Int -> Vector Int -> Vector Int
forall a. Unbox a => Vector a -> Vector Int -> Vector a
V.backpermute Vector Int
ins Vector Int
ins
        where
          p :: Vector Int
p = Vector Int -> Vector Int
shortcut_all
            (Vector Int -> Vector Int) -> Vector Int -> Vector Int
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector (Int, Int) -> Vector Int
forall a. Unbox a => Vector a -> Vector (Int, a) -> Vector a
V.update (Int -> Int -> Vector Int
forall a. (Unbox a, Enum a) => a -> a -> Vector a
V.enumFromTo Int
0 (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)) Vector (Int, Int)
es

          (Vector (Int, Int)
es',Vector Int
i) = Vector Int -> Vector (Int, Int) -> (Vector (Int, Int), Vector Int)
forall {b}.
(Num b, Unbox b) =>
Vector Int -> Vector (Int, Int) -> (Vector (b, b), Vector Int)
compress Vector Int
p Vector (Int, Int)
es
          r :: Vector Int
r = Vector (Int, Int) -> Int -> Vector Int
concomp Vector (Int, Int)
es' (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
V.length Vector Int
i)
          ins :: Vector Int
ins = Vector Int -> Vector Int -> Vector Int -> Vector Int
forall a. Unbox a => Vector a -> Vector Int -> Vector a -> Vector a
V.update_ Vector Int
p Vector Int
i
              (Vector Int -> Vector Int) -> Vector Int -> Vector Int
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Int -> Vector Int
forall a. Unbox a => Vector a -> Vector Int -> Vector a
V.backpermute Vector Int
i Vector Int
r

      enumerate :: Vector Bool -> Vector b
enumerate Vector Bool
bs = (b -> b -> b) -> b -> Vector b -> Vector b
forall a b.
(Unbox a, Unbox b) =>
(a -> b -> a) -> a -> Vector b -> Vector a
V.prescanl' b -> b -> b
forall a. Num a => a -> a -> a
(+) b
0 (Vector b -> Vector b) -> Vector b -> Vector b
forall a b. (a -> b) -> a -> b
$ (Bool -> b) -> Vector Bool -> Vector b
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
V.map (\Bool
b -> if Bool
b then b
1 else b
0) Vector Bool
bs

      pack_index :: Vector Bool -> Vector Int
pack_index Vector Bool
bs = ((Int, Bool) -> Int) -> Vector (Int, Bool) -> Vector Int
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
V.map (Int, Bool) -> Int
forall a b. (a, b) -> a
fst
                    (Vector (Int, Bool) -> Vector Int)
-> (Vector (Int, Bool) -> Vector (Int, Bool))
-> Vector (Int, Bool)
-> Vector Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Int, Bool) -> Bool) -> Vector (Int, Bool) -> Vector (Int, Bool)
forall a. Unbox a => (a -> Bool) -> Vector a -> Vector a
V.filter (Int, Bool) -> Bool
forall a b. (a, b) -> b
snd
                    (Vector (Int, Bool) -> Vector Int)
-> Vector (Int, Bool) -> Vector Int
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Bool -> Vector (Int, Bool)
forall a b.
(Unbox a, Unbox b) =>
Vector a -> Vector b -> Vector (a, b)
V.zip (Int -> Int -> Vector Int
forall a. (Unbox a, Enum a) => a -> a -> Vector a
V.enumFromTo Int
0 (Vector Bool -> Int
forall a. Unbox a => Vector a -> Int
V.length Vector Bool
bs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) Vector Bool
bs

      shortcut_all :: Vector Int -> Vector Int
shortcut_all Vector Int
p | Vector Int
p Vector Int -> Vector Int -> Bool
forall a. Eq a => a -> a -> Bool
== Vector Int
pp   = Vector Int
pp
                     | Bool
otherwise = Vector Int -> Vector Int
shortcut_all Vector Int
pp
        where
          pp :: Vector Int
pp = Vector Int -> Vector Int -> Vector Int
forall a. Unbox a => Vector a -> Vector Int -> Vector a
V.backpermute Vector Int
p Vector Int
p

      compress :: Vector Int -> Vector (Int, Int) -> (Vector (b, b), Vector Int)
compress Vector Int
p Vector (Int, Int)
es = (Vector (b, b)
new_es, Vector Bool -> Vector Int
pack_index Vector Bool
roots)
        where
          (Vector Int
e1,Vector Int
e2) = Vector (Int, Int) -> (Vector Int, Vector Int)
forall a b.
(Unbox a, Unbox b) =>
Vector (a, b) -> (Vector a, Vector b)
V.unzip Vector (Int, Int)
es
          es' :: Vector (Int, Int)
es' = ((Int, Int) -> (Int, Int))
-> Vector (Int, Int) -> Vector (Int, Int)
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
V.map (\(Int
x,Int
y) -> if Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
y then (Int
y,Int
x) else (Int
x,Int
y))
              (Vector (Int, Int) -> Vector (Int, Int))
-> (Vector (Int, Int) -> Vector (Int, Int))
-> Vector (Int, Int)
-> Vector (Int, Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Int, Int) -> Bool) -> Vector (Int, Int) -> Vector (Int, Int)
forall a. Unbox a => (a -> Bool) -> Vector a -> Vector a
V.filter (\(Int
x,Int
y) -> Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
y)
              (Vector (Int, Int) -> Vector (Int, Int))
-> Vector (Int, Int) -> Vector (Int, Int)
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Int -> Vector (Int, Int)
forall a b.
(Unbox a, Unbox b) =>
Vector a -> Vector b -> Vector (a, b)
V.zip (Vector Int -> Vector Int -> Vector Int
forall a. Unbox a => Vector a -> Vector Int -> Vector a
V.backpermute Vector Int
p Vector Int
e1) (Vector Int -> Vector Int -> Vector Int
forall a. Unbox a => Vector a -> Vector Int -> Vector a
V.backpermute Vector Int
p Vector Int
e2)

          roots :: Vector Bool
roots = (Int -> Int -> Bool) -> Vector Int -> Vector Int -> Vector Bool
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
V.zipWith Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
(==) Vector Int
p (Int -> Int -> Vector Int
forall a. (Unbox a, Enum a) => a -> a -> Vector a
V.enumFromTo Int
0 (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
V.length Vector Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
          labels :: Vector b
labels = Vector Bool -> Vector b
forall {b}. (Unbox b, Num b) => Vector Bool -> Vector b
enumerate Vector Bool
roots
          (Vector Int
e1',Vector Int
e2') = Vector (Int, Int) -> (Vector Int, Vector Int)
forall a b.
(Unbox a, Unbox b) =>
Vector (a, b) -> (Vector a, Vector b)
V.unzip Vector (Int, Int)
es'
          new_es :: Vector (b, b)
new_es = Vector b -> Vector b -> Vector (b, b)
forall a b.
(Unbox a, Unbox b) =>
Vector a -> Vector b -> Vector (a, b)
V.zip (Vector b -> Vector Int -> Vector b
forall a. Unbox a => Vector a -> Vector Int -> Vector a
V.backpermute Vector b
labels Vector Int
e1') (Vector b -> Vector Int -> Vector b
forall a. Unbox a => Vector a -> Vector Int -> Vector a
V.backpermute Vector b
labels Vector Int
e2')