Improve the Fiber and Model stream context #312

pedrobslisboa · 2025-10-11T17:40:46Z

Description

This PR refactors the way we handle the context for Model and Fiber. Instead of 2 different contexts for the Stream, we will have the same context type for both.

type 'a context = {
  push : 'a -> unit;
  close : unit -> unit;
  mutable index : int;
  mutable pending : int;
  env : env;
  debug : bool;
  }

The context controls async and sync push as the index and pending, so we no longer need to worry about it. We declare the value to push and use push/push_async, and the context will handle the rest:

  let push chunk ~context =
    let index = context.index in
    context.push (chunk index);
    context.index <- context.index + 1;
    index

  let push_async promise ~context =
    let index = context.index in
    context.index <- context.index + 1;
    context.pending <- context.pending + 1;
    Lwt.async (fun () ->
        let%lwt chunk = promise in
        context.pending <- context.pending - 1;
        context.push (chunk index);
        if context.pending = 0 then context.close ();
        Lwt.return ());
    index

The concept of inversion of control was applied at render_element_to_html and client_values_to_json. Then we can use the client_values_to_jsonatrender_element_to_html, which reduces the amount of code to maintain, also keeping it way simpler to understand:

Before

After

| Client_component { import_module; import_name; props; client } ->
    let context = Fiber.get_context fiber in
    let lwt_props =
      Lwt_list.map_p
        (fun (name, value) ->
          match (value : React.client_value) with
          | Element element ->
              let%lwt _html, model = render_element_to_html ~fiber element in
              Lwt.return (name, model)
          | Promise (promise, value_to_json) ->
              let index = Fiber.use_index fiber in
              let async : Html.element Lwt.t =
                let%lwt value = promise in
                let json = value_to_json value in
                let ret = chunk_script (Model.model_to_chunk index json) in
                Lwt.return ret
              in
              context.pending <- context.pending + 1;
              Lwt.async (fun () ->
                  let%lwt html = async in
                  context.push html;
                  context.pending <- context.pending - 1;
                  if context.pending = 0 then context.close ();
                  Lwt.return ());
              Lwt.return (name, `String (Model.promise_value index))
          | Json json -> Lwt.return (name, json)
          | Error error ->
              let index = Fiber.use_index fiber in
              let error_json =
                Model.make_error_json 
                     ~env:context.env 
                     ~stack:error.stack 
                     ~message:error.message 
                     ~digest:error.digest
              in
              context.push (error_chunk_script index error_json);
              Lwt.return (name, `String (Model.error_value index))
          | Function action ->
              let index = Fiber.use_index fiber in
              let html =
                chunk_script 
                         (Model.model_to_chunk index 
                         (`Assoc [ ("id", `String action.id); ("bound", `Null) ]))
              in
              context.push html;
              Lwt.return (name, `String (Model.action_value index)))
        props
    in
    let lwt_html = client_to_html ~fiber (client ()) in
    let index = Fiber.use_index fiber in

    let ref : json = Model.component_ref 
                                 ~module_:import_module 
                                 ~name:import_name in
    context.push (client_reference_chunk_script index ref);
    let%lwt html, props = Lwt.both lwt_html lwt_props in
    let model = Model.node ~tag:(Model.ref_value index) ~props [] in
    Lwt.return (html, model)

| Client_component { import_module; import_name; props; client } ->
    let context = get_context fiber in
    let props =
      let on_push value = on_push (Model value) in
      Model.client_values_to_json ~context:fiber.context ~on_push props
    in
    let%lwt html = client_to_html ~fiber ~on_push (client ()) in
    let ref : json = Model.component_ref 
                                 ~module_:import_module 
                                 ~name:import_name in
    let index = Stream.push ~context (on_push (Model (Component_ref ref))) in
    let model = Model.node ~tag:(Model.ref_value index) ~props [] in
    Lwt.return (html, model)

davesnx

Great job

packages/reactDom/test/test_RSC_model.ml