Add intmap-replace!.

* module/language/cps/intmap.scm (intmap-replace!): New interface.

* module/language/cps2/dce.scm (elide-type-checks):
* module/language/cps2/simplify.scm (transform-conts):
* module/language/cps2/utils.scm (intmap-map):
* module/language/tree-il/compile-cps2.scm (cps-convert/thunk): Use intmap-replace!.

module/language/cps/intmap.scm

@@ -43,6 +43,7 @@
             intmap-add
             intmap-add!
             intmap-replace
+            intmap-replace!
             intmap-remove
             intmap-ref
             intmap-next
@@ -285,20 +286,61 @@
     (($ <transient-intmap>)
      (intmap-add (persistent-intmap map) i val meet))))
 
+(define* (intmap-replace! map i val #:optional (meet (lambda (old new) new)))
+  "Like intmap-add!, but requires that @var{i} was present in the map
+already, and always calls the meet procedure."
+  (define (not-found)
+    (error "not found" i))
+  (define (ensure-branch! root idx)
+    (let ((edit (vector-ref root *edit-index*))
+          (v (vector-ref root idx)))
+      (when (absent? v) (not-found))
+      (let ((v* (writable-branch v edit)))
+        (unless (eq? v v*)
+          (vector-set! root idx v*))
+        v*)))
+  (define (adjoin! i shift root)
+    (let* ((shift (- shift *branch-bits*))
+           (idx (logand (ash i (- shift)) *branch-mask*)))
+      (if (zero? shift)
+          (let ((node (vector-ref root idx)))
+            (when (absent? node) (not-found))
+            (vector-set! root idx (meet node val)))
+          (adjoin! i shift (ensure-branch! root idx)))))
+  (match map
+    (($ <transient-intmap> min shift root edit)
+     (assert-readable! edit)
+     (cond
+      ((< i 0)
+       ;; The power-of-two spanning trick doesn't work across 0.
+       (error "Intmaps can only map non-negative integers." i))
+      ((and (present? root) (<= min i) (< i (+ min (ash 1 shift))))
+       (if (zero? shift)
+           (set-transient-intmap-root! map (meet root val))
+           (let ((root* (writable-branch root edit)))
+             (unless (eq? root root*)
+               (set-transient-intmap-root! map root*))
+             (adjoin! (- i min) shift root*))))
+      (else
+       (not-found)))
+     map)
+    (($ <intmap>)
+     (intmap-add! (transient-intmap map) i val meet))))
+
 (define* (intmap-replace map i val #:optional (meet (lambda (old new) new)))
   "Like intmap-add, but requires that @var{i} was present in the map
 already, and always calls the meet procedure."
-  (define (not-found i)
+  (define (not-found)
     (error "not found" i))
   (define (adjoin i shift root)
     (if (zero? shift)
         (if (absent? root)
-            (not-found i)
+            (not-found)
             (meet root val))
         (let* ((shift (- shift *branch-bits*))
                (idx (logand (ash i (- shift)) *branch-mask*)))
           (if (absent? root)
-              (not-found i)
+              (not-found)
               (let* ((node (vector-ref root idx))
                      (node* (adjoin i shift node)))
                 (if (eq? node node*)
@@ -316,8 +358,7 @@ already, and always calls the meet procedure."
          (if (eq? root old-root)
              map
              (make-intmap min shift root))))
-      (else
-       (not-found i))))
+      (else (not-found))))
     (($ <transient-intmap>)
      (intmap-replace (persistent-intmap map) i val meet))))
 

module/language/cps2/dce.scm

@@ -44,8 +44,7 @@ KFUN where we can prove that no assertion will be raised at run-time."
   (let ((types (infer-types conts kfun)))
     (define (visit-primcall effects fx label name args)
       (if (primcall-types-check? types label name args)
-          (intmap-add! effects label (logand fx (lognot &type-check))
-                       (lambda (old new) new))
+          (intmap-replace! effects label (logand fx (lognot &type-check)))
           effects))
     (persistent-intmap
      (intmap-fold (lambda (label types effects)

module/language/cps2/simplify.scm

@@ -58,7 +58,7 @@
                   (let ((v* (f k v)))
                     (if (equal? v v*)
                         out
-                        (intmap-add! out k v* (lambda (old new) new)))))
+                        (intmap-replace! out k v*))))
                 conts
                 conts)))
 

module/language/cps2/utils.scm

@@ -106,9 +106,7 @@
 
 (define (intmap-map proc map)
   (persistent-intmap
-   (intmap-fold (lambda (k v out)
-                  (intmap-add! out k (proc k v)
-                               (lambda (old new) new)))
+   (intmap-fold (lambda (k v out) (intmap-replace! out k (proc k v)))
                 map
                 map)))
 

module/language/tree-il/compile-cps2.scm

@@ -915,8 +915,7 @@ integer."
       ($ ((lambda (cps)
             (let ((init (build-cont
                           ($kfun (tree-il-src exp) '() init ktail kclause))))
-              (with-cps (persistent-intmap (intmap-add! cps kinit init
-                                                        (lambda (old new) new)))
+              (with-cps (persistent-intmap (intmap-replace! cps kinit init))
                 kinit))))))))
 
 (define *comp-module* (make-fluid))