From 344ef8558b2c5e66c9b844d3944218165ef7de85 Mon Sep 17 00:00:00 2001
From: Greg Johnston <greg.johnston@gmail.com>
Date: Sun, 4 Feb 2024 20:56:29 -0500
Subject: [PATCH] feat: modular `SharedContext` for hydration

---
 Cargo.toml                       |   1 +
 hydration_context/Cargo.toml     |  18 +++++
 hydration_context/src/hydrate.rs |  74 +++++++++++++++++
 hydration_context/src/lib.rs     |  91 +++++++++++++++++++++
 hydration_context/src/ssr.rs     | 133 +++++++++++++++++++++++++++++++
 5 files changed, 317 insertions(+)
 create mode 100644 hydration_context/Cargo.toml
 create mode 100644 hydration_context/src/hydrate.rs
 create mode 100644 hydration_context/src/lib.rs
 create mode 100644 hydration_context/src/ssr.rs

diff --git a/Cargo.toml b/Cargo.toml
index 598f7f758e..5322b42bd9 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -5,6 +5,7 @@ members = [
   "or_poisoned",
 
   # core
+  "hydration_context",
   "leptos",
   "leptos_dom",
   "leptos_config",
diff --git a/hydration_context/Cargo.toml b/hydration_context/Cargo.toml
new file mode 100644
index 0000000000..e570b1d4c1
--- /dev/null
+++ b/hydration_context/Cargo.toml
@@ -0,0 +1,18 @@
+[package]
+name = "hydration_context"
+edition = "2021"
+version.workspace = true
+
+[dependencies]
+or_poisoned = { workspace = true }
+futures = "0.3"
+serde = { version = "1", features = ["derive"] }
+wasm-bindgen = { version = "0.2", optional = true }
+js-sys = { version = "0.3", optional = true }
+
+[features]
+browser = ["dep:wasm-bindgen", "dep:js-sys"]
+
+[package.metadata.docs.rs]
+all-features = true
+rustdoc-args = ["--cfg", "docsrs"]
diff --git a/hydration_context/src/hydrate.rs b/hydration_context/src/hydrate.rs
new file mode 100644
index 0000000000..7362368dc5
--- /dev/null
+++ b/hydration_context/src/hydrate.rs
@@ -0,0 +1,74 @@
+use super::{SerializedDataId, SharedContext};
+use crate::{PinnedFuture, PinnedStream};
+use core::fmt::Debug;
+use js_sys::Array;
+use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
+use wasm_bindgen::prelude::wasm_bindgen;
+
+#[wasm_bindgen]
+extern "C" {
+    static __RESOLVED_RESOURCES: Array;
+}
+
+#[derive(Default)]
+/// The shared context that should be used in the browser while hydrating.
+pub struct HydrateSharedContext {
+    id: AtomicUsize,
+    is_hydrating: AtomicBool,
+}
+
+impl HydrateSharedContext {
+    /// Creates a new shared context for hydration in the browser.
+    pub fn new() -> Self {
+        Self {
+            id: AtomicUsize::new(0),
+            is_hydrating: AtomicBool::new(true),
+        }
+    }
+
+    /// Creates a new shared context for hydration in the browser.
+    ///
+    /// This defaults to a mode in which the app is not hydrated, but allows you to opt into
+    /// hydration for certain portions using [`SharedContext::set_is_hydrating`].
+    pub fn new_islands() -> Self {
+        Self {
+            id: AtomicUsize::new(0),
+            is_hydrating: AtomicBool::new(false),
+        }
+    }
+}
+
+impl Debug for HydrateSharedContext {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("HydrateSharedContext").finish()
+    }
+}
+
+impl SharedContext for HydrateSharedContext {
+    fn next_id(&self) -> SerializedDataId {
+        let id = self.id.fetch_add(1, Ordering::Relaxed);
+        SerializedDataId(id)
+    }
+
+    fn write_async(&self, _id: SerializedDataId, _fut: PinnedFuture<String>) {}
+
+    fn read_data(&self, id: &SerializedDataId) -> Option<String> {
+        __RESOLVED_RESOURCES.get(id.0 as u32).as_string()
+    }
+
+    fn await_data(&self, _id: &SerializedDataId) -> Option<String> {
+        todo!()
+    }
+
+    fn pending_data(&self) -> Option<PinnedStream<String>> {
+        None
+    }
+
+    fn get_is_hydrating(&self) -> bool {
+        self.is_hydrating.load(Ordering::Relaxed)
+    }
+
+    fn set_is_hydrating(&self, is_hydrating: bool) {
+        self.is_hydrating.store(true, Ordering::Relaxed)
+    }
+}
diff --git a/hydration_context/src/lib.rs b/hydration_context/src/lib.rs
new file mode 100644
index 0000000000..07c17df4d6
--- /dev/null
+++ b/hydration_context/src/lib.rs
@@ -0,0 +1,91 @@
+//! Isomorphic web applications that run on the server to render HTML, then add interactivity in
+//! the client, need to accomplish two tasks:
+//! 1. Send HTML from the server, so that the client can "hydrate" it in the browser by adding
+//!    event listeners and setting up other interactivity.
+//! 2. Send data that was loaded on the server to the client, so that the client "hydrates" with
+//!    the same data with which the server rendered HTML.
+//!
+//! This crate helps with the second part of this process. It provides a [`SharedContext`] type
+//! that allows you to store data on the server, and then extract the same data in the client.
+
+#![deny(missing_docs)]
+#![forbid(unsafe_code)]
+
+#[cfg(feature = "browser")]
+#[cfg_attr(docsrs, doc(cfg(feature = "browser")))]
+mod hydrate;
+mod ssr;
+use futures::Stream;
+#[cfg(feature = "browser")]
+pub use hydrate::*;
+use serde::{Deserialize, Serialize};
+pub use ssr::*;
+use std::{fmt::Debug, future::Future, pin::Pin};
+
+/// Type alias for a boxed [`Future`].
+pub type PinnedFuture<T> = Pin<Box<dyn Future<Output = T> + Send + Sync>>;
+/// Type alias for a boxed [`Future`] that is `!Send`.
+pub type PinnedLocalFuture<T> = Pin<Box<dyn Future<Output = T>>>;
+/// Type alias for a boxed [`Stream`].
+pub type PinnedStream<T> = Pin<Box<dyn Stream<Item = T> + Send + Sync>>;
+
+#[derive(
+    Clone, Debug, PartialEq, Eq, Hash, Default, Deserialize, Serialize,
+)]
+#[serde(transparent)]
+/// A unique identifier for a piece of data that will be serialized
+/// from the server to the client.
+pub struct SerializedDataId(usize);
+
+/// Information that will be shared between the server and the client.
+pub trait SharedContext: Debug {
+    /// Returns the next in a series of IDs that is unique to a particular request and response.
+    ///
+    /// This should not be used as a global unique ID mechanism. It is specific to the process
+    /// of serializing and deserializing data from the server to the browser as part of an HTTP
+    /// response.
+    fn next_id(&self) -> SerializedDataId;
+
+    /// The given [`Future`] should resolve with some data that can be serialized
+    /// from the server to the client. This will be polled as part of the process of
+    /// building the HTTP response, *not* when it is first created.
+    ///
+    /// In browser implementations, this should be a no-op.
+    fn write_async(&self, id: SerializedDataId, fut: PinnedFuture<String>);
+
+    /// Reads the current value of some data from the shared context, if it has been
+    /// sent from the server. This returns the serialized data as a `String` that should
+    /// be deserialized using [`Serializable::de`].
+    ///
+    /// On the server and in client-side rendered implementations, this should
+    /// always return [`None`].
+    fn read_data(&self, id: &SerializedDataId) -> Option<String>;
+
+    /// Returns a [`Future`] that resolves with a `String` that should
+    /// be deserialized using [`Serializable::de`] once the given piece of server
+    /// data has resolved.
+    ///
+    /// On the server and in client-side rendered implementations, this should
+    /// return a [`Future`] that is immediately ready with [`None`].
+    fn await_data(&self, id: &SerializedDataId) -> Option<String>;
+
+    /// Returns some [`Stream`] of HTML that contains JavaScript `<script>` tags defining
+    /// all values being serialized from the server to the client, with their serialized values
+    /// and any boilerplate needed to notify a running application that they exist; or `None`.
+    ///
+    /// In browser implementations, this return `None`.
+    fn pending_data(&self) -> Option<PinnedStream<String>>;
+
+    /// Returns `true` if you are currently in a part of the application tree that should be
+    /// hydrated.
+    ///
+    /// For example, in an app with "islands," this should be `true` inside islands and
+    /// false elsewhere.
+    fn get_is_hydrating(&self) -> bool;
+
+    /// Sets whether you are currently in a part of the application tree that should be hydrated.
+    ///
+    /// For example, in an app with "islands," this should be `true` inside islands and
+    /// false elsewhere.
+    fn set_is_hydrating(&self, is_hydrating: bool);
+}
diff --git a/hydration_context/src/ssr.rs b/hydration_context/src/ssr.rs
new file mode 100644
index 0000000000..9524a1110c
--- /dev/null
+++ b/hydration_context/src/ssr.rs
@@ -0,0 +1,133 @@
+use super::{SerializedDataId, SharedContext};
+use crate::{PinnedFuture, PinnedStream};
+use futures::{
+    stream::{self, FuturesUnordered},
+    StreamExt,
+};
+use or_poisoned::OrPoisoned;
+use std::{
+    fmt::{Debug, Write},
+    mem,
+    sync::{
+        atomic::{AtomicBool, AtomicUsize, Ordering},
+        RwLock,
+    },
+};
+
+#[derive(Default)]
+/// The shared context that should be used on the server side.
+pub struct SsrSharedContext {
+    id: AtomicUsize,
+    is_hydrating: AtomicBool,
+    sync_buf: RwLock<Vec<ResolvedData>>,
+    async_buf: RwLock<Vec<(SerializedDataId, PinnedFuture<String>)>>,
+}
+
+impl SsrSharedContext {
+    /// Creates a new shared context for rendering HTML on the server.
+    pub fn new() -> Self {
+        Self {
+            is_hydrating: AtomicBool::new(true),
+            ..Default::default()
+        }
+    }
+
+    /// Creates a new shared context for rendering HTML on the server in "islands" mode.
+    ///
+    /// This defaults to a mode in which the app is not hydrated, but allows you to opt into
+    /// hydration for certain portions using [`SharedContext::set_is_hydrating`].
+    pub fn new_islands() -> Self {
+        Self {
+            is_hydrating: AtomicBool::new(false),
+            ..Default::default()
+        }
+    }
+}
+
+impl Debug for SsrSharedContext {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("SsrSharedContext")
+            .field("id", &self.id)
+            .field("is_hydrating", &self.is_hydrating)
+            .field("sync_buf", &self.sync_buf)
+            .field("async_buf", &self.async_buf.read().or_poisoned().len())
+            .finish()
+    }
+}
+
+impl SharedContext for SsrSharedContext {
+    fn next_id(&self) -> SerializedDataId {
+        let id = self.id.fetch_add(1, Ordering::Relaxed);
+        SerializedDataId(id)
+    }
+
+    fn write_async(&self, id: SerializedDataId, fut: PinnedFuture<String>) {
+        self.async_buf.write().or_poisoned().push((id, fut))
+    }
+
+    fn pending_data(&self) -> Option<PinnedStream<String>> {
+        let sync_data = mem::take(&mut *self.sync_buf.write().or_poisoned());
+        let async_data = mem::take(&mut *self.async_buf.write().or_poisoned());
+
+        // 1) initial, synchronous setup chunk
+        let mut initial_chunk = String::new();
+        // resolved synchronous resources
+        initial_chunk.push_str("__RESOLVED_RESOURCES=[");
+        for resolved in sync_data {
+            resolved.write_to_buf(&mut initial_chunk);
+            initial_chunk.push(',');
+        }
+        initial_chunk.push_str("];");
+
+        // pending async resources
+        initial_chunk.push_str("__PENDING_RESOURCES=[");
+        for (id, _) in &async_data {
+            write!(&mut initial_chunk, "{},", id.0).unwrap();
+        }
+        initial_chunk.push_str("];");
+
+        // resolvers
+        initial_chunk.push_str("__RESOURCE_RESOLVERS=[];");
+
+        // 2) async resources as they resolve
+        let async_data = async_data
+            .into_iter()
+            .map(|(id, data)| async move {
+                let data = data.await;
+                format!("__RESOLVED_RESOURCES[{}] = {data:?};", id.0)
+            })
+            .collect::<FuturesUnordered<_>>();
+
+        let stream =
+            stream::once(async move { initial_chunk }).chain(async_data);
+        Some(Box::pin(stream))
+    }
+
+    fn read_data(&self, _id: &SerializedDataId) -> Option<String> {
+        None
+    }
+
+    fn await_data(&self, _id: &SerializedDataId) -> Option<String> {
+        None
+    }
+
+    fn get_is_hydrating(&self) -> bool {
+        self.is_hydrating.load(Ordering::Relaxed)
+    }
+
+    fn set_is_hydrating(&self, is_hydrating: bool) {
+        self.is_hydrating.store(is_hydrating, Ordering::Relaxed)
+    }
+}
+
+#[derive(Debug)]
+struct ResolvedData(SerializedDataId, String);
+
+impl ResolvedData {
+    pub fn write_to_buf(&self, buf: &mut String) {
+        let ResolvedData(id, ser) = self;
+        // escapes < to prevent it being interpreted as another opening HTML tag
+        let ser = ser.replace('<', "\\u003c");
+        write!(buf, "{}: {:?}", id.0, ser).unwrap();
+    }
+}