From 283aca83c2f1af99caa5b861d7a09e04353750e1 Mon Sep 17 00:00:00 2001
From: Barak Michener <me@barakmich.com>
Date: Fri, 30 Oct 2015 17:16:40 -0400
Subject: [PATCH] Expose LabelContext in gremlin

---
 graph/path/morphism_apply_functions.go | 9 ++++++---
 graph/path/path.go                     | 9 ++++++++-
 query/gremlin/build_iterator.go        | 6 ++++++
 query/gremlin/gremlin_test.go          | 7 +++++++
 query/gremlin/traversals.go            | 1 +
 5 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/graph/path/morphism_apply_functions.go b/graph/path/morphism_apply_functions.go
index 185aec2..0a9db5e 100644
--- a/graph/path/morphism_apply_functions.go
+++ b/graph/path/morphism_apply_functions.go
@@ -16,6 +16,7 @@ package path
 
 import (
 	"fmt"
+	"reflect"
 
 	"github.com/google/cayley/graph"
 	"github.com/google/cayley/graph/iterator"
@@ -161,25 +162,27 @@ func bothMorphism(tags []string, via ...interface{}) morphism {
 	}
 }
 
-func labelContextMorphism(via ...interface{}) morphism {
+func labelContextMorphism(tags []string, via ...interface{}) morphism {
 	var path *Path
 	if len(via) == 0 {
 		path = nil
 	} else {
 		path = buildViaPath(nil, via...)
+		path = path.Tag(tags...)
 	}
 	return morphism{
 		Name: "label_context",
 		Reversal: func(ctx *context) (morphism, *context) {
 			out := ctx.copy()
 			ctx.labelSet = path
-			return labelContextMorphism(via...), &out
+			return labelContextMorphism(tags, via...), &out
 		},
 		Apply: func(qs graph.QuadStore, in graph.Iterator, ctx *context) (graph.Iterator, *context) {
 			out := ctx.copy()
 			out.labelSet = path
 			return in, &out
 		},
+		tags: tags,
 	}
 }
 
@@ -364,7 +367,7 @@ func buildViaPath(qs graph.QuadStore, via ...interface{}) *Path {
 		case string:
 			return StartPath(qs, p)
 		default:
-			panic(fmt.Sprint("Invalid type passed to buildViaPath. ", p))
+			panic(fmt.Sprintln("Invalid type passed to buildViaPath.", reflect.TypeOf(v), p))
 		}
 	}
 	var strings []string
diff --git a/graph/path/path.go b/graph/path/path.go
index 25ba54a..03e0df0 100644
--- a/graph/path/path.go
+++ b/graph/path/path.go
@@ -274,7 +274,14 @@ func (p *Path) Has(via interface{}, nodes ...string) *Path {
 // LabelContext restricts the following operations (such as In, Out) to only
 // traverse edges that match the given set of labels.
 func (p *Path) LabelContext(via ...interface{}) *Path {
-	p.stack = append(p.stack, labelContextMorphism(via...))
+	p.stack = append(p.stack, labelContextMorphism(nil, via...))
+	return p
+}
+
+// LabelContextWithTags is exactly like LabelContext, except it tags the value
+// of the label used in the traversal with the tags provided.
+func (p *Path) LabelContextWithTags(tags []string, via ...interface{}) *Path {
+	p.stack = append(p.stack, labelContextMorphism(tags, via...))
 	return p
 }
 
diff --git a/query/gremlin/build_iterator.go b/query/gremlin/build_iterator.go
index 511a83e..221a0f7 100644
--- a/query/gremlin/build_iterator.go
+++ b/query/gremlin/build_iterator.go
@@ -163,6 +163,12 @@ func buildPathFromObject(obj *otto.Object) *path.Path {
 		return p.InPredicates()
 	case "OutPredicates":
 		return p.OutPredicates()
+	case "LabelContext":
+		labels, tags, ok := getViaData(obj)
+		if !ok {
+			return nil
+		}
+		return p.LabelContextWithTags(tags, labels...)
 	default:
 		panic(fmt.Sprint("Unimplemented Gremlin function", gremlinType))
 	}
diff --git a/query/gremlin/gremlin_test.go b/query/gremlin/gremlin_test.go
index f7468cf..5194669 100644
--- a/query/gremlin/gremlin_test.go
+++ b/query/gremlin/gremlin_test.go
@@ -278,6 +278,13 @@ var testQueries = []struct {
 		`,
 		expect: []string{"are", "follows", "status"},
 	},
+	{
+		message: "traverse using LabelContext",
+		query: `
+			g.V("greg").LabelContext("smart_graph").Out("status").All()
+		`,
+		expect: []string{"smart_person"},
+	},
 }
 
 func runQueryGetTag(g []quad.Quad, query string, tag string) []string {
diff --git a/query/gremlin/traversals.go b/query/gremlin/traversals.go
index b5304e4..36d67cd 100644
--- a/query/gremlin/traversals.go
+++ b/query/gremlin/traversals.go
@@ -42,6 +42,7 @@ var traversals = []string{
 	"Difference",
 	"InPredicates",
 	"OutPredicates",
+	"LabelContext",
 }
 
 func (wk *worker) embedTraversals(env *otto.Otto, obj *otto.Object) {