examples/complex: and bugfixes

pass/alloc.go

@@ -46,11 +46,23 @@ func (a *Allocator) AddInterferenceSet(r reg.Register, s reg.Set) {
 }
 
 func (a *Allocator) AddInterference(x, y reg.Register) {
-	a.add(x)
-	a.add(y)
+	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(r reg.Register) {
+	v, ok := r.(reg.Virtual)
+	if !ok {
+		return
+	}
+	if _, found := a.possible[v]; found {
+		return
+	}
+	a.possible[v] = a.registersofsize(v.Bytes())
+}
+
 func (a *Allocator) Allocate() (reg.Allocation, error) {
 	for a.remaining() > 0 {
 		if err := a.update(); err != nil {
@@ -65,15 +77,6 @@ func (a *Allocator) Allocate() (reg.Allocation, error) {
 	return a.allocation, nil
 }
 
-// add adds a register.
-func (a *Allocator) add(r reg.Register) {
-	v, ok := r.(reg.Virtual)
-	if !ok {
-		return
-	}
-	a.possible[v] = a.registersofsize(v.Bytes())
-}
-
 // update possible allocations based on edges.
 func (a *Allocator) update() error {
 	var rem []*edge

pass/reg.go

@@ -57,11 +57,11 @@ func Liveness(fn *avo.Function) error {
 }
 
 func AllocateRegisters(fn *avo.Function) error {
-	// Build one allocator per register kind and record register interferences.
+	// Populate allocators (one per kind).
 	as := map[reg.Kind]*Allocator{}
 	for _, i := range fn.Instructions() {
-		for _, d := range i.OutputRegisters() {
-			k := d.Kind()
+		for _, r := range i.Registers() {
+			k := r.Kind()
 			if _, found := as[k]; !found {
 				a, err := NewAllocatorForKind(k)
 				if err != nil {
@@ -69,7 +69,14 @@ func AllocateRegisters(fn *avo.Function) error {
 				}
 				as[k] = a
 			}
+			as[k].Add(r)
+		}
+	}
 
+	// Record register interferences.
+	for _, i := range fn.Instructions() {
+		for _, d := range i.OutputRegisters() {
+			k := d.Kind()
 			out := i.LiveOut.OfKind(k)
 			out.Discard(d)
 			as[k].AddInterferenceSet(d, out)