From c762502f6e764cd69fde7813d4e812b2aa7b2f70 Mon Sep 17 00:00:00 2001
From: Christian Schwarz <me@cschwarz.com>
Date: Sun, 7 May 2017 14:07:50 +0200
Subject: [PATCH] zfs: implement DatasetPath top down walk

---
 zfs/datasetpath_visitor.go      | 122 ++++++++++++++++++++++++++++++++
 zfs/datasetpath_visitor_test.go |  95 +++++++++++++++++++++++++
 2 files changed, 217 insertions(+)
 create mode 100644 zfs/datasetpath_visitor.go
 create mode 100644 zfs/datasetpath_visitor_test.go

diff --git a/zfs/datasetpath_visitor.go b/zfs/datasetpath_visitor.go
new file mode 100644
index 0000000..02abd88
--- /dev/null
+++ b/zfs/datasetpath_visitor.go
@@ -0,0 +1,122 @@
+package zfs
+
+type DatasetPathForest struct {
+	roots []*datasetPathTree
+}
+
+func NewDatasetPathForest() *DatasetPathForest {
+	return &DatasetPathForest{
+		make([]*datasetPathTree, 0),
+	}
+}
+
+func (f *DatasetPathForest) Add(p DatasetPath) {
+	if len(p) <= 0 {
+		panic("dataset path too short. must have length > 0")
+	}
+
+	// Find its root
+	var root *datasetPathTree
+	for _, r := range f.roots {
+		if r.Add(p) {
+			root = r
+			break
+		}
+	}
+	if root == nil {
+		root = newDatasetPathTree(p)
+		f.roots = append(f.roots, root)
+	}
+}
+
+type DatasetPathVisit struct {
+	Path DatasetPath
+	// If true, the dataset referenced by Path was not in the list of datasets to traverse
+	FilledIn bool
+}
+
+type DatasetPathsVisitor func(v DatasetPathVisit) (visitChildTree bool)
+
+// Traverse a list of DatasetPaths top down, i.e. given a set of datasets with same
+// path prefix, those with shorter prefix are traversed first.
+// If there are gaps, i.e. the intermediary component a/b bewtween a and a/b/c,
+// those gaps are still visited but the FilledIn property of the visit is set to true.
+func (f *DatasetPathForest) WalkTopDown(visitor DatasetPathsVisitor) {
+
+	for _, r := range f.roots {
+		r.WalkTopDown([]string{}, visitor)
+	}
+
+}
+
+/* PRIVATE IMPLEMENTATION */
+
+type datasetPathTree struct {
+	Component string
+	FilledIn  bool
+	Children  []*datasetPathTree
+}
+
+func (t *datasetPathTree) Add(p DatasetPath) bool {
+
+	if len(p) == 0 {
+		return true
+	}
+
+	if p[0] == t.Component {
+
+		remainder := p[1:]
+		if len(remainder) == 0 {
+			t.FilledIn = false
+			return true
+		}
+
+		for _, c := range t.Children {
+			if c.Add(remainder) {
+				return true
+			}
+		}
+
+		t.Children = append(t.Children, newDatasetPathTree(remainder))
+		return true
+
+	} else {
+
+		return false
+
+	}
+
+}
+
+func (t *datasetPathTree) WalkTopDown(parent DatasetPath, visitor DatasetPathsVisitor) {
+
+	this := append(parent, t.Component)
+
+	visitChildTree := visitor(DatasetPathVisit{this, t.FilledIn})
+
+	if visitChildTree {
+		for _, c := range t.Children {
+			c.WalkTopDown(this, visitor)
+		}
+	}
+
+}
+
+func newDatasetPathTree(initial DatasetPath) (t *datasetPathTree) {
+	t = &datasetPathTree{}
+	var cur *datasetPathTree
+	cur = t
+	for i, comp := range initial {
+		cur.Component = comp
+		cur.FilledIn = true
+		cur.Children = make([]*datasetPathTree, 0, 1)
+		if i == len(initial)-1 {
+			cur.FilledIn = false // last component is not filled in
+			break
+		}
+		child := &datasetPathTree{}
+		cur.Children = append(cur.Children, child)
+		cur = child
+	}
+	return t
+}
diff --git a/zfs/datasetpath_visitor_test.go b/zfs/datasetpath_visitor_test.go
new file mode 100644
index 0000000..b1b3cad
--- /dev/null
+++ b/zfs/datasetpath_visitor_test.go
@@ -0,0 +1,95 @@
+package zfs
+
+import (
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestNewDatasetPathTree(t *testing.T) {
+
+	r := newDatasetPathTree(toDatasetPath("pool1/foo/bar"))
+
+	assert.Equal(t, "pool1", r.Component)
+	assert.True(t, len(r.Children) == 1)
+	cr := r.Children[0]
+	assert.Equal(t, "foo", cr.Component)
+	assert.True(t, len(cr.Children) == 1)
+	ccr := cr.Children[0]
+	assert.Equal(t, "bar", ccr.Component)
+
+}
+
+type visitRecorder struct {
+	visits []DatasetPathVisit
+}
+
+func makeVisitRecorder() (v DatasetPathsVisitor, rec *visitRecorder) {
+	rec = &visitRecorder{
+		visits: make([]DatasetPathVisit, 0),
+	}
+	v = func(v DatasetPathVisit) error {
+		rec.visits = append(rec.visits, v)
+		return nil
+	}
+	return
+}
+
+func buildForest(paths []DatasetPath) (f *DatasetPathForest) {
+	f = NewDatasetPathForest()
+	for _, p := range paths {
+		f.Add(p)
+	}
+	return
+}
+
+func TestDatasetPathForestWalkTopDown(t *testing.T) {
+
+	paths := []DatasetPath{
+		toDatasetPath("pool1"),
+		toDatasetPath("pool1/foo/bar"),
+		toDatasetPath("pool1/foo/bar/looloo"),
+		toDatasetPath("pool2/test/bar"),
+	}
+
+	v, rec := makeVisitRecorder()
+
+	buildForest(paths).WalkTopDown(v)
+
+	expectedVisists := []DatasetPathVisit{
+		{toDatasetPath("pool1"), false},
+		{toDatasetPath("pool1/foo"), true},
+		{toDatasetPath("pool1/foo/bar"), false},
+		{toDatasetPath("pool1/foo/bar/looloo"), false},
+		{toDatasetPath("pool2"), true},
+		{toDatasetPath("pool2/test"), true},
+		{toDatasetPath("pool2/test/bar"), false},
+	}
+	assert.Equal(t, expectedVisists, rec.visits)
+
+}
+
+func TestDatasetPathWalkTopDownWorksUnordered(t *testing.T) {
+
+	paths := []DatasetPath{
+		toDatasetPath("pool1"),
+		toDatasetPath("pool1/foo/bar/looloo"),
+		toDatasetPath("pool1/foo/bar"),
+		toDatasetPath("pool1/bang/baz"),
+	}
+
+	v, rec := makeVisitRecorder()
+
+	buildForest(paths).WalkTopDown(v)
+
+	expectedVisists := []DatasetPathVisit{
+		{toDatasetPath("pool1"), false},
+		{toDatasetPath("pool1/foo"), true},
+		{toDatasetPath("pool1/foo/bar"), false},
+		{toDatasetPath("pool1/foo/bar/looloo"), false},
+		{toDatasetPath("pool1/bang"), true},
+		{toDatasetPath("pool1/bang/baz"), false},
+	}
+
+	assert.Equal(t, expectedVisists, rec.visits)
+
+}