Pop the values from the stack ASAP

spectec/src/il/eq.ml

@@ -138,6 +138,8 @@ and eq_prem prem1 prem2 =
   | IfPr e1, IfPr e2 -> eq_exp e1 e2
   | IterPr (prem1, e1), IterPr (prem2, e2) ->
     eq_prem prem1 prem2 && eq_iterexp e1 e2
+  | LetPr (e1, e1', ids1), LetPr (e2, e2', ids2) ->
+    eq_exp e1 e2 && eq_exp e1' e2' && ids1 = ids2
   | _, _ -> prem1.it = prem2.it
 
 

spectec/src/il2al/animate.ml

@@ -53,7 +53,26 @@ let is_pop env row =
 
 (* iteratively select pop, condition and assignment premises,
  * effectively sorting the premises as a result. *)
-let rec select_tight prems acc env fb =
+let rec select_pops prems acc env fb =
+  match prems with
+  | [] -> Some acc
+  | _ ->
+    let (pops, non_pops) = List.partition (is_pop env) prems in
+    match pops with
+    | [] ->
+      select_tight prems acc env fb
+    | _ ->
+      let pops' = List.map unwrap pops in
+      let new_env = pops
+        |> List.map targets
+        |> List.concat
+        |> List.fold_left (fun env x ->
+          union env { empty with varid = Set.singleton x }
+        ) env
+      in
+      select_pops non_pops (acc @ pops') new_env fb
+
+and select_tight prems acc env fb =
   match prems with
   | [] -> Some acc
   | _ ->
@@ -234,7 +253,7 @@ let animate_prems known_vars prems =
 
   (* 2. Reorder *)
   let best' = ref (-1, []) in
-  match List.find_map (fun cand -> select_tight cand other known_vars best') candidates with
+  match List.find_map (fun cand -> select_pops cand other known_vars best') candidates with
   | None ->
     if (not k_fail) then
       let unhandled_prems = Lib.List.drop (fst !best') (snd !best') in

spectec/src/il2al/encode.ml

@@ -64,10 +64,11 @@ let rec stack_to_list e =
   (* List.map (fun e -> { e with it = ListE [e] }) es *)
   | _ -> [ e ]
 
-let rec drop_until f xs =
+let rec drop_until' acc f xs =
   match xs with
-  | [] -> []
-  | hd :: tl -> if f hd then xs else drop_until f tl
+  | [] -> acc, []
+  | hd :: tl -> if f hd then acc, xs else drop_until' (hd :: acc) f tl
+let drop_until = drop_until' []
 
 let free_ids e =
   (free_exp e)
@@ -90,12 +91,21 @@ let input_vars = [
   "stack0";
   "ctxt";
   "state";
+  "unused";
 ]
 
 (* Encode stack *)
 
 let encode_inner_stack context_opt stack =
-  let es = stack_to_list stack |> List.rev |> drop_until is_case in
+  let dropped, es = stack_to_list stack |> List.rev |> drop_until is_case in
+
+  let unused_prems =
+    if List.length dropped = 0 then
+      []
+    else
+      let unused = TupE dropped $$ no_region % (mk_varT "unusedT") in
+      [LetPr (unused, mk_varE "unused" "unusedT", free_ids unused) $ no_region]
+  in
 
   match es with
   | [] ->
@@ -104,7 +114,7 @@ let encode_inner_stack context_opt stack =
       match context_opt with
       | None -> assert false
       | Some e ->
-        Some (LetPr (e, mk_varE "input" "inputT", free_ids e) $ e.at), []
+        Some (LetPr (e, mk_varE "input" "inputT", free_ids e) $ e.at), unused_prems
     )
   | _ ->
     (* ASSUMPTION: The top of the stack should be now the target instruction *)
@@ -126,7 +136,7 @@ let encode_inner_stack context_opt stack =
       IfPr (CmpE (EqOp, lhs, rhs) $$ e.at % (BoolT $ no_region)) $ e.at
     ) operands in
 
-    None, prem :: prems
+    None, prem :: prems @ unused_prems
 
 let encode_stack stack =
   match stack.it with

spectec/src/il2al/il2il.ml

@@ -18,11 +18,17 @@ type rgroup = (exp * exp * (prem list)) phrase list
 
 let unified_prefix = "u"
 let _unified_idx = ref 0
-let init_unified_idx () = _unified_idx := 0
+let _unified_idx_cache = ref None
+let init_unified_idx () = _unified_idx := 0; _unified_idx_cache := None
+let soft_init_unified_idx () =
+  match !_unified_idx_cache with
+  | None -> _unified_idx_cache := Some (!_unified_idx)
+  | Some n -> _unified_idx := n
 let get_unified_idx () = let i = !_unified_idx in _unified_idx := (i+1); i
 let gen_new_unified ty = (Al.Al_util.typ_to_var_name ty) ^ "_" ^ unified_prefix ^ (string_of_int (get_unified_idx())) $ no_region
 let is_unified_id id = String.split_on_char '_' id |> Util.Lib.List.last |> String.starts_with ~prefix:unified_prefix
 
+
 let rec overlap e1 e2 = if eq_exp e1 e2 then e1 else
   let replace_it it = { e1 with it = it } in
   match e1.it, e2.it with
@@ -235,23 +241,30 @@ let unify_enc premss encs =
 
   List.map2 (apply_template_to_prems template) premss idxs
 
-let is_enc pr =
+let is_encoded_ctxt pr =
   match pr.it with
   | LetPr (_, e, _) ->
     (match e.note.it with
-    | VarT (id, []) -> List.mem id.it ["stackT"; "inputT"; "contextT"]
+    | VarT (id, []) -> List.mem id.it ["inputT"; "stackT"; "contextT"]
+    | _ -> false)
+  | _ -> false
+let is_encoded_pop_or_winstr pr =
+  match pr.it with
+  | LetPr (_, e, _) ->
+    (match e.note.it with
+    | VarT (id, []) -> List.mem id.it ["inputT"; "stackT"]
     | _ -> false)
   | _ -> false
 
-let rec extract_encs' cnt =
+let rec extract_encs' pred cnt =
   function
   | [] -> []
   | hd :: tl ->
-    if is_enc hd then
-      (cnt, hd) :: extract_encs' (cnt + 1) tl
+    if pred hd then
+      (cnt, hd) :: extract_encs' pred (cnt + 1) tl
     else
-      extract_encs' (cnt + 1) tl
-let extract_encs = extract_encs' 0
+      extract_encs' pred (cnt + 1) tl
+let extract_encs pred = extract_encs' pred 0
 
 let has_identical_rhs iprem1 iprem2 =
   let rhs1 = iprem1 |> snd |> rhs_of_prem in
@@ -281,25 +294,23 @@ let rec filter_unifiable encss =
 
 let replace_prems r prems =
   let (lhs, rhs, _prems) = r.it in
-  (*
-  List.iter (fun p -> print_endline (Il.Print.string_of_prem p)) _prems;
-  print_endline "->";
-  List.iter (fun p -> print_endline (Il.Print.string_of_prem p)) prems;
-  print_endline "";
-  *)
   { r with it = (lhs, rhs, prems) }
 
-let unify_rgroup rgroup =
-  init_unified_idx();
-
+let unify_rgroup pred input_vars rgroup =
   let premss = List.map (fun g -> let (_, _, prems) = g.it in prems) rgroup in
-  let encss = List.map extract_encs premss in
+  let encss = List.map (extract_encs pred) premss in
   let unifiable_encss = filter_unifiable encss in
   let new_premss = List.fold_left unify_enc (lift premss) unifiable_encss |> unlift in
-  let animated_premss = List.map (Animate.animate_prems {empty with varid = Set.of_list Encode.input_vars}) new_premss in
+  let animated_premss = List.map (Animate.animate_prems {empty with varid = Set.of_list (input_vars @ Encode.input_vars)}) new_premss in
 
   List.map2 replace_prems rgroup animated_premss
 
+let unify_ctxt input_vars rgroup =
+  soft_init_unified_idx();
+  unify_rgroup is_encoded_ctxt input_vars rgroup
+let unify_pop_and_winstr rgroup =
+  init_unified_idx();
+  unify_rgroup is_encoded_pop_or_winstr [] rgroup
 
 (** 3. Functions **)
 

spectec/src/il2al/il2il.mli

@@ -3,5 +3,6 @@ open Il.Ast
 type rgroup = (exp * exp * (prem list)) Util.Source.phrase list
 
 val unify_rules : rule list -> rule list
-val unify_rgroup : rgroup -> rgroup
+val unify_pop_and_winstr : rgroup -> rgroup
+val unify_ctxt : string list -> rgroup -> rgroup
 val unify_defs : clause list -> clause list

spectec/src/il2al/translate.ml

@@ -792,6 +792,7 @@ and handle_special_lhs lhs rhs free_ids =
   (* Handle encoded premises generated by `encode.ml` *)
   | _ when (Il.Print.string_of_typ rhs.note) = "inputT" -> []
   | _ when (Il.Print.string_of_typ rhs.note) = "stateT" -> []
+  | _ when (Il.Print.string_of_typ rhs.note) = "unusedT" -> []
   | TupE [e; _stack] when (Il.Print.string_of_typ rhs.note) = "stackT" ->
     let args = args_of_call rhs in
     let pop_num = List.hd args |> arg2expr in
@@ -1068,8 +1069,9 @@ let rec add eq k v = function (* add a (k,v) to an assoc list *)
 
 let extract_winstr r =
   let (_, _, prems) = r.it in
-  List.find_opt is_winstr_prem prems
-  |> Option.map lhs_of_prem (* TODO: Collect helper functions into one place *)
+  match List.find_opt is_winstr_prem prems with
+  | Some p -> lhs_of_prem p (* TODO: Collect helper functions into one place *)
+  | None -> error r.at "Failed to extract the target wasm instruction"
 
 let extract_context r =
   let (_, _, prems) = r.it in
@@ -1085,26 +1087,17 @@ let exit_context context_opt instrs =
   | None -> instrs
   | Some instr -> instr :: instrs
 
-let remove_pop_after_otherwise prems =
-  match prems with
-  | hd :: tl when hd.it = Il.ElsePr -> hd :: Lib.List.filter_not is_pop tl;
-  | _ -> prems
-
 (* `reduction` -> `instr list` *)
 let translate_reduction ?(context_opt=None) reduction =
   let _, rhs, prems = reduction.it in
 
-  (* ASSUMPTION: ElsePr is only usable if all LHS are identical,
-     meaning all of pops after ElsePr can be/should be removed *)
-  let prems' = remove_pop_after_otherwise prems in
-
   (* Translate rhs *)
   translate_rhs rhs
   (* Exit context *)
   |> exit_context context_opt
   |> Transpile.insert_nop
   (* Translate premises *)
-  |> translate_prems prems'
+  |> translate_prems prems
 
 
 let translate_context_winstr winstr =
@@ -1170,27 +1163,56 @@ let translate_context ctx =
     exitI atom ~at:at
   | _ -> [ yetI "TODO: translate_context" ~at:at ], yetI "TODO: translate_context"
 
-let rec translate_rgroup' rgroup =
-  let subgroups =
-    rgroup
-    |> group_by_context
+let extract_pops rgroup =
+  (* Helpers *)
+  let rec extract_pops' acc prems premss =
+    match prems with
+    | hd :: tl when is_pop hd ->
+      let partitions = List.map (List.partition (Il.Eq.eq_prem hd)) premss in
+      let fsts = List.map fst partitions in
+      let snds = List.map snd partitions in
+
+      if List.for_all (fun l -> List.length l = 1) fsts then
+        extract_pops' (hd :: acc) tl snds
+      else
+        List.rev acc, prems :: premss
+    | _ -> List.rev acc, prems :: premss
+  in
+
+  let get_prems r =
+    let (_, _, prems) = r.it in
+    prems
   in
+  let set_prems r prems =
+    let (lhs, rhs, _) = r.it in
+    { r with it = (lhs, rhs, prems) }
+  in
+  (* End of helpers *)
 
-  (* TODO *)
-  let winstr = ref (Il.NatE (Z.of_int 42) $$ no_region % topT) in
+  let hd = List.hd rgroup in
+  let tl = List.tl rgroup in
 
-  let blocks = List.map (fun (ctxt, subgroup) ->
-    let u_group = Il2il.unify_rgroup subgroup in
+  let extracted, new_premss = extract_pops' [] (get_prems hd) (List.map get_prems tl) in
+  extracted, List.map2 set_prems rgroup new_premss
 
-    (* TODO: This should be done only once *)
-    begin match extract_winstr (List.hd u_group) with
-    | Some wi -> winstr := wi
-    | None -> () end;
+
+let rec translate_rgroup' rgroup =
+  let pops, rgroup' = extract_pops rgroup in
+  let subgroups = group_by_context rgroup' in
+
+  let blocks = List.map (fun (ctxt, subgroup) ->
+    let popped_vars = List.concat_map (fun p ->
+      match p.it with
+      | Il.LetPr (_, _, ids) -> ids
+      | _ -> assert false
+    ) pops in
+    let u_group = Il2il.unify_ctxt popped_vars subgroup in
 
     match ctxt with
     (* Normal case *)
     | None ->
-      let inner_pop_instrs = translate_context_winstr !winstr in
+      let winstr = extract_winstr (List.hd u_group) in
+      let inner_pop_instrs = translate_context_winstr winstr in
       let blocks = List.map translate_reduction u_group in
       let instrs =
         match blocks with
@@ -1221,13 +1243,16 @@ let rec translate_rgroup' rgroup =
       ]
   ) subgroups in
 
-  !winstr, Transpile.merge_blocks blocks
+  translate_prems pops (Transpile.merge_blocks blocks)
 
 (* Main translation for reduction rules
- * `rgroup` -> `Backend-prose.Algo` *)
+ * `rgroup` -> `Al.Algo` *)
 
 and translate_rgroup (instr_name, rel_id, rgroup) =
-  let winstr, instrs = translate_rgroup' rgroup in
+  let unified_rgroup = Il2il.unify_pop_and_winstr rgroup in
+
+  let winstr = extract_winstr (List.hd unified_rgroup) in
+  let instrs = translate_rgroup' unified_rgroup in
 
   let name =
     match case_of_case winstr with