Source file recover_match.ml

open Wax_lang
open Ast

(* Recover a [match] from the nested type-test ladder that
   {!Ast_utils.lower_match} emits (and that hand-written GC code uses): a stack
   of blocks, the innermost holding a threaded [br_on_cast]/[br_on_null] chain on
   the scrutinee, each test branching out to its arm's block, with the arm body
   after that block and a void [escape] block past them all; the [default]
   follows the [escape] block as trailing code.

     'escape: do {
       'Lₙ₋₁: do { … 'L₀: do { _ = br_on_cast 'L₀ … (br_on_cast 'L₁ … (… v)); br 'escape }
                        <bind/​drop block 'L₀>; arm₀ } … <bind/​drop>; armₙ₋₁ }
       …trailing default…
     ⇒ match v { …armᵢ… _ => { …trailing… } }

   Because every arm body leaves the [match] (diverges), absorbing the trailing
   statements into the default preserves semantics — they only run on the
   no-match (fall-through) path. Bound cast arms surface their binding as a
   [local.set] ([Set] / a fused [let]); {!Sink_let} sinks a local used across
   several arms to their common-ancestor ladder block, where it surfaces as a
   leading declaration. The fold hoists those out before the [match] (a
   declaration an arm rebinds is dropped instead, since re-lowering reintroduces
   it). A Wax [match] thus round-trips through the binary. Meant to run on
   {!Sink_let.module_} output.

   We also recover a [match] from the *flat* [br_on_cast_fail] chain that
   hand-written GC code more often uses — one discarded block per arm rather than
   the nested ladder; see {!collect_arms}. That round trip is not byte-for-byte
   (re-lowering emits the ladder), only semantically faithful. *)

let is_block i = match i.desc with Block _ -> true | _ -> false

let is_chain i =
  match i.desc with Br_on_cast _ | Br_on_null _ -> true | _ -> false

