add merge insert to keep a minimum nodes that cover all the space

brute.go

@@ -43,6 +43,11 @@ func (b *bruteRanger) Insert(entry RangerEntry) error {
 	return nil
 }
 
+// MergeInsert not implement in bruteRanger
+func (b *bruteRanger) MergeInsert(_ RangerEntry) error {
+	return nil
+}
+
 // Remove removes a RangerEntry identified by given network from ranger.
 func (b *bruteRanger) Remove(network net.IPNet) (RangerEntry, error) {
 	networks, err := b.getEntriesByVersion(network.IP)

cidranger.go

@@ -85,6 +85,7 @@ func NewBasicRangerEntry(ipNet net.IPNet) RangerEntry {
 // Ranger is an interface for cidr block containment lookups.
 type Ranger interface {
 	Insert(entry RangerEntry) error
+	MergeInsert(entry RangerEntry) error
 	Remove(network net.IPNet) (RangerEntry, error)
 	Contains(ip net.IP) (bool, error)
 	ContainingNetworks(ip net.IP) ([]RangerEntry, error)

trie.go

@@ -87,6 +87,13 @@ func (p *prefixTrie) Insert(entry RangerEntry) error {
 	return err
 }
 
+// MergeInsert inserts a RangerEntry into prefix trie, and apply merge if possible
+func (p *prefixTrie) MergeInsert(entry RangerEntry) error {
+	network := rnet.NewNetwork(entry.Network())
+	_, err := p.mergeInsert(network.Number, network, entry)
+	return err
+}
+
 // Remove removes RangerEntry identified by given network from trie.
 func (p *prefixTrie) Remove(network net.IPNet) (RangerEntry, error) {
 	entry, err := p.remove(rnet.NewNetwork(network))
@@ -252,6 +259,94 @@ func (p *prefixTrie) insert(network rnet.Network, entry RangerEntry) (bool, erro
 	return existingChild.insert(network, entry)
 }
 
+func (p *prefixTrie) mergeInsert(number rnet.NetworkNumber, network rnet.Network, entry RangerEntry) (bool, error) {
+
+	if p.network.Equal(network) {
+		p.entry = entry
+		if p.childrenCount() > 0 {
+			// has child, which can be merged
+			p.children = p.children[:0]
+			return true, nil
+		}
+		// no child, no merge
+		return false, nil
+	}
+
+	if p.hasEntry() && p.network.Contains(number) {
+		// p contains network already, no need to insert
+		return false, nil
+	}
+
+	bit, err := p.targetBitFromIP(network.Number)
+	if err != nil {
+		return false, err
+	}
+
+	existingChild := p.children[bit]
+	// no child on this direction
+	if existingChild == nil {
+		ones, _ := network.IPNet.Mask.Size()
+		if uint(ones) != p.numBitsSkipped+1 {
+			// new network are not adjacent from p, just insert
+			leaf := newPathprefixTrie(network, uint(ones))
+			leaf.entry = entry
+			p.appendTrie(bit, leaf)
+			return false, nil
+		}
+
+		// new network are adjacent from p, check p's other child
+		if p.children[^bit&1] != nil && p.children[^bit&1].numBitsSkipped == p.numBitsSkipped+1 {
+			// reset the child
+			p.children = p.children[:0]
+			if p.entry == nil {
+				p.entry = NewBasicRangerEntry(p.network.IPNet)
+			}
+			return true, nil
+		}
+
+		// the other child are not meet the requirement for merge
+		p.appendTrie(bit, newEntryTrie(network, entry))
+		return false, nil
+	}
+
+	// there's child on this direction
+	// Check whether it is necessary to insert additional path prefix between current trie and existing child,
+	// in the case that inserted network diverges on its path to existing child.
+	lcb, err := network.LeastCommonBitPosition(existingChild.network)
+	divergingBitPos := int(lcb) - 1
+	if divergingBitPos > existingChild.targetBitPosition() {
+		pathPrefix := newPathprefixTrie(network, p.totalNumberOfBits()-lcb)
+		err := p.insertPrefix(bit, pathPrefix, existingChild)
+		if err != nil {
+			return false, err
+		}
+		// Update new child
+		existingChild = pathPrefix
+	}
+
+	merged, err := existingChild.mergeInsert(number, network, entry)
+	if merged {
+		// child merged, check if p need to merge as well
+		return p.mergeChildren(), err
+	}
+	return merged, err
+}
+
+func (p *prefixTrie) mergeChildren() bool {
+	for _, child := range p.children {
+		if child == nil || child.numBitsSkipped != p.numBitsSkipped+1 {
+			return false
+		}
+	}
+
+	// this can be merged
+	p.children = p.children[:0]
+	if p.entry == nil {
+		p.entry = NewBasicRangerEntry(p.network.IPNet)
+	}
+	return true
+}
+
 func (p *prefixTrie) appendTrie(bit uint32, prefix *prefixTrie) {
 	p.children[bit] = prefix
 	prefix.parent = p

trie_test.go

@@ -2,6 +2,7 @@ package cidranger
 
 import (
 	"encoding/binary"
+	"fmt"
 	"math/rand"
 	"net"
 	"runtime"
@@ -116,6 +117,50 @@ func TestPrefixTrieString(t *testing.T) {
 	assert.Equal(t, expected, trie.String())
 }
 
+func TestPrefixTrieString2(t *testing.T) {
+	trie := newPrefixTree(rnet.IPv4).(*prefixTrie)
+	inserts := []string{"8.8.8.1/25", "8.8.9.1/25", "8.8.10.1/24"}
+	for _, insert := range inserts {
+		_, network, _ := net.ParseCIDR(insert)
+		trie.Insert(NewBasicRangerEntry(*network))
+		fmt.Printf("%s\n\n", trie.String())
+	}
+}
+
+func TestPrefixTrieMergeInsert(t *testing.T) {
+	trie := newPrefixTree(rnet.IPv4).(*prefixTrie)
+	inserts := []string{
+		"8.8.8.1/25",
+		"8.8.9.1/25",
+		"8.8.10.1/24",
+		"8.8.8.128/25",
+		"8.8.9.128/25",
+		"8.8.11.128/25",
+		"8.8.11.0/25",
+	}
+	for _, insert := range inserts {
+		_, network, _ := net.ParseCIDR(insert)
+		trie.MergeInsert(NewBasicRangerEntry(*network))
+		fmt.Printf("%s\n\n", trie.String())
+	}
+}
+
+func TestPrefixTrieMergeInsert2(t *testing.T) {
+	trie := newPrefixTree(rnet.IPv4).(*prefixTrie)
+	inserts := []string{
+		"8.8.8.128/29",
+		"8.8.8.128/29",
+		"8.8.9.0/24",
+		"8.8.8.0/23",
+		"8.8.9.0/24",
+	}
+	for _, insert := range inserts {
+		_, network, _ := net.ParseCIDR(insert)
+		trie.MergeInsert(NewBasicRangerEntry(*network))
+		fmt.Printf("%s\n\n", trie.String())
+	}
+}
+
 func TestPrefixTrieRemove(t *testing.T) {
 	cases := []struct {
 		version                      rnet.IPVersion

version.go

@@ -29,6 +29,16 @@ func (v *versionedRanger) Insert(entry RangerEntry) error {
 	return ranger.Insert(entry)
 }
 
+// MergeInsert inserts a RangerEntry into prefix trie, and apply merge if possible
+func (v *versionedRanger) MergeInsert(entry RangerEntry) error {
+	network := entry.Network()
+	ranger, err := v.getRangerForIP(network.IP)
+	if err != nil {
+		return err
+	}
+	return ranger.MergeInsert(entry)
+}
+
 func (v *versionedRanger) Remove(network net.IPNet) (RangerEntry, error) {
 	ranger, err := v.getRangerForIP(network.IP)
 	if err != nil {