#lang at-exp racket

(require ; racket/fixnum
         racket/future
         scribble/srcdoc
         (for-doc scribble/manual
                  scribble-math/dollar))

(require racket/require
         (filtered-in
          (λ (name) (regexp-replace #rx"unsafe-" name ""))
          racket/unsafe/ops)
         (only-in
          racket/fixnum
          make-fxvector
          fxvector?))

(provide

 (proc-doc
  fxvector-futures-sort!
  (->i ((unsorted fxvector?))
       ((compare procedure?)
        #:progress-proc (progress-proc (or/c procedure? false?))
        #:progress-sleep (progress-sleep flonum?))
       (res fxvector?))
  ((λ (a b) (fx< a b)) #f 0.1)
  @{

    Sorts fxvector in place. If progress-proc is not #f, calls it
every progress-sleep seconds.

    The procedure uses merge sort with @${n} merge operations. Its
overall algorithmic time complexity is @${O(n\cdot\log_2 n)} and
memory complexity is @${O(n)}. The actual memory consumption is @${8n}
in addition to input of the same size on 64bit platforms.

The implementation uses all available processors.

If a custom compare function is provided, it should be a lambda term
and not a reference to some other function. For example, providing fx<
as compare blocks running in parallel, but using the default compare
function as is provides support for maximum parallelism.

    })

 (rename-out (futures-depth fxvector-util-futures-depth)
             (num-futures fxvector-util-num-futures))
 
 )

(define processors (processor-count))
(define futures-depth (inexact->exact (ceiling (log processors 2))))
(define num-futures (expt 2 futures-depth))

(define (fxvector-futures-sort! unsorted
                                (compare (λ (a b) (fx< a b)))
                                #:progress-proc (progress-proc #f)
                                #:progress-sleep (progress-sleep 0.1))
  (define unsorted-length (fxvector-length unsorted))
  (define scratchpad (make-fxvector unsorted-length))

  (define total-merges (fx- unsorted-length 1))

  ; Create progresses vector, start progress thread and create stop
  ; procedure
  (define (create-progresses+stop)
    (define progresses (make-fxvector num-futures))
    (define progress-running #t)
    (define progress-thread
      (thread
       (λ ()
         (let loop ()
           (when progress-running
             (define progress
               (let loop ((i 0)
                          (acc 0))
                 (if (fx< i num-futures)
                     (loop (fx+ i 1)
                           (fx+ acc (fxvector-ref progresses i)))
                     acc)))
             (progress-proc progress total-merges)
             (sleep progress-sleep)
             (loop))))))

    ; Stops the loop and waits for the thread
    (define (progress-stop)
      (set! progress-running #f)
      (thread-wait progress-thread)
      (progress-proc total-merges total-merges))

    ; Must be assigned at once
    (values progresses progress-stop))

  ; Create progresses vector and stop proc if required
  (define-values
    (progresses progress-stop)
    (if progress-proc
        (create-progresses+stop)
        (values #f #f)))

  ; Generic merge stage
  (define (merge! from1 to1/from2 to2 depth fid)
    (when progress-proc
      (fxvector-set! progresses fid (fx+ (fxvector-ref progresses fid) 1)))
    (define src (if (fx= (fxand depth 1) 1) unsorted scratchpad))
    (define dst (if (fx= (fxand depth 1) 1) scratchpad unsorted))
    (let loop1 ((i from1)
                (j to1/from2)
                (k from1))
      ; Always select correct element based on compare function
      (cond
        ((and (fx< i to1/from2)
              (fx< j to2))
         (define v1 (fxvector-ref src i))
         (define v2 (fxvector-ref src j))
         (cond
           ((compare v1 v2)
            (fxvector-set! dst k v1)
            (loop1 (fx+ i 1) j (fx+ k 1)))
           (else
            (fxvector-set! dst k v2)
            (loop1 i (fx+ j 1) (fx+ k 1)))))
        (else
         (let loop2 ((i i)
                     (j j)
                     (k k))
           ; Finish from the first part
           (cond
             ((fx< i to1/from2)
              (fxvector-set! dst k (fxvector-ref src i))
              (loop2 (fx+ i 1) j (fx+ k 1)))
             (else
              (let loop3 ((i i)
                          (j j)
                          (k k))
                ; Finish from the second part
                (when (fx< j to2)
                  (fxvector-set! dst k (fxvector-ref src j))
                  (loop3 i (fx+ j 1) (fx+ k 1)))))))))))
  
  (define (sort-step from to (depth 0) (fid 0))
    (define cnt (fx- to from))
    (cond ((fx> cnt 2)
           ; >2 means we do a proper split/merge
           ; this is the only part to leverage futures
           (define cnt2 (fxrshift cnt 1))
           (define from1 from)
           (define to1/from2 (fx+ from cnt2))
           (define to2 to)
           (cond ((fx< depth futures-depth)
                  (let ((f1 (future
                             (λ ()
                               (sort-step from1 to1/from2
                                          (fx+ depth 1)
                                          fid)
                               #f)))
                        (f2 (future
                             (λ ()
                               (sort-step to1/from2 to2
                                          (fx+ depth 1)
                                          (fxior fid (fxlshift 1 depth)))
                               #f))))
                    (or (touch f1)
                        (touch f2))))
                 (else
                  (sort-step from1 to1/from2 (fx+ depth 1) fid)
                  (sort-step to1/from2 to2 (fx+ depth 1) fid)
                  ))
           (merge! from1 to1/from2 to2 depth fid))
          ((fx= cnt 2)
           (when progress-proc
             (fxvector-set! progresses fid (fx+ (fxvector-ref progresses fid) 1)))
           ; =2 - swap in-place or from the other
           (define v1 (fxvector-ref unsorted from))
           (define v2 (fxvector-ref unsorted (fx+ from 1)))
           (define dst (if (fx= (fxand depth 1) 1) scratchpad unsorted))
           (define v1first (compare v1 v2))
           (fxvector-set! dst
                          from
                          (if v1first v1 v2))
           (fxvector-set! dst
                          (fx+ from 1)
                          (if v1first v2 v1)))
          ((fx= cnt 1)
           ; =1 - only copy if it must go to scratchpad
           (when (fx= (fxand depth 1) 1)
             (fxvector-set! scratchpad from (fxvector-ref unsorted from))))
          ))

  ; Start with whole fxvector
  (sort-step 0 unsorted-length)

  ; Stop progress thread
  (when progress-proc
    (progress-stop))

  ; Return result
  unsorted)