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

(define-struct multimethod
  (make-signature
   default-signature
   ref-hierarchy
   methods
   preferred-methods
   [last-hierarchy #:mutable]
   [cached-methods #:mutable])
  #:property prop:procedure
  (make-keyword-procedure
   (λ (kw-args kw-vals m . args)
     (keyword-apply
       (find-method
        m (keyword-apply
            (multimethod-make-signature m)
            kw-args kw-vals args))
       kw-args kw-vals args))))

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

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

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

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

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

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

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

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

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

(define (find-method m signature)
  (let ([h ((multimethod-ref-hierarchy m))])
    ;; Reset the cache, if necessary
    (unless (equal? h (multimethod-last-hierarchy m))
      (set-multimethod-cached-methods! m
        #hash())
      (set-multimethod-last-hierarchy! m
        h))
    (match m
      [(struct multimethod (make-signature
                            default-signature
                            _
                            methods preferred-methods
                            _
                            cached-methods))
       (local [;; Define a cache updater
               (define (update-cache! method)
                 (set-multimethod-cached-methods! m
                   (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? signature-b (cdr candidates))))))]
         ;; Now start looking for a matching method:
         (call-with-amb-prompt
          (λ ()
            (let ([signature (amb signature (multimethod-default-signature m))])
              (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. If the first
                  ;;       match is as good as the second, signal an error.
                  [(list-rest (cons signature-a method-a)
                              (cons signature-b method-b)
                              more)
                   (if (and (better? signature-a signature-b)
                            (not (better? signature-b signature-a)))
                       method-a
                       (raise (make-exn:fail:multimethod
                               (format
                                "~a: ambiguous methods for signature ~e"
                                (object-name m) signature)
                               (current-continuation-marks)
                               m 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
                     "~a: no method for signature ~e"
                     (object-name m) signature)
                    (current-continuation-marks)
                    m signature)))))])))

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

(define-struct pmultimethod
  (parameter)
  #:property prop:procedure
  (make-keyword-procedure
   (λ (kw-args kw-vals pm . args)
     (keyword-apply (pmultimethod->multimethod pm)
       kw-args kw-vals args))))

(define (pmultimethod->multimethod pm)
  ((pmultimethod-parameter pm)))

(define (make-pmultimethod* m)
  (make-pmultimethod (make-parameter m)))

(provide/contract
 [pmultimethod? (any/c . -> . boolean?)]
 [pmultimethod-parameter (pmultimethod? . -> . (parameter/c multimethod?))]
 [pmultimethod->multimethod (pmultimethod? . -> . multimethod?)]
 [rename make-pmultimethod* make-pmultimethod
         (multimethod? . -> . pmultimethod?)])

(define-syntax define-multi
  (syntax-rules (::)
    [(define-multi (name . args) :: signature options ...)
     (define-multi name :: (λ args signature)
       options ...)]
    [(define-multi name :: signature options ...)
     (define name
       (make-pmultimethod* (make-multimethod* signature options ...)))]))

(define-syntax parameterize-multi
  (syntax-rules ()
    [(parameterize-multi (name ...)
       expr ...)
     (parameterize ([(pmultimethod-parameter name) (pmultimethod->multimethod name)]
                    ...)
       expr ...)]))

(define-syntax define-method
  (syntax-rules (::)
    [(define-method (name . args) #:default
       expr ...)
     (define-method name #:default
       (λ args
         expr ...))]
    [(define-method name #:default method)
     (let ([p (pmultimethod-parameter name)])
       (p (set-method (p) (multimethod-default-signature (p)) method)))]
    [(define-method (name . args) :: signature
       expr ...)
     (define-method name :: `signature
       (λ args
         expr ...))]
    [(define-method name :: signature method)
     (let ([p (pmultimethod-parameter name)])
       (p (set-method (p) signature method)))]))

(provide
 define-multi
 parameterize-multi
 define-method)