diff --git a/lib/OptimizeMiniRust.ml b/lib/OptimizeMiniRust.ml
index 494d77e3..375525a2 100644
--- a/lib/OptimizeMiniRust.ml
+++ b/lib/OptimizeMiniRust.ml
@@ -787,3 +787,72 @@ let infer_mut_borrows files =
     | x -> x
     ) decls
   ) (List.rev files)
+
+
+(* Transformations occuring on the MiniRust AST, after translation and mutability inference *)
+
+(* Loop unrolling introduces an implicit binder that does not interact well
+   with the substitutions occurring in mutability inference.
+   We perform it after the translation *)
+let unroll_loops = object
+  inherit [_] map_expr as super
+  method! visit_For _ b e1 e2 =
+    let e2 = super#visit_expr () e2 in
+
+    match e1 with
+    | Range (Some (Constant (UInt32, init as k_init) as e_init), Some (Constant (UInt32, max)), false)
+      when (
+        let init = int_of_string init in
+        let max = int_of_string max in
+        let n_loops = max - init in
+        n_loops <= !Options.unroll_loops
+      ) ->
+        let init = int_of_string init in
+        let max = int_of_string max in
+        let n_loops = max - init in
+
+        if n_loops = 0 then Unit
+
+        else if n_loops = 1 then subst e_init 0 e2
+
+        else Call (Name ["krml"; "unroll_for!"], [], [
+          Constant (CInt, string_of_int n_loops);
+          ConstantString b.name;
+          Constant k_init;
+          Constant (UInt32, "1");
+          e2
+        ])
+
+    | _ -> For (b, e1, e2)
+end
+
+let remove_trailing_unit = object
+  inherit [_] map_expr as super
+  method! visit_Let _ b e1 e2 =
+    let e1 = super#visit_expr () e1 in
+    let e2 = super#visit_expr () e2 in
+    match e2 with
+    | Unit -> e1
+    | _ -> Let (b, e1, e2)
+end
+
+let map_funs f_map files =
+  let files =
+    List.fold_left (fun files (filename, decls) ->
+      let decls = List.fold_left (fun decls decl ->
+        match decl with
+        | Function ({ body; _ } as f) ->
+          let body = f_map () body in
+          Function {f with body} :: decls
+        | _ -> decl :: decls
+      ) [] decls in
+      let decls = List.rev decls in
+      (filename, decls) :: files
+    ) [] files
+  in
+  List.rev files
+
+let simplify_minirust files =
+  let files = map_funs unroll_loops#visit_expr files in
+  let files = map_funs remove_trailing_unit#visit_expr files in
+  files
diff --git a/lib/SimplifyRust.ml b/lib/SimplifyRust.ml
index 30a5f89e..132b8ebe 100644
--- a/lib/SimplifyRust.ml
+++ b/lib/SimplifyRust.ml
@@ -2,8 +2,6 @@
 
 open Ast
 
-(* Simplifications occuring on Ast *)
-
 (* Doesn't make sense in Rust *)
 let remove_push_pop_frame = object
   inherit [_] map
@@ -15,73 +13,3 @@ let simplify_ast files =
   let files = remove_push_pop_frame#visit_files () files in
   files
 
-
-(* Transformations occuring on the MiniRust AST, after translation and mutability inference *)
-
-open MiniRust
-
-(* Loop unrolling introduces an implicit binder that does not interact well
-   with the substitutions occurring in mutability inference.
-   We perform it after the translation *)
-let unroll_loops = object
-  inherit [_] map_expr as super
-  method! visit_For _ b e1 e2 =
-    let e2 = super#visit_expr () e2 in 
-
-    match e1 with
-    | Range (Some (Constant (UInt32, init as k_init) as e_init), Some (Constant (UInt32, max)), false)
-      when (
-        let init = int_of_string init in
-        let max = int_of_string max in
-        let n_loops = max - init in
-        n_loops <= !Options.unroll_loops
-      ) ->
-        let init = int_of_string init in
-        let max = int_of_string max in
-        let n_loops = max - init in 
-
-        if n_loops = 0 then Unit
-
-        else if n_loops = 1 then subst e_init 0 e2
-
-        else Call (Name ["krml"; "unroll_for!"], [], [
-          Constant (CInt, string_of_int n_loops);
-          ConstantString b.name;
-          Constant k_init;
-          Constant (UInt32, "1");
-          e2
-        ])
-
-    | _ -> For (b, e1, e2)
-end 
-
-let remove_trailing_unit = object
-  inherit [_] map_expr as super
-  method! visit_Let _ b e1 e2 =
-    let e1 = super#visit_expr () e1 in
-    let e2 = super#visit_expr () e2 in
-    match e2 with
-    | Unit -> e1
-    | _ -> Let (b, e1, e2)
-end
-
-let map_funs f_map files =
-  let files =
-    List.fold_left (fun files (filename, decls) ->
-      let decls = List.fold_left (fun decls decl ->
-        match decl with
-        | Function ({ body; _ } as f) ->
-          let body = f_map () body in
-          Function {f with body} :: decls
-        | _ -> decl :: decls
-      ) [] decls in
-      let decls = List.rev decls in
-      (filename, decls) :: files
-    ) [] files
-  in
-  List.rev files
-
-let simplify_minirust files =
-  let files = map_funs unroll_loops#visit_expr files in
-  let files = map_funs remove_trailing_unit#visit_expr files in
-  files
diff --git a/src/Karamel.ml b/src/Karamel.ml
index 98d58be3..a7568766 100644
--- a/src/Karamel.ml
+++ b/src/Karamel.ml
@@ -751,7 +751,7 @@ Supported options:|}
       print PrintAst.print_files files;
     let files = AstToMiniRust.translate_files files in
     let files = OptimizeMiniRust.infer_mut_borrows files in
-    let files = SimplifyRust.simplify_minirust files in
+    let files = OptimizeMiniRust.simplify_minirust files in
     OutputRust.write_all files
 
   else