ast: Optimize set/object lookups and Number comparison.

Lookups for sets and objects use their own type specific equality
checks instead of relying on the general purpose Compare. This
improves the performance about 25%. The duplication of code is
unavoidable to avoid the allocations, unfortunately.

Number comparison is now allocation free for integers. This because
Int64() of json.Number uses allocation free strconv.ParseInt
internally. This shows up in the additional 60% improvement of set
membership benchmarks as they use integer keys.

Most notable improvements in benchmarks:

benchmark                                        old ns/op      new ns/op      delta
BenchmarkObjectLookup/5-8                        37.3           31.9           -14.48%
BenchmarkObjectLookup/50-8                       44.7           39.7           -11.19%
BenchmarkObjectLookup/500-8                      46.6           41.7           -10.52%
BenchmarkObjectLookup/5000-8                     43.6           38.4           -11.93%
BenchmarkSetIntersection/5-8                     2309           928            -59.81%
BenchmarkSetIntersection/50-8                    24562          9121           -62.87%
BenchmarkSetIntersection/500-8                   260022         93373          -64.09%
BenchmarkSetIntersection/5000-8                  2796645        1035045        -62.99%
BenchmarkSetIntersectionDifferentSize/4-8        1774           781            -55.98%
BenchmarkSetIntersectionDifferentSize/50-8       1847           875            -52.63%
BenchmarkSetIntersectionDifferentSize/500-8      1849           864            -53.27%
BenchmarkSetIntersectionDifferentSize/5000-8     1884           890            -52.76%
BenchmarkSetMembership/5-8                       339            55.1           -83.75%
BenchmarkSetMembership/50-8                      383            65.3           -82.95%
BenchmarkSetMembership/500-8                     403            70.5           -82.51%
BenchmarkSetMembership/5000-8                    418            75.1           -82.03%
BenchmarkConcurrency1-8                          219100992      218699478      -0.18%
BenchmarkConcurrency2-8                          121519274      119096088      -1.99%
BenchmarkConcurrency4-8                          88765327       78934589       -11.07%
BenchmarkConcurrency8-8                          93624724       82432804       -11.95%
BenchmarkConcurrency4Readers1Writer-8            90931759       80533879       -11.43%
BenchmarkConcurrency8Writers-8                   237043373      233709595      -1.41%

Signed-off-by: Teemu Koponen <koponen@styra.com>

ast/compare.go

@@ -5,6 +5,7 @@
 package ast
 
 import (
+	"encoding/json"
 	"fmt"
 	"math/big"
 )
@@ -84,6 +85,18 @@ func Compare(a, b interface{}) int {
 		}
 		return 1
 	case Number:
+		if ai, err := json.Number(a).Int64(); err == nil {
+			if bi, err := json.Number(b.(Number)).Int64(); err == nil {
+				if ai == bi {
+					return 0
+				}
+				if ai < bi {
+					return -1
+				}
+				return 1
+			}
+		}
+
 		bigA, ok := new(big.Float).SetString(string(a))
 		if !ok {
 			panic("illegal value")

ast/compare_test.go

@@ -27,6 +27,9 @@ func TestCompare(t *testing.T) {
 		// Numbers
 		{"0", "1", -1},
 		{"1", "0", 1},
+		{"0", "0", 0},
+		{"0", "1.5", -1},
+		{"1.5", "0", 1},
 
 		// Object comparisons are consistent
 		{`{1: 2, 3: 4}`, `{4: 3, 1: 2}`, -1},

ast/term.go

@@ -8,6 +8,7 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
+	"math/big"
 	"net/url"
 	"regexp"
 	"sort"
@@ -377,6 +378,8 @@ func (term *Term) Equal(other *Term) bool {
 		return v.Equal(other.Value)
 	case Boolean:
 		return v.Equal(other.Value)
+	case Number:
+		return v.Equal(other.Value)
 	case String:
 		return v.Equal(other.Value)
 	case Var:
@@ -1426,8 +1429,48 @@ func (s *set) Slice() []*Term {
 
 func (s *set) insert(x *Term) {
 	hash := x.Hash()
+	var equal func(v Value) bool
+
+	switch x := x.Value.(type) {
+	case Null, Boolean, String, Var:
+		equal = func(y Value) bool { return x == y }
+	case Number:
+		if xi, err := json.Number(x).Int64(); err == nil {
+			equal = func(y Value) bool {
+				if y, ok := y.(Number); ok {
+					if yi, err := json.Number(y).Int64(); err == nil {
+						return xi == yi
+					}
+				}
+
+				return false
+			}
+			break
+		}
+
+		a, ok := new(big.Float).SetString(string(x))
+		if !ok {
+			panic("illegal value")
+		}
+
+		equal = func(b Value) bool {
+			if b, ok := b.(Number); ok {
+				b, ok := new(big.Float).SetString(string(b))
+				if !ok {
+					panic("illegal value")
+				}
+
+				return a.Cmp(b) == 0
+			}
+
+			return false
+		}
+	default:
+		equal = func(y Value) bool { return Compare(x, y) == 0 }
+	}
+
 	for curr, ok := s.elems[hash]; ok; {
-		if Compare(curr, x) == 0 {
+		if equal(curr.Value) {
 			return
 		}
 
@@ -1441,8 +1484,48 @@ func (s *set) insert(x *Term) {
 
 func (s *set) get(x *Term) *Term {
 	hash := x.Hash()
+	var equal func(v Value) bool
+
+	switch x := x.Value.(type) {
+	case Null, Boolean, String, Var:
+		equal = func(y Value) bool { return x == y }
+	case Number:
+		if xi, err := json.Number(x).Int64(); err == nil {
+			equal = func(y Value) bool {
+				if y, ok := y.(Number); ok {
+					if yi, err := json.Number(y).Int64(); err == nil {
+						return xi == yi
+					}
+				}
+
+				return false
+			}
+			break
+		}
+
+		a, ok := new(big.Float).SetString(string(x))
+		if !ok {
+			panic("illegal value")
+		}
+
+		equal = func(b Value) bool {
+			if b, ok := b.(Number); ok {
+				b, ok := new(big.Float).SetString(string(b))
+				if !ok {
+					panic("illegal value")
+				}
+
+				return a.Cmp(b) == 0
+			}
+
+			return false
+		}
+	default:
+		equal = func(y Value) bool { return Compare(x, y) == 0 }
+	}
+
 	for curr, ok := s.elems[hash]; ok; {
-		if Compare(curr, x) == 0 {
+		if equal(curr.Value) {
 			return curr
 		}
 
@@ -1747,8 +1830,49 @@ func (obj object) String() string {
 
 func (obj *object) get(k *Term) *objectElem {
 	hash := k.Hash()
+
+	var equal func(v Value) bool
+
+	switch x := k.Value.(type) {
+	case Null, Boolean, String, Var:
+		equal = func(y Value) bool { return x == y }
+	case Number:
+		if xi, err := json.Number(x).Int64(); err == nil {
+			equal = func(y Value) bool {
+				if y, ok := y.(Number); ok {
+					if yi, err := json.Number(y).Int64(); err == nil {
+						return xi == yi
+					}
+				}
+
+				return false
+			}
+			break
+		}
+
+		a, ok := new(big.Float).SetString(string(x))
+		if !ok {
+			panic("illegal value")
+		}
+
+		equal = func(b Value) bool {
+			if b, ok := b.(Number); ok {
+				b, ok := new(big.Float).SetString(string(b))
+				if !ok {
+					panic("illegal value")
+				}
+
+				return a.Cmp(b) == 0
+			}
+
+			return false
+		}
+	default:
+		equal = func(y Value) bool { return Compare(x, y) == 0 }
+	}
+
 	for curr := obj.elems[hash]; curr != nil; curr = curr.next {
-		if Compare(curr.key, k) == 0 {
+		if equal(curr.key.Value) {
 			return curr
 		}
 	}
@@ -1758,8 +1882,48 @@ func (obj *object) get(k *Term) *objectElem {
 func (obj *object) insert(k, v *Term) {
 	hash := k.Hash()
 	head := obj.elems[hash]
+	var equal func(v Value) bool
+
+	switch x := k.Value.(type) {
+	case Null, Boolean, String, Var:
+		equal = func(y Value) bool { return x == y }
+	case Number:
+		if xi, err := json.Number(x).Int64(); err == nil {
+			equal = func(y Value) bool {
+				if y, ok := y.(Number); ok {
+					if yi, err := json.Number(y).Int64(); err == nil {
+						return xi == yi
+					}
+				}
+
+				return false
+			}
+			break
+		}
+
+		a, ok := new(big.Float).SetString(string(x))
+		if !ok {
+			panic("illegal value")
+		}
+
+		equal = func(b Value) bool {
+			if b, ok := b.(Number); ok {
+				b, ok := new(big.Float).SetString(string(b))
+				if !ok {
+					panic("illegal value")
+				}
+
+				return a.Cmp(b) == 0
+			}
+
+			return false
+		}
+	default:
+		equal = func(y Value) bool { return Compare(x, y) == 0 }
+	}
+
 	for curr := head; curr != nil; curr = curr.next {
-		if Compare(curr.key, k) == 0 {
+		if equal(curr.key.Value) {
 			curr.value = v
 			return
 		}

ast/term_test.go

@@ -84,6 +84,40 @@ func TestInterfaceToValue(t *testing.T) {
 
 }
 
+func TestObjectInsertGetLen(t *testing.T) {
+	tests := []struct {
+		insert   [][2]string
+		expected map[string]string
+	}{
+		{[][2]string{{`null`, `value1`}, {`null`, `value2`}}, map[string]string{`null`: `value2`}},
+		{[][2]string{{`false`, `value`}, {`true`, `value1`}, {`true`, `value2`}}, map[string]string{`false`: `value`, `true`: `value2`}},
+		{[][2]string{{`0`, `value`}, {`1`, `value1`}, {`1`, `value2`}, {`1.5`, `value`}}, map[string]string{`0`: `value`, `1`: `value2`, `1.5`: `value`}},
+		{[][2]string{{`"string"`, `value1`}, {`"string"`, `value2`}}, map[string]string{`"string"`: `value2`}},
+		{[][2]string{{`["other"]`, `value1`}, {`["other"]`, `value2`}}, map[string]string{`["other"]`: `value2`}},
+	}
+
+	for _, tc := range tests {
+		o := NewObject()
+		for _, kv := range tc.insert {
+			o.Insert(MustParseTerm(kv[0]), MustParseTerm(kv[1]))
+
+			if v := o.Get(MustParseTerm(kv[0])); v == nil || !MustParseTerm(kv[1]).Equal(v) {
+				t.Errorf("Expected the object to contain %v", v)
+			}
+		}
+
+		if o.Len() != len(tc.expected) {
+			t.Errorf("Expected the object to have %v entries", len(tc.expected))
+		}
+
+		for k, v := range tc.expected {
+			if x := o.Get(MustParseTerm(k)); x == nil || !MustParseTerm(v).Equal(x) {
+				t.Errorf("Expected the object to contain %v", k)
+			}
+		}
+	}
+}
+
 func TestObjectSetOperations(t *testing.T) {
 
 	a := MustParseTerm(`{"a": "b", "c": "d"}`).Value.(Object)
@@ -522,6 +556,41 @@ func TestSetMap(t *testing.T) {
 
 }
 
+func TestSetAddContainsLen(t *testing.T) {
+	tests := []struct {
+		add      []string
+		expected []string
+	}{
+		{[]string{`null`, `null`}, []string{`null`}},
+		{[]string{`true`, `true`, `false`}, []string{`true`, `false`}},
+		{[]string{`0`, `1`, `1`, `1.5`}, []string{`0`, `1`, `1.5`}},
+		{[]string{`"string"`, `"string"`}, []string{`"string"`}},
+		{[]string{`["other"]`, `["other"]`}, []string{`["other"]`}},
+	}
+
+	for _, tc := range tests {
+		s := NewSet()
+		for _, v := range tc.add {
+			x := MustParseTerm(v)
+			s.Add(x)
+
+			if !s.Contains(x) {
+				t.Errorf("Expected the set to contain %v", v)
+			}
+		}
+
+		if s.Len() != len(tc.expected) {
+			t.Errorf("Expected the set to have %v entries", len(tc.expected))
+		}
+
+		for _, v := range tc.expected {
+			if !s.Contains(MustParseTerm(v)) {
+				t.Errorf("Expected the set to contain %v", v)
+			}
+		}
+	}
+}
+
 func TestSetOperations(t *testing.T) {
 
 	tests := []struct {