From 1d2e175c37527fa68518617f1e0aaa73c7a54bf6 Mon Sep 17 00:00:00 2001
From: Roman Strashkin <roman.strashkin@gmail.com>
Date: Fri, 19 Jun 2015 18:01:39 +0300
Subject: [PATCH] added ability to iterate over all indexes and use
 index.Iterate() instead of ReadDir() to walk over the graph

Signed-off-by: Roman Strashkin <roman.strashkin@gmail.com>
---
 truncindex/truncindex.go      | 10 ++++++++++
 truncindex/truncindex_test.go | 23 +++++++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/truncindex/truncindex.go b/truncindex/truncindex.go
index 9aae5c0..8d8bee0 100644
--- a/truncindex/truncindex.go
+++ b/truncindex/truncindex.go
@@ -108,3 +108,13 @@ func (idx *TruncIndex) Get(s string) (string, error) {
 	}
 	return "", fmt.Errorf("no such id: %s", s)
 }
+
+// Iterates over all stored IDs, and passes each of them to the given handler
+func (idx *TruncIndex) Iterate(handler func(id string)) {
+	idx.RLock()
+	defer idx.RUnlock()
+	idx.trie.Visit(func(prefix patricia.Prefix, item patricia.Item) error {
+		handler(string(prefix))
+		return nil
+	})
+}
diff --git a/truncindex/truncindex_test.go b/truncindex/truncindex_test.go
index 7253f6a..cc7bc01 100644
--- a/truncindex/truncindex_test.go
+++ b/truncindex/truncindex_test.go
@@ -96,6 +96,29 @@ func TestTruncIndex(t *testing.T) {
 	assertIndexGet(t, index, id[:7], id, false)
 	assertIndexGet(t, index, id[:15], id, false)
 	assertIndexGet(t, index, id, id, false)
+
+	assertIndexIterate(t)
+}
+
+func assertIndexIterate(t *testing.T) {
+	ids := []string{
+		"19b36c2c326ccc11e726eee6ee78a0baf166ef96",
+		"28b36c2c326ccc11e726eee6ee78a0baf166ef96",
+		"37b36c2c326ccc11e726eee6ee78a0baf166ef96",
+		"46b36c2c326ccc11e726eee6ee78a0baf166ef96",
+	}
+
+	index := NewTruncIndex(ids)
+
+	index.Iterate(func(targetId string) {
+		for _, id := range ids {
+			if targetId == id {
+				return
+			}
+		}
+
+		t.Fatalf("An unknown ID '%s'", targetId)
+	})
 }
 
 func assertIndexGet(t *testing.T, index *TruncIndex, input, expectedResult string, expectError bool) {