blob: 87f87c538e16361a1122dd2fa3bbf4e719cc3b94 [file] [log] [blame]
khenaidooac637102019-01-14 15:44:34 -05001package goraph
2
3import (
4 "bytes"
5 "encoding/json"
6 "fmt"
7 "io"
8 "sync"
9)
10
11// ID is unique identifier.
12type ID interface {
13 // String returns the string ID.
14 String() string
15}
16
17type StringID string
18
19func (s StringID) String() string {
20 return string(s)
21}
22
23// Node is vertex. The ID must be unique within the graph.
24type Node interface {
25 // ID returns the ID.
26 ID() ID
27 String() string
28}
29
30type node struct {
31 id string
32}
33
34var nodeCnt uint64
35
36func NewNode(id string) Node {
37 return &node{
38 id: id,
39 }
40}
41
42func (n *node) ID() ID {
43 return StringID(n.id)
44}
45
46func (n *node) String() string {
47 return n.id
48}
49
50// Edge connects between two Nodes.
51type Edge interface {
52 Source() Node
53 Target() Node
54 Weight() float64
55 String() string
56}
57
58// edge is an Edge from Source to Target.
59type edge struct {
60 src Node
61 tgt Node
62 wgt float64
63}
64
65func NewEdge(src, tgt Node, wgt float64) Edge {
66 return &edge{
67 src: src,
68 tgt: tgt,
69 wgt: wgt,
70 }
71}
72
73func (e *edge) Source() Node {
74 return e.src
75}
76
77func (e *edge) Target() Node {
78 return e.tgt
79}
80
81func (e *edge) Weight() float64 {
82 return e.wgt
83}
84
85func (e *edge) String() string {
86 return fmt.Sprintf("%s -- %.3f -→ %s\n", e.src, e.wgt, e.tgt)
87}
88
89type EdgeSlice []Edge
90
91func (e EdgeSlice) Len() int { return len(e) }
92func (e EdgeSlice) Less(i, j int) bool { return e[i].Weight() < e[j].Weight() }
93func (e EdgeSlice) Swap(i, j int) { e[i], e[j] = e[j], e[i] }
94
95// Graph describes the methods of graph operations.
96// It assumes that the identifier of a Node is unique.
97// And weight values is float64.
98type Graph interface {
99 // Init initializes a Graph.
100 Init()
101
102 // GetNodeCount returns the total number of nodes.
103 GetNodeCount() int
104
105 // GetNode finds the Node. It returns nil if the Node
106 // does not exist in the graph.
107 GetNode(id ID) Node
108
109 // GetNodes returns a map from node ID to
110 // empty struct value. Graph does not allow duplicate
111 // node ID or name.
112 GetNodes() map[ID]Node
113
114 // AddNode adds a node to a graph, and returns false
115 // if the node already existed in the graph.
116 AddNode(nd Node) bool
117
118 // DeleteNode deletes a node from a graph.
119 // It returns true if it got deleted.
120 // And false if it didn't get deleted.
121 DeleteNode(id ID) bool
122
123 // AddEdge adds an edge from nd1 to nd2 with the weight.
124 // It returns error if a node does not exist.
125 AddEdge(id1, id2 ID, weight float64) error
126
127 // ReplaceEdge replaces an edge from id1 to id2 with the weight.
128 ReplaceEdge(id1, id2 ID, weight float64) error
129
130 // DeleteEdge deletes an edge from id1 to id2.
131 DeleteEdge(id1, id2 ID) error
132
133 // GetWeight returns the weight from id1 to id2.
134 GetWeight(id1, id2 ID) (float64, error)
135
136 // GetSources returns the map of parent Nodes.
137 // (Nodes that come towards the argument vertex.)
138 GetSources(id ID) (map[ID]Node, error)
139
140 // GetTargets returns the map of child Nodes.
141 // (Nodes that go out of the argument vertex.)
142 GetTargets(id ID) (map[ID]Node, error)
143
144 // String describes the Graph.
145 String() string
146}
147
148// graph is an internal default graph type that
149// implements all methods in Graph interface.
150type graph struct {
151 mu sync.RWMutex // guards the following
152
153 // idToNodes stores all nodes.
154 idToNodes map[ID]Node
155
156 // nodeToSources maps a Node identifer to sources(parents) with edge weights.
157 nodeToSources map[ID]map[ID]float64
158
159 // nodeToTargets maps a Node identifer to targets(children) with edge weights.
160 nodeToTargets map[ID]map[ID]float64
161}
162
163// newGraph returns a new graph.
164func newGraph() *graph {
165 return &graph{
166 idToNodes: make(map[ID]Node),
167 nodeToSources: make(map[ID]map[ID]float64),
168 nodeToTargets: make(map[ID]map[ID]float64),
169 //
170 // without this
171 // panic: assignment to entry in nil map
172 }
173}
174
175// NewGraph returns a new graph.
176func NewGraph() Graph {
177 return newGraph()
178}
179
180func (g *graph) Init() {
181 // (X) g = newGraph()
182 // this only updates the pointer
183 //
184 //
185 // (X) *g = *newGraph()
186 // assignment copies lock value
187
188 g.idToNodes = make(map[ID]Node)
189 g.nodeToSources = make(map[ID]map[ID]float64)
190 g.nodeToTargets = make(map[ID]map[ID]float64)
191}
192
193func (g *graph) GetNodeCount() int {
194 g.mu.RLock()
195 defer g.mu.RUnlock()
196
197 return len(g.idToNodes)
198}
199
200func (g *graph) GetNode(id ID) Node {
201 g.mu.RLock()
202 defer g.mu.RUnlock()
203
204 return g.idToNodes[id]
205}
206
207func (g *graph) GetNodes() map[ID]Node {
208 g.mu.RLock()
209 defer g.mu.RUnlock()
210
211 return g.idToNodes
212}
213
214func (g *graph) unsafeExistID(id ID) bool {
215 _, ok := g.idToNodes[id]
216 return ok
217}
218
219func (g *graph) AddNode(nd Node) bool {
220 g.mu.Lock()
221 defer g.mu.Unlock()
222
223 if g.unsafeExistID(nd.ID()) {
224 return false
225 }
226
227 id := nd.ID()
228 g.idToNodes[id] = nd
229 return true
230}
231
232func (g *graph) DeleteNode(id ID) bool {
233 g.mu.Lock()
234 defer g.mu.Unlock()
235
236 if !g.unsafeExistID(id) {
237 return false
238 }
239
240 delete(g.idToNodes, id)
241
242 delete(g.nodeToTargets, id)
243 for _, smap := range g.nodeToTargets {
244 delete(smap, id)
245 }
246
247 delete(g.nodeToSources, id)
248 for _, smap := range g.nodeToSources {
249 delete(smap, id)
250 }
251
252 return true
253}
254
255func (g *graph) AddEdge(id1, id2 ID, weight float64) error {
256 g.mu.Lock()
257 defer g.mu.Unlock()
258
259 if !g.unsafeExistID(id1) {
260 return fmt.Errorf("%s does not exist in the graph.", id1)
261 }
262 if !g.unsafeExistID(id2) {
263 return fmt.Errorf("%s does not exist in the graph.", id2)
264 }
265
266 if _, ok := g.nodeToTargets[id1]; ok {
267 if v, ok2 := g.nodeToTargets[id1][id2]; ok2 {
268 g.nodeToTargets[id1][id2] = v + weight
269 } else {
270 g.nodeToTargets[id1][id2] = weight
271 }
272 } else {
273 tmap := make(map[ID]float64)
274 tmap[id2] = weight
275 g.nodeToTargets[id1] = tmap
276 }
277 if _, ok := g.nodeToSources[id2]; ok {
278 if v, ok2 := g.nodeToSources[id2][id1]; ok2 {
279 g.nodeToSources[id2][id1] = v + weight
280 } else {
281 g.nodeToSources[id2][id1] = weight
282 }
283 } else {
284 tmap := make(map[ID]float64)
285 tmap[id1] = weight
286 g.nodeToSources[id2] = tmap
287 }
288
289 return nil
290}
291
292func (g *graph) ReplaceEdge(id1, id2 ID, weight float64) error {
293 g.mu.Lock()
294 defer g.mu.Unlock()
295
296 if !g.unsafeExistID(id1) {
297 return fmt.Errorf("%s does not exist in the graph.", id1)
298 }
299 if !g.unsafeExistID(id2) {
300 return fmt.Errorf("%s does not exist in the graph.", id2)
301 }
302
303 if _, ok := g.nodeToTargets[id1]; ok {
304 g.nodeToTargets[id1][id2] = weight
305 } else {
306 tmap := make(map[ID]float64)
307 tmap[id2] = weight
308 g.nodeToTargets[id1] = tmap
309 }
310 if _, ok := g.nodeToSources[id2]; ok {
311 g.nodeToSources[id2][id1] = weight
312 } else {
313 tmap := make(map[ID]float64)
314 tmap[id1] = weight
315 g.nodeToSources[id2] = tmap
316 }
317 return nil
318}
319
320func (g *graph) DeleteEdge(id1, id2 ID) error {
321 g.mu.Lock()
322 defer g.mu.Unlock()
323
324 if !g.unsafeExistID(id1) {
325 return fmt.Errorf("%s does not exist in the graph.", id1)
326 }
327 if !g.unsafeExistID(id2) {
328 return fmt.Errorf("%s does not exist in the graph.", id2)
329 }
330
331 if _, ok := g.nodeToTargets[id1]; ok {
332 if _, ok := g.nodeToTargets[id1][id2]; ok {
333 delete(g.nodeToTargets[id1], id2)
334 }
335 }
336 if _, ok := g.nodeToSources[id2]; ok {
337 if _, ok := g.nodeToSources[id2][id1]; ok {
338 delete(g.nodeToSources[id2], id1)
339 }
340 }
341 return nil
342}
343
344func (g *graph) GetWeight(id1, id2 ID) (float64, error) {
345 g.mu.RLock()
346 defer g.mu.RUnlock()
347
348 if !g.unsafeExistID(id1) {
349 return 0, fmt.Errorf("%s does not exist in the graph.", id1)
350 }
351 if !g.unsafeExistID(id2) {
352 return 0, fmt.Errorf("%s does not exist in the graph.", id2)
353 }
354
355 if _, ok := g.nodeToTargets[id1]; ok {
356 if v, ok := g.nodeToTargets[id1][id2]; ok {
357 return v, nil
358 }
359 }
360 return 0.0, fmt.Errorf("there is no edge from %s to %s", id1, id2)
361}
362
363func (g *graph) GetSources(id ID) (map[ID]Node, error) {
364 g.mu.RLock()
365 defer g.mu.RUnlock()
366
367 if !g.unsafeExistID(id) {
368 return nil, fmt.Errorf("%s does not exist in the graph.", id)
369 }
370
371 rs := make(map[ID]Node)
372 if _, ok := g.nodeToSources[id]; ok {
373 for n := range g.nodeToSources[id] {
374 rs[n] = g.idToNodes[n]
375 }
376 }
377 return rs, nil
378}
379
380func (g *graph) GetTargets(id ID) (map[ID]Node, error) {
381 g.mu.RLock()
382 defer g.mu.RUnlock()
383
384 if !g.unsafeExistID(id) {
385 return nil, fmt.Errorf("%s does not exist in the graph.", id)
386 }
387
388 rs := make(map[ID]Node)
389 if _, ok := g.nodeToTargets[id]; ok {
390 for n := range g.nodeToTargets[id] {
391 rs[n] = g.idToNodes[n]
392 }
393 }
394 return rs, nil
395}
396
397func (g *graph) String() string {
398 g.mu.RLock()
399 defer g.mu.RUnlock()
400
401 buf := new(bytes.Buffer)
402 for id1, nd1 := range g.idToNodes {
403 nmap, _ := g.GetTargets(id1)
404 for id2, nd2 := range nmap {
405 weight, _ := g.GetWeight(id1, id2)
406 fmt.Fprintf(buf, "%s -- %.3f -→ %s\n", nd1, weight, nd2)
407 }
408 }
409 return buf.String()
410}
411
412// NewGraphFromJSON returns a new Graph from a JSON file.
413// Here's the sample JSON data:
414//
415// {
416// "graph_00": {
417// "S": {
418// "A": 100,
419// "B": 14,
420// "C": 200
421// },
422// "A": {
423// "S": 15,
424// "B": 5,
425// "D": 20,
426// "T": 44
427// },
428// "B": {
429// "S": 14,
430// "A": 5,
431// "D": 30,
432// "E": 18
433// },
434// "C": {
435// "S": 9,
436// "E": 24
437// },
438// "D": {
439// "A": 20,
440// "B": 30,
441// "E": 2,
442// "F": 11,
443// "T": 16
444// },
445// "E": {
446// "B": 18,
447// "C": 24,
448// "D": 2,
449// "F": 6,
450// "T": 19
451// },
452// "F": {
453// "D": 11,
454// "E": 6,
455// "T": 6
456// },
457// "T": {
458// "A": 44,
459// "D": 16,
460// "F": 6,
461// "E": 19
462// }
463// },
464// }
465//
466func NewGraphFromJSON(rd io.Reader, graphID string) (Graph, error) {
467 js := make(map[string]map[string]map[string]float64)
468 dec := json.NewDecoder(rd)
469 for {
470 if err := dec.Decode(&js); err == io.EOF {
471 break
472 } else if err != nil {
473 return nil, err
474 }
475 }
476 if _, ok := js[graphID]; !ok {
477 return nil, fmt.Errorf("%s does not exist", graphID)
478 }
479 gmap := js[graphID]
480
481 g := newGraph()
482 for id1, mm := range gmap {
483 nd1 := g.GetNode(StringID(id1))
484 if nd1 == nil {
485 nd1 = NewNode(id1)
486 if ok := g.AddNode(nd1); !ok {
487 return nil, fmt.Errorf("%s already exists", nd1)
488 }
489 }
490 for id2, weight := range mm {
491 nd2 := g.GetNode(StringID(id2))
492 if nd2 == nil {
493 nd2 = NewNode(id2)
494 if ok := g.AddNode(nd2); !ok {
495 return nil, fmt.Errorf("%s already exists", nd2)
496 }
497 }
498 g.ReplaceEdge(nd1.ID(), nd2.ID(), weight)
499 }
500 }
501
502 return g, nil
503}