sub/ipalloc/ipalloc.go

package ipalloc

import (
	"fmt"
	"log"
	"net/netip"
	"slices"
	"strings"
	"sync"
)

type tree struct {
	prefix     netip.Prefix
	lo         *tree
	hi         *tree
	parent     *tree
	minfreelen int
}

func (t *tree) Prefix() netip.Prefix {
	return t.prefix
}
func (t *tree) Leaf() bool {
	return t != nil && t.lo == nil && t.hi == nil
}
func (t *tree) Full() bool {
	return t != nil && t.lo != nil && t.hi != nil
}
func (t *tree) Walk(f func(*tree) bool) {
	if t == nil {
		return
	}
	t.lo.Walk(f)
	if !f(t) {
		return
	}
	t.hi.Walk(f)
}
func (t *tree) haveSpace() bool {
	return t != nil && !t.Full() && (t.parent == nil || !t.Leaf())
}

func (t *tree) Find(p netip.Prefix) *tree {
	best := t.FindBestPrefix(p)
	if best != nil && best.prefix == p {
		return best
	}
	return nil
}
func (t *tree) FindBestPrefix(p netip.Prefix) *tree {
	if t == nil {
		return nil
	}
	if t.prefix == p {
		return t
	}
	if !t.prefix.Overlaps(p) {
		return nil
	}
	if best := t.lo.FindBestPrefix(p); best != nil {
		return best
	}
	if best := t.hi.FindBestPrefix(p); best != nil {
		return best
	}
	return t
}
func (t *tree) FindBestPrefixLen(bits int) *tree {
	if t == nil {
		return nil
	}
	if t.Leaf() {
		// already taken
		return nil
	}
	if t.minfreelen > bits {
		return nil
	}
	if t.prefix.Bits() >= bits {
		return nil
	}

	var best *tree

	lobest := t.lo.FindBestPrefixLen(bits)
	hibest := t.hi.FindBestPrefixLen(bits)
	switch {
	case lobest == nil:
		best = hibest
	case hibest == nil:
		best = lobest
	case lobest.minfreelen >= hibest.minfreelen:
		best = lobest
	default:
		best = hibest
	}

	if best == nil {
		best = t
	}

	return best
}

func (t *tree) Alloc(bits int) *tree {
	//log.Println("Searching space for", bits)
	best := t.FindBestPrefixLen(bits)
	if best == nil {
		best = t
	}
	if best.Full() {
		return nil
	}

	tofix := best

	// the first insert can be into a new hi branch
	if best.lo != nil {
		best.hi = &tree{
			prefix:     SplitPrefixHi(best.prefix),
			parent:     best,
			minfreelen: best.prefix.Bits() + 2,
		}
		best = best.hi
	}

	// the rest of the inserts are always into the lo branch
	for best.prefix.Bits() < bits {
		best.lo = &tree{
			prefix:     SplitPrefixLo(best.prefix),
			parent:     best,
			minfreelen: best.prefix.Bits() + 2,
		}
		best = best.lo
	}

	best.fixMinFree()
	tofix.fixMinFreeAll()

	return best
}

func (t *tree) fixMinFree() {
	if t == nil {
		return
	}
	if t.Leaf() {
		t.minfreelen = t.prefix.Addr().BitLen() + 1
		return
	}
	lofree := t.lo.calcfreelen(t)
	hifree := t.hi.calcfreelen(t)
	minfreelen := max(min(lofree, hifree), t.prefix.Bits()+1)

	//if t.minfreelen != minfreelen {
	//	log.Printf("lo:%v, hi:%v\n", t.lo == nil, t.hi == nil)
	//	log.Printf("Fixing space for %q, lo: %d, hi: %d, old: %d, new: %d", t.prefix, lofree, hifree, t.minfreelen, minfreelen)
	//}

	t.minfreelen = minfreelen
}
func (t *tree) fixMinFreeAll() {
	if t == nil {
		return
	}

	fixfree := t

	for fixfree != nil {
		fixfree.fixMinFree()
		fixfree = fixfree.parent
	}
}

func (t *tree) calcfreelen(parent *tree) int {
	if t == nil {
		return parent.prefix.Bits() + 1
	}
	return t.minfreelen
}

func (t *tree) Insert(p netip.Prefix) *tree {
	t = t.FindBestPrefix(p)

	if !t.haveSpace() {
		return nil
	}

	for t.prefix.Bits() < p.Bits() {
		if lo := SplitPrefixLo(t.prefix); lo.Overlaps(p) {
			t.lo = &tree{prefix: lo, parent: t}
			t = t.lo
		} else if hi := SplitPrefixHi(t.prefix); hi.Overlaps(p) {
			t.hi = &tree{prefix: hi, parent: t}
			t = t.hi
		} else {
			return nil
		}
	}
	t.fixMinFreeAll()
	return t
}

type ErrNotFound struct{}

func (e ErrNotFound) Error() string {
	return "prefix not found"
}

type ErrNotEmpty struct{}

func (e ErrNotEmpty) Error() string {
	return "prefix not empty"
}

func (t *tree) Dealloc(p netip.Prefix) error {
	match := t.Find(p)
	if match == nil {
		return ErrNotFound{}
	}
	if !match.Leaf() {
		return ErrNotEmpty{}
	}

	for match != nil {
		if match.Leaf() && match.parent != nil {
			if match.parent.lo == match {
				match.parent.lo = nil
			}
			if match.parent.hi == match {
				match.parent.hi = nil
			}
		}

		match.fixMinFree()
		match = match.parent
	}

	return nil
}

func (t *tree) Dot() string {
	var buf strings.Builder

	fmt.Fprintln(&buf, "strict digraph {")

	t.Walk(func(t *tree) bool {
		printEdge := func(parent, child *tree) {
			if child == nil {
				return
			}
			fmt.Fprintf(&buf, "\t%q -> %q;\n", parent.prefix, child.prefix)
		}

		fmt.Fprintf(&buf, "\t%q [shape=none label=<<table border=\"0\" cellspacing=\"0\" cellborder=\"1\"><tr><td>prefix</td><td>%s</td></tr><tr><td>minfreelen</td><td>%d</td></tr></table>>];\n", t.prefix, t.prefix, t.minfreelen)

		printEdge(t, t.lo)
		printEdge(t, t.hi)

		return true
	})
	fmt.Fprintln(&buf, "}")

	return buf.String()
}

type DB struct {
	sync.Mutex
	provisions     *tree
	allowedLengths []int
}

func NewDB(prefix string, allowedLengths ...int) *DB {
	p := netip.MustParsePrefix(prefix)
	if canonical := p.Masked(); p.Addr() != canonical.Addr() {
		log.Fatalf("Prefix %q is not in canonical form, use: %q", p, canonical)
	}
	return &DB{provisions: &tree{prefix: p}, allowedLengths: allowedLengths}
}

func (db *DB) AllowedLengths() []int {
	return slices.Clip(db.allowedLengths)
}
func (db *DB) Used() []netip.Prefix {
	db.Lock()
	defer db.Unlock()
	var ret []netip.Prefix
	db.provisions.Walk(func(t *tree) bool {
		if t.Leaf() && t.parent != nil {
			ret = append(ret, t.prefix)
		}
		return true
	})

	return ret
}
func (db *DB) UsedGrouped() [][]netip.Prefix {
	db.Lock()
	defer db.Unlock()
	bins := make(map[int][]netip.Prefix)
	db.provisions.Walk(func(t *tree) bool {
		if t.Leaf() && t.parent != nil {
			p := t.prefix
			bins[p.Bits()] = append(bins[p.Bits()], p)
		}
		return true
	})

	var ret [][]netip.Prefix
	for _, bin := range bins {
		ret = append(ret, bin)
	}

	slices.SortFunc(ret, func(ps1, ps2 []netip.Prefix) int {
		return ps2[0].Bits() - ps1[0].Bits()
	})

	return ret
}
func (db *DB) All() []*tree {
	db.Lock()
	defer db.Unlock()
	var ret []*tree
	db.provisions.Walk(func(t *tree) bool {
		ret = append(ret, t)
		return true
	})
	return ret
}
func (db *DB) Dot() string {
	return db.provisions.Dot()
}
func (db *DB) Alloc(bits int) (*tree, error) {
	db.Lock()
	defer db.Unlock()
	t := db.provisions.Alloc(bits)
	if t == nil {
		return nil, fmt.Errorf("Allocation not possible for size %d", bits)
	}
	return t, nil
}
func (db *DB) Insert(prefix netip.Prefix) error {
	return nil
}
func (db *DB) Dealloc(prefix netip.Prefix) error {
	db.Lock()
	defer db.Unlock()
	return db.provisions.Dealloc(prefix)
}
func (db *DB) Free() ([]netip.Prefix, error) {
	db.Lock()
	defer db.Unlock()
	return nil, nil
}
func (db *DB) Find(p netip.Prefix) bool {
	db.Lock()
	defer db.Unlock()
	p = p.Masked()
	return db.provisions.Find(p) != nil
}
func (db *DB) Root() netip.Prefix {
	db.Lock()
	defer db.Unlock()
	return db.provisions.prefix
}

//func (db *DB) Provision(p netip.Prefix) error {
//	db.Lock()
//	defer db.Unlock()
//	p = p.Masked()
//	if db.provisions.Provision(p) == nil {
//		return fmt.Errorf("Overlapping provisions: %s", p)
//	}
//	return nil
//}

type Treemap struct {
	ViewBox Rect
	Regions []Region
}
type Region struct {
	Prefix netip.Prefix
	Rect   Rect
}

func (db *DB) Treemap() Treemap {
	var ret Treemap
	root := db.provisions.prefix
	ret.ViewBox = prefixShape(root)
	// leave room for shadows
	ret.ViewBox.X--
	ret.ViewBox.Y--
	ret.ViewBox.W += 2
	ret.ViewBox.H += 2

	for _, p := range db.Used() {
		ret.Regions = append(ret.Regions, Region{Prefix: p, Rect: placeInto(p, root)})
	}

	return ret
}

type Rect struct {
	X, Y, W, H int
}

func prefixShape(p netip.Prefix) Rect {
	bits := p.Addr().BitLen() - p.Bits()

	return Rect{
		Y: 0,
		X: 0,
		H: 1 << (bits / 2), // TODO: handle overflow for normal ipv6 subnet sizes :)
		W: 1 << ((bits + 1) / 2),
	}
}

func placeInto(p, into netip.Prefix) Rect {
	return placeIntoRect(p, into, prefixShape(into))
}
func placeIntoRect(p, into netip.Prefix, r Rect) Rect {
	if p == into {
		return r
	}
	if p.Bits() < into.Bits() {
		panic("unreachable")
	}

	if lo := SplitPrefixLo(into); lo.Overlaps(p) {
		var subr Rect
		if r.W > r.H {
			subr.Y = r.Y
			subr.X = r.X
			subr.H = r.H
			subr.W = r.W / 2
		} else {
			subr.Y = r.Y
			subr.X = r.X
			subr.H = r.H / 2
			subr.W = r.W
		}
		return placeIntoRect(p, lo, subr)
	} else if hi := SplitPrefixHi(into); hi.Overlaps(p) {
		var subr Rect
		if r.W > r.H {
			subr.Y = r.Y
			subr.X = (r.X + (r.X + r.W)) / 2
			subr.H = r.H
			subr.W = r.W / 2
		} else {
			subr.Y = (r.Y + (r.Y + r.H)) / 2
			subr.X = r.X
			subr.H = r.H / 2
			subr.W = r.W
		}
		return placeIntoRect(p, hi, subr)
	}

	panic("unreachable")
}

func prepareSplit(p netip.Prefix) netip.Prefix {
	if p.IsSingleIP() {
		panic(fmt.Sprint("can't split single ip prefix ", p))
	}
	return p.Masked()
}
func SplitPrefixLo(p netip.Prefix) netip.Prefix {
	p = prepareSplit(p)
	return netip.PrefixFrom(p.Addr(), p.Bits()+1)
}
func SplitPrefixHi(p netip.Prefix) netip.Prefix {
	p = prepareSplit(p)

	bs := p.Addr().AsSlice()
	off := p.Bits() / 8
	bit := p.Bits() % 8
	bs[off] |= 0x80 >> bit
	hiaddr, ok := netip.AddrFromSlice(bs)
	if !ok {
		log.Fatal("can't use slice as addr", bs)
	}
	return netip.PrefixFrom(hiaddr, p.Bits()+1)
}
func splitPrefix(p netip.Prefix) (netip.Prefix, netip.Prefix) {
	return SplitPrefixLo(p), SplitPrefixHi(p)
}