(* Parse the threaded chain (innermost operand the scrutinee). Returns the tests
   in source order — innermost test ([L₀]) first — each as its label and the
   pattern's cast target ([`Cast]) or null ([`Null]); the binding (if any) comes
   from the wrapping block's consume, not the chain. *)
let rec chain_tests e =
  match e.desc with
  | Br_on_cast (l, rt, operand) ->
      let tests, scrut = chain_tests operand in
      (tests @ [ (l, `Cast rt) ], scrut)
  | Br_on_null (l, operand) ->
      let tests, scrut = chain_tests operand in
      (tests @ [ (l, `Null) ], scrut)
  | _ -> ([], e)

(* Split off leading bare local declarations [let x;]. {!Sink_let} sinks a local
   used across several arms to their common-ancestor ladder block, where it shows
   up here as a leading declaration; the fold hoists those out before the match
   (they are *not* arm bindings — those are fused into the consume — so dropping
   them would unbind the local). Returns the declarations and the rest. *)
let split_decls stmts =
  let rec aux acc = function
    | ({ desc = Let ([ (Some _, _) ], None); _ } as d) :: rest ->
        aux (d :: acc) rest
    | rest -> (List.rev acc, rest)
  in
  aux [] stmts

(* Strip a ladder block's consume of its inner block, returning any leading
   declarations to hoist, the binding ([Some (x, set)] for a bound cast arm),
   that inner block, and the trailing arm body. A [null] arm consumes nothing (a
   bare block), an unbound cast drops the block (an anonymous [Let],
   [_ = block]), a bound cast binds it (a fused [let], [set = false], or a [Set]
   of a local declared elsewhere, [set = true]). *)
let consume_step stmts =
  let decls, stmts = split_decls stmts in
  match stmts with
  | { desc = Let ([ (Some x, _) ], Some inner); _ } :: body when is_block inner
    ->
      Some (decls, Some (x, false), inner, body)
  | { desc = Set (x, _, inner); _ } :: body when is_block inner ->
      Some (decls, Some (x, true), inner, body)
  | { desc = Let ([ (None, _) ], Some inner); _ } :: body when is_block inner ->
      Some (decls, None, inner, body)
  (* The bare-block arm is the NULL arm, which consumes nothing — so its block
     must be void. Without the guard, a RESULT-carrying arm whose binding was
     separated from the block (a [nop] between the [end] and the [local.set],
     say) was folded as if it were a null arm: the binding's value vanished
     into the rebuilt [match] and the emitted Wax no longer typed, on a module
     the validator accepts (a recover-shapes near-miss finding). Left alone,
     the unrecovered print of the same shape is fine — the trailing hole
     reconnects to the block's value through the statement. *)
  | ({ desc = Block { typ; _ }; _ } as inner) :: body
    when typ.params = [||] && typ.results = [||] ->
      Some (decls, None, inner, body)
  | _ -> None

(* Descend the ladder from block [blk]. Returns the wrapper levels (outer→inner,
   one per arm: the block label, the arm's binding, and its body), the innermost
   block's label, the chain, the escape label, and the declarations to hoist. *)
let rec descend blk =
  match blk.desc with
  | Block { label = Some lbl; block = { desc = body; _ }; _ } -> (
      let decls0, body = split_decls body in
      match body with
      | [
       { desc = Let ([ (None, _) ], Some chain); _ };
       { desc = Br (escape, None); _ };
      ]
        when is_chain chain ->
          Some ([], lbl, chain, escape, decls0)
      | _ -> (
          match consume_step body with
          | Some (decls1, binding, inner, arm_body) -> (
              match descend inner with
              | Some (levels, inner_lbl, chain, escape, decls) ->
                  Some
                    ( (lbl, binding, arm_body) :: levels,
                      inner_lbl,
                      chain,
                      escape,
                      decls0 @ decls1 @ decls )
              | None -> None)
          | None -> None))
  | _ -> None

(* --- Flat [br_on_cast_fail] chain --------------------------------------- *)

(* Hand-written GC code (and pre-ladder Wax output) takes a value apart with a
   flat run of discarded blocks rather than the nested ladder {!descend} folds:

     _ = 'L: do S { let x = br_on_cast_fail 'L &T v; body }   (a bound cast arm)
     _ = 'L: do S { _ = br_on_cast_fail 'L &T v; body }       (an unbound cast arm)
     _ = 'L: do S { br_on_non_null 'L v; body }               (a null arm)

   Each block re-reads the same scrutinee [v] and, on a failed test, branches to
   its own label (forwarding [v], type [S]) and is dropped, falling through to
   the next block; on success the (optionally bound) narrowed value falls into
   [body], which diverges. So such a run is a [match] on [v], the trailing
   statements its default. Re-lowering a recovered match emits the nested ladder,
   not this flat chain, so the round trip is not byte-for-byte — but it is
   semantically equivalent (the scrutinee is side-effect-free, see
   {!same_scrut}). *)

(* Structural equality of scrutinee expressions, ignoring source locations: the
   side-effect-free forms a re-read scrutinee may take. Anything else compares
   unequal, so the run simply does not fold. *)
let rec same_scrut a b =
  match (a.desc, b.desc) with
  | Get x, Get y -> x.desc = y.desc
  | Null, Null -> true
  | Int s, Int t -> s = t
  | NonNull e, NonNull f -> same_scrut e f
  | Cast (e, s), Cast (f, t) -> s = t && same_scrut e f
  | Test (e, s), Test (f, t) -> s = t && same_scrut e f
  | StructGet (e, x), StructGet (f, y) -> x.desc = y.desc && same_scrut e f
  | ArrayGet (e, i), ArrayGet (f, j) -> same_scrut e f && same_scrut i j
  | _ -> false

(* Whether control cannot fall off the end of [body] — a conservative,
   syntactic check (the last statement is a clear terminator, or an [if]/[match]
   whose every branch diverges). A flat-chain arm is a genuine [match] arm only
   when its success path leaves the [match]; folding a non-diverging body (the
   block's [do S] result is produced and dropped instead) would be wrong. *)
let rec diverges_instr i =
  match i.desc with
  | Return _ | Br _ | Br_table _ | Unreachable | Throw _ | ThrowRef _
  | TailCall _ ->
      true
  | If { if_block; else_block = Some else_block; _ } ->
      diverges_list if_block.desc && diverges_list else_block.desc
  | Match { arms; default; _ } ->
      List.for_all
        (fun (_, (b : (_ instr list, _) Ast.annotated)) -> diverges_list b.desc)
        arms
      && diverges_list default.desc
  | Loop { block; _ } ->
      (* A loop whose body always branches (back to the loop or out) never falls
         through to the statement after it. *)
      diverges_list block.desc
  | _ -> false

and diverges_list l =
  match List.rev l with [] -> false | last :: _ -> diverges_instr last

(* Whether [i] branches to a label named [name] bound outside it. A nested
   construct rebinding [name] shadows it within its body (nested ladders reuse
   the same label names), so only the parts evaluated outside that body are
   searched there. *)
let rec branches_out name (i : location instr) =
  List.exists
    (fun (l : label) -> l.desc = name)
    (Recover_trycatch.target_labels i.desc)
  ||
  match i.desc with
  | Block { label = Some l; _ }
  | Loop { label = Some l; _ }
  | TryTable { label = Some l; _ }
  | Try { label = Some l; _ }
  | TryCatch { label = Some l; _ }
    when l.desc = name ->
      false
  | If { label = Some l; cond; _ } when l.desc = name -> branches_out name cond
  | While { label = Some l; cond; step; _ } when l.desc = name -> (
      branches_out name cond
      || match step with Some s -> branches_out name s | None -> false)
  | _ -> branches_out_list name (Ast_utils.sub_instrs i)

and branches_out_list name l = List.exists (branches_out name) l

(* Recognise one flat-chain arm block. Returns its pattern, scrutinee, body, and
   whether a bound cast carries its binding itself (a fused [let x = …],
   [`Fused]) or names a local declared just before the block ([`Decl x], which
   the fold drops). The body must diverge (leave the [match]), and neither it
   nor the scrutinee may branch to the block's own label: the [match] replaces the block, so such a
   branch (a second [br_on_cast_fail 'L] in the same block, say) would be left
   unbound. *)
let arm_block stmt =
  match stmt.desc with
  | Let
      ( [ (None, _) ],
        Some
          {
            desc =
              Block
                { label = Some self; typ; block = { desc = test :: body; _ } };
            _;
          } )
    when typ.params = [||]
         && Array.length typ.results = 1
         && diverges_list body
         && not (branches_out_list self.desc body) -> (
      match
        match test.desc with
        | Let
            ([ (Some x, _) ], Some { desc = Br_on_cast_fail (l, rt, scrut); _ })
          when l.desc = self.desc ->
            Some (MatchCast (Some x, rt), scrut, body, `Fused)
        | Set (x, _, { desc = Br_on_cast_fail (l, rt, scrut); _ })
          when l.desc = self.desc ->
            Some (MatchCast (Some x, rt), scrut, body, `Decl x)
        | Let ([ (None, _) ], Some { desc = Br_on_cast_fail (l, rt, scrut); _ })
          when l.desc = self.desc ->
            Some (MatchCast (None, rt), scrut, body, `Fused)
        | Br_on_non_null (l, scrut) when l.desc = self.desc ->
            Some (MatchNull, scrut, body, `Fused)
        | _ -> None
      with
      (* The scrutinee moves out of the block too. *)
      | Some (_, scrut, _, _) when branches_out self.desc scrut -> None
      | arm -> arm)
  | _ -> None

let compat scrut s =
  match scrut with None -> true | Some s0 -> same_scrut s0 s

(* Collect a maximal run of flat-chain arms from the front of [stmts], threading
   the shared scrutinee. Returns the arms, the shared scrutinee, bare local
   declarations to hoist before the match, and the remaining (default)
   statements.

   A bare local declaration [let x;] between arms is either the binding for the
   next [`Decl]-form arm (dropped — the recovered arm re-declares it) or an
   unrelated local that {!Sink_let} floated into the run (hoisted before the
   match, where it stays in scope for every arm; this only fires when more arms
   follow, so a trailing declaration stays in the default). *)
let rec collect_arms scrut stmts =
  let take pat body s rest =
    let scrut = match scrut with None -> Some s | some -> some in
    let arms, scrut, hoisted, rest = collect_arms scrut rest in
    ((pat, body) :: arms, scrut, hoisted, rest)
  in
  match stmts with
  | ({ desc = Let ([ (Some x, _) ], None); _ } as decl) :: rest -> (
      match rest with
      | blk :: rest'
        when match arm_block blk with
             | Some (_, s, _, `Decl y) -> y.desc = x.desc && compat scrut s
             | _ -> false ->
          let pat, s, body =
            match arm_block blk with
            | Some (pat, s, body, _) -> (pat, s, body)
            | None -> assert false
          in
          (* The arm's binding replaces the local [x], but only within its
             body. A use in the scrutinee, a later arm or the default would be
             left reading a local the arm no longer writes (reached again past
             a loop back-edge, say), so the run stops before this arm. *)
          if Sink_let.occurs_list x.desc (s :: rest') then ([], scrut, [], stmts)
          else take pat body s rest'
      | _ ->
          let arms, scrut, hoisted, trailing = collect_arms scrut rest in
          if arms = [] then ([], scrut, [], stmts)
          else (arms, scrut, decl :: hoisted, trailing))
  | blk :: rest -> (
      match arm_block blk with
      | Some (pat, s, body, `Fused) when compat scrut s -> take pat body s rest
      | _ -> ([], scrut, [], stmts))
  | [] -> ([], scrut, [], stmts)

let rec rewrite_instr ~faithful (i : location instr) : location instr =
  let d = rewrite_desc ~faithful i.desc in
  if d == i.desc then i else { i with desc = d }

(* Fold the [escape] block [i] (and the trailing statements after it, which
   become the default) into a [match]. *)
and try_fold ~faithful (i : location instr) (trailing : location instr list) :
    (location instr list * location instr) option =
  match descend i with
  | None -> None
  | Some (levels, inner_lbl, chain, escape, decls) ->
      let tests, scrut = chain_tests chain in
      let n = List.length tests in
      (* Each test branches to its arm's block; descending nests them
         innermost→outermost as the chain orders them, with the escape block
         outermost. Checking that pins the fold to a genuine ladder. *)
      let block_labels =
        inner_lbl :: List.rev_map (fun (l, _, _) -> l) levels
      in
      let rec take k = function
        | x :: r when k > 0 -> x :: take (k - 1) r
        | _ -> []
      in
      let label_names = List.map (fun (l : label) -> l.desc) block_labels in
      let distinct =
        List.length (List.sort_uniq compare label_names)
        = List.length label_names
      in
      (* An arm body is moved out of the ladder blocks into the [match], which
         binds none of their labels, so a body still branching to one (a
         [br_if 'Lᵢ] to a later arm's consume, a [br 'escape] to the default)
         would be left dangling. The scrutinee moves out of them too. *)
      let bodies_free =
        not
          (List.exists
             (fun body ->
               List.exists (fun name -> branches_out_list name body) label_names)
             ([ scrut ] :: List.map (fun (_, _, body) -> body) levels))
      in
      let chain_ok =
        List.length levels = n
        && List.map (fun ((l : label), _) -> l.desc) tests = take n label_names
        &&
        match List.rev block_labels with
        | last :: _ -> last.desc = escape.desc
        | [] -> false
      in
      if n < 1 || (not distinct) || (not chain_ok) || not bodies_free then None
      else
        let arm (_, pat_kind) (_, binding, body) =
          let pat =
            match (pat_kind, binding) with
            | `Cast rt, Some (x, _) -> Some (MatchCast (Some x, rt))
            | `Cast rt, None -> Some (MatchCast (None, rt))
            | `Null, None -> Some MatchNull
            | `Null, Some _ -> None
          in
          Option.map
            (fun pat -> (pat, no_loc (rewrite_list ~faithful body)))
            pat
        in
        let arms = List.map2 arm tests (List.rev levels) in
        if List.exists Option.is_none arms then None
        else
          let arms = List.filter_map Fun.id arms in
          (* A local an arm's binding replaces must not be used outside the
             arms rebinding it (in the scrutinee, another arm or the default):
             that use would read a local the arm no longer writes. So such an
             arm is folded only when the local is declared within the ladder
             (so no use lies beyond it) and has no use outside. *)
          let used_outside x =
            Sink_let.occurs_list x (scrut :: trailing)
            || List.exists
                 (fun (p, (b : (_ instr list, _) Ast.annotated)) ->
                   (match p with
                     | MatchCast (Some y, _) -> y.desc <> x
                     | _ -> true)
                   && Sink_let.occurs_list x b.desc)
                 arms
          in
          let declared x =
            List.exists
              (fun d ->
                match d.desc with
                | Let ([ (Some y, _) ], None) -> y.desc = x
                | _ -> false)
              decls
          in
          let replaced_ok =
            List.for_all
              (fun (_, binding, _) ->
                match binding with
                | Some ((x : ident), true) ->
                    declared x.desc && not (used_outside x.desc)
                | _ -> true)
              levels
          in
          (* A declaration whose name an arm rebinds is redundant (re-lowering
             reintroduces it) unless the local is still used outside the arms
             rebinding it (then it is a different local that a fused binding
             only shadows). The rest are genuine locals to hoist. *)
          let hoisted =
            List.filter
              (fun d ->
                match d.desc with
                | Let ([ (Some x, _) ], None) ->
                    (not
                       (List.exists
                          (fun (p, _) ->
                            match p with
                            | MatchCast (Some y, _) -> y.desc = x.desc
                            | _ -> false)
                          arms))
                    || used_outside x.desc
                | _ -> true)
              decls
          in
          if not replaced_ok then None
          else
            Some
              ( hoisted,
                {
                  i with
                  desc =
                    Match
                      {
                        scrutinee = rewrite_instr ~faithful scrut;
                        arms;
                        default = no_loc (rewrite_list ~faithful trailing);
                      };
                } )

and rewrite_list ~faithful stmts =
  match stmts with
  | [] -> []
  | i :: rest -> (
      (* First the nested ladder (the shape {!Ast_utils.lower_match} emits),
         then a flat [br_on_cast_fail] chain — folded even for a single arm (a
         lone downcast-or-branch reads as a one-arm [match]). The flat-chain arm
         re-lowers to the nested ladder rather than the original flat chain, so
         [--faithful] ([faithful]) skips it while keeping the exact-inverse
         nested-ladder fold. *)
      match try_fold ~faithful i rest with
      | Some (hoisted, m) -> List.map (rewrite_instr ~faithful) hoisted @ [ m ]
      | None -> (
          match
            if faithful then ([], None, [], []) else collect_arms None stmts
          with
          | (_ :: _ as arms), Some scrut, hoisted, trailing ->
              List.map (rewrite_instr ~faithful) hoisted
              @ [
                  {
                    i with
                    desc =
                      Match
                        {
                          scrutinee = rewrite_instr ~faithful scrut;
                          arms =
                            List.map
                              (fun (p, b) ->
                                (p, no_loc (rewrite_list ~faithful b)))
                              arms;
                          default = no_loc (rewrite_list ~faithful trailing);
                        };
                  };
                ]
          | _ ->
              let i' = rewrite_instr ~faithful i
              and rest' = rewrite_list ~faithful rest in
              if i' == i && rest' == rest then stmts else i' :: rest'))

and rewrite_desc ~faithful (desc : location instr_desc) : location instr_desc =
  (* Purely structural recursion; the folding lives in [try_fold]/[rewrite_list].
     Share-preserving, so an untouched subtree is not rebuilt. *)
  Ast_utils.map_desc ~instr:(rewrite_instr ~faithful)
    ~block:(rewrite_list ~faithful) desc

let rec field_desc ~faithful (f : location modulefield) =
  let map_fields =
    List.map (fun (a : (location modulefield, _) Ast.annotated) ->
        { a with desc = field_desc ~faithful a.desc })
  in
  match f with
  | Func ({ body = label, instrs; _ } as r) ->
      Func { r with body = (label, rewrite_list ~faithful instrs) }
  | Conditional ({ then_fields; else_fields; _ } as r) ->
      Conditional
        {
          r with
          then_fields = { then_fields with desc = map_fields then_fields.desc };
          else_fields =
            Option.map
              (fun (b :
                     ( (location modulefield, location) Ast.annotated list,
                       location )
                     Ast.annotated) -> { b with desc = map_fields b.desc })
              else_fields;
        }
  | ( Type _ | Type_alias _ | Module_annotation _ | Import _ | Import_group _
    | Global _ | Tag _ | Memory _ | Data _ | Table _ | Elem _ ) as f ->
      f

let module_ ?(faithful = false) (m : location module_) : location module_ =
  List.map
    (fun (a : (location modulefield, location) Ast.annotated) ->
      { a with desc = field_desc ~faithful a.desc })
    m