mealy-decompose/src/Merger.hs

module Merger where

import Partition

import Control.Monad (replicateM)
import Control.Monad.IO.Class (liftIO)
import Control.Monad.Trans.State.Strict
import Data.Bifunctor (first)
import Data.List (minimumBy)
import Data.Map.Strict qualified as Map
import Data.Ord (comparing)
import Data.Set qualified as Set

data MergerStats = MergerStats
  { numberOfComponents :: Int
  , maximalComponent :: Int
  , totalSize :: Int
  }
  deriving (Eq, Ord, Read, Show)

data MergerAction = Stop | Continue
  deriving (Eq, Ord, Read, Show)

type MergerStrategy = MergerStats -> MergerAction

heuristicMerger :: Ord o => [(o, Partition)] -> MergerStrategy -> IO [([o], Partition)]
heuristicMerger components strategy = do
  projmap <- evalStateT (loop 2) (Map.fromList (fmap (first pure) components))
  return $ Map.assocs projmap
 where
  score ps p3 = numBlocks p3 - sum (fmap numBlocks ps)
  combine ops =
    let os = fmap fst ops
        ps = fmap snd ops
        p3 = foldr1 commonRefinement ps
     in ((os, p3), score ps p3)
  isSortedOn f ls = and $ zipWith (\a b -> f a < f b) ls (drop 1 ls)
  allCombs n projs = fmap combine . filter (isSortedOn fst) $ replicateM n projs
  minComb n projs = minimumBy (comparing snd) (allCombs n projs)
  safeStrategy ms@MergerStats{..}
    | numberOfComponents <= 1 = Stop
    | otherwise = strategy ms

  loop n = do
    projmap <- get
    let numberOfComponents = Map.size projmap
        componentSizes = fmap numBlocks . Map.elems $ projmap
        maximalComponent = maximum componentSizes
        totalSize = sum componentSizes
        ms = MergerStats{..}
    liftIO . print $ ms
    case safeStrategy ms of
      Stop -> return projmap
      Continue -> do
        let ((os, p3), _) = minComb n (Map.assocs projmap)
            o3 = mconcat os
            newProjmap = Map.insert o3 p3 . Map.withoutKeys projmap $ Set.fromList os
        put newProjmap
        if Map.size newProjmap < n
          then loop (Map.size newProjmap)
          else loop n