Change order of prose's instructions

spectec/src/il2al/animate.ml

@@ -107,12 +107,20 @@ let rec recover_iterexp' iterexp pr =
   | IterPr (pr, (iter, xes)) -> IterPr (recover_iterexp iterexp pr, (iter, xes |> List.map (fun (x, e) -> (x, new_ e))))
 and recover_iterexp iterexp pr = Source.map (recover_iterexp' iterexp) pr
 
+(* is this assign premise a if-let? *)
+let is_cond_assign prem =
+  match prem.it with
+  | LetPr ({it = CaseE (_, _); _}, _, _) -> true
+  | _ -> false
 
-(* is this assign premise encoded premise for pop? *)
+(* is this assign premise encoded premise for popping one value? *)
 let is_pop env row =
   is_assign env row &&
   match (unwrap row).it with
-  | LetPr (_, rhs, _) -> Il.Print.string_of_typ rhs.note = "stackT"
+  | LetPr (_, {it = CallE (_, {it = ExpA n; _} :: _); note; _}, _) when Il.Print.string_of_typ note = "stackT" ->
+    (match n.it with
+    | NatE i -> Z.equal i (Z.one)
+    | _ -> false)
   | _ -> false
 
 (* iteratively select pop, condition and assignment premises,
@@ -161,7 +169,11 @@ and select_assign prems acc env fb =
         fb := (len, acc @ List.map unwrap non_assigns);
       None
     | _ ->
-      let assigns' = List.map unwrap assigns in
+      let cond_assigns, non_cond_assigns =
+        List.map unwrap assigns
+        |> List.partition is_cond_assign
+      in
+      let assigns' = cond_assigns @ non_cond_assigns in
       let new_env = assigns
         |> List.map targets
         |> List.concat

spectec/src/il2al/translate.ml

@@ -107,11 +107,11 @@ let rec insert_instrs target il =
     h @ [ ifI (cond, insert_instrs (Transpile.insert_nop target) il' , []) ]
   | _ -> il @ target
 
-(* Insert `target` at the else-branch of innermost if instruction *)
-let rec insert_instrs_to_else target il =
+(* Insert `target` at the last instruction *)
+let insert_to_last target il =
   match Util.Lib.List.split_last_opt il with
-  | Some (h, { it = IfI (cond, il1, il2); _ }) ->
-    h @ [ ifI (cond, il1, insert_instrs_to_else target il2) ]
+  | Some (hd, tl) ->
+    hd @ Transpile.merge_blocks [[tl]; [otherwiseI target]]
   | _ -> il @ target
 
 let has_branch = List.exists (fun i ->
@@ -1169,8 +1169,8 @@ let rec translate_rgroup' (rule: rule_def) =
     | None, b -> b
     | Some b, [] -> b
     | Some b1, b2 ->
-      (* Assert: b1 must have the else-less IfI as inner most instruction *)
-      insert_instrs_to_else b2 (Transpile.flatten_if b1)
+      (* Assert: last instruction of b1 must be else-less IfI *)
+      insert_to_last b2 b1
   in
 
   translate_prems pops body_instrs

spectec/src/il2al/unify.ml

@@ -354,13 +354,42 @@ let unify_rule_def (rule: rule_def) : rule_def =
   in
   (instr_name, rel_id, new_clauses) $ rule.at
 
+let has_substring str sub =
+  let len_s = String.length str in
+  let len_sub = String.length sub in
+  let rec aux i =
+    if i > len_s - len_sub then false
+    else if String.sub str i len_sub = sub then true
+    else aux (i + 1)
+  in
+  aux 0
+
+let reorder_unified_args args prems =
+  (* Helpers *)
+  let has_uarg_on_rhs p =
+    match p.it with
+    | LetPr (_, {it = VarE id; _}, _) -> is_unified_id id.it
+    | _ -> false
+  in
+  let on_rhs p a =
+    match a.it with
+    | ExpA e -> has_substring (e |> Il.Print.string_of_exp) (rhs_of_prem p |> Il.Print.string_of_exp)
+    | _ -> false
+  in
+  let index_of p = List.find_index (on_rhs p) args |> Option.get in
+  let cmp p1 p2 = index_of p1 - index_of p2 in
+
+  let uprems, prems = List.partition has_uarg_on_rhs prems in
+  let uprems' = List.sort cmp uprems in
+  uprems' @ prems
 
 let apply_template_to_def template def =
   match def.it with
   | DefD (binds, lhs, rhs, prems) ->
     let new_prems, new_binds = collect_unified_args template lhs in
     let animated_prems = Animate.animate_prems (free_list free_arg template) new_prems in
-    DefD (binds @ new_binds, template, rhs, (animated_prems @ prems) |> prioritize_else) $ def.at
+    let reordered_prems = reorder_unified_args template animated_prems in
+    DefD (binds @ new_binds, template, rhs, (reordered_prems @ prems) |> prioritize_else) $ def.at
 
 let unify_defs defs =
   init_unified_idx();