Commits

Moritz Heidkamp committed 1d19b6d

Add map-ref-in and map-update-in

Comments (0)

Files changed (2)

persistent-hash-map.scm

  map-assoc
  map-dissoc
  map-ref
+ map-ref-in
+ map-update-in
  map-size
  map-contains?
  map-assoc!
   (->map alist caar cdar cdr))
 
 (: map-ref (hash-map * #!optional * --> *))
+(define +not-found+
+  (list 'not-found))
+
 (define (map-ref map key #!optional not-found)
   (-lookup map key not-found))
 
+(: map-ref-in (hash-map list #!optional * -> *))
+(define (map-ref-in map keys #!optional not-found)
+  (if (null? keys)
+      map
+      (let ((val (map-ref map (car keys) +not-found+)))
+        (cond
+         ((eq? +not-found+ val)
+          not-found)
+         ((null? (cdr keys))
+          val)
+         (else
+          (map-ref-in val (cdr keys) not-found))))))
+
 (define (map-keyvals-proc proc)
   (rec (self map key val . more)
        (let ((res (proc map key val)))
 (define map-size
   -count)
 
-(define +not-found+
-  (list 'not-found))
-
 (: map-contains? (hash-map * -> boolean))
 (define (map-contains? map key)
   (not (eq? +not-found+ (map-ref map key +not-found+))))
                (map->transient-map m1)
                m2)))
 
+(: map-update-in (persistent-hash-map list (* #!rest * -> *) #!rest * -> persistent-hash-map))
+(define (map-update-in map keys proc . args)
+  (when (null? keys)
+    (error 'map-update-in "Need at least one key"))
+  (let* ((key  (car keys))
+         (smap (map-ref map key +empty-persistent-hash-map+)))
+    (if (null? (cdr keys))
+        (map-assoc map key (apply proc smap args))
+        (map-assoc map key (apply map-update-in smap (cdr keys) proc args)))))
+
 (define (print-map type map out)
   (display "#<" out)
   (display type out)
-(use test srfi-1 persistent-hash-map)
+(use test srfi-1 clojurian-syntax persistent-hash-map)
 
 (test-begin)
 
     (test 4 (map-ref m 3))
     (test 6 (map-ref m 5))))
 
+(test-group "map-ref-in, map-update-in"
+  (let ((m (-> (persistent-map 'foo (persistent-map 'bar 1))
+               (map-update-in '(foo bar) + 1)
+               (map-update-in '(foo baz) map-assoc 'qux 9))))
+    (test 2 (map-ref-in m '(foo bar)))
+    (test 9 (map-ref-in m '(foo baz qux)))))
+
 (test-end)
 
 (test-exit)