ast: {Input,Output}Registers()

operand/checks.go

@@ -6,6 +6,22 @@ import (
 	"github.com/mmcloughlin/avo/reg"
 )
 
+// Pure type assertion checks:
+
+// IsRegister returns whether op has type reg.Register.
+func IsRegister(op Op) bool { _, ok := op.(reg.Register); return ok }
+
+// IsMem returns whether op has type Mem.
+func IsMem(op Op) bool { _, ok := op.(Mem); return ok }
+
+// IsImm returns whether op has type Imm.
+func IsImm(op Op) bool { _, ok := op.(Imm); return ok }
+
+// IsRel returns whether op has type Rel.
+func IsRel(op Op) bool { _, ok := op.(Rel); return ok }
+
+// Checks corresponding to specific operand types in the Intel Manual:
+
 // Is1 returns true if op is the immediate constant 1.
 func Is1(op Op) bool {
 	i, ok := op.(Imm)

operand/types.go

@@ -12,8 +12,8 @@ type Op interface {
 
 type Mem struct {
 	Disp  int
-	Base  reg.Physical
-	Index reg.Physical
+	Base  reg.Register
+	Index reg.Register
 	Scale uint8
 }
 
@@ -50,3 +50,23 @@ type LabelRef string
 func (l LabelRef) Asm() string {
 	return string(l)
 }
+
+// Registers returns the list of all operands involved in the given operand.
+func Registers(op Op) []reg.Register {
+	switch op := op.(type) {
+	case reg.Register:
+		return []reg.Register{op}
+	case Mem:
+		var r []reg.Register
+		if op.Base != nil {
+			r = append(r, op.Base)
+		}
+		if op.Index != nil {
+			r = append(r, op.Index)
+		}
+		return r
+	case Imm, Rel, LabelRef:
+		return nil
+	}
+	panic("unknown operand type")
+}

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)
+		}
+	}
+}