Source file atomics.ml

(* Single source of truth for the atomic memory instructions (the threads
   proposal): the mapping between an [Ast.atomicop], its WAT mnemonic, its
   [0xFE]-prefix sub-opcode, its natural alignment and its stack signature.
   [Atomic.fence] is handled separately (it has no memory operand). *)

open Ast

let type_str = function `I32 -> "i32" | `I64 -> "i64"
let width_str = function `I8 -> "8" | `I16 -> "16" | `I32 -> "32"

let rmw_str = function
  | AtomicAdd -> "add"
  | AtomicSub -> "sub"
  | AtomicAnd -> "and"
  | AtomicOr -> "or"
  | AtomicXor -> "xor"
  | AtomicXchg -> "xchg"
  | AtomicCmpxchg -> "cmpxchg"

(* The WAT mnemonic of an atomic op. *)
let name = function
  | AtomicNotify -> "memory.atomic.notify"
  | AtomicWait `I32 -> "memory.atomic.wait32"
  | AtomicWait `I64 -> "memory.atomic.wait64"
  | AtomicLoad (t, None) -> type_str t ^ ".atomic.load"
  | AtomicLoad (t, Some w) -> type_str t ^ ".atomic.load" ^ width_str w ^ "_u"
  | AtomicStore (t, None) -> type_str t ^ ".atomic.store"
  | AtomicStore (t, Some w) -> type_str t ^ ".atomic.store" ^ width_str w
  | AtomicRmw (op, t, None) -> type_str t ^ ".atomic.rmw." ^ rmw_str op
  | AtomicRmw (op, t, Some w) ->
      type_str t ^ ".atomic.rmw" ^ width_str w ^ "." ^ rmw_str op ^ "_u"

(* The (type, width) sequence shared by every load / store / rmw sub-opcode
   block, in binary order: i32-full, i64-full, then the narrow accesses. *)
let variants =
  [
    (`I32, None);
    (`I64, None);
    (`I32, Some `I8);
    (`I32, Some `I16);
    (`I64, Some `I8);
    (`I64, Some `I16);
    (`I64, Some `I32);
  ]

let rmw_ops =
  [
    AtomicAdd;
    AtomicSub;
    AtomicAnd;
    AtomicOr;
    AtomicXor;
    AtomicXchg;
    AtomicCmpxchg;
  ]

(* Every [(sub-opcode, op)] pair, generated from the regular block layout. *)
let table =
  [ (0x00, AtomicNotify); (0x01, AtomicWait `I32); (0x02, AtomicWait `I64) ]
  @ List.mapi (fun i (t, w) -> (0x10 + i, AtomicLoad (t, w))) variants
  @ List.mapi (fun i (t, w) -> (0x17 + i, AtomicStore (t, w))) variants
  @ List.concat
      (List.mapi
         (fun j op ->
           let base = 0x1E + (j * 7) in
           List.mapi (fun i (t, w) -> (base + i, AtomicRmw (op, t, w))) variants)
         rmw_ops)

let all = List.map snd table
let by_opcode = Hashtbl.create 128
let () = List.iter (fun (code, op) -> Hashtbl.replace by_opcode code op) table
let opcode op = fst (List.find (fun (_, o) -> o = op) table)
let of_opcode code = Hashtbl.find_opt by_opcode code

(* The Wax surface: a method name carries the *access width* only
   ([atomic_load16], [atomic_rmw_add8]); the i32/i64 value type is resolved from
   the operand and result types during typing, mirroring the plain scalar
   accesses ([load16(p) as i64_u]). A name therefore denotes a *family* of
   concrete ops; [atomic_wait32]/[atomic_wait64] and [atomic_notify] resolve
   from the name alone. *)

type width = [ `W8 | `W16 | `W32 | `W64 ]

type family =
  | Load of width
  | Store of width
  | Rmw of atomic_rmwop * width
  | Wait of [ `I32 | `I64 ]
  | Notify

let widths : width list = [ `W8; `W16; `W32; `W64 ]

let width_name : width -> string = function
  | `W8 -> "8"
  | `W16 -> "16"
  | `W32 -> "32"
  | `W64 -> "64"

(* Number of bytes a family accesses — the width from the name, independent of
   which i32/i64 value type the operands select; its base-2 logarithm is the
   required (exact) alignment. *)
let width_bytes : width -> int = function
  | `W8 -> 1
  | `W16 -> 2
  | `W32 -> 4
  | `W64 -> 8

let family_bytes = function
  | Load w | Store w | Rmw (_, w) -> width_bytes w
  | Wait `I32 -> 4
  | Wait `I64 -> 8
  | Notify -> 4

(* Every Wax method family, in completion order: loads, stores, RMWs, then
   wait/notify. *)
let families =
  List.map (fun w -> Load w) widths
  @ List.map (fun w -> Store w) widths
  @ List.concat_map (fun op -> List.map (fun w -> Rmw (op, w)) widths) rmw_ops
  @ [ Wait `I32; Wait `I64; Notify ]

(* The Wax method spelling on a memory receiver: [atomic_<op><width>]. *)
let method_name = function
  | Load w -> "atomic_load" ^ width_name w
  | Store w -> "atomic_store" ^ width_name w
  | Rmw (op, w) -> "atomic_rmw_" ^ rmw_str op ^ width_name w
  | Wait `I32 -> "atomic_wait32"
  | Wait `I64 -> "atomic_wait64"
  | Notify -> "atomic_notify"

let by_method = Hashtbl.create 64

let () =
  List.iter (fun f -> Hashtbl.replace by_method (method_name f) f) families

let of_method_name n = Hashtbl.find_opt by_method n

(* The family a concrete op belongs to (its width is the access width). *)
let family op =
  let width t w : width =
    match w with
    | Some `I8 -> `W8
    | Some `I16 -> `W16
    | Some `I32 -> `W32
    | None -> ( match t with `I32 -> `W32 | `I64 -> `W64)
  in
  match op with
  | AtomicNotify -> Notify
  | AtomicWait t -> Wait t
  | AtomicLoad (t, w) -> Load (width t w)
  | AtomicStore (t, w) -> Store (width t w)
  | AtomicRmw (op, t, w) -> Rmw (op, width t w)

(* Number of bytes accessed, whose base-2 logarithm is the required (exact)
   alignment. *)
let access_bytes = function
  | AtomicNotify | AtomicWait `I32 -> 4
  | AtomicWait `I64 -> 8
  | AtomicLoad (t, w) | AtomicStore (t, w) | AtomicRmw (_, t, w) -> (
      match w with
      | Some `I8 -> 1
      | Some `I16 -> 2
      | Some `I32 -> 4
      | None -> ( match t with `I32 -> 4 | `I64 -> 8))

let natural_align_log2 op =
  match access_bytes op with 1 -> 0 | 2 -> 1 | 4 -> 2 | _ -> 3

(* Stack signature after the address operand (which has the memory's address
   type): the remaining operands and the results. *)
let signature op : [ `I32 | `I64 ] list * [ `I32 | `I64 ] list =
  match op with
  | AtomicNotify -> ([ `I32 ], [ `I32 ])
  | AtomicWait `I32 -> ([ `I32; `I64 ], [ `I32 ])
  | AtomicWait `I64 -> ([ `I64; `I64 ], [ `I32 ])
  | AtomicLoad (t, _) -> ([], [ t ])
  | AtomicStore (t, _) -> ([ t ], [])
  | AtomicRmw (AtomicCmpxchg, t, _) -> ([ t; t ], [ t ])
  | AtomicRmw (_, t, _) -> ([ t ], [ t ])