Set operations

svn path=/scripts/rpki/resource_set.py; revision=798

1 files changed, 150 insertions, 3 deletions

diff --git a/scripts/rpki/resource_set.py b/scripts/rpki/resource_set.py
index 57900d10..8a3d1751 100644
--- a/scripts/rpki/resource_set.py
+++ b/scripts/rpki/resource_set.py
@@ -62,17 +62,33 @@ class resource_range_ipv6(resource_range_ip):
   """
   pass
 
+def rsplit(rset, that):
+  this = rset.pop(0)
+  cell_type = type(this.min)
+  assert type(this) is type(that) and type(this.max) is cell_type and type(that.min) is cell_type and type(that.max) is cell_type
+  if this.min < that.min:
+    rset.insert(0, type(this)(this.min, cell_type(that.min - 1)))
+    rset.insert(1, type(this)(that.min, this.max))
+  else:
+    assert this.max > that.max
+    rset.insert(0, type(this)(this.min, that.max))
+    rset.insert(1, type(this)(cell_type(that.max + 1), this.max))
+
 class resource_set(list):
   """
-  Generic resource set.  List type containing resource ranges.
-  You probably don't want to use this type directly.
+  Generic resource set.  List type containing resource ranges.  You
+  probably don't want to use this type directly.
   """
 
   def __init__(self, ini=None):
     if isinstance(ini, str) and len(ini):
       self.extend(map(self.parse_str, ini.split(",")))
-    if isinstance(ini, tuple):
+    elif isinstance(ini, tuple):
       self.parse_tuple(ini)
+    elif isinstance(ini, list):
+      self.extend(ini)
+    else:
+      assert ini is None
     self.sort()
     if __debug__:
       for i in range(0, len(self) - 1):
@@ -81,6 +97,96 @@ class resource_set(list):
   def __str__(self):
     return ",".join(map(str, self))
 
+  def comm(self, other):
+    """
+    Like comm(1), sort of.  Returns a tuple of three resource sets:
+    resources only in self, resources only in other, and resources in
+    both.  Used (not very efficiently) as the basis for most set
+    operations on resource sets.
+    """
+    assert type(self) is type(other)
+    set1 = [i for i in self]
+    set2 = [i for i in other]
+    only1, only2, both = [], [], []
+    while set1 or set2:
+      if set1 and (not set2 or set1[0].max < set2[0].min):
+        only1.append(set1.pop(0))
+      elif set2 and (not set1 or set2[0].max < set1[0].min):
+        only2.append(set2.pop(0))
+      elif set1[0].min < set2[0].min:
+        rsplit(set1, set2[0])
+      elif set2[0].min < set1[0].min:
+        rsplit(set2, set1[0])
+      elif set1[0].max < set2[0].max:
+        rsplit(set2, set1[0])
+      elif set2[0].max < set1[0].max:
+        rsplit(set1, set2[0])
+      else:
+        assert set1[0].min == set2[0].min and set1[0].max == set2[0].max
+        both.append(set1.pop(0))
+        set2.pop(0)
+    return type(self)(only1), type(self)(only2), type(self)(both)
+
+  def union(self, other):
+    """
+    Set union for resource sets.
+    """
+    assert type(self) is type(other)
+    set1 = [i for i in self]
+    set2 = [i for i in other]
+    result = []
+    while set1 or set2:
+      if set1 and (not set2 or set1[0].max < set2[0].min):
+        result.append(set1.pop(0))
+      elif set2 and (not set1 or set2[0].max < set1[0].min):
+        result.append(set2.pop(0))
+      else:
+        this = set1.pop(0)
+        that = set2.pop(0)
+        assert type(this) is type(that)
+        if this.min < that.min:
+          min = this.min
+        else:
+          min = that.min
+        if this.max > that.max:
+          max = this.max
+        else:
+          max = that.max
+        result.append(type(this)(min, max))
+    return type(self)(result)
+
+  def intersection(self, other):
+    """
+    Set intersection for resource sets.
+    """
+    return self.comm(other)[2]
+
+  def difference(self, other):
+    """
+    Set difference for resource sets.
+    """
+    return self.comm(other)[0]
+
+  def symmetric_difference(self, other):
+    """
+    Set symmetric difference (XOR) for resource sets.
+    """
+    com = self.comm(other)
+    return com[0].union(com[1])
+
+  def contains(self, item):
+    """
+    Set membership test for resource sets.
+    """
+    for i in self:
+      if isinstance(item, type(i)) and i.min <= item.min and i.max >= item.max:
+        return True
+      elif isinstance(item, type(i.min)) and i.min <= item and i.max >= item:
+        return True
+      else:
+        assert isinstance(item, (type(i), type(i.min)))
+    return False
+
 class resource_set_as(resource_set):
   """
   ASN resource set.
@@ -185,3 +291,44 @@ def parse_extensions(exts):
           assert v6 is None
           v6 = resource_set_ipv6(fam[1])
   return as, v4, v6
+
+# Test suite for set operations.  This will probably go away eventually
+
+if __name__ == "__main__":
+
+  def test(t, s1, s2):
+    print
+    r1 = t(s1)
+    r2 = t(s2)
+    print "x:  ", r1
+    print "y:  ", r2
+    v1 = r1.comm(r2)
+    v2 = r2.comm(r1)
+    assert v1[0] == v2[1] and v1[1] == v2[0] and v1[2] == v2[2]
+    for i in r1: assert r1.contains(i) and r1.contains(i.min) and r1.contains(i.max)
+    for i in r2: assert r2.contains(i) and r2.contains(i.min) and r2.contains(i.max)
+    for i in v1[0]: assert r1.contains(i) and not r2.contains(i)
+    for i in v1[1]: assert not r1.contains(i) and r2.contains(i)
+    for i in v1[2]: assert r1.contains(i) and r2.contains(i)
+    v1 = r1.union(r2)
+    v2 = r2.union(r1)
+    assert v1 == v2
+    print "x|y:", v1
+    v1 = r1.difference(r2)
+    v2 = r2.difference(r1)
+    print "x-y:", v1
+    print "y-x:", v2
+    v1 = r1.symmetric_difference(r2)
+    v2 = r2.symmetric_difference(r1)
+    assert v1 == v2
+    print "x^y:", v1
+    v1 = r1.intersection(r2)
+    v2 = r2.intersection(r1)
+    assert v1 == v2
+    print "x&y:", v1
+
+  print "Testing set operations on resource sets"
+  test(resource_set_as, "1,2,3,4,5,6,11,12,13,14,15", "1,2,3,4,5,6,111,121,131,141,151")
+  test(resource_set_ipv4, "10.0.0.44/32,10.6.0.2/32", "10.3.0.0/24,10.0.0.77/32")
+  test(resource_set_ipv4, "10.0.0.44/32,10.6.0.2/32", "10.0.0.0/24")
+  test(resource_set_ipv4, "10.0.0.0/24", "10.3.0.0/24,10.0.0.77/32")