Updates to function unfolding to support scattered definitions

1 files changed, 36 insertions, 22 deletions

diff --git a/src/optimize.ml b/src/optimize.ml
index 1fc2fbe8..b0d05bef 100644
--- a/src/optimize.ml
+++ b/src/optimize.ml
@@ -52,43 +52,57 @@ open Ast
 open Ast_util
 open Rewriter
 
+let rec split_at_function' id defs acc =
+  match defs with
+  | [] -> None
+  | ([def], env) :: defs when is_fundef id def -> Some (acc, (def, env), defs)
+  | (def, env) :: defs -> split_at_function' id defs ((def, env) :: acc)
+
+let split_at_function id defs =
+  match split_at_function' id defs [] with
+  | None -> None
+  | Some (pre_defs, def, post_defs) ->
+     Some (List.rev pre_defs, def, post_defs)
+
 let recheck (Defs defs) =
   let defs = Type_check.check_with_envs Type_check.initial_env defs in
 
   let rec find_optimizations = function
-    | ([DEF_pragma ("optimize", pragma, p_l)], env)
-      :: ([DEF_spec vs as def1], _)
-      :: ([DEF_fundef fdef as def2], _)
-      :: defs ->
+    | ([DEF_pragma ("optimize", pragma, p_l)], env) :: ([DEF_spec vs as def1], _) :: defs ->
        let id = id_of_val_spec vs in
        let args = Str.split (Str.regexp " +") (String.trim pragma) in
        begin match args with
        | ["unroll"; n]->
           let n = int_of_string n in
+          begin match split_at_function id defs with
+          | Some (intervening_defs, ((DEF_fundef fdef as def2, _)), defs) ->
+             let rw_app subst (fn, args) =
+               if Id.compare id fn = 0 then E_app (subst, args) else E_app (fn, args)
+             in
+             let rw_exp subst = { id_exp_alg with e_app = rw_app subst } in
+             let rw_defs subst = { rewriters_base with rewrite_exp = (fun _ -> fold_exp (rw_exp subst)) } in
 
-          let rw_app subst (fn, args) =
-            if Id.compare id fn = 0 then E_app (subst, args) else E_app (fn, args)
-          in
-          let rw_exp subst = { id_exp_alg with e_app = rw_app subst } in
-          let rw_defs subst = { rewriters_base with rewrite_exp = (fun _ -> fold_exp (rw_exp subst)) } in
-
-          let specs = ref [def1] in
-          let bodies = ref [rewrite_def (rw_defs (append_id id "_unroll_1")) def2] in
+             let specs = ref [def1] in
+             let bodies = ref [rewrite_def (rw_defs (append_id id "_unroll_1")) def2] in
 
-          for i = 1 to n do
-            let current_id = append_id id ("_unroll_" ^ string_of_int i) in
-            let next_id = if i = n then current_id else append_id id ("_unroll_" ^ string_of_int (i + 1)) in
-            (* Create a valspec for the new unrolled function *)
-            specs := !specs @ [DEF_spec (rename_valspec current_id vs)];
-            (* Then duplicate it's function body and make it call the next unrolled function *)
-            bodies := !bodies @ [rewrite_def (rw_defs next_id) (DEF_fundef (rename_fundef current_id fdef))]
-          done;
+             for i = 1 to n do
+               let current_id = append_id id ("_unroll_" ^ string_of_int i) in
+               let next_id = if i = n then current_id else append_id id ("_unroll_" ^ string_of_int (i + 1)) in
+               (* Create a valspec for the new unrolled function *)
+               specs := !specs @ [DEF_spec (rename_valspec current_id vs)];
+               (* Then duplicate it's function body and make it call the next unrolled function *)
+               bodies := !bodies @ [rewrite_def (rw_defs next_id) (DEF_fundef (rename_fundef current_id fdef))]
+             done;
 
-          !specs @ !bodies @ find_optimizations defs
+             !specs @ List.concat (List.map fst intervening_defs) @ !bodies @ find_optimizations defs
 
+          | _ ->
+             Reporting.warn "Could not find function body for unroll pragma at " p_l "";
+             def1 :: find_optimizations defs
+          end
        | _ ->
           Reporting.warn "Unrecognised optimize pragma at" p_l "";
-          def1 :: def2 :: find_optimizations defs
+          def1 :: find_optimizations defs
        end
 
     | (defs, _) :: defs' ->