ast: {Input,Output}Registers()

operand/types_test.go

@@ -1,6 +1,7 @@
 package operand
 
 import (
+	"reflect"
 	"testing"
 
 	"github.com/mmcloughlin/avo/reg"
@@ -26,3 +27,23 @@ func TestMemAsm(t *testing.T) {
 		}
 	}
 }
+
+func TestRegisters(t *testing.T) {
+	cases := []struct {
+		Op     Op
+		Expect []reg.Register
+	}{
+		{reg.R11, []reg.Register{reg.R11}},
+		{Mem{Base: reg.EAX}, []reg.Register{reg.EAX}},
+		{Mem{Base: reg.RBX, Index: reg.R10}, []reg.Register{reg.RBX, reg.R10}},
+		{Imm(42), nil},
+		{Rel(42), nil},
+		{LabelRef("idk"), nil},
+	}
+	for _, c := range cases {
+		got := Registers(c.Op)
+		if !reflect.DeepEqual(got, c.Expect) {
+			t.Errorf("Registers(%#v) = %#v expected %#v", c.Op, got, c.Expect)
+		}
+	}
+}