1// Copyright ©2014 The Gonum Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5package simple
6
7import (
8	"fmt"
9
10	"gonum.org/v1/gonum/graph"
11	"gonum.org/v1/gonum/graph/internal/uid"
12	"gonum.org/v1/gonum/graph/iterator"
13)
14
15var (
16	dg *DirectedGraph
17
18	_ graph.Graph       = dg
19	_ graph.Directed    = dg
20	_ graph.NodeAdder   = dg
21	_ graph.NodeRemover = dg
22	_ graph.EdgeAdder   = dg
23	_ graph.EdgeRemover = dg
24)
25
26// DirectedGraph implements a generalized directed graph.
27type DirectedGraph struct {
28	nodes map[int64]graph.Node
29	from  map[int64]map[int64]graph.Edge
30	to    map[int64]map[int64]graph.Edge
31
32	nodeIDs uid.Set
33}
34
35// NewDirectedGraph returns a DirectedGraph.
36func NewDirectedGraph() *DirectedGraph {
37	return &DirectedGraph{
38		nodes: make(map[int64]graph.Node),
39		from:  make(map[int64]map[int64]graph.Edge),
40		to:    make(map[int64]map[int64]graph.Edge),
41
42		nodeIDs: uid.NewSet(),
43	}
44}
45
46// AddNode adds n to the graph. It panics if the added node ID matches an existing node ID.
47func (g *DirectedGraph) AddNode(n graph.Node) {
48	if _, exists := g.nodes[n.ID()]; exists {
49		panic(fmt.Sprintf("simple: node ID collision: %d", n.ID()))
50	}
51	g.nodes[n.ID()] = n
52	g.from[n.ID()] = make(map[int64]graph.Edge)
53	g.to[n.ID()] = make(map[int64]graph.Edge)
54	g.nodeIDs.Use(n.ID())
55}
56
57// Edge returns the edge from u to v if such an edge exists and nil otherwise.
58// The node v must be directly reachable from u as defined by the From method.
59func (g *DirectedGraph) Edge(uid, vid int64) graph.Edge {
60	edge, ok := g.from[uid][vid]
61	if !ok {
62		return nil
63	}
64	return edge
65}
66
67// Edges returns all the edges in the graph.
68func (g *DirectedGraph) Edges() graph.Edges {
69	var edges []graph.Edge
70	for _, u := range g.nodes {
71		for _, e := range g.from[u.ID()] {
72			edges = append(edges, e)
73		}
74	}
75	if len(edges) == 0 {
76		return graph.Empty
77	}
78	return iterator.NewOrderedEdges(edges)
79}
80
81// From returns all nodes in g that can be reached directly from n.
82func (g *DirectedGraph) From(id int64) graph.Nodes {
83	if _, ok := g.from[id]; !ok {
84		return graph.Empty
85	}
86
87	from := make([]graph.Node, len(g.from[id]))
88	i := 0
89	for vid := range g.from[id] {
90		from[i] = g.nodes[vid]
91		i++
92	}
93	if len(from) == 0 {
94		return graph.Empty
95	}
96	return iterator.NewOrderedNodes(from)
97}
98
99// HasEdgeBetween returns whether an edge exists between nodes x and y without
100// considering direction.
101func (g *DirectedGraph) HasEdgeBetween(xid, yid int64) bool {
102	if _, ok := g.from[xid][yid]; ok {
103		return true
104	}
105	_, ok := g.from[yid][xid]
106	return ok
107}
108
109// HasEdgeFromTo returns whether an edge exists in the graph from u to v.
110func (g *DirectedGraph) HasEdgeFromTo(uid, vid int64) bool {
111	if _, ok := g.from[uid][vid]; !ok {
112		return false
113	}
114	return true
115}
116
117// NewEdge returns a new Edge from the source to the destination node.
118func (g *DirectedGraph) NewEdge(from, to graph.Node) graph.Edge {
119	return &Edge{F: from, T: to}
120}
121
122// NewNode returns a new unique Node to be added to g. The Node's ID does
123// not become valid in g until the Node is added to g.
124func (g *DirectedGraph) NewNode() graph.Node {
125	if len(g.nodes) == 0 {
126		return Node(0)
127	}
128	if int64(len(g.nodes)) == uid.Max {
129		panic("simple: cannot allocate node: no slot")
130	}
131	return Node(g.nodeIDs.NewID())
132}
133
134// Node returns the node with the given ID if it exists in the graph,
135// and nil otherwise.
136func (g *DirectedGraph) Node(id int64) graph.Node {
137	return g.nodes[id]
138}
139
140// Nodes returns all the nodes in the graph.
141func (g *DirectedGraph) Nodes() graph.Nodes {
142	if len(g.nodes) == 0 {
143		return graph.Empty
144	}
145	nodes := make([]graph.Node, len(g.nodes))
146	i := 0
147	for _, n := range g.nodes {
148		nodes[i] = n
149		i++
150	}
151	return iterator.NewOrderedNodes(nodes)
152}
153
154// RemoveEdge removes the edge with the given end point IDs from the graph, leaving the terminal
155// nodes. If the edge does not exist it is a no-op.
156func (g *DirectedGraph) RemoveEdge(fid, tid int64) {
157	if _, ok := g.nodes[fid]; !ok {
158		return
159	}
160	if _, ok := g.nodes[tid]; !ok {
161		return
162	}
163
164	delete(g.from[fid], tid)
165	delete(g.to[tid], fid)
166}
167
168// RemoveNode removes the node with the given ID from the graph, as well as any edges attached
169// to it. If the node is not in the graph it is a no-op.
170func (g *DirectedGraph) RemoveNode(id int64) {
171	if _, ok := g.nodes[id]; !ok {
172		return
173	}
174	delete(g.nodes, id)
175
176	for from := range g.from[id] {
177		delete(g.to[from], id)
178	}
179	delete(g.from, id)
180
181	for to := range g.to[id] {
182		delete(g.from[to], id)
183	}
184	delete(g.to, id)
185
186	g.nodeIDs.Release(id)
187}
188
189// SetEdge adds e, an edge from one node to another. If the nodes do not exist, they are added
190// and are set to the nodes of the edge otherwise.
191// It will panic if the IDs of the e.From and e.To are equal.
192func (g *DirectedGraph) SetEdge(e graph.Edge) {
193	var (
194		from = e.From()
195		fid  = from.ID()
196		to   = e.To()
197		tid  = to.ID()
198	)
199
200	if fid == tid {
201		panic("simple: adding self edge")
202	}
203
204	if _, ok := g.nodes[fid]; !ok {
205		g.AddNode(from)
206	} else {
207		g.nodes[fid] = from
208	}
209	if _, ok := g.nodes[tid]; !ok {
210		g.AddNode(to)
211	} else {
212		g.nodes[tid] = to
213	}
214
215	g.from[fid][tid] = e
216	g.to[tid][fid] = e
217}
218
219// To returns all nodes in g that can reach directly to n.
220func (g *DirectedGraph) To(id int64) graph.Nodes {
221	if _, ok := g.from[id]; !ok {
222		return graph.Empty
223	}
224
225	to := make([]graph.Node, len(g.to[id]))
226	i := 0
227	for uid := range g.to[id] {
228		to[i] = g.nodes[uid]
229		i++
230	}
231	if len(to) == 0 {
232		return graph.Empty
233	}
234	return iterator.NewOrderedNodes(to)
235}
236