rewrite.ml — Coverage report

(* SPDX-License-Identifier: AGPL-3.0-or-later *)
(* Copyright © 2021-2024 OCamlPro *)
(* Written by the Owi programmers *)

open Syntax

let rewrite_heap_type (assigned : Assigned.t) (ht : Text.heap_type) :
  Binary.heap_type Result.t =
  match ht with
  | Text.TypeUse id ->
    let* id = Assigned.find_type assigned id in
    Ok (Binary.TypeUse id)
  | Any_ht -> Ok Any_ht
  | None_ht -> Ok None_ht
  | Func_ht -> Ok Func_ht
  | NoFunc_ht -> Ok NoFunc_ht
  | Exn_ht -> Ok Exn_ht
  | NoExn_ht -> Ok NoExn_ht
  | Extern_ht -> Ok Extern_ht
  | NoExtern_ht -> Ok NoExtern_ht

let rewrite_val_type (assigned : Assigned.t) (vt : Text.val_type) :
  Binary.val_type Result.t =
  match vt with
  | Text.Ref_type (n, ht) ->
    let* ht = rewrite_heap_type assigned ht in
    Ok (Binary.Ref_type (n, ht))
  | Num_type nt -> Ok (Num_type nt)

let rewrite_func_type (assigned : Assigned.t) ((param, res) : Text.func_type) :
  Binary.func_type Result.t =
  let* param' =
    list_map
      (fun (n, vt) ->
        let* vt = rewrite_val_type assigned vt in
        Ok (n, vt) )
      param
  in
  let* res' = list_map (rewrite_val_type assigned) res in
  Ok (param', res')

let rewrite_block_type (assigned : Assigned.t) (block_type : Text.block_type) :
  (Text.param list * Binary.block_type) Result.t =
  match block_type with
  | Bt_ind id ->
    let* idx = Assigned.find_type assigned id in
    let* ((params, _) as t) =
      match Assigned.get_type assigned idx with
      | None -> Error (`Unknown_type id)
      | Some v -> Ok v
    in
    let* t = rewrite_func_type assigned t in
    Ok (params, Binary.Bt_raw (Some idx, t))
  | Bt_raw (_, ((params, _) as func_type)) ->
    let idx = Assigned.find_raw_type assigned func_type in
    let* func_type = rewrite_func_type assigned func_type in
    Ok (params, Binary.Bt_raw (Some idx, func_type))

let rewrite_memarg ({ offset; align } : Text.memarg) : Binary.memarg Result.t =
  let* offset =
    match offset with
    | None -> Ok Int64.zero
    | Some offset -> (
      match Int64.of_string offset with
      | None -> Error (`Msg "offset")
      | Some n -> Ok n )
  in
  let+ align =
    match align with
    | None -> Ok Int32.zero
    | Some align -> (
      match Int32.of_string align with
      | None -> Error `Alignment_too_large
      | Some n ->
        if Int32.eq n 0l || Int32.(ne (logand n (sub n 1l)) 0l) then
          Error `Alignment_too_large
        else Ok (Int32.div n 2l) )
  in
  Binary.{ offset; align }

let rewrite_expr (assigned : Assigned.t) (locals : Text.param list)
  (iexpr : Text.expr Annotated.t) : Binary.expr Annotated.t Result.t =
  (* block_ids handling *)
  let block_id_to_raw (loop_count, block_ids) id =
    let* id =
      match id with
      | Text.Text id -> begin
        match
          List.find_index
            (function Some id' -> String.equal id id' | None -> false)
            block_ids
        with
        | None -> Error (`Unknown_label (Text.Text id))
        | Some id -> Ok id
      end
      | Raw id -> Ok id
    in
    (* this is > and not >= because you can `br 0` without any block to target the function *)
    if id > List.length block_ids + loop_count then
      Error (`Unknown_label (Text.Raw id))
    else Ok id
  in

  (* block_types handling *)
  let block_ty_opt_rewrite = function
    | Some bt ->
      let+ _, bt = rewrite_block_type assigned bt in
      Some bt
    | None -> Ok None
  in

  let seen_locals = Hashtbl.create 64 in

  (* Fill locals *)
  let* (_ : int) =
    list_fold_left
      (fun next_free_int ((name, _type) : Text.param) ->
        match name with
        | None -> Ok (next_free_int + 1)
        | Some name ->
          if Hashtbl.mem seen_locals name then Error (`Duplicate_local name)
          else begin
            Hashtbl.add seen_locals name next_free_int;
            Ok (next_free_int + 1)
          end )
      0 locals
  in

  let find_local : Text.indice -> Binary.indice = function
    | Raw i -> i
    | Text name -> (
      match Hashtbl.find_opt seen_locals name with
      | None -> assert false
      | Some id -> id )
  in

  let rec body (loop_count, block_ids) (instr : Text.instr Annotated.t) :
    Binary.instr Result.t =
    match instr.Annotated.raw with
    | Br_table (ids, id) ->
      let block_id_to_raw = block_id_to_raw (loop_count, block_ids) in
      let* ids = array_map block_id_to_raw ids in
      let+ id = block_id_to_raw id in
      Binary.Br_table (ids, id)
    | Br_if id ->
      let+ id = block_id_to_raw (loop_count, block_ids) id in
      Binary.Br_if id
    | Br id ->
      let+ id = block_id_to_raw (loop_count, block_ids) id in
      Binary.Br id
    | Br_on_null id ->
      let+ id = block_id_to_raw (loop_count, block_ids) id in
      Binary.Br_on_null id
    | Br_on_non_null id ->
      let+ id = block_id_to_raw (loop_count, block_ids) id in
      Binary.Br_on_non_null id
    | Call id ->
      let+ id = Assigned.find_func assigned id in
      Binary.Call id
    | Return_call id ->
      let+ id = Assigned.find_func assigned id in
      Binary.Return_call id
    | Local_set id ->
      let id = find_local id in
      Ok (Binary.Local_set id)
    | Local_get id ->
      let id = find_local id in
      Ok (Binary.Local_get id)
    | Local_tee id ->
      let id = find_local id in
      Ok (Binary.Local_tee id)
    | If_else (id, bt, e1, e2) ->
      let* bt = block_ty_opt_rewrite bt in
      let block_ids = id :: block_ids in
      let* e1 = expr e1 (loop_count, block_ids) in
      let+ e2 = expr e2 (loop_count, block_ids) in
      Binary.If_else (id, bt, e1, e2)
    | Loop (id, bt, e) ->
      let* bt = block_ty_opt_rewrite bt in
      let+ e = expr e (loop_count + 1, id :: block_ids) in
      Binary.Loop (id, bt, e)
    | Block (id, bt, e) ->
      let* bt = block_ty_opt_rewrite bt in
      let+ e = expr e (loop_count, id :: block_ids) in
      Binary.Block (id, bt, e)
    | Call_indirect (tbl_i, bt) ->
      let* tbl_i = Assigned.find_table assigned tbl_i in
      let+ _, bt = rewrite_block_type assigned bt in
      Binary.Call_indirect (tbl_i, bt)
    | Return_call_indirect (tbl_i, bt) ->
      let* tbl_i = Assigned.find_table assigned tbl_i in
      let+ _, bt = rewrite_block_type assigned bt in
      Binary.Return_call_indirect (tbl_i, bt)
    | Call_ref t ->
      let+ t = Assigned.find_type assigned t in
      Binary.Call_ref t
    | Return_call_ref bt ->
      let+ _, bt = rewrite_block_type assigned bt in
      Binary.Return_call_ref bt
    | Global_set id ->
      let+ idx = Assigned.find_global assigned id in
      Binary.Global_set idx
    | Global_get id ->
      let* idx = Assigned.find_global assigned id in
      Ok (Binary.Global_get idx)
    | Ref_func id ->
      let+ id = Assigned.find_func assigned id in
      Binary.Ref_func id
    | Table_size id ->
      let+ id = Assigned.find_table assigned id in
      Binary.Table_size id
    | Table_get id ->
      let+ id = Assigned.find_table assigned id in
      Binary.Table_get id
    | Table_set id ->
      let+ id = Assigned.find_table assigned id in
      Binary.Table_set id
    | Table_grow id ->
      let+ id = Assigned.find_table assigned id in
      Binary.Table_grow id
    | Table_init (i, i') ->
      let* table = Assigned.find_table assigned i in
      let+ elem = Assigned.find_elem assigned i' in
      Binary.Table_init (table, elem)
    | Table_fill id ->
      let+ id = Assigned.find_table assigned id in
      Binary.Table_fill id
    | Table_copy (i, i') ->
      let* table = Assigned.find_table assigned i in
      let+ table' = Assigned.find_table assigned i' in
      Binary.Table_copy (table, table')
    | Memory_init (memidx, dataidx) ->
      let* memidx = Assigned.find_memory assigned memidx in
      let+ dataidx = Assigned.find_data assigned dataidx in
      Binary.Memory_init (memidx, dataidx)
    | Data_drop id ->
      let+ id = Assigned.find_data assigned id in
      Binary.Data_drop id
    | Elem_drop id ->
      let+ id = Assigned.find_elem assigned id in
      Binary.Elem_drop id
    | Select typ -> begin
      match typ with
      | None -> Ok (Binary.Select None)
      | Some [ t ] ->
        let+ t = rewrite_val_type assigned t in
        Binary.Select (Some [ t ])
      | Some [] | Some (_ :: _ :: _) -> Error `Invalid_result_arity
    end
    | I_unop (nn, op) -> Ok (Binary.I_unop (nn, op))
    | I_binop (nn, op) -> Ok (I_binop (nn, op))
    | I_testop (nn, op) -> Ok (Binary.I_testop (nn, op))
    | I_relop (nn, op) -> Ok (I_relop (nn, op))
    | F_unop (nn, op) -> Ok (Binary.F_unop (nn, op))
    | F_relop (nn, op) -> Ok (F_relop (nn, op))
    | I32_wrap_i64 -> Ok Binary.I32_wrap_i64
    | F_reinterpret_i (nn1, nn2) -> Ok (F_reinterpret_i (nn1, nn2))
    | I_reinterpret_f (nn1, nn2) -> Ok (I_reinterpret_f (nn1, nn2))
    | I64_extend_i32 sx -> Ok (Binary.I64_extend_i32 sx)
    | I64_extend32_s -> Ok Binary.I64_extend32_s
    | F32_demote_f64 -> Ok Binary.F32_demote_f64
    | I_extend8_s nn -> Ok (I_extend8_s nn)
    | I_extend16_s nn -> Ok (I_extend16_s nn)
    | F64_promote_f32 -> Ok Binary.F64_promote_f32
    | F_convert_i (nn1, nn2, sx) -> Ok (Binary.F_convert_i (nn1, nn2, sx))
    | I_trunc_f (nn1, nn2, sx) -> Ok (Binary.I_trunc_f (nn1, nn2, sx))
    | I_trunc_sat_f (nn1, nn2, sx) -> Ok (Binary.I_trunc_sat_f (nn1, nn2, sx))
    | Ref_as_non_null -> Ok Binary.Ref_as_non_null
    | Ref_is_null -> Ok Binary.Ref_is_null
    | F_binop (nn, op) -> Ok (Binary.F_binop (nn, op))
    | F32_const v -> Ok (Binary.F32_const v)
    | F64_const v -> Ok (Binary.F64_const v)
    | I32_const v -> Ok (Binary.I32_const v)
    | I64_const v -> Ok (Binary.I64_const v)
    | V128_const v -> Ok (Binary.V128_const v)
    | Unreachable -> Ok Binary.Unreachable
    | Drop -> Ok Binary.Drop
    | Nop -> Ok Binary.Nop
    | Return -> Ok Binary.Return
    | Extern_externalize -> Ok Binary.Extern_externalize
    | Extern_internalize -> Ok Binary.Extern_internalize
    | I_load8 (id, nn, sx, memarg) ->
      let* memarg = rewrite_memarg memarg in
      let+ id = Assigned.find_memory assigned id in
      Binary.I_load8 (id, nn, sx, memarg)
    | I_store8 (id, nn, memarg) ->
      let* memarg = rewrite_memarg memarg in
      let+ id = Assigned.find_memory assigned id in
      Binary.I_store8 (id, nn, memarg)
    | I_load16 (id, nn, sx, memarg) ->
      let* memarg = rewrite_memarg memarg in
      let+ id = Assigned.find_memory assigned id in
      Binary.I_load16 (id, nn, sx, memarg)
    | I_store16 (id, nn, memarg) ->
      let* memarg = rewrite_memarg memarg in
      let+ id = Assigned.find_memory assigned id in
      Binary.I_store16 (id, nn, memarg)
    | I64_load32 (id, sx, memarg) ->
      let* memarg = rewrite_memarg memarg in
      let+ id = Assigned.find_memory assigned id in
      Binary.I64_load32 (id, sx, memarg)
    | I64_store32 (id, memarg) ->
      let* memarg = rewrite_memarg memarg in
      let+ id = Assigned.find_memory assigned id in
      Binary.I64_store32 (id, memarg)
    | I_load (id, nn, memarg) ->
      let* memarg = rewrite_memarg memarg in
      let+ id = Assigned.find_memory assigned id in
      Binary.I_load (id, nn, memarg)
    | F_load (id, nn, memarg) ->
      let* memarg = rewrite_memarg memarg in
      let+ id = Assigned.find_memory assigned id in
      Binary.F_load (id, nn, memarg)
    | F_store (id, nn, memarg) ->
      let* memarg = rewrite_memarg memarg in
      let+ id = Assigned.find_memory assigned id in
      Binary.F_store (id, nn, memarg)
    | I_store (id, nn, memarg) ->
      let* memarg = rewrite_memarg memarg in
      let+ id = Assigned.find_memory assigned id in
      Binary.I_store (id, nn, memarg)
    | Memory_copy (id1, id2) ->
      let* id1 = Assigned.find_memory assigned id1 in
      let+ id2 = Assigned.find_memory assigned id2 in
      Binary.Memory_copy (id1, id2)
    | Memory_size id ->
      let+ id = Assigned.find_memory assigned id in
      Binary.Memory_size id
    | Memory_fill id ->
      let+ id = Assigned.find_memory assigned id in
      Binary.Memory_fill id
    | Memory_grow id ->
      let+ id = Assigned.find_memory assigned id in
      Binary.Memory_grow id
    | V_ibinop (shape, op) -> Ok (Binary.V_ibinop (shape, op))
    | Ref_null t ->
      let* t = rewrite_heap_type assigned t in
      Ok (Binary.Ref_null t)
  and expr (e : Text.expr Annotated.t) (loop_count, block_ids) :
    Binary.expr Annotated.t Result.t =
    let+ e =
      list_map
        (fun i ->
          let+ i = body (loop_count, block_ids) i in
          Annotated.dummy i )
        e.Annotated.raw
    in
    Annotated.dummy e
  in
  expr iexpr (0, [])

let rewrite_table_limits ({ is_i64; min; max } : Text.limits) :
  Binary.Table.Type.limits Result.t =
  if is_i64 then
    let* min =
      try Ok (Int64.of_string_exn min) with Failure _ -> Error `Table_size
    in
    let* max =
      try Ok (Option.map Int64.of_string_exn max)
      with Failure _ -> Error `Table_size
    in
    Ok (Binary.Table.Type.I64 { min; max })
  else
    let* min =
      try Ok (Int32.of_string_exn min) with Failure _ -> Error `Table_size
    in
    let* max =
      try Ok (Option.map Int32.of_string_exn max)
      with Failure _ -> Error `Table_size
    in
    Ok (Binary.Table.Type.I32 { min; max })

let rewrite_table (assigned : Assigned.t)
  ({ id; typ = limits, (null, ht); init } : Text.Table.t) :
  Binary.Table.t Result.t =
  match init with
  | None ->
    let* ht = rewrite_heap_type assigned ht in
    let+ limits = rewrite_table_limits limits in
    { Binary.Table.id; typ = (limits, (null, ht)); init = None }
  | Some e ->
    let* ht = rewrite_heap_type assigned ht in
    let* e = rewrite_expr assigned [] e in
    let+ limits = rewrite_table_limits limits in
    { Binary.Table.id; typ = (limits, (null, ht)); init = Some e }

let rewrite_memory_limits ({ is_i64; min; max } : Text.limits) :
  Binary.Mem.Type.limits Result.t =
  if is_i64 then
    let* min =
      match int_of_string_opt min with
      | Some min -> Ok min
      | None -> Error `Constant_out_of_range
    in
    let* max =
      match max with
      | None -> Ok None
      | Some max -> (
        match int_of_string_opt max with
        | Some max -> Ok (Some max)
        | None -> Error `Constant_out_of_range )
    in
    Ok (Binary.Mem.Type.I64 { min; max })
  else
    let* min =
      try Ok (Int32.of_string_exn min)
      with Failure _ -> Error `Constant_out_of_range
    in
    let* max =
      try Ok (Option.map Int32.of_string_exn max)
      with Failure _ -> Error `Constant_out_of_range
    in
    Ok (Binary.Mem.Type.I32 { min; max })

let rewrite_memory ((id, limits) : Text.Mem.t) : Binary.Mem.t Result.t =
  let+ limits = rewrite_memory_limits limits in
  (id, limits)

let rewrite_global (assigned : Assigned.t) (global : Text.Global.t) :
  Binary.Global.t Result.t =
  let mut, vt = global.typ in
  let* vt = rewrite_val_type assigned vt in
  let+ init = rewrite_expr assigned [] global.init in
  { Binary.Global.id = global.id; init; typ = (mut, vt) }

let rewrite_elem (assigned : Assigned.t) (elem : Text.Elem.t) :
  Binary.Elem.t Result.t =
  let* (mode : Binary.Elem.Mode.t) =
    match elem.mode with
    | Declarative -> Ok Binary.Elem.Mode.Declarative
    | Passive -> Ok Passive
    | Active (None, _expr) -> assert false
    | Active (Some id, expr) ->
      let* indice = Assigned.find_table assigned id in
      let+ expr = rewrite_expr assigned [] expr in
      Binary.Elem.Mode.Active (Some indice, expr)
  in
  let* init = list_map (rewrite_expr assigned []) elem.init in
  let nullable, ht = elem.typ in
  let+ ht = rewrite_heap_type assigned ht in
  { Binary.Elem.init
  ; mode
  ; id = elem.id
  ; typ = (nullable, ht)
  ; explicit_typ = elem.explicit_typ
  }

let rewrite_data (assigned : Assigned.t) (data : Text.Data.t) :
  Binary.Data.t Result.t =
  let+ mode =
    match data.mode with
    | Passive -> Ok Binary.Data.Mode.Passive
    | Active (None, _expr) -> assert false
    | Active (Some indice, expr) ->
      let* indice = Assigned.find_memory assigned indice in
      let+ expr = rewrite_expr assigned [] expr in
      Binary.Data.Mode.Active (indice, expr)
  in
  { Binary.Data.mode; id = data.id; init = data.init }

let rewrite_export find assigned (exports : Grouped.opt_export Array.t) :
  Binary.Export.t Array.t Result.t =
  array_map
    (fun { Grouped.name; id } ->
      match find assigned id with
      | Error _ -> Error (`Unknown_export id)
      | Ok id -> Ok { Binary.Export.name; id } )
    exports

let rewrite_exports (modul : Grouped.t) (assigned : Assigned.t) :
  Binary.Module.Exports.t Result.t =
  let* global =
    rewrite_export Assigned.find_global assigned modul.global_exports
  in
  let* mem = rewrite_export Assigned.find_memory assigned modul.mem_exports in
  let* table =
    rewrite_export Assigned.find_table assigned modul.table_exports
  in
  let* func = rewrite_export Assigned.find_func assigned modul.func_exports in
  let+ tag = rewrite_export Assigned.find_tag assigned modul.tag_exports in
  { Binary.Module.Exports.global; mem; table; func; tag }

let rewrite_func (assigned : Assigned.t)
  ({ id; type_f; locals; body; _ } : Text.Func.t) : Binary.Func.t Result.t =
  let* params, type_f = rewrite_block_type assigned type_f in
  let* body = rewrite_expr assigned (params @ locals) body in
  let+ locals : Binary.param_type =
    list_map
      (fun (_n, vt) ->
        let* vt = rewrite_val_type assigned vt in
        Ok (None, vt) )
      locals
  in
  { Binary.Func.body; type_f; id; locals }

let rewrite_tag (assigned : Assigned.t) ({ id; typ } : Text.Tag.t) :
  Binary.Tag.t Result.t =
  let+ _, typ = rewrite_block_type assigned typ in
  Binary.Tag.{ id; typ }

let rewrite_types (assigned : Assigned.t) (ft : Text.func_type) :
  Binary.Typedef.t Result.t =
  let* ft = rewrite_func_type assigned ft in
  Ok (None, ft)

let modul (modul : Grouped.t) (assigned : Assigned.t) : Binary.Module.t Result.t
    =
  Log.debug (fun m -> m "rewriting    ...");
  let* global =
    let f_local g = rewrite_global assigned g in
    let f_imported (m, val_type) =
      let+ val_type = rewrite_val_type assigned val_type in
      (m, val_type)
    in
    array_map (Origin.monadic_map ~f_local ~f_imported) modul.global
  in
  let* table =
    let f_local g = rewrite_table assigned g in
    let f_imported (l, (n, ht)) =
      let* l = rewrite_table_limits l in
      let+ ht = rewrite_heap_type assigned ht in
      (l, (n, ht))
    in
    array_map (Origin.monadic_map ~f_local ~f_imported) modul.table
  in
  let* mem =
    let f_local g = rewrite_memory g in
    let f_imported = rewrite_memory_limits in
    array_map (Origin.monadic_map ~f_local ~f_imported) modul.mem
  in
  let* elem = array_map (rewrite_elem assigned) modul.elem in
  let* data = array_map (rewrite_data assigned) modul.data in
  let* exports = rewrite_exports modul assigned in
  let* func =
    let f_imported bt =
      let+ _, rt = rewrite_block_type assigned bt in
      rt
    in
    let f_local = rewrite_func assigned in
    let runtime = Origin.monadic_map ~f_local ~f_imported in
    array_map runtime modul.func
  in
  let* tag =
    let f_imported bt =
      let+ _, rt = rewrite_block_type assigned bt in
      rt
    in
    let f_local = rewrite_tag assigned in
    let runtime = Origin.monadic_map ~f_local ~f_imported in
    array_map runtime modul.tag
  in
  let* types =
    array_map (rewrite_types assigned) (Assigned.get_types assigned)
  in
  let+ start =
    match modul.start with
    | None -> Ok None
    | Some id ->
      let+ id = Assigned.find_func assigned id in
      Some id
  in

  { Binary.Module.id = modul.id
  ; mem
  ; table
  ; types
  ; global
  ; elem
  ; data
  ; exports
  ; func
  ; tag
  ; start
  ; custom = []
  }
src/compile/rewrite.ml

95.04%
