Library zoo_parabs.pool__code

From zoo Require Import
  prelude.
From zoo.language Require Import
  typeclasses
  notations.
From zoo_parabs Require Import
  ws_hub_std.
From zoo_std Require Import
  ivar_4
  array
  domain.
From zoo_parabs Require Import
  pool__types.
From zoo Require Import
  options.

Definition pool ٠ max_round_noyield : val :=
  1024.

Definition pool ٠ max_round_yield : val :=
  32.

Definition pool ٠ context : val :=
  fun : "sz" "hub" "id" ⇒
    ("sz", "hub", "id").

Definition pool ٠ context_main : val :=
  fun : "t" ⇒
    pool ٠ context "t".{size } "t".{hub } 0.

Definition pool ٠ execute : val :=
  fun : "ctx" "job" ⇒
    "job" "ctx".

Definition pool ٠ worker : val :=
  rec : "worker" "ctx" ⇒
    match :
      ws_hub_std ٠ pop_steal
        "ctx".<context_hub >
        "ctx".<context_id >
        pool ٠ max_round_noyield
        pool ٠ max_round_yield
    with
    | None ⇒
        ()
    | Some "job" ⇒
        pool ٠ execute "ctx" "job" ;;
        "worker" "ctx"
    end.

Definition pool ٠ create : val :=
  fun : "sz" ⇒
    let : "hub" := ws_hub_std ٠ create ("sz" + 1) in
    ws_hub_std ٠ block "hub" 0 ;;
    let : "domains" :=
      array ٠ unsafe_initi
        "sz"
        (fun : "i" ⇒
           domain ٠ spawn
             (fun : ≠ ⇒ pool ٠ worker (pool ٠ context "sz" "hub" ("i" + 1))))
    in
    { "sz", "hub", "domains", () }.

Definition pool ٠ run_on : val :=
  fun : "t" "task" ⇒
    ws_hub_std ٠ unblock "t".{hub } 0 ;;
    let : "res" := pool ٠ execute (pool ٠ context_main "t") "task" in
    ws_hub_std ٠ block "t".{hub } 0 ;;
    "res".

Definition pool ٠ close : val :=
  fun : "t" ⇒
    ws_hub_std ٠ close "t".{hub } ;;
    ws_hub_std ٠ unblock "t".{hub } 0 ;;
    pool ٠ worker (pool ٠ context_main "t") ;;
    array ٠ iter domain ٠ join "t".{domains }.

Definition pool ٠ run : val :=
  fun : "num_worker" "task" ⇒
    let : "t" := pool ٠ create "num_worker" in
    let : "res" := pool ٠ run_on "t" "task" in
    pool ٠ close "t" ;;
    "res".

Definition pool ٠ size : val :=
  fun : "ctx" ⇒
    "ctx".<context_size >.

Definition pool ٠ async : val :=
  fun : "ctx" "task" ⇒
    ws_hub_std ٠ push "ctx".<context_hub > "ctx".<context_id > "task".

Definition pool ٠ wait₀ : val :=
  rec : "wait" "ctx" "notification" "pred" ⇒
    match :
      ws_hub_std ٠ pop_steal_until
        "ctx".<context_hub >
        "ctx".<context_id >
        pool ٠ max_round_noyield
        pool ٠ max_round_yield
        "notification"
        "pred"
    with
    | None ⇒
        ()
    | Some "job" ⇒
        pool ٠ execute "ctx" "job" ;;
        "wait" "ctx" "notification" "pred"
    end.

Definition pool ٠ wait : val :=
  fun : "ctx" "notification" "pred" ⇒
    let : "notification_registered" := ref false in
    let : "notification" "notify" :=
      if : ¬ !"notification_registered" then (
        "notification_registered" <- true ;;
        "notification" "notify"
      )
    in
    pool ٠ wait₀ "ctx" "notification" "pred".

Definition pool ٠ wait_ivar : val :=
  fun : "ctx" "ivar" ⇒
    pool ٠ wait
      "ctx"
      (fun : "notify" ⇒
         ivar_4 ٠ wait "ivar" (fun : "_ctx" "_v" ⇒ "notify" ()) ;;
         ())
      (fun : ≠ ⇒ ivar_4 ٠ is_set "ivar").