guix: utils: Add fold-tree and fold-tree-leaves.

* guix/utils.scm (fold-tree, fold-tree-leaves): New functions.
* tests/utils.scm: Add tests for them.

2 files changed, 67 insertions, 1 deletions

diff --git a/guix/utils.scm b/guix/utils.scm
index 700a191d71..b61ff2477d 100644
--- a/guix/utils.scm
+++ b/guix/utils.scm
@@ -1,6 +1,7 @@
 ;;; GNU Guix --- Functional package management for GNU
 ;;; Copyright © 2012, 2013, 2014 Ludovic Courtès <ludo@gnu.org>
 ;;; Copyright © 2013 Mark H Weaver <mhw@netris.org>
+;;; Copyright © 2014 Eric Bavier <bavier@member.fsf.org>
 ;;;
 ;;; This file is part of GNU Guix.
 ;;;
@@ -72,6 +73,8 @@
             call-with-temporary-output-file
             with-atomic-file-output
             fold2
+            fold-tree
+            fold-tree-leaves
 
             filtered-port
             compressed-port
@@ -649,6 +652,36 @@ output port, and PROC's result is returned."
              (lambda (result1 result2)
                (fold2 proc result1 result2 (cdr lst1) (cdr lst2)))))))))
 
+(define (fold-tree proc init children roots)
+  "Call (PROC NODE RESULT) for each node in the tree that is reachable from
+ROOTS, using INIT as the initial value of RESULT.  The order in which nodes
+are traversed is not specified, however, each node is visited only once, based
+on an eq? check.  Children of a node to be visited are generated by
+calling (CHILDREN NODE), the result of which should be a list of nodes that
+are connected to NODE in the tree, or '() or #f if NODE is a leaf node."
+  (let loop ((result init)
+             (seen vlist-null)
+             (lst roots))
+    (match lst
+      (() result)
+      ((head . tail)
+       (if (not (vhash-assq head seen))
+           (loop (proc head result)
+                 (vhash-consq head #t seen)
+                 (match (children head)
+                   ((or () #f) tail)
+                   (children (append tail children))))
+           (loop result seen tail))))))
+
+(define (fold-tree-leaves proc init children roots)
+  "Like fold-tree, but call (PROC NODE RESULT) only for leaf nodes."
+  (fold-tree
+   (lambda (node result)
+     (match (children node)
+       ((or () #f) (proc node result))
+       (else result)))
+   init children roots))
+
 
 ;;;
 ;;; Source location.
diff --git a/tests/utils.scm b/tests/utils.scm
index 8ad399f75c..611867ca09 100644
--- a/tests/utils.scm
+++ b/tests/utils.scm
@@ -1,5 +1,6 @@
 ;;; GNU Guix --- Functional package management for GNU
 ;;; Copyright © 2012, 2013, 2014 Ludovic Courtès <ludo@gnu.org>
+;;; Copyright © 2014 Eric Bavier <bavier@member.fsf.org>
 ;;;
 ;;; This file is part of GNU Guix.
 ;;;
@@ -25,7 +26,8 @@
   #:use-module (srfi srfi-64)
   #:use-module (rnrs bytevectors)
   #:use-module (rnrs io ports)
-  #:use-module (ice-9 match))
+  #:use-module (ice-9 match)
+  #:use-module (ice-9 vlist))
 
 (define temp-file
   (string-append "t-utils-" (number->string (getpid))))
@@ -118,6 +120,37 @@
                '(0 1 2 3)))
     list))
 
+(let* ((tree (alist->vhash
+              '((0 2 3) (1 3 4) (2) (3 5 6) (4 6) (5) (6))
+              hashq))
+       (add-one (lambda (_ r) (1+ r)))
+       (tree-lookup (lambda (n) (cdr (vhash-assq n tree)))))
+  (test-equal "fold-tree, single root"
+    5 (fold-tree add-one 0 tree-lookup '(0)))
+  (test-equal "fold-tree, two roots"
+    7 (fold-tree add-one 0 tree-lookup '(0 1)))
+  (test-equal "fold-tree, sum"
+    16 (fold-tree + 0 tree-lookup '(0)))
+  (test-equal "fold-tree, internal"
+    18 (fold-tree + 0 tree-lookup '(3 4)))
+  (test-equal "fold-tree, cons"
+    '(1 3 4 5 6)
+    (sort (fold-tree cons '() tree-lookup '(1)) <))
+  (test-equal "fold-tree, overlapping paths"
+    '(1 3 4 5 6)
+    (sort (fold-tree cons '() tree-lookup '(1 4)) <))
+  (test-equal "fold-tree, cons, two roots"
+    '(0 2 3 4 5 6)
+    (sort (fold-tree cons '() tree-lookup '(0 4)) <))
+  (test-equal "fold-tree-leaves, single root"
+    2 (fold-tree-leaves add-one 0 tree-lookup '(1)))
+  (test-equal "fold-tree-leaves, single root, sum"
+    11 (fold-tree-leaves + 0 tree-lookup '(1)))
+  (test-equal "fold-tree-leaves, two roots"
+    3 (fold-tree-leaves add-one 0 tree-lookup '(0 1)))
+  (test-equal "fold-tree-leaves, two roots, sum"
+    13 (fold-tree-leaves + 0 tree-lookup '(0 1))))
+
 (test-assert "filtered-port, file"
   (let* ((file  (search-path %load-path "guix.scm"))
          (input (open-file file "r0b")))