Implement IPSet.Invert and IPSet.Intersect.

Fixes #123.

Signed-off-by: David Anderson <[email protected]>

ipset.go

@@ -22,6 +22,22 @@ type IPSet struct {
 	out []IPRange
 }
 
+// toInOnly updates s to clear s.out, by merging any s.out into s.in.
+func (s *IPSet) toInOnly() {
+	if len(s.out) > 0 {
+		s.in = s.Ranges()
+		s.out = nil
+	}
+}
+
+// Clone returns a copy of s that shares no memory with s.
+func (s *IPSet) Clone() *IPSet {
+	return &IPSet{
+		in:  append([]IPRange(nil), s.in...),
+		out: append([]IPRange(nil), s.out...),
+	}
+}
+
 // Add adds ip to the set s.
 func (s *IPSet) Add(ip IP) { s.AddRange(IPRange{ip, ip}) }
 
@@ -35,10 +51,7 @@ func (s *IPSet) AddRange(r IPRange) {
 	}
 	// If there are any removals (s.out), then we need to compact the set
 	// first to get the order right.
-	if len(s.out) > 0 {
-		s.in = s.Ranges()
-		s.out = nil
-	}
+	s.toInOnly()
 	s.in = append(s.in, r)
 }
 
@@ -99,6 +112,23 @@ func (s *IPSet) RemoveSet(b *IPSet) {
 	}
 }
 
+// Invert updates s to contain the opposite of is current contents.
+func (s *IPSet) Invert() {
+	s.toInOnly()
+	s.out = s.in
+	s.in = []IPRange{
+		IPPrefix{IP: IPv4(0, 0, 0, 0), Bits: 0}.Range(),
+		IPPrefix{IP: IPv6Unspecified(), Bits: 0}.Range(),
+	}
+}
+
+// Intersect updates s to the set intersection of s and b.
+func (s *IPSet) Intersect(b *IPSet) {
+	b = b.Clone()
+	b.Invert()
+	s.RemoveSet(b)
+}
+
 func discardf(format string, args ...interface{}) {}
 
 // debugf is reassigned by tests.

ipset_test.go

@@ -226,6 +226,82 @@ func TestIPSet(t *testing.T) {
 			},
 			wantPrefixes: pxv("10.0.0.0/30"),
 		},
+		{
+			name: "invert_empty",
+			f: func(s *IPSet) {
+				s.Invert()
+			},
+			wantRanges: []IPRange{
+				{mustIP("0.0.0.0"), mustIP("255.255.255.255")},
+				{mustIP("::"), mustIP("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff")},
+			},
+			wantPrefixes: pxv("0.0.0.0/0", "::/0"),
+		},
+		{
+			name: "invert_full",
+			f: func(s *IPSet) {
+				s.AddPrefix(mustIPPrefix("0.0.0.0/0"))
+				s.AddPrefix(mustIPPrefix("::/0"))
+				s.Invert()
+			},
+			wantRanges: []IPRange{},
+			wantPrefixes: pxv(),
+		},
+		{
+			name: "invert_partial",
+			f: func(s *IPSet) {
+				s.AddRange(IPRange{mustIP("1.1.1.1"), mustIP("2.2.2.2")})
+				s.Add(mustIP("3.3.3.3"))
+				s.AddPrefix(mustIPPrefix("4.4.4.0/24"))
+				s.Add(mustIP("1::1"))
+				s.Invert()
+			},
+			wantRanges: []IPRange{
+				{mustIP("0.0.0.0"), mustIP("1.1.1.0")},
+				{mustIP("2.2.2.3"), mustIP("3.3.3.2")},
+				{mustIP("3.3.3.4"), mustIP("4.4.3.255")},
+				{mustIP("4.4.5.0"), mustIP("255.255.255.255")},
+				{mustIP("::"), mustIP("1::")},
+				{mustIP("1::2"), mustIP("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff")},
+			},
+		},
+		{
+			name: "intersect",
+			f: func(s *IPSet) {
+				var t IPSet
+				t.AddRange(IPRange{mustIP("2.2.2.2"), mustIP("3.3.3.3")})
+
+				s.AddRange(IPRange{mustIP("1.1.1.1"), mustIP("4.4.4.4")})
+				s.Intersect(&t)
+			},
+			wantRanges: []IPRange{
+				{mustIP("2.2.2.2"), mustIP("3.3.3.3")},
+			},
+		},
+		{
+			name: "intersect_disjoint",
+			f: func(s *IPSet) {
+				var t IPSet
+				t.AddRange(IPRange{mustIP("1.1.1.1"), mustIP("2.2.2.2")})
+
+				s.AddRange(IPRange{mustIP("3.3.3.3"), mustIP("4.4.4.4")})
+				s.Intersect(&t)
+			},
+			wantRanges: []IPRange{},
+		},
+		{
+			name: "intersect_partial",
+			f: func(s *IPSet) {
+				var t IPSet
+				t.AddRange(IPRange{mustIP("1.1.1.1"), mustIP("3.3.3.3")})
+
+				s.AddRange(IPRange{mustIP("2.2.2.2"), mustIP("4.4.4.4")})
+				s.Intersect(&t)
+			},
+			wantRanges: []IPRange{
+				{mustIP("2.2.2.2"), mustIP("3.3.3.3")},
+			},
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
@@ -236,7 +312,6 @@ func TestIPSet(t *testing.T) {
 			got := s.Ranges()
 			t.Run("ranges", func(t *testing.T) {
 				for _, v := range got {
-					debugf("%s -> %s", v.From, v.To)
 					if !v.Valid() {
 						t.Errorf("invalid IPRange in result: %s -> %s", v.From, v.To)
 					}