(module memoize mzscheme

  ;; TODO: make the implementation thread-safe

  (define undefined (letrec ([x x]) x))

  ;; extracts the argument list from a procedure definition via LAMBDA or DEFINE
  (define-for-syntax (lambda-or-define-args lambda-or-define args)
    (cond
      [(module-identifier=? lambda-or-define #'lambda) args]
      [(module-identifier=? lambda-or-define #'define)
       (with-syntax ([(name . rest) args])
         #'rest)]
      [else (raise-syntax-error 'memoize "internal error" lambda-or-define)]))

  ;; creates a local binding using LET or internal DEFINE depending on context
  (define-syntax local-declaration
    (syntax-rules (lambda define)
      [(_ lambda ([var binding] ...) body)
       (let ([var binding] ...) body)]
      [(_ define ([var binding] ...) body)
       (begin (define var binding) ... body)]))

  (define-syntax (memo-lambda-or-define/memo stx)
    (syntax-case stx ()
      [(_ lambda-or-define build-hash-table hash-code args body0 body1 ...)
       (syntax-case (lambda-or-define-args #'lambda-or-define #'args) ()
         [()
          #'(local-declaration lambda-or-define ([cached? #f] [cache undefined])
              (lambda-or-define args
                (unless cached?
                  (set! cache (let () body0 body1 ...))
                  (set! cached? #t))
                cache))]
         [(arg)
          #'(local-declaration lambda-or-define ([cache build-hash-table])
              (lambda-or-define args
                (hash-table-get cache arg (lambda ()
                                            (let ([ans (let () body0 body1 ...)])
                                              (hash-table-put! cache arg ans)
                                              ans)))))]
         [multiple-args
          (let loop ([arg-exp #'multiple-args] [fixed-args null])
            (syntax-case arg-exp ()
              [(arg . rest)
               (loop #'rest (cons #'arg fixed-args))]
              [rest
               (with-syntax ([list-exp (let loop ([fixed-args fixed-args]
                                                  [result (syntax-case #'rest ()
                                                            [() #''()] ;; PLT v4 doesn't allow () as an expression form
                                                            [_ #'rest])])
                                         (if (null? fixed-args)
                                             result
                                             (loop (cdr fixed-args)
                                                   #`(cons #,(car fixed-args) #,result))))])
                 ;; NOTE: once we've created the (equal?-based) hash code,
                 ;;       we can just store it in an eq?-based hash table.
                 #'(local-declaration lambda-or-define ([cache (make-hash-table)])
                     (lambda-or-define args
                       (let* ([arg-list list-exp]
                              [key (apply bitwise-xor (map hash-code arg-list))]
                              [alist (hash-table-get cache key (lambda () null))])
                         ;; BUG: this should be using a custom comparison function
                         ;;      that uses eq/equal on the individual args
                         (cond
                           [(assoc arg-list alist) => cdr]
                           [else (let ([ans (let () body0 body1 ...)])
                                   (hash-table-put! cache key (cons (cons arg-list ans) alist))
                                   ans)])))))]))])]))

  (define-syntax memo-lambda
    (syntax-rules ()
      [(_ args body0 body1 ...)
       (memo-lambda-or-define/memo lambda (make-hash-table) eq-hash-code args body0 body1 ...)]))

  (define-syntax memo-lambda*
    (syntax-rules ()
      [(_ args body0 body1 ...)
       (memo-lambda-or-define/memo lambda (make-hash-table 'equal) equal-hash-code args body0 body1 ...)]))

  (define-syntax define/memo
    (syntax-rules ()
      [(_ args body0 body1 ...)
       (memo-lambda-or-define/memo define (make-hash-table) eq-hash-code args body0 body1 ...)]))

  (define-syntax define/memo*
    (syntax-rules ()
      [(_ args body0 body1 ...)
       (memo-lambda-or-define/memo define (make-hash-table 'equal) equal-hash-code args body0 body1 ...)]))

  (provide define/memo memo-lambda define/memo* memo-lambda*))