module Futhark.Optimise.MemoryBlockMerging.GreedyColoringTests
  ( tests,
  )
where

import Control.Arrow ((***))
import Data.Function ((&))
import Data.Map qualified as M
import Data.Set qualified as S
import Futhark.Optimise.MemoryBlockMerging.GreedyColoring qualified as GreedyColoring
import Test.Tasty
import Test.Tasty.HUnit

tests :: TestTree
tests :: TestTree
tests =
  String -> [TestTree] -> TestTree
testGroup
    String
"GreedyColoringTests"
    [TestTree
psumTest, TestTree
allIntersect, TestTree
emptyGraph, TestTree
noIntersections, TestTree
differentSpaces]

psumTest :: TestTree
psumTest :: TestTree
psumTest =
  String -> Assertion -> TestTree
testCase String
"psumTest"
    (Assertion -> TestTree) -> Assertion -> TestTree
forall a b. (a -> b) -> a -> b
$ String
-> ([(Int, String)], [(Int, Int)])
-> ([(Int, String)], [(Int, Int)])
-> Assertion
forall a.
(Eq a, Show a, HasCallStack) =>
String -> a -> a -> Assertion
assertEqual
      String
"Color simple 1-2-3 using two colors"
      ( [(Int
0, String
"shared"), (Int
1, String
"shared")] :: [(Int, String)],
        [(Int
1 :: Int, Int
0), (Int
2, Int
1), (Int
3, Int
0)]
      )
    (([(Int, String)], [(Int, Int)]) -> Assertion)
-> ([(Int, String)], [(Int, Int)]) -> Assertion
forall a b. (a -> b) -> a -> b
$ (Map Int String -> [(Int, String)]
forall k a. Map k a -> [(k, a)]
M.toList (Map Int String -> [(Int, String)])
-> (Map Int Int -> [(Int, Int)])
-> (Map Int String, Map Int Int)
-> ([(Int, String)], [(Int, Int)])
forall b c b' c'. (b -> c) -> (b' -> c') -> (b, b') -> (c, c')
forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** Map Int Int -> [(Int, Int)]
forall k a. Map k a -> [(k, a)]
M.toList)
    ((Map Int String, Map Int Int) -> ([(Int, String)], [(Int, Int)]))
-> (Map Int String, Map Int Int) -> ([(Int, String)], [(Int, Int)])
forall a b. (a -> b) -> a -> b
$ Map Int String -> Graph Int -> (Map Int String, Map Int Int)
forall a space.
(Ord a, Ord space) =>
Map a space -> Graph a -> (Map Int space, Coloring a)
GreedyColoring.colorGraph
      ([(Int, String)] -> Map Int String
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Int
1, String
"shared"), (Int
2, String
"shared"), (Int
3, String
"shared")])
    (Graph Int -> (Map Int String, Map Int Int))
-> Graph Int -> (Map Int String, Map Int Int)
forall a b. (a -> b) -> a -> b
$ [(Int, Int)] -> Graph Int
forall a. Ord a => [a] -> Set a
S.fromList [(Int
1, Int
2), (Int
2, Int
3)]

allIntersect :: TestTree
allIntersect :: TestTree
allIntersect =
  String -> Assertion -> TestTree
testCase String
"allIntersect"
    (Assertion -> TestTree) -> Assertion -> TestTree
forall a b. (a -> b) -> a -> b
$ String
-> ([(Int, String)], [(Int, Int)])
-> ([(Int, String)], [(Int, Int)])
-> Assertion
forall a.
(Eq a, Show a, HasCallStack) =>
String -> a -> a -> Assertion
assertEqual
      String
"Color a graph where all values intersect"
      ( [(Int
0, String
"shared"), (Int
1, String
"shared"), (Int
2, String
"shared")] :: [(Int, String)],
        [(Int
1 :: Int, Int
2), (Int
2, Int
1), (Int
3, Int
0)]
      )
    (([(Int, String)], [(Int, Int)]) -> Assertion)
-> ([(Int, String)], [(Int, Int)]) -> Assertion
forall a b. (a -> b) -> a -> b
$ (Map Int String -> [(Int, String)]
forall k a. Map k a -> [(k, a)]
M.toList (Map Int String -> [(Int, String)])
-> (Map Int Int -> [(Int, Int)])
-> (Map Int String, Map Int Int)
-> ([(Int, String)], [(Int, Int)])
forall b c b' c'. (b -> c) -> (b' -> c') -> (b, b') -> (c, c')
forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** Map Int Int -> [(Int, Int)]
forall k a. Map k a -> [(k, a)]
M.toList)
    ((Map Int String, Map Int Int) -> ([(Int, String)], [(Int, Int)]))
-> (Map Int String, Map Int Int) -> ([(Int, String)], [(Int, Int)])
forall a b. (a -> b) -> a -> b
$ Map Int String -> Graph Int -> (Map Int String, Map Int Int)
forall a space.
(Ord a, Ord space) =>
Map a space -> Graph a -> (Map Int space, Coloring a)
GreedyColoring.colorGraph
      ([(Int, String)] -> Map Int String
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Int
1, String
"shared"), (Int
2, String
"shared"), (Int
3, String
"shared")])
    (Graph Int -> (Map Int String, Map Int Int))
-> Graph Int -> (Map Int String, Map Int Int)
forall a b. (a -> b) -> a -> b
$ [(Int, Int)] -> Graph Int
forall a. Ord a => [a] -> Set a
S.fromList [(Int
1, Int
2), (Int
2, Int
3), (Int
1, Int
3)]

emptyGraph :: TestTree
emptyGraph :: TestTree
emptyGraph =
  String -> Assertion -> TestTree
testCase String
"emptyGraph"
    (Assertion -> TestTree) -> Assertion -> TestTree
forall a b. (a -> b) -> a -> b
$ String
-> ([(Int, Char)], [(Int, Int)])
-> ([(Int, Char)], [(Int, Int)])
-> Assertion
forall a.
(Eq a, Show a, HasCallStack) =>
String -> a -> a -> Assertion
assertEqual
      String
"Color an empty graph"
      ([] :: [(Int, Char)], [] :: [(Int, Int)])
    (([(Int, Char)], [(Int, Int)]) -> Assertion)
-> ([(Int, Char)], [(Int, Int)]) -> Assertion
forall a b. (a -> b) -> a -> b
$ (Map Int Char -> [(Int, Char)]
forall k a. Map k a -> [(k, a)]
M.toList (Map Int Char -> [(Int, Char)])
-> (Map Int Int -> [(Int, Int)])
-> (Map Int Char, Map Int Int)
-> ([(Int, Char)], [(Int, Int)])
forall b c b' c'. (b -> c) -> (b' -> c') -> (b, b') -> (c, c')
forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** Map Int Int -> [(Int, Int)]
forall k a. Map k a -> [(k, a)]
M.toList)
    ((Map Int Char, Map Int Int) -> ([(Int, Char)], [(Int, Int)]))
-> (Map Int Char, Map Int Int) -> ([(Int, Char)], [(Int, Int)])
forall a b. (a -> b) -> a -> b
$ Map Int Char -> Graph Int -> (Map Int Char, Map Int Int)
forall a space.
(Ord a, Ord space) =>
Map a space -> Graph a -> (Map Int space, Coloring a)
GreedyColoring.colorGraph Map Int Char
forall k a. Map k a
M.empty
    (Graph Int -> (Map Int Char, Map Int Int))
-> Graph Int -> (Map Int Char, Map Int Int)
forall a b. (a -> b) -> a -> b
$ [(Int, Int)] -> Graph Int
forall a. Ord a => [a] -> Set a
S.fromList []

noIntersections :: TestTree
noIntersections :: TestTree
noIntersections =
  Map Int String -> Graph Int -> (Map Int String, Map Int Int)
forall a space.
(Ord a, Ord space) =>
Map a space -> Graph a -> (Map Int space, Coloring a)
GreedyColoring.colorGraph
    ([(Int, String)] -> Map Int String
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Int
1, String
"shared"), (Int
2, String
"shared"), (Int
3, String
"shared")])
    ([(Int, Int)] -> Graph Int
forall a. Ord a => [a] -> Set a
S.fromList [])
    (Map Int String, Map Int Int)
-> ((Map Int String, Map Int Int)
    -> ([(Int, String)], [(Int, Int)]))
-> ([(Int, String)], [(Int, Int)])
forall a b. a -> (a -> b) -> b
& Map Int String -> [(Int, String)]
forall k a. Map k a -> [(k, a)]
M.toList (Map Int String -> [(Int, String)])
-> (Map Int Int -> [(Int, Int)])
-> (Map Int String, Map Int Int)
-> ([(Int, String)], [(Int, Int)])
forall b c b' c'. (b -> c) -> (b' -> c') -> (b, b') -> (c, c')
forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** Map Int Int -> [(Int, Int)]
forall k a. Map k a -> [(k, a)]
M.toList
    ([(Int, String)], [(Int, Int)])
-> (([(Int, String)], [(Int, Int)]) -> Assertion) -> Assertion
forall a b. a -> (a -> b) -> b
& String
-> ([(Int, String)], [(Int, Int)])
-> ([(Int, String)], [(Int, Int)])
-> Assertion
forall a.
(Eq a, Show a, HasCallStack) =>
String -> a -> a -> Assertion
assertEqual
      String
"Color nodes with no intersections"
      ( [(Int
0, String
"shared")] :: [(Int, String)],
        [(Int
1, Int
0), (Int
2, Int
0), (Int
3, Int
0)] :: [(Int, Int)]
      )
    Assertion -> (Assertion -> TestTree) -> TestTree
forall a b. a -> (a -> b) -> b
& String -> Assertion -> TestTree
testCase String
"noIntersections"

differentSpaces :: TestTree
differentSpaces :: TestTree
differentSpaces =
  Map Int String -> Graph Int -> (Map Int String, Map Int Int)
forall a space.
(Ord a, Ord space) =>
Map a space -> Graph a -> (Map Int space, Coloring a)
GreedyColoring.colorGraph
    ([(Int, String)] -> Map Int String
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Int
1, String
"a"), (Int
2, String
"b"), (Int
3, String
"c")])
    ([(Int, Int)] -> Graph Int
forall a. Ord a => [a] -> Set a
S.fromList [])
    (Map Int String, Map Int Int)
-> ((Map Int String, Map Int Int)
    -> ([(Int, String)], [(Int, Int)]))
-> ([(Int, String)], [(Int, Int)])
forall a b. a -> (a -> b) -> b
& Map Int String -> [(Int, String)]
forall k a. Map k a -> [(k, a)]
M.toList (Map Int String -> [(Int, String)])
-> (Map Int Int -> [(Int, Int)])
-> (Map Int String, Map Int Int)
-> ([(Int, String)], [(Int, Int)])
forall b c b' c'. (b -> c) -> (b' -> c') -> (b, b') -> (c, c')
forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** Map Int Int -> [(Int, Int)]
forall k a. Map k a -> [(k, a)]
M.toList
    ([(Int, String)], [(Int, Int)])
-> (([(Int, String)], [(Int, Int)]) -> Assertion) -> Assertion
forall a b. a -> (a -> b) -> b
& String
-> ([(Int, String)], [(Int, Int)])
-> ([(Int, String)], [(Int, Int)])
-> Assertion
forall a.
(Eq a, Show a, HasCallStack) =>
String -> a -> a -> Assertion
assertEqual
      String
"Color nodes with no intersections but in different spaces"
      ( [(Int
0, String
"c"), (Int
1, String
"b"), (Int
2, String
"a")] :: [(Int, String)],
        [(Int
1, Int
2), (Int
2, Int
1), (Int
3, Int
0)] :: [(Int, Int)]
      )
    Assertion -> (Assertion -> TestTree) -> TestTree
forall a b. a -> (a -> b) -> b
& String -> Assertion -> TestTree
testCase String
"differentSpaces"