faster VECTOR-SUBSEQ*
[sbcl.git] / src / code / sort.lisp
index 3087025..06cd2c6 100644 (file)
 ;;; allows the compiler to make enough optimizations that it might be
 ;;; worth the (large) cost in space.
 (declaim (maybe-inline sort))
-(defun sort (sequence predicate &key key)
+(defun sort (sequence predicate &rest args &key key)
   #!+sb-doc
   "Destructively sort SEQUENCE. PREDICATE should return non-NIL if
    ARG1 is to precede ARG2."
+  (declare (truly-dynamic-extent args))
   (let ((predicate-fun (%coerce-callable-to-fun predicate)))
-    (typecase sequence
-      (list
-       (stable-sort-list sequence
-                         predicate-fun
-                         (if key (%coerce-callable-to-fun key) #'identity)))
-      (vector
-       (let ((key-fun-or-nil (and key (%coerce-callable-to-fun key))))
-         (with-array-data ((vector (the vector sequence))
-                           (start 0)
-                           (end (length sequence)))
-           (sort-vector vector start end predicate-fun key-fun-or-nil)))
-       sequence)
-      (t
-       (error 'simple-type-error
-              :datum sequence
-              :expected-type 'sequence
-              :format-control "~S is not a sequence."
-              :format-arguments (list sequence))))))
+    (seq-dispatch sequence
+      (stable-sort-list sequence
+                        predicate-fun
+                        (if key (%coerce-callable-to-fun key) #'identity))
+      (let ((key-fun-or-nil (and key (%coerce-callable-to-fun key))))
+        (with-array-data ((vector (the vector sequence))
+                          (start)
+                          (end)
+                          :check-fill-pointer t)
+          (sort-vector vector start end predicate-fun key-fun-or-nil))
+        sequence)
+      (apply #'sb!sequence:sort sequence predicate args))))
 \f
 ;;;; stable sorting
-
-(defun stable-sort (sequence predicate &key key)
+(defun stable-sort (sequence predicate &rest args &key key)
   #!+sb-doc
   "Destructively sort SEQUENCE. PREDICATE should return non-NIL if
    ARG1 is to precede ARG2."
+  (declare (truly-dynamic-extent args))
   (let ((predicate-fun (%coerce-callable-to-fun predicate)))
-    (typecase sequence
-      (simple-vector
-       (stable-sort-simple-vector sequence
-                                  predicate-fun
-                                  (and key (%coerce-callable-to-fun key))))
-      (list
-       (stable-sort-list sequence
-                         predicate-fun
-                         (if key (%coerce-callable-to-fun key) #'identity)))
-      (vector
-       (stable-sort-vector sequence
-                           predicate-fun
-                           (and key (%coerce-callable-to-fun key))))
-      (t
-       (error 'simple-type-error
-              :datum sequence
-              :expected-type 'sequence
-              :format-control "~S is not a sequence."
-              :format-arguments (list sequence))))))
-  \f
+    (seq-dispatch sequence
+      (stable-sort-list sequence
+                        predicate-fun
+                        (if key (%coerce-callable-to-fun key) #'identity))
+      (if (typep sequence 'simple-vector)
+          (stable-sort-simple-vector sequence
+                                     predicate-fun
+                                     (and key (%coerce-callable-to-fun key)))
+          (stable-sort-vector sequence
+                              predicate-fun
+                              (and key (%coerce-callable-to-fun key))))
+      (apply #'sb!sequence:stable-sort sequence predicate args))))
+\f
 ;;; FUNCALL-USING-KEY saves us a function call sometimes.
 (eval-when (:compile-toplevel :execute)
   (sb!xc:defmacro funcall2-using-key (pred key one two)
 ;;; (NIL).
 (sb!xc:defmacro vector-merge-sort (vector pred key vector-ref)
   (with-unique-names
-      (vector-len n direction unsorted start-1 end-1 end-2 temp temp-len i)
+      (vector-len n direction unsorted start-1 end-1 end-2 temp i)
     `(let* ((,vector-len (length (the vector ,vector)))
-            (,n 1)        ; bottom-up size of contiguous runs to be merged
-            (,direction t) ; t vector --> temp    nil temp --> vector
-            (,temp *merge-sort-temp-vector*)
-            (,temp-len (length ,temp))
-            (,unsorted 0)  ; unsorted..vector-len are the elements that need
-                           ; to be merged for a given n
-            (,start-1 0))  ; one n-len subsequence to be merged with the next
-       (declare (fixnum ,vector-len ,n ,temp-len ,unsorted ,start-1)
+            (,n 1)            ; bottom-up size of contiguous runs to be merged
+            (,direction t)    ; t vector --> temp    nil temp --> vector
+            (,temp (make-array ,vector-len))
+            (,unsorted 0)   ; unsorted..vector-len are the elements that need
+                                        ; to be merged for a given n
+            (,start-1 0))   ; one n-len subsequence to be merged with the next
+       (declare (fixnum ,vector-len ,n ,unsorted ,start-1)
                 (simple-vector ,temp))
-       (if (> ,vector-len ,temp-len)
-           (setf ,temp (make-array (max ,vector-len
-                                        (min most-positive-fixnum
-                                             (+ ,temp-len ,temp-len))))
-                 *merge-sort-temp-vector* ,temp))
-       ;; rebind, in case PRED or KEY calls STABLE-SORT
-       (let ((*merge-sort-temp-vector* (vector)))
+       (loop
+         ;; for each n, we start taking n-runs from the start of the vector
+         (setf ,unsorted 0)
          (loop
-            ;; for each n, we start taking n-runs from the start of the vector
-            (setf ,unsorted 0)
-            (loop
-               (setf ,start-1 ,unsorted)
-               (let ((,end-1 (+ ,start-1 ,n)))
-                 (declare (fixnum ,end-1))
-                 (cond ((< ,end-1 ,vector-len)
-                        ;; there are enough elements for a second run
-                        (let ((,end-2 (+ ,end-1 ,n)))
-                          (declare (fixnum ,end-2))
-                          (if (> ,end-2 ,vector-len) (setf ,end-2 ,vector-len))
-                          (setf ,unsorted ,end-2)
-                          (if ,direction
-                              (stable-sort-merge-vectors*
-                               ,vector ,temp
-                               ,start-1 ,end-1 ,end-2 ,pred ,key ,vector-ref svref)
-                              (stable-sort-merge-vectors*
-                               ,temp ,vector
-                               ,start-1 ,end-1 ,end-2 ,pred ,key svref ,vector-ref))
-                          (if (= ,unsorted ,vector-len) (return))))
-                       ;; if there is only one run, copy those elements to the end
-                       (t (if ,direction
-                              (do ((,i ,start-1 (1+ ,i)))
-                                  ((= ,i ,vector-len))
-                                (declare (fixnum ,i))
-                               (setf (svref ,temp ,i)
-                                     (,vector-ref ,vector ,i)))
-                             (do ((,i ,start-1 (1+ ,i)))
-                                 ((= ,i ,vector-len))
-                               (declare (fixnum ,i))
-                               (setf (,vector-ref ,vector ,i)
-                                     (svref ,temp ,i))))
-                         (return)))))
-           ;; If the inner loop only executed once, then there were only enough
-           ;; elements for two subsequences given n, so all the elements have
-           ;; been merged into one list. Start-1 will have remained 0 upon exit.
-           (when (zerop ,start-1)
-             (if ,direction
-                 ;; if we just merged into the temporary, copy it all back
-                 ;; to the given vector.
-                 (dotimes (,i ,vector-len)
-                   (setf (,vector-ref ,vector ,i)
-                         (svref ,temp ,i))))
-             (return ,vector))
-           (setf ,n (ash ,n 1))         ; (* 2 n)
-           (setf ,direction (not ,direction)))))))
+           (setf ,start-1 ,unsorted)
+           (let ((,end-1 (+ ,start-1 ,n)))
+             (declare (fixnum ,end-1))
+             (cond ((< ,end-1 ,vector-len)
+                    ;; there are enough elements for a second run
+                    (let ((,end-2 (+ ,end-1 ,n)))
+                      (declare (fixnum ,end-2))
+                      (if (> ,end-2 ,vector-len) (setf ,end-2 ,vector-len))
+                      (setf ,unsorted ,end-2)
+                      (if ,direction
+                          (stable-sort-merge-vectors*
+                           ,vector ,temp
+                           ,start-1 ,end-1 ,end-2 ,pred ,key ,vector-ref svref)
+                          (stable-sort-merge-vectors*
+                           ,temp ,vector
+                           ,start-1 ,end-1 ,end-2 ,pred ,key svref ,vector-ref))
+                      (if (= ,unsorted ,vector-len) (return))))
+                   ;; if there is only one run, copy those elements to the end
+                   (t (if ,direction
+                          (do ((,i ,start-1 (1+ ,i)))
+                              ((= ,i ,vector-len))
+                            (declare (fixnum ,i))
+                            (setf (svref ,temp ,i) (,vector-ref ,vector ,i)))
+                          (do ((,i ,start-1 (1+ ,i)))
+                              ((= ,i ,vector-len))
+                            (declare (fixnum ,i))
+                            (setf (,vector-ref ,vector ,i) (svref ,temp ,i))))
+                      (return)))))
+         ;; If the inner loop only executed once, then there were only enough
+         ;; elements for two subsequences given n, so all the elements have
+         ;; been merged into one list. Start-1 will have remained 0 upon exit.
+         (when (zerop ,start-1)
+           (when ,direction
+             ;; if we just merged into the temporary, copy it all back
+             ;; to the given vector.
+             (dotimes (,i ,vector-len)
+               (setf (,vector-ref ,vector ,i) (svref ,temp ,i))))
+           ;; Kill the new vector to prevent garbage from being retained.
+           (%shrink-vector ,temp 0)
+           (return ,vector))
+         (setf ,n (ash ,n 1))           ; (* 2 n)
+         (setf ,direction (not ,direction))))))
 
 ) ; EVAL-when
 
-;;; temporary vector for stable sorting vectors, allocated for each new thread
-(defvar *merge-sort-temp-vector* (make-array 50))
-
-(declaim (simple-vector *merge-sort-temp-vector*))
-
 (defun stable-sort-simple-vector (vector pred key)
   (declare (type simple-vector vector)
            (type function pred)
               (vector-2 (coerce sequence2 'vector))
               (length-1 (length vector-1))
               (length-2 (length vector-2))
-              (result (make-sequence result-type
-                                     (+ length-1 length-2))))
+              (result (make-sequence result-type (+ length-1 length-2))))
          (declare (vector vector-1 vector-2)
                   (fixnum length-1 length-2))
          (if (and (simple-vector-p result)
                             result predicate key svref)
              (merge-vectors vector-1 length-1 vector-2 length-2
                             result predicate key aref))))
+      ((and (csubtypep type (specifier-type 'sequence))
+            (find-class result-type nil))
+       (let* ((vector-1 (coerce sequence1 'vector))
+              (vector-2 (coerce sequence2 'vector))
+              (length-1 (length vector-1))
+              (length-2 (length vector-2))
+              (temp (make-array (+ length-1 length-2)))
+              (result (make-sequence result-type (+ length-1 length-2))))
+         (declare (vector vector-1 vector-2) (fixnum length-1 length-2))
+         (merge-vectors vector-1 length-1 vector-2 length-2
+                        temp predicate key aref)
+         (replace result temp)
+         result))
       (t (bad-sequence-type-error result-type)))))