mealy-decompose/app/Main.hs

module Main where

import DotParser
import DotWriter
import Mealy
import MealyRefine
import Merger
import Partition
import Preorder

import Control.Monad (forM_)
import Data.Bifunctor
import Data.List (sort, sortOn, intercalate)
import Data.List.Ordered (nubSort)
import Data.Map.Strict qualified as Map
import Data.Maybe (mapMaybe)
import Data.Set qualified as Set
import Data.Tuple (swap)
import System.Environment
import Text.Megaparsec

converseRelation :: (Ord a, Ord b) => Map.Map a b -> Map.Map b [a]
converseRelation m = Map.fromListWith (++) . fmap (second pure . swap) . Map.assocs $ m

{-
  Hacked together, you can view the result with:

    tred relation.dot | dot -Tpng -G"rankdir=BT" > relation.png

  tred is the graphviz tool to remove transitive edges. And the rankdir
  attribute flips the graph upside down.
-}
main :: IO ()
main = do
  -- Read dot file
  [dotFile] <- getArgs
  print dotFile
  transitions <- mapMaybe (parseMaybe parseTransFull) . lines <$> readFile dotFile

  -- convert to mealy
  let machine = convertToMealy transitions

  -- print some basic info
  putStrLn $ (show . length $ states machine) <> " states, " <> (show . length $ inputs machine) <> " inputs and " <> (show . length $ outputs machine) <> " outputs"
  putStrLn "Small sample:"
  print . take 4 . states $ machine
  print . take 4 . inputs $ machine
  print . take 4 . outputs $ machine

  -- -- DEBUG OUTPUT
  -- forM_ (states machine) (\s -> do
  --   print s
  --   forM_ (inputs machine) (\i -> do
  --     putStr "  "
  --     let (o, t) = behaviour machine s i
  --     putStrLn $ "--" <> (show i) <> "/" <> (show o) <> "->" <> (show t)
  --     )
  --   )

  let printPartition p = putStrLn $ "number of states = " <> show (numBlocks p)

  -- Minimise input, so we know the actual number of states
  printPartition (refineMealy (mealyMachineToEncoding machine))
  putStrLn ""

  -- Then compute each projection
  -- I did some manual preprocessing, these are the only interesting bits
  let -- outs = ["10", "10-O9", "2.2", "3.0", "3.1", "3.10", "3.12", "3.13", "3.14", "3.16", "3.17", "3.18", "3.19", "3.2", "3.20", "3.21", "3.3", "3.4", "3.6", "3.7", "3.8", "3.9", "5.0", "5.1", "5.12", "5.13", "5.17", "5.2", "5.21", "5.23", "5.6", "5.7", "5.8", "5.9", "quiescence"]
      outs = outputs machine
      (projections0, state2idx) = allProjections machine outs
      projections = zip outs $ fmap refineMealy projections0

  -- Print number of states of each projection
  forM_ projections (\(o, partition) -> do
    putStr $ o <> " -> "
    printPartition partition
    )

  -- First we check for equivalent partitions, so that we skip redundant work.
  let preord p1 p2 = toPreorder (comparePartitions p1 p2)
      (equiv, uniqPartitions) = equivalenceClasses preord projections

  putStrLn ""
  putStrLn "Representatives"
  print . fmap fst $ uniqPartitions

  putStrLn ""
  putStrLn "Equivalences"
  forM_ (Map.assocs equiv) (\(o2, o1) -> do
    putStrLn $ "  " <> (show o2) <> " == " <> (show o1)
    )

  -- Then we compare each pair of partitions. We only keep the finest
  -- partitions, since the coarse ones don't provide value to us.
  let (topMods, downSets) = maximalElements preord uniqPartitions
      foo (a, b) = (numBlocks b, a)

  putStrLn ""
  putStrLn "Top modules"
  forM_ (reverse . sort . fmap foo $ topMods) (\(b, o) -> do
    putStrLn $ "  " <> (show o) <> " has size " <> (show b)
    )

  -- Then we try to combine paritions, so that we don't end up with
  -- too many components. (Which would be too big to be useful.)
  let strategy MergerStats{..}
        | numberOfComponents <= 4 = Stop
        | otherwise = Continue

  projmap <- heuristicMerger topMods strategy

  -- Now we are going to output the components we found.
  let equivInv = converseRelation equiv
      projmapN = zip projmap [1 :: Int ..]

  forM_ projmapN (\((os, p), i) -> do
    let name = intercalate "x" os
        filename = "component" <> show i <> ".dot"

        osWithRel = concat $ os:[Map.findWithDefault [] o downSets | o <- os]
        osWithRelAndEquiv = concat $ osWithRel:[Map.findWithDefault [] o equivInv | o <- osWithRel]
        componentOutputs = Set.fromList osWithRelAndEquiv
        proj = projectToComponent (flip Set.member componentOutputs) machine
        -- Sanity check: compute partition again
        partition = refineMealy . mealyMachineToEncoding $ proj

    putStrLn $ ""
    putStrLn $ "Component " <> show os
    putStrLn $ "Correct? " <> show (comparePartitions p partition)
    putStrLn $ "Size = " <> show (numBlocks p)
    putStrLn $ "Output in file " <> filename

    let MealyMachine{..} = proj
        -- We enumerate all transitions in the full automaton
        transitions = [(s, i, o, t) | s <- states, i <- inputs, let (o, t) = behaviour s i]
        -- This is the quotient map, from state to block
        state2block = blockOfState p . (state2idx Map.!)
        -- We apply this to each transition, and then nubSort the duplicates away
        transitionsBlocks = nubSort [(state2block s, i, o, state2block t) | (s, i, o, t) <- transitions]
        -- The initial state should be first
        initialBlock = state2block initialState
        -- Sorting on "/= initialBlock" puts the initialBlock in front
        initialFirst = sortOn (\(s,_,_,_) -> s /= initialBlock) transitionsBlocks

        -- So far so good, `initialFirst` could serve as our output
        -- But we do one more optimisation on the machine
        -- We remove inputs, on which the machine does nothing
        deadInputs0 = Map.fromListWith (++) . fmap (\(s,i,o,t) -> (i, [(s,o,t)])) $ initialFirst
        deadInputs = Map.keysSet . Map.filter (all (\(s,o,t) -> s == t && o == Nothing)) $ deadInputs0
        result = filter (\(_,i,_,_) -> i `Set.notMember` deadInputs) initialFirst

        -- Convert to a file
        content = toString . mealyToDot name $ result

    putStrLn $ "Dead inputs = " <> show (Set.size deadInputs)

    writeFile filename content
    )

  return ()