Add optimised RouteSet data type

rig/routing_table/__init__.py

@@ -1,5 +1,5 @@
 # Basic routing table datastructures
-from rig.routing_table.entries import RoutingTableEntry, Routes
+from rig.routing_table.entries import RoutingTableEntry, Routes, RouteSet
 
 # Common exceptions produced by algorithms in this module
 from rig.routing_table.exceptions import (MinimisationFailedError,

rig/routing_table/entries.py

@@ -8,6 +8,61 @@
 from collections import namedtuple
 
 
+class RouteSet(object):
+    __slots__ = ["_entries"]
+
+    def __init__(self, entries=tuple(), _entries=0x0):
+        """Create a new Route Set.
+
+        Parameters
+        ----------
+        entries : Route, ...
+            Elements to include in the set.
+        """
+        # Initialise the set
+        self._entries = _entries
+
+        # Include any items from the set
+        self.update(entries)
+
+    def add(self, elem):
+        self._entries |= (1 << (elem if elem is not None else 31))
+
+    def update(self, elems):
+        for e in elems:
+            self.add(e)
+
+    def __contains__(self, elem):
+        return self._entries & (1 << (elem if elem is not None else 31))
+
+    def __eq__(self, other):
+        if isinstance(other, RouteSet):
+            return self._entries == other._entries
+        else:
+            return set(self) == set(other)
+
+    def __int__(self):
+        return self._entries
+
+    def __ior__(self, other):
+        self.update(other)
+        return self
+
+    def __iter__(self):
+        for r in Routes:
+            if (1 << r) & self._entries:
+                yield r
+
+        if (1 << 31) & self._entries:
+            yield None
+
+    def __len__(self):
+        return sum(1 for _ in self)
+
+    def __sub__(self, elems):
+        return RouteSet(_entries=self._entries & ~RouteSet(elems)._entries)
+
+
 class RoutingTableEntry(namedtuple("RoutingTableEntry",
                                    "route key mask sources")):
     """Named tuple representing a single routing entry in a SpiNNaker routing
@@ -31,7 +86,7 @@ class RoutingTableEntry(namedtuple("RoutingTableEntry",
     """
     def __new__(cls, route, key, mask, sources={None}):
         return super(RoutingTableEntry, cls).__new__(
-            cls, frozenset(route), key, mask, set(sources)
+            cls, RouteSet(route), key, mask, RouteSet(sources)
         )
 
     def __str__(self):

rig/routing_table/utils.py

@@ -1,6 +1,7 @@
 from collections import defaultdict, namedtuple, OrderedDict
 
-from rig.routing_table import RoutingTableEntry, MultisourceRouteError
+from rig.routing_table import (
+    RoutingTableEntry, MultisourceRouteError, RouteSet)
 from six import iteritems
 import warnings
 
@@ -69,7 +70,8 @@ def routing_tree_to_tables(routes, net_keys):
             else:
                 # Otherwise create a new route set
                 route_sets[(x, y)][(key, mask)] = \
-                    InOutPair({in_direction}, set(out_directions))
+                    InOutPair(RouteSet({in_direction}),
+                              RouteSet(out_directions))
 
     # Construct the routing tables from the route sets
     routing_tables = defaultdict(list)

tests/routing_table/test_entries.py

@@ -2,7 +2,74 @@
 
 from rig.links import Links
 
-from rig.routing_table import Routes, RoutingTableEntry
+from rig.routing_table import Routes, RoutingTableEntry, RouteSet
+
+
+class TestRouteSet(object):
+    def test_or_update_and_contains_and_len(self):
+        # Create an empty route set
+        rs = RouteSet()
+        assert len(rs) == 0
+
+        # Include east
+        rs |= {Routes.east}
+        assert len(rs) == 1
+        assert set(rs) == {Routes.east}
+        assert Routes.east in rs
+
+        # Include north east
+        rs |= {Routes.north_east}
+        assert len(rs) == 2
+        assert set(rs) == {Routes.east, Routes.north_east}
+        assert Routes.north_east in rs
+        assert not Routes.north in rs
+
+        # Shouldn't do anything
+        rs |= {Routes.north_east}
+        assert len(rs) == 2
+        assert set(rs) == {Routes.east, Routes.north_east}
+        assert Routes.north_east in rs
+        assert not Routes.south in rs
+        assert not None in rs
+
+        # Check that None works
+        rs |= {Routes.south, None}
+        assert len(rs) == 4
+        assert set(rs) == {Routes.east, Routes.north_east, Routes.south, None}
+        assert None in rs
+
+        rs |= Routes
+        assert len(rs) == 25
+
+    def test_remove(self):
+        rs = RouteSet(Routes)
+
+        rs -= {Routes.north, Routes.south}
+        assert len(rs) == 22
+        assert Routes.north not in rs
+        assert Routes.south not in rs
+
+    def test_equality(self):
+        assert RouteSet({Routes.north}) == RouteSet({Routes.north})
+        assert RouteSet({None}) != RouteSet()
+        assert RouteSet({None}) != RouteSet({Routes.south_west})
+        assert RouteSet({None}) == RouteSet({None})
+        assert RouteSet({Routes.north}) != RouteSet({Routes.south})
+        assert RouteSet(Routes) != RouteSet()
+
+        assert RouteSet({Routes.north}) == {Routes.north}
+        assert RouteSet({None}) != set()
+        assert RouteSet({None}) != {Routes.south_west}
+        assert RouteSet({None}) == {None}
+        assert RouteSet({Routes.north}) != {Routes.south}
+        assert RouteSet(Routes) != set()
+
+    def test_int(self):
+        for r in Routes:
+            rs = RouteSet({r})
+            assert int(rs) == 1 << r
+
+        assert int(RouteSet({None})) == 0x80000000
 
 
 class TestRoutingTableEntry(object):