diff --git a/valmari/partition.ml b/valmari/partition.ml
index 9027291..efc5620 100644
--- a/valmari/partition.ml
+++ b/valmari/partition.ml
@@ -9,11 +9,11 @@ type 'a loc_array = 'a array
 type 'a t = {
   mutable set_count: set;
   element  : 'a Finite.elt loc_array;
-  location : loc array;
-  set_of   : set array;
-  first    : loc set_array;
-  past     : loc set_array;
-  marked   : int set_array;
+  location : loc array; (* L *)
+  set_of   : set array; (* S *)
+  first    : loc set_array; (* F *)
+  past     : loc set_array; (* P *)
+  marked   : int set_array; (* M *)
   mutable worklist: set list;
 }
 
@@ -106,20 +106,17 @@ let split t =
     ) worklist
 
 let discard_unmarked t =
-  let worklist = t.worklist in
   t.worklist <- [];
-  List.iter (fun set ->
-      let first_unmarked = t.first.(set) + t.marked.(set) in
-      if first_unmarked < t.past.(set) then (
-        for i = first_unmarked to t.past.(set) - 1 do
-          let elt = (t.element.(i) : _ Finite.elt :> int) in
-          (*prerr_endline ("discarding " ^ string_of_int elt);*)
-          t.set_of.(elt) <- -1
-        done;
-        t.past.(set) <- first_unmarked;
-      );
-      t.marked.(set) <- 0
-    ) worklist
+  for set = 0 to t.set_count - 1 do
+    let first_unmarked = t.first.(set) + t.marked.(set) in
+    for i = first_unmarked to t.past.(set) - 1 do
+      let elt = (t.element.(i) : _ Finite.elt :> int) in
+      (*prerr_endline ("discarding " ^ string_of_int elt);*)
+      t.set_of.(elt) <- -1
+    done;
+    t.past.(set) <- first_unmarked;
+    t.marked.(set) <- 0
+  done
 
 let discard t f =
   for set = 0 to t.set_count - 1 do
@@ -136,8 +133,15 @@ let set_count t = t.set_count
 let set_of (t : 'a t) elt = t.set_of.((elt : 'a Finite.elt :> int))
 
 let choose t set =
+  assert (t.first.(set) < t.past.(set));
   t.element.(t.first.(set))
 
+let choose_opt t set =
+  if t.first.(set) < t.past.(set) then
+    Some t.element.(t.first.(set))
+  else
+    None
+
 let iter_elements t set f =
   for i = t.first.(set) to t.past.(set) - 1 do
     f t.element.(i)
@@ -160,3 +164,8 @@ let clear_marks t =
   t.worklist <- [];
   List.iter (fun set -> t.marked.(set) <- 0) worklist
 
+let is_first t n =
+  let n = (n : 'n Finite.elt :> int) in
+  let s = t.set_of.(n) in
+  let loc = t.location.(n) in
+  (s > -1 && loc = t.first.(s) && loc < t.past.(s))
diff --git a/valmari/partition.mli b/valmari/partition.mli
index 3dca0f2..b7badbd 100644
--- a/valmari/partition.mli
+++ b/valmari/partition.mli
@@ -1,21 +1,87 @@
 open Strong
+
+(* An intermediate datastructure used by Valmari automata minimization
+   algorithm for efficiently representing incremental refinements of a set
+   partition. *)
+
 type set = int
+(** Each set is identified by an integer *)
+
+type 'n t
+(** A partitioning structure for a set of cardinal 'n \
+    (encoded as a Strong.Natural) *)
 
-type 'a t
 val create :
-  ?partition:('a Finite.elt -> 'a Finite.elt -> int) ->
-  'a Finite.set -> 'a t
-
-val mark : 'a t -> 'a Finite.elt -> unit
-val split : 'a t -> unit
-val discard_unmarked : 'a t -> unit
-val discard : 'a t -> ('a Finite.elt -> bool) -> unit
-
-val set_count : 'a t -> int
-val set_of : 'a t -> 'a Finite.elt -> set
-val choose : 'a t -> set -> 'a Finite.elt
-val iter_elements : 'a t -> set -> ('a Finite.elt -> unit) -> unit
-val iter_marked_elements : 'a t -> set -> ('a Finite.elt -> unit) -> unit
-
-val clear_marks : 'a t -> unit
-val marked_sets : 'a t -> set list
+  ?partition:('n Finite.elt -> 'n Finite.elt -> int) ->
+  'n Finite.set -> 'n t
+(** [create ?partition n] create a fresh partitioning data structure for a set
+    of cardinal [n].
+    If [partition] is not provided, the datastructure is initialized with a
+    single subset that contains all elements.
+    Otherwise, [partition] must be a total ordering function and elements that
+    can be distinguished are put in different subsets.
+*)
+
+val mark : 'n t -> 'n Finite.elt -> unit
+(** [mark part elt] marks the element [elt] as active.
+    The datastructure manages an active set by marking a certain number of
+    elements, and then applying an operation to all of them at once.
+*)
+
+val split : 'n t -> unit
+(** Put marked elements in different sets.
+    That is, each input set is split in two subsets one with the marked and one
+    with the unmarked elements.
+    Active set is reset after (no elements are marked).
+*)
+
+val discard_unmarked : 'n t -> unit
+(** Elements that are not marked are removed from the partition (they will be
+    ignored by future operations).
+    In practice, they are considered as belonging to set [-1] (which can be
+    observed by [set_of] function), and this [-1] set is not counted by the
+    [set_count] function.
+    Active set is reset after (no elements are marked).
+*)
+
+val discard : 'n t -> ('n Finite.elt -> bool) -> unit
+(** [discard part f] calls the function [f] for each element in the set
+    and discard it if the function returns [true].
+    Active set must be empty before and is reset after (no elements are marked).
+*)
+
+val set_count : 'n t -> int
+(** Number of sets in the current partition *)
+
+val set_of : 'n t -> 'n Finite.elt -> set
+(** [set_of part elt] returns the index of the set that contains element [elt].
+    Result is between [0] and [set_of part - 1] unless the element has been
+    discarded, in which case it is [-1]. *)
+
+val choose : 'n t -> set -> 'n Finite.elt
+(** [choose part set] returns an arbitrary element that belongs to set [set].
+    [set] must be between [0] and [set_of part - 1].
+*)
+
+val choose_opt : 'n t -> set -> 'n Finite.elt option
+(** [choose part set] returns an arbitrary element that belongs to set [set].
+    [set] must be between [0] and [set_of part - 1].
+*)
+
+val iter_elements : 'n t -> set -> ('n Finite.elt -> unit) -> unit
+(** [iter_elements part set f] applies function [f] to each element that
+    currently belongs to set [set].
+*)
+
+val iter_marked_elements : 'n t -> set -> ('n Finite.elt -> unit) -> unit
+(** [iter_marked_elements part set f] applies function [f] to each element that
+    currently belongs to set [set] and is marked.
+*)
+
+val clear_marks : 'n t -> unit
+(** Remove all marks (reset the active set) *)
+
+val marked_sets : 'n t -> set list
+(** Returns all sets that have marked elements. *)
+
+val is_first : 'n t -> 'n Finite.elt -> bool
diff --git a/valmari/test/test.ml b/valmari/test/test.ml
index 583d55f..f89b1ef 100644
--- a/valmari/test/test.ml
+++ b/valmari/test/test.ml
@@ -43,12 +43,15 @@ let () =
       let (_, _, d) = trans_table.((t : Transitions.n Finite.elt :> int)) in
       d
 
-    let initials =
-      [|Finite.Elt.of_int States.n initial_state|]
+    let initials f = f (Finite.Elt.of_int States.n initial_state)
 
-    let finals = Array.init final_state_count
-        (fun _i -> Scanf.bscanf ic "%d\n"
-            (Finite.Elt.of_int States.n))
+    let finals =
+      let finals =
+        Array.init final_state_count
+          (fun _i -> Scanf.bscanf ic "%d\n"
+              (Finite.Elt.of_int States.n));
+      in
+      fun f -> Array.iter f finals
 
     let refinements ~refine:_ = ()
   end in
diff --git a/valmari/valmari.ml b/valmari/valmari.ml
index 9d707ad..2885b57 100644
--- a/valmari/valmari.ml
+++ b/valmari/valmari.ml
@@ -11,13 +11,14 @@ module type DFA = sig
   val source : transitions Fin.elt -> states Fin.elt
   val target : transitions Fin.elt -> states Fin.elt
 
-  val initials : states Fin.elt array
-  val finals : states Fin.elt array
 end
 
 module type INPUT = sig
   include DFA
 
+  val initials : (states Fin.elt -> unit) -> unit
+  val finals : (states Fin.elt -> unit) -> unit
+
   val refinements :
     refine:(iter:((states Fin.elt -> unit) -> unit) -> unit) -> unit
 end
@@ -68,6 +69,9 @@ module Minimize
 sig
   include DFA with type label = Label.t
 
+  val initials : states Fin.elt array
+  val finals : states Fin.elt array
+
   val transport_state :
     In.states Fin.elt -> states Fin.elt option
   val transport_transition :
@@ -84,7 +88,7 @@ end = struct
 
   (* Remove states unreachable from initial state *)
   let () =
-    Array.iter (Partition.mark blocks) In.initials;
+    In.initials (Partition.mark blocks);
     let transitions_source =
       index_transitions In.states In.transitions In.source in
     discard_unreachable blocks transitions_source In.target
@@ -93,14 +97,14 @@ end = struct
   let transitions_targeting =
     index_transitions In.states In.transitions In.target
 
-  (* Remove states unreachable from final states *)
+  (* Remove states which cannot reach any final state *)
   let () =
-    Array.iter (Partition.mark blocks) In.finals;
+    In.finals (Partition.mark blocks);
     discard_unreachable blocks transitions_targeting In.source
 
   (* Split final states *)
   let () =
-    Array.iter (Partition.mark blocks) In.finals;
+    In.finals (Partition.mark blocks);
     Partition.split blocks
 
   (* Split explicitely refined states *)
@@ -148,36 +152,45 @@ end = struct
   module Transitions = Fin.Array.Of_array(struct
       type a = In.transitions Fin.elt
       let table =
-        match Partition.set_count cords with
+        let count = ref 0 in
+        Fin.Set.iter In.transitions (fun tr ->
+            if Partition.is_first blocks (In.source tr) &&
+               Partition.set_of blocks (In.target tr) > -1
+            then incr count
+          );
+        match !count with
         | 0 -> [||]
-        | count ->
-          let count' = ref 0 in
-          for i = 0 to count - 1 do
-            let elt = Partition.choose cords i in
-            if Partition.set_of blocks (In.target elt) > -1 then
-              incr count';
-          done;
-          let table = Array.make !count' (Partition.choose cords 0) in
-          let count' = ref 0 in
-          for i = 0 to count - 1 do
-            let elt = Partition.choose cords i in
-            if Partition.set_of blocks (In.target elt) > -1 then (
-              table.(!count') <- elt;
-              incr count';
+        | n -> Array.make n (Fin.Elt.of_int In.transitions 0)
+
+      let () =
+        let count = ref 0 in
+        Fin.Set.iter In.transitions (fun tr ->
+            if Partition.is_first blocks (In.source tr) &&
+               Partition.set_of blocks (In.target tr) > -1
+            then (
+              let index = !count in
+              incr count;
+              table.(index) <- tr
             )
-          done;
-          table
+          );
     end)
   type transitions = Transitions.n
   let transitions = Transitions.n
 
   type label = Label.t
 
-  let transport_state_unsafe state =
-    Fin.Elt.of_int states (Partition.set_of blocks state)
+  let transport_state_unsafe =
+    let table =
+      Fin.Array.init In.states (Partition.set_of blocks)
+    in
+    Fin.Array.get table
 
-  let represent_state state =
-    Partition.choose blocks (state : states Fin.elt :> int)
+  let represent_state =
+    let table =
+      Fin.Array.init states
+        (fun st -> Partition.choose blocks (st : states Fin.elt :> int))
+    in
+    Fin.Array.get table
 
   let represent_transition transition =
     Fin.(Transitions.table.(transition))
@@ -186,28 +199,36 @@ end = struct
     In.label (represent_transition transition)
 
   let source transition =
-    transport_state_unsafe (In.source (represent_transition transition))
+    Fin.Elt.of_int states
+      (transport_state_unsafe (In.source (represent_transition transition)))
 
   let target transition =
-    transport_state_unsafe (In.target (represent_transition transition))
+    Fin.Elt.of_int states
+      (transport_state_unsafe (In.target (represent_transition transition)))
 
   let initials =
-    Array.map transport_state_unsafe In.initials
+    In.initials (Partition.mark blocks);
+    let sets = Partition.marked_sets blocks in
+    Partition.clear_marks blocks;
+    Array.map (Fin.Elt.of_int states) (Array.of_list sets)
 
   let finals =
-    Array.iter (Partition.mark blocks) In.finals;
+    In.finals (Partition.mark blocks);
     let sets = Partition.marked_sets blocks in
     Partition.clear_marks blocks;
     Array.map (Fin.Elt.of_int states) (Array.of_list sets)
 
   let transport_state state =
-    match Partition.set_of blocks state with
+    match transport_state_unsafe state with
     | -1 -> None
     | n -> Some (Fin.Elt.of_int states n)
 
-  let transport_transition transition =
-    match Partition.set_of cords transition with
-    | -1 -> None
-    | n -> Some (Fin.Elt.of_int transitions n)
+  let transport_transition =
+    let table = Fin.Array.make In.transitions None in
+    Fin.Array.iteri (fun tr trin ->
+        assert (Fin.Array.get table trin = None);
+        Fin.Array.set table trin (Some tr);
+      ) Transitions.table;
+    Fin.Array.get table
 
 end
diff --git a/valmari/valmari.mli b/valmari/valmari.mli
index 6f07f52..4e3a91e 100644
--- a/valmari/valmari.mli
+++ b/valmari/valmari.mli
@@ -1,25 +1,61 @@
 open Strong.Finite
 
+(** Valmari is an automata minimization algorithm, described in
+    "Fast brief practical DFA minimization"
+    https://dl.acm.org/doi/10.1016/j.ipl.2011.12.004 *)
+
 module type DFA = sig
+
   type states
   val states : states set
+  (** The set of DFA nodes *)
+
   type transitions
   val transitions : transitions set
+  (** The set of DFA transitions *)
 
   type label
+  (** The type of labels that annotate transitions *)
+
   val label  : transitions elt -> label
+  (** Get the label associated with a transition *)
+
   val source : transitions elt -> states elt
-  val target : transitions elt -> states elt
+  (** Get the source state of the transition *)
 
-  val initials : states elt array
-  val finals : states elt array
+  val target : transitions elt -> states elt
+  (** Get the target state of the transition *)
 end
 
 module type INPUT = sig
   include DFA
 
+  val initials : (states elt -> unit) -> unit
+  (** Iterate on initial states *)
+
+  val finals : (states elt -> unit) -> unit
+  (** Iterate final states *)
+
   val refinements :
     refine:(iter:((states elt -> unit) -> unit) -> unit) -> unit
+    (** The minimization algorithms operate on a DFA plus an optional initial
+        refinement (state that must be distinguished, because of some external
+        properties not observable from the labelled transitions alone).
+
+        If no refinements are needed, the minimum implementation is just:
+          [let refinements ~refine:_ = ()]
+
+        Otherwise, the [refinements] function should invoke the [refine]
+        function for each set of equivalent states and call the [iter] for each
+        equivalent state.
+
+        E.g if our automata has 5 states, and states 2 and 3 have tag A while
+        states 4 and 5 have tag B, we will do:
+
+        let refinements ~refine =
+          refine (fun ~iter -> iter [2; 3]);
+          refine (fun ~iter -> iter [4; 5])
+    *)
 end
 
 module Minimize
@@ -27,6 +63,8 @@ module Minimize
     (In: INPUT with type label := Label.t) :
 sig
   include DFA with type label = Label.t
+  val initials : states elt array
+  val finals : states elt array
 
   val transport_state : In.states elt -> states elt option
   val transport_transition : In.transitions elt -> transitions elt option