;;; Copyright 2024 Masaya Tojo <masaya@tojo.tokyo>
;;;
;;; Licensed under the Apache License, Version 2.0 (the "License");
;;; you may not use this file except in compliance with the License.
;;; You may obtain a copy of the License at
;;;
;;;     http://www.apache.org/licenses/LICENSE-2.0
;;;
;;; Unless required by applicable law or agreed to in writing, software
;;; distributed under the License is distributed on an "AS IS" BASIS,
;;; WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
;;; See the License for the specific language governing permissions and
;;; limitations under the License.

(define-library (qklib memoization)
  (export memoize memoize1)
  (import (scheme base)
          (scheme case-lambda)
          (srfi 128))
  (cond-expand
    ((library (srfi srfi-146 hash))     ; for guile-srfi-146
     (import (srfi srfi-146 hash)))
    ((library (srfi 146 hash))
     (import (srfi 146 hash))))
  (begin
    (define memoize
      (case-lambda
        ((f)
         (memoize f (make-default-comparator)))
        ((f comparator)
         (let ((memo (hashmap comparator)))
           (let memo-fix ((f f))
             (lambda (args)
               (hashmap-ref memo args
                            (lambda ()
                              (define result (apply (f (memo-fix f)) args))
                              (set! memo (hashmap-set! memo args result))
                              result))))))))

    (define memoize1
      (case-lambda
        ((f)
         (memoize1 f (make-default-comparator)))
        ((f comparator)
         (let ((memo (hashmap comparator)))
           (let memo-fix ((f f))
             (lambda (x)
               (hashmap-ref memo x
                            (lambda ()
                              (define result ((f (memo-fix f)) x))
                              (set! memo (hashmap-set! memo x result))
                              result))))))))
    ))