1 files changed, 190 insertions, 0 deletions
diff --git a/vendor/github.com/mmcloughlin/avo/pass/alloc.go b/vendor/github.com/mmcloughlin/avo/pass/alloc.go
new file mode 100644
index 0000000..fc7773a
--- /dev/null
+++ b/vendor/github.com/mmcloughlin/avo/pass/alloc.go
@@ -0,0 +1,190 @@
+package pass
+
+import (
+	"errors"
+	"math"
+	"sort"
+
+	"github.com/mmcloughlin/avo/reg"
+)
+
+// edge is an edge of the interference graph, indicating that registers X and Y
+// must be in non-conflicting registers.
+type edge struct {
+	X, Y reg.ID
+}
+
+// Allocator is a graph-coloring register allocator.
+type Allocator struct {
+	registers  []reg.ID
+	allocation reg.Allocation
+	edges      []*edge
+	possible   map[reg.ID][]reg.ID
+}
+
+// NewAllocator builds an allocator for the given physical registers.
+func NewAllocator(rs []reg.Physical) (*Allocator, error) {
+	// Set of IDs, excluding restricted registers.
+	idset := map[reg.ID]bool{}
+	for _, r := range rs {
+		if (r.Info() & reg.Restricted) != 0 {
+			continue
+		}
+		idset[r.ID()] = true
+	}
+
+	if len(idset) == 0 {
+		return nil, errors.New("no allocatable registers")
+	}
+
+	// Produce slice of unique register IDs.
+	var ids []reg.ID
+	for id := range idset {
+		ids = append(ids, id)
+	}
+	sort.Slice(ids, func(i, j int) bool { return ids[i] < ids[j] })
+
+	return &Allocator{
+		registers:  ids,
+		allocation: reg.NewEmptyAllocation(),
+		possible:   map[reg.ID][]reg.ID{},
+	}, nil
+}
+
+// NewAllocatorForKind builds an allocator for the given kind of registers.
+func NewAllocatorForKind(k reg.Kind) (*Allocator, error) {
+	f := reg.FamilyOfKind(k)
+	if f == nil {
+		return nil, errors.New("unknown register family")
+	}
+	return NewAllocator(f.Registers())
+}
+
+// AddInterferenceSet records that r interferes with every register in s. Convenience wrapper around AddInterference.
+func (a *Allocator) AddInterferenceSet(r reg.Register, s reg.MaskSet) {
+	for id, mask := range s {
+		if (r.Mask() & mask) != 0 {
+			a.AddInterference(r.ID(), id)
+		}
+	}
+}
+
+// AddInterference records that x and y must be assigned to non-conflicting physical registers.
+func (a *Allocator) AddInterference(x, y reg.ID) {
+	a.Add(x)
+	a.Add(y)
+	a.edges = append(a.edges, &edge{X: x, Y: y})
+}
+
+// Add adds a register to be allocated. Does nothing if the register has already been added.
+func (a *Allocator) Add(v reg.ID) {
+	if !v.IsVirtual() {
+		return
+	}
+	if _, found := a.possible[v]; found {
+		return
+	}
+	a.possible[v] = a.possibleregisters(v)
+}
+
+// Allocate allocates physical registers.
+func (a *Allocator) Allocate() (reg.Allocation, error) {
+	for {
+		if err := a.update(); err != nil {
+			return nil, err
+		}
+
+		if a.remaining() == 0 {
+			break
+		}
+
+		v := a.mostrestricted()
+		if err := a.alloc(v); err != nil {
+			return nil, err
+		}
+	}
+	return a.allocation, nil
+}
+
+// update possible allocations based on edges.
+func (a *Allocator) update() error {
+	var rem []*edge
+	for _, e := range a.edges {
+		x := a.allocation.LookupDefault(e.X)
+		y := a.allocation.LookupDefault(e.Y)
+		switch {
+		case x.IsVirtual() && y.IsVirtual():
+			rem = append(rem, e)
+			continue
+		case x.IsPhysical() && y.IsPhysical():
+			if x == y {
+				return errors.New("impossible register allocation")
+			}
+		case x.IsPhysical() && y.IsVirtual():
+			a.discardconflicting(y, x)
+		case x.IsVirtual() && y.IsPhysical():
+			a.discardconflicting(x, y)
+		default:
+			panic("unreachable")
+		}
+	}
+	a.edges = rem
+
+	return nil
+}
+
+// mostrestricted returns the virtual register with the least possibilities.
+func (a *Allocator) mostrestricted() reg.ID {
+	n := int(math.MaxInt32)
+	var v reg.ID
+	for w, p := range a.possible {
+		// On a tie, choose the smallest ID in numeric order. This avoids
+		// non-deterministic allocations due to map iteration order.
+		if len(p) < n || (len(p) == n && w < v) {
+			n = len(p)
+			v = w
+		}
+	}
+	return v
+}
+
+// discardconflicting removes registers from vs possible list that conflict with p.
+func (a *Allocator) discardconflicting(v, p reg.ID) {
+	a.possible[v] = filterregisters(a.possible[v], func(r reg.ID) bool {
+		return r != p
+	})
+}
+
+// alloc attempts to allocate a register to v.
+func (a *Allocator) alloc(v reg.ID) error {
+	ps := a.possible[v]
+	if len(ps) == 0 {
+		return errors.New("failed to allocate registers")
+	}
+	p := ps[0]
+	a.allocation[v] = p
+	delete(a.possible, v)
+	return nil
+}
+
+// remaining returns the number of unallocated registers.
+func (a *Allocator) remaining() int {
+	return len(a.possible)
+}
+
+// possibleregisters returns all allocate-able registers for the given virtual.
+func (a *Allocator) possibleregisters(v reg.ID) []reg.ID {
+	return filterregisters(a.registers, func(r reg.ID) bool {
+		return v.Kind() == r.Kind()
+	})
+}
+
+func filterregisters(in []reg.ID, predicate func(reg.ID) bool) []reg.ID {
+	var rs []reg.ID
+	for _, r := range in {
+		if predicate(r) {
+			rs = append(rs, r)
+		}
+	}
+	return rs
+}