(module memoize mzscheme

  (define-syntax memo-lambda
    (syntax-rules ()
      [(_ () body0 body1 ...)
       (let* ([undefined (gensym)]
              [cached undefined])
         (lambda ()
           (when (eq? cached undefined)
             (set! cached (begin body0 body1 ...)))
           cached))]
      [(_ (arg) body0 body1 ...)
       (let ([cache (make-hash-table)])
         (lambda (arg)
           (hash-table-get cache arg (lambda ()
                                       (let ([ans (begin body0 body1 ...)])
                                         (hash-table-put! cache arg ans)
                                         ans)))))]
      [(_ (arg ...) body0 body1 ...)
       (let ([cache (make-hash-table)])
         (lambda (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 body0 body1 ...)])
                       (hash-table-put! cache key (cons (cons args ans) alist))
                       ans)]))))]))

  (define-syntax define/memo
    (syntax-rules ()
      [(_ (name) body0 body1 ...)
       (begin
         (define undefined (gensym))
         (define cached undefined)
         (define (name)
           (when (eq? cached undefined)
             (set! cached (begin body0 body1 ...)))
           cached))]
      [(_ (name arg) body0 body1 ...)
       (begin
         (define cache (make-hash-table))
         (define (name arg)
           (hash-table-get cache arg (lambda ()
                                       (let ([ans (begin body0 body1 ...)])
                                         (hash-table-put! cache arg ans)
                                         ans)))))]
      [(_ (name arg ...) body0 body1 ...)
       (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 body0 body1 ...)])
                       (hash-table-put! cache key (cons (cons args ans) alist))
                       ans)]))))]))

  (provide define/memo memo-lambda))