#lang racket


(require ; racket/fixnum
         racket/future
         future-visualizer)

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

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

(define (fxvector-gen3-sort! unsorted)
  (define unsorted-length (fxvector-length unsorted))
  (define scratchpad (make-fxvector unsorted-length))
  ;(define total-merges (- unsorted-length 2))

  (define (merge! from1 to1/from2 to2 depth)
    (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
           ((fx< 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))
    (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 (fxquotient cnt 2))
           (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))
                               #f)))
                        (f2 (future
                             (λ ()
                               (sort-step to1/from2 to2 (fx+ depth 1))
                               #f))))
                    (or (touch f1)
                        (touch f2))))
                 (else
                  (sort-step from1 to1/from2 (fx+ depth 1))
                  (sort-step to1/from2 to2 (fx+ depth 1))
                  ))
           (merge! from1 to1/from2 to2 depth))
          ((fx= cnt 2)
           ; =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))
           ; COMPARE
           (define v1first (fx< 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))))
          ))
  
  (sort-step 0 unsorted-length)
  
  unsorted)

(define num-nums (expt 2 28))

(random-seed 1)
(define testvec2 (make-fxvector num-nums))
(let ((start (current-inexact-milliseconds)))
  (for ((i num-nums))
    (fxvector-set! testvec2 i (random 1000000000)))
  (define end (current-inexact-milliseconds))
  (displayln (~a "Generated in " (- end start) " ms")))

;(visualize-futures
 (let ((start (current-inexact-milliseconds)))
   (fxvector-gen3-sort! testvec2)
   (define end (current-inexact-milliseconds))
   (displayln (~a "Sorting time = " (- end start) " ms")))
; #f)