    -- | 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 =
      ((a, a) -> Neighbors a -> Neighbors a)
    -> Neighbors a -> Set (a, a) -> Neighbors a
    forall a b. (a -> b -> b) -> b -> Set a -> b
    S.foldr
        ( \(a
    x, a
    y) Neighbors a
    acc ->
            Neighbors a
    acc
              Neighbors a -> (Neighbors a -> Neighbors a) -> Neighbors a
    forall a b. a -> (a -> b) -> b
    & (Set a -> Set a -> Set a)
    -> a -> Set a -> Neighbors a -> Neighbors a
    forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
    M.insertWith Set a -> Set a -> Set a
    forall a. Ord a => Set a -> Set a -> Set a
    S.union a
    x (a -> Set a
    forall a. a -> Set a
    S.singleton a
    y)
              Neighbors a -> (Neighbors a -> Neighbors a) -> Neighbors a
    forall a b. a -> (a -> b) -> b
    & (Set a -> Set a -> Set a)
    -> a -> Set a -> Neighbors a -> Neighbors a
    forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
    M.insertWith Set a -> Set a -> Set a
    forall a. Ord a => Set a -> Set a -> Set a
    S.union a
    y (a -> Set a
    forall a. a -> Set a
    S.singleton a
    x)
        )
        Neighbors a
    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 Int -> Set Int -> Bool
    forall a. Ord a => a -> Set a -> Bool
    `S.member` Set Int
    xs, Map Int space
    spaces Map Int space -> Int -> Maybe space
    forall k a. Ord k => Map k a -> k -> Maybe a
    M.!? Int
    i) of
        (Bool
    False, Just space
    sp') | space
    sp' space -> space -> Bool
    forall a. Eq a => a -> a -> Bool
    == space
    sp -> (Map Int space
    spaces, Int
    i)
        (Bool
    False, Maybe space
    Nothing) -> (Int -> space -> Map Int space -> Map Int space
    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)
    _ -> Map Int space -> Set Int -> Int -> space -> (Map Int space, Int)
    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 Int -> Int -> Int
    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 =
            (a -> Set Int) -> Set a -> Set Int
    forall m a. Monoid m => (a -> m) -> Set a -> m
    forall (t :: * -> *) m a.
    (Foldable t, Monoid m) =>
    (a -> m) -> t a -> m
    foldMap (Set Int -> (Int -> Set Int) -> Maybe Int -> Set Int
    forall b a. b -> (a -> b) -> Maybe a -> b
    maybe Set Int
    forall a. Set a
    S.empty Int -> Set Int
    forall a. a -> Set a
    S.singleton (Maybe Int -> Set Int) -> (a -> Maybe Int) -> a -> Set Int
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Coloring a
    coloring M.!?)) (Set a -> Set Int) -> Set a -> Set Int
    forall a b. (a -> b) -> a -> b
    $
              Set a -> Maybe (Set a) -> Set a
    forall a. a -> Maybe a -> a
    fromMaybe Set a
    forall a. Monoid a => a
    mempty (Neighbors a
    nbs Neighbors a -> a -> Maybe (Set a)
    forall k a. Ord k => Map k a -> k -> Maybe a
    M.!? a
    x)
          (Map Int space
    spaces', Int
    color) = Map Int space -> Set Int -> Int -> space -> (Map Int space, Int)
    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', a -> Int -> Coloring a -> Coloring a
    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 = [(a, space)] -> Set (a, space)
    forall a. Ord a => [a] -> Set a
    S.fromList ([(a, space)] -> Set (a, space)) -> [(a, space)] -> Set (a, space)
    forall a b. (a -> b) -> a -> b
    $ Map a space -> [(a, space)]
    forall k a. Map k a -> [(k, a)]
    M.toList Map a space
    spaces
          nbs :: Neighbors a
    nbs = Graph a -> Neighbors a
    forall a. Ord a => Graph a -> Neighbors a
    neighbors Graph a
    graph
       in ((a, space)
     -> (Map Int space, Coloring a) -> (Map Int space, Coloring a))
    -> (Map Int space, Coloring a)
    -> Set (a, space)
    -> (Map Int space, Coloring a)
    forall a b. (a -> b -> b) -> b -> Set a -> b
    S.foldr (Neighbors a
    -> (a, space)
    -> (Map Int space, Coloring a)
    -> (Map Int space, Coloring a)
    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) (Map Int space, Coloring a)
    forall a. Monoid a => a
    mempty Set (a, space)
    nodes