From 5d5045fdd6fa986fdd7bf41bae2076537b5e91ed Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Thu, 14 Jan 2016 13:52:01 -0800
Subject: [PATCH] dag: add HasVertex, HasEdge

---
 dag/graph.go      | 10 ++++++++++
 dag/graph_test.go | 26 ++++++++++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/dag/graph.go b/dag/graph.go
index 2572096eda..5178648d26 100644
--- a/dag/graph.go
+++ b/dag/graph.go
@@ -48,6 +48,16 @@ func (g *Graph) Edges() []Edge {
 	return result
 }
 
+// HasVertex checks if the given Vertex is present in the graph.
+func (g *Graph) HasVertex(v Vertex) bool {
+	return g.vertices.Include(v)
+}
+
+// HasEdge checks if the given Edge is present in the graph.
+func (g *Graph) HasEdge(e Edge) bool {
+	return g.edges.Include(e)
+}
+
 // Add adds a vertex to the graph. This is safe to call multiple time with
 // the same Vertex.
 func (g *Graph) Add(v Vertex) Vertex {
diff --git a/dag/graph_test.go b/dag/graph_test.go
index eb3e40b3af..7acd4a8310 100644
--- a/dag/graph_test.go
+++ b/dag/graph_test.go
@@ -98,6 +98,32 @@ func TestGraph_hashcode(t *testing.T) {
 	}
 }
 
+func TestGraphHasVertex(t *testing.T) {
+	var g Graph
+	g.Add(1)
+
+	if !g.HasVertex(1) {
+		t.Fatal("should have 1")
+	}
+	if g.HasVertex(2) {
+		t.Fatal("should not have 2")
+	}
+}
+
+func TestGraphHasEdge(t *testing.T) {
+	var g Graph
+	g.Add(1)
+	g.Add(2)
+	g.Connect(BasicEdge(1, 2))
+
+	if !g.HasEdge(BasicEdge(1, 2)) {
+		t.Fatal("should have 1,2")
+	}
+	if g.HasVertex(BasicEdge(2, 3)) {
+		t.Fatal("should not have 2,3")
+	}
+}
+
 type hashVertex struct {
 	code interface{}
 }