;;; Guile Knots
;;; Copyright © 2020 Christopher Baines <mail@cbaines.net>
;;;
;;; This file is part of Guile Knots.
;;;
;;; The Guile Knots is free software; you can redistribute it and/or
;;; modify it under the terms of the GNU General Public License as
;;; published by the Free Software Foundation; either version 3 of the
;;; License, or (at your option) any later version.
;;;
;;; The Guile Knots is distributed in the hope that it will be useful,
;;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
;;; General Public License for more details.
;;;
;;; You should have received a copy of the GNU General Public License
;;; along with the guix-data-service.  If not, see
;;; <http://www.gnu.org/licenses/>.

(define-module (knots parallelism)
  #:use-module (srfi srfi-1)
  #:use-module (srfi srfi-71)
  #:use-module (ice-9 match)
  #:use-module (fibers)
  #:use-module (fibers channels)
  #:use-module (fibers operations)
  #:export (fibers-batch-map
            fibers-map

            fibers-map-with-progress

            fibers-batch-for-each
            fibers-for-each

            fibers-parallel
            fibers-let))

;; Like split-at, but don't care about the order of the resulting lists, and
;; don't error if the list is shorter than i elements
(define (split-at* lst i)
  (let lp ((l lst) (n i) (acc '()))
    (if (or (<= n 0) (null? l))
        (values (reverse! acc) l)
        (lp (cdr l) (- n 1) (cons (car l) acc)))))

;; As this can be called with lists with tens of thousands of items in them,
;; batch the
(define (get-batch batch-size lists)
  (let ((split-lists
         (map (lambda (lst)
                (let ((batch rest (split-at* lst batch-size)))
                  (cons batch rest)))
              lists)))
    (values (map car split-lists)
            (map cdr split-lists))))

(define (defer-to-parallel-fiber thunk)
  (let ((reply (make-channel)))
    (spawn-fiber
     (lambda ()
       (with-exception-handler
           (lambda (exn)
             (put-message reply (cons 'exception exn)))
         (lambda ()
           (call-with-values
               (lambda ()
                 (with-throw-handler #t
                   thunk
                   (lambda _
                     (backtrace))))
             (lambda vals
               (put-message reply vals))))
         #:unwind? #t))
     #:parallel? #t)
    reply))

(define (fetch-result-of-defered-thunks . reply-channels)
  (let ((responses (map get-message
                        reply-channels)))
    (map
     (match-lambda
       (('exception . exn)
        (raise-exception exn))
       (result
        (apply values result)))
     responses)))

(define (fibers-batch-map proc batch-size . lists)
  (let loop ((lists lists)
             (result '()))
    (let ((batch
           rest
           (get-batch batch-size lists)))
      (if (any null? batch)
          result
          (let ((response-channels
                 (apply map
                        (lambda args
                          (defer-to-parallel-fiber
                            (lambda ()
                              (apply proc args))))
                        batch)))
            (loop rest
                  (append! result
                           (apply fetch-result-of-defered-thunks
                                  response-channels))))))))

(define (fibers-map proc . lists)
  (apply fibers-batch-map proc 20 lists))

(define (fibers-batch-for-each proc batch-size . lists)
  (let loop ((lists lists))
    (let ((batch
           rest
           (get-batch batch-size lists)))
      (if (any null? batch)
          *unspecified*
          (let ((response-channels
                 (apply map
                        (lambda args
                          (defer-to-parallel-fiber
                            (lambda ()
                              (apply proc args))))
                        batch)))
            (apply fetch-result-of-defered-thunks
                   response-channels)
            (loop rest))))))

(define (fibers-for-each proc . lists)
  (apply fibers-batch-for-each proc 20 lists))

(define-syntax fibers-parallel
  (lambda (x)
    (syntax-case x ()
      ((_ e0 ...)
       (with-syntax (((tmp0 ...) (generate-temporaries (syntax (e0 ...)))))
         #'(let ((tmp0 (defer-to-parallel-fiber
                         (lambda ()
                           e0)))
                 ...)
             (apply values (fetch-result-of-defered-thunks tmp0 ...))))))))

(define-syntax-rule (fibers-let ((v e) ...) b0 b1 ...)
  (call-with-values
      (lambda () (fibers-parallel e ...))
    (lambda (v ...)
      b0 b1 ...)))

(define* (fibers-map-with-progress proc lists #:key report)
  (let loop ((channels-to-results
              (apply map
                     (lambda args
                       (cons (defer-to-parallel-fiber
                               (lambda ()
                                 (apply proc args)))
                             #f))
                     lists)))
    (let ((active-channels
           (filter-map car channels-to-results)))
      (when report
        (report (apply map
                       list
                       (map cdr channels-to-results)
                       lists)))
      (if (null? active-channels)
          (map
           (match-lambda
             ((#f . ('exception . exn))
              (raise-exception exn))
             ((#f . ('result . val))
              val))
           channels-to-results)
          (loop
           (perform-operation
            (apply
             choice-operation
             (filter-map
              (lambda (p)
                (match p
                  ((channel . _)
                   (if channel
                       (wrap-operation
                        (get-operation channel)
                        (lambda (result)
                          (map (match-lambda
                                 ((c . r)
                                  (if (eq? channel c)
                                      (cons #f
                                            (match result
                                              (('exception . exn)
                                               result)
                                              (_
                                               (cons 'result result))))
                                      (cons c r))))
                               channels-to-results)))
                       #f))))
              channels-to-results))))))))