From e5b47c6c53e1cbdccf0803fc7d38c80689183f92 Mon Sep 17 00:00:00 2001
From: Torin Sandall <torinsandall@gmail.com>
Date: Mon, 21 Jun 2021 04:14:18 -0400
Subject: [PATCH] types: Use binary search for static property lookup (#3567)

The static properties are stored in sorted order so perform a binary
search to lookup instead of scanning all keys.

Signed-off-by: Torin Sandall <torinsandall@gmail.com>
---
 types/types.go            | 13 +++++++++----
 types/types_bench_test.go | 40 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 4 deletions(-)
 create mode 100644 types/types_bench_test.go

diff --git a/types/types.go b/types/types.go
index 7d78061d84..05f13c5e5c 100644
--- a/types/types.go
+++ b/types/types.go
@@ -346,16 +346,21 @@ func (t *Object) MarshalJSON() ([]byte, error) {
 
 // Select returns the type of the named property.
 func (t *Object) Select(name interface{}) Type {
-	for _, p := range t.static {
-		if util.Compare(p.Key, name) == 0 {
-			return p.Value
-		}
+
+	pos := sort.Search(len(t.static), func(x int) bool {
+		return util.Compare(t.static[x].Key, name) >= 0
+	})
+
+	if pos < len(t.static) && util.Compare(t.static[pos].Key, name) == 0 {
+		return t.static[pos].Value
 	}
+
 	if t.dynamic != nil {
 		if Contains(t.dynamic.Key, TypeOf(name)) {
 			return t.dynamic.Value
 		}
 	}
+
 	return nil
 }
 
diff --git a/types/types_bench_test.go b/types/types_bench_test.go
new file mode 100644
index 0000000000..21680a7ed1
--- /dev/null
+++ b/types/types_bench_test.go
@@ -0,0 +1,40 @@
+// Copyright 2021 The OPA Authors.  All rights reserved.
+// Use of this source code is governed by an Apache2
+// license that can be found in the LICENSE file.
+
+package types
+
+import (
+	"encoding/json"
+	"fmt"
+	"testing"
+)
+
+func BenchmarkSelect(b *testing.B) {
+	sizes := []int{1000, 10000, 100000}
+	for _, size := range sizes {
+		b.Run(fmt.Sprint(size), func(b *testing.B) {
+			tpe := generateType(size)
+			runSelectBenchmark(b, tpe, json.Number(fmt.Sprint(size-1)))
+		})
+	}
+}
+
+func runSelectBenchmark(b *testing.B, tpe Type, key interface{}) {
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		if result := Select(tpe, key); result != nil {
+			if Compare(result, N) != 0 {
+				b.Fatal("expected number type")
+			}
+		}
+	}
+}
+
+func generateType(n int) Type {
+	static := make([]*StaticProperty, n)
+	for i := 0; i < n; i++ {
+		static[i] = NewStaticProperty(json.Number(fmt.Sprint(i)), N)
+	}
+	return NewObject(static, nil)
+}