Flat map statements into blocks, and use in Aslp

lib/lang/block.ml

@@ -106,11 +106,14 @@ let map ~phi f (b : ('v, 'e) t) : ('vv, 'ee) t =
   { stmts = Vector.map f b.stmts; phis = phi b.phis; attrib = b.attrib }
 
 (** Modify stmt list by creating a mutable copy of the underlying vector *)
-let fmap_stmts_copy (f : (('a, 'a, 'b) Stmt.t, 'c) Vector.t -> unit) b =
+let fmap_stmts_copy (f : (('a, 'a, 'b) Stmt.t, Vector.rw) Vector.t -> unit) b =
   let v = Vector.copy b.stmts in
   f v;
   { b with stmts = Vector.freeze v }
 
+let clear_stmts (b : ('v, 'e) t) =
+  { b with stmts = Vector.create () |> Vector.freeze }
+
 (** prepend statements to block statement list (copies underlying vector) *)
 let prepend_stmts (b : ('v, 'e) t) (nstmts : ('v, 'v, 'e) Stmt.t list) :
     ('v, 'e) t =

lib/lang/procedure.ml

@@ -403,8 +403,7 @@ let decl_block_exn p name ?(phis = [])
   let p = add_block p id ~phis ~stmts ~successors ~attrib () in
   (p, id)
 
-let modify_block p id
-    (f : (Var.t, BasilExpr.t) Block.t -> (Var.t, BasilExpr.t) Block.t) =
+let modify_block p id (f : _ Block.t -> _ Block.t) =
   let open Edge in
   let open G in
   let block = f (find_block p id) in
@@ -456,7 +455,10 @@ let replace_edge p id (block : (Var.t, BasilExpr.t) Block.t) =
     [to_] is modified to additionally have the original outgoing edges of
     [from_].
 
-    This includes any edges to [Return] nodes, if they exist on [from_]. *)
+    This includes any edges to [Return] nodes, if they exist on [from_].
+
+    TODO: how does this interact with phi nodes?? probably impossible to handle.
+*)
 let transplant_outgoing_edges p ~from ~to_ : _ t =
   let replace_outgoing_uses ~from ~to_ g =
     G.fold_succ_e
@@ -478,6 +480,21 @@ let transplant_incoming_edges p ~from ~to_ : _ t =
   in
   p |> map_graph (replace_incoming_uses ~from:(Begin from) ~to_:(Begin to_))
 
+(** Replaces uses of the old block ID with the new [(first, last)] block IDs.
+    The incoming edges to [old] will be redirected to [first] and the outgoing
+    edges of [old] will be rebased to originate from [last].
+
+    [first] and [last] may be the same. One or both of [first]/[last] may be the
+    same as [old]. If neither is the same as [old], [old] will be removed from
+    the procedure. *)
+let replace_block ~old ~new_:(new_first, new_last) proc =
+  proc
+  |> transplant_incoming_edges ~from:old ~to_:new_first
+  |> transplant_outgoing_edges ~from:old ~to_:new_last
+  |>
+  if ID.equal old new_first || ID.equal old new_last then Fun.id
+  else Fun.flip remove_block old
+
 let lookup_local_decl p v =
   Var.Decls.find_opt (local_decls p) v
   |> Option.or_lazy ~else_:(fun () ->
@@ -649,6 +666,122 @@ let flat_map_stmts_topo_rev ?visit rewriter p =
           update_block p bid (Block.flat_map ~rev:true ~phi:Fun.id rewriter b))
     p blocks
 
+type ('v, 'e) mapped_stmt =
+  [ `Stmts of (Var.t, Var.t, BasilExpr.t) Stmt.t list
+    (** Zero or more straight-line statements. *)
+  | `Blocks of (Var.t, BasilExpr.t) Block.t list
+    (** Zero or more straight-line blocks. *)
+  | `Graph of ID.t * ID.t * ('v, 'e) t
+    (** Multi-block subgraph. [`Graph (begin, end, proc)] represents
+        control-flow entering at [begin] and exiting at the [end], and with any
+        control-flow between them.
+
+        [proc] is the procedure modified to include fresh blocks [first] and
+        [end], as well as any other blocks or control-flow edges {i between}
+        them. [proc] should be unchanged aside from the addition of fresh blocks
+        and control-flow between them. [begin] should have no predcessors, and
+        dominate all new blocks, including [end]. [end] should have no
+        successors.
+
+        [begin] and [end] may be the same block. *) ]
+(** Program fragment to replace a statement during {!cfg_concatmap_block}:
+    either a list of sequential statements, list of sequential blocks, {i or} a
+    the procedure including an additional cfg fragment bounded by fresh entry
+    and exit block ids. *)
+
+(** [cfg_concatmap_block ~f bid proc] takes a function [f] from statement to a
+    code fragment, and replaces block [bid] in [proc] with the concatenated
+    result of mapping its statements through through [f]. [f] may return either
+    a statement list, block list, or [proc] modified to include a new CFG
+    fragment bounded by a begining or end block. See {!type-mapped_stmt} for
+    details about [f]'s return type.
+
+    {b Returns} [(first, last, proc)] where [proc] is the updated procedure and
+    [first] / [last] is the first / last block of the concatenated output
+    program. *)
+let cfg_concatmap_block ~(f : proc:_ t -> _ Stmt.t -> _ mapped_stmt) base_bid
+    proc =
+  let existing_bids = lazy (IDSet.of_iter (Iter.map fst (iter_blocks proc))) in
+  let is_fresh = fun bid -> not (IDSet.mem bid (Lazy.force existing_bids)) in
+
+  (* Applies the [f] function while recording the current block ID for
+     [`Stmts] insertion, if available.
+
+     Also normalises the output into either:
+     - an iter of already-inserted [(first,last)] bookend IDs, or
+     - a [(bid, stmts)] which will be inserted later into the specified [bid]. *)
+  let apply_f ~proc (cur_stmts_bid : ID.t option) stmt :
+      (_ t * ID.t option)
+      * ((ID.t * ID.t) Iter.t, ID.t * _ Stmt.t Iter.t) Either.t =
+    match f ~proc stmt with
+    | `Blocks [] | `Stmts [] -> ((proc, cur_stmts_bid), Left Iter.empty)
+    | `Blocks bs ->
+        let[@warning "+missing-record-field-pattern"] proc, bids =
+          List.fold_map
+            (fun proc { Block.attrib; stmts; phis } ->
+              fresh_block proc ~attrib ~phis ~stmts:(CCVector.to_list stmts) ())
+            proc bs
+        in
+        ((proc, None), Left (Iter.map Pair.dup (Iter.of_list bids)))
+    | `Graph (first, last, proc) ->
+        if is_fresh first && is_fresh last then
+          ((proc, None), Left (Iter.singleton (first, last)))
+        else failwith "cfg_concatmap_block: `Graph blocks should be fresh"
+    | `Stmts stmts ->
+        let proc, bid =
+          match cur_stmts_bid with
+          | None -> fresh_block proc ~name:(ID.name base_bid) ~stmts:[] ()
+          | Some bid -> (proc, bid)
+        in
+        ((proc, Some bid), Right (bid, Iter.of_list stmts))
+  in
+
+  (* Map, while generating block names for bare statements returned by the mapping function. *)
+  let (proc, _), mapped =
+    find_block proc base_bid |> Block.stmts_iter |> Iter.to_list
+    |> List.fold_map (fun (proc, b) -> apply_f ~proc b) (proc, Some base_bid)
+  in
+  (* Collects adjacent bare statements into a basic block, and inserts those statements. *)
+  let proc = modify_block proc base_bid Block.clear_stmts in
+  let proc, block_id_pairs =
+    Extras.group_succ_either mapped
+    |> List.fold_flat_map
+         (fun proc -> function
+           | Either.Left (hd, tl) -> (proc, Iter.(to_list (append_l (hd :: tl))))
+           | Either.Right ((bid, hd), rest) ->
+               let stmts =
+                 Iter.append hd (Iter.flat_map snd (Iter.of_list rest))
+                 |> CCVector.of_iter |> CCVector.freeze
+               in
+               ( modify_block proc bid (fun b -> { b with stmts }),
+                 [ (bid, bid) ] ))
+         proc
+  in
+  (* Transplant predecessors and successors of the original block as needed. *)
+  let first, last =
+    ( List.head_opt block_id_pairs |> Option.map_or fst ~default:base_bid,
+      List.last_opt block_id_pairs |> Option.map_or snd ~default:base_bid )
+  in
+  let proc =
+    proc
+    |> transplant_outgoing_edges ~from:base_bid ~to_:last
+    |>
+    if not ID.(equal first base_bid) then
+      add_goto ~from:base_bid ~targets:[ first ]
+    else Fun.id
+  in
+  (* Insert gotos between mapped blocks. This must happen after transplanting
+     so we do not transplant these edges. *)
+  let proc =
+    List.combine_gen block_id_pairs (List.drop 1 block_id_pairs)
+    |> Iter.of_gen
+    |> Iter.fold
+         (fun proc ((_, prev), (next, _)) ->
+           add_goto proc ~from:prev ~targets:[ next ])
+         proc
+  in
+  (first, last, proc)
+
 let pretty_spec show_var show_expr (p : ('a, 'b) proc_spec) =
   let open Containers_pp in
   let ml f v = if List.is_empty v then [] else [ f v ] in

lib/lang/program.ml

@@ -93,9 +93,9 @@ let pretty_declaration d =
   | Procedure { definition } -> pretty_proc definition
 
 (*match definition with
-      | Some d -> 
+      | Some d ->
       let param, rt = Types.uncurry (Var.typ binding) in
-      let param = 
+      let param =
       text "let " ^ text (Var.name binding) ^ text (Var.to_decl_string_il binding)
       | None -> text @@ Var.to_decl_string_il binding)
       *)
@@ -284,6 +284,16 @@ let flat_map_decls f p =
          else { prog with declarations = IDMap.remove i prog.declarations })
        p
 
+(** Iterates over global variables in the given program, including both read and
+    assigned variables. Order is unspecified and may have duplicates. *)
+let referenced_vars_of_prog =
+  procs
+  %> Iter.flat_map
+       (snd %> Procedure.iter_blocks
+       %> Iter.flat_map (fun (_, b) ->
+           Iter.append (Block.read_vars_iter b) (Block.assigned_vars_iter b)))
+  %> Iter.filter Var.is_global
+
 let pretty_to_chan chan (p : t) =
   let p = prog_pretty p in
   flush chan;

lib/lang/program.mli

@@ -108,6 +108,10 @@ val decl_global :
 
 val decl_typ : ?attrib:'a Types.StringMap.t -> t -> Types.t -> t
 
+val referenced_vars_of_prog : t -> Var.t Iter.t
+(** Iterates over global variables in the given program, including both read and
+    assigned variables. Order is unspecified and may have duplicates. *)
+
 val create_single_proc :
   ?name:string -> unit -> t * (Var.t, Expr.BasilExpr.t) Procedure.t
 

lib/transforms/aslp/aslp.ml

@@ -37,18 +37,15 @@ let lift_code_block (module I : Bincaml_ibi.IBI) ~address =
 
 (** {1 Interfacing with Bincaml IR} *)
 
-and error_attrib_key = ".error"
-
 (** Extracts the opcode, address, and attribute from the given Bincaml
     statement, if it is an {!Lang.Stmt.Intrinsic.Aarch64Eval} intrinsic call.
     Otherwise, returns [None].
 
     Raises an exception if an {!Lang.Stmt.Intrinsic.Aarch64Eval} intrinsic call
     has an unexpected structure. *)
-let aarch64_intrin_of_stmt ?(include_failed = false) ?default_address :
+let aarch64_intrin_of_stmt ?default_address :
     Program.stmt -> (Bitvec.t * Bitvec.t * Attrib.attrib_map) option = function
-  | Stmt.Instr_IntrinCall { attrib; lhs; name = Aarch64Eval; args }
-    when include_failed || not (StringMap.mem error_attrib_key attrib) -> (
+  | Stmt.Instr_IntrinCall { attrib; lhs; name = Aarch64Eval; args } -> (
       let args =
         match (args, default_address) with
         | [ x ], Some default -> [ x; Expr.BasilExpr.bvconst default ]
@@ -69,23 +66,13 @@ let aarch64_intrin_of_stmt ?(include_failed = false) ?default_address :
   | _ -> None
 
 (** Inverse of {!aarch64_intrin_of_stmt}. *)
-let stmt_of_aarch64_intrin :
+let stmt_of_aarch64_intrin ?error :
     Bitvec.t * Bitvec.t * Attrib.attrib_map -> Program.stmt =
  fun (opcode, address, attrib) ->
-  let args = List.map Expr.BasilExpr.bvconst [ opcode; address ] in
+  let attrib = Option.fold (Fun.flip (StringMap.add ".error")) attrib error in
+  let args = Expr.BasilExpr.[ bvconst opcode; bvconst address ] in
   Stmt.Instr_IntrinCall { attrib; lhs = []; name = Aarch64Eval; args }
 
-(** Extracts the next Aarch64 intrinsic from the given list of statements,
-    returning [Some (before, intrin, after)] if there exists an intrinsic. *)
-let next_aarch64_stmt stmts =
-  let open CCOption.Infix in
-  let before, rest =
-    CCList.take_drop_while (Option.is_none % aarch64_intrin_of_stmt) stmts
-  in
-  let* hd, after = Aslp_util_internal.uncons rest in
-  let* intrin = aarch64_intrin_of_stmt hd in
-  Some (before, intrin, after)
-
 (** Returns the Bincaml global variable representing heap memory. *)
 let aarch64_mem_of_prog prog =
   Program.get_decl_by_name "$mem" prog |> function
@@ -109,69 +96,48 @@ let insert_one_diamond ~proc dia =
     with_ids |> Diamond.iter_backwards
     |> Iter.fold
          (fun proc (id, successors, st) ->
-           let assume = Aslp_state.assume_of_aslp_block st in
-           let stmts = Option.to_list assume @ CCVector.to_list st.stmts in
+           let stmts = Aslp_state.stmts_of_aslp_block st in
            Procedure.add_block proc id ~stmts ~successors ())
          proc
   and (first, _, _), (last, _, _) = Diamond.(first with_ids, last with_ids) in
 
   (first, last, proc)
 
-(** Transforms the next {!Lang.Stmt.Intrinsic.Aarch64Eval} intrinsic within the
-    given block ID within the given procedure.
-
-    The first intrinsic appearing within the block's statement list, if any,
-    will be transformed. The block will be split at this point. Any statements
-    after the intrinsic (and any successor edges) will be moved to the last
-    block of the freshly-inserted ASLp blocks.
-
-    If a change happened, [Some] will be returned with the updated procedure,
-    along with the ID of the last block of the ASLp output (containing the
-    suffix of the original block's statements). If no intrinsic exists, [None]
-    will be returned.
+(** Maps the given statement through the ASLp lifter transform, if applicable.
+    If the statement is a {!Lang.Stmt.Intrinsic.Aarch64Eval} intrinsic, then
+    returns [Left] with the lifted blocks. Otherwise, or if an error occurs,
+    returns [Right] with the original statement.
 
     If an error occurs while lifting through ASLp, an error attribute will be
     attached to the intrinsic (and it is excluded from subsequent calls). *)
-let transform_one_stmt (module I : Bincaml_ibi.IBI) ~proc bid =
-  let b = Procedure.get_block proc bid |> Option.get_exn_or "block not found" in
-
-  match next_aarch64_stmt (Vector.to_list b.stmts) with
-  | None -> None
-  | Some (before, (opcode, address, intrin_attrib), after) -> (
+let map_one_stmt (module I : Bincaml_ibi.IBI) ~proc stmt :
+    _ Procedure.mapped_stmt =
+  match aarch64_intrin_of_stmt stmt with
+  | None -> `Stmts [ stmt ]
+  | Some (opcode, address, attrib) -> (
       match lift_opcode (module I) ~address opcode with
       | diamond ->
           let aslp_first, aslp_last, proc = insert_one_diamond ~proc diamond in
-
-          let proc =
-            proc
-            |> Procedure.modify_block' ~id:bid ~f:(fun b ->
-                { b with stmts = Vector.of_list before })
-            |> Procedure.modify_block' ~id:aslp_first ~f:(fun b ->
-                { b with attrib = intrin_attrib })
-            |> Procedure.modify_block' ~id:aslp_last ~f:(fun b ->
-                Block.fmap_stmts_copy (Fun.flip CCVector.append_list after) b)
-            |> Procedure.transplant_outgoing_edges ~from:bid ~to_:aslp_last
-            |> Procedure.add_goto ~from:bid ~targets:[ aslp_first ]
-          in
-          Some (proc, aslp_last)
-      | exception exn ->
-          let exn = `String (Printexc.to_string exn) in
-          let attrib = StringMap.add error_attrib_key exn intrin_attrib in
-          let intrin_stmt = stmt_of_aarch64_intrin (opcode, address, attrib) in
-          let stmts = Vector.of_list (before @ (intrin_stmt :: after)) in
-          Some (Procedure.modify_block proc bid (fun b -> { b with stmts }), bid)
-      )
+          `Graph
+            ( aslp_first,
+              aslp_last,
+              Procedure.modify_block proc aslp_first (fun b ->
+                  { b with attrib }) )
+      | exception error ->
+          let error = `String (Printexc.to_string error) in
+          let intrin = (opcode, address, attrib) in
+          `Stmts [ stmt_of_aarch64_intrin ~error intrin ])
 
 (** Transforms the {!Lang.Stmt.Intrinsic.Aarch64Eval} intrinsics within the
     given block ID within the given procedure.
 
     Inserts control-flow edges between successive instructions within the block,
     and emits an assertion for the ITE expression representing the final [PC]
     value. *)
-let rec transform_block (module I : Bincaml_ibi.IBI) ~proc bid =
-  match transform_one_stmt (module I) ~proc bid with
-  | Some (proc, bid) -> (transform_block [@tailcall]) (module I) ~proc bid
-  | None -> proc
+let transform_block (module I : Bincaml_ibi.IBI) ~proc bid =
+  let f = map_one_stmt (module I) in
+  let _, _, proc = Procedure.cfg_concatmap_block bid ~f proc in
+  proc
 
 (** Transforms the {!Lang.Stmt.Intrinsic.Aarch64Eval} intrinsics of all blocks
     within the given procedure. *)
@@ -188,8 +154,7 @@ let add_aarch64_global_declarations ?(add_all = false) prog =
     if add_all then Fun.const true
     else
       Fun.flip VarSet.mem
-        (Aslp_util_internal.referenced_vars_of_prog prog
-        |> Iter.to_set (module VarSet))
+        (Program.referenced_vars_of_prog prog |> Iter.to_set (module VarSet))
   in
 
   Lazy.force Aslp_lexpr.global_vars
@@ -227,7 +192,7 @@ let apply_stmt_addresses_from_block (block : _ Block.t) =
 
   let apply_one address stmt =
     let address' = Bitvec.(add address (of_int ~size:64 4)) in
-    match aarch64_intrin_of_stmt ~include_failed:true ~default_address stmt with
+    match aarch64_intrin_of_stmt ~default_address stmt with
     | Some (opcode, a, attr) when Bitvec.equal a default_address ->
         (address', stmt_of_aarch64_intrin (opcode, address, attr))
     | Some _ -> (address', stmt) (* increment address *)

lib/transforms/aslp/aslp_state.ml

@@ -229,6 +229,9 @@ let assume_of_aslp_block = function
   | { assume = body; _ } ->
       Some (Stmt.Instr_Assume { attrib = Attrib.empty; body; branch = false })
 
+let stmts_of_aslp_block ({ stmts } as st) =
+  Option.to_list (assume_of_aslp_block st) @ CCVector.to_list stmts
+
 (** {1 Formatters} *)
 
 let show_aslp_block = show_aslp_block