diff --git a/compiler/heptagon/transformations/automata_mapfold.ml b/compiler/heptagon/transformations/automata_mapfold.ml
index 0839854..2a5f4e1 100644
--- a/compiler/heptagon/transformations/automata_mapfold.ml
+++ b/compiler/heptagon/transformations/automata_mapfold.ml
@@ -170,3 +170,4 @@ let program p =
                with eq = eq; block = block; node_dec = node_dec } in
   let p, _ = Hept_mapfold.program_it funs ([],[]) p in
     p
+    { p with p_types = !state_type_dec_list @ p.p_types }