Performance improvements, tests, and benchmarks

2 files changed, 334 insertions, 70 deletions

diff --git a/capset.go b/capset.go
index 9a74038..140c488 100644
--- a/capset.go
+++ b/capset.go
@@ -7,26 +7,24 @@ package main
 import "log"
 import "os"
 import "fmt"
+import "bufio"
+import "errors"
 
-func is_set(cards []string) bool {
-	if len(cards) != 3 {
-		log.Fatal("Not the right number of cards.")
+var deck [81]string
+var card2IntMap map[string]int
+var useMap = true
+
+func init() {
+	for i := range deck {
+		deck[i] = computeInt2Card(80 - i)
 	}
-	for i := 0; i < 4; i++ {
-		if cards[0][i] == cards[1][i] && cards[0][i] == cards[2][i] {
-			continue
-		}
-		if cards[0][i] != cards[1][i] &&
-			cards[0][i] != cards[2][i] &&
-			cards[1][i] != cards[2][i] {
-			continue
-		}
-		return false
+	card2IntMap = make(map[string]int)
+	for i, v := range deck {
+		card2IntMap[v] = i
 	}
-	return true
 }
 
-func int2card(num int) string {
+func computeInt2Card(num int) string {
 	if num < 0 {
 		log.Fatal("Number must be nonnegative.")
 	}
@@ -41,91 +39,266 @@ func int2card(num int) string {
 	return string(strlist[:])
 }
 
-func find_index(deck []string, item string) int {
-	for i, v := range deck {
-		if v == item {
-			return i
+func card2int(item string) int {
+	if useMap {
+		i, ok := card2IntMap[item]
+		if !ok {
+			return -1
+		}
+		return i
+	} else {
+		for i, v := range deck {
+			if v == item {
+				return i
+			}
+		}
+		return -1
+	}
+}
+
+type capset struct {
+	tableau []string
+	maxim   int
+	hasSet  bool
+}
+
+func (cs *capset) Scan(state fmt.ScanState, verb rune) error {
+	if verb != 'v' {
+		return errors.New("unexpected verb")
+	}
+	r, _, err := state.ReadRune()
+	if err != nil {
+		return err
+	}
+	if r != '[' {
+		return errors.New("didn't start with [")
+	}
+	cs.tableau = cs.tableau[:0]
+	for end := false; !end; {
+		token, err := state.Token(true, nil)
+		if err != nil {
+			return err
+		}
+		if len(token) == 0 {
+			return errors.New("Premature end")
+		}
+		if token[len(token)-1] == ']' {
+			end = true
+			token = token[:len(token)-1]
 		}
+		if len(token) > 0 {
+			card := string(token)
+			index := card2int(card)
+			if index == -1 {
+				return errors.New(fmt.Sprintf("Invalid card: %v", card))
+			}
+			// Always use same strings in hope underlying strcmp short circuits
+			card = deck[index]
+			cs.tableau = append(cs.tableau, card)
+		}
+	}
+	cs.updateTableauHasSetCached()
+	// Loop until it isn't a set, to make sure we didn't load a tableau with lots
+	// of sets.
+	for cs.hasSet {
+		cs.incrementTableau(true)
 	}
-	panic("Somehow couldn't find index")
+	return nil
 }
 
-func list_pop(arr *[]string) string {
-	v := (*arr)[len(*arr)-1]
-	*arr = (*arr)[:len(*arr)-1]
+func (cs *capset) popTableau() string {
+	v := cs.tableau[len(cs.tableau)-1]
+	cs.tableau = cs.tableau[:len(cs.tableau)-1]
 	return v
 }
 
-func increment_tableau(tableau []string, deck []string, pop bool) []string {
-	index := find_index(deck, tableau[len(tableau)-1]) + 1
+func (cs *capset) incrementTableau(noAssume bool) bool {
+	pop := cs.hasSet
+	index := card2int(cs.tableau[len(cs.tableau)-1]) + 1
 	if index != len(deck) {
 		if pop {
-			list_pop(&tableau)
+			cs.popTableau()
 		}
-		tableau = append(tableau, deck[index])
-		return tableau
+		cs.tableau = append(cs.tableau, deck[index])
+		cs.updateTableauHasSet(noAssume)
+		return true
 	}
-	list_pop(&tableau)
-	index = find_index(deck, list_pop(&tableau)) + 1
-	if len(tableau) == 1 {
+	cs.popTableau()
+	index = card2int(cs.popTableau()) + 1
+	if len(cs.tableau) == 1 {
 		fmt.Println("incrementing second card to index", index)
 		if index+18 > len(deck) {
-			os.Exit(0)
+			return false
 		}
 	}
-	tableau = append(tableau, deck[index])
-	return tableau
+	cs.tableau = append(cs.tableau, deck[index])
+	cs.updateTableauHasSet(noAssume)
+	return true
 }
 
-func check_tableau(tableau []string) bool {
-	if len(tableau) == 2 {
-		return true
+func isSet(cards [3]string) bool {
+	for i := 0; i < 4; i++ {
+		if cards[0][i] == cards[1][i] && cards[0][i] == cards[2][i] {
+			continue
+		}
+		if cards[0][i] != cards[1][i] &&
+			cards[0][i] != cards[2][i] &&
+			cards[1][i] != cards[2][i] {
+			continue
+		}
+		return false
+	}
+	return true
+}
+
+func (cs *capset) updateTableauHasSet(noAssume bool) {
+	if noAssume {
+		cs.updateTableauHasSetCached()
+	} else {
+		cs.updateTableauHasSetOnlyNew()
+	}
+}
+
+// TODO(ejona): remove
+func (cs *capset) updateTableauHasSetSimple() {
+	if len(cs.tableau) <= 2 {
+		cs.hasSet = false
+		return
 	}
 	var comb [3]string
-	for i := 0; i < len(tableau); i++ {
-		comb[0] = tableau[i]
-		for j := i + 1; j < len(tableau); j++ {
-			comb[1] = tableau[j]
-			for k := j + 1; k < len(tableau); k++ {
-				comb[2] = tableau[k]
-				if is_set(comb[:]) {
-					return false
+	for i := 0; i < len(cs.tableau); i++ {
+		comb[0] = cs.tableau[i]
+		for j := i + 1; j < len(cs.tableau); j++ {
+			comb[1] = cs.tableau[j]
+			for k := j + 1; k < len(cs.tableau); k++ {
+				comb[2] = cs.tableau[k]
+				if isSet(comb) {
+					cs.hasSet = true
+					return
 				}
 			}
 		}
 	}
-	return true
+	cs.hasSet = false
+}
+
+func (cs *capset) updateTableauHasSetCached() {
+	var same [4]bool
+	for i := 0; i < len(cs.tableau); i++ {
+		for j := i + 1; j < len(cs.tableau); j++ {
+			for x := 0; x < 4; x++ {
+				same[x] = cs.tableau[i][x] == cs.tableau[j][x]
+			}
+		kloop:
+			for k := j + 1; k < len(cs.tableau); k++ {
+				for x := 0; x < 4; x++ {
+					if same[x] {
+						if cs.tableau[i][x] != cs.tableau[k][x] {
+							continue kloop
+						}
+					} else {
+						if cs.tableau[i][x] == cs.tableau[k][x] ||
+							cs.tableau[j][x] == cs.tableau[k][x] {
+							continue kloop
+						}
+					}
+				}
+				cs.hasSet = true
+				return
+			}
+		}
+	}
+	cs.hasSet = false
+}
+
+func (cs *capset) updateTableauHasSetOnlyNew() {
+	if len(cs.tableau) <= 2 {
+		cs.hasSet = false
+		return
+	}
+	// We try sets such that only the rightmost item could possibly be in a set.
+	i := len(cs.tableau) - 1
+	var same [4]bool
+	for j := 0; j < i; j++ {
+		for x := 0; x < 4; x++ {
+			same[x] = cs.tableau[i][x] == cs.tableau[j][x]
+		}
+	kloop:
+		for k := j + 1; k < i; k++ {
+			for x := 0; x < 4; x++ {
+				if same[x] {
+					if cs.tableau[i][x] != cs.tableau[k][x] {
+						continue kloop
+					}
+				} else {
+					if cs.tableau[i][x] == cs.tableau[k][x] ||
+						cs.tableau[j][x] == cs.tableau[k][x] {
+						continue kloop
+					}
+				}
+			}
+			cs.hasSet = true
+			return
+		}
+	}
+	cs.hasSet = false
+}
+
+func (cs *capset) findNextCapset() bool {
+	for cs.incrementTableau(false) {
+		if cs.hasSet {
+			continue
+		}
+		if len(cs.tableau) < cs.maxim {
+			continue
+		}
+		cs.maxim = len(cs.tableau)
+		return true
+	}
+	return false
+}
+
+func loadFromSave() *capset {
+	savefile, err := os.Open("./capset.out")
+	if err != nil {
+		return nil
+	}
+	defer savefile.Close()
+	scanner := bufio.NewScanner(savefile)
+	for scanner.Scan() {
+	}
+	var cs capset
+	_, err = fmt.Sscanln(scanner.Text(), &cs)
+	if err != nil {
+		return nil
+	}
+	return &cs
 }
 
 func main() {
-	outfile, err := os.OpenFile("./capset.out", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
+	cs := loadFromSave()
+	if cs == nil {
+		cs = &capset{tableau: []string{"2222"}}
+	}
+	cs.maxim = 20
+	fmt.Println("Starting processing with", cs.tableau)
+
+	outfile, err := os.OpenFile("./capset.out",
+		os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0644)
 	if err != nil {
 		log.Fatal(err)
 	}
 	defer outfile.Close()
-	deck := [81]string{}
-	for i := range deck {
-		deck[i] = int2card(80-i)
-	}
-	tableau := []string{deck[0], deck[1]}
-	maxim := 20
-	for {
-		if check_tableau(tableau) {
-			if len(tableau) >= maxim {
-				fmt.Println(tableau)
-				_, err = fmt.Fprintln(outfile, tableau)
-				if err != nil {
-					log.Println(err)
-				}
-				err = outfile.Sync()
-				if err != nil {
-					log.Println(err)
-				}
-				maxim = len(tableau)
-			}
-			tableau = increment_tableau(tableau, deck[:], false)
-		} else {
-			tableau = increment_tableau(tableau, deck[:], true)
+
+	for cs.findNextCapset() {
+		fmt.Println(cs.tableau)
+		_, err = fmt.Fprintln(outfile, cs.tableau)
+		if err != nil {
+			log.Println(err)
+		}
+		err = outfile.Sync()
+		if err != nil {
+			log.Println(err)
 		}
 	}
 }
diff --git a/capset_test.go b/capset_test.go
new file mode 100644
index 0000000..5382aa8
--- /dev/null
+++ b/capset_test.go
@@ -0,0 +1,91 @@
+package main
+
+import "testing"
+import "reflect"
+import "fmt"
+
+var capsettests = []struct {
+	maxim         int
+	start, finish []string
+}{
+	{20, []string{"2222", "2221", "2212", "2211", "2122", "2121", "2112", "2111",
+		"1222", "1221", "1210"},
+		[]string{"2222", "2221", "2212", "2211", "2122", "2121", "2112", "2111",
+			"1222", "1221", "1210", "1200", "1120", "1020", "0210", "0120", "0112",
+			"0111", "0100", "0010"}},
+}
+
+func TestFindNextCapset(t *testing.T) {
+	for _, tt := range capsettests {
+		cs := capset{tableau: tt.start, maxim: tt.maxim}
+		cs.findNextCapset()
+		expected := capset{tableau: tt.finish, maxim: len(tt.finish)}
+		if !reflect.DeepEqual(cs, expected) {
+			t.Fatalf("Got \n%v but expected \n%v", cs, expected)
+		}
+	}
+}
+
+var scantests = []struct {
+	in  string
+	out []string
+	err bool
+}{
+	{"", nil, true},
+	{"[]", []string{}, false},
+	{"[0000]", []string{"0000"}, false},
+	{"[0000 2222]", []string{"0000", "2222"}, false},
+	{"[2]", nil, true},
+	{"not it", nil, true},
+	{"[0000", nil, true},
+	{"[  0000   2222  ]", []string{"0000", "2222"}, false},
+}
+
+func TestScan(t *testing.T) {
+	for _, tt := range scantests {
+		cs := capset{tableau: []string{}}
+		if testing.Verbose() {
+			fmt.Println("Trying input", tt.in)
+		}
+		_, err := fmt.Sscanln(tt.in, &cs)
+		if tt.err {
+			if err == nil {
+				t.Errorf("Sscanln(%v) => %v, want an error", tt.in, cs.tableau)
+			}
+		} else {
+			if err != nil {
+				t.Errorf("Sscanln(%v). Unexpected error %v", tt.in, err)
+			} else {
+				if !reflect.DeepEqual(cs.tableau, tt.out) {
+					t.Errorf("Sscanln(%v) => %v, want %v", tt.in, cs.tableau, tt.out)
+				}
+			}
+		}
+	}
+}
+
+func testNoAllocs(t *testing.T) {
+	cs := capset{tableau: []string{"2222"}}
+	result := testing.AllocsPerRun(10, func() { cs.incrementTableau(false) })
+	if result > 0 {
+		t.Errorf("Too many allocs: %v", result)
+	}
+}
+
+func BenchmarkFindNextCapset(b *testing.B) {
+	cs := capset{tableau: []string{"2222"}}
+	useMap = true
+	for i := 0; i < b.N; i++ {
+		cs.incrementTableau(false)
+	}
+}
+
+func BenchmarkFindNextCapsetNoMap(b *testing.B) {
+	cs := capset{tableau: []string{"2222"}}
+	useMap = false
+	for i := 0; i < b.N; i++ {
+		cs.incrementTableau(false)
+	}
+	useMap = true
+}
+