From 0b0020153e39abfd3766cba2b0104425700459ce Mon Sep 17 00:00:00 2001
From: Tony Garnock-Jones <tonygarnockjones@gmail.com>
Date: Fri, 18 Jul 2014 11:26:06 -0700
Subject: [PATCH] Constant-time sorted-map-size

---
 minimart/memoize.rkt    | 52 +++++++++++++++++++++++++++++++++++++++++
 minimart/sorted-map.rkt | 21 ++++++++++-------
 2 files changed, 64 insertions(+), 9 deletions(-)
 create mode 100644 minimart/memoize.rkt

diff --git a/minimart/memoize.rkt b/minimart/memoize.rkt
new file mode 100644
index 0000000..94f1c40
--- /dev/null
+++ b/minimart/memoize.rkt
@@ -0,0 +1,52 @@
+#lang racket/base
+;; Poor-man's memoization.
+
+(provide memoize1)
+
+(define sentinel (cons #f #f))
+
+(define (memoize1 f)
+  (define results (make-weak-hash))
+  (lambda (arg)
+    (hash-ref results arg (lambda ()
+			    (define val (f arg))
+			    (hash-set! results arg val)
+			    val))))
+
+(module+ test
+  (require rackunit)
+
+  (define call-counter 0)
+
+  (define (raw x)
+    (set! call-counter (+ call-counter 1))
+    (gensym 'raw-result))
+
+  (define cooked (memoize1 raw))
+
+  ;; These tests will *likely* pass, but if garbage collection strikes
+  ;; at an inopportune moment, they may fail.
+
+  (collect-garbage)
+
+  (define v (cons 1 2))
+
+  (check-equal? call-counter 0)
+  (check-eq? (cooked v) (cooked v))
+  (check-equal? call-counter 1)
+
+  (set! v (cons 1 2))
+
+  (check-equal? call-counter 1)
+  (check-equal? (cooked v) (cooked v))
+  (check-equal? call-counter 1)
+
+  (set! v (cons 1 2))
+
+  (collect-garbage)
+  (collect-garbage)
+  (collect-garbage)
+
+  (check-equal? call-counter 1)
+  (check-equal? (cooked v) (cooked v))
+  (check-equal? call-counter 2))
diff --git a/minimart/sorted-map.rkt b/minimart/sorted-map.rkt
index 83e66c4..f2c952c 100644
--- a/minimart/sorted-map.rkt
+++ b/minimart/sorted-map.rkt
@@ -4,7 +4,7 @@
 ;;
 ;; Modified by Tony Garnock-Jones, July 2014:
 ;;  - trees are hashconsed
-;;  - sorted-map-size is made constant-time (TODO)
+;;  - sorted-map-size is made constant-time
 
 (provide (struct-out sorted-map)
 	 sorted-map-empty
@@ -20,7 +20,8 @@
 	 sorted-map-delete)
 
 (require "canonicalize.rkt")
-  
+(require "memoize.rkt")
+
 ; A purely functional sorted-map library.
 
 ; Provides logarithmic insert, update, get & delete.
@@ -323,14 +324,16 @@
 
 
 ; Returns the size of the sorted map:
-;; tonyg 20140718 TODO: make this O(1) for every smap
-(define (sorted-map-size smap)
-  (match smap
-    [(T! l r)   (+ 1 (sorted-map-size l) 
-                     (sorted-map-size r))]
-    [(L!)           0]))
+;; tonyg 20140718 this is memoized to run in O(1) for every smap
+(define sorted-map-size
+  (memoize1
+   (lambda (smap)
+     (match smap
+       [(T! l r)   (+ 1 (sorted-map-size l)
+		        (sorted-map-size r))]
+       [(L!)       0]))))
+
 
-  
 ; Returns the maxium (key . value) pair:
 (define/match (sorted-map-max node)
   [(T! _ k v (L!))   (cons k v)]