(module weak-map mzscheme
  (require (lib "contract.ss"))
  
  ;; Basically a copy-and-paste of the example of
  ;; ephemerons from the mzscheme reference manual.
  ;; Provides a weak map where the keys are held weakly.
  ;; If the key is only weakly reachable, the item is
  ;; removed from the map.
  
  (define-struct weak-map (ht))
  (provide struct:weak-map weak-map? weak-map)
  
  (provide/contract [rename -make-weak-map make-weak-map
                            (case->
                             (-> weak-map?)
                             ((one-of/c 'equal) . -> . weak-map?))])
  ;; This weak map is like a weak hash table, but
  ;; without the key-in-value problem:
  (define -make-weak-map
    (case-lambda
      [()
       (make-weak-map (make-hash-table 'weak))]
      [(flag)
       (make-weak-map (make-hash-table flag 'weak))]))
  
  
  
  (provide/contract [weak-map-put! (weak-map? any/c any/c . -> . any)])
  (define (weak-map-put! m k v)
    (hash-table-put! (weak-map-ht m) k (make-ephemeron k (box v))))
  
  
  
  (provide/contract [weak-map-get (case->
                                   (weak-map? any/c . -> . any)
                                   (weak-map? any/c any/c . -> . any))])
  (define weak-map-get
    (case-lambda
      [(m k)
       (weak-map-get
        m k (lambda () (error 'weak-map-get "no value found for key ~a" k)))]
      [(m k thunk-or-value)
       (let ([v (hash-table-get (weak-map-ht m) k (lambda () #f))])
         (if v
             (let ([v (ephemeron-value v)])
               (if v
                   (unbox v)
                   (if (procedure? thunk-or-value)
                       (thunk-or-value)
                       thunk-or-value)))
             (if (procedure? thunk-or-value)
                 (thunk-or-value)
                 thunk-or-value)))]))
  
  
  
  (provide/contract [weak-map-for-each (weak-map? (any/c any/c . -> . any) . -> . any)])
  (define (weak-map-for-each m f)
    (hash-table-for-each
     (weak-map-ht m)
     (lambda (k weak-v)
       (let ([v (ephemeron-value weak-v)])
         (when v
           (f k (unbox v)))))))
  
  
  
  (provide/contract [weak-map-map (weak-map? (any/c any/c . -> . any) . -> . (listof any/c))])
  (define (weak-map-map m f)
    (define results '())
    (weak-map-for-each
     m
     (lambda (k v)
       (set! results (cons (f k v) results))))
    (reverse! results)))