Track loop returns in addition to continue/break.

Besides `continue` and `break`, another way that a for loop can
exit early is via a `return` statement. This requires slightly
deeper inspection, because a `return` can be nested inside a
loop or switch.

Change-Id: I0a0e27d7f38dd51ffacb57a7a1f90657008d640b
Reviewed-on: https://skia-review.googlesource.com/c/skia/+/634457
Reviewed-by: Brian Osman <brianosman@google.com>
Auto-Submit: John Stiles <johnstiles@google.com>
diff --git a/gn/sksl.gni b/gn/sksl.gni
index 99be65c..0ce3d68 100644
--- a/gn/sksl.gni
+++ b/gn/sksl.gni
@@ -113,9 +113,9 @@
   "$_src/sksl/analysis/SkSLCanExitWithoutReturningValue.cpp",
   "$_src/sksl/analysis/SkSLCheckProgramStructure.cpp",
   "$_src/sksl/analysis/SkSLFinalizationChecks.cpp",
+  "$_src/sksl/analysis/SkSLGetLoopControlFlowInfo.cpp",
   "$_src/sksl/analysis/SkSLGetLoopUnrollInfo.cpp",
   "$_src/sksl/analysis/SkSLGetReturnComplexity.cpp",
-  "$_src/sksl/analysis/SkSLHasContinueOrBreak.cpp",
   "$_src/sksl/analysis/SkSLHasSideEffects.cpp",
   "$_src/sksl/analysis/SkSLIsConstantExpression.cpp",
   "$_src/sksl/analysis/SkSLIsDynamicallyUniformExpression.cpp",
diff --git a/public.bzl b/public.bzl
index bcd1f2b..1ed5150 100644
--- a/public.bzl
+++ b/public.bzl
@@ -1451,9 +1451,9 @@
     "src/sksl/analysis/SkSLCanExitWithoutReturningValue.cpp",
     "src/sksl/analysis/SkSLCheckProgramStructure.cpp",
     "src/sksl/analysis/SkSLFinalizationChecks.cpp",
+    "src/sksl/analysis/SkSLGetLoopControlFlowInfo.cpp",
     "src/sksl/analysis/SkSLGetLoopUnrollInfo.cpp",
     "src/sksl/analysis/SkSLGetReturnComplexity.cpp",
-    "src/sksl/analysis/SkSLHasContinueOrBreak.cpp",
     "src/sksl/analysis/SkSLHasSideEffects.cpp",
     "src/sksl/analysis/SkSLIsConstantExpression.cpp",
     "src/sksl/analysis/SkSLIsDynamicallyUniformExpression.cpp",
diff --git a/src/sksl/SkSLAnalysis.h b/src/sksl/SkSLAnalysis.h
index f580267..b875098 100644
--- a/src/sksl/SkSLAnalysis.h
+++ b/src/sksl/SkSLAnalysis.h
@@ -119,14 +119,16 @@
 bool StatementWritesToVariable(const Statement& stmt, const Variable& var);
 
 /**
- * Returns true if the passed-in block contains a `continue` or `break` that could directly affect
- * its control flow. (A `continue` or `break` nested inside an inner loop/switch does not count.)
+ * Detects if the passed-in block contains a `continue`, `break` or `return` that could directly
+ * affect its control flow. (A `continue` or `break` nested inside an inner loop/switch will not
+ * affect the loop, but a `return` will.)
  */
-struct ContinueOrBreakInfo {
+struct LoopControlFlowInfo {
     bool fHasContinue = false;
     bool fHasBreak = false;
+    bool fHasReturn = false;
 };
-ContinueOrBreakInfo HasContinueOrBreak(const Statement& stmt);
+LoopControlFlowInfo GetLoopControlFlowInfo(const Statement& stmt);
 
 /**
  * Returns true if the expression can be assigned-into. Pass `info` if you want to know the
diff --git a/src/sksl/analysis/BUILD.bazel b/src/sksl/analysis/BUILD.bazel
index 32bfb77..a434348 100644
--- a/src/sksl/analysis/BUILD.bazel
+++ b/src/sksl/analysis/BUILD.bazel
@@ -8,9 +8,9 @@
     "SkSLCanExitWithoutReturningValue.cpp",
     "SkSLCheckProgramStructure.cpp",
     "SkSLFinalizationChecks.cpp",
+    "SkSLGetLoopControlFlowInfo.cpp",
     "SkSLGetLoopUnrollInfo.cpp",
     "SkSLGetReturnComplexity.cpp",
-    "SkSLHasContinueOrBreak.cpp",
     "SkSLHasSideEffects.cpp",
     "SkSLIsConstantExpression.cpp",
     "SkSLIsDynamicallyUniformExpression.cpp",
diff --git a/src/sksl/analysis/SkSLGetLoopControlFlowInfo.cpp b/src/sksl/analysis/SkSLGetLoopControlFlowInfo.cpp
new file mode 100644
index 0000000..65c9e5e4
--- /dev/null
+++ b/src/sksl/analysis/SkSLGetLoopControlFlowInfo.cpp
@@ -0,0 +1,79 @@
+/*
+ * Copyright 2023 Google LLC
+ *
+ * Use of this source code is governed by a BSD-style license that can be
+ * found in the LICENSE file.
+ */
+
+#include "src/sksl/SkSLAnalysis.h"
+
+#include "include/private/SkSLIRNode.h"
+#include "include/private/SkSLStatement.h"
+#include "src/sksl/analysis/SkSLProgramVisitor.h"
+
+namespace SkSL {
+
+class Expression;
+
+namespace Analysis {
+namespace {
+
+class LoopControlFlowVisitor : public ProgramVisitor {
+public:
+    LoopControlFlowVisitor() {}
+
+    bool visitExpression(const Expression& expr) override {
+        // We can avoid processing expressions entirely.
+        return false;
+    }
+
+    bool visitStatement(const Statement& stmt) override {
+        switch (stmt.kind()) {
+            case Statement::Kind::kContinue:
+                // A continue only affects the control flow of the loop if it's not nested inside
+                // another looping structure. (Inside a switch, SkSL disallows continue entirely.)
+                fResult.fHasContinue |= (fDepth == 0);
+                break;
+
+            case Statement::Kind::kBreak:
+                // A break only affects the control flow of the loop if it's not nested inside
+                // another loop/switch structure.
+                fResult.fHasBreak |= (fDepth == 0);
+                break;
+
+            case Statement::Kind::kReturn:
+                // A return will abort the loop's control flow no matter how deeply it is nested.
+                fResult.fHasReturn = true;
+                break;
+
+            case Statement::Kind::kFor:
+            case Statement::Kind::kDo:
+            case Statement::Kind::kSwitch: {
+                ++fDepth;
+                bool done = ProgramVisitor::visitStatement(stmt);
+                --fDepth;
+                return done;
+            }
+
+            default:
+                return ProgramVisitor::visitStatement(stmt);
+        }
+
+        // If we've already found everything we're hunting for, we can stop searching early.
+        return fResult.fHasContinue && fResult.fHasBreak && fResult.fHasReturn;
+    }
+
+    LoopControlFlowInfo fResult;
+    int fDepth = 0;
+};
+
+}  // namespace
+
+LoopControlFlowInfo GetLoopControlFlowInfo(const Statement& stmt) {
+    LoopControlFlowVisitor visitor;
+    visitor.visitStatement(stmt);
+    return visitor.fResult;
+}
+
+}  // namespace Analysis
+}  // namespace SkSL
diff --git a/src/sksl/analysis/SkSLHasContinueOrBreak.cpp b/src/sksl/analysis/SkSLHasContinueOrBreak.cpp
deleted file mode 100644
index 466740f..0000000
--- a/src/sksl/analysis/SkSLHasContinueOrBreak.cpp
+++ /dev/null
@@ -1,68 +0,0 @@
-/*
- * Copyright 2023 Google LLC
- *
- * Use of this source code is governed by a BSD-style license that can be
- * found in the LICENSE file.
- */
-
-#include "src/sksl/SkSLAnalysis.h"
-
-#include "include/private/SkSLIRNode.h"
-#include "include/private/SkSLStatement.h"
-#include "src/sksl/analysis/SkSLProgramVisitor.h"
-
-namespace SkSL {
-
-class Expression;
-
-namespace Analysis {
-namespace {
-
-class HasContinueOrBreakVisitor : public ProgramVisitor {
-public:
-    HasContinueOrBreakVisitor() {}
-
-    bool visitExpression(const Expression& expr) override {
-        // We can avoid processing expressions entirely.
-        return false;
-    }
-
-    bool visitStatement(const Statement& stmt) override {
-        switch (stmt.kind()) {
-            case Statement::Kind::kContinue:
-                fResult.fHasContinue = true;
-                return fResult.fHasContinue && fResult.fHasBreak;
-
-            case Statement::Kind::kBreak:
-                fResult.fHasBreak = true;
-                return fResult.fHasContinue && fResult.fHasBreak;
-
-            case Statement::Kind::kFor:
-            case Statement::Kind::kDo:
-            case Statement::Kind::kSwitch: {
-                // A continue or break inside of a for-statement or do-statement would not affect
-                // the control flow of the outer statement; they would affect the inner loop
-                // instead. A break inside of a switch-case statement escapes the switch, and SkSL
-                // does not allow continues inside of a switch at all, so we don't need to visit
-                // switches either.
-                return false;
-            }
-
-            default:
-                return ProgramVisitor::visitStatement(stmt);
-        }
-    }
-
-    ContinueOrBreakInfo fResult;
-};
-
-}  // namespace
-
-ContinueOrBreakInfo HasContinueOrBreak(const Statement& stmt) {
-    HasContinueOrBreakVisitor visitor;
-    visitor.visitStatement(stmt);
-    return visitor.fResult;
-}
-
-}  // namespace Analysis
-}  // namespace SkSL
diff --git a/src/sksl/codegen/SkSLRasterPipelineCodeGenerator.cpp b/src/sksl/codegen/SkSLRasterPipelineCodeGenerator.cpp
index 52ce363..c151b3f 100644
--- a/src/sksl/codegen/SkSLRasterPipelineCodeGenerator.cpp
+++ b/src/sksl/codegen/SkSLRasterPipelineCodeGenerator.cpp
@@ -859,7 +859,7 @@
     // Create a dedicated slot for continue-mask storage.
     SlotRange previousContinueMask = fCurrentContinueMask;
 
-    Analysis::ContinueOrBreakInfo loopInfo = Analysis::HasContinueOrBreak(*f.statement());
+    Analysis::LoopControlFlowInfo loopInfo = Analysis::GetLoopControlFlowInfo(*f.statement());
     if (loopInfo.fHasContinue) {
         fCurrentContinueMask =
                 fProgramSlots.createSlots(this->makeMaskName("for-loop continue mask"),