@@ -1230,45 +1230,61 @@ void ConstraintSystem::removeResultBuilderTransform(AnyFunctionRef fn) {
1230
1230
ASSERT (erased);
1231
1231
}
1232
1232
1233
- namespace {
1234
- class ReturnStmtFinder : public ASTWalker {
1235
- std::vector<ReturnStmt *> ReturnStmts;
1233
+ // / Walks the given brace statement and calls the given function reference on
1234
+ // / every occurrence of an explicit `return` statement.
1235
+ // /
1236
+ // / \param callback A function reference that takes a `return` statement and
1237
+ // / returns a boolean value indicating whether to abort the walk.
1238
+ // /
1239
+ // / \returns `true` if the walk was aborted, `false` otherwise.
1240
+ static bool walkExplicitReturnStmts (const BraceStmt *BS,
1241
+ function_ref<bool (ReturnStmt *)> callback) {
1242
+ class Walker : public ASTWalker {
1243
+ function_ref<bool (ReturnStmt *)> callback;
1244
+
1245
+ public:
1246
+ Walker (decltype (Walker::callback) callback) : callback(callback) {}
1247
+
1248
+ MacroWalking getMacroWalkingBehavior () const override {
1249
+ return MacroWalking::Arguments;
1250
+ }
1236
1251
1237
- public:
1238
- static std::vector<ReturnStmt *> find (const BraceStmt *BS) {
1239
- ReturnStmtFinder finder;
1240
- const_cast <BraceStmt *>(BS)->walk (finder);
1241
- return std::move (finder.ReturnStmts );
1242
- }
1252
+ PreWalkResult<Expr *> walkToExprPre (Expr *E) override {
1253
+ return Action::SkipNode (E);
1254
+ }
1243
1255
1244
- MacroWalking getMacroWalkingBehavior () const override {
1245
- return MacroWalking::Arguments;
1246
- }
1256
+ PreWalkResult<Stmt *> walkToStmtPre (Stmt *S) override {
1257
+ if (S->isImplicit ()) {
1258
+ return Action::SkipNode (S);
1259
+ }
1247
1260
1248
- PreWalkResult<Expr *> walkToExprPre (Expr *E) override {
1249
- return Action::SkipNode (E);
1250
- }
1261
+ auto *returnStmt = dyn_cast<ReturnStmt>(S);
1262
+ if (!returnStmt) {
1263
+ return Action::Continue (S);
1264
+ }
1251
1265
1252
- PreWalkResult<Stmt *> walkToStmtPre (Stmt *S) override {
1253
- // If we see a return statement, note it..
1254
- auto *returnStmt = dyn_cast<ReturnStmt>(S);
1255
- if (!returnStmt || returnStmt->isImplicit ())
1256
- return Action::Continue (S);
1266
+ if (callback (returnStmt)) {
1267
+ return Action::Stop ();
1268
+ }
1257
1269
1258
- ReturnStmts. push_back (returnStmt);
1259
- return Action::SkipNode (S);
1260
- }
1270
+ // Skip children & post walk and continue.
1271
+ return Action::SkipNode (S);
1272
+ }
1261
1273
1262
- // / Ignore patterns.
1263
- PreWalkResult<Pattern *> walkToPatternPre (Pattern *pat) override {
1264
- return Action::SkipNode (pat);
1265
- }
1266
- };
1267
- } // end anonymous namespace
1274
+ // / Ignore patterns.
1275
+ PreWalkResult<Pattern *> walkToPatternPre (Pattern *pat) override {
1276
+ return Action::SkipNode (pat);
1277
+ }
1278
+ };
1279
+
1280
+ Walker walker (callback);
1281
+
1282
+ return const_cast <BraceStmt *>(BS)->walk (walker) == nullptr ;
1283
+ }
1268
1284
1269
1285
bool BraceHasExplicitReturnStmtRequest::evaluate (Evaluator &evaluator,
1270
1286
const BraceStmt *BS) const {
1271
- return ! ReturnStmtFinder::find (BS). empty ( );
1287
+ return walkExplicitReturnStmts (BS, [](ReturnStmt *) { return true ; } );
1272
1288
}
1273
1289
1274
1290
bool AnyFunctionRef::bodyHasExplicitReturnStmt () const {
@@ -1287,7 +1303,14 @@ std::vector<ReturnStmt *> TypeChecker::findReturnStatements(AnyFunctionRef fn) {
1287
1303
return std::vector<ReturnStmt *>();
1288
1304
}
1289
1305
1290
- return ReturnStmtFinder::find (fn.getBody ());
1306
+ std::vector<ReturnStmt *> results;
1307
+
1308
+ walkExplicitReturnStmts (fn.getBody (), [&results](ReturnStmt *RS) {
1309
+ results.push_back (RS);
1310
+ return false ;
1311
+ });
1312
+
1313
+ return results;
1291
1314
}
1292
1315
1293
1316
ResultBuilderOpSupport TypeChecker::checkBuilderOpSupport (
0 commit comments