pass: support priority in register allocator (#180)

Updates #156

pass/alloc.go

@@ -17,6 +17,7 @@ type edge struct {
 // Allocator is a graph-coloring register allocator.
 type Allocator struct {
 	registers  []reg.ID
+	priority   map[reg.ID]int
 	allocation reg.Allocation
 	edges      []*edge
 	possible   map[reg.ID][]reg.ID
@@ -42,13 +43,16 @@ func NewAllocator(rs []reg.Physical) (*Allocator, error) {
 	for id := range idset {
 		ids = append(ids, id)
 	}
-	sort.Slice(ids, func(i, j int) bool { return ids[i] < ids[j] })
 
-	return &Allocator{
+	a := &Allocator{
 		registers:  ids,
+		priority:   map[reg.ID]int{},
 		allocation: reg.NewEmptyAllocation(),
 		possible:   map[reg.ID][]reg.ID{},
-	}, nil
+	}
+	a.sortregisters()
+
+	return a, nil
 }
 
 // NewAllocatorForKind builds an allocator for the given kind of registers.
@@ -60,6 +64,25 @@ func NewAllocatorForKind(k reg.Kind) (*Allocator, error) {
 	return NewAllocator(f.Registers())
 }
 
+// SetPriority sets the priority of the given regiser to p. Higher priority
+// registers are preferred in allocations. By default all registers have 0
+// priority. Priority will only apply to subsequent Add() calls, therefore
+// typically all SetPriority calls should happen at allocator initialization.
+func (a *Allocator) SetPriority(id reg.ID, p int) {
+	a.priority[id] = p
+	a.sortregisters()
+}
+
+// sortregisters sorts the list of available registers: higher priority first,
+// falling back to sorting by ID.
+func (a *Allocator) sortregisters() {
+	sort.Slice(a.registers, func(i, j int) bool {
+		ri, rj := a.registers[i], a.registers[j]
+		pi, pj := a.priority[ri], a.priority[rj]
+		return (pi > pj) || (pi == pj && ri < rj)
+	})
+}
+
 // 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 {