solve async tree stuff

Author: elliott-with-the-longest-name-on-github

packages/svelte/src/internal/server/index.js

@@ -168,7 +168,7 @@ export async function render_async(component, options = {}) {
 		for (const cleanup of async_on_destroy) cleanup();
 		async_on_destroy = prev_on_destroy;
 
-		let { head, body } = await payload.collect_async();
+		let { head, body } = await payload;
 		head += payload.global.head.title.value;
 
 		for (const { hash, code } of payload.global.css) {
@@ -625,10 +625,10 @@ export function maybe_selected(payload, value) {
  * @returns {void}
  */
 export function valueless_option(payload, children) {
-	var i = payload.out.length;
+	var i = payload.length;
 
 	// prior to children, `payload` has some combination of string/unresolved payload that ends in `<option ...>`
-	payload.child((payload) => children(payload));
+	payload.child(children);
 
 	// post-children, `payload` has child content, possibly also with some number of hydration comments.
 	// we can compact this last chunk of content to see if it matches the select value...

packages/svelte/src/internal/server/payload.js

@@ -1,20 +1,29 @@
 /** @typedef {'head' | 'body'} PayloadType */
 /** @typedef {{ [key in PayloadType]: string }} AccumulatedContent */
 /** @typedef {{ start: number, end: number, fn: (content: AccumulatedContent) => AccumulatedContent | Promise<AccumulatedContent> }} Compaction */
+/**
+ * @template T
+ * @typedef {T | Promise<T>} MaybePromise<T>
+ */
 
 /**
  * Payloads are basically a tree of `string | Payload`s, where each `Payload` in the tree represents
  * work that may or may not have completed. A payload can be {@link collect}ed to aggregate the
  * content from itself and all of its children, but this will throw if any of the children are
- * performing asynchronous work. A payload can also be collected asynchronously with
- * {@link collect_async}, which will wait for all children to complete before collecting their
- * contents.
+ * performing asynchronous work. To asynchronously collect a payload, just `await` it.
  *
  * The `string` values within a payload are always associated with the {@link type} of that payload. To switch types,
  * call {@link child} with a different `type` argument.
  */
 export class Payload {
 	/**
+	 * The contents of the payload.
+	 * @type {(string | Payload)[]}
+	 */
+	#out = [];
+
+	/**
+	 * The type of string content that this payload is accumulating.
 	 * @type {PayloadType}
 	 */
 	type;
@@ -23,17 +32,12 @@ export class Payload {
 	parent;
 
 	/**
-	 * The contents of the payload.
-	 * @type {(string | Payload)[]}
-	 */
-	out = [];
-
-	/**
-	 * A promise that resolves when this payload's blocking asynchronous work is done.
-	 * If this promise is not resolved, it is not safe to collect the payload from `out`.
-	 * @type {Promise<void> | undefined}
+	 * Asynchronous work associated with this payload. `initial` is the promise from the function
+	 * this payload was passed to (if that function was async), and `followup` is any any additional
+	 * work from `compact` calls that needs to complete prior to collecting this payload's content.
+	 * @type {{ initial: Promise<void> | undefined, followup: Promise<void>[] | undefined }}
 	 */
-	promise;
+	promises = { initial: undefined, followup: undefined };
 
 	/**
 	 * State which is associated with the content tree as a whole.
@@ -65,72 +69,62 @@ export class Payload {
 
 	/**
 	 * Create a child payload. The child payload inherits the state from the parent,
-	 * but has its own `out` array and `promise` property. The child payload is automatically
-	 * inserted into the parent payload's `out` array.
-	 * @param {(tree: Payload) => void | Promise<void>} render
+	 * but has its own content.
+	 * @param {(tree: Payload) => MaybePromise<void>} render
 	 * @param {PayloadType} [type]
 	 * @returns {void}
 	 */
 	child(render, type) {
 		const child = new Payload(this.global, this.local, this, type);
-		this.out.push(child);
+		this.#out.push(child);
 		const result = render(child);
 		if (result instanceof Promise) {
-			child.promise = result;
+			child.promises.initial = result;
 		}
 	}
 
-	/** @param {string} content */
+	/**
+	 * @param {(value: { head: string, body: string }) => void} onfulfilled
+	 */
+	async then(onfulfilled) {
+		const content = await Payload.#collect_content([this], this.type);
+		return onfulfilled(content);
+	}
+
+	/**
+	 * @param {string} content
+	 */
 	push(content) {
-		this.out.push(content);
+		this.#out.push(content);
 	}
 
 	/**
 	 * Compact everything between `start` and `end` into a single payload, then call `fn` with the result of that payload.
 	 * The compacted payload will be sync if all of the children are sync and {@link fn} is sync, otherwise it will be async.
-	 * @param {{ start: number, end?: number, fn: (content: AccumulatedContent) => AccumulatedContent | Promise<AccumulatedContent> }} args
+	 * @param {{ start: number, end?: number, fn: (content: AccumulatedContent) => AccumulatedContent }} args
 	 */
-	compact({ start, end = this.out.length, fn }) {
+	compact({ start, end = this.#out.length, fn }) {
 		const child = new Payload(this.global, this.local, this);
-		const to_compact = this.out.splice(start, end - start, child);
-		const promises = Payload.#collect_promises(to_compact, []);
-
-		const push_result = () => {
-			const res = fn(Payload.#collect_content(to_compact, this.type));
-			if (res instanceof Promise) {
-				const promise = res.then((resolved) => {
-					Payload.#push_accumulated_content(child, resolved);
-				});
-				return promise;
-			} else {
-				Payload.#push_accumulated_content(child, res);
-			}
-		};
-
-		if (promises.length > 0) {
-			// we have to wait for the accumulated work associated with all pruned branches to complete,
-			// then we can accumulate their content to compact it.
-			child.promise = Promise.all(promises).then(push_result);
+		const to_compact = this.#out.splice(start, end - start, child);
+		const content = Payload.#collect_content(to_compact, this.type);
+
+		if (content instanceof Promise) {
+			const followup = content
+				.then((content) => fn(content))
+				.then((transformed_content) =>
+					Payload.#push_accumulated_content(child, transformed_content)
+				);
+			(this.promises.followup ??= []).push(followup);
 		} else {
-			push_result();
+			Payload.#push_accumulated_content(child, fn(content));
 		}
 	}
 
 	/**
 	 * @returns {number[]}
 	 */
 	get_path() {
-		return this.parent ? [...this.parent.get_path(), this.parent.out.indexOf(this)] : [];
-	}
-
-	/**
-	 * Waits for all child payloads to finish their blocking asynchronous work, then returns the generated content.
-	 * @returns {Promise<AccumulatedContent>}
-	 */
-	async collect_async() {
-		// TODO: Should probably use `Promise.allSettled` here just so we can report detailed errors
-		await Promise.all(Payload.#collect_promises(this.out, this.promise ? [this.promise] : []));
-		return Payload.#collect_content(this.out, this.type);
+		return this.parent ? [...this.parent.get_path(), this.parent.#out.indexOf(this)] : [];
 	}
 
 	/**
@@ -139,19 +133,19 @@ export class Payload {
 	 * @returns {AccumulatedContent}
 	 */
 	collect() {
-		const promises = Payload.#collect_promises(this.out, this.promise ? [this.promise] : []);
-		if (promises.length > 0) {
+		const content = Payload.#collect_content(this.#out, this.type);
+		if (content instanceof Promise) {
 			// TODO is there a good way to report where this is? Probably by using some sort of loc or stack trace in `child` creation.
 			throw new Error('Encountered an asynchronous component while rendering synchronously');
 		}
 
-		return Payload.#collect_content(this.out, this.type);
+		return content;
 	}
 
 	copy() {
 		const copy = new Payload(this.global, this.local, this.parent, this.type);
-		copy.out = this.out.map((item) => (typeof item === 'string' ? item : item.copy()));
-		copy.promise = this.promise;
+		copy.#out = this.#out.map((item) => (typeof item === 'string' ? item : item.copy()));
+		copy.promises = this.promises;
 		return copy;
 	}
 
@@ -161,45 +155,94 @@ export class Payload {
 	subsume(other) {
 		this.global.subsume(other.global);
 		this.local = other.local;
-		this.out = other.out.map((item) => {
+		this.#out = other.#out.map((item) => {
 			if (typeof item !== 'string') {
 				item.subsume(item);
 			}
 			return item;
 		});
-		this.promise = other.promise;
+		this.promises = other.promises;
 		this.type = other.type;
 	}
 
+	get length() {
+		return this.#out.length;
+	}
+
 	/**
+	 * Collect all of the code from the `out` array and return it as a string, or a promise resolving to a string.
 	 * @param {(string | Payload)[]} items
-	 * @param {Promise<void>[]} promises
-	 * @returns {Promise<void>[]}
+	 * @param {PayloadType} current_type
+	 * @param {AccumulatedContent} content
+	 * @returns {MaybePromise<AccumulatedContent>}
 	 */
-	static #collect_promises(items, promises) {
+	static #collect_content(items, current_type, content = { head: '', body: '' }) {
+		/** @type {MaybePromise<AccumulatedContent>[]} */
+		const segments = [];
+		let has_async = false;
+
+		const flush = () => {
+			if (content.head || content.body) {
+				segments.push(content);
+				content = { head: '', body: '' };
+			}
+		};
+
 		for (const item of items) {
-			if (typeof item === 'string') continue;
-			if (item.promise) {
-				promises.push(item.promise);
+			if (typeof item === 'string') {
+				content[current_type] += item;
+			} else {
+				flush();
+
+				if (item.promises.initial) {
+					has_async = true;
+					segments.push(
+						Payload.#collect_content_async([item], current_type, { head: '', body: '' })
+					);
+				} else {
+					const sub = Payload.#collect_content(item.#out, item.type, { head: '', body: '' });
+					if (sub instanceof Promise) {
+						has_async = true;
+					}
+					segments.push(sub);
+				}
 			}
-			Payload.#collect_promises(item.out, promises);
 		}
-		return promises;
+
+		flush();
+
+		if (has_async) {
+			return Promise.all(segments).then((content_array) =>
+				Payload.#squash_accumulated_content(content_array)
+			);
+		}
+
+		// No async segments — combine synchronously
+		return Payload.#squash_accumulated_content(/** @type {AccumulatedContent[]} */ (segments));
 	}
 
 	/**
 	 * Collect all of the code from the `out` array and return it as a string.
 	 * @param {(string | Payload)[]} items
 	 * @param {PayloadType} current_type
 	 * @param {AccumulatedContent} content
-	 * @returns {AccumulatedContent}
+	 * @returns {Promise<AccumulatedContent>}
 	 */
-	static #collect_content(items, current_type, content = { head: '', body: '' }) {
+	static async #collect_content_async(items, current_type, content = { head: '', body: '' }) {
 		for (const item of items) {
 			if (typeof item === 'string') {
 				content[current_type] += item;
 			} else {
-				Payload.#collect_content(item.out, item.type, content);
+				if (item.promises.initial) {
+					// this represents the async function that's modifying this payload.
+					// we can't do anything until it's done and we know our `out` array is complete.
+					await item.promises.initial;
+				}
+				for (const followup of item.promises.followup ?? []) {
+					// this is sequential because `compact` could synchronously queue up additional followup work
+					await followup;
+				}
+				await Payload.#collect_content_async(item.#out, item.type, content);
 			}
 		}
 		return content;
@@ -214,9 +257,24 @@ export class Payload {
 			if (!content) continue;
 			const child = new Payload(tree.global, tree.local, tree, /** @type {PayloadType} */ (type));
 			child.push(content);
-			tree.out.push(child);
+			tree.#out.push(child);
 		}
 	}
+
+	/**
+	 * @param {AccumulatedContent[]} content_array
+	 * @returns {AccumulatedContent}
+	 */
+	static #squash_accumulated_content(content_array) {
+		return content_array.reduce(
+			(acc, content) => {
+				acc.head += content.head;
+				acc.body += content.body;
+				return acc;
+			},
+			{ head: '', body: '' }
+		);
+	}
 }
 
 export class TreeState {