(module memoize mzscheme

  (define-syntax define/memo
    (syntax-rules ()
      [(_ (name) exp0 exp1 ...)
       (begin
         (define undefined (gensym))
         (define cached undefined)
         (define (name)
           (if (eq? cached undefined)
               (set! cached (begin exp0 exp1 ...))
               cached)))]
      [(_ (name arg) exp0 exp1 ...)
       (begin
         (define cache (make-hash-table))
         (define (name arg)
           (hash-table-get cache arg (lambda ()
                                       (let ([ans (begin exp0 exp1 ...)])
                                         (hash-table-put! cache arg ans)
                                         ans)))))]
      [(_ (name arg ...) exp0 exp1 ...)
       (begin
         (define cache (make-hash-table))
         (define (name arg ...)
           (let ([args (list arg ...)]
                 [key (bitwise-xor (eq-hash-code arg) ...)]
                 [alist (hash-table-get cache key (lambda () null))])
             (cond
               [(assoc args alist) => cdr]
               [else (let ([ans (begin exp0 exp1 ...)])
                       (hash-table-put! cache key (cons (cons args ans) alist))
                       ans)]))))]))

  (define-syntax memo-lambda
    (syntax-rules ()
      [(_ (arg) body ...)
       (let ([memo (make-hash-table)]
             [f (lambda (arg) body ...)])
         (lambda (arg)
           (hash-table-get memo arg (lambda ()
                                      (let ([ans (f arg)])
                                        (hash-table-put! memo arg ans)
                                        ans)))))]))

  (define-syntax memo-lambda*
    (syntax-rules ()
      [(_ (arg ...) body ...)
       (let ([memo (make-hash-table)]
             [f (lambda (arg ...) body ...)])
         (lambda (arg ...)
           (let* ([args (list arg ...)]
                  [key (bitwise-xor (eq-hash-code arg) ...)]
                  [alist (hash-table-get memo key (lambda () null))])
             (cond
               [(assoc args alist) => cdr]
               [else (let ([ans (f arg ...)])
                       (hash-table-put! memo key (cons (cons args ans) alist))
                       ans)]))))]))


  (provide memo-lambda memo-lambda*))