From cdd7083c8694696d92d71754e25d01d4d4fb6a04 Mon Sep 17 00:00:00 2001
From: primo-ppcg <mike.tryczak@gmail.com>
Date: Tue, 15 Aug 2023 11:15:38 +0700
Subject: [PATCH] special case common `sort` usages

---
 src/boot/boot.janet | 37 +++++++++++++++++++++++++------------
 1 file changed, 25 insertions(+), 12 deletions(-)

diff --git a/src/boot/boot.janet b/src/boot/boot.janet
index 9701f1b0..af4f4853 100644
--- a/src/boot/boot.janet
+++ b/src/boot/boot.janet
@@ -805,21 +805,31 @@
 ###
 ###
 
-(defn- median-of-three [a b c]
-  (if (not= (> a b) (> a c))
-    a
-    (if (not= (> b a) (> b c)) b c)))
+(defmacro- median-of-three
+  [x y z]
+  ~(if (<= ,x ,y)
+     (if (<= ,y ,z) ,y (if (<= ,z ,x) ,x ,z))
+     (if (<= ,z ,y) ,y (if (<= ,x ,z) ,x ,z))))
+
+(defmacro- sort-partition-template
+  [ind before? left right pivot]
+  ~(do
+     (while (,before? (in ,ind ,left) ,pivot) (++ ,left))
+     (while (,before? ,pivot (in ,ind ,right)) (-- ,right))))
 
 (defn- sort-help [a lo hi before?]
   (when (< lo hi)
-    (def pivot
-      (median-of-three (in a hi) (in a lo)
-                       (in a (math/floor (/ (+ lo hi) 2)))))
+    (def [x y z] [(in a lo)
+                  (in a (div (+ lo hi) 2))
+                  (in a hi)])
+    (def pivot (median-of-three x y z))
     (var left lo)
     (var right hi)
     (while true
-      (while (before? (in a left) pivot) (++ left))
-      (while (before? pivot (in a right)) (-- right))
+      (case before?
+        < (sort-partition-template a < left right pivot)
+        > (sort-partition-template a > left right pivot)
+        (sort-partition-template a before? left right pivot))
       (when (<= left right)
         (def tmp (in a left))
         (set (a left) (in a right))
@@ -827,8 +837,10 @@
         (++ left)
         (-- right))
       (if (>= left right) (break)))
-    (sort-help a lo right before?)
-    (sort-help a left hi before?))
+    (if (< lo right)
+      (sort-help a lo right before?))
+    (if (< left hi)
+      (sort-help a left hi before?)))
   a)
 
 (defn sort
@@ -836,7 +848,8 @@
   If a `before?` comparator function is provided, sorts elements using that,
   otherwise uses `<`.``
   [ind &opt before?]
-  (sort-help ind 0 (- (length ind) 1) (or before? <)))
+  (default before? <)
+  (sort-help ind 0 (- (length ind) 1) before?))
 
 (defn sort-by
   ``Sorts `ind` in-place by calling a function `f` on each element and