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/iterator"
12	"gonum.org/v1/gonum/graph/set/uid"
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.nodeIDs.Use(n.ID())
53}
54
55// Edge returns the edge from u to v if such an edge exists and nil otherwise.
56// The node v must be directly reachable from u as defined by the From method.
57func (g *DirectedGraph) Edge(uid, vid int64) graph.Edge {
58	edge, ok := g.from[uid][vid]
59	if !ok {
60		return nil
61	}
62	return edge
63}
64
65// Edges returns all the edges in the graph.
66func (g *DirectedGraph) Edges() graph.Edges {
67	var edges []graph.Edge
68	for _, u := range g.nodes {
69		for _, e := range g.from[u.ID()] {
70			edges = append(edges, e)
71		}
72	}
73	if len(edges) == 0 {
74		return graph.Empty
75	}
76	return iterator.NewOrderedEdges(edges)
77}
78
79// From returns all nodes in g that can be reached directly from n.
80//
81// The returned graph.Nodes is only valid until the next mutation of
82// the receiver.
83func (g *DirectedGraph) From(id int64) graph.Nodes {
84	if len(g.from[id]) == 0 {
85		return graph.Empty
86	}
87	return iterator.NewNodesByEdge(g.nodes, g.from[id])
88}
89
90// HasEdgeBetween returns whether an edge exists between nodes x and y without
91// considering direction.
92func (g *DirectedGraph) HasEdgeBetween(xid, yid int64) bool {
93	if _, ok := g.from[xid][yid]; ok {
94		return true
95	}
96	_, ok := g.from[yid][xid]
97	return ok
98}
99
100// HasEdgeFromTo returns whether an edge exists in the graph from u to v.
101func (g *DirectedGraph) HasEdgeFromTo(uid, vid int64) bool {
102	if _, ok := g.from[uid][vid]; !ok {
103		return false
104	}
105	return true
106}
107
108// NewEdge returns a new Edge from the source to the destination node.
109func (g *DirectedGraph) NewEdge(from, to graph.Node) graph.Edge {
110	return Edge{F: from, T: to}
111}
112
113// NewNode returns a new unique Node to be added to g. The Node's ID does
114// not become valid in g until the Node is added to g.
115func (g *DirectedGraph) NewNode() graph.Node {
116	if len(g.nodes) == 0 {
117		return Node(0)
118	}
119	if int64(len(g.nodes)) == uid.Max {
120		panic("simple: cannot allocate node: no slot")
121	}
122	return Node(g.nodeIDs.NewID())
123}
124
125// Node returns the node with the given ID if it exists in the graph,
126// and nil otherwise.
127func (g *DirectedGraph) Node(id int64) graph.Node {
128	return g.nodes[id]
129}
130
131// Nodes returns all the nodes in the graph.
132//
133// The returned graph.Nodes is only valid until the next mutation of
134// the receiver.
135func (g *DirectedGraph) Nodes() graph.Nodes {
136	if len(g.nodes) == 0 {
137		return graph.Empty
138	}
139	return iterator.NewNodes(g.nodes)
140}
141
142// RemoveEdge removes the edge with the given end point IDs from the graph, leaving the terminal
143// nodes. If the edge does not exist it is a no-op.
144func (g *DirectedGraph) RemoveEdge(fid, tid int64) {
145	if _, ok := g.nodes[fid]; !ok {
146		return
147	}
148	if _, ok := g.nodes[tid]; !ok {
149		return
150	}
151
152	delete(g.from[fid], tid)
153	delete(g.to[tid], fid)
154}
155
156// RemoveNode removes the node with the given ID from the graph, as well as any edges attached
157// to it. If the node is not in the graph it is a no-op.
158func (g *DirectedGraph) RemoveNode(id int64) {
159	if _, ok := g.nodes[id]; !ok {
160		return
161	}
162	delete(g.nodes, id)
163
164	for from := range g.from[id] {
165		delete(g.to[from], id)
166	}
167	delete(g.from, id)
168
169	for to := range g.to[id] {
170		delete(g.from[to], id)
171	}
172	delete(g.to, id)
173
174	g.nodeIDs.Release(id)
175}
176
177// SetEdge adds e, an edge from one node to another. If the nodes do not exist, they are added
178// and are set to the nodes of the edge otherwise.
179// It will panic if the IDs of the e.From and e.To are equal.
180func (g *DirectedGraph) SetEdge(e graph.Edge) {
181	var (
182		from = e.From()
183		fid  = from.ID()
184		to   = e.To()
185		tid  = to.ID()
186	)
187
188	if fid == tid {
189		panic("simple: adding self edge")
190	}
191
192	if _, ok := g.nodes[fid]; !ok {
193		g.AddNode(from)
194	} else {
195		g.nodes[fid] = from
196	}
197	if _, ok := g.nodes[tid]; !ok {
198		g.AddNode(to)
199	} else {
200		g.nodes[tid] = to
201	}
202
203	if fm, ok := g.from[fid]; ok {
204		fm[tid] = e
205	} else {
206		g.from[fid] = map[int64]graph.Edge{tid: e}
207	}
208	if tm, ok := g.to[tid]; ok {
209		tm[fid] = e
210	} else {
211		g.to[tid] = map[int64]graph.Edge{fid: e}
212	}
213}
214
215// To returns all nodes in g that can reach directly to n.
216//
217// The returned graph.Nodes is only valid until the next mutation of
218// the receiver.
219func (g *DirectedGraph) To(id int64) graph.Nodes {
220	if len(g.to[id]) == 0 {
221		return graph.Empty
222	}
223	return iterator.NewNodesByEdge(g.nodes, g.to[id])
224}
225