From 009f26bb9d658f5c7269c77440c0c3625a50fb26 Mon Sep 17 00:00:00 2001
From: Sebastien Bourdeauducq <sebastien@milkymist.org>
Date: Fri, 8 Jun 2012 22:49:49 +0200
Subject: [PATCH] flow/network: refactor graph

---
 examples/dataflow.py     |  2 +-
 examples/dataflow_dma.py | 11 +++++------
 examples/dataflow_sim.py |  6 ++----
 migen/flow/composer.py   |  8 +++-----
 migen/flow/network.py    | 25 +++++++++++++------------
 5 files changed, 24 insertions(+), 28 deletions(-)

diff --git a/examples/dataflow.py b/examples/dataflow.py
index aab603f9a..5c3d5fc0e 100644
--- a/examples/dataflow.py
+++ b/examples/dataflow.py
@@ -6,7 +6,7 @@ from migen.flow.ala import *
 from migen.flow.network import *
 from migen.flow.composer import *
 
-g = nx.MultiDiGraph()
+g = DataFlowGraph()
 a1 = make_composable(g, Add(BV(16)))
 a2 = make_composable(g, Add(BV(16)))
 a3 = make_composable(g, Add(BV(16)))
diff --git a/examples/dataflow_dma.py b/examples/dataflow_dma.py
index 183c79ffc..b5f24a7cf 100644
--- a/examples/dataflow_dma.py
+++ b/examples/dataflow_dma.py
@@ -1,4 +1,3 @@
-import networkx as nx
 from random import Random
 
 from migen.fhdl import verilog
@@ -43,9 +42,9 @@ def test_reader():
 	adrgen = SimActor(adrgen_gen(), ("address", Source, [("a", BV(30))]))
 	reader = dma_wishbone.Reader()
 	dumper = SimActor(dumper_gen(), ("data", Sink, [("d", BV(32))]))
-	g = nx.MultiDiGraph()
-	add_connection(g, adrgen, reader)
-	add_connection(g, reader, dumper)
+	g = DataFlowGraph()
+	g.add_connection(adrgen, reader)
+	g.add_connection(reader, dumper)
 	comp = CompositeActor(g)
 	
 	peripheral = MyPeripheral()
@@ -73,8 +72,8 @@ def test_writer():
 	print("*** Testing writer")
 	trgen = SimActor(trgen_gen(), ("address_data", Source, [("a", BV(30)), ("d", BV(32))]))
 	writer = dma_wishbone.Writer()
-	g = nx.MultiDiGraph()
-	add_connection(g, trgen, writer)
+	g = DataFlowGraph()
+	g.add_connection(trgen, writer)
 	comp = CompositeActor(g)
 	
 	peripheral = MyPeripheral()
diff --git a/examples/dataflow_sim.py b/examples/dataflow_sim.py
index 3426a4fc9..98d79652e 100644
--- a/examples/dataflow_sim.py
+++ b/examples/dataflow_sim.py
@@ -1,5 +1,3 @@
-import networkx as nx
-
 from migen.fhdl.structure import *
 from migen.flow.actor import *
 from migen.flow.network import *
@@ -21,8 +19,8 @@ def sink_gen():
 def main():
 	source = SimActor(source_gen(), ("source", Source, [("value", BV(32))]))
 	sink = SimActor(sink_gen(), ("sink", Sink, [("value", BV(32))]))
-	g = nx.MultiDiGraph()
-	add_connection(g, source, sink)
+	g = DataFlowGraph()
+	g.add_connection(source, sink)
 	comp = CompositeActor(g)
 	def end_simulation(s):
 		s.interrupt = source.done
diff --git a/migen/flow/composer.py b/migen/flow/composer.py
index fe63d4da2..b9237dfb7 100644
--- a/migen/flow/composer.py
+++ b/migen/flow/composer.py
@@ -1,5 +1,3 @@
-import networkx as nx
-
 from migen.flow.actor import *
 from migen.flow.ala import *
 from migen.flow.plumbing import *
@@ -16,9 +14,9 @@ def _simple_binary(a, b, actor_class):
 	signed = signal_self.bv.signed and signal_other.bv.signed
 	actor = actor_class(BV(width, signed))
 	combinator = Combinator(actor.token("operands").layout(), ["a"], ["b"])
-	add_connection(a.dfg, combinator, actor)
-	add_connection(a.dfg, a.actor, combinator, a.endp, "sink0")
-	add_connection(a.dfg, b.actor, combinator, b.endp, "sink1")
+	a.dfg.add_connection(combinator, actor)
+	a.dfg.add_connection(a.actor, combinator, a.endp, "sink0")
+	a.dfg.add_connection(b.actor, combinator, b.endp, "sink1")
 	return make_composable(a.dfg, actor)
 
 class ComposableSource():
diff --git a/migen/flow/network.py b/migen/flow/network.py
index 5eb98ac7f..600384436 100644
--- a/migen/flow/network.py
+++ b/migen/flow/network.py
@@ -1,9 +1,21 @@
-import networkx as nx
+from networkx import MultiDiGraph
 
 from migen.fhdl.structure import *
 from migen.flow.actor import *
 from migen.corelogic.misc import optree
 
+class DataFlowGraph(MultiDiGraph):
+	def add_connection(self, source_node, sink_node, source_ep=None, sink_ep=None):
+		if source_ep is None:
+			source_eps = source_node.sources()
+			assert(len(source_eps) == 1)
+			source_ep = source_eps[0]
+		if sink_ep is None:
+			sink_eps = sink_node.sinks()
+			assert(len(sink_eps) == 1)
+			sink_ep = sink_eps[0]
+		self.add_edge(source_node, sink_node, source=source_ep, sink=sink_ep)
+
 class CompositeActor(Actor):
 	def __init__(self, dfg): # TODO: endpoints
 		self.dfg = dfg
@@ -16,14 +28,3 @@ class CompositeActor(Actor):
 		others = sum([node.get_fragment() for node in self.dfg], Fragment())
 		busy = Fragment([self.busy.eq(optree("|", [node.busy for node in self.dfg]))])
 		return this + others + busy
-
-def add_connection(dfg, source_node, sink_node, source_ep=None, sink_ep=None):
-	if source_ep is None:
-		source_eps = source_node.sources()
-		assert(len(source_eps) == 1)
-		source_ep = source_eps[0]
-	if sink_ep is None:
-		sink_eps = sink_node.sinks()
-		assert(len(sink_eps) == 1)
-		sink_ep = sink_eps[0]
-	dfg.add_edge(source_node, sink_node, source=source_ep, sink=sink_ep)