feat: implement message selectors

Author: leostera

riot/riot.mli

@@ -268,7 +268,12 @@ val wait_pids : Pid.t list -> unit
 exception Receive_timeout
 exception Syscall_timeout
 
-val receive : ?after:int64 -> ?ref:unit Ref.t -> unit -> Message.t
+val receive :
+  ?selector:(Message.t -> [ `select of 'msg | `skip ]) ->
+  ?after:int64 ->
+  ?ref:unit Ref.t ->
+  unit ->
+  'msg
 (** [receive ()] will return the first message in the process mailbox.
 
     This function will suspend a process that has an empty mailbox, and the
@@ -283,6 +288,11 @@ val receive : ?after:int64 -> ?ref:unit Ref.t -> unit -> Message.t
 
     ### Selective Receive
 
+    If a `selector` was passed, the `selector` function will be used to select
+    if a message will be picked or if it will be skipped.
+
+    ### Receive with Refs
+
     If a `ref` was passed, then `[receive ~ref ()]` will skip all messages
     created before the creation of this `Ref.t` value, and will only return
     newer messages.

riot/runtime/core/proc_effect.ml

@@ -5,8 +5,9 @@ type _ Effect.t +=
   | Receive : {
       ref : 'a Ref.t option;
       timeout : Timeout.t;
+      selector : Message.t -> [ `select of 'msg | `skip ];
     }
-      -> Message.t Effect.t
+      -> 'msg Effect.t
   [@@unboxed]
 
 type _ Effect.t += Yield : unit Effect.t [@@unboxed]

riot/runtime/import.ml

@@ -17,11 +17,13 @@ let syscall ?timeout name interest source cb =
   Effect.perform (Proc_effect.Syscall { name; interest; source; timeout });
   cb ()
 
-let receive ?after ?ref () =
+let default_selector msg = `select msg
+
+let receive ?(selector = default_selector) ?after ?ref () =
   let timeout =
     match after with None -> `infinity | Some after -> `after after
   in
-  Effect.perform (Proc_effect.Receive { ref; timeout })
+  Effect.perform (Proc_effect.Receive { ref; timeout; selector })
 
 let yield () = Effect.perform Proc_effect.Yield
 let random () = (_get_sch ()).rnd

riot/runtime/scheduler/scheduler.ml

@@ -95,8 +95,8 @@ module Scheduler = struct
           add_to_run_queue sch proc)
       pool.schedulers
 
-  let handle_receive k pool sch (proc : Process.t) (ref : 'a Ref.t option)
-      timeout =
+  let handle_receive k pool sch (proc : Process.t) ~(ref : 'a Ref.t option)
+      ~timeout ~selector =
     Trace.handle_receive_span @@ fun () ->
     let open Proc_state in
     (* When a timeout is specified, we want to create it in the timer
@@ -179,21 +179,35 @@ module Scheduler = struct
              but we already have removed the monitor, then we will simply ignore this message. *)
           | ( _,
               Some
-                Message.
-                  {
-                    msg = Process.Messages.Monitor (Process_down mon_pid) as msg;
-                    _;
-                  } ) ->
+                (Message.
+                   {
+                     msg =
+                       Process.Messages.Monitor (Process_down mon_pid) as msg;
+                     _;
+                   } as envelope) ) ->
               Process.clear_receive_timeout proc;
-              if Process.is_monitored_by_pid proc mon_pid then k (Continue msg)
+              if Process.is_monitored_by_pid proc mon_pid then (
+                match selector msg with
+                | `select msg ->
+                    Process.clear_receive_timeout proc;
+                    k (Continue msg)
+                | `skip ->
+                    Process.add_to_save_queue proc envelope;
+                    go (fuel - 1))
               else go (fuel - 1)
           (* lastly, if we have a ref and the mesasge is newer than the ref, and
              when we don't have a ref, we just pop the message and continue with it
           *)
-          | _, Some Message.{ msg; _ } ->
-              Process.clear_receive_timeout proc;
-              k (Continue msg)
+          | _, Some msg -> (
+              match selector Message.(msg.msg) with
+              | `select msg ->
+                  Process.clear_receive_timeout proc;
+                  k (Continue msg)
+              | `skip ->
+                  Process.add_to_save_queue proc msg;
+                  go (fuel - 1))
       in
+
       go fuel
 
   let handle_syscall k pool (sch : t) (proc : Process.t) name interest source
@@ -263,9 +277,10 @@ module Scheduler = struct
     let perform : type a b. (a, b) step_callback =
      fun k eff ->
       match eff with
+      | Receive { ref; timeout; selector } ->
+          handle_receive k pool sch proc ~ref ~timeout ~selector
       | Syscall { name; interest; source; timeout } ->
           handle_syscall k pool sch proc name interest source timeout
-      | Receive { ref; timeout } -> handle_receive k pool sch proc ref timeout
       | Yield ->
           Log.trace (fun f ->
               f "Process %a: yielding" Pid.pp (Process.pid proc));