#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-method 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! method)
                 (set-overloads-cached-methods! os
                   (hash-set cached-methods signature method))
                 method)
               ;; 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)))))]
         ;; 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) Try the direct method table and update the cache
               (update-cache!
                (hash-ref methods signature amb-fail))
               ;; (3) Search for a method and update the cache
               (update-cache!
                (match (sort
                        (append
                         ;; (3.1) Try direct matches for signature ancestors
                         (for/list ([candidate (in-hash-keys
                                                (ancestors h signature))])
                           (cons candidate
                                 (hash-ref methods candidate amb-fail)))
                         ;; (3.2) Try any derived signatures
                         (if (or (class? signature) (interface? signature)
                                 (dict? signature))
                             (for/list ([(candidate method) (in-hash methods)]
                                        #:when (derived? h signature candidate))
                               (cons candidate method))
                             null))
                        ;; (3.3) Sort by quality of match and preference
                        better? #:key car)
                  ;; (3.4) If there are multiple matches, but the first one is
                  ;;       truly better than the others, use it. Signal an error
                  ;;       otherwise.
                  [(list-rest (cons signature method) more)
                   (if (for/and ([candidate (in-list more)])
                         (let ([candidate (car candidate)])
                           (and (better? signature candidate)
                                (not (better? candidate signature)))))
                       method
                       (raise (make-exn:fail:multimethod
                               (format
                                "find-method: ambiguous methods for signature ~e"
                                signature)
                               (current-continuation-marks)
                               os signature)))]
                  ;; (3.5) If there is only one match, use it.
                  [(list (cons signature method))
                   method]
                  ;; (3.6) 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-method: no method for signature ~e"
                     signature)
                    (current-continuation-marks)
                    os signature)))))])))

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