#lang scheme
(require
 (planet murphy/amb:1:1/amb)
 "hierarchy.ss")

(define-struct overloads
  (default-signature
   ref-hierarchy
   methods
   preferred-methods
   [last-hierarchy #:mutable]
   [cached-methods #:mutable]))

(define (make-overloads* #:default [default-signature #f]
                         #:hierarchy [ref-hierarchy global-hierarchy])
  (make-overloads
   default-signature
   ref-hierarchy
   #hash() #hash()
   (ref-hierarchy)
   #hash()))

(provide/contract
 [overloads? (any/c . -> . boolean?)]
 [overloads-default-signature (overloads? . -> . any/c)]
 [rename make-overloads* make-overloads
         (() (#:default any/c #:hierarchy (-> hierarchy?)) . ->* . overloads?)])

(define (set-method os signature method)
  (match os
    [(struct overloads (default-signature
                        ref-hierarchy
                        methods preferred-methods
                        _ _))
     (make-overloads
      default-signature
      ref-hierarchy
      (hash-set methods signature method)
      preferred-methods
      (ref-hierarchy)
      #hash())]))

(define (remove-method os signature)
  (match os
    [(struct overloads (default-signature
                        ref-hierarchy
                        methods preferred-methods
                        _ _))
     (make-overloads
      default-signature
      ref-hierarchy
      (hash-remove methods signature)
      preferred-methods
      (ref-hierarchy)
      #hash())]))

(define (prefer-method os signature-a signature-b)
  (match os
    [(struct overloads (default-signature
                        ref-hierarchy
                        methods preferred-methods
                        _ _))
     (make-overloads
      default-signature
      ref-hierarchy
      methods
      (hash-set preferred-methods (cons signature-a signature-b)
        #t)
      (ref-hierarchy)
      #hash())]))

(define (unprefer-method os signature-a signature-b)
  (match os
    [(struct overloads (default-signature
                        ref-hierarchy
                        methods preferred-methods
                        _ _))
     (make-overloads
      default-signature
      ref-hierarchy
      methods
      (hash-remove preferred-methods (cons signature-a signature-b))
      (ref-hierarchy)
      #hash())]))

(provide/contract
 [set-method (overloads? any/c procedure? . -> . overloads?)]
 [remove-method (overloads? any/c . -> . overloads?)]
 [prefer-method (overloads? any/c any/c . -> . overloads?)]
 [unprefer-method (overloads? any/c any/c . -> . overloads?)])

(define-struct (exn:fail:multimethod exn:fail)
  (overloads signature)
  #:transparent)

(provide/contract
 [struct (exn:fail:multimethod exn:fail)
   ([message string?]
    [continuation-marks continuation-mark-set?]
    [overloads overloads?]
    [signature any/c])])

(define (find-methods os signature)
  (let ([h ((overloads-ref-hierarchy os))])
    ;; Reset the cache, if necessary
    (unless (equal? h (overloads-last-hierarchy os))
      (set-overloads-cached-methods! os
        #hash())
      (set-overloads-last-hierarchy! os
        h))
    (match os
      [(struct overloads (default-signature
                          _
                          methods preferred-methods
                          _
                          cached-methods))
       (local [;; Define a cache updater
               (define (update-cache! methods)
                 (set-overloads-cached-methods! os
                   (hash-set cached-methods signature methods))
                 methods)
               ;; Define a comparison for inheritance or preference order
               (define (better? signature-a signature-b)
                 (or (derived? signature-a signature-b)
                     (hash-ref
                      preferred-methods (cons signature-a signature-b) #f)
                     (for/or ([candidates (in-hash-keys preferred-methods)])
                       (and (derived? signature-a (car candidates))
                            (derived? (cdr candidates) signature-b)))))
               (define (strictly-better? signature-a signature-b)
                 (and (better? signature-a signature-b)
                      (not (better? signature-b signature-a))))
               (define (strictly-sorted? head rest)
                 (let ([signature (car head)])
                   (for/and ([candidate (in-list rest)])
                     (strictly-better? signature (car candidate)))))]
         ;; Now start looking for a matching method:
         (call-with-amb-prompt
          (λ ()
            (let ([signature (amb signature (overloads-default-signature os))])
              (amb
               ;; (1) Try the cache
               (hash-ref cached-methods signature amb-fail)
               ;; (2) Search for a method and update the cache
               (update-cache!
                (match (sort
                        (append
                         ;; (2.1) Try the direct method table
                         (cond
                           [(hash-ref methods signature #f)
                            => (λ (method)
                                 (list (cons signature method)))]
                           [else
                            null])
                         ;; (2.2) Try direct matches for signature ancestors
                         (for/list ([candidate (in-hash-keys
                                                (ancestors h signature))]
                                    #:when #t
                                    [method (in-value
                                             (hash-ref methods candidate #f))]
                                    #:when method)
                           (cons candidate method))
                         ;; (2.3) Try any derived signatures
                         (if (or (class? signature) (interface? signature)
                                 (dict? signature))
                             (for/list ([(candidate method) (in-hash methods)]
                                        #:when (and (not
                                                     (equal? signature candidate))
                                                    (derived?
                                                     h signature candidate)))
                               (cons candidate method))
                             null))
                        ;; (2.4) Sort by quality of match and preference
                        better? #:key car)
                  ;; (2.5) If there is only one match, use it.
                  [(list (cons signature method))
                   (list method)]
                  ;; (2.6) If there are multiple matches, limit the list to its
                  ;;       strictly sorted head. Signal an error if that head is
                  ;;       empty.
                  [(list-rest head rest)
                   (let ([methods (let more ([head head] [rest rest])
                                    (if (strictly-sorted? head rest)
                                        (cons (cdr head)
                                              (if (pair? rest)
                                                  (more (car rest) (cdr rest))
                                                  null))
                                        null))])
                     (if (null? methods)
                         (raise (make-exn:fail:multimethod
                                 (format
                                  "find-methods: ambiguous methods for signature ~e, possible matches ~e"
                                  signature
                                  (let ([signature (car head)]
                                        [candidates (map car rest)])
                                    (cons
                                     signature
                                     (filter
                                      (λ (candidate)
                                        (not
                                         (strictly-better? signature candidate)))
                                      candidates))))
                                 (current-continuation-marks)
                                 os signature))
                         methods))]
                  ;; (2.7) If there are no matches, continue searching.
                  [(list)
                   (amb)])))))
          ;; If no matching method was found at all, signal an error.
          (λ ()
            (raise (make-exn:fail:multimethod
                    (format
                     "find-methods: no method for signature ~e"
                     signature)
                    (current-continuation-marks)
                    os signature)))))])))

(provide/contract
 [find-methods (overloads? any/c . -> . (listof procedure?))])