-- | Provides a greedy graph-coloring algorithm.
module Futhark.Optimise.MemoryBlockMerging.GreedyColoring (colorGraph, Coloring) where

import Data.Function ((&))
import Data.Map qualified as M
import Data.Maybe (fromMaybe)
import Data.Set qualified as S
import Futhark.Analysis.Interference qualified as Interference

-- | A map of values to their color, identified by an integer.
type Coloring a = M.Map a Int

-- | A map of values to the set "neighbors" in the graph
type Neighbors a = M.Map a (S.Set a)

-- | Computes the neighbor map of a graph.
neighbors :: Ord a => Interference.Graph a -> Neighbors a
neighbors :: forall a. Ord a => Graph a -> Neighbors a
neighbors =
  forall a b. (a -> b -> b) -> b -> Set a -> b
S.foldr
    ( \(a
x, a
y) Neighbors a
acc ->
        Neighbors a
acc
          forall a b. a -> (a -> b) -> b
& forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith forall a. Ord a => Set a -> Set a -> Set a
S.union a
x (forall a. a -> Set a
S.singleton a
y)
          forall a b. a -> (a -> b) -> b
& forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith forall a. Ord a => Set a -> Set a -> Set a
S.union a
y (forall a. a -> Set a
S.singleton a
x)
    )
    forall k a. Map k a
M.empty

firstAvailable :: Eq space => M.Map Int space -> S.Set Int -> Int -> space -> (M.Map Int space, Int)
firstAvailable :: forall space.
Eq space =>
Map Int space -> Set Int -> Int -> space -> (Map Int space, Int)
firstAvailable Map Int space
spaces Set Int
xs Int
i space
sp =
  case (Int
i forall a. Ord a => a -> Set a -> Bool
`S.member` Set Int
xs, Map Int space
spaces forall k a. Ord k => Map k a -> k -> Maybe a
M.!? Int
i) of
    (Bool
False, Just space
sp') | space
sp' forall a. Eq a => a -> a -> Bool
== space
sp -> (Map Int space
spaces, Int
i)
    (Bool
False, Maybe space
Nothing) -> (forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Int
i space
sp Map Int space
spaces, Int
i)
    (Bool, Maybe space)
_ -> forall space.
Eq space =>
Map Int space -> Set Int -> Int -> space -> (Map Int space, Int)
firstAvailable Map Int space
spaces Set Int
xs (Int
i forall a. Num a => a -> a -> a
+ Int
1) space
sp

colorNode ::
  (Ord a, Eq space) =>
  Neighbors a ->
  (a, space) ->
  (M.Map Int space, Coloring a) ->
  (M.Map Int space, Coloring a)
colorNode :: forall a space.
(Ord a, Eq space) =>
Neighbors a
-> (a, space)
-> (Map Int space, Coloring a)
-> (Map Int space, Coloring a)
colorNode Neighbors a
nbs (a
x, space
sp) (Map Int space
spaces, Coloring a
coloring) =
  let nb_colors :: Set Int
nb_colors =
        forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. Set a
S.empty forall a. a -> Set a
S.singleton forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Coloring a
coloring M.!?)) forall a b. (a -> b) -> a -> b
$
          forall a. a -> Maybe a -> a
fromMaybe forall a. Monoid a => a
mempty (Neighbors a
nbs forall k a. Ord k => Map k a -> k -> Maybe a
M.!? a
x)
      (Map Int space
spaces', Int
color) = forall space.
Eq space =>
Map Int space -> Set Int -> Int -> space -> (Map Int space, Int)
firstAvailable Map Int space
spaces Set Int
nb_colors Int
0 space
sp
   in (Map Int space
spaces', forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
x Int
color Coloring a
coloring)

-- | Graph coloring that takes into account the @space@ of values. Two values
-- can only share the same color if they live in the same space. The result is
-- map from each color to a space and a map from each value in the input graph
-- to it's new color.
colorGraph ::
  (Ord a, Ord space) =>
  M.Map a space ->
  Interference.Graph a ->
  (M.Map Int space, Coloring a)
colorGraph :: forall a space.
(Ord a, Ord space) =>
Map a space -> Graph a -> (Map Int space, Coloring a)
colorGraph Map a space
spaces Graph a
graph =
  let nodes :: Set (a, space)
nodes = forall a. Ord a => [a] -> Set a
S.fromList forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [(k, a)]
M.toList Map a space
spaces
      nbs :: Neighbors a
nbs = forall a. Ord a => Graph a -> Neighbors a
neighbors Graph a
graph
   in forall a b. (a -> b -> b) -> b -> Set a -> b
S.foldr (forall a space.
(Ord a, Eq space) =>
Neighbors a
-> (a, space)
-> (Map Int space, Coloring a)
-> (Map Int space, Coloring a)
colorNode Neighbors a
nbs) forall a. Monoid a => a
mempty Set (a, space)
nodes