diff --git a/toolchain/parser/parse_node_kind.def b/toolchain/parser/parse_node_kind.def
index 355c9cfbfd01..16dffa020bb2 100644
--- a/toolchain/parser/parse_node_kind.def
+++ b/toolchain/parser/parse_node_kind.def
@@ -47,9 +47,11 @@ CARBON_PARSE_NODE_KIND(WhileStatement)
 CARBON_PARSE_NODE_KIND(Condition)
 CARBON_PARSE_NODE_KIND(ConditionEnd)
 CARBON_PARSE_NODE_KIND(ContinueStatement)
+CARBON_PARSE_NODE_KIND(ContinueStatementStart)
 CARBON_PARSE_NODE_KIND(BreakStatement)
+CARBON_PARSE_NODE_KIND(BreakStatementStart)
 CARBON_PARSE_NODE_KIND(ReturnStatement)
-CARBON_PARSE_NODE_KIND(StatementEnd)
+CARBON_PARSE_NODE_KIND(ReturnStatementStart)
 CARBON_PARSE_NODE_KIND(ForStatement)
 CARBON_PARSE_NODE_KIND(ForHeader)
 CARBON_PARSE_NODE_KIND(ForHeaderEnd)
diff --git a/toolchain/parser/parser.cpp b/toolchain/parser/parser.cpp
index ea0dc4492e2d..b3e494a878e6 100644
--- a/toolchain/parser/parser.cpp
+++ b/toolchain/parser/parser.cpp
@@ -1356,12 +1356,12 @@ auto Parser::HandleStatementState() -> void {
   switch (PositionKind()) {
     case TokenKind::Break(): {
       PushState(ParserState::StatementBreakFinish());
-      ++position_;
+      AddLeafNode(ParseNodeKind::BreakStatementStart(), Consume());
       break;
     }
     case TokenKind::Continue(): {
       PushState(ParserState::StatementContinueFinish());
-      ++position_;
+      AddLeafNode(ParseNodeKind::ContinueStatementStart(), Consume());
       break;
     }
     case TokenKind::For(): {
@@ -1398,13 +1398,11 @@ auto Parser::HandleStatementState() -> void {
 }
 
 auto Parser::HandleStatementBreakFinishState() -> void {
-  HandleStatementKeywordFinish(TokenKind::Break(),
-                               ParseNodeKind::BreakStatement());
+  HandleStatementKeywordFinish(ParseNodeKind::BreakStatement());
 }
 
 auto Parser::HandleStatementContinueFinishState() -> void {
-  HandleStatementKeywordFinish(TokenKind::Continue(),
-                               ParseNodeKind::ContinueStatement());
+  HandleStatementKeywordFinish(ParseNodeKind::ContinueStatement());
 }
 
 auto Parser::HandleStatementForHeaderState() -> void {
@@ -1529,37 +1527,38 @@ auto Parser::HandleStatementIfElseBlockFinishState() -> void {
           state.has_error);
 }
 
-auto Parser::HandleStatementKeywordFinish(TokenKind token_kind,
-                                          ParseNodeKind node_kind) -> void {
+auto Parser::HandleStatementKeywordFinish(ParseNodeKind node_kind) -> void {
   auto state = PopState();
 
-  if (!ConsumeAndAddLeafNodeIf(TokenKind::Semi(),
-                               ParseNodeKind::StatementEnd())) {
+  auto semi = ConsumeIf(TokenKind::Semi());
+  if (!semi) {
     CARBON_DIAGNOSTIC(ExpectedSemiAfter, Error, "Expected `;` after `{0}`.",
                       TokenKind);
-    emitter_.Emit(*position_, ExpectedSemiAfter, token_kind);
-    if (auto semi_token = SkipPastLikelyEnd(state.token)) {
-      AddLeafNode(ParseNodeKind::StatementEnd(), *semi_token,
-                  /*has_error=*/true);
+    emitter_.Emit(*position_, ExpectedSemiAfter, tokens_.GetKind(state.token));
+    state.has_error = true;
+    // Recover to the next semicolon if possible, otherwise indicate the
+    // keyword for the error.
+    semi = SkipPastLikelyEnd(state.token);
+    if (!semi) {
+      semi = state.token;
     }
   }
-  AddNode(node_kind, state.token, state.subtree_start, state.has_error);
+  AddNode(node_kind, *semi, state.subtree_start, state.has_error);
 }
 
 auto Parser::HandleStatementReturnState() -> void {
   auto state = PopState();
-
   state.state = ParserState::StatementReturnFinish();
   PushState(state);
-  ++position_;
+
+  AddLeafNode(ParseNodeKind::ReturnStatementStart(), Consume());
   if (!PositionIs(TokenKind::Semi())) {
     PushState(ParserState::Expression());
   }
 }
 
 auto Parser::HandleStatementReturnFinishState() -> void {
-  HandleStatementKeywordFinish(TokenKind::Return(),
-                               ParseNodeKind::ReturnStatement());
+  HandleStatementKeywordFinish(ParseNodeKind::ReturnStatement());
 }
 
 auto Parser::HandleStatementScopeLoopState() -> void {
diff --git a/toolchain/parser/parser.h b/toolchain/parser/parser.h
index 4cacbe116d29..d54e4428c4ac 100644
--- a/toolchain/parser/parser.h
+++ b/toolchain/parser/parser.h
@@ -265,8 +265,7 @@ class Parser {
   auto HandlePattern(PatternKind pattern_kind) -> void;
 
   // Handles the `;` after a keyword statement.
-  auto HandleStatementKeywordFinish(TokenKind token_kind,
-                                    ParseNodeKind node_kind) -> void;
+  auto HandleStatementKeywordFinish(ParseNodeKind node_kind) -> void;
 
   // Handles VarAs(RequireSemicolon|NoSemicolon).
   auto HandleVar(bool require_semicolon) -> void;
diff --git a/toolchain/parser/testdata/basics/return.carbon b/toolchain/parser/testdata/basics/return.carbon
deleted file mode 100644
index 7f0695bf11fd..000000000000
--- a/toolchain/parser/testdata/basics/return.carbon
+++ /dev/null
@@ -1,47 +0,0 @@
-// Part of the Carbon Language project, under the Apache License v2.0 with LLVM
-// Exceptions. See /LICENSE for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-// AUTOUPDATE
-// RUN: %{carbon-run-parser}
-// CHECK:STDOUT: [
-// CHECK:STDOUT:     {kind: 'FunctionIntroducer', text: 'fn'},
-// CHECK:STDOUT:     {kind: 'DeclaredName', text: 'F'},
-// CHECK:STDOUT:       {kind: 'ParameterListEnd', text: ')'},
-// CHECK:STDOUT:     {kind: 'ParameterList', text: '(', subtree_size: 2},
-// CHECK:STDOUT:   {kind: 'FunctionDefinitionStart', text: '{', subtree_size: 5},
-// CHECK:STDOUT:       {kind: 'NameReference', text: 'c'},
-// CHECK:STDOUT:       {kind: 'ConditionEnd', text: ')'},
-// CHECK:STDOUT:     {kind: 'Condition', text: '(', subtree_size: 3},
-// CHECK:STDOUT:       {kind: 'CodeBlockStart', text: '{'},
-// CHECK:STDOUT:         {kind: 'StatementEnd', text: ';'},
-// CHECK:STDOUT:       {kind: 'ReturnStatement', text: 'return', subtree_size: 2},
-// CHECK:STDOUT:     {kind: 'CodeBlock', text: '}', subtree_size: 4},
-// CHECK:STDOUT:   {kind: 'IfStatement', text: 'if', subtree_size: 8},
-// CHECK:STDOUT: {kind: 'FunctionDefinition', text: '}', subtree_size: 14},
-// CHECK:STDOUT:     {kind: 'FunctionIntroducer', text: 'fn'},
-// CHECK:STDOUT:     {kind: 'DeclaredName', text: 'G'},
-// CHECK:STDOUT:         {kind: 'DeclaredName', text: 'x'},
-// CHECK:STDOUT:         {kind: 'NameReference', text: 'Foo'},
-// CHECK:STDOUT:       {kind: 'PatternBinding', text: ':', subtree_size: 3},
-// CHECK:STDOUT:       {kind: 'ParameterListEnd', text: ')'},
-// CHECK:STDOUT:     {kind: 'ParameterList', text: '(', subtree_size: 5},
-// CHECK:STDOUT:       {kind: 'NameReference', text: 'Foo'},
-// CHECK:STDOUT:     {kind: 'ReturnType', text: '->', subtree_size: 2},
-// CHECK:STDOUT:   {kind: 'FunctionDefinitionStart', text: '{', subtree_size: 10},
-// CHECK:STDOUT:     {kind: 'NameReference', text: 'x'},
-// CHECK:STDOUT:     {kind: 'StatementEnd', text: ';'},
-// CHECK:STDOUT:   {kind: 'ReturnStatement', text: 'return', subtree_size: 3},
-// CHECK:STDOUT: {kind: 'FunctionDefinition', text: '}', subtree_size: 14},
-// CHECK:STDOUT: {kind: 'FileEnd', text: ''},
-// CHECK:STDOUT: ]
-
-// NOTE: Move to its own directory when more tests are added.
-fn F() {
-  if (c) {
-    return;
-  }
-}
-fn G(x: Foo) -> Foo {
-  return x;
-}
diff --git a/toolchain/parser/testdata/function/definition/with_return_type.carbon b/toolchain/parser/testdata/function/definition/with_return_type.carbon
index 4f6207afd434..13018141ec97 100644
--- a/toolchain/parser/testdata/function/definition/with_return_type.carbon
+++ b/toolchain/parser/testdata/function/definition/with_return_type.carbon
@@ -12,9 +12,9 @@
 // CHECK:STDOUT:       {kind: 'Literal', text: 'f64'},
 // CHECK:STDOUT:     {kind: 'ReturnType', text: '->', subtree_size: 2},
 // CHECK:STDOUT:   {kind: 'FunctionDefinitionStart', text: '{', subtree_size: 7},
+// CHECK:STDOUT:     {kind: 'ReturnStatementStart', text: 'return'},
 // CHECK:STDOUT:     {kind: 'Literal', text: '42'},
-// CHECK:STDOUT:     {kind: 'StatementEnd', text: ';'},
-// CHECK:STDOUT:   {kind: 'ReturnStatement', text: 'return', subtree_size: 3},
+// CHECK:STDOUT:   {kind: 'ReturnStatement', text: ';', subtree_size: 3},
 // CHECK:STDOUT: {kind: 'FunctionDefinition', text: '}', subtree_size: 11},
 // CHECK:STDOUT: {kind: 'FileEnd', text: ''},
 // CHECK:STDOUT: ]
diff --git a/toolchain/parser/testdata/return/basic.carbon b/toolchain/parser/testdata/return/basic.carbon
new file mode 100644
index 000000000000..a3fe512c802e
--- /dev/null
+++ b/toolchain/parser/testdata/return/basic.carbon
@@ -0,0 +1,21 @@
+// Part of the Carbon Language project, under the Apache License v2.0 with LLVM
+// Exceptions. See /LICENSE for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+// AUTOUPDATE
+// RUN: %{carbon-run-parser}
+// CHECK:STDOUT: [
+// CHECK:STDOUT:     {kind: 'FunctionIntroducer', text: 'fn'},
+// CHECK:STDOUT:     {kind: 'DeclaredName', text: 'F'},
+// CHECK:STDOUT:       {kind: 'ParameterListEnd', text: ')'},
+// CHECK:STDOUT:     {kind: 'ParameterList', text: '(', subtree_size: 2},
+// CHECK:STDOUT:   {kind: 'FunctionDefinitionStart', text: '{', subtree_size: 5},
+// CHECK:STDOUT:     {kind: 'ReturnStatementStart', text: 'return'},
+// CHECK:STDOUT:   {kind: 'ReturnStatement', text: ';', subtree_size: 2},
+// CHECK:STDOUT: {kind: 'FunctionDefinition', text: '}', subtree_size: 8},
+// CHECK:STDOUT: {kind: 'FileEnd', text: ''},
+// CHECK:STDOUT: ]
+
+fn F() {
+  return;
+}
diff --git a/toolchain/parser/testdata/return/expr.carbon b/toolchain/parser/testdata/return/expr.carbon
new file mode 100644
index 000000000000..6f3921054719
--- /dev/null
+++ b/toolchain/parser/testdata/return/expr.carbon
@@ -0,0 +1,22 @@
+// Part of the Carbon Language project, under the Apache License v2.0 with LLVM
+// Exceptions. See /LICENSE for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+// AUTOUPDATE
+// RUN: %{carbon-run-parser}
+// CHECK:STDOUT: [
+// CHECK:STDOUT:     {kind: 'FunctionIntroducer', text: 'fn'},
+// CHECK:STDOUT:     {kind: 'DeclaredName', text: 'F'},
+// CHECK:STDOUT:       {kind: 'ParameterListEnd', text: ')'},
+// CHECK:STDOUT:     {kind: 'ParameterList', text: '(', subtree_size: 2},
+// CHECK:STDOUT:   {kind: 'FunctionDefinitionStart', text: '{', subtree_size: 5},
+// CHECK:STDOUT:     {kind: 'ReturnStatementStart', text: 'return'},
+// CHECK:STDOUT:     {kind: 'NameReference', text: 'x'},
+// CHECK:STDOUT:   {kind: 'ReturnStatement', text: ';', subtree_size: 3},
+// CHECK:STDOUT: {kind: 'FunctionDefinition', text: '}', subtree_size: 9},
+// CHECK:STDOUT: {kind: 'FileEnd', text: ''},
+// CHECK:STDOUT: ]
+
+fn F() {
+  return x;
+}
diff --git a/toolchain/parser/testdata/return/fail_expr_no_semi.carbon b/toolchain/parser/testdata/return/fail_expr_no_semi.carbon
new file mode 100644
index 000000000000..a51db502a6cd
--- /dev/null
+++ b/toolchain/parser/testdata/return/fail_expr_no_semi.carbon
@@ -0,0 +1,23 @@
+// Part of the Carbon Language project, under the Apache License v2.0 with LLVM
+// Exceptions. See /LICENSE for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+// AUTOUPDATE
+// RUN: %{not} %{carbon-run-parser}
+// CHECK:STDOUT: [
+// CHECK:STDOUT:     {kind: 'FunctionIntroducer', text: 'fn'},
+// CHECK:STDOUT:     {kind: 'DeclaredName', text: 'F'},
+// CHECK:STDOUT:       {kind: 'ParameterListEnd', text: ')'},
+// CHECK:STDOUT:     {kind: 'ParameterList', text: '(', subtree_size: 2},
+// CHECK:STDOUT:   {kind: 'FunctionDefinitionStart', text: '{', subtree_size: 5},
+// CHECK:STDOUT:     {kind: 'ReturnStatementStart', text: 'return'},
+// CHECK:STDOUT:     {kind: 'NameReference', text: 'x'},
+// CHECK:STDOUT:   {kind: 'ReturnStatement', text: 'return', has_error: yes, subtree_size: 3},
+// CHECK:STDOUT: {kind: 'FunctionDefinition', text: '}', subtree_size: 9},
+// CHECK:STDOUT: {kind: 'FileEnd', text: ''},
+// CHECK:STDOUT: ]
+
+fn F() {
+  return x
+// CHECK:STDERR: {{.*}}/toolchain/parser/testdata/return/fail_expr_no_semi.carbon:[[@LINE+1]]:1: Expected `;` after `return`.
+}
diff --git a/toolchain/parser/testdata/return/fail_no_semi.carbon b/toolchain/parser/testdata/return/fail_no_semi.carbon
new file mode 100644
index 000000000000..a5d44794f8c7
--- /dev/null
+++ b/toolchain/parser/testdata/return/fail_no_semi.carbon
@@ -0,0 +1,23 @@
+// Part of the Carbon Language project, under the Apache License v2.0 with LLVM
+// Exceptions. See /LICENSE for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+// AUTOUPDATE
+// RUN: %{not} %{carbon-run-parser}
+// CHECK:STDOUT: [
+// CHECK:STDOUT:     {kind: 'FunctionIntroducer', text: 'fn'},
+// CHECK:STDOUT:     {kind: 'DeclaredName', text: 'F'},
+// CHECK:STDOUT:       {kind: 'ParameterListEnd', text: ')'},
+// CHECK:STDOUT:     {kind: 'ParameterList', text: '(', subtree_size: 2},
+// CHECK:STDOUT:   {kind: 'FunctionDefinitionStart', text: '{', subtree_size: 5},
+// CHECK:STDOUT:     {kind: 'ReturnStatementStart', text: 'return'},
+// CHECK:STDOUT:   {kind: 'ReturnStatement', text: 'return', has_error: yes, subtree_size: 2},
+// CHECK:STDOUT: {kind: 'FunctionDefinition', text: '}', subtree_size: 8},
+// CHECK:STDOUT: {kind: 'FileEnd', text: ''},
+// CHECK:STDOUT: ]
+
+fn F() {
+  return
+// CHECK:STDERR: {{.*}}/toolchain/parser/testdata/return/fail_no_semi.carbon:[[@LINE+2]]:1: Expected expression.
+// CHECK:STDERR: {{.*}}/toolchain/parser/testdata/return/fail_no_semi.carbon:[[@LINE+1]]:1: Expected `;` after `return`.
+}
diff --git a/toolchain/parser/testdata/while/basic.carbon b/toolchain/parser/testdata/while/basic.carbon
index 7b5b0d349c07..1860e96362f0 100644
--- a/toolchain/parser/testdata/while/basic.carbon
+++ b/toolchain/parser/testdata/while/basic.carbon
@@ -18,16 +18,16 @@
 // CHECK:STDOUT:           {kind: 'ConditionEnd', text: ')'},
 // CHECK:STDOUT:         {kind: 'Condition', text: '(', subtree_size: 3},
 // CHECK:STDOUT:           {kind: 'CodeBlockStart', text: '{'},
-// CHECK:STDOUT:             {kind: 'StatementEnd', text: ';'},
-// CHECK:STDOUT:           {kind: 'BreakStatement', text: 'break', subtree_size: 2},
+// CHECK:STDOUT:             {kind: 'BreakStatementStart', text: 'break'},
+// CHECK:STDOUT:           {kind: 'BreakStatement', text: ';', subtree_size: 2},
 // CHECK:STDOUT:         {kind: 'CodeBlock', text: '}', subtree_size: 4},
 // CHECK:STDOUT:       {kind: 'IfStatement', text: 'if', subtree_size: 8},
 // CHECK:STDOUT:           {kind: 'NameReference', text: 'c'},
 // CHECK:STDOUT:           {kind: 'ConditionEnd', text: ')'},
 // CHECK:STDOUT:         {kind: 'Condition', text: '(', subtree_size: 3},
 // CHECK:STDOUT:           {kind: 'CodeBlockStart', text: '{'},
-// CHECK:STDOUT:             {kind: 'StatementEnd', text: ';'},
-// CHECK:STDOUT:           {kind: 'ContinueStatement', text: 'continue', subtree_size: 2},
+// CHECK:STDOUT:             {kind: 'ContinueStatementStart', text: 'continue'},
+// CHECK:STDOUT:           {kind: 'ContinueStatement', text: ';', subtree_size: 2},
 // CHECK:STDOUT:         {kind: 'CodeBlock', text: '}', subtree_size: 4},
 // CHECK:STDOUT:       {kind: 'IfStatement', text: 'if', subtree_size: 8},
 // CHECK:STDOUT:     {kind: 'CodeBlock', text: '}', subtree_size: 18},
diff --git a/toolchain/parser/testdata/while/fail_no_semi.carbon b/toolchain/parser/testdata/while/fail_no_semi.carbon
new file mode 100644
index 000000000000..e2cee81ef08d
--- /dev/null
+++ b/toolchain/parser/testdata/while/fail_no_semi.carbon
@@ -0,0 +1,50 @@
+// Part of the Carbon Language project, under the Apache License v2.0 with LLVM
+// Exceptions. See /LICENSE for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+// AUTOUPDATE
+// RUN: %{not} %{carbon-run-parser}
+// CHECK:STDOUT: [
+// CHECK:STDOUT:     {kind: 'FunctionIntroducer', text: 'fn'},
+// CHECK:STDOUT:     {kind: 'DeclaredName', text: 'F'},
+// CHECK:STDOUT:       {kind: 'ParameterListEnd', text: ')'},
+// CHECK:STDOUT:     {kind: 'ParameterList', text: '(', subtree_size: 2},
+// CHECK:STDOUT:   {kind: 'FunctionDefinitionStart', text: '{', subtree_size: 5},
+// CHECK:STDOUT:       {kind: 'NameReference', text: 'a'},
+// CHECK:STDOUT:       {kind: 'ConditionEnd', text: ')'},
+// CHECK:STDOUT:     {kind: 'Condition', text: '(', subtree_size: 3},
+// CHECK:STDOUT:       {kind: 'CodeBlockStart', text: '{'},
+// CHECK:STDOUT:           {kind: 'NameReference', text: 'b'},
+// CHECK:STDOUT:           {kind: 'ConditionEnd', text: ')'},
+// CHECK:STDOUT:         {kind: 'Condition', text: '(', subtree_size: 3},
+// CHECK:STDOUT:           {kind: 'CodeBlockStart', text: '{'},
+// CHECK:STDOUT:             {kind: 'BreakStatementStart', text: 'break'},
+// CHECK:STDOUT:           {kind: 'BreakStatement', text: 'break', has_error: yes, subtree_size: 2},
+// CHECK:STDOUT:         {kind: 'CodeBlock', text: '}', subtree_size: 4},
+// CHECK:STDOUT:       {kind: 'IfStatement', text: 'if', subtree_size: 8},
+// CHECK:STDOUT:           {kind: 'NameReference', text: 'c'},
+// CHECK:STDOUT:           {kind: 'ConditionEnd', text: ')'},
+// CHECK:STDOUT:         {kind: 'Condition', text: '(', subtree_size: 3},
+// CHECK:STDOUT:           {kind: 'CodeBlockStart', text: '{'},
+// CHECK:STDOUT:             {kind: 'ContinueStatementStart', text: 'continue'},
+// CHECK:STDOUT:           {kind: 'ContinueStatement', text: 'continue', has_error: yes, subtree_size: 2},
+// CHECK:STDOUT:         {kind: 'CodeBlock', text: '}', subtree_size: 4},
+// CHECK:STDOUT:       {kind: 'IfStatement', text: 'if', subtree_size: 8},
+// CHECK:STDOUT:     {kind: 'CodeBlock', text: '}', subtree_size: 18},
+// CHECK:STDOUT:   {kind: 'WhileStatement', text: 'while', subtree_size: 22},
+// CHECK:STDOUT: {kind: 'FunctionDefinition', text: '}', subtree_size: 28},
+// CHECK:STDOUT: {kind: 'FileEnd', text: ''},
+// CHECK:STDOUT: ]
+
+fn F() {
+  while (a) {
+    if (b) {
+      break
+    // CHECK:STDERR: {{.*}}/toolchain/parser/testdata/while/fail_no_semi.carbon:[[@LINE+1]]:5: Expected `;` after `break`.
+    }
+    if (c) {
+      continue
+    // CHECK:STDERR: {{.*}}/toolchain/parser/testdata/while/fail_no_semi.carbon:[[@LINE+1]]:5: Expected `;` after `continue`.
+    }
+  }
+}
diff --git a/toolchain/parser/testdata/while/fail_unbraced.carbon b/toolchain/parser/testdata/while/fail_unbraced.carbon
index cf8451883f58..e4726be0bb8e 100644
--- a/toolchain/parser/testdata/while/fail_unbraced.carbon
+++ b/toolchain/parser/testdata/while/fail_unbraced.carbon
@@ -15,8 +15,8 @@
 // CHECK:STDOUT:       {kind: 'ConditionEnd', text: ')'},
 // CHECK:STDOUT:     {kind: 'Condition', text: '(', subtree_size: 3},
 // CHECK:STDOUT:       {kind: 'CodeBlockStart', text: 'break', has_error: yes},
-// CHECK:STDOUT:         {kind: 'StatementEnd', text: ';'},
-// CHECK:STDOUT:       {kind: 'BreakStatement', text: 'break', subtree_size: 2},
+// CHECK:STDOUT:         {kind: 'BreakStatementStart', text: 'break'},
+// CHECK:STDOUT:       {kind: 'BreakStatement', text: ';', subtree_size: 2},
 // CHECK:STDOUT:     {kind: 'CodeBlock', text: 'break', has_error: yes, subtree_size: 4},
 // CHECK:STDOUT:   {kind: 'WhileStatement', text: 'while', subtree_size: 8},
 // CHECK:STDOUT: {kind: 'FunctionDefinition', text: '}', subtree_size: 14},
diff --git a/toolchain/semantics/semantics_parse_tree_handler.cpp b/toolchain/semantics/semantics_parse_tree_handler.cpp
index 9095ec895a13..c7a7e9fe4d67 100644
--- a/toolchain/semantics/semantics_parse_tree_handler.cpp
+++ b/toolchain/semantics/semantics_parse_tree_handler.cpp
@@ -106,7 +106,7 @@ auto SemanticsParseTreeHandler::Build() -> void {
       case ParseNodeKind::DeclaredName():
       case ParseNodeKind::FunctionIntroducer():
       case ParseNodeKind::ParameterListEnd():
-      case ParseNodeKind::StatementEnd(): {
+      case ParseNodeKind::ReturnStatementStart(): {
         // The token has no action, but we still track it for the stack.
         Push(parse_node);
         break;
@@ -255,14 +255,13 @@ auto SemanticsParseTreeHandler::HandleParameterList(ParseTree::Node parse_node)
 
 auto SemanticsParseTreeHandler::HandleReturnStatement(
     ParseTree::Node parse_node) -> void {
-  Pop(ParseNodeKind::StatementEnd());
-
-  // TODO: Restructure ReturnStatement so that we can do this without
-  // looking at the subtree size.
-  if (parse_tree_->node_subtree_size(parse_node) == 2) {
+  if (parse_tree_->node_kind(node_stack_.back().parse_node) ==
+      ParseNodeKind::ReturnStatementStart()) {
+    Pop(ParseNodeKind::ReturnStatementStart());
     Push(parse_node, SemanticsNode::MakeReturn());
   } else {
     auto arg = PopWithResult();
+    Pop(ParseNodeKind::ReturnStatementStart());
     Push(parse_node, SemanticsNode::MakeReturnExpression(arg));
   }
 }