only wrap awaits in `$.save` when necessary

packages/svelte/src/compiler/phases/2-analyze/index.js

@@ -1,5 +1,5 @@
 /** @import { Expression, Node, Program } from 'estree' */
-/** @import { Binding, AST, ValidatedCompileOptions, ValidatedModuleCompileOptions } from '#compiler' */
+/** @import { Binding, AST, ValidatedCompileOptions, ValidatedModuleCompileOptions, ExpressionMetadata } from '#compiler' */
 /** @import { AnalysisState, Visitors } from './types' */
 /** @import { Analysis, ComponentAnalysis, Js, ReactiveStatement, Template } from '../types' */
 import { walk } from 'zimmerframe';
@@ -76,6 +76,10 @@ import { UseDirective } from './visitors/UseDirective.js';
 import { VariableDeclarator } from './visitors/VariableDeclarator.js';
 import is_reference from 'is-reference';
 import { mark_subtree_dynamic } from './visitors/shared/fragment.js';
+import { find_last_await, is_last_evaluated_expression } from './utils/awaits.js';
+
+/** @type {Array<ExpressionMetadata | null>} */
+const metadata_stack = [];
 
 /**
  * @type {Visitors}
@@ -127,9 +131,23 @@ const visitors = {
 
 		ignore_map.set(node, structuredClone(ignore_stack));
 
+		metadata_stack.push(state.expression);
+
 		const scope = state.scopes.get(node);
 		next(scope !== undefined && scope !== state.scope ? { ...state, scope } : state);
 
+		metadata_stack.pop();
+
+		// if this node set `state.expression`, now that we've visited it we can determine
+		// which `await` expressions need to be wrapped in `$.save(...)`
+		if (state.expression && metadata_stack[metadata_stack.length - 1] === null) {
+			for (const { path, node } of state.expression.awaits) {
+				if (!is_last_evaluated_expression(path, node)) {
+					state.analysis.context_preserving_awaits.add(node);
+				}
+			}
+		}
+
 		if (ignores.length > 0) {
 			pop_ignore();
 		}

packages/svelte/src/compiler/phases/2-analyze/utils/awaits.js

@@ -0,0 +1,70 @@
+/** @import { Expression, Property, SpreadElement } from 'estree' */
+/** @import { AST } from '#compiler' */
+
+/**
+ *
+ * @param {AST.SvelteNode[]} path
+ * @param {Expression | SpreadElement | Property} node
+ */
+export function is_last_evaluated_expression(path, node) {
+	let i = path.length;
+
+	while (i--) {
+		const parent = /** @type {Expression | Property | SpreadElement} */ (path[i]);
+
+		// @ts-expect-error we could probably use a neater/more robust mechanism
+		if (parent.metadata) {
+			return true;
+		}
+
+		switch (parent.type) {
+			case 'ArrayExpression':
+				if (node !== parent.elements.at(-1)) return false;
+				break;
+
+			case 'AssignmentExpression':
+			case 'BinaryExpression':
+			case 'LogicalExpression':
+				if (node === parent.left) return false; // TODO is this right for assignment expressions?
+				break;
+
+			case 'CallExpression':
+			case 'NewExpression':
+				if (node !== parent.arguments.at(-1)) return false;
+				break;
+
+			case 'ConditionalExpression':
+				if (node === parent.test) return false;
+				break;
+
+			case 'MemberExpression':
+				if (parent.computed && node === parent.object) return false;
+				break;
+
+			case 'ObjectExpression':
+				if (node !== parent.properties.at(-1)) return false;
+				break;
+
+			case 'Property':
+				if (node === parent.key) return false;
+				break;
+
+			case 'SequenceExpression':
+				if (node !== parent.expressions.at(-1)) return false;
+				break;
+
+			case 'TaggedTemplateExpression':
+				if (node !== parent.quasi.expressions.at(-1)) return false;
+				break;
+
+			case 'TemplateLiteral':
+				if (node !== parent.expressions.at(-1)) return false;
+				break;
+
+			default:
+				return false;
+		}
+
+		node = parent;
+	}
+}

packages/svelte/src/compiler/phases/2-analyze/visitors/AwaitExpression.js

@@ -8,28 +8,17 @@ import * as e from '../../../errors.js';
  */
 export function AwaitExpression(node, context) {
 	const tla = context.state.ast_type === 'instance' && context.state.function_depth === 1;
+
+	if (tla) {
+		context.state.analysis.context_preserving_awaits.add(node);
+	}
+
 	let suspend = tla;
-	let preserve_context = tla;
 
 	if (context.state.expression) {
+		context.state.expression.awaits.push({ node, path: context.path.slice() });
 		context.state.expression.has_await = true;
 		suspend = true;
-
-		// wrap the expression in `(await $.save(...)).restore()` if necessary,
-		// i.e. whether anything could potentially be read _after_ the await
-		let i = context.path.length;
-		while (i--) {
-			const parent = context.path[i];
-
-			// stop walking up when we find a node with metadata, because that
-			// means we've hit the template node containing the expression
-			// @ts-expect-error we could probably use a neater/more robust mechanism
-			if (parent.metadata) break;
-
-			// TODO make this more accurate — we don't need to call suspend
-			// if this is the last thing that could be read
-			preserve_context = true;
-		}
 	}
 
 	if (suspend) {
@@ -42,9 +31,5 @@ export function AwaitExpression(node, context) {
 		}
 	}
 
-	if (preserve_context) {
-		context.state.analysis.context_preserving_awaits.add(node);
-	}
-
 	context.next();
 }

packages/svelte/src/compiler/phases/nodes.js

@@ -67,7 +67,8 @@ export function create_expression_metadata() {
 		has_call: false,
 		has_member_expression: false,
 		has_assignment: false,
-		has_await: false
+		has_await: false,
+		awaits: []
 	};
 }
 

packages/svelte/src/compiler/types/index.d.ts

@@ -5,6 +5,7 @@ import type { ICompileDiagnostic } from '../utils/compile_diagnostic.js';
 import type { StateCreationRuneName } from '../../utils.js';
 import type {
 	AssignmentExpression,
+	AwaitExpression,
 	CallExpression,
 	PrivateIdentifier,
 	PropertyDefinition
@@ -298,6 +299,8 @@ export interface ExpressionMetadata {
 	has_member_expression: boolean;
 	/** True if the expression includes an assignment or an update */
 	has_assignment: boolean;
+	/** An array of await expressions, so we can figure out which ones need context preservation */
+	awaits: Array<{ node: AwaitExpression; path: AST.SvelteNode[] }>;
 }
 
 export interface StateField {

packages/svelte/src/compiler/utils/builders.js

@@ -1,4 +1,5 @@
 /** @import * as ESTree from 'estree' */
+import { walk } from 'zimmerframe';
 import { regex_is_valid_identifier } from '../phases/patterns.js';
 import { sanitize_template_string } from './sanitize_template_string.js';
 
@@ -432,8 +433,20 @@ export function thunk(expression, async = false) {
  * @returns {ESTree.Expression}
  */
 export function unthunk(expression) {
+	// optimize `async () => await x()`, but not `async () => await x(await y)`
 	if (expression.async && expression.body.type === 'AwaitExpression') {
-		return unthunk(arrow(expression.params, expression.body.argument));
+		let has_await = false;
+
+		walk(expression.body.argument, null, {
+			AwaitExpression(node, context) {
+				has_await = true;
+				context.stop();
+			}
+		});
+
+		if (!has_await) {
+			return unthunk(arrow(expression.params, expression.body.argument));
+		}
 	}
 
 	if (