From 4dcfed6e1655529dfb5dee45bf266f7afbc1af66 Mon Sep 17 00:00:00 2001
From: Michael McLoughlin <mmcloughlin@gmail.com>
Date: Sun, 25 Nov 2018 18:25:51 -0800
Subject: [PATCH] add instruction arities function

---
 internal/inst/table_test.go | 21 ++++++++++++++++++++-
 internal/inst/types.go      | 15 +++++++++++++++
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/internal/inst/table_test.go b/internal/inst/table_test.go
index 9529ce4..f3d5373 100644
--- a/internal/inst/table_test.go
+++ b/internal/inst/table_test.go
@@ -2,6 +2,7 @@ package inst_test
 
 import (
 	"io/ioutil"
+	"reflect"
 	"strings"
 	"testing"
 
@@ -40,7 +41,6 @@ func TestInstructionProperties(t *testing.T) {
 			t.Errorf("instruction %s has no forms", i.Opcode)
 		}
 	}
-
 }
 
 func TestAssembles(t *testing.T) {
@@ -61,6 +61,25 @@ func TestLookup(t *testing.T) {
 	}
 }
 
+func TestInstructionArities(t *testing.T) {
+	cases := map[string][]int{
+		"AESDEC":    {2},
+		"EXTRACTPS": {3},
+		"SHRQ":      {2, 3},
+		"VMOVHPD":   {2, 3},
+	}
+	for opcode, expect := range cases {
+		i, ok := inst.Lookup(opcode)
+		if !ok {
+			t.Fatalf("could not find %s", opcode)
+		}
+		got := i.Arities()
+		if !reflect.DeepEqual(got, expect) {
+			t.Errorf("arity of %s is %v expected %v", opcode, got, expect)
+		}
+	}
+}
+
 func TestStdLibOpcodes(t *testing.T) {
 	b, err := ioutil.ReadFile("testdata/stdlibopcodes.txt")
 	if err != nil {
diff --git a/internal/inst/types.go b/internal/inst/types.go
index 3c49aeb..1126270 100644
--- a/internal/inst/types.go
+++ b/internal/inst/types.go
@@ -1,5 +1,7 @@
 package inst
 
+import "sort"
+
 type Instruction struct {
 	Opcode  string
 	AliasOf string
@@ -7,6 +9,19 @@ type Instruction struct {
 	Forms   []Form
 }
 
+func (i Instruction) Arities() []int {
+	s := map[int]bool{}
+	for _, f := range i.Forms {
+		s[len(f.Operands)] = true
+	}
+	a := make([]int, 0, len(s))
+	for n := range s {
+		a = append(a, n)
+	}
+	sort.Ints(a)
+	return a
+}
+
 type Form struct {
 	ISA              []string
 	Operands         []Operand