Refactored some code

app/Main.hs

@@ -3,13 +3,12 @@ module Main where
 import DotParser
 import Mealy
 import MealyRefine
+import Merger
 import Partition
 
-import Control.Monad (forM_, when, replicateM)
-import Control.Monad.IO.Class (liftIO)
+import Control.Monad (forM_, when)
 import Control.Monad.Trans.State.Strict
-import Data.Function (on)
-import Data.List (minimumBy, maximum, sort, intercalate)
+import Data.List (sort)
 import Data.Map.Strict qualified as Map
 import Data.Maybe (mapMaybe)
 import Data.Set qualified as Set
@@ -139,37 +138,12 @@ main = do
     putStrLn $ "  " <> (show o) <> " has size " <> (show b)
     )
 
+  let strategy MergerStats{..}
+        | numberOfComponents <= 4 = Stop
+        | otherwise = Continue
 
-  -- Now let's combine components to minimise the total size
-  let totalSize = sum (fmap (numBlocks . snd) topMods)
+  projmap <- heuristicMerger topMods strategy
 
-  putStrLn ""
-  putStrLn $ "num = " <> show (length topMods) <> ", size = " <> show totalSize
-
-  let 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 (compare `on` snd) (allCombs n projs)
-
-  let loop 1 = return ()
-      loop n = do
-        (projmap, currentSize) <- get
-        -- liftIO . print . fmap numBlocks . Map.elems $ projmap
-        let ((os, p3), gain) = minComb n (Map.assocs projmap)
-            o3 = intercalate "x" os
-            newSize = currentSize + gain
-            newProjmap = Map.insert o3 p3 . Map.withoutKeys projmap $ Set.fromList os
-        liftIO . putStrLn $ show o3 <> " -> num = " <> show (Map.size newProjmap) <> ", size = " <> show newSize <> ", max = " <> show (maximum . fmap numBlocks . Map.elems $ newProjmap)
-        put (newProjmap, newSize)
-        if Map.size newProjmap < n
-          then loop (Map.size newProjmap)
-          else loop n
-
-  putStrLn "2"
-  _ <- execStateT (loop 2) (Map.fromList topMods, totalSize)
+  print . fmap fst $ projmap
 
   return ()

mealy-decompose.cabal

@@ -28,6 +28,7 @@ library
     DotParser,
     Mealy,
     MealyRefine,
+    Merger,
     Partition
   build-depends:
     vector

src/Merger.hs

@@ -0,0 +1,60 @@
+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.Function (on)
+import Data.List (minimumBy)
+import Data.Map.Strict qualified as Map
+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 (compare `on` 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