1 files changed, 33 insertions, 1 deletions

diff --git a/grailsort.c b/grailsort.c
index 9caf28f..2e0131a 100644
--- a/grailsort.c
+++ b/grailsort.c
@@ -83,6 +83,39 @@ void grailsort(void *unsorted, size_t nmel, size_t width, cmpfun cmp)
 	grailsort(buf, acount * blknmel + bufnmel, width, cmp);
 	grailsort(b,   bcount * blknmel,           width, cmp);
 
+	struct counts snapshot = counts[CURRENT];
+	int tmp = cmp(tail(a, acount*blknmel, width), b) <= 0;
+	add_counts(counts + CHECK_SORTED, &snapshot);
+	if (tmp) return;
+
+	snapshot = counts[CURRENT];
+	tmp = cmp(base, tail(b, bcount*blknmel, width)) >= 0;
+	if (tmp) {
+		rotate(base, nmel*width, acount*bwidth + bufnmel*width);
+		add_counts(counts + CHECK_SORTED, &snapshot);
+		return;
+	}
+	add_counts(counts + CHECK_SORTED, &snapshot);
+
+	/* steal some small elements from b for the buffer,
+	 * so the buffer has the bufnmel smallest elements */
+	size_t usebuf = 0, stealb = 0;
+	snapshot = counts[CURRENT];
+	do {
+		if (cmp(buf + usebuf*width, b + stealb*width) <= 0) {
+			usebuf++;
+		} else {
+			stealb++;
+		}
+	} while (usebuf + stealb < bufnmel);
+	if (stealb) {
+		swap(a - stealb*width, b, stealb*width);
+		add_counts(counts + STEAL_BUF, &snapshot);
+		grailsort(b, bcount * blknmel, width, cmp);
+	} else {
+		add_counts(counts + STEAL_BUF, &snapshot);
+	}
+
 	/* sort all the a and b blocks together by their head elements */
 	grailsort(a, blocks, bwidth, cmp);
 
@@ -102,7 +135,6 @@ void grailsort(void *unsorted, size_t nmel, size_t width, cmpfun cmp)
 	}
 
 	grailsort(buf, bufnmel, width, cmp);
-	distribute_buffer(buf, bufnmel, nmel - bufnmel, width, cmp);
 
 #if 0
 	assert_sorted(base, nmel, width, cmp);