diff --git a/LICENSE b/LICENSE
index 45644ff..dba13ed 100644
--- a/LICENSE
+++ b/LICENSE
@@ -1,23 +1,21 @@
-              GNU GENERAL PUBLIC LICENSE
-                Version 3, 29 June 2007
+                    GNU AFFERO GENERAL PUBLIC LICENSE
+                       Version 3, 19 November 2007
 
  Copyright (C) 2007 Free Software Foundation, Inc. <http://fsf.org/>
  Everyone is permitted to copy and distribute verbatim copies
  of this license document, but changing it is not allowed.
 
-                     Preamble
+                            Preamble
 
-  The GNU General Public License is a free, copyleft license for
-software and other kinds of works.
+  The GNU Affero General Public License is a free, copyleft license for
+software and other kinds of works, specifically designed to ensure
+cooperation with the community in the case of network server software.
 
   The licenses for most software and other practical works are designed
 to take away your freedom to share and change the works.  By contrast,
-the GNU General Public License is intended to guarantee your freedom to
+our General Public Licenses are intended to guarantee your freedom to
 share and change all versions of a program--to make sure it remains free
-software for all its users.  We, the Free Software Foundation, use the
-GNU General Public License for most of our software; it applies also to
-any other work released this way by its authors.  You can apply it to
-your programs, too.
+software for all its users.
 
   When we speak of free software, we are referring to freedom, not
 price.  Our General Public Licenses are designed to make sure that you
@@ -26,53 +24,43 @@ them if you wish), that you receive source code or can get it if you
 want it, that you can change the software or use pieces of it in new
 free programs, and that you know you can do these things.
 
-  To protect your rights, we need to prevent others from denying you
-these rights or asking you to surrender the rights.  Therefore, you have
-certain responsibilities if you distribute copies of the software, or if
-you modify it: responsibilities to respect the freedom of others.
+  Developers that use our General Public Licenses protect your rights
+with two steps: (1) assert copyright on the software, and (2) offer
+you this License which gives you legal permission to copy, distribute
+and/or modify the software.
 
-  For example, if you distribute copies of such a program, whether
-gratis or for a fee, you must pass on to the recipients the same
-freedoms that you received.  You must make sure that they, too, receive
-or can get the source code.  And you must show them these terms so they
-know their rights.
+  A secondary benefit of defending all users' freedom is that
+improvements made in alternate versions of the program, if they
+receive widespread use, become available for other developers to
+incorporate.  Many developers of free software are heartened and
+encouraged by the resulting cooperation.  However, in the case of
+software used on network servers, this result may fail to come about.
+The GNU General Public License permits making a modified version and
+letting the public access it on a server without ever releasing its
+source code to the public.
 
-  Developers that use the GNU GPL protect your rights with two steps:
-(1) assert copyright on the software, and (2) offer you this License
-giving you legal permission to copy, distribute and/or modify it.
+  The GNU Affero General Public License is designed specifically to
+ensure that, in such cases, the modified source code becomes available
+to the community.  It requires the operator of a network server to
+provide the source code of the modified version running there to the
+users of that server.  Therefore, public use of a modified version, on
+a publicly accessible server, gives the public access to the source
+code of the modified version.
 
-  For the developers' and authors' protection, the GPL clearly explains
-that there is no warranty for this free software.  For both users' and
-authors' sake, the GPL requires that modified versions be marked as
-changed, so that their problems will not be attributed erroneously to
-authors of previous versions.
-
-  Some devices are designed to deny users access to install or run
-modified versions of the software inside them, although the manufacturer
-can do so.  This is fundamentally incompatible with the aim of
-protecting users' freedom to change the software.  The systematic
-pattern of such abuse occurs in the area of products for individuals to
-use, which is precisely where it is most unacceptable.  Therefore, we
-have designed this version of the GPL to prohibit the practice for those
-products.  If such problems arise substantially in other domains, we
-stand ready to extend this provision to those domains in future versions
-of the GPL, as needed to protect the freedom of users.
-
-  Finally, every program is threatened constantly by software patents.
-States should not allow patents to restrict development and use of
-software on general-purpose computers, but in those that do, we wish to
-avoid the special danger that patents applied to a free program could
-make it effectively proprietary.  To prevent this, the GPL assures that
-patents cannot be used to render the program non-free.
+  An older license, called the Affero General Public License and
+published by Affero, was designed to accomplish similar goals.  This is
+a different license, not a version of the Affero GPL, but Affero has
+released a new version of the Affero GPL which permits relicensing under
+this license.
 
   The precise terms and conditions for copying, distribution and
 modification follow.
 
-                TERMS AND CONDITIONS
+                       TERMS AND CONDITIONS
 
   0. Definitions.
 
-  "This License" refers to version 3 of the GNU General Public License.
+  "This License" refers to version 3 of the GNU Affero General Public License.
 
   "Copyright" also means copyright-like laws that apply to other kinds of
 works, such as semiconductor masks.
@@ -549,35 +537,45 @@ to collect a royalty for further conveying from those to whom you convey
 the Program, the only way you could satisfy both those terms and this
 License would be to refrain entirely from conveying the Program.
 
-  13. Use with the GNU Affero General Public License.
+  13. Remote Network Interaction; Use with the GNU General Public License.
+
+  Notwithstanding any other provision of this License, if you modify the
+Program, your modified version must prominently offer all users
+interacting with it remotely through a computer network (if your version
+supports such interaction) an opportunity to receive the Corresponding
+Source of your version by providing access to the Corresponding Source
+from a network server at no charge, through some standard or customary
+means of facilitating copying of software.  This Corresponding Source
+shall include the Corresponding Source for any work covered by version 3
+of the GNU General Public License that is incorporated pursuant to the
+following paragraph.
 
   Notwithstanding any other provision of this License, you have
 permission to link or combine any covered work with a work licensed
-under version 3 of the GNU Affero General Public License into a single
+under version 3 of the GNU General Public License into a single
 combined work, and to convey the resulting work.  The terms of this
 License will continue to apply to the part which is the covered work,
-but the special requirements of the GNU Affero General Public License,
-section 13, concerning interaction through a network will apply to the
-combination as such.
+but the work with which it is combined will remain governed by version
+3 of the GNU General Public License.
 
   14. Revised Versions of this License.
 
   The Free Software Foundation may publish revised and/or new versions of
-the GNU General Public License from time to time.  Such new versions will
-be similar in spirit to the present version, but may differ in detail to
+the GNU Affero General Public License from time to time.  Such new versions
+will be similar in spirit to the present version, but may differ in detail to
 address new problems or concerns.
 
   Each version is given a distinguishing version number.  If the
-Program specifies that a certain numbered version of the GNU General
+Program specifies that a certain numbered version of the GNU Affero General
 Public License "or any later version" applies to it, you have the
 option of following the terms and conditions either of that numbered
 version or of any later version published by the Free Software
 Foundation.  If the Program does not specify a version number of the
-GNU General Public License, you may choose any version ever published
+GNU Affero General Public License, you may choose any version ever published
 by the Free Software Foundation.
 
   If the Program specifies that a proxy can decide which future
-versions of the GNU General Public License can be used, that proxy's
+versions of the GNU Affero General Public License can be used, that proxy's
 public statement of acceptance of a version permanently authorizes you
 to choose that version for the Program.
 
@@ -618,9 +616,9 @@ an absolute waiver of all civil liability in connection with the
 Program, unless a warranty or assumption of liability accompanies a
 copy of the Program in return for a fee.
 
-              END OF TERMS AND CONDITIONS
+                     END OF TERMS AND CONDITIONS
 
-     How to Apply These Terms to Your New Programs
+            How to Apply These Terms to Your New Programs
 
   If you develop a new program, and you want it to be of the greatest
 possible use to the public, the best way to achieve this is to make it
@@ -635,40 +633,29 @@ the "copyright" line and a pointer to where the full notice is found.
     Copyright (C) <year>  <name of author>
 
     This program is free software: you can redistribute it and/or modify
-    it under the terms of the GNU General Public License as published by
+    it under the terms of the GNU Affero General Public License as published by
     the Free Software Foundation, either version 3 of the License, or
     (at your option) any later version.
 
     This program is distributed in the hope that it will be useful,
     but WITHOUT ANY WARRANTY; without even the implied warranty of
     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-    GNU General Public License for more details.
+    GNU Affero General Public License for more details.
 
-    You should have received a copy of the GNU General Public License
+    You should have received a copy of the GNU Affero General Public License
     along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 Also add information on how to contact you by electronic and paper mail.
 
-  If the program does terminal interaction, make it output a short
-notice like this when it starts in an interactive mode:
-
-    <program>  Copyright (C) <year>  <name of author>
-    This program comes with ABSOLUTELY NO WARRANTY; for details type `show w'.
-    This is free software, and you are welcome to redistribute it
-    under certain conditions; type `show c' for details.
-
-The hypothetical commands `show w' and `show c' should show the appropriate
-parts of the General Public License.  Of course, your program's commands
-might be different; for a GUI interface, you would use an "about box".
+  If your software can interact with users remotely through a computer
+network, you should also make sure that it provides a way for users to
+get its source.  For example, if your program is a web application, its
+interface could display a "Source" link that leads users to an archive
+of the code.  There are many ways you could offer source, and different
+solutions will be better for different programs; see section 13 for the
+specific requirements.
 
   You should also get your employer (if you work as a programmer) or school,
 if any, to sign a "copyright disclaimer" for the program, if necessary.
-For more information on this, and how to apply and follow the GNU GPL, see
+For more information on this, and how to apply and follow the GNU AGPL, see
 <http://www.gnu.org/licenses/>.
-
-  The GNU General Public License does not permit incorporating your program
-into proprietary programs.  If your program is a subroutine library, you
-may consider it more useful to permit linking proprietary applications with
-the library.  If this is what you want to do, use the GNU Lesser General
-Public License instead of this License.  But first, please read
-<http://www.gnu.org/philosophy/why-not-lgpl.html>.
diff --git a/app/LStarMain.hs b/app/LStarMain.hs
index 120f3c9..cbed3ff 100644
--- a/app/LStarMain.hs
+++ b/app/LStarMain.hs
@@ -1,7 +1,7 @@
 module Main where
 
 import Bisimulation (bisimulation2)
-import DotParser
+import DotParser (readDotFile)
 import LStar
 import Mealy
 
@@ -9,9 +9,7 @@ import Control.Monad (when)
 import Control.Monad.Trans.Class
 import Control.Monad.Trans.State.Strict
 import Data.Map.Strict qualified as Map
-import Data.Maybe (mapMaybe)
 import System.Environment
-import Text.Megaparsec
 
 debugOutput :: Bool
 debugOutput = False
@@ -25,10 +23,9 @@ main :: IO ()
 main = do
   [dotFile] <- getArgs
   print dotFile
-  transitions <- mapMaybe (parseMaybe parseTransFull) . lines <$> readFile dotFile
+  machine <- readDotFile dotFile
 
   let
-    machine = convertToMealy transitions
     alphabet = inputs machine
     tInit = initialState machine
     tOut s i = fst (behaviour machine s i)
diff --git a/app/Main.hs b/app/Main.hs
index 51f3020..97f09b1 100644
--- a/app/Main.hs
+++ b/app/Main.hs
@@ -1,86 +1,72 @@
+{-# LANGUAGE OverloadedStrings #-}
 {-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}
 
-{-# HLINT ignore "Avoid reverse" #-}
 module Main where
 
-import DotParser
+import Data.Partition
+import Data.Preorder
+import DotParser (readDotFile)
 import DotWriter
 import Mealy
 import MealyRefine
 import Merger
-import Data.Partition
-import Data.Preorder
 
-import Control.Monad (forM_)
 import Data.Bifunctor
-import Data.List (intercalate, sort, sortOn)
+import Data.List (sortOn)
 import Data.List.Ordered (nubSort)
 import Data.Map.Strict qualified as Map
-import Data.Maybe (isNothing, mapMaybe)
+import Data.Maybe (isNothing)
 import Data.Set qualified as Set
+import Data.Text qualified as T
+import Data.Text.IO qualified as T
 import Data.Tuple (swap)
 import System.Environment
-import Text.Megaparsec
 
+-- | This functions inverts a map. In the new map the values are lists.
 converseRelation :: Ord b => Map.Map a b -> Map.Map b [a]
 converseRelation = Map.fromListWith (++) . fmap (second pure . swap) . Map.assocs
 
+-- TODO: use Data.Text here
 myWriteFile :: FilePath -> String -> IO ()
 myWriteFile filename = writeFile ("results/" ++ filename)
 
-{-
-  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
-  machine <- convertToMealy . mapMaybe (parseMaybe parseTransFull) . lines <$> readFile dotFile
+  machine <- readDotFile dotFile
 
   -- 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)
 
-  let printPartition p = putStrLn $ "number of states = " <> show (numBlocks p)
+  let
+    outputFuns = [(i, fun) | i <- inputs machine, let fun s = fst (behaviour machine s i)]
+    reverseTransitionMaps i = Map.fromListWith (++) [(t, [s]) | s <- states machine, let t = snd (behaviour machine s i)]
+    reverseFuns = [(i, fun) | i <- inputs machine, let mm = reverseTransitionMaps i, let fun s = Map.findWithDefault [] s mm]
 
   -- Minimise input, so we know the actual number of states
-  printPartition (refineMealy (mealyMachineToEncoding machine))
+  printPartition (refineMealy3 outputFuns reverseFuns (states 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
+    mappedOutputFuns o = [(i, (o ==) . f) | (i, f) <- outputFuns]
+    projections = [(o, refineMealy3 (mappedOutputFuns o) reverseFuns (states machine)) | o <- outs]
 
   -- Print number of states of each projection
-  forM_
-    projections
+  mapM_
     ( \(o, partition) -> do
-        putStr $ o <> " -> "
+        T.putStr o
+        putStr " -> "
         printPartition partition
     )
+    projections
 
   -- First we check for equivalent partitions, so that we skip redundant work.
   let
@@ -92,11 +78,11 @@ main = do
 
   putStrLn ""
   putStrLn "Equivalences"
-  forM_
-    (Map.assocs equiv)
+  mapM_
     ( \(o2, o1) -> do
         putStrLn $ "  " <> show o2 <> " == " <> show o1
     )
+    (Map.assocs equiv)
 
   -- Then we compare each pair of partitions. We only keep the finest
   -- partitions, since the coarse ones don't provide value to us.
@@ -106,11 +92,11 @@ main = do
 
   putStrLn ""
   putStrLn "Top modules"
-  forM_
-    (reverse . sort . fmap foo $ topMods)
+  mapM_
     ( \(b, o) -> do
         putStrLn $ "  " <> show o <> " has size " <> show b
     )
+    (sortOn (negate . fst) . fmap foo $ topMods)
 
   -- 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.)
@@ -120,19 +106,21 @@ main = do
 
   projmap <- heuristicMerger topMods strategy
 
+  print projmap
+
   -- Now we are going to output the components we found.
   let
     equivInv = converseRelation equiv
     projmapN = zip projmap [1 :: Int ..]
     action ((os, p), componentIdx) = do
       let
-        name = intercalate "x" os
+        name = T.intercalate "x" os
         osWithRel = concat $ os : [Map.findWithDefault [] o downSets | o <- os]
         osWithRelAndEquiv = concat $ osWithRel : [Map.findWithDefault [] o equivInv | o <- osWithRel]
         componentOutputs = Set.fromList osWithRelAndEquiv
         proj = projectToComponent (`Set.member` componentOutputs) machine
         -- Sanity check: compute partition again
-        partition = refineMealy . mealyMachineToEncoding $ proj
+        partition = refineMealy2 proj
 
       putStrLn ""
       putStrLn $ "Component " <> show os
@@ -142,9 +130,7 @@ main = do
       do
         let
           filename = "partition_" <> show componentIdx <> ".dot"
-          idx2State = Map.map head . converseRelation $ state2idx
-          stateBlocks = fmap (fmap (idx2State Map.!)) . toBlocks $ partition
-          content = unlines . fmap unwords $ stateBlocks
+          content = T.unpack . T.unlines . fmap T.unwords . toBlocks $ p
 
         putStrLn $ "Output (partition) in file " <> filename
         myWriteFile filename content
@@ -155,7 +141,7 @@ main = do
           -- 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.!)
+          state2block = (Map.!) (getPartition p)
           -- 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
@@ -165,7 +151,7 @@ main = do
 
           -- Convert to a file
           filename1 = "component_" <> show componentIdx <> ".dot"
-          content1 = toString . mealyToDot name $ initialFirst
+          content1 = toString . mealyToDot (T.unpack name) $ initialFirst
 
           -- So far so good, `initialFirst` could serve as our output
           -- But we do one more optimisation on the machine
@@ -176,7 +162,7 @@ main = do
 
           -- Convert to a file
           filename2 = "component_reduced_" <> show componentIdx <> ".dot"
-          content2 = toString . mealyToDot name $ result
+          content2 = toString . mealyToDot (T.unpack name) $ result
 
         putStrLn $ "Output (reduced machine) in file " <> filename1
         myWriteFile filename1 content1
diff --git a/app/Playground.hs b/app/Playground.hs
index a951428..50466ac 100644
--- a/app/Playground.hs
+++ b/app/Playground.hs
@@ -1,22 +1,21 @@
 module Main where
 
 import Bisimulation (bisimulation2)
-import Data.UnionFind
-import DotParser (convertToMealy, parseTransFull)
-import Mealy (MealyMachine (..), outputFunction, transitionFunction)
 import Data.Partition (numBlocks)
-import SplittingTree (PRState (..), getPartition, initialPRState, refine)
-import StateIdentifiers (stateIdentifierFor)
 import Data.Trie qualified as Trie
+import Data.UnionFind
+import DotParser (readDotFile)
+import Mealy (MealyMachine (..), outputFunction, transitionFunction)
+import SplittingTree (initialPRState, refine)
+import StateIdentifiers (stateIdentifierFor)
 
-import Control.Monad.Trans.State (execStateT)
+import Control.Monad.Trans.State (evalStateT)
 import Data.List qualified as List
 import Data.Map.Strict qualified as Map
-import Data.Maybe (isJust, mapMaybe)
+import Data.Maybe (isJust)
 import Data.Set qualified as Set
-import MealyRefine
+import Data.Text.IO qualified as T
 import System.Environment (getArgs)
-import Text.Megaparsec (parseMaybe)
 
 main :: IO ()
 main = do
@@ -34,16 +33,16 @@ mainHSI args = case args of
  where
   run dotFile = do
     print dotFile
-    transitions <- mapMaybe (parseMaybe parseTransFull) . lines <$> readFile dotFile
+    machine <- readDotFile dotFile
 
     -- convert to mealy
     let
-      MealyMachine{..} = convertToMealy transitions
+      MealyMachine{..} = machine
       outputFuns = [(i, fun) | i <- inputs, let fun s = fst (behaviour s i)]
       reverseTransitionMaps i = Map.fromListWith (++) [(t, [s]) | s <- states, let t = snd (behaviour s i)]
       reverseFuns = [(i, fun) | i <- inputs, let m = reverseTransitionMaps i, let fun s = Map.findWithDefault [] s m]
 
-    PRState{..} <- execStateT (refine print outputFuns reverseFuns) (initialPRState states)
+    (partition, splittingTree) <- evalStateT (refine print outputFuns reverseFuns) (initialPRState states)
 
     putStrLn "\nPARTITION"
     print partition
@@ -85,22 +84,22 @@ mainInputDecomp args = case args of
  where
   run dotFile = do
     print dotFile
-    transitions <- mapMaybe (parseMaybe parseTransFull) . lines <$> readFile dotFile
+    model <- readDotFile dotFile
 
-    let model = convertToMealy transitions
-        composition i j = interleavingComposition [i] [j] model
-        bisim i j =
-          let compo = composition i j
-           in bisimulation2
-                [i, j]
-                (outputFunction model)
-                (transitionFunction model)
-                (initialState model)
-                (outputFunction compo)
-                (transitionFunction compo)
-                (initialState compo)
-        dependent i j = isJust $ bisim i j
-        dependentPairs = [(i, j) | i <- inputs model, j <- inputs model, j > i, dependent i j]
+    let
+      composition i j = interleavingComposition [i] [j] model
+      bisim i j =
+        let compo = composition i j
+         in bisimulation2
+              [i, j]
+              (outputFunction model)
+              (transitionFunction model)
+              (initialState model)
+              (outputFunction compo)
+              (transitionFunction compo)
+              (initialState compo)
+      dependent i j = isJust $ bisim i j
+      dependentPairs = [(i, j) | i <- inputs model, j <- inputs model, j > i, dependent i j]
 
     print $ length (states model)
     print $ length (inputs model)
@@ -129,21 +128,24 @@ mainInputDecomp args = case args of
       n -> putStrLn ("MAYBE DECOMPOSABLE: " ++ show n ++ " classes")
 
 -- Used to determine whether Copar is faster than SplittingTree (it is).
+-- Copar is almost twice as fast on ESM, but SplittingTree is faster on a
+-- BRP benchmark. I guess, theoretically, Copar should be faster generally.
 mainRefine :: [String] -> IO ()
 mainRefine args = case args of
   [dotFile, copar] -> run dotFile (read copar)
   _ -> putStrLn "Please provide a dot file and Boolean"
  where
   run dotFile copar = do
-    m <- convertToMealy . mapMaybe (parseMaybe parseTransFull) . lines <$> readFile dotFile
-    putStrLn $ "file parsed, initial state = " <> initialState m
+    m <- readDotFile dotFile
+    putStr $ "file parsed, initial state = "
+    T.putStrLn $ initialState m
     if copar
       then runCopar m
       else runSplittingTree m
 
-  runCopar m =
-    let printPartition p = putStrLn $ "Done " <> show (numBlocks p)
-     in printPartition (refineMealy (mealyMachineToEncoding m))
+  runCopar _ = error "no longer supported"
+  --  let printPartition p = putStrLn $ "Done " <> show (numBlocks p)
+  --   in printPartition (refineMealy (mealyMachineToEncoding m))
 
   runSplittingTree MealyMachine{..} = do
     let
@@ -151,5 +153,5 @@ mainRefine args = case args of
       reverseTransitionMaps i = Map.fromListWith (++) [(t, [s]) | s <- states, let t = snd (behaviour s i)]
       reverseFuns = [(i, fun) | i <- inputs, let mm = reverseTransitionMaps i, let fun s = Map.findWithDefault [] s mm]
 
-    PRState{..} <- execStateT (refine (\_ -> pure ()) outputFuns reverseFuns) (initialPRState states)
-    putStrLn $ "Done" <> show (Map.size (getPartition partition))
+    (partition, _splittingTree) <- evalStateT (refine (\_ -> pure ()) outputFuns reverseFuns) (initialPRState states)
+    putStrLn $ "Done" <> show (numBlocks partition)
diff --git a/app/RandomGen.hs b/app/RandomGen.hs
index 9471a25..d27313d 100644
--- a/app/RandomGen.hs
+++ b/app/RandomGen.hs
@@ -3,11 +3,11 @@
 
 module Main where
 
+import Data.Partition (Block (..))
 import SplittingTree
 
 import Control.Monad.Random.Strict
 import Control.Monad.Trans.State (execStateT)
-import Data.Coerce
 import Data.List (sortOn)
 import Data.Map qualified as Map
 import Data.Set qualified as Set
@@ -91,11 +91,11 @@ main = do
 
   -- print
   let
-    toBlock s = getPartition partition Map.! s
+    toBlock s = getBarePartition partition Map.! s
     allTransitions = [(toBlock s, i, o, toBlock t) | s <- states, i <- inputs, let o = outpf s i, let t = trans s i]
     uniqueTransitions = sortOn (\(s, _, _, _) -> s /= toBlock init) . Set.toList . Set.fromList $ allTransitions
     showLabel i o = "[label=\"" <> [i] <> "/" <> [o] <> "\"]"
-    showTransition s i o t = "s" <> show (coerce s :: Int) <> " -> " <> "s" <> show (coerce t :: Int) <> " " <> showLabel i o
+    showTransition (Block s) i o (Block t) = "s" <> show s <> " -> " <> "s" <> show t <> " " <> showLabel i o
 
   putStrLn "digraph g {"
   mapM_ (\(s, i, o, t) -> putStrLn ("  " <> showTransition s i o t)) uniqueTransitions
diff --git a/mealy-decompose.cabal b/mealy-decompose.cabal
index c987e8c..8fe0c3e 100644
--- a/mealy-decompose.cabal
+++ b/mealy-decompose.cabal
@@ -1,23 +1,24 @@
 cabal-version:      2.2
 name:               mealy-decompose
-version:            0.1.0.0
-license:            GPL-3.0-only
+version:            0.2.0.0
+license:            AGPL-3.0-only
 license-file:       LICENSE
 author:             Joshua Moerman
 maintainer:         joshua.moerman@ou.nl
-copyright:          Joshua Moerman (C) 2023
+copyright:          (c) 2024 Joshua Moerman, Open Universiteit
 build-type:         Simple
 
 common stuff
   build-depends:
-    base ^>=4.19.0.0,
+    base >= 4.15,
     containers,
+    unordered-containers,
     data-ordlist,
     megaparsec,
+    text,
     transformers
   default-language: GHC2021
-  default-extensions:
-    RecordWildCards
+  default-extensions: RecordWildCards
 
 library
   import: stuff
@@ -36,9 +37,6 @@ library
     Merger,
     SplittingTree,
     StateIdentifiers
-  build-depends:
-    copar,
-    vector
 
 executable mealy-decompose
   import: stuff
diff --git a/src/Data/Partition.hs b/src/Data/Partition.hs
index 3b5bb00..32e346d 100644
--- a/src/Data/Partition.hs
+++ b/src/Data/Partition.hs
@@ -1,55 +1,59 @@
 {-# LANGUAGE PackageImports #-}
 
-module Data.Partition (
-  -- $partitions
-  module Data.Partition,
-) where
+module Data.Partition where
 
 import Data.Preorder
 
-import Control.Monad.Trans.State.Strict (get, put, runState)
-import Data.Coerce (coerce)
+import Control.Monad.Trans.State.Strict as State
+import Data.Map.Merge.Strict qualified as Map
 import Data.Map.Strict qualified as Map
-import Data.Partition.Common (Block (..))
-import Data.Vector qualified as V
-import "copar" Data.Partition (Partition (..), blockOfState, numStates, toBlocks)
+import Data.Tuple (swap)
+import Data.List (groupBy, sortOn)
+import Data.Function (on)
 
--- $partitions
---
--- This module re-exports the `Data.Partition` module from the `copar` library,
--- and adds some additional functions for working with partitions. A partition
--- on a set of type @a@ is represented as a map @a -> `Block`@, where a `Block`
--- is a unique identifier (integer) for a set of elements.
---
--- In this module, we define
---
--- * `commonRefinement` to compute the common refinement of two partitions.
--- * `isRefinementOf2` to check if one partition is a refinement of another.
--- * `isEquivalent` to check if two partitions are equal.
--- * `comparePartitions` to compare two partitions in the partition lattice.
---
--- Partitions form a lattice (the so-called /partition lattice/), where the
--- partial order is given by the refinement relation. We put the finest
--- partition at the top, and the coarsest at the bottom. (This is the opposite
--- of the convection used on wikipedia.)
+newtype Block = Block Int
+  deriving (Eq, Ord, Show, Enum, Num)
+
+-- | A partition is represented by a finite map of type 's -> Block'. Two
+-- elements mapped to the same block are equivalent. Note that a permutation
+-- of the blocks will not change the partition, but it does change the
+-- underlying representation. (That is why I haven't given it an Eq instance
+-- yet.)
+data Partition s = Partition
+  { getPartition :: Map.Map s Block
+  , numBlocks :: Block
+  }
+  deriving Show
+
+-- | Constraint for using partitions. Currently this is 'Ord'. But it might
+-- change to '(Eq s, Hashable s)' in the future.
+type Partitionable s = Ord s
+
+-- | Determines whether two elements are equivalent in the partition.
+sameBlock :: Partitionable s => Partition s -> s -> s -> Bool
+sameBlock (Partition m _) s t = m Map.! s == m Map.! t
+
+-- | Returns a list of all the blocks of the partition.
+toBlocks :: Partition s -> [[s]]
+toBlocks = fmap (fmap snd) . groupBy ((==) `on` fst) . sortOn fst . fmap swap . Map.assocs . getPartition
 
 -- | Returns the common refinement of two partitions. This is the coarsest
 -- partition which is finer than either input, i.e., the lowest upper bound.
-commonRefinement :: Partition -> Partition -> Partition
-commonRefinement p1 p2 =
-  let n = numStates p1
-      sa1 = (stateAssignment p1 V.!)
-      sa2 = (stateAssignment p2 V.!)
-      blockAtIdx i = do
-        (m, b) <- get
-        let key = (sa1 i, sa2 i)
-        case Map.lookup key m of
-          Just b0 -> return b0
-          Nothing -> do
-            put (Map.insert key b m, succ b)
-            return b
-      (vect, (_, nextBlock)) = runState (V.generateM n blockAtIdx) (Map.empty, 0)
-   in Partition{numBlocks = coerce nextBlock, stateAssignment = vect}
+-- Runs in O(n), where n is the number of elements of the partition. Note
+-- that both partitions must be on the same set of elements.
+commonRefinement :: Partitionable s => Partition s -> Partition s -> Partition s
+commonRefinement (Partition p1 _) (Partition p2 _) =
+  let (p, (_, b)) = State.runState (mergeFun p1 p2) (Map.empty, Block 0)
+   in Partition p b
+ where
+  mergeFun = Map.mergeA Map.dropMissing Map.dropMissing (Map.zipWithAMatched checkPair)
+  checkPair _ b1 b2 = do
+    (m, n) <- get
+    case Map.lookup (b1, b2) m of
+      Just k -> pure k
+      Nothing -> do
+        put (Map.insert (b1, b2) n m, succ n)
+        pure n
 
 -- | This function checks whether one partition is a refinement of the other.
 -- This function already appears in the `copar` library, but the one here is
@@ -57,27 +61,26 @@ commonRefinement p1 p2 =
 
 -- Could be made faster by doing what commonRefinement is doing but
 -- stopping early. But it's fast enough for now, so I won't bother.
-isRefinementOf2 :: Partition -> Partition -> Bool
+isRefinementOf2 :: Partitionable s => Partition s -> Partition s -> Bool
 isRefinementOf2 refined original = comparePartitions refined original == GT'
 
 -- | Checks whether two partitions are equal as partitions. Note that the `Eq`
 -- instance on partitions checks for equality of the state assignments, not
 -- whether the partitions are equal as partitions.
-isEquivalent :: Partition -> Partition -> Bool
+isEquivalent :: Partitionable s => Partition s -> Partition s -> Bool
 isEquivalent p1 p2 = comparePartitions p1 p2 == EQ'
 
 -- | Compares two partitions. Returns `EQ'` if the partitions are equal, `GT'`
 -- if the first partition is a refinement of the second, `LT'` if the first
 -- partition is a coarsening of the second, and `IC'` if the partitions are
 -- incomparable.
-comparePartitions :: Partition -> Partition -> PartialOrdering
-comparePartitions p1 p2
-  | p1 == p2 = EQ'
+comparePartitions :: Partitionable s => Partition s -> Partition s -> PartialOrdering
+comparePartitions p1@(Partition m1 b1) p2@(Partition m2 b2)
+  | m1 == m2 = EQ'
   | otherwise =
-      let glb = commonRefinement p1 p2
-          n1 = numBlocks p1
-          n2 = numBlocks p2
-          n3 = numBlocks glb
+      let (Partition _ n3) = commonRefinement p1 p2
+          n1 = b1
+          n2 = b2
        in case (n1 == n3, n2 == n3) of
             (True, True) -> EQ'
             (True, False) -> GT'
diff --git a/src/DotParser.hs b/src/DotParser.hs
index f7f28a3..9bf0820 100644
--- a/src/DotParser.hs
+++ b/src/DotParser.hs
@@ -1,43 +1,65 @@
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE PartialTypeSignatures #-}
+{-# OPTIONS_GHC -Wno-partial-type-signatures #-}
+
 module DotParser where
 
-import Data.Char (isAlphaNum)
-import Data.List.Ordered qualified as OrdList
-import Data.Map.Strict qualified as Map
-import Data.Void (Void)
 import Mealy
+
+import Data.Char (isAlphaNum)
+import Data.HashMap.Strict qualified as Map
+import Data.HashSet qualified as Set
+import Data.Maybe (mapMaybe)
+import Data.Text qualified as T
+import Data.Text.IO qualified as T
+import Data.Void (Void)
 import Text.Megaparsec
-import Text.Megaparsec.Char
+import Text.Megaparsec.Char (hspace, hspace1, space)
 import Text.Megaparsec.Char.Lexer qualified as L
 
-{-
-   Parser for Dot files generated by the RERS LearnLib learner. This is not
-   a fully fledged parser. It is specific to our models.
+-- * Main function(s)
 
-   Really only parses a single transition. We just collect all succesfull
-   transitions. This gives all transitions.
+-- $moduleDoc
+--
+-- Parser for Dot files generated by LearnLib (and other tools). This is not
+-- a fully fledged parser for dot files. It is specific to our encoding of
+-- Mealy machines. The parser works line-based, so transitions need to be
+-- on a single line. It assumes the first transitions belongs to the initial
+-- state. The parser ignores any line which is not a transition.
 
-   Usage:
-     transitions <- mapMaybe (parseMaybe parseTransFull) . lines <$> readFile dotFile
--}
+-- | Read a dot file and convert it to a Mealy machine.
+readDotFile :: FilePath -> IO (MealyMachine T.Text T.Text T.Text)
+readDotFile dotFile = convertToMealy . mapMaybe (parseMaybe lineP) . T.lines <$> T.readFile dotFile
+ where
+  -- Parses a full line, ignoring whitespace, until the end of the line
+  lineP = hspace *> parseTrans <* space <* eof
 
-type Stat = String
-type Input = String
-type Output = String
-type Trans = (Stat, Stat, Input, Output)
+-- * Internals
 
-type Parser = Parsec Void String
+-- | Type of tokens we accept is 'T.Text', but it could be any type which
+-- is compatible with Megaparsec.
+type Toks = T.Text
 
+-- | A transition is a tuple of state, successor state, input label, and
+-- output label
+type Trans = (Toks, Toks, Toks, Toks)
+
+-- | Our parser does not have any custom error messages, and always consumes
+-- a stream as 'T.Text'.
+type Parser = Parsec Void Toks
+
+-- | Parse a single transition.
 parseTrans :: Parser Trans
 parseTrans = assoc <$> identifierQ <* symbol "->" <*> identifierQ <*> brackets parseLabel <* optional (symbol ";")
  where
   -- defines whitespace and lexemes
-  sc = L.space space1 empty empty
+  sc = L.space hspace1 empty empty
   lexeme = L.lexeme sc
   symbol = L.symbol sc
-  -- state, input, output is any string of alphaNumChar's
-  isAlphaNumExtra c = isAlphaNum c || c == '_' || c == '+' || c == '.' || c == ',' || c == '-' || c == '(' || c == ')'
-  alphaNumCharExtra = satisfy isAlphaNumExtra <?> "alphanumeric character or extra"
-  identifier = lexeme (some alphaNumCharExtra)
+  -- state, input, output is any string of alphaNumChar's (and some additional characters)
+  isAlphaNumExtra c = isAlphaNum c || ('(' <= c && c <= '.') || c == '_'
+  alphaNumCharExtra = takeWhile1P (Just "alphanumeric character or extra") isAlphaNumExtra
+  identifier = lexeme alphaNumCharExtra
   identifierQ = identifier <|> between (symbol "\"") (symbol "\"") identifier
   -- The label has the shape [label="i/o"]
   brackets = between (symbol "[") (symbol "]")
@@ -45,24 +67,21 @@ parseTrans = assoc <$> identifierQ <* symbol "->" <*> identifierQ <*> brackets p
   -- re-associate different parts of the parser
   assoc from to (i, o) = (from, to, i, o)
 
-parseTransFull :: Parser Trans
-parseTransFull = space *> parseTrans <* eof
-
-convertToMealy :: [Trans] -> MealyMachine String String String
+-- | Convert a list of transitions to a Mealy machine, assumes the first
+-- transition belongs to the initial state.
+convertToMealy :: [Trans] -> MealyMachine T.Text T.Text T.Text
 convertToMealy l =
   MealyMachine
-    { states = states
-    , inputs = ins
-    , outputs = outs
+    { states = fmap (allStrs Map.!) . Set.toList $ states
+    , inputs = Set.toList $ ins
+    , outputs = Set.toList $ outs
     , behaviour = curry (base Map.!)
-    , initialState = (\(a, _, _, _) -> a) . head $ l
+    , initialState = (allStrs Map.!) . (\(a, _, _, _) -> a) . head $ l
     }
  where
-  -- \^ Assumption: first transition in the file belongs to the initial state
-
-  froms = OrdList.nubSort . fmap (\(a, _, _, _) -> a) $ l
-  tos = OrdList.nubSort . fmap (\(_, a, _, _) -> a) $ l
-  ins = OrdList.nubSort . fmap (\(_, _, i, _) -> i) $ l
-  outs = OrdList.nubSort . fmap (\(_, _, _, o) -> o) $ l
-  states = froms `OrdList.union` tos
-  base = Map.fromList . fmap (\(from, to, i, o) -> ((from, i), (o, to))) $ l
+  states = Set.fromList . concatMap (\(a, b, _, _) -> [a, b]) $ l
+  ins = Set.fromList . fmap (\(_, _, i, _) -> i) $ l
+  outs = Set.fromList . fmap (\(_, _, _, o) -> o) $ l
+  -- We put some effort in sharing string values
+  allStrs = Map.mapWithKey (\k _ -> k) . Set.toMap . Set.unions $ [states, ins, outs]
+  base = Map.fromList . fmap (\(from, to, i, o) -> ((allStrs Map.! from, allStrs Map.! i), (allStrs Map.! o, allStrs Map.! to))) $ l
diff --git a/src/DotWriter.hs b/src/DotWriter.hs
index 4d9140b..780f593 100644
--- a/src/DotWriter.hs
+++ b/src/DotWriter.hs
@@ -1,7 +1,10 @@
 module DotWriter where
 
-import Data.Monoid (Endo(..))
-import Data.Partition.Common (Block(..))
+import Data.Monoid (Endo (..))
+import Data.Partition (Block (..))
+import Data.Text qualified as T
+
+-- TODO: use Data.Text here instead of strings
 
 type StringBuilder = Endo String
 
@@ -17,6 +20,9 @@ class ToDot s where
 instance ToDot String where
   toDot = string
 
+instance ToDot T.Text where
+  toDot = string . T.unpack
+
 instance ToDot a => ToDot (Maybe a) where
   -- should be chosen not to conflict with possible outputs
   toDot Nothing = string "nil"
@@ -28,13 +34,21 @@ instance ToDot Block where
 
 transitionToDot :: (ToDot s, ToDot i, ToDot o) => (s, i, o, s) -> StringBuilder
 transitionToDot (s, i, o, t) =
-  toDot s <> string " -> " <> toDot t
-  <> string " [label=\"" <> toDot i <> string " / " <> toDot o <> string "\"]"
+  toDot s
+    <> string " -> "
+    <> toDot t
+    <> string " [label=\""
+    <> toDot i
+    <> string " / "
+    <> toDot o
+    <> string "\"]"
 
 mealyToDot :: (ToDot s, ToDot i, ToDot o) => String -> [(s, i, o, s)] -> StringBuilder
 mealyToDot name transitions =
-  string "digraph " <> string name <> string " {\n"
-  <> foldMap transitionToDotSep transitions
-  <> string "}\n"
-  where
-    transitionToDotSep t = string "  " <> transitionToDot t <> string "\n"
+  string "digraph "
+    <> string name
+    <> string " {\n"
+    <> foldMap transitionToDotSep transitions
+    <> string "}\n"
+ where
+  transitionToDotSep t = string "  " <> transitionToDot t <> string "\n"
diff --git a/src/MealyRefine.hs b/src/MealyRefine.hs
index 9c51f96..072ff59 100644
--- a/src/MealyRefine.hs
+++ b/src/MealyRefine.hs
@@ -1,19 +1,15 @@
+{-# LANGUAGE PackageImports #-}
+
 module MealyRefine where
 
-import Mealy
 import Data.Partition (Partition)
+import Mealy
+import SplittingTree qualified as ST
 
-import Control.Monad.ST (runST)
-import Copar.Algorithm (refine)
-import Copar.Functors.Polynomial (PolyF1 (..), Polynomial)
-import Copar.RefinementInterface (F1, Label)
-import Data.Bool (bool)
-import Data.CoalgebraEncoding (Encoding (..))
+import Control.Monad.Trans.State (evalState)
+import Data.Functor.Identity (Identity (..))
 import Data.List.Ordered (nubSort)
 import Data.Map.Strict qualified as Map
-import Data.Proxy (Proxy (..))
-import Data.Vector qualified
-import Data.Vector.Unboxed qualified
 
 project :: Ord u => (o -> u) -> MealyMachine s i o -> MealyMachine s i u
 project f MealyMachine{..} =
@@ -38,6 +34,8 @@ projectToComponent oPred = project oMaybe
 -- among those mealy machines, I do this in one function. The static parts
 -- are converted only once. Only "eStructure" (the state-labels) are different
 -- for each projection.
+
+{-
 allProjections :: (Ord s, Ord i, Eq o) => MealyMachine s i o -> [o] -> ([Encoding (Label Polynomial) (F1 Polynomial)], Map.Map s Int)
 allProjections MealyMachine{..} outs = (fmap mkEncoding outs, state2idx)
  where
@@ -74,8 +72,8 @@ allProjections MealyMachine{..} outs = (fmap mkEncoding outs, state2idx)
       }
 
 -- Refine a encoded mealy machine
-refineMealy :: Encoding (Label Polynomial) (F1 Polynomial) -> Partition
-refineMealy machine = runST $ refine (Proxy @Polynomial) machine True
+refineMealy :: Encoding (Label Polynomial) (F1 Polynomial) -> Copar.Partition
+refineMealy machine = runST $ Copar.refine (Proxy @Polynomial) machine True
 
 mealyMachineToEncoding :: (Ord s, Ord i, Ord o) => MealyMachine s i o -> Encoding (Label Polynomial) (F1 Polynomial)
 mealyMachineToEncoding MealyMachine{..} =
@@ -105,3 +103,18 @@ mealyMachineToEncoding MealyMachine{..} =
       eEdgesLabel = Data.Vector.generate (numStates * numInputs) (snd . stateInputIndex)
       eEdgesTo = Data.Vector.Unboxed.generate (numStates * numInputs) ((state2idx Map.!) . snd . (\(s, i) -> behaviour (idx2state Map.! s) (idx2input Map.! i)) . stateInputIndex)
    in Encoding{..}
+-}
+
+refineMealy2 :: (Ord s, Ord i, Ord o) => MealyMachine s i o -> Partition s
+refineMealy2 MealyMachine{..} =
+  let
+    outputFuns = [(i, fun) | i <- inputs, let fun s = fst (behaviour s i)]
+    reverseTransitionMaps i = Map.fromListWith (++) [(t, [s]) | s <- states, let t = snd (behaviour s i)]
+    reverseFuns = [(i, fun) | i <- inputs, let mm = reverseTransitionMaps i, let fun s = Map.findWithDefault [] s mm]
+   in
+    refineMealy3 outputFuns reverseFuns states
+
+refineMealy3 :: (Ord o, Ord s) => [(i, s -> o)] -> [(i, s -> [s])] -> [s] -> Partition s
+refineMealy3 outputFuns reverseFuns states =
+  let (partition, _splittingTree) = evalState (ST.refine (\_ -> Identity ()) outputFuns reverseFuns) (ST.initialPRState states)
+   in partition
diff --git a/src/Merger.hs b/src/Merger.hs
index 6024a0b..edb604c 100644
--- a/src/Merger.hs
+++ b/src/Merger.hs
@@ -13,17 +13,17 @@ import Data.Set qualified as Set
 
 data MergerStats = MergerStats
   { numberOfComponents :: Int
-  , maximalComponent :: Int
-  , totalSize :: Int
+  , maximalComponent :: Block
+  , totalSize :: Block
   }
-  deriving (Eq, Ord, Read, Show)
+  deriving (Eq, Ord, Show)
 
 data MergerAction = Stop | Continue
-  deriving (Eq, Ord, Read, Show)
+  deriving (Eq, Ord, Show)
 
 type MergerStrategy = MergerStats -> MergerAction
 
-heuristicMerger :: Ord o => [(o, Partition)] -> MergerStrategy -> IO [([o], Partition)]
+heuristicMerger :: (Ord o, Ord s) => [(o, Partition s)] -> MergerStrategy -> IO [([o], Partition s)]
 heuristicMerger components strategy = do
   projmap <- evalStateT (loop 2) (Map.fromList (fmap (first pure) components))
   return $ Map.assocs projmap
diff --git a/src/SplittingTree.hs b/src/SplittingTree.hs
index e0fa5f2..c2814fb 100644
--- a/src/SplittingTree.hs
+++ b/src/SplittingTree.hs
@@ -3,6 +3,8 @@
 
 module SplittingTree where
 
+import Data.Partition (Block (..), Partition (..))
+
 import Control.Monad.Trans.Class
 import Control.Monad.Trans.State
 import Data.Coerce (coerce)
@@ -10,19 +12,8 @@ import Data.Foldable (traverse_)
 import Data.List (sortOn)
 import Data.Map.Strict qualified as Map
 
-newtype Block = Block Int
-  deriving (Eq, Ord, Read, Show, Enum)
-
--- A partition is represented by a map s -> Block. Two elements mapped to the
--- same block are equivalent. Note that a permutation of the blocks will
--- not change the partition, but it does change the underlying representation.
--- (That is why I haven't given it an Eq instance yet.)
-newtype Partition s = Partition {getPartition :: Map.Map s Block}
-  deriving (Read, Show)
-
--- Determines whether two elements are equivalent in the partition.
-sameBlock :: Ord s => Partition s -> s -> s -> Bool
-sameBlock (Partition m) s t = m Map.! s == m Map.! t
+newtype BarePartition s = BarePartition {getBarePartition :: Map.Map s Block}
+  deriving Show
 
 -- In the splitting tree, we record the splits we have made during partition
 -- refinement. The leafs correspond to the blocks in the partition, and the
@@ -56,7 +47,7 @@ data Splitter s i o = Splitter
 
 -- The data structure used during partition refinement.
 data PRState s i o = PRState
-  { partition :: Partition s
+  { partition :: BarePartition s
   , nextBlockId :: Block
   , splittingTree :: SplittingTree s i o
   , nextNodeId :: InnerNode
@@ -98,7 +89,7 @@ genNextNodeId = do
 
 refineWithSplitter :: (Monad m, Ord o, Ord s) => i -> (s -> [s]) -> Splitter s i o -> StateT (PRState s i o) m [Splitter s i o]
 refineWithSplitter action rev Splitter{..} = do
-  currentPartition <- getPartition <$> gets partition
+  currentPartition <- getBarePartition <$> gets partition
   currentSplittingTree <- gets splittingTree
 
   let
@@ -190,7 +181,7 @@ refineWithSplitter action rev Splitter{..} = do
 
 refineWithOutput :: (Monad m, Ord o, Ord s) => i -> (s -> o) -> StateT (PRState s i o) m [Splitter s i o]
 refineWithOutput action out = do
-  currentPartition <- getPartition <$> gets partition
+  currentPartition <- getBarePartition <$> gets partition
   currentSplittingTree <- gets splittingTree
 
   let
@@ -247,7 +238,7 @@ refineWithOutput action out = do
 initialPRState :: Ord s => [s] -> PRState s i o
 initialPRState ls =
   PRState
-    { partition = Partition . Map.fromList $ [(s, Block 0) | s <- ls]
+    { partition = BarePartition . Map.fromList $ [(s, Block 0) | s <- ls]
     , nextBlockId = Block 1
     , splittingTree =
         SplittingTree
@@ -265,7 +256,7 @@ refineWithAllOutputs ls = concat <$> traverse (uncurry refineWithOutput) ls
 refineWithSplitterAllInputs :: (Monad m, Ord o, Ord s) => [(i, s -> [s])] -> Splitter s i o -> StateT (PRState s i o) m [Splitter s i o]
 refineWithSplitterAllInputs ls splitter = concat <$> traverse (\(i, rev) -> refineWithSplitter i rev splitter) ls
 
-refine :: (Monad m, Ord o, Ord s) => ([i] -> m ()) -> [(i, s -> o)] -> [(i, s -> [s])] -> StateT (PRState s i o) m ()
+refine :: (Monad m, Ord o, Ord s) => ([i] -> m ()) -> [(i, s -> o)] -> [(i, s -> [s])] -> StateT (PRState s i o) m (Partition s, SplittingTree s i o)
 refine ping outputs transitionsReverse = do
   initialQueue <- refineWithAllOutputs outputs
 
@@ -277,3 +268,18 @@ refine ping outputs transitionsReverse = do
       loop (splitters <> newQueue)
 
   loop initialQueue
+  PRState{..} <- get
+  pure (cleanupP partition, splittingTree)
+
+cleanupP :: Ord s => BarePartition s -> Partition s
+cleanupP (BarePartition m) =
+  let (p, (_, b)) = runState (Map.traverseWithKey update m) (Map.empty, Block 0)
+   in Partition p b
+ where
+  update _ v = do
+    (m2, n) <- get
+    case Map.lookup v m2 of
+      Just k -> pure k
+      Nothing -> do
+        put (Map.insert v n m2, succ n)
+        pure n
diff --git a/src/StateIdentifiers.hs b/src/StateIdentifiers.hs
index 5ad9229..6152cfe 100644
--- a/src/StateIdentifiers.hs
+++ b/src/StateIdentifiers.hs
@@ -1,7 +1,8 @@
 module StateIdentifiers where
 
-import SplittingTree
+import Data.Partition (Partition (..))
 import Data.Trie qualified as Trie
+import SplittingTree (SplittingTree (..))
 
 import Data.Map.Strict qualified as Map