Implement Iterator

Author: tylertreat-wf

trie/ctrie/ctrie.go

@@ -339,7 +339,36 @@ func (c *Ctrie) ReadOnlySnapshot() *Ctrie {
 	}
 }
 
-//func (c *Ctrie) Iterator()
+// Iterator returns a channel which yields the Entries of the Ctrie.
+func (c *Ctrie) Iterator() <-chan *Entry {
+	ch := make(chan *Entry)
+	snapshot := c.ReadOnlySnapshot()
+	go func() {
+		traverse(snapshot.root, ch)
+		close(ch)
+	}()
+	return ch
+}
+
+func traverse(i *iNode, ch chan<- *Entry) {
+	switch {
+	case i.main.cNode != nil:
+		for _, br := range i.main.cNode.array {
+			switch b := br.(type) {
+			case *iNode:
+				traverse(b, ch)
+			case *sNode:
+				ch <- b.Entry
+			}
+		}
+	case i.main.lNode != nil:
+		for _, e := range i.main.lNode.Map(func(sn interface{}) interface{} {
+			return sn.(*sNode).Entry
+		}) {
+			ch <- e.(*Entry)
+		}
+	}
+}
 
 func (c *Ctrie) assertReadWrite() {
 	if c.readOnly {

trie/ctrie/ctrie_test.go

@@ -253,6 +253,36 @@ func TestSnapshot(t *testing.T) {
 	assert.Equal(0, val)
 }
 
+func TestIterator(t *testing.T) {
+	assert := assert.New(t)
+	ctrie := New(nil)
+	for i := 0; i < 10; i++ {
+		ctrie.Insert([]byte(strconv.Itoa(i)), i)
+	}
+	expected := map[string]int{
+		"0": 0,
+		"1": 1,
+		"2": 2,
+		"3": 3,
+		"4": 4,
+		"5": 5,
+		"6": 6,
+		"7": 7,
+		"8": 8,
+		"9": 9,
+	}
+
+	count := 0
+	for entry := range ctrie.Iterator() {
+		exp, ok := expected[string(entry.Key)]
+		if assert.True(ok) {
+			assert.Equal(exp, entry.Value)
+		}
+		count++
+	}
+	assert.Equal(len(expected), count)
+}
+
 func BenchmarkInsert(b *testing.B) {
 	ctrie := New(nil)
 	b.ResetTimer()