From 633828c8818959b28457c28e2887eb81ce1b9600 Mon Sep 17 00:00:00 2001
From: Denys Smirnov <denis.smirnov.91@gmail.com>
Date: Sat, 16 Dec 2017 18:37:15 +0200
Subject: [PATCH] shape: implement helper to walk the shape tree; resolves #651

---
 graph/shape/shape.go      | 65 ++++++++++++++++++++++++++++++++++++++-
 graph/shape/shape_test.go | 33 ++++++++++++++++++++
 2 files changed, 97 insertions(+), 1 deletion(-)

diff --git a/graph/shape/shape.go b/graph/shape/shape.go
index 557c46614..926831a6c 100644
--- a/graph/shape/shape.go
+++ b/graph/shape/shape.go
@@ -1,11 +1,13 @@
 package shape
 
 import (
+	"reflect"
+	"regexp"
+
 	"github.com/cayleygraph/cayley/clog"
 	"github.com/cayleygraph/cayley/graph"
 	"github.com/cayleygraph/cayley/graph/iterator"
 	"github.com/cayleygraph/cayley/quad"
-	"regexp"
 )
 
 // Shape represent a query tree shape.
@@ -25,6 +27,10 @@ type Optimizer interface {
 	OptimizeShape(s Shape) (Shape, bool)
 }
 
+// WalkFunc is used to visit all shapes in the tree.
+// If false is returned, branch will not be traversed further.
+type WalkFunc func(Shape) bool
+
 type resolveValues struct {
 	qs graph.QuadStore
 }
@@ -67,6 +73,63 @@ func Optimize(s Shape, qs graph.QuadStore) (Shape, bool) {
 	return s, opt
 }
 
+var rtShape = reflect.TypeOf((*Shape)(nil)).Elem()
+
+// Walk calls provided function for each shape in the tree.
+func Walk(s Shape, fnc WalkFunc) {
+	if s == nil {
+		return
+	}
+	if !fnc(s) {
+		return
+	}
+	walkReflect(reflect.ValueOf(s), fnc)
+}
+
+func walkReflect(rv reflect.Value, fnc WalkFunc) {
+	rt := rv.Type()
+	switch rv.Kind() {
+	case reflect.Slice:
+		if rt.Elem().ConvertibleTo(rtShape) {
+			// all element are shapes - call function on each of them
+			for i := 0; i < rv.Len(); i++ {
+				Walk(rv.Index(i).Interface().(Shape), fnc)
+			}
+		} else {
+			// elements are not shapes, but might contain them
+			for i := 0; i < rv.Len(); i++ {
+				walkReflect(rv.Index(i), fnc)
+			}
+		}
+	case reflect.Map:
+		keys := rv.MapKeys()
+		if rt.Elem().ConvertibleTo(rtShape) {
+			// all element are shapes - call function on each of them
+			for _, k := range keys {
+				Walk(rv.MapIndex(k).Interface().(Shape), fnc)
+			}
+		} else {
+			// elements are not shapes, but might contain them
+			for _, k := range keys {
+				walkReflect(rv.MapIndex(k), fnc)
+			}
+		}
+	case reflect.Struct:
+		// visit all fields
+		for i := 0; i < rt.NumField(); i++ {
+			f := rt.Field(i)
+			// if field is of shape type - call function on it
+			// we skip anonymous fields because they were already visited as part of the parent
+			if !f.Anonymous && f.Type.ConvertibleTo(rtShape) {
+				Walk(rv.Field(i).Interface().(Shape), fnc)
+				continue
+			}
+			// it might be a struct/map/slice field, so we need to go deeper
+			walkReflect(rv.Field(i), fnc)
+		}
+	}
+}
+
 // InternalQuad is an internal representation of quad index in QuadStore.
 type InternalQuad struct {
 	Subject   graph.Value
diff --git a/graph/shape/shape_test.go b/graph/shape/shape_test.go
index 92f640cb2..293bc04be 100644
--- a/graph/shape/shape_test.go
+++ b/graph/shape/shape_test.go
@@ -15,6 +15,7 @@
 package shape_test
 
 import (
+	"reflect"
 	"testing"
 
 	"github.com/cayleygraph/cayley/graph"
@@ -22,6 +23,7 @@ import (
 	. "github.com/cayleygraph/cayley/graph/shape"
 	"github.com/cayleygraph/cayley/quad"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func intVal(v int) graph.Value {
@@ -325,3 +327,34 @@ func TestOptimize(t *testing.T) {
 		})
 	}
 }
+
+func TestWalk(t *testing.T) {
+	var s Shape = NodesFrom{
+		Dir: quad.Subject,
+		Quads: Quads{
+			{Dir: quad.Subject, Values: Fixed{intVal(1)}},
+			{Dir: quad.Predicate, Values: Fixed{intVal(2)}},
+			{
+				Dir: quad.Object,
+				Values: QuadsAction{
+					Result: quad.Subject,
+					Filter: map[quad.Direction]graph.Value{
+						quad.Predicate: intVal(2),
+					},
+				},
+			},
+		},
+	}
+	var types []string
+	Walk(s, func(s Shape) bool {
+		types = append(types, reflect.TypeOf(s).String())
+		return true
+	})
+	require.Equal(t, []string{
+		"shape.NodesFrom",
+		"shape.Quads",
+		"shape.Fixed",
+		"shape.Fixed",
+		"shape.QuadsAction",
+	}, types)
+}