From cb8ceed2af12fb9eebec57bb4beabb67dced1a55 Mon Sep 17 00:00:00 2001
From: Karoy Lorentey <klorentey@apple.com>
Date: Wed, 18 Nov 2020 14:37:39 -0800
Subject: [PATCH] =?UTF-8?q?[stdlib][NFC]=20Make=20partition(by:)=E2=80=99s?=
 =?UTF-8?q?=20implementation=20easier=20to=20follow?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Also, fix a case in the description of the loop invariant
---
 stdlib/public/core/CollectionAlgorithms.swift | 39 ++++++++-----------
 1 file changed, 17 insertions(+), 22 deletions(-)

diff --git a/stdlib/public/core/CollectionAlgorithms.swift b/stdlib/public/core/CollectionAlgorithms.swift
index 529713b59a6..14c8083c151 100644
--- a/stdlib/public/core/CollectionAlgorithms.swift
+++ b/stdlib/public/core/CollectionAlgorithms.swift
@@ -401,35 +401,30 @@ extension MutableCollection where Self: BidirectionalCollection {
   ) rethrows -> Index {
     var lo = startIndex
     var hi = endIndex
+    while true {
+      // Invariants at this point:
+      //
+      // * `lo <= hi`
+      // * all elements in `startIndex ..< lo` belong in the first partition
+      // * all elements in `hi ..< endIndex` belong in the second partition
 
-    // 'Loop' invariants (at start of Loop, all are true):
-    // * lo < hi
-    // * predicate(self[i]) == false, for i in startIndex ..< lo
-    // * predicate(self[i]) == true, for i in hi ..< endIndex
+      // Find next element from `lo` that may not be in the right place.
+      while true {
+        guard lo < hi else { return lo }
+        if try belongsInSecondPartition(self[lo]) { break }
+        formIndex(after: &lo)
+      }
 
-    Loop: while true {
-      FindLo: repeat {
-        while lo < hi {
-          if try belongsInSecondPartition(self[lo]) { break FindLo }
-          formIndex(after: &lo)
-        }
-        break Loop
-      } while false
-
-      FindHi: repeat {
+      // Find next element down from `hi` that we can swap `lo` with.
+      while true {
         formIndex(before: &hi)
-        while lo < hi {
-          if try !belongsInSecondPartition(self[hi]) { break FindHi }
-          formIndex(before: &hi)
-        }
-        break Loop
-      } while false
+        guard lo < hi else { return lo }
+        if try !belongsInSecondPartition(self[hi]) { break }
+      }
 
       swapAt(lo, hi)
       formIndex(after: &lo)
     }
-
-    return lo
   }
 }