jsonschema: using hashing for uniqueItems

Determine if a slice contains unique items by hashing them.
This should take linear time on average, better than the previous
quadratic algorithm.

Change-Id: I8dc95bb6f29d802bdbc64bc7f7e698c71eae5ce7
Reviewed-on: https://go-review.googlesource.com/c/tools/+/669455
Reviewed-by: Alan Donovan <adonovan@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>

Author: jba

internal/mcp/internal/jsonschema/util.go

@@ -6,11 +6,15 @@ package jsonschema
 
 import (
 	"bytes"
+	"cmp"
+	"encoding/binary"
 	"encoding/json"
 	"fmt"
+	"hash/maphash"
 	"math"
 	"math/big"
 	"reflect"
+	"slices"
 )
 
 // Equal reports whether two Go values representing JSON values are equal according
@@ -126,14 +130,93 @@ func equalValue(x, y reflect.Value) bool {
 		return x.String() == y.String()
 	case reflect.Bool:
 		return x.Bool() == y.Bool()
-	case reflect.Complex64, reflect.Complex128:
-		return x.Complex() == y.Complex()
 	// Ints, uints and floats handled in jsonNumber, at top of function.
 	default:
 		panic(fmt.Sprintf("unsupported kind: %s", x.Kind()))
 	}
 }
 
+// hashValue adds v to the data hashed by h. v must not have cycles.
+// hashValue panics if the value contains functions or channels, or maps whose
+// key type is not string.
+// It ignores unexported fields of structs.
+// Calls to hashValue with the equal values (in the sense
+// of [Equal]) result in the same sequence of values written to the hash.
+func hashValue(h *maphash.Hash, v reflect.Value) {
+	// TODO: replace writes of basic types with WriteComparable in 1.24.
+
+	writeUint := func(u uint64) {
+		var buf [8]byte
+		binary.BigEndian.PutUint64(buf[:], u)
+		h.Write(buf[:])
+	}
+
+	var write func(reflect.Value)
+	write = func(v reflect.Value) {
+		if r, ok := jsonNumber(v); ok {
+			// We want 1.0 and 1 to hash the same.
+			// big.Rats are always normalized, so they will be.
+			// We could do this more efficiently by handling the int and float cases
+			// separately, but that's premature.
+			writeUint(uint64(r.Sign() + 1))
+			h.Write(r.Num().Bytes())
+			h.Write(r.Denom().Bytes())
+			return
+		}
+		switch v.Kind() {
+		case reflect.Invalid:
+			h.WriteByte(0)
+		case reflect.String:
+			h.WriteString(v.String())
+		case reflect.Bool:
+			if v.Bool() {
+				h.WriteByte(1)
+			} else {
+				h.WriteByte(0)
+			}
+		case reflect.Complex64, reflect.Complex128:
+			c := v.Complex()
+			writeUint(math.Float64bits(real(c)))
+			writeUint(math.Float64bits(imag(c)))
+		case reflect.Array, reflect.Slice:
+			// Although we could treat []byte more efficiently,
+			// JSON values are unlikely to contain them.
+			writeUint(uint64(v.Len()))
+			for i := range v.Len() {
+				write(v.Index(i))
+			}
+		case reflect.Interface, reflect.Pointer:
+			write(v.Elem())
+		case reflect.Struct:
+			t := v.Type()
+			for i := range t.NumField() {
+				if sf := t.Field(i); sf.IsExported() {
+					write(v.FieldByIndex(sf.Index))
+				}
+			}
+		case reflect.Map:
+			if v.Type().Key().Kind() != reflect.String {
+				panic("map with non-string key")
+			}
+			// Sort the keys so the hash is deterministic.
+			keys := v.MapKeys()
+			// Write the length. That distinguishes between, say, two consecutive
+			// maps with disjoint keys from one map that has the items of both.
+			writeUint(uint64(len(keys)))
+			slices.SortFunc(keys, func(x, y reflect.Value) int { return cmp.Compare(x.String(), y.String()) })
+			for _, k := range keys {
+				write(k)
+				write(v.MapIndex(k))
+			}
+		// Ints, uints and floats handled in jsonNumber, at top of function.
+		default:
+			panic(fmt.Sprintf("unsupported kind: %s", v.Kind()))
+		}
+	}
+
+	write(v)
+}
+
 // jsonNumber converts a numeric value or a json.Number to a [big.Rat].
 // If v is not a number, it returns nil, false.
 func jsonNumber(v reflect.Value) (*big.Rat, bool) {

internal/mcp/internal/jsonschema/util_test.go

@@ -6,6 +6,7 @@ package jsonschema
 
 import (
 	"encoding/json"
+	"hash/maphash"
 	"reflect"
 	"testing"
 )
@@ -71,3 +72,56 @@ func TestJSONType(t *testing.T) {
 
 	}
 }
+
+func TestHash(t *testing.T) {
+	x := map[string]any{
+		"s": []any{1, "foo", nil, true},
+		"f": 2.5,
+		"m": map[string]any{
+			"n":      json.Number("123.456"),
+			"schema": &Schema{Type: "integer", UniqueItems: true},
+		},
+		"c": 1.2 + 3.4i,
+		"n": nil,
+	}
+
+	seed := maphash.MakeSeed()
+
+	hash := func(x any) uint64 {
+		var h maphash.Hash
+		h.SetSeed(seed)
+		hashValue(&h, reflect.ValueOf(x))
+		return h.Sum64()
+	}
+
+	want := hash(x)
+	// Run several times to verify consistency.
+	for range 10 {
+		if got := hash(x); got != want {
+			t.Errorf("hash values differ: %d vs. %d", got, want)
+		}
+	}
+
+	// Check mathematically equal values.
+	nums := []any{
+		5,
+		uint(5),
+		5.0,
+		json.Number("5"),
+		json.Number("5.00"),
+	}
+	for i, n := range nums {
+		if i == 0 {
+			want = hash(n)
+		} else if got := hash(n); got != want {
+			t.Errorf("hashes differ between %v (%[1]T) and %v (%[2]T)", nums[0], n)
+		}
+	}
+
+	// Check that a bare JSON `null` is OK.
+	var null any
+	if err := json.Unmarshal([]byte(`null`), &null); err != nil {
+		t.Fatal(err)
+	}
+	_ = hash(null)
+}

internal/mcp/internal/jsonschema/validate.go

@@ -6,6 +6,7 @@ package jsonschema
 
 import (
 	"fmt"
+	"hash/maphash"
 	"math"
 	"math/big"
 	"reflect"
@@ -297,16 +298,32 @@ func (st *state) validate(instance reflect.Value, schema *Schema, callerAnns *an
 			}
 		}
 		if schema.UniqueItems {
-			// Determine uniqueness with O(n²) comparisons.
-			// TODO: optimize via hashing.
-			for i := range instance.Len() {
-				for j := i + 1; j < instance.Len(); j++ {
-					if equalValue(instance.Index(i), instance.Index(j)) {
-						return fmt.Errorf("uniqueItems: array items %d and %d are equal", i, j)
+			if instance.Len() > 1 {
+				// Hash each item and compare the hashes.
+				// If two hashes differ, the items differ.
+				// If two hashes are the same, compare the collisions for equality.
+				// (The same logic as hash table lookup.)
+				// TODO(jba): Use container/hash.Map when it becomes available (https://go.dev/issue/69559),
+				hashes := map[uint64][]int{} // from hash to indices
+				seed := maphash.MakeSeed()
+				for i := range instance.Len() {
+					item := instance.Index(i)
+					var h maphash.Hash
+					h.SetSeed(seed)
+					hashValue(&h, item)
+					hv := h.Sum64()
+					if sames := hashes[hv]; len(sames) > 0 {
+						for _, j := range sames {
+							if equalValue(item, instance.Index(j)) {
+								return fmt.Errorf("uniqueItems: array items %d and %d are equal", i, j)
+							}
+						}
 					}
+					hashes[hv] = append(hashes[hv], i)
 				}
 			}
 		}
+
 		// https://json-schema.org/draft/2020-12/json-schema-core#section-11.2
 		if schema.UnevaluatedItems != nil && !anns.allItems {
 			// Apply this subschema to all items in the array that haven't been successfully validated.