From 42e38b04bfd5759e065897db9cc70975f4537ebc Mon Sep 17 00:00:00 2001 From: TristonianJones Date: Thu, 6 Aug 2026 14:40:40 -0700 Subject: [PATCH 1/3] Support aggregate size computations over list, maps, and structs --- common/types/BUILD.bazel | 1 + common/types/aggregate_sizer.go | 29 ++ common/types/list.go | 33 +- common/types/list_test.go | 64 ++++ common/types/map.go | 37 ++- common/types/map_test.go | 95 ++++++ common/types/native.go | 17 + common/types/native_test.go | 82 +++++ common/types/object.go | 14 +- common/types/object_test.go | 23 ++ common/types/optional.go | 8 + common/types/optional_test.go | 19 ++ common/types/overflow.go | 24 ++ common/types/util.go | 275 ++++++++++++++++ common/types/util_test.go | 560 +++++++++++++++++++++++++++++++- 15 files changed, 1266 insertions(+), 15 deletions(-) create mode 100644 common/types/aggregate_sizer.go diff --git a/common/types/BUILD.bazel b/common/types/BUILD.bazel index 4ecc40031..b8d6926fb 100644 --- a/common/types/BUILD.bazel +++ b/common/types/BUILD.bazel @@ -8,6 +8,7 @@ package( go_library( name = "go_default_library", srcs = [ + "aggregate_sizer.go", "any_value.go", "bool.go", "bytes.go", diff --git a/common/types/aggregate_sizer.go b/common/types/aggregate_sizer.go new file mode 100644 index 000000000..ae046a0c2 --- /dev/null +++ b/common/types/aggregate_sizer.go @@ -0,0 +1,29 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package types + +// AggregateSizer calculates the recursive element size of values. +type AggregateSizer interface { + // AggregateSize returns the size of the input value, if known. + // Otherwise, a unit size of 1 is returned. + AggregateSize(val any) uint32 +} + +// AggregateSizeVisitor interface for ref.Val implementations capable of returning +// their total recursive element count. +type AggregateSizeVisitor interface { + // AggregateSize returns the total count of nested atomic elements (capped at math.MaxUint32). + AggregateSize(sizer AggregateSizer) uint32 +} diff --git a/common/types/list.go b/common/types/list.go index 028770ed6..483f71261 100644 --- a/common/types/list.go +++ b/common/types/list.go @@ -109,15 +109,10 @@ func NewMutableList(adapter Adapter) traits.MutableLister { // The `Adapter` enables native type to CEL type conversions. type baseList struct { Adapter - value any - - // size indicates the number of elements within the list. - // Since objects are immutable the size of a list is static. - size int - - // get returns a value at the specified integer index. - // The index is guaranteed to be checked against the list index range. - get func(int) any + value any + size int + aggSize uint32 + get func(int) any } // Add implements the traits.Adder interface method. @@ -269,6 +264,19 @@ func (l *baseList) Size() ref.Val { return Int(l.size) } +// AggregateSize implements the AggregateSizeVisitor interface method. +func (l *baseList) AggregateSize(sizer AggregateSizer) uint32 { + if l.aggSize != 0 { + return l.aggSize + } + total := uint32(1) + for i := range l.size { + total = safeAddUint32(total, sizer.AggregateSize(l.get(i))) + } + l.aggSize = total + return total +} + // Type implements the ref.Val interface method. func (l *baseList) Type() ref.Type { return ListType @@ -322,11 +330,13 @@ func (l *mutableList) Add(other ref.Val) ref.Val { case *mutableList: l.mutableValues = append(l.mutableValues, otherList.mutableValues...) l.size += len(otherList.mutableValues) + l.aggSize = 0 case traits.Lister: for i := IntZero; i < otherList.Size().(Int); i++ { l.size++ l.mutableValues = append(l.mutableValues, otherList.Get(i)) } + l.aggSize = 0 default: return MaybeNoSuchOverloadErr(otherList) } @@ -480,6 +490,11 @@ func (l *concatList) Size() ref.Val { return l.cachedSize } +// AggregateSize implements the AggregateSizeVisitor interface method. +func (l *concatList) AggregateSize(sizer AggregateSizer) uint32 { + return safeAddUint32(sizer.AggregateSize(l.prevList), sizer.AggregateSize(l.nextList)) +} + // String converts the concatenated list to a human-readable string. func (l *concatList) String() string { var sb strings.Builder diff --git a/common/types/list_test.go b/common/types/list_test.go index ca134b716..a962a3f1b 100644 --- a/common/types/list_test.go +++ b/common/types/list_test.go @@ -930,3 +930,67 @@ func TestConcatListSizeCached(t *testing.T) { } } } + +func TestListCalculateSize(t *testing.T) { + adapter := DefaultTypeAdapter + + // List literal: [1, [3, 4], [[7, 8], [9, 10]]] + l1 := NewRefValList(adapter, []ref.Val{Int(3), Int(4)}) + l2_1 := NewRefValList(adapter, []ref.Val{Int(7), Int(8)}) + l2_2 := NewRefValList(adapter, []ref.Val{Int(9), Int(10)}) + l2 := NewRefValList(adapter, []ref.Val{l2_1, l2_2}) + nested := NewRefValList(adapter, []ref.Val{Int(1), l1, l2}) + + tests := []struct { + name string + val ref.Val + want uint32 + }{ + { + name: "empty_list", + val: NewRefValList(adapter, []ref.Val{}), + want: 1, + }, + { + name: "flat_list", + val: l1, + want: 3, + }, + { + name: "nested_list", + val: nested, + want: 12, + }, + { + name: "concat_list", + val: l1.Add(l2_1), + want: 6, + }, + { + name: "string_list", + val: NewStringList(adapter, []string{"hello", "world"}), + want: 11, + }, + { + name: "dynamic_list", + val: NewDynamicList(adapter, []any{int64(1), []int64{3, 4}}), + want: 5, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + sizer, ok := tc.val.(AggregateSizeVisitor) + if !ok { + t.Fatalf("expected AggregateSizeVisitor implementation for %T", tc.val) + } + if got := sizer.AggregateSize(NewSizeCalculator()); got != tc.want { + t.Errorf("got aggregate size %d, want %d", got, tc.want) + } + // Caching check (memoized aggSize) + if got := sizer.AggregateSize(NewSizeCalculator()); got != tc.want { + t.Errorf("memoized AggregateSize() got %d, want %d", got, tc.want) + } + }) + } +} diff --git a/common/types/map.go b/common/types/map.go index e4d6f7657..093098998 100644 --- a/common/types/map.go +++ b/common/types/map.go @@ -142,8 +142,8 @@ type baseMap struct { // value is the native Go value upon which the map type operators. value any - // size is the number of entries in the map. - size int + size int + aggSize uint32 } // Contains implements the traits.Container interface method. @@ -303,6 +303,23 @@ func (m *baseMap) Size() ref.Val { return Int(m.size) } +// AggregateSize implements the AggregateSizeVisitor interface method. +func (m *baseMap) AggregateSize(sizer AggregateSizer) uint32 { + if m.aggSize != 0 { + return m.aggSize + } + var total uint32 = 1 + it := m.Iterator() + for it.HasNext() == True { + key := it.Next() + val, _ := m.Find(key) + total = safeAddUint32(total, sizer.AggregateSize(key)) + total = safeAddUint32(total, sizer.AggregateSize(val)) + } + m.aggSize = total + return total +} + // String converts the map into a human-readable string. func (m *baseMap) String() string { var sb strings.Builder @@ -380,6 +397,8 @@ func (m *mutableMap) Insert(k, v ref.Val) ref.Val { return NewErr("insert failed: key %v already exists", k) } m.mutableValues[k] = v + m.size++ + m.aggSize = 0 return m } @@ -909,6 +928,20 @@ func (m *protoMap) Size() ref.Val { return Int(m.value.Len()) } +// AggregateSize implements the AggregateSizeVisitor interface method. +func (m *protoMap) AggregateSize(sizer AggregateSizer) uint32 { + if m.value == nil { + return 0 + } + total := uint32(1) + m.value.Range(func(k protoreflect.MapKey, v protoreflect.Value) bool { + total = safeAddUint32(total, sizer.AggregateSize(k)) + total = safeAddUint32(total, sizer.AggregateSize(v)) + return true + }) + return total +} + // Type implements the ref.Val interface method. func (m *protoMap) Type() ref.Type { return MapType diff --git a/common/types/map_test.go b/common/types/map_test.go index 81989120d..f085b7e67 100644 --- a/common/types/map_test.go +++ b/common/types/map_test.go @@ -1252,3 +1252,98 @@ func (m proxyLegacyMap) Iterator() traits.Iterator { func (m proxyLegacyMap) Size() ref.Val { return m.proxy.Size() } + +func TestMapCalculateSize(t *testing.T) { + adapter := DefaultTypeAdapter + + // Setup helper data + l := NewRefValList(adapter, []ref.Val{Int(2), Int(3)}) + refValMap := NewRefValMap(adapter, map[ref.Val]ref.Val{ + String("a"): Int(1), + String("b"): l, + }) + + ifaceMap := NewStringInterfaceMap(adapter, map[string]any{ + "a": int64(1), + "b": []any{int64(2), int64(3)}, + }) + + mutMap := NewMutableMap(adapter, map[ref.Val]ref.Val{ + String("a"): Int(1), + String("b"): l, + }) + // Initial evaluation before insert to test aggSize reset + _ = mutMap.(AggregateSizeVisitor).AggregateSize(NewSizeCalculator()) + mutMap.Insert(String("c"), Int(4)) + + reg, err := NewRegistry(&proto3pb.TestAllTypes{}) + if err != nil { + t.Fatalf("NewRegistry() failed: %v", err) + } + msg := &proto3pb.TestAllTypes{ + MapStringString: map[string]string{ + "a": "b", + "c": "d", + }, + } + pbMsg := reg.NativeToValue(msg).(traits.Indexer) + pm := pbMsg.Get(String("map_string_string")).(traits.Mapper) + + tests := []struct { + name string + val ref.Val + want uint32 + }{ + { + name: "empty_ref_val_map", + val: NewRefValMap(adapter, map[ref.Val]ref.Val{}), + want: 1, + }, + { + name: "ref_val_map_nested", + val: refValMap, + want: 7, + }, + { + name: "string_interface_map", + val: ifaceMap, + want: 7, + }, + { + name: "string_string_map", + val: NewStringStringMap(adapter, map[string]string{"k1": "v1", "k2": "v2"}), + want: 9, + }, + { + name: "mutable_map_after_insert", + val: mutMap, + want: 9, + }, + { + name: "proto_map", + val: pm, + want: 5, + }, + { + name: "nil_proto_map", + val: &protoMap{}, + want: 0, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + sizer, ok := tc.val.(AggregateSizeVisitor) + if !ok { + t.Fatalf("expected AggregateSizeVisitor implementation for %T", tc.val) + } + if got := sizer.AggregateSize(NewSizeCalculator()); got != tc.want { + t.Errorf("got aggregate size %d, want %d", got, tc.want) + } + // Caching check (memoized aggSize) + if got := sizer.AggregateSize(NewSizeCalculator()); got != tc.want { + t.Errorf("memoized AggregateSize() got %d, want %d", got, tc.want) + } + }) + } +} diff --git a/common/types/native.go b/common/types/native.go index 802abdff4..4020e46ac 100644 --- a/common/types/native.go +++ b/common/types/native.go @@ -358,6 +358,23 @@ func (o *nativeObj) Value() any { return o.val } +// AggregateSize implements the AggregateSizeVisitor interface method. +func (o *nativeObj) AggregateSize(sizer AggregateSizer) uint32 { + refVal := reflect.Indirect(o.refValue) + if !refVal.IsValid() { + return 0 + } + total := uint32(1) + for _, fieldType := range o.valType.fieldsByName { + fieldValue := refVal.FieldByIndex(fieldType.Index) + if !fieldValue.IsValid() || fieldValue.IsZero() { + continue + } + total = safeAddUint32(total, sizer.AggregateSize(fieldValue)) + } + return total +} + func newNativeTypes(rawType reflect.Type, fieldNameHandler NativeTypesFieldNameHandler) ([]*NativeType, error) { nt, err := newNativeType(rawType, fieldNameHandler) if err != nil { diff --git a/common/types/native_test.go b/common/types/native_test.go index 161337538..c18dfc8ea 100644 --- a/common/types/native_test.go +++ b/common/types/native_test.go @@ -1303,6 +1303,88 @@ func TestNativeToValueDelegatesUnregisteredStructs(t *testing.T) { } } +func TestNativeObjectCalculateSize(t *testing.T) { + env, err := cel.NewEnv( + ext.NativeTypes( + reflect.TypeOf(TestAllTypes{}), + reflect.TypeOf(TestNestedType{}), + ), + ) + if err != nil { + t.Fatalf("cel.NewEnv() failed: %v", err) + } + adapter := env.CELTypeAdapter() + + tests := []struct { + name string + val any + want uint32 + }{ + { + name: "empty_struct", + val: &TestNestedType{}, + want: 1, // 1 (container) + }, + { + name: "struct_with_scalar_and_list", + val: &TestNestedType{ + NestedListVal: []string{"a", "b", "c"}, + }, + want: 5, // 1 (root struct) + ["a", "b", "c"] (1 list container + 3 elements = 4) = 5 + }, + { + name: "struct_with_nested_map", + val: &TestNestedType{ + NestedMapVal: map[int64]bool{1: true, 2: false}, + }, + want: 6, // 1 (root struct) + map (1 container + (1+1) + (1+1) = 5) = 6 + }, + { + name: "nested_struct", + val: &TestAllTypes{ + StringVal: "hello", + NestedVal: &TestNestedType{ + NestedListVal: []string{"a", "b"}, + }, + }, + // 1 (root struct) + "hello"(5) + NestedVal(1 container + ["a", "b"](1+2=3) = 4) = 10 + want: 10, + }, + { + name: "bytes_and_time", + val: &TestAllTypes{ + BytesVal: []byte("test"), + DurationVal: time.Second, + TimestampVal: time.Unix(100, 0), + }, + want: 7, + }, + { + name: "slice_of_structs", + val: &TestAllTypes{ + ListVal: []*TestNestedType{ + {NestedListVal: []string{"x"}}, + {NestedListVal: []string{"y", "z"}}, + }, + }, + want: 9, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + val := adapter.NativeToValue(tc.val) + sizer, ok := val.(types.AggregateSizeVisitor) + if !ok { + t.Fatalf("expected types.AggregateSizeVisitor implementation for %T", val) + } + if got := sizer.AggregateSize(types.NewSizeCalculator()); got != tc.want { + t.Errorf("got aggregate size %d, want %d", got, tc.want) + } + }) + } +} + func BenchmarkNativeTypesEval(b *testing.B) { benchmarks := []struct { name string diff --git a/common/types/object.go b/common/types/object.go index bb2a09e87..1d45d4e88 100644 --- a/common/types/object.go +++ b/common/types/object.go @@ -167,9 +167,17 @@ func (o *protoObj) Value() any { return o.value } -type protoObjField struct { - fd protoreflect.FieldDescriptor - v protoreflect.Value +// AggregateSize implements the AggregateSizeVisitor interface method. +func (o *protoObj) AggregateSize(sizer AggregateSizer) uint32 { + if o.value == nil { + return 0 + } + total := uint32(1) + o.value.ProtoReflect().Range(func(fd protoreflect.FieldDescriptor, v protoreflect.Value) bool { + total = safeAddUint32(total, sizer.AggregateSize(v)) + return true + }) + return total } func (o *protoObj) format(sb *strings.Builder) { diff --git a/common/types/object_test.go b/common/types/object_test.go index b2e2207ea..3412d9188 100644 --- a/common/types/object_test.go +++ b/common/types/object_test.go @@ -257,3 +257,26 @@ func TestProtoObjectConvertToType(t *testing.T) { t.Error("identity type conversion failed") } } + +func TestProtoObjectCalculateSize(t *testing.T) { + msg := &exprpb.ParsedExpr{ + SourceInfo: &exprpb.SourceInfo{ + LineOffsets: []int32{1, 2, 3}, + }, + } + reg := newTestRegistry(t, ProtoTypeDefs(msg)) + objVal := reg.NativeToValue(msg) + sizer, ok := objVal.(AggregateSizeVisitor) + if !ok { + t.Fatalf("expected AggregateSizeVisitor implementation for protoObj") + } + // 1 (protoObj container) + SourceInfo field (1 container + 1 list container + 3 list elements = 5) = 6 + if got := sizer.AggregateSize(NewSizeCalculator()); got != 6 { + t.Errorf("got aggregate size %d, want 6", got) + } + + nilObj := &protoObj{} + if got := nilObj.AggregateSize(NewSizeCalculator()); got != 0 { + t.Errorf("got nil protoObj aggregate size %d, want 0", got) + } +} diff --git a/common/types/optional.go b/common/types/optional.go index 0d861823d..5d27f2c2d 100644 --- a/common/types/optional.go +++ b/common/types/optional.go @@ -120,3 +120,11 @@ func (o *Optional) Value() any { } return o.value.Value() } + +// AggregateSize implements the AggregateSizeVisitor interface method. +func (o *Optional) AggregateSize(sizer AggregateSizer) uint32 { + if !o.HasValue() { + return 0 + } + return safeAddUint32(1, sizer.AggregateSize(o.value)) +} diff --git a/common/types/optional_test.go b/common/types/optional_test.go index 89f28d7e7..f41b77b6c 100644 --- a/common/types/optional_test.go +++ b/common/types/optional_test.go @@ -185,3 +185,22 @@ func TestOptionalValue(t *testing.T) { t.Errorf("OptionalNone.Value() got %v, wanted nil", OptionalNone.Value()) } } + +func TestOptionalCalculateSize(t *testing.T) { + calc := NewSizeCalculator() + none := OptionalNone + if sizer, ok := any(none).(AggregateSizeVisitor); !ok || sizer.AggregateSize(calc) != 0 { + t.Errorf("expected 0 for OptionalNone") + } + + someScalar := OptionalOf(Int(42)) + if sizer, ok := any(someScalar).(AggregateSizeVisitor); !ok || sizer.AggregateSize(calc) != 2 { + t.Errorf("got %d for OptionalOf(scalar), want 2", sizer.AggregateSize(calc)) + } + + l := NewRefValList(DefaultTypeAdapter, []ref.Val{Int(1), Int(2)}) + someList := OptionalOf(l) + if sizer, ok := any(someList).(AggregateSizeVisitor); !ok || sizer.AggregateSize(calc) != 4 { + t.Errorf("got %d for OptionalOf(list of 2), want 4", sizer.AggregateSize(calc)) + } +} diff --git a/common/types/overflow.go b/common/types/overflow.go index dcb66ef59..d56fb1dd0 100644 --- a/common/types/overflow.go +++ b/common/types/overflow.go @@ -427,3 +427,27 @@ func uint64ToInt64Lossless(v uint64) (int64, bool) { i, err := uint64ToInt64Checked(v) return i, err == nil } + +func safeAddUint32(a, b uint32) uint32 { + if math.MaxUint32-a < b { + return math.MaxUint32 + } + return a + b +} + +func safeUint32FromInt(n int) uint32 { + if n < 0 { + return 0 + } + if uint64(n) > math.MaxUint32 { + return math.MaxUint32 + } + return uint32(n) +} + +func safeUint32FromBoxedInt(v Int) uint32 { + if v < 0 || v > math.MaxUint32 { + return math.MaxUint32 + } + return uint32(v) +} diff --git a/common/types/util.go b/common/types/util.go index 71662eee3..c43e773cd 100644 --- a/common/types/util.go +++ b/common/types/util.go @@ -15,7 +15,15 @@ package types import ( + "reflect" + "time" + "unicode/utf8" + + "google.golang.org/protobuf/proto" + "google.golang.org/protobuf/reflect/protoreflect" + "github.com/google/cel-go/common/types/ref" + "github.com/google/cel-go/common/types/traits" ) // IsUnknownOrError returns whether the input element ref.Val is an ErrType or UnknownType. @@ -46,3 +54,270 @@ func Equal(lhs ref.Val, rhs ref.Val) ref.Val { } return lhs.Equal(rhs) } + +// SizeCalculator calculates the recursive element size of values. +type SizeCalculator struct { + version int +} + +// NewSizeCalculator returns a new SizeCalculator for the specified version. +func NewSizeCalculator() *SizeCalculator { + return &SizeCalculator{version: 0} +} + +// Version returns the calculation version. +func (s *SizeCalculator) Version() int { + return s.version +} + +// AggregateSize returns the size of the input value, if known. +// Otherwise, a unit size of 1 is returned. +func (s *SizeCalculator) AggregateSize(val any) uint32 { + switch v := val.(type) { + case AggregateSizeVisitor: + return v.AggregateSize(s) + case traits.Mapper: + total := uint32(1) + it := v.Iterator() + for it.HasNext() == True { + key := it.Next() + val, _ := v.Find(key) + total = safeAddUint32(total, s.AggregateSize(key)) + total = safeAddUint32(total, s.AggregateSize(val)) + } + return total + case traits.Lister: + total := uint32(1) + it := v.Iterator() + for it.HasNext() == True { + total = safeAddUint32(total, s.AggregateSize(it.Next())) + } + return total + case String: + return safeUint32FromInt(utf8.RuneCountInString(string(v))) + case Bytes: + return safeUint32FromInt(len(v)) + case traits.Sizer: + return safeUint32FromBoxedInt(v.Size().(Int)) + case Bool, Int, Uint, Double, Duration, Timestamp, Null, *Type, *Err, *Unknown: + return 1 + case ref.Val: + return s.AggregateSize(v.Value()) + case protoreflect.Value: + return getProtoValueAggregateSize(s, v) + case protoreflect.MapKey: + return getProtoValueAggregateSize(s, v.Value()) + case protoreflect.Message: + return getProtoMessageAggregateSize(s, v) + case protoreflect.List: + return getProtoListAggregateSize(s, v) + case protoreflect.Map: + return getProtoMapAggregateSize(s, v) + case proto.Message: + if v == nil { + return 0 + } + return getProtoMessageAggregateSize(s, v.ProtoReflect()) + case reflect.Value: + return getReflectValueAggregateSize(s, v) + case string: + return safeUint32FromInt(utf8.RuneCountInString(v)) + case []byte: + return safeUint32FromInt(len(v)) + case int, int8, int16, int32, int64, + uint, uint8, uint16, uint32, uint64, + float32, float64, bool, time.Time, time.Duration, nil: + return 1 + default: + return getReflectValueAggregateSize(s, reflect.ValueOf(val)) + } +} + +func getProtoValueAggregateSize(s AggregateSizer, v protoreflect.Value) uint32 { + switch val := v.Interface().(type) { + case string: + return safeUint32FromInt(utf8.RuneCountInString(val)) + case []byte: + return safeUint32FromInt(len(val)) + case protoreflect.Message: + return getProtoMessageAggregateSize(s, val) + case protoreflect.List: + return getProtoListAggregateSize(s, val) + case protoreflect.Map: + return getProtoMapAggregateSize(s, val) + default: + return 1 + } +} + +func getProtoFieldAggregateSize(s AggregateSizer, fd protoreflect.FieldDescriptor, v protoreflect.Value) uint32 { + if fd.IsMap() { + m := v.Map() + total := uint32(1) + valKind := fd.MapValue().Kind() + m.Range(func(k protoreflect.MapKey, val protoreflect.Value) bool { + total = safeAddUint32(total, getProtoMapKeyAggregateSize(fd.MapKey(), k)) + total = safeAddUint32(total, getProtoFieldByKindAggregateSize(s, valKind, val)) + return true + }) + return total + } + if fd.IsList() { + l := v.List() + elemKind := fd.Kind() + total := safeAddUint32(1, safeUint32FromInt(l.Len())) + if elemKind == protoreflect.StringKind { + total = 1 + for i := 0; i < l.Len(); i++ { + total = safeAddUint32(total, safeUint32FromInt(utf8.RuneCountInString(l.Get(i).String()))) + } + } else if elemKind == protoreflect.BytesKind { + total = 1 + for i := 0; i < l.Len(); i++ { + total = safeAddUint32(total, safeUint32FromInt(len(l.Get(i).Bytes()))) + } + } else if elemKind == protoreflect.MessageKind || elemKind == protoreflect.GroupKind { + total = 1 + for i := 0; i < l.Len(); i++ { + total = safeAddUint32(total, getProtoMessageAggregateSize(s, l.Get(i).Message())) + } + } + return total + } + return getProtoFieldByKindAggregateSize(s, fd.Kind(), v) +} + +func getProtoMapKeyAggregateSize(keyDesc protoreflect.FieldDescriptor, k protoreflect.MapKey) uint32 { + if keyDesc.Kind() == protoreflect.StringKind { + return safeUint32FromInt(utf8.RuneCountInString(k.String())) + } + return 1 +} + +func getProtoFieldByKindAggregateSize(s AggregateSizer, kind protoreflect.Kind, v protoreflect.Value) uint32 { + switch kind { + case protoreflect.StringKind: + return safeUint32FromInt(utf8.RuneCountInString(v.String())) + case protoreflect.BytesKind: + return safeUint32FromInt(len(v.Bytes())) + case protoreflect.MessageKind, protoreflect.GroupKind: + return getProtoMessageAggregateSize(s, v.Message()) + default: + return 1 + } +} + +func getProtoMessageAggregateSize(s AggregateSizer, m protoreflect.Message) uint32 { + if !m.IsValid() { + return 0 + } + total := uint32(1) + m.Range(func(fd protoreflect.FieldDescriptor, v protoreflect.Value) bool { + total = safeAddUint32(total, getProtoFieldAggregateSize(s, fd, v)) + return true + }) + return total +} + +func getProtoListAggregateSize(s AggregateSizer, l protoreflect.List) uint32 { + if !l.IsValid() { + return 0 + } + total := uint32(1) + for i := 0; i < l.Len(); i++ { + total = safeAddUint32(total, getProtoValueAggregateSize(s, l.Get(i))) + } + return total +} + +func getProtoMapAggregateSize(s AggregateSizer, m protoreflect.Map) uint32 { + if !m.IsValid() { + return 0 + } + total := uint32(1) + m.Range(func(k protoreflect.MapKey, v protoreflect.Value) bool { + total = safeAddUint32(total, getProtoValueAggregateSize(s, k.Value())) + total = safeAddUint32(total, getProtoValueAggregateSize(s, v)) + return true + }) + return total +} + +func getReflectValueAggregateSize(s AggregateSizer, fieldVal reflect.Value) uint32 { + if !fieldVal.IsValid() { + return 0 + } + switch fieldVal.Kind() { + case reflect.String: + return safeUint32FromInt(utf8.RuneCountInString(fieldVal.String())) + case reflect.Slice, reflect.Array: + elemType := fieldVal.Type().Elem() + if elemType.Kind() == reflect.Uint8 { + return safeUint32FromInt(fieldVal.Len()) + } + total := safeAddUint32(1, safeUint32FromInt(fieldVal.Len())) + switch elemType.Kind() { + case reflect.String: + total = 1 + for i := 0; i < fieldVal.Len(); i++ { + total = safeAddUint32(total, safeUint32FromInt(utf8.RuneCountInString(fieldVal.Index(i).String()))) + } + case reflect.Struct, reflect.Pointer, reflect.Slice, reflect.Array, reflect.Map, reflect.Interface: + total = 1 + for i := 0; i < fieldVal.Len(); i++ { + total = safeAddUint32(total, getReflectValueAggregateSize(s, fieldVal.Index(i))) + } + } + return total + case reflect.Map: + total := uint32(1) + iter := fieldVal.MapRange() + for iter.Next() { + total = safeAddUint32(total, getReflectValueAggregateSize(s, iter.Key())) + total = safeAddUint32(total, getReflectValueAggregateSize(s, iter.Value())) + } + return total + case reflect.Pointer, reflect.Interface: + if fieldVal.IsNil() { + return 0 + } + if fieldVal.CanInterface() { + if sizer, ok := fieldVal.Interface().(AggregateSizeVisitor); ok { + return sizer.AggregateSize(s) + } + if sizer, ok := fieldVal.Interface().(traits.Sizer); ok { + return safeUint32FromBoxedInt(sizer.Size().(Int)) + } + } + return getReflectValueAggregateSize(s, fieldVal.Elem()) + case reflect.Struct: + if fieldVal.Type() == timestampType || fieldVal.Type() == durationType { + return 1 + } + if fieldVal.CanInterface() { + if sizer, ok := fieldVal.Interface().(AggregateSizeVisitor); ok { + return sizer.AggregateSize(s) + } + if sizer, ok := fieldVal.Interface().(traits.Sizer); ok { + return safeUint32FromBoxedInt(sizer.Size().(Int)) + } + } + total := uint32(1) + t := fieldVal.Type() + numFields := fieldVal.NumField() + for i := 0; i < numFields; i++ { + f := t.Field(i) + if !f.IsExported() { + continue + } + fVal := fieldVal.Field(i) + if !fVal.IsValid() || fVal.IsZero() { + continue + } + total = safeAddUint32(total, getReflectValueAggregateSize(s, fVal)) + } + return total + default: + return 1 + } +} diff --git a/common/types/util_test.go b/common/types/util_test.go index b10b3e84c..96b2fdf1a 100644 --- a/common/types/util_test.go +++ b/common/types/util_test.go @@ -14,7 +14,250 @@ package types -import "testing" +import ( + "fmt" + "math" + "reflect" + "testing" + "time" + + "google.golang.org/protobuf/reflect/protoreflect" + + "github.com/google/cel-go/common/types/ref" + proto3pb "github.com/google/cel-go/test/proto3pb" +) + +func TestCalculateSize(t *testing.T) { + adapter := DefaultTypeAdapter + + tests := []struct { + name string + val any + want uint32 + }{ + { + name: "aggregate_sizer_list", + val: NewRefValList(adapter, []ref.Val{Int(1), Int(2)}), + want: 3, + }, + { + name: "sizer_string", + val: String("hello"), + want: 5, + }, + { + name: "sizer_bytes", + val: Bytes("world"), + want: 5, + }, + { + name: "err_val", + val: NewErr("test error"), + want: 1, + }, + { + name: "unknown_val", + val: &Unknown{}, + want: 1, + }, + { + name: "type_val", + val: IntType, + want: 1, + }, + { + name: "null_val", + val: NullValue, + want: 1, + }, + { + name: "scalar_ref_val_int", + val: Int(42), + want: 1, + }, + { + name: "scalar_ref_val_double", + val: Double(1.5), + want: 1, + }, + { + name: "scalar_ref_val_bool", + val: True, + want: 1, + }, + { + name: "scalar_ref_val_timestamp", + val: Timestamp{Time: time.Unix(100, 0)}, + want: 1, + }, + { + name: "scalar_ref_val_duration", + val: Duration{Duration: time.Second}, + want: 1, + }, + { + name: "proto_value_string", + val: protoreflect.ValueOfString("hello"), + want: 5, + }, + { + name: "proto_value_bytes", + val: protoreflect.ValueOfBytes([]byte("world")), + want: 5, + }, + { + name: "proto_value_int", + val: protoreflect.ValueOfInt32(42), + want: 1, + }, + { + name: "proto_map_key", + val: protoreflect.MapKey(protoreflect.ValueOfString("key")), + want: 3, + }, + { + name: "proto_message", + val: &proto3pb.TestAllTypes{SingleString: "hello"}, + want: 6, // 1 (root) + 5 (string) = 6 + }, + { + name: "proto_message_with_list_and_map", + val: &proto3pb.TestAllTypes{ + RepeatedString: []string{"a", "b"}, + MapStringString: map[string]string{"k": "v"}, + }, + want: 7, + }, + { + name: "protoreflect_message", + val: (&proto3pb.TestAllTypes{SingleInt64: 10}).ProtoReflect(), + want: 2, // 1 (root) + 1 (int64) = 2 + }, + { + name: "protoreflect_list", + val: (&proto3pb.TestAllTypes{RepeatedString: []string{"a", "b"}}).ProtoReflect().Get((&proto3pb.TestAllTypes{}).ProtoReflect().Descriptor().Fields().ByName("repeated_string")).List(), + want: 3, // 1 (container) + 1("a") + 1("b") = 3 + }, + { + name: "protoreflect_map", + val: (&proto3pb.TestAllTypes{MapStringString: map[string]string{"k": "v"}}).ProtoReflect().Get((&proto3pb.TestAllTypes{}).ProtoReflect().Descriptor().Fields().ByName("map_string_string")).Map(), + want: 3, // 1 (container) + 1("k") + 1("v") = 3 + }, + { + name: "nil_proto_message", + val: (*proto3pb.TestAllTypes)(nil), + want: 0, + }, + { + name: "reflect_value", + val: reflect.ValueOf("reflected"), + want: 9, + }, + { + name: "native_string", + val: "hello", + want: 5, + }, + { + name: "native_bytes", + val: []byte("world"), + want: 5, + }, + { + name: "native_int", + val: 42, + want: 1, + }, + { + name: "native_float", + val: 3.14, + want: 1, + }, + { + name: "native_bool", + val: true, + want: 1, + }, + { + name: "native_time", + val: time.Now(), + want: 1, + }, + { + name: "native_duration", + val: time.Hour, + want: 1, + }, + { + name: "native_nil", + val: nil, + want: 1, + }, + { + name: "custom_struct", + val: struct{ Name string }{"cel"}, + want: 4, // 1 (root) + 3 ("cel") = 4 + }, + { + name: "custom_lister", + val: proxyLegacyList{proxy: NewRefValList(DefaultTypeAdapter, []ref.Val{String("a"), String("b")})}, + want: 3, // 1 (container) + 1 ("a") + 1 ("b") = 3 + }, + { + name: "custom_mapper", + val: interopFoldableMap{Mapper: NewStringStringMap(DefaultTypeAdapter, map[string]string{"key": "val"})}, + want: 7, // 1 (container) + 3 ("key") + 3 ("val") = 7 + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + calculator := NewSizeCalculator() + if got := calculator.AggregateSize(tc.val); got != tc.want { + t.Errorf("AggregateSize(%v) got %d, want %d", tc.val, got, tc.want) + } + }) + } +} + +func TestSafeUint32Helpers(t *testing.T) { + // safeAddUint32 + if got := safeAddUint32(10, 20); got != 30 { + t.Errorf("safeAddUint32(10, 20) got %d, want 30", got) + } + if got := safeAddUint32(math.MaxUint32-5, 10); got != math.MaxUint32 { + t.Errorf("safeAddUint32(overflow) got %d, want MaxUint32", got) + } + + // safeUint32FromInt + if got := safeUint32FromInt(42); got != 42 { + t.Errorf("safeUint32FromInt(42) got %d, want 42", got) + } + if got := safeUint32FromInt(-1); got != 0 { + t.Errorf("safeUint32FromInt(-1) got %d, want 0", got) + } + if got := safeUint32FromInt(int(uint64(math.MaxUint32) + 100)); got != math.MaxUint32 { + t.Errorf("safeUint32FromInt(overflow) got %d, want MaxUint32", got) + } + + // safeUint32FromBoxedInt + if got := safeUint32FromBoxedInt(Int(42)); got != 42 { + t.Errorf("safeUint32FromBoxedInt(42) got %d, want 42", got) + } + if got := safeUint32FromBoxedInt(Int(-1)); got != math.MaxUint32 { + t.Errorf("safeUint32FromBoxedInt(-1) got %d, want MaxUint32", got) + } + if got := safeUint32FromBoxedInt(Int(int64(math.MaxUint32) + 100)); got != math.MaxUint32 { + t.Errorf("safeUint32FromBoxedInt(overflow) got %d, want MaxUint32", got) + } +} + +func TestNativeObjCalculateSizeNil(t *testing.T) { + nilNative := &nativeObj{} + if got := nilNative.AggregateSize(NewSizeCalculator()); got != 0 { + t.Errorf("nil nativeObj.AggregateSize() got %d, want 0", got) + } +} func BenchmarkIsUnknownOrError(b *testing.B) { err := NewErr("test") @@ -25,3 +268,318 @@ func BenchmarkIsUnknownOrError(b *testing.B) { } } } + +type nestedNative struct { + NestedList []string + NestedMap map[string]int +} + +type rootNative struct { + Name string + Count int + Children []nestedNative +} + +func BenchmarkCalculateSizeAmortized(b *testing.B) { + adapter := DefaultTypeAdapter + + flatList := NewRefValList(adapter, []ref.Val{Int(1), Int(2), Int(3), Int(4), Int(5), Int(6), Int(7), Int(8)}) + nestedList := NewRefValList(adapter, []ref.Val{ + String("hello"), + NewRefValList(adapter, []ref.Val{String("nested1"), String("nested2")}), + NewRefValMap(adapter, map[ref.Val]ref.Val{String("k1"): String("v1")}), + }) + flatMap := NewRefValMap(adapter, map[ref.Val]ref.Val{ + String("k1"): Int(1), + String("k2"): Int(2), + String("k3"): Int(3), + }) + protoMsg := &proto3pb.TestAllTypes{ + SingleString: "hello world", + SingleInt64: 42, + RepeatedString: []string{"first", "second", "third"}, + MapStringString: map[string]string{ + "key1": "value1", + "key2": "value2", + }, + } + nativeData := &rootNative{ + Name: "parent", + Count: 100, + Children: []nestedNative{ + {NestedList: []string{"a", "b", "c"}, NestedMap: map[string]int{"k1": 1, "k2": 2}}, + {NestedList: []string{"d", "e"}, NestedMap: map[string]int{"k3": 3}}, + }, + } + nativeVal := adapter.NativeToValue(nativeData) + + benchmarks := []struct { + name string + val any + }{ + {name: "scalar_int", val: Int(42)}, + {name: "scalar_string", val: String("hello world this is a test string")}, + {name: "native_string", val: "hello world this is a test string"}, + {name: "native_bytes", val: []byte("hello world this is a test string")}, + {name: "list_flat", val: flatList}, + {name: "list_nested", val: nestedList}, + {name: "map_flat", val: flatMap}, + {name: "custom_list_flat", val: proxyLegacyList{proxy: flatList}}, + {name: "custom_list_nested", val: proxyLegacyList{proxy: nestedList}}, + {name: "custom_map_flat", val: interopFoldableMap{Mapper: flatMap}}, + {name: "proto_message", val: protoMsg}, + {name: "proto_obj", val: adapter.NativeToValue(protoMsg)}, + {name: "native_obj", val: nativeVal}, + {name: "native_struct", val: nativeData}, + } + + for _, bm := range benchmarks { + b.Run(bm.name, func(b *testing.B) { + b.ReportAllocs() + calculator := NewSizeCalculator() + for i := 0; i < b.N; i++ { + _ = calculator.AggregateSize(bm.val) + } + }) + } +} + +func BenchmarkCalculateSizeDynamic(b *testing.B) { + adapter := DefaultTypeAdapter + + benchmarks := []struct { + name string + valFn func() any + }{ + { + name: "list_flat", + valFn: func() any { + return NewRefValList(adapter, []ref.Val{Int(1), Int(2), Int(3), Int(4), Int(5), Int(6), Int(7), Int(8)}) + }, + }, + { + name: "list_nested", + valFn: func() any { + return NewRefValList(adapter, []ref.Val{ + String("hello"), + NewRefValList(adapter, []ref.Val{String("nested1"), String("nested2")}), + NewRefValMap(adapter, map[ref.Val]ref.Val{String("k1"): String("v1")}), + }) + }, + }, + { + name: "map_flat", + valFn: func() any { + return NewRefValMap(adapter, map[ref.Val]ref.Val{ + String("k1"): Int(1), + String("k2"): Int(2), + String("k3"): Int(3), + }) + }, + }, + { + name: "custom_list_flat", + valFn: func() any { + return proxyLegacyList{proxy: NewRefValList(adapter, []ref.Val{Int(1), Int(2), Int(3), Int(4), Int(5), Int(6), Int(7), Int(8)})} + }, + }, + { + name: "custom_list_nested", + valFn: func() any { + return proxyLegacyList{proxy: NewRefValList(adapter, []ref.Val{ + String("hello"), + NewRefValList(adapter, []ref.Val{String("nested1"), String("nested2")}), + NewRefValMap(adapter, map[ref.Val]ref.Val{String("k1"): String("v1")}), + })} + }, + }, + { + name: "custom_map_flat", + valFn: func() any { + return interopFoldableMap{Mapper: NewRefValMap(adapter, map[ref.Val]ref.Val{ + String("k1"): Int(1), + String("k2"): Int(2), + String("k3"): Int(3), + })} + }, + }, + { + name: "proto_obj", + valFn: func() any { + return adapter.NativeToValue(&proto3pb.TestAllTypes{ + SingleString: "hello world", + SingleInt64: 42, + RepeatedString: []string{"first", "second", "third"}, + MapStringString: map[string]string{ + "key1": "value1", + "key2": "value2", + }, + }) + }, + }, + { + name: "native_obj", + valFn: func() any { + return adapter.NativeToValue(&rootNative{ + Name: "parent", + Count: 100, + Children: []nestedNative{ + {NestedList: []string{"a", "b", "c"}, NestedMap: map[string]int{"k1": 1, "k2": 2}}, + {NestedList: []string{"d", "e"}, NestedMap: map[string]int{"k3": 3}}, + }, + }) + }, + }, + } + + for _, bm := range benchmarks { + b.Run(bm.name, func(b *testing.B) { + b.ReportAllocs() + calculator := NewSizeCalculator() + for i := 0; i < b.N; i++ { + val := bm.valFn() + _ = calculator.AggregateSize(val) + } + }) + } +} + +func BenchmarkCalculateSizeScaled(b *testing.B) { + adapter := DefaultTypeAdapter + sizes := []int{10, 100, 1000} + + for _, size := range sizes { + // Prepare list elements + listElems := make([]ref.Val, size) + for i := 0; i < size; i++ { + listElems[i] = Int(i) + } + builtinList := NewRefValList(adapter, listElems) + customList := proxyLegacyList{proxy: builtinList} + + // Prepare map entries + mapEntries := make(map[ref.Val]ref.Val, size) + for i := 0; i < size; i++ { + mapEntries[String(fmt.Sprintf("k%d", i))] = Int(i) + } + builtinMap := NewRefValMap(adapter, mapEntries) + customMap := interopFoldableMap{Mapper: builtinMap} + + // Amortized (repeated calculation on memoized vs unmemoized custom instance) + b.Run(fmt.Sprintf("Amortized/builtin_list/N=%d", size), func(b *testing.B) { + b.ReportAllocs() + calc := NewSizeCalculator() + for i := 0; i < b.N; i++ { + _ = calc.AggregateSize(builtinList) + } + }) + b.Run(fmt.Sprintf("Amortized/custom_list/N=%d", size), func(b *testing.B) { + b.ReportAllocs() + calc := NewSizeCalculator() + for i := 0; i < b.N; i++ { + _ = calc.AggregateSize(customList) + } + }) + b.Run(fmt.Sprintf("Amortized/builtin_map/N=%d", size), func(b *testing.B) { + b.ReportAllocs() + calc := NewSizeCalculator() + for i := 0; i < b.N; i++ { + _ = calc.AggregateSize(builtinMap) + } + }) + b.Run(fmt.Sprintf("Amortized/custom_map/N=%d", size), func(b *testing.B) { + b.ReportAllocs() + calc := NewSizeCalculator() + for i := 0; i < b.N; i++ { + _ = calc.AggregateSize(customMap) + } + }) + + // First-time / Uncached calculation + b.Run(fmt.Sprintf("FirstTime/builtin_list/N=%d", size), func(b *testing.B) { + b.ReportAllocs() + calc := NewSizeCalculator() + for i := 0; i < b.N; i++ { + l := NewRefValList(adapter, listElems) + _ = calc.AggregateSize(l) + } + }) + b.Run(fmt.Sprintf("FirstTime/custom_list/N=%d", size), func(b *testing.B) { + b.ReportAllocs() + calc := NewSizeCalculator() + for i := 0; i < b.N; i++ { + l := proxyLegacyList{proxy: NewRefValList(adapter, listElems)} + _ = calc.AggregateSize(l) + } + }) + b.Run(fmt.Sprintf("FirstTime/builtin_map/N=%d", size), func(b *testing.B) { + b.ReportAllocs() + calc := NewSizeCalculator() + for i := 0; i < b.N; i++ { + m := NewRefValMap(adapter, mapEntries) + _ = calc.AggregateSize(m) + } + }) + b.Run(fmt.Sprintf("FirstTime/custom_map/N=%d", size), func(b *testing.B) { + b.ReportAllocs() + calc := NewSizeCalculator() + for i := 0; i < b.N; i++ { + m := interopFoldableMap{Mapper: NewRefValMap(adapter, mapEntries)} + _ = calc.AggregateSize(m) + } + }) + } + + // Benchmark nested tree complexity (Depth x Width) + depths := []int{2, 3} + width := 10 + for _, depth := range depths { + builtinNested := createNestedList(adapter, depth, width) + customNested := createNestedCustomList(adapter, depth, width) + + b.Run(fmt.Sprintf("Complexity/builtin_nested_list/Depth=%d_Width=%d", depth, width), func(b *testing.B) { + b.ReportAllocs() + calc := NewSizeCalculator() + for i := 0; i < b.N; i++ { + _ = calc.AggregateSize(builtinNested) + } + }) + b.Run(fmt.Sprintf("Complexity/custom_nested_list/Depth=%d_Width=%d", depth, width), func(b *testing.B) { + b.ReportAllocs() + calc := NewSizeCalculator() + for i := 0; i < b.N; i++ { + _ = calc.AggregateSize(customNested) + } + }) + } +} + +func createNestedList(adapter Adapter, depth, width int) ref.Val { + if depth <= 1 { + elems := make([]ref.Val, width) + for i := 0; i < width; i++ { + elems[i] = Int(i) + } + return NewRefValList(adapter, elems) + } + elems := make([]ref.Val, width) + for i := 0; i < width; i++ { + elems[i] = createNestedList(adapter, depth-1, width) + } + return NewRefValList(adapter, elems) +} + +func createNestedCustomList(adapter Adapter, depth, width int) ref.Val { + if depth <= 1 { + elems := make([]ref.Val, width) + for i := 0; i < width; i++ { + elems[i] = Int(i) + } + return proxyLegacyList{proxy: NewRefValList(adapter, elems)} + } + elems := make([]ref.Val, width) + for i := 0; i < width; i++ { + elems[i] = createNestedCustomList(adapter, depth-1, width) + } + return proxyLegacyList{proxy: NewRefValList(adapter, elems)} +} From 10498063197e8bca583fae40331d29d1ceb3bbef Mon Sep 17 00:00:00 2001 From: TristonianJones Date: Fri, 14 Aug 2026 14:36:18 -0700 Subject: [PATCH 2/3] Depth and traversal limits for aggregate size calculations --- common/types/BUILD.bazel | 2 + common/types/overflow.go | 5 +- common/types/size_calc.go | 396 ++++++++++++++++++++ common/types/size_calc_test.go | 661 +++++++++++++++++++++++++++++++++ common/types/util.go | 275 -------------- common/types/util_test.go | 527 +------------------------- 6 files changed, 1062 insertions(+), 804 deletions(-) create mode 100644 common/types/size_calc.go create mode 100644 common/types/size_calc_test.go diff --git a/common/types/BUILD.bazel b/common/types/BUILD.bazel index b8d6926fb..9a2290f3e 100644 --- a/common/types/BUILD.bazel +++ b/common/types/BUILD.bazel @@ -29,6 +29,7 @@ go_library( "overflow.go", "provider.go", "regex.go", + "size_calc.go", "string.go", "struct.go", "timestamp.go", @@ -77,6 +78,7 @@ go_test( "optional_test.go", "provider_test.go", "regex_test.go", + "size_calc_test.go", "string_test.go", "timestamp_test.go", "types_test.go", diff --git a/common/types/overflow.go b/common/types/overflow.go index d56fb1dd0..49b15377f 100644 --- a/common/types/overflow.go +++ b/common/types/overflow.go @@ -436,10 +436,7 @@ func safeAddUint32(a, b uint32) uint32 { } func safeUint32FromInt(n int) uint32 { - if n < 0 { - return 0 - } - if uint64(n) > math.MaxUint32 { + if n < 0 || uint64(n) > math.MaxUint32 { return math.MaxUint32 } return uint32(n) diff --git a/common/types/size_calc.go b/common/types/size_calc.go new file mode 100644 index 000000000..b3689acbd --- /dev/null +++ b/common/types/size_calc.go @@ -0,0 +1,396 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package types + +import ( + "math" + "reflect" + "time" + "unicode/utf8" + + "google.golang.org/protobuf/proto" + "google.golang.org/protobuf/reflect/protoreflect" + + "github.com/google/cel-go/common/types/ref" + "github.com/google/cel-go/common/types/traits" +) + +const ( + defaultSizeCalculatorMaxDepth = 5 + defaultSizeCalculatorMaxTraversal = 10000 +) + +// SizeCalculatorOption configures a SizeCalculator instance. +type SizeCalculatorOption func(*SizeCalculator) + +// SizeCalculatorMaxDepth sets the maximum object depth limit before saturating to math.MaxUint32. +func SizeCalculatorMaxDepth(depth int) SizeCalculatorOption { + return func(s *SizeCalculator) { + s.maxDepth = depth + } +} + +// SizeCalculatorMaxTraversal sets the maximum object traversal limit before saturating to math.MaxUint32. +func SizeCalculatorMaxTraversal(traversal int) SizeCalculatorOption { + return func(s *SizeCalculator) { + s.maxTraversal = traversal + } +} + +// SizeCalculator calculates the recursive element size of values. +type SizeCalculator struct { + version int + maxDepth int + maxTraversal int +} + +// NewSizeCalculator returns a new SizeCalculator configured with optional SizeCalculatorOption settings. +func NewSizeCalculator(opts ...SizeCalculatorOption) *SizeCalculator { + s := &SizeCalculator{ + version: 0, + maxDepth: defaultSizeCalculatorMaxDepth, + maxTraversal: defaultSizeCalculatorMaxTraversal, + } + for _, opt := range opts { + opt(s) + } + return s +} + +// Version returns the calculation version. +func (s *SizeCalculator) Version() int { + return s.version +} + +type sizeContext struct { + calc *SizeCalculator + depth int + traversalCount *int +} + +func (c *sizeContext) childContext() *sizeContext { + return &sizeContext{ + calc: c.calc, + depth: c.depth + 1, + traversalCount: c.traversalCount, + } +} + +func (c *sizeContext) visitNode() bool { + *c.traversalCount++ + if *c.traversalCount > c.calc.maxTraversal || c.depth > c.calc.maxDepth { + return false + } + return true +} + +// AggregateSize returns the size of the input value, if known. +// Otherwise, a unit size of 1 is returned. +func (s *SizeCalculator) AggregateSize(val any) uint32 { + traversals := 0 + ctx := &sizeContext{ + calc: s, + depth: 1, + traversalCount: &traversals, + } + return ctx.AggregateSize(val) +} + +// AggregateSize implements the ref.Val interface and allows for the generation of nested +// child context values which are necessary for correct traversal count tracking. +func (c *sizeContext) AggregateSize(val any) uint32 { + if !c.visitNode() { + return math.MaxUint32 + } + switch v := val.(type) { + case AggregateSizeVisitor: + return v.AggregateSize(c.childContext()) + case traits.Mapper: + total := uint32(1) + it := v.Iterator() + childCtx := c.childContext() + for it.HasNext() == True { + key := it.Next() + val, _ := v.Find(key) + total = safeAddUint32(total, childCtx.AggregateSize(key)) + total = safeAddUint32(total, childCtx.AggregateSize(val)) + } + return total + case traits.Lister: + total := uint32(1) + it := v.Iterator() + childCtx := c.childContext() + for it.HasNext() == True { + total = safeAddUint32(total, childCtx.AggregateSize(it.Next())) + } + return total + case String: + return safeUint32FromInt(utf8.RuneCountInString(string(v))) + case Bytes: + return safeUint32FromInt(len(v)) + case traits.Sizer: + return safeUint32FromBoxedInt(v.Size().(Int)) + case Bool, Int, Uint, Double, Duration, Timestamp, Null, *Type, *Err, *Unknown: + return 1 + case ref.Val: + return c.AggregateSize(v.Value()) + case protoreflect.Value: + return getProtoValueAggregateSize(c, v) + case protoreflect.MapKey: + return getProtoValueAggregateSize(c, v.Value()) + case protoreflect.Message: + return getProtoMessageAggregateSize(c, v) + case protoreflect.List: + return getProtoListAggregateSize(c, v) + case protoreflect.Map: + return getProtoMapAggregateSize(c, v) + case proto.Message: + if v == nil { + return 0 + } + return getProtoMessageAggregateSize(c, v.ProtoReflect()) + case reflect.Value: + return getReflectValueAggregateSize(c, v) + case string: + return safeUint32FromInt(utf8.RuneCountInString(v)) + case []byte: + return safeUint32FromInt(len(v)) + case int, int8, int16, int32, int64, + uint, uint8, uint16, uint32, uint64, + float32, float64, bool, time.Time, time.Duration, nil: + return 1 + default: + return getReflectValueAggregateSize(c, reflect.ValueOf(val)) + } +} + +func getProtoValueAggregateSize(c *sizeContext, v protoreflect.Value) uint32 { + switch val := v.Interface().(type) { + case string: + return c.AggregateSize(val) + case []byte: + return c.AggregateSize(val) + case protoreflect.Message: + return getProtoMessageAggregateSize(c.childContext(), val) + case protoreflect.List: + return getProtoListAggregateSize(c.childContext(), val) + case protoreflect.Map: + return getProtoMapAggregateSize(c.childContext(), val) + default: + if !c.visitNode() { + return math.MaxUint32 + } + return 1 + } +} + +func getProtoFieldAggregateSize(c *sizeContext, fd protoreflect.FieldDescriptor, v protoreflect.Value) uint32 { + if !c.visitNode() { + return math.MaxUint32 + } + childCtx := c.childContext() + if fd.IsMap() { + m := v.Map() + total := uint32(1) + valKind := fd.MapValue().Kind() + m.Range(func(k protoreflect.MapKey, val protoreflect.Value) bool { + total = safeAddUint32(total, getProtoMapKeyAggregateSize(childCtx, fd.MapKey(), k)) + total = safeAddUint32(total, getProtoFieldByKindAggregateSize(childCtx, valKind, val)) + return true + }) + return total + } + if fd.IsList() { + l := v.List() + elemKind := fd.Kind() + total := safeAddUint32(1, safeUint32FromInt(l.Len())) + switch elemKind { + case protoreflect.StringKind: + total = 1 + for i := 0; i < l.Len(); i++ { + total = safeAddUint32(total, childCtx.AggregateSize(l.Get(i).String())) + } + case protoreflect.BytesKind: + total = 1 + for i := 0; i < l.Len(); i++ { + total = safeAddUint32(total, childCtx.AggregateSize(l.Get(i).Bytes())) + } + case protoreflect.MessageKind, protoreflect.GroupKind: + total = 1 + for i := 0; i < l.Len(); i++ { + total = safeAddUint32(total, getProtoMessageAggregateSize(childCtx, l.Get(i).Message())) + } + } + return total + } + return getProtoFieldByKindAggregateSize(childCtx, fd.Kind(), v) +} + +func getProtoMapKeyAggregateSize(c *sizeContext, keyDesc protoreflect.FieldDescriptor, k protoreflect.MapKey) uint32 { + if keyDesc.Kind() == protoreflect.StringKind { + return c.AggregateSize(k.String()) + } + if !c.visitNode() { + return math.MaxUint32 + } + return 1 +} + +func getProtoFieldByKindAggregateSize(c *sizeContext, kind protoreflect.Kind, v protoreflect.Value) uint32 { + switch kind { + case protoreflect.StringKind: + return c.AggregateSize(v.String()) + case protoreflect.BytesKind: + return c.AggregateSize(v.Bytes()) + case protoreflect.MessageKind, protoreflect.GroupKind: + return getProtoMessageAggregateSize(c, v.Message()) + default: + if !c.visitNode() { + return math.MaxUint32 + } + return 1 + } +} + +func getProtoMessageAggregateSize(c *sizeContext, m protoreflect.Message) uint32 { + if !m.IsValid() { + return 0 + } + if !c.visitNode() { + return math.MaxUint32 + } + childCtx := c.childContext() + total := uint32(1) + m.Range(func(fd protoreflect.FieldDescriptor, v protoreflect.Value) bool { + total = safeAddUint32(total, getProtoFieldAggregateSize(childCtx, fd, v)) + return true + }) + return total +} + +func getProtoListAggregateSize(c *sizeContext, l protoreflect.List) uint32 { + if !l.IsValid() { + return 0 + } + if !c.visitNode() { + return math.MaxUint32 + } + childCtx := c.childContext() + total := uint32(1) + for i := 0; i < l.Len(); i++ { + total = safeAddUint32(total, getProtoValueAggregateSize(childCtx, l.Get(i))) + } + return total +} + +func getProtoMapAggregateSize(c *sizeContext, m protoreflect.Map) uint32 { + if !m.IsValid() { + return 0 + } + if !c.visitNode() { + return math.MaxUint32 + } + childCtx := c.childContext() + total := uint32(1) + m.Range(func(k protoreflect.MapKey, v protoreflect.Value) bool { + total = safeAddUint32(total, getProtoValueAggregateSize(childCtx, k.Value())) + total = safeAddUint32(total, getProtoValueAggregateSize(childCtx, v)) + return true + }) + return total +} + +func getReflectValueAggregateSize(c *sizeContext, fieldVal reflect.Value) uint32 { + if !fieldVal.IsValid() { + return 0 + } + if !c.visitNode() { + return math.MaxUint32 + } + childCtx := c.childContext() + switch fieldVal.Kind() { + case reflect.String: + return safeUint32FromInt(utf8.RuneCountInString(fieldVal.String())) + case reflect.Slice, reflect.Array: + elemType := fieldVal.Type().Elem() + if elemType.Kind() == reflect.Uint8 { + return safeUint32FromInt(fieldVal.Len()) + } + total := safeAddUint32(1, safeUint32FromInt(fieldVal.Len())) + switch elemType.Kind() { + case reflect.String: + total = 1 + for i := 0; i < fieldVal.Len(); i++ { + total = safeAddUint32(total, childCtx.AggregateSize(fieldVal.Index(i).String())) + } + case reflect.Struct, reflect.Pointer, reflect.Slice, reflect.Array, reflect.Map, reflect.Interface: + total = 1 + for i := 0; i < fieldVal.Len(); i++ { + total = safeAddUint32(total, getReflectValueAggregateSize(childCtx, fieldVal.Index(i))) + } + } + return total + case reflect.Map: + total := uint32(1) + iter := fieldVal.MapRange() + for iter.Next() { + total = safeAddUint32(total, getReflectValueAggregateSize(childCtx, iter.Key())) + total = safeAddUint32(total, getReflectValueAggregateSize(childCtx, iter.Value())) + } + return total + case reflect.Pointer, reflect.Interface: + if fieldVal.IsNil() { + return 0 + } + if fieldVal.CanInterface() { + if sizer, ok := fieldVal.Interface().(AggregateSizeVisitor); ok { + return sizer.AggregateSize(childCtx) + } + if sizer, ok := fieldVal.Interface().(traits.Sizer); ok { + return safeUint32FromBoxedInt(sizer.Size().(Int)) + } + } + return getReflectValueAggregateSize(c, fieldVal.Elem()) + case reflect.Struct: + if fieldVal.Type() == timestampType || fieldVal.Type() == durationType { + return 1 + } + if fieldVal.CanInterface() { + if sizer, ok := fieldVal.Interface().(AggregateSizeVisitor); ok { + return sizer.AggregateSize(childCtx) + } + if sizer, ok := fieldVal.Interface().(traits.Sizer); ok { + return safeUint32FromBoxedInt(sizer.Size().(Int)) + } + } + total := uint32(1) + t := fieldVal.Type() + numFields := fieldVal.NumField() + for i := 0; i < numFields; i++ { + f := t.Field(i) + if !f.IsExported() { + continue + } + fVal := fieldVal.Field(i) + if !fVal.IsValid() || fVal.IsZero() { + continue + } + total = safeAddUint32(total, getReflectValueAggregateSize(childCtx, fVal)) + } + return total + default: + return 1 + } +} diff --git a/common/types/size_calc_test.go b/common/types/size_calc_test.go new file mode 100644 index 000000000..d649387a5 --- /dev/null +++ b/common/types/size_calc_test.go @@ -0,0 +1,661 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package types + +import ( + "fmt" + "math" + "reflect" + "testing" + "time" + + "google.golang.org/protobuf/reflect/protoreflect" + + "github.com/google/cel-go/common/types/ref" + proto3pb "github.com/google/cel-go/test/proto3pb" +) + +func TestCalculateSize(t *testing.T) { + adapter := DefaultTypeAdapter + + tests := []struct { + name string + val any + want uint32 + }{ + { + name: "aggregate_sizer_list", + val: NewRefValList(adapter, []ref.Val{Int(1), Int(2)}), + want: 3, + }, + { + name: "sizer_string", + val: String("hello"), + want: 5, + }, + { + name: "sizer_bytes", + val: Bytes("world"), + want: 5, + }, + { + name: "err_val", + val: NewErr("test error"), + want: 1, + }, + { + name: "unknown_val", + val: &Unknown{}, + want: 1, + }, + { + name: "type_val", + val: IntType, + want: 1, + }, + { + name: "null_val", + val: NullValue, + want: 1, + }, + { + name: "scalar_ref_val_int", + val: Int(42), + want: 1, + }, + { + name: "scalar_ref_val_double", + val: Double(1.5), + want: 1, + }, + { + name: "scalar_ref_val_bool", + val: True, + want: 1, + }, + { + name: "scalar_ref_val_timestamp", + val: Timestamp{Time: time.Unix(100, 0)}, + want: 1, + }, + { + name: "scalar_ref_val_duration", + val: Duration{Duration: time.Second}, + want: 1, + }, + { + name: "proto_value_string", + val: protoreflect.ValueOfString("hello"), + want: 5, + }, + { + name: "proto_value_bytes", + val: protoreflect.ValueOfBytes([]byte("world")), + want: 5, + }, + { + name: "proto_value_int", + val: protoreflect.ValueOfInt32(42), + want: 1, + }, + { + name: "proto_map_key", + val: protoreflect.MapKey(protoreflect.ValueOfString("key")), + want: 3, + }, + { + name: "proto_message", + val: &proto3pb.TestAllTypes{SingleString: "hello"}, + want: 6, // 1 (root) + 5 (string) = 6 + }, + { + name: "proto_message_with_list_and_map", + val: &proto3pb.TestAllTypes{ + RepeatedString: []string{"a", "b"}, + MapStringString: map[string]string{"k": "v"}, + }, + want: 7, + }, + { + name: "protoreflect_message", + val: (&proto3pb.TestAllTypes{SingleInt64: 10}).ProtoReflect(), + want: 2, // 1 (root) + 1 (int64) = 2 + }, + { + name: "protoreflect_list", + val: (&proto3pb.TestAllTypes{RepeatedString: []string{"a", "b"}}).ProtoReflect().Get((&proto3pb.TestAllTypes{}).ProtoReflect().Descriptor().Fields().ByName("repeated_string")).List(), + want: 3, // 1 (container) + 1("a") + 1("b") = 3 + }, + { + name: "protoreflect_map", + val: (&proto3pb.TestAllTypes{MapStringString: map[string]string{"k": "v"}}).ProtoReflect().Get((&proto3pb.TestAllTypes{}).ProtoReflect().Descriptor().Fields().ByName("map_string_string")).Map(), + want: 3, // 1 (container) + 1("k") + 1("v") = 3 + }, + { + name: "nil_proto_message", + val: (*proto3pb.TestAllTypes)(nil), + want: 0, + }, + { + name: "reflect_value", + val: reflect.ValueOf("reflected"), + want: 9, + }, + { + name: "native_string", + val: "hello", + want: 5, + }, + { + name: "native_bytes", + val: []byte("world"), + want: 5, + }, + { + name: "native_int", + val: 42, + want: 1, + }, + { + name: "native_float", + val: 3.14, + want: 1, + }, + { + name: "native_bool", + val: true, + want: 1, + }, + { + name: "native_time", + val: time.Now(), + want: 1, + }, + { + name: "native_duration", + val: time.Hour, + want: 1, + }, + { + name: "native_nil", + val: nil, + want: 1, + }, + { + name: "custom_struct", + val: struct{ Name string }{"cel"}, + want: 4, // 1 (root) + 3 ("cel") = 4 + }, + { + name: "custom_lister", + val: proxyLegacyList{proxy: NewRefValList(DefaultTypeAdapter, []ref.Val{String("a"), String("b")})}, + want: 3, // 1 (container) + 1 ("a") + 1 ("b") = 3 + }, + { + name: "custom_mapper", + val: interopFoldableMap{Mapper: NewStringStringMap(DefaultTypeAdapter, map[string]string{"key": "val"})}, + want: 7, // 1 (container) + 3 ("key") + 3 ("val") = 7 + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + calculator := NewSizeCalculator() + if got := calculator.AggregateSize(tc.val); got != tc.want { + t.Errorf("AggregateSize(%v) got %d, want %d", tc.val, got, tc.want) + } + }) + } +} + +func TestNativeObjCalculateSizeNil(t *testing.T) { + nilNative := &nativeObj{} + if got := nilNative.AggregateSize(NewSizeCalculator()); got != 0 { + t.Errorf("nil nativeObj.AggregateSize() got %d, want 0", got) + } +} + +func TestSizeCalculatorOptions(t *testing.T) { + adapter := DefaultTypeAdapter + + var makeNestedList func(depth int) ref.Val + makeNestedList = func(depth int) ref.Val { + if depth <= 1 { + return NewRefValList(adapter, []ref.Val{Int(1)}) + } + return NewRefValList(adapter, []ref.Val{makeNestedList(depth - 1)}) + } + + t.Run("maxDepth default limit 5", func(t *testing.T) { + calc := NewSizeCalculator() + list5 := makeNestedList(4) + if got := calc.AggregateSize(list5); got == math.MaxUint32 { + t.Errorf("AggregateSize for depth 5 got MaxUint32, want calculated size") + } + + list6 := makeNestedList(5) + if got := calc.AggregateSize(list6); got != math.MaxUint32 { + t.Errorf("AggregateSize for depth 6 got %d, want MaxUint32", got) + } + }) + + t.Run("maxDepth custom option", func(t *testing.T) { + calc := NewSizeCalculator(SizeCalculatorMaxDepth(2)) + list2 := makeNestedList(1) + if got := calc.AggregateSize(list2); got == math.MaxUint32 { + t.Errorf("AggregateSize for depth 2 got MaxUint32, want calculated size") + } + + list3 := makeNestedList(2) + if got := calc.AggregateSize(list3); got != math.MaxUint32 { + t.Errorf("AggregateSize for depth 3 got %d, want MaxUint32", got) + } + }) + + t.Run("maxTraversal default limit 10000", func(t *testing.T) { + calc := NewSizeCalculator() + smallElems := make([]ref.Val, 100) + for i := 0; i < 100; i++ { + smallElems[i] = Int(i) + } + smallList := NewRefValList(adapter, smallElems) + if got := calc.AggregateSize(smallList); got == math.MaxUint32 { + t.Errorf("AggregateSize for 100 elements got MaxUint32, want calculated size") + } + + largeElems := make([]ref.Val, 10001) + for i := 0; i < 10001; i++ { + largeElems[i] = Int(i) + } + largeList := NewRefValList(adapter, largeElems) + if got := calc.AggregateSize(largeList); got != math.MaxUint32 { + t.Errorf("AggregateSize for 10001 elements got %d, want MaxUint32", got) + } + }) + + t.Run("maxTraversal custom option", func(t *testing.T) { + calc := NewSizeCalculator(SizeCalculatorMaxTraversal(5)) + list4 := NewRefValList(adapter, []ref.Val{Int(1), Int(2), Int(3), Int(4)}) + if got := calc.AggregateSize(list4); got == math.MaxUint32 { + t.Errorf("AggregateSize for 5 nodes got MaxUint32, want calculated size") + } + + list5 := NewRefValList(adapter, []ref.Val{Int(1), Int(2), Int(3), Int(4), Int(5)}) + if got := calc.AggregateSize(list5); got != math.MaxUint32 { + t.Errorf("AggregateSize for 6 nodes got %d, want MaxUint32", got) + } + }) + + t.Run("proto depth limit", func(t *testing.T) { + calc := NewSizeCalculator(SizeCalculatorMaxDepth(2)) + msg := &proto3pb.TestAllTypes{ + RepeatedNestedMessage: []*proto3pb.TestAllTypes_NestedMessage{ + {Bb: 42}, + }, + } + if got := calc.AggregateSize(msg); got != math.MaxUint32 { + t.Errorf("AggregateSize for proto nested msg got %d, want MaxUint32", got) + } + }) + + t.Run("cel map depth limit", func(t *testing.T) { + calc := NewSizeCalculator(SizeCalculatorMaxDepth(2)) + nestedMap := NewRefValMap(adapter, map[ref.Val]ref.Val{ + String("k"): NewRefValMap(adapter, map[ref.Val]ref.Val{ + String("subk"): String("subv"), + }), + }) + if got := calc.AggregateSize(nestedMap); got != math.MaxUint32 { + t.Errorf("AggregateSize for nested map got %d, want MaxUint32", got) + } + }) + + t.Run("native struct depth limit", func(t *testing.T) { + calc := NewSizeCalculator(SizeCalculatorMaxDepth(2)) + type Level3 struct{ Val string } + type Level2 struct{ L3 Level3 } + type Level1 struct{ L2 Level2 } + + obj := Level1{L2: Level2{L3: Level3{Val: "deep"}}} + if got := calc.AggregateSize(obj); got != math.MaxUint32 { + t.Errorf("AggregateSize for native struct depth > 2 got %d, want MaxUint32", got) + } + }) + + t.Run("native map depth limit", func(t *testing.T) { + calc := NewSizeCalculator(SizeCalculatorMaxDepth(2)) + m := map[string]map[string]string{ + "outer": {"inner": "val"}, + } + if got := calc.AggregateSize(m); got != math.MaxUint32 { + t.Errorf("AggregateSize for native nested map depth > 2 got %d, want MaxUint32", got) + } + }) +} + +type nestedNative struct { + NestedList []string + NestedMap map[string]int +} + +type rootNative struct { + Name string + Count int + Children []nestedNative +} + +func BenchmarkCalculateSizeAmortized(b *testing.B) { + adapter := DefaultTypeAdapter + + flatList := NewRefValList(adapter, []ref.Val{Int(1), Int(2), Int(3), Int(4), Int(5), Int(6), Int(7), Int(8)}) + nestedList := NewRefValList(adapter, []ref.Val{ + String("hello"), + NewRefValList(adapter, []ref.Val{String("nested1"), String("nested2")}), + NewRefValMap(adapter, map[ref.Val]ref.Val{String("k1"): String("v1")}), + }) + flatMap := NewRefValMap(adapter, map[ref.Val]ref.Val{ + String("k1"): Int(1), + String("k2"): Int(2), + String("k3"): Int(3), + }) + protoMsg := &proto3pb.TestAllTypes{ + SingleString: "hello world", + SingleInt64: 42, + RepeatedString: []string{"first", "second", "third"}, + MapStringString: map[string]string{ + "key1": "value1", + "key2": "value2", + }, + } + nativeData := &rootNative{ + Name: "parent", + Count: 100, + Children: []nestedNative{ + {NestedList: []string{"a", "b", "c"}, NestedMap: map[string]int{"k1": 1, "k2": 2}}, + {NestedList: []string{"d", "e"}, NestedMap: map[string]int{"k3": 3}}, + }, + } + nativeVal := adapter.NativeToValue(nativeData) + + benchmarks := []struct { + name string + val any + }{ + {name: "scalar_int", val: Int(42)}, + {name: "scalar_string", val: String("hello world this is a test string")}, + {name: "native_string", val: "hello world this is a test string"}, + {name: "native_bytes", val: []byte("hello world this is a test string")}, + {name: "list_flat", val: flatList}, + {name: "list_nested", val: nestedList}, + {name: "map_flat", val: flatMap}, + {name: "custom_list_flat", val: proxyLegacyList{proxy: flatList}}, + {name: "custom_list_nested", val: proxyLegacyList{proxy: nestedList}}, + {name: "custom_map_flat", val: interopFoldableMap{Mapper: flatMap}}, + {name: "proto_message", val: protoMsg}, + {name: "proto_obj", val: adapter.NativeToValue(protoMsg)}, + {name: "native_obj", val: nativeVal}, + {name: "native_struct", val: nativeData}, + } + + for _, bm := range benchmarks { + b.Run(bm.name, func(b *testing.B) { + b.ReportAllocs() + calculator := NewSizeCalculator() + for i := 0; i < b.N; i++ { + _ = calculator.AggregateSize(bm.val) + } + }) + } +} + +func BenchmarkCalculateSizeDynamic(b *testing.B) { + adapter := DefaultTypeAdapter + + benchmarks := []struct { + name string + valFn func() any + }{ + { + name: "list_flat", + valFn: func() any { + return NewRefValList(adapter, []ref.Val{Int(1), Int(2), Int(3), Int(4), Int(5), Int(6), Int(7), Int(8)}) + }, + }, + { + name: "list_nested", + valFn: func() any { + return NewRefValList(adapter, []ref.Val{ + String("hello"), + NewRefValList(adapter, []ref.Val{String("nested1"), String("nested2")}), + NewRefValMap(adapter, map[ref.Val]ref.Val{String("k1"): String("v1")}), + }) + }, + }, + { + name: "map_flat", + valFn: func() any { + return NewRefValMap(adapter, map[ref.Val]ref.Val{ + String("k1"): Int(1), + String("k2"): Int(2), + String("k3"): Int(3), + }) + }, + }, + { + name: "custom_list_flat", + valFn: func() any { + return proxyLegacyList{proxy: NewRefValList(adapter, []ref.Val{Int(1), Int(2), Int(3), Int(4), Int(5), Int(6), Int(7), Int(8)})} + }, + }, + { + name: "custom_list_nested", + valFn: func() any { + return proxyLegacyList{proxy: NewRefValList(adapter, []ref.Val{ + String("hello"), + NewRefValList(adapter, []ref.Val{String("nested1"), String("nested2")}), + NewRefValMap(adapter, map[ref.Val]ref.Val{String("k1"): String("v1")}), + })} + }, + }, + { + name: "custom_map_flat", + valFn: func() any { + return interopFoldableMap{Mapper: NewRefValMap(adapter, map[ref.Val]ref.Val{ + String("k1"): Int(1), + String("k2"): Int(2), + String("k3"): Int(3), + })} + }, + }, + { + name: "proto_obj", + valFn: func() any { + return adapter.NativeToValue(&proto3pb.TestAllTypes{ + SingleString: "hello world", + SingleInt64: 42, + RepeatedString: []string{"first", "second", "third"}, + MapStringString: map[string]string{ + "key1": "value1", + "key2": "value2", + }, + }) + }, + }, + { + name: "native_obj", + valFn: func() any { + return adapter.NativeToValue(&rootNative{ + Name: "parent", + Count: 100, + Children: []nestedNative{ + {NestedList: []string{"a", "b", "c"}, NestedMap: map[string]int{"k1": 1, "k2": 2}}, + {NestedList: []string{"d", "e"}, NestedMap: map[string]int{"k3": 3}}, + }, + }) + }, + }, + } + + for _, bm := range benchmarks { + b.Run(bm.name, func(b *testing.B) { + b.ReportAllocs() + calculator := NewSizeCalculator() + for i := 0; i < b.N; i++ { + val := bm.valFn() + _ = calculator.AggregateSize(val) + } + }) + } +} + +func BenchmarkCalculateSizeScaled(b *testing.B) { + adapter := DefaultTypeAdapter + sizes := []int{10, 100, 1000} + + for _, size := range sizes { + // Prepare list elements + listElems := make([]ref.Val, size) + for i := 0; i < size; i++ { + listElems[i] = Int(i) + } + builtinList := NewRefValList(adapter, listElems) + customList := proxyLegacyList{proxy: builtinList} + + // Prepare map entries + mapEntries := make(map[ref.Val]ref.Val, size) + for i := 0; i < size; i++ { + mapEntries[String(fmt.Sprintf("k%d", i))] = Int(i) + } + builtinMap := NewRefValMap(adapter, mapEntries) + customMap := interopFoldableMap{Mapper: builtinMap} + + // Amortized (repeated calculation on memoized vs unmemoized custom instance) + b.Run(fmt.Sprintf("Amortized/builtin_list/N=%d", size), func(b *testing.B) { + b.ReportAllocs() + calc := NewSizeCalculator() + for i := 0; i < b.N; i++ { + _ = calc.AggregateSize(builtinList) + } + }) + b.Run(fmt.Sprintf("Amortized/custom_list/N=%d", size), func(b *testing.B) { + b.ReportAllocs() + calc := NewSizeCalculator() + for i := 0; i < b.N; i++ { + _ = calc.AggregateSize(customList) + } + }) + b.Run(fmt.Sprintf("Amortized/builtin_map/N=%d", size), func(b *testing.B) { + b.ReportAllocs() + calc := NewSizeCalculator() + for i := 0; i < b.N; i++ { + _ = calc.AggregateSize(builtinMap) + } + }) + b.Run(fmt.Sprintf("Amortized/custom_map/N=%d", size), func(b *testing.B) { + b.ReportAllocs() + calc := NewSizeCalculator() + for i := 0; i < b.N; i++ { + _ = calc.AggregateSize(customMap) + } + }) + + // First-time / Uncached calculation + b.Run(fmt.Sprintf("FirstTime/builtin_list/N=%d", size), func(b *testing.B) { + b.ReportAllocs() + calc := NewSizeCalculator() + for i := 0; i < b.N; i++ { + l := NewRefValList(adapter, listElems) + _ = calc.AggregateSize(l) + } + }) + b.Run(fmt.Sprintf("FirstTime/custom_list/N=%d", size), func(b *testing.B) { + b.ReportAllocs() + calc := NewSizeCalculator() + for i := 0; i < b.N; i++ { + l := proxyLegacyList{proxy: NewRefValList(adapter, listElems)} + _ = calc.AggregateSize(l) + } + }) + b.Run(fmt.Sprintf("FirstTime/builtin_map/N=%d", size), func(b *testing.B) { + b.ReportAllocs() + calc := NewSizeCalculator() + for i := 0; i < b.N; i++ { + m := NewRefValMap(adapter, mapEntries) + _ = calc.AggregateSize(m) + } + }) + b.Run(fmt.Sprintf("FirstTime/custom_map/N=%d", size), func(b *testing.B) { + b.ReportAllocs() + calc := NewSizeCalculator() + for i := 0; i < b.N; i++ { + m := interopFoldableMap{Mapper: NewRefValMap(adapter, mapEntries)} + _ = calc.AggregateSize(m) + } + }) + } + + // Benchmark nested tree complexity (Depth x Width) + depths := []int{2, 3} + width := 10 + for _, depth := range depths { + builtinNested := createNestedList(adapter, depth, width) + customNested := createNestedCustomList(adapter, depth, width) + + b.Run(fmt.Sprintf("Complexity/builtin_nested_list/Depth=%d_Width=%d", depth, width), func(b *testing.B) { + b.ReportAllocs() + calc := NewSizeCalculator() + for i := 0; i < b.N; i++ { + _ = calc.AggregateSize(builtinNested) + } + }) + b.Run(fmt.Sprintf("Complexity/custom_nested_list/Depth=%d_Width=%d", depth, width), func(b *testing.B) { + b.ReportAllocs() + calc := NewSizeCalculator() + for i := 0; i < b.N; i++ { + _ = calc.AggregateSize(customNested) + } + }) + } +} + +func createNestedList(adapter Adapter, depth, width int) ref.Val { + if depth <= 1 { + elems := make([]ref.Val, width) + for i := 0; i < width; i++ { + elems[i] = Int(i) + } + return NewRefValList(adapter, elems) + } + elems := make([]ref.Val, width) + for i := 0; i < width; i++ { + elems[i] = createNestedList(adapter, depth-1, width) + } + return NewRefValList(adapter, elems) +} + +func createNestedCustomList(adapter Adapter, depth, width int) ref.Val { + if depth <= 1 { + elems := make([]ref.Val, width) + for i := 0; i < width; i++ { + elems[i] = Int(i) + } + return proxyLegacyList{proxy: NewRefValList(adapter, elems)} + } + elems := make([]ref.Val, width) + for i := 0; i < width; i++ { + elems[i] = createNestedCustomList(adapter, depth-1, width) + } + return proxyLegacyList{proxy: NewRefValList(adapter, elems)} +} diff --git a/common/types/util.go b/common/types/util.go index c43e773cd..71662eee3 100644 --- a/common/types/util.go +++ b/common/types/util.go @@ -15,15 +15,7 @@ package types import ( - "reflect" - "time" - "unicode/utf8" - - "google.golang.org/protobuf/proto" - "google.golang.org/protobuf/reflect/protoreflect" - "github.com/google/cel-go/common/types/ref" - "github.com/google/cel-go/common/types/traits" ) // IsUnknownOrError returns whether the input element ref.Val is an ErrType or UnknownType. @@ -54,270 +46,3 @@ func Equal(lhs ref.Val, rhs ref.Val) ref.Val { } return lhs.Equal(rhs) } - -// SizeCalculator calculates the recursive element size of values. -type SizeCalculator struct { - version int -} - -// NewSizeCalculator returns a new SizeCalculator for the specified version. -func NewSizeCalculator() *SizeCalculator { - return &SizeCalculator{version: 0} -} - -// Version returns the calculation version. -func (s *SizeCalculator) Version() int { - return s.version -} - -// AggregateSize returns the size of the input value, if known. -// Otherwise, a unit size of 1 is returned. -func (s *SizeCalculator) AggregateSize(val any) uint32 { - switch v := val.(type) { - case AggregateSizeVisitor: - return v.AggregateSize(s) - case traits.Mapper: - total := uint32(1) - it := v.Iterator() - for it.HasNext() == True { - key := it.Next() - val, _ := v.Find(key) - total = safeAddUint32(total, s.AggregateSize(key)) - total = safeAddUint32(total, s.AggregateSize(val)) - } - return total - case traits.Lister: - total := uint32(1) - it := v.Iterator() - for it.HasNext() == True { - total = safeAddUint32(total, s.AggregateSize(it.Next())) - } - return total - case String: - return safeUint32FromInt(utf8.RuneCountInString(string(v))) - case Bytes: - return safeUint32FromInt(len(v)) - case traits.Sizer: - return safeUint32FromBoxedInt(v.Size().(Int)) - case Bool, Int, Uint, Double, Duration, Timestamp, Null, *Type, *Err, *Unknown: - return 1 - case ref.Val: - return s.AggregateSize(v.Value()) - case protoreflect.Value: - return getProtoValueAggregateSize(s, v) - case protoreflect.MapKey: - return getProtoValueAggregateSize(s, v.Value()) - case protoreflect.Message: - return getProtoMessageAggregateSize(s, v) - case protoreflect.List: - return getProtoListAggregateSize(s, v) - case protoreflect.Map: - return getProtoMapAggregateSize(s, v) - case proto.Message: - if v == nil { - return 0 - } - return getProtoMessageAggregateSize(s, v.ProtoReflect()) - case reflect.Value: - return getReflectValueAggregateSize(s, v) - case string: - return safeUint32FromInt(utf8.RuneCountInString(v)) - case []byte: - return safeUint32FromInt(len(v)) - case int, int8, int16, int32, int64, - uint, uint8, uint16, uint32, uint64, - float32, float64, bool, time.Time, time.Duration, nil: - return 1 - default: - return getReflectValueAggregateSize(s, reflect.ValueOf(val)) - } -} - -func getProtoValueAggregateSize(s AggregateSizer, v protoreflect.Value) uint32 { - switch val := v.Interface().(type) { - case string: - return safeUint32FromInt(utf8.RuneCountInString(val)) - case []byte: - return safeUint32FromInt(len(val)) - case protoreflect.Message: - return getProtoMessageAggregateSize(s, val) - case protoreflect.List: - return getProtoListAggregateSize(s, val) - case protoreflect.Map: - return getProtoMapAggregateSize(s, val) - default: - return 1 - } -} - -func getProtoFieldAggregateSize(s AggregateSizer, fd protoreflect.FieldDescriptor, v protoreflect.Value) uint32 { - if fd.IsMap() { - m := v.Map() - total := uint32(1) - valKind := fd.MapValue().Kind() - m.Range(func(k protoreflect.MapKey, val protoreflect.Value) bool { - total = safeAddUint32(total, getProtoMapKeyAggregateSize(fd.MapKey(), k)) - total = safeAddUint32(total, getProtoFieldByKindAggregateSize(s, valKind, val)) - return true - }) - return total - } - if fd.IsList() { - l := v.List() - elemKind := fd.Kind() - total := safeAddUint32(1, safeUint32FromInt(l.Len())) - if elemKind == protoreflect.StringKind { - total = 1 - for i := 0; i < l.Len(); i++ { - total = safeAddUint32(total, safeUint32FromInt(utf8.RuneCountInString(l.Get(i).String()))) - } - } else if elemKind == protoreflect.BytesKind { - total = 1 - for i := 0; i < l.Len(); i++ { - total = safeAddUint32(total, safeUint32FromInt(len(l.Get(i).Bytes()))) - } - } else if elemKind == protoreflect.MessageKind || elemKind == protoreflect.GroupKind { - total = 1 - for i := 0; i < l.Len(); i++ { - total = safeAddUint32(total, getProtoMessageAggregateSize(s, l.Get(i).Message())) - } - } - return total - } - return getProtoFieldByKindAggregateSize(s, fd.Kind(), v) -} - -func getProtoMapKeyAggregateSize(keyDesc protoreflect.FieldDescriptor, k protoreflect.MapKey) uint32 { - if keyDesc.Kind() == protoreflect.StringKind { - return safeUint32FromInt(utf8.RuneCountInString(k.String())) - } - return 1 -} - -func getProtoFieldByKindAggregateSize(s AggregateSizer, kind protoreflect.Kind, v protoreflect.Value) uint32 { - switch kind { - case protoreflect.StringKind: - return safeUint32FromInt(utf8.RuneCountInString(v.String())) - case protoreflect.BytesKind: - return safeUint32FromInt(len(v.Bytes())) - case protoreflect.MessageKind, protoreflect.GroupKind: - return getProtoMessageAggregateSize(s, v.Message()) - default: - return 1 - } -} - -func getProtoMessageAggregateSize(s AggregateSizer, m protoreflect.Message) uint32 { - if !m.IsValid() { - return 0 - } - total := uint32(1) - m.Range(func(fd protoreflect.FieldDescriptor, v protoreflect.Value) bool { - total = safeAddUint32(total, getProtoFieldAggregateSize(s, fd, v)) - return true - }) - return total -} - -func getProtoListAggregateSize(s AggregateSizer, l protoreflect.List) uint32 { - if !l.IsValid() { - return 0 - } - total := uint32(1) - for i := 0; i < l.Len(); i++ { - total = safeAddUint32(total, getProtoValueAggregateSize(s, l.Get(i))) - } - return total -} - -func getProtoMapAggregateSize(s AggregateSizer, m protoreflect.Map) uint32 { - if !m.IsValid() { - return 0 - } - total := uint32(1) - m.Range(func(k protoreflect.MapKey, v protoreflect.Value) bool { - total = safeAddUint32(total, getProtoValueAggregateSize(s, k.Value())) - total = safeAddUint32(total, getProtoValueAggregateSize(s, v)) - return true - }) - return total -} - -func getReflectValueAggregateSize(s AggregateSizer, fieldVal reflect.Value) uint32 { - if !fieldVal.IsValid() { - return 0 - } - switch fieldVal.Kind() { - case reflect.String: - return safeUint32FromInt(utf8.RuneCountInString(fieldVal.String())) - case reflect.Slice, reflect.Array: - elemType := fieldVal.Type().Elem() - if elemType.Kind() == reflect.Uint8 { - return safeUint32FromInt(fieldVal.Len()) - } - total := safeAddUint32(1, safeUint32FromInt(fieldVal.Len())) - switch elemType.Kind() { - case reflect.String: - total = 1 - for i := 0; i < fieldVal.Len(); i++ { - total = safeAddUint32(total, safeUint32FromInt(utf8.RuneCountInString(fieldVal.Index(i).String()))) - } - case reflect.Struct, reflect.Pointer, reflect.Slice, reflect.Array, reflect.Map, reflect.Interface: - total = 1 - for i := 0; i < fieldVal.Len(); i++ { - total = safeAddUint32(total, getReflectValueAggregateSize(s, fieldVal.Index(i))) - } - } - return total - case reflect.Map: - total := uint32(1) - iter := fieldVal.MapRange() - for iter.Next() { - total = safeAddUint32(total, getReflectValueAggregateSize(s, iter.Key())) - total = safeAddUint32(total, getReflectValueAggregateSize(s, iter.Value())) - } - return total - case reflect.Pointer, reflect.Interface: - if fieldVal.IsNil() { - return 0 - } - if fieldVal.CanInterface() { - if sizer, ok := fieldVal.Interface().(AggregateSizeVisitor); ok { - return sizer.AggregateSize(s) - } - if sizer, ok := fieldVal.Interface().(traits.Sizer); ok { - return safeUint32FromBoxedInt(sizer.Size().(Int)) - } - } - return getReflectValueAggregateSize(s, fieldVal.Elem()) - case reflect.Struct: - if fieldVal.Type() == timestampType || fieldVal.Type() == durationType { - return 1 - } - if fieldVal.CanInterface() { - if sizer, ok := fieldVal.Interface().(AggregateSizeVisitor); ok { - return sizer.AggregateSize(s) - } - if sizer, ok := fieldVal.Interface().(traits.Sizer); ok { - return safeUint32FromBoxedInt(sizer.Size().(Int)) - } - } - total := uint32(1) - t := fieldVal.Type() - numFields := fieldVal.NumField() - for i := 0; i < numFields; i++ { - f := t.Field(i) - if !f.IsExported() { - continue - } - fVal := fieldVal.Field(i) - if !fVal.IsValid() || fVal.IsZero() { - continue - } - total = safeAddUint32(total, getReflectValueAggregateSize(s, fVal)) - } - return total - default: - return 1 - } -} diff --git a/common/types/util_test.go b/common/types/util_test.go index 96b2fdf1a..4d16495ee 100644 --- a/common/types/util_test.go +++ b/common/types/util_test.go @@ -15,211 +15,10 @@ package types import ( - "fmt" "math" - "reflect" "testing" - "time" - - "google.golang.org/protobuf/reflect/protoreflect" - - "github.com/google/cel-go/common/types/ref" - proto3pb "github.com/google/cel-go/test/proto3pb" ) -func TestCalculateSize(t *testing.T) { - adapter := DefaultTypeAdapter - - tests := []struct { - name string - val any - want uint32 - }{ - { - name: "aggregate_sizer_list", - val: NewRefValList(adapter, []ref.Val{Int(1), Int(2)}), - want: 3, - }, - { - name: "sizer_string", - val: String("hello"), - want: 5, - }, - { - name: "sizer_bytes", - val: Bytes("world"), - want: 5, - }, - { - name: "err_val", - val: NewErr("test error"), - want: 1, - }, - { - name: "unknown_val", - val: &Unknown{}, - want: 1, - }, - { - name: "type_val", - val: IntType, - want: 1, - }, - { - name: "null_val", - val: NullValue, - want: 1, - }, - { - name: "scalar_ref_val_int", - val: Int(42), - want: 1, - }, - { - name: "scalar_ref_val_double", - val: Double(1.5), - want: 1, - }, - { - name: "scalar_ref_val_bool", - val: True, - want: 1, - }, - { - name: "scalar_ref_val_timestamp", - val: Timestamp{Time: time.Unix(100, 0)}, - want: 1, - }, - { - name: "scalar_ref_val_duration", - val: Duration{Duration: time.Second}, - want: 1, - }, - { - name: "proto_value_string", - val: protoreflect.ValueOfString("hello"), - want: 5, - }, - { - name: "proto_value_bytes", - val: protoreflect.ValueOfBytes([]byte("world")), - want: 5, - }, - { - name: "proto_value_int", - val: protoreflect.ValueOfInt32(42), - want: 1, - }, - { - name: "proto_map_key", - val: protoreflect.MapKey(protoreflect.ValueOfString("key")), - want: 3, - }, - { - name: "proto_message", - val: &proto3pb.TestAllTypes{SingleString: "hello"}, - want: 6, // 1 (root) + 5 (string) = 6 - }, - { - name: "proto_message_with_list_and_map", - val: &proto3pb.TestAllTypes{ - RepeatedString: []string{"a", "b"}, - MapStringString: map[string]string{"k": "v"}, - }, - want: 7, - }, - { - name: "protoreflect_message", - val: (&proto3pb.TestAllTypes{SingleInt64: 10}).ProtoReflect(), - want: 2, // 1 (root) + 1 (int64) = 2 - }, - { - name: "protoreflect_list", - val: (&proto3pb.TestAllTypes{RepeatedString: []string{"a", "b"}}).ProtoReflect().Get((&proto3pb.TestAllTypes{}).ProtoReflect().Descriptor().Fields().ByName("repeated_string")).List(), - want: 3, // 1 (container) + 1("a") + 1("b") = 3 - }, - { - name: "protoreflect_map", - val: (&proto3pb.TestAllTypes{MapStringString: map[string]string{"k": "v"}}).ProtoReflect().Get((&proto3pb.TestAllTypes{}).ProtoReflect().Descriptor().Fields().ByName("map_string_string")).Map(), - want: 3, // 1 (container) + 1("k") + 1("v") = 3 - }, - { - name: "nil_proto_message", - val: (*proto3pb.TestAllTypes)(nil), - want: 0, - }, - { - name: "reflect_value", - val: reflect.ValueOf("reflected"), - want: 9, - }, - { - name: "native_string", - val: "hello", - want: 5, - }, - { - name: "native_bytes", - val: []byte("world"), - want: 5, - }, - { - name: "native_int", - val: 42, - want: 1, - }, - { - name: "native_float", - val: 3.14, - want: 1, - }, - { - name: "native_bool", - val: true, - want: 1, - }, - { - name: "native_time", - val: time.Now(), - want: 1, - }, - { - name: "native_duration", - val: time.Hour, - want: 1, - }, - { - name: "native_nil", - val: nil, - want: 1, - }, - { - name: "custom_struct", - val: struct{ Name string }{"cel"}, - want: 4, // 1 (root) + 3 ("cel") = 4 - }, - { - name: "custom_lister", - val: proxyLegacyList{proxy: NewRefValList(DefaultTypeAdapter, []ref.Val{String("a"), String("b")})}, - want: 3, // 1 (container) + 1 ("a") + 1 ("b") = 3 - }, - { - name: "custom_mapper", - val: interopFoldableMap{Mapper: NewStringStringMap(DefaultTypeAdapter, map[string]string{"key": "val"})}, - want: 7, // 1 (container) + 3 ("key") + 3 ("val") = 7 - }, - } - - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - calculator := NewSizeCalculator() - if got := calculator.AggregateSize(tc.val); got != tc.want { - t.Errorf("AggregateSize(%v) got %d, want %d", tc.val, got, tc.want) - } - }) - } -} - func TestSafeUint32Helpers(t *testing.T) { // safeAddUint32 if got := safeAddUint32(10, 20); got != 30 { @@ -233,8 +32,8 @@ func TestSafeUint32Helpers(t *testing.T) { if got := safeUint32FromInt(42); got != 42 { t.Errorf("safeUint32FromInt(42) got %d, want 42", got) } - if got := safeUint32FromInt(-1); got != 0 { - t.Errorf("safeUint32FromInt(-1) got %d, want 0", got) + if got := safeUint32FromInt(-1); got != math.MaxUint32 { + t.Errorf("safeUint32FromInt(-1) got %d, want MaxUint32", got) } if got := safeUint32FromInt(int(uint64(math.MaxUint32) + 100)); got != math.MaxUint32 { t.Errorf("safeUint32FromInt(overflow) got %d, want MaxUint32", got) @@ -252,13 +51,6 @@ func TestSafeUint32Helpers(t *testing.T) { } } -func TestNativeObjCalculateSizeNil(t *testing.T) { - nilNative := &nativeObj{} - if got := nilNative.AggregateSize(NewSizeCalculator()); got != 0 { - t.Errorf("nil nativeObj.AggregateSize() got %d, want 0", got) - } -} - func BenchmarkIsUnknownOrError(b *testing.B) { err := NewErr("test") unk := &Unknown{} @@ -268,318 +60,3 @@ func BenchmarkIsUnknownOrError(b *testing.B) { } } } - -type nestedNative struct { - NestedList []string - NestedMap map[string]int -} - -type rootNative struct { - Name string - Count int - Children []nestedNative -} - -func BenchmarkCalculateSizeAmortized(b *testing.B) { - adapter := DefaultTypeAdapter - - flatList := NewRefValList(adapter, []ref.Val{Int(1), Int(2), Int(3), Int(4), Int(5), Int(6), Int(7), Int(8)}) - nestedList := NewRefValList(adapter, []ref.Val{ - String("hello"), - NewRefValList(adapter, []ref.Val{String("nested1"), String("nested2")}), - NewRefValMap(adapter, map[ref.Val]ref.Val{String("k1"): String("v1")}), - }) - flatMap := NewRefValMap(adapter, map[ref.Val]ref.Val{ - String("k1"): Int(1), - String("k2"): Int(2), - String("k3"): Int(3), - }) - protoMsg := &proto3pb.TestAllTypes{ - SingleString: "hello world", - SingleInt64: 42, - RepeatedString: []string{"first", "second", "third"}, - MapStringString: map[string]string{ - "key1": "value1", - "key2": "value2", - }, - } - nativeData := &rootNative{ - Name: "parent", - Count: 100, - Children: []nestedNative{ - {NestedList: []string{"a", "b", "c"}, NestedMap: map[string]int{"k1": 1, "k2": 2}}, - {NestedList: []string{"d", "e"}, NestedMap: map[string]int{"k3": 3}}, - }, - } - nativeVal := adapter.NativeToValue(nativeData) - - benchmarks := []struct { - name string - val any - }{ - {name: "scalar_int", val: Int(42)}, - {name: "scalar_string", val: String("hello world this is a test string")}, - {name: "native_string", val: "hello world this is a test string"}, - {name: "native_bytes", val: []byte("hello world this is a test string")}, - {name: "list_flat", val: flatList}, - {name: "list_nested", val: nestedList}, - {name: "map_flat", val: flatMap}, - {name: "custom_list_flat", val: proxyLegacyList{proxy: flatList}}, - {name: "custom_list_nested", val: proxyLegacyList{proxy: nestedList}}, - {name: "custom_map_flat", val: interopFoldableMap{Mapper: flatMap}}, - {name: "proto_message", val: protoMsg}, - {name: "proto_obj", val: adapter.NativeToValue(protoMsg)}, - {name: "native_obj", val: nativeVal}, - {name: "native_struct", val: nativeData}, - } - - for _, bm := range benchmarks { - b.Run(bm.name, func(b *testing.B) { - b.ReportAllocs() - calculator := NewSizeCalculator() - for i := 0; i < b.N; i++ { - _ = calculator.AggregateSize(bm.val) - } - }) - } -} - -func BenchmarkCalculateSizeDynamic(b *testing.B) { - adapter := DefaultTypeAdapter - - benchmarks := []struct { - name string - valFn func() any - }{ - { - name: "list_flat", - valFn: func() any { - return NewRefValList(adapter, []ref.Val{Int(1), Int(2), Int(3), Int(4), Int(5), Int(6), Int(7), Int(8)}) - }, - }, - { - name: "list_nested", - valFn: func() any { - return NewRefValList(adapter, []ref.Val{ - String("hello"), - NewRefValList(adapter, []ref.Val{String("nested1"), String("nested2")}), - NewRefValMap(adapter, map[ref.Val]ref.Val{String("k1"): String("v1")}), - }) - }, - }, - { - name: "map_flat", - valFn: func() any { - return NewRefValMap(adapter, map[ref.Val]ref.Val{ - String("k1"): Int(1), - String("k2"): Int(2), - String("k3"): Int(3), - }) - }, - }, - { - name: "custom_list_flat", - valFn: func() any { - return proxyLegacyList{proxy: NewRefValList(adapter, []ref.Val{Int(1), Int(2), Int(3), Int(4), Int(5), Int(6), Int(7), Int(8)})} - }, - }, - { - name: "custom_list_nested", - valFn: func() any { - return proxyLegacyList{proxy: NewRefValList(adapter, []ref.Val{ - String("hello"), - NewRefValList(adapter, []ref.Val{String("nested1"), String("nested2")}), - NewRefValMap(adapter, map[ref.Val]ref.Val{String("k1"): String("v1")}), - })} - }, - }, - { - name: "custom_map_flat", - valFn: func() any { - return interopFoldableMap{Mapper: NewRefValMap(adapter, map[ref.Val]ref.Val{ - String("k1"): Int(1), - String("k2"): Int(2), - String("k3"): Int(3), - })} - }, - }, - { - name: "proto_obj", - valFn: func() any { - return adapter.NativeToValue(&proto3pb.TestAllTypes{ - SingleString: "hello world", - SingleInt64: 42, - RepeatedString: []string{"first", "second", "third"}, - MapStringString: map[string]string{ - "key1": "value1", - "key2": "value2", - }, - }) - }, - }, - { - name: "native_obj", - valFn: func() any { - return adapter.NativeToValue(&rootNative{ - Name: "parent", - Count: 100, - Children: []nestedNative{ - {NestedList: []string{"a", "b", "c"}, NestedMap: map[string]int{"k1": 1, "k2": 2}}, - {NestedList: []string{"d", "e"}, NestedMap: map[string]int{"k3": 3}}, - }, - }) - }, - }, - } - - for _, bm := range benchmarks { - b.Run(bm.name, func(b *testing.B) { - b.ReportAllocs() - calculator := NewSizeCalculator() - for i := 0; i < b.N; i++ { - val := bm.valFn() - _ = calculator.AggregateSize(val) - } - }) - } -} - -func BenchmarkCalculateSizeScaled(b *testing.B) { - adapter := DefaultTypeAdapter - sizes := []int{10, 100, 1000} - - for _, size := range sizes { - // Prepare list elements - listElems := make([]ref.Val, size) - for i := 0; i < size; i++ { - listElems[i] = Int(i) - } - builtinList := NewRefValList(adapter, listElems) - customList := proxyLegacyList{proxy: builtinList} - - // Prepare map entries - mapEntries := make(map[ref.Val]ref.Val, size) - for i := 0; i < size; i++ { - mapEntries[String(fmt.Sprintf("k%d", i))] = Int(i) - } - builtinMap := NewRefValMap(adapter, mapEntries) - customMap := interopFoldableMap{Mapper: builtinMap} - - // Amortized (repeated calculation on memoized vs unmemoized custom instance) - b.Run(fmt.Sprintf("Amortized/builtin_list/N=%d", size), func(b *testing.B) { - b.ReportAllocs() - calc := NewSizeCalculator() - for i := 0; i < b.N; i++ { - _ = calc.AggregateSize(builtinList) - } - }) - b.Run(fmt.Sprintf("Amortized/custom_list/N=%d", size), func(b *testing.B) { - b.ReportAllocs() - calc := NewSizeCalculator() - for i := 0; i < b.N; i++ { - _ = calc.AggregateSize(customList) - } - }) - b.Run(fmt.Sprintf("Amortized/builtin_map/N=%d", size), func(b *testing.B) { - b.ReportAllocs() - calc := NewSizeCalculator() - for i := 0; i < b.N; i++ { - _ = calc.AggregateSize(builtinMap) - } - }) - b.Run(fmt.Sprintf("Amortized/custom_map/N=%d", size), func(b *testing.B) { - b.ReportAllocs() - calc := NewSizeCalculator() - for i := 0; i < b.N; i++ { - _ = calc.AggregateSize(customMap) - } - }) - - // First-time / Uncached calculation - b.Run(fmt.Sprintf("FirstTime/builtin_list/N=%d", size), func(b *testing.B) { - b.ReportAllocs() - calc := NewSizeCalculator() - for i := 0; i < b.N; i++ { - l := NewRefValList(adapter, listElems) - _ = calc.AggregateSize(l) - } - }) - b.Run(fmt.Sprintf("FirstTime/custom_list/N=%d", size), func(b *testing.B) { - b.ReportAllocs() - calc := NewSizeCalculator() - for i := 0; i < b.N; i++ { - l := proxyLegacyList{proxy: NewRefValList(adapter, listElems)} - _ = calc.AggregateSize(l) - } - }) - b.Run(fmt.Sprintf("FirstTime/builtin_map/N=%d", size), func(b *testing.B) { - b.ReportAllocs() - calc := NewSizeCalculator() - for i := 0; i < b.N; i++ { - m := NewRefValMap(adapter, mapEntries) - _ = calc.AggregateSize(m) - } - }) - b.Run(fmt.Sprintf("FirstTime/custom_map/N=%d", size), func(b *testing.B) { - b.ReportAllocs() - calc := NewSizeCalculator() - for i := 0; i < b.N; i++ { - m := interopFoldableMap{Mapper: NewRefValMap(adapter, mapEntries)} - _ = calc.AggregateSize(m) - } - }) - } - - // Benchmark nested tree complexity (Depth x Width) - depths := []int{2, 3} - width := 10 - for _, depth := range depths { - builtinNested := createNestedList(adapter, depth, width) - customNested := createNestedCustomList(adapter, depth, width) - - b.Run(fmt.Sprintf("Complexity/builtin_nested_list/Depth=%d_Width=%d", depth, width), func(b *testing.B) { - b.ReportAllocs() - calc := NewSizeCalculator() - for i := 0; i < b.N; i++ { - _ = calc.AggregateSize(builtinNested) - } - }) - b.Run(fmt.Sprintf("Complexity/custom_nested_list/Depth=%d_Width=%d", depth, width), func(b *testing.B) { - b.ReportAllocs() - calc := NewSizeCalculator() - for i := 0; i < b.N; i++ { - _ = calc.AggregateSize(customNested) - } - }) - } -} - -func createNestedList(adapter Adapter, depth, width int) ref.Val { - if depth <= 1 { - elems := make([]ref.Val, width) - for i := 0; i < width; i++ { - elems[i] = Int(i) - } - return NewRefValList(adapter, elems) - } - elems := make([]ref.Val, width) - for i := 0; i < width; i++ { - elems[i] = createNestedList(adapter, depth-1, width) - } - return NewRefValList(adapter, elems) -} - -func createNestedCustomList(adapter Adapter, depth, width int) ref.Val { - if depth <= 1 { - elems := make([]ref.Val, width) - for i := 0; i < width; i++ { - elems[i] = Int(i) - } - return proxyLegacyList{proxy: NewRefValList(adapter, elems)} - } - elems := make([]ref.Val, width) - for i := 0; i < width; i++ { - elems[i] = createNestedCustomList(adapter, depth-1, width) - } - return proxyLegacyList{proxy: NewRefValList(adapter, elems)} -} From 98b642ab0f2860ed2e9953da0f8487891e63cd21 Mon Sep 17 00:00:00 2001 From: TristonianJones Date: Fri, 14 Aug 2026 16:50:09 -0700 Subject: [PATCH 3/3] Enhancements to use Foldable and reduce duplication in native size computations --- common/types/aggregate_sizer.go | 14 +++ common/types/map.go | 14 +-- common/types/size_calc.go | 163 +++++++++----------------------- common/types/size_calc_test.go | 133 ++++++++++++++++++++++++++ 4 files changed, 196 insertions(+), 128 deletions(-) diff --git a/common/types/aggregate_sizer.go b/common/types/aggregate_sizer.go index ae046a0c2..d56b8bc4d 100644 --- a/common/types/aggregate_sizer.go +++ b/common/types/aggregate_sizer.go @@ -27,3 +27,17 @@ type AggregateSizeVisitor interface { // AggregateSize returns the total count of nested atomic elements (capped at math.MaxUint32). AggregateSize(sizer AggregateSizer) uint32 } + +// Helper for computing aggregate sizes of traits.Foldable types. +type foldableAggregateSizer struct { + sizer AggregateSizer + total uint32 +} + +// FoldEntry implements the traits.FoldEntry interface method and counts the aggregate size +// keys and values. +func (f *foldableAggregateSizer) FoldEntry(k, v any) bool { + f.total = safeAddUint32(f.total, f.sizer.AggregateSize(k)) + f.total = safeAddUint32(f.total, f.sizer.AggregateSize(v)) + return true +} diff --git a/common/types/map.go b/common/types/map.go index 093098998..dc502323f 100644 --- a/common/types/map.go +++ b/common/types/map.go @@ -308,16 +308,10 @@ func (m *baseMap) AggregateSize(sizer AggregateSizer) uint32 { if m.aggSize != 0 { return m.aggSize } - var total uint32 = 1 - it := m.Iterator() - for it.HasNext() == True { - key := it.Next() - val, _ := m.Find(key) - total = safeAddUint32(total, sizer.AggregateSize(key)) - total = safeAddUint32(total, sizer.AggregateSize(val)) - } - m.aggSize = total - return total + f := foldableAggregateSizer{sizer: sizer, total: 1} + m.Fold(&f) + m.aggSize = f.total + return f.total } // String converts the map into a human-readable string. diff --git a/common/types/size_calc.go b/common/types/size_calc.go index b3689acbd..ab6cb4576 100644 --- a/common/types/size_calc.go +++ b/common/types/size_calc.go @@ -80,15 +80,12 @@ type sizeContext struct { traversalCount *int } -func (c *sizeContext) childContext() *sizeContext { - return &sizeContext{ - calc: c.calc, - depth: c.depth + 1, - traversalCount: c.traversalCount, - } +func (c sizeContext) childContext() sizeContext { + c.depth++ + return c } -func (c *sizeContext) visitNode() bool { +func (c sizeContext) visitNode() bool { *c.traversalCount++ if *c.traversalCount > c.calc.maxTraversal || c.depth > c.calc.maxDepth { return false @@ -100,7 +97,7 @@ func (c *sizeContext) visitNode() bool { // Otherwise, a unit size of 1 is returned. func (s *SizeCalculator) AggregateSize(val any) uint32 { traversals := 0 - ctx := &sizeContext{ + ctx := sizeContext{ calc: s, depth: 1, traversalCount: &traversals, @@ -110,13 +107,17 @@ func (s *SizeCalculator) AggregateSize(val any) uint32 { // AggregateSize implements the ref.Val interface and allows for the generation of nested // child context values which are necessary for correct traversal count tracking. -func (c *sizeContext) AggregateSize(val any) uint32 { +func (c sizeContext) AggregateSize(val any) uint32 { if !c.visitNode() { return math.MaxUint32 } switch v := val.(type) { case AggregateSizeVisitor: return v.AggregateSize(c.childContext()) + case traits.Foldable: + f := foldableAggregateSizer{sizer: c.childContext(), total: 1} + v.Fold(&f) + return f.total case traits.Mapper: total := uint32(1) it := v.Iterator() @@ -136,10 +137,6 @@ func (c *sizeContext) AggregateSize(val any) uint32 { total = safeAddUint32(total, childCtx.AggregateSize(it.Next())) } return total - case String: - return safeUint32FromInt(utf8.RuneCountInString(string(v))) - case Bytes: - return safeUint32FromInt(len(v)) case traits.Sizer: return safeUint32FromBoxedInt(v.Size().(Int)) case Bool, Int, Uint, Double, Duration, Timestamp, Null, *Type, *Err, *Unknown: @@ -147,9 +144,9 @@ func (c *sizeContext) AggregateSize(val any) uint32 { case ref.Val: return c.AggregateSize(v.Value()) case protoreflect.Value: - return getProtoValueAggregateSize(c, v) + return c.AggregateSize(v.Interface()) case protoreflect.MapKey: - return getProtoValueAggregateSize(c, v.Value()) + return c.AggregateSize(v.Value().Interface()) case protoreflect.Message: return getProtoMessageAggregateSize(c, v) case protoreflect.List: @@ -176,95 +173,21 @@ func (c *sizeContext) AggregateSize(val any) uint32 { } } -func getProtoValueAggregateSize(c *sizeContext, v protoreflect.Value) uint32 { - switch val := v.Interface().(type) { - case string: - return c.AggregateSize(val) - case []byte: - return c.AggregateSize(val) - case protoreflect.Message: - return getProtoMessageAggregateSize(c.childContext(), val) - case protoreflect.List: - return getProtoListAggregateSize(c.childContext(), val) - case protoreflect.Map: - return getProtoMapAggregateSize(c.childContext(), val) - default: - if !c.visitNode() { - return math.MaxUint32 - } - return 1 - } -} - -func getProtoFieldAggregateSize(c *sizeContext, fd protoreflect.FieldDescriptor, v protoreflect.Value) uint32 { +func getProtoFieldAggregateSize(c sizeContext, fd protoreflect.FieldDescriptor, v protoreflect.Value) uint32 { if !c.visitNode() { return math.MaxUint32 } childCtx := c.childContext() if fd.IsMap() { - m := v.Map() - total := uint32(1) - valKind := fd.MapValue().Kind() - m.Range(func(k protoreflect.MapKey, val protoreflect.Value) bool { - total = safeAddUint32(total, getProtoMapKeyAggregateSize(childCtx, fd.MapKey(), k)) - total = safeAddUint32(total, getProtoFieldByKindAggregateSize(childCtx, valKind, val)) - return true - }) - return total + return getProtoMapAggregateSize(childCtx, v.Map()) } if fd.IsList() { - l := v.List() - elemKind := fd.Kind() - total := safeAddUint32(1, safeUint32FromInt(l.Len())) - switch elemKind { - case protoreflect.StringKind: - total = 1 - for i := 0; i < l.Len(); i++ { - total = safeAddUint32(total, childCtx.AggregateSize(l.Get(i).String())) - } - case protoreflect.BytesKind: - total = 1 - for i := 0; i < l.Len(); i++ { - total = safeAddUint32(total, childCtx.AggregateSize(l.Get(i).Bytes())) - } - case protoreflect.MessageKind, protoreflect.GroupKind: - total = 1 - for i := 0; i < l.Len(); i++ { - total = safeAddUint32(total, getProtoMessageAggregateSize(childCtx, l.Get(i).Message())) - } - } - return total + return getProtoListAggregateSize(childCtx, v.List()) } - return getProtoFieldByKindAggregateSize(childCtx, fd.Kind(), v) + return childCtx.AggregateSize(v.Interface()) } -func getProtoMapKeyAggregateSize(c *sizeContext, keyDesc protoreflect.FieldDescriptor, k protoreflect.MapKey) uint32 { - if keyDesc.Kind() == protoreflect.StringKind { - return c.AggregateSize(k.String()) - } - if !c.visitNode() { - return math.MaxUint32 - } - return 1 -} - -func getProtoFieldByKindAggregateSize(c *sizeContext, kind protoreflect.Kind, v protoreflect.Value) uint32 { - switch kind { - case protoreflect.StringKind: - return c.AggregateSize(v.String()) - case protoreflect.BytesKind: - return c.AggregateSize(v.Bytes()) - case protoreflect.MessageKind, protoreflect.GroupKind: - return getProtoMessageAggregateSize(c, v.Message()) - default: - if !c.visitNode() { - return math.MaxUint32 - } - return 1 - } -} - -func getProtoMessageAggregateSize(c *sizeContext, m protoreflect.Message) uint32 { +func getProtoMessageAggregateSize(c sizeContext, m protoreflect.Message) uint32 { if !m.IsValid() { return 0 } @@ -280,7 +203,7 @@ func getProtoMessageAggregateSize(c *sizeContext, m protoreflect.Message) uint32 return total } -func getProtoListAggregateSize(c *sizeContext, l protoreflect.List) uint32 { +func getProtoListAggregateSize(c sizeContext, l protoreflect.List) uint32 { if !l.IsValid() { return 0 } @@ -289,13 +212,13 @@ func getProtoListAggregateSize(c *sizeContext, l protoreflect.List) uint32 { } childCtx := c.childContext() total := uint32(1) - for i := 0; i < l.Len(); i++ { - total = safeAddUint32(total, getProtoValueAggregateSize(childCtx, l.Get(i))) + for i := range l.Len() { + total = safeAddUint32(total, childCtx.AggregateSize(l.Get(i).Interface())) } return total } -func getProtoMapAggregateSize(c *sizeContext, m protoreflect.Map) uint32 { +func getProtoMapAggregateSize(c sizeContext, m protoreflect.Map) uint32 { if !m.IsValid() { return 0 } @@ -305,14 +228,14 @@ func getProtoMapAggregateSize(c *sizeContext, m protoreflect.Map) uint32 { childCtx := c.childContext() total := uint32(1) m.Range(func(k protoreflect.MapKey, v protoreflect.Value) bool { - total = safeAddUint32(total, getProtoValueAggregateSize(childCtx, k.Value())) - total = safeAddUint32(total, getProtoValueAggregateSize(childCtx, v)) + total = safeAddUint32(total, childCtx.AggregateSize(k.Value().Interface())) + total = safeAddUint32(total, childCtx.AggregateSize(v.Interface())) return true }) return total } -func getReflectValueAggregateSize(c *sizeContext, fieldVal reflect.Value) uint32 { +func getReflectValueAggregateSize(c sizeContext, fieldVal reflect.Value) uint32 { if !fieldVal.IsValid() { return 0 } @@ -354,33 +277,22 @@ func getReflectValueAggregateSize(c *sizeContext, fieldVal reflect.Value) uint32 if fieldVal.IsNil() { return 0 } - if fieldVal.CanInterface() { - if sizer, ok := fieldVal.Interface().(AggregateSizeVisitor); ok { - return sizer.AggregateSize(childCtx) - } - if sizer, ok := fieldVal.Interface().(traits.Sizer); ok { - return safeUint32FromBoxedInt(sizer.Size().(Int)) - } + if sz, ok := checkCustomSizer(childCtx, fieldVal); ok { + return sz } return getReflectValueAggregateSize(c, fieldVal.Elem()) case reflect.Struct: if fieldVal.Type() == timestampType || fieldVal.Type() == durationType { return 1 } - if fieldVal.CanInterface() { - if sizer, ok := fieldVal.Interface().(AggregateSizeVisitor); ok { - return sizer.AggregateSize(childCtx) - } - if sizer, ok := fieldVal.Interface().(traits.Sizer); ok { - return safeUint32FromBoxedInt(sizer.Size().(Int)) - } + if sz, ok := checkCustomSizer(childCtx, fieldVal); ok { + return sz } total := uint32(1) t := fieldVal.Type() numFields := fieldVal.NumField() - for i := 0; i < numFields; i++ { - f := t.Field(i) - if !f.IsExported() { + for i := range numFields { + if !t.Field(i).IsExported() { continue } fVal := fieldVal.Field(i) @@ -394,3 +306,18 @@ func getReflectValueAggregateSize(c *sizeContext, fieldVal reflect.Value) uint32 return 1 } } + +func checkCustomSizer(c sizeContext, fieldVal reflect.Value) (uint32, bool) { + if !fieldVal.CanInterface() { + return 0, false + } + + switch sizer := fieldVal.Interface().(type) { + case AggregateSizeVisitor: + return sizer.AggregateSize(c), true + case traits.Sizer: + return safeUint32FromBoxedInt(sizer.Size().(Int)), true + default: + return 0, false + } +} diff --git a/common/types/size_calc_test.go b/common/types/size_calc_test.go index d649387a5..43a4b064e 100644 --- a/common/types/size_calc_test.go +++ b/common/types/size_calc_test.go @@ -24,6 +24,8 @@ import ( "google.golang.org/protobuf/reflect/protoreflect" "github.com/google/cel-go/common/types/ref" + "github.com/google/cel-go/common/types/traits" + proto3pb "github.com/google/cel-go/test/proto3pb" ) @@ -208,6 +210,31 @@ func TestCalculateSize(t *testing.T) { val: interopFoldableMap{Mapper: NewStringStringMap(DefaultTypeAdapter, map[string]string{"key": "val"})}, want: 7, // 1 (container) + 3 ("key") + 3 ("val") = 7 }, + { + name: "custom_pure_mapper", + val: customPureMapper{Mapper: NewStringStringMap(DefaultTypeAdapter, map[string]string{"key": "val"})}, + want: 7, // 1 (container) + 3 ("key") + 3 ("val") = 7 + }, + { + name: "custom_sizer_struct_field", + val: struct{ Sizer traits.Sizer }{Sizer: customSizerVal(42)}, + want: 43, // 1 (struct container) + 42 (custom sizer) = 43 + }, + { + name: "custom_visitor_struct_field", + val: struct{ Visitor customVisitorVal }{Visitor: customVisitorVal{Val: 1}}, + want: 101, // 1 (struct container) + 100 (custom visitor) = 101 + }, + { + name: "custom_sizer_pointer", + val: newCustomSizerPtr(42), + want: 42, + }, + { + name: "custom_visitor_pointer", + val: &customVisitorVal{Val: 1}, + want: 100, + }, } for _, tc := range tests { @@ -220,6 +247,36 @@ func TestCalculateSize(t *testing.T) { } } +type customPureMapper struct { + traits.Mapper +} + +type customSizerVal int + +func (c customSizerVal) Size() ref.Val { + return Int(c) +} + +type customSizerPtr struct { + val int +} + +func (c *customSizerPtr) Size() ref.Val { + return Int(c.val) +} + +func newCustomSizerPtr(v int) *customSizerPtr { + return &customSizerPtr{val: v} +} + +type customVisitorVal struct { + Val int +} + +func (c customVisitorVal) AggregateSize(sizer AggregateSizer) uint32 { + return 100 +} + func TestNativeObjCalculateSizeNil(t *testing.T) { nilNative := &nativeObj{} if got := nilNative.AggregateSize(NewSizeCalculator()); got != 0 { @@ -343,6 +400,66 @@ func TestSizeCalculatorOptions(t *testing.T) { t.Errorf("AggregateSize for native nested map depth > 2 got %d, want MaxUint32", got) } }) + + t.Run("maxTraversal map limit", func(t *testing.T) { + calc := NewSizeCalculator(SizeCalculatorMaxTraversal(3)) + m := NewRefValMap(adapter, map[ref.Val]ref.Val{ + String("k1"): String("v1"), + String("k2"): String("v2"), + }) + if got := calc.AggregateSize(m); got != math.MaxUint32 { + t.Errorf("AggregateSize for map traversal > 3 got %d, want MaxUint32", got) + } + }) + + t.Run("maxTraversal proto limit", func(t *testing.T) { + calc := NewSizeCalculator(SizeCalculatorMaxTraversal(2)) + msg := &proto3pb.TestAllTypes{ + SingleString: "hello", + SingleInt64: 42, + } + if got := calc.AggregateSize(msg); got != math.MaxUint32 { + t.Errorf("AggregateSize for proto traversal > 2 got %d, want MaxUint32", got) + } + }) + + t.Run("maxTraversal native struct limit", func(t *testing.T) { + calc := NewSizeCalculator(SizeCalculatorMaxTraversal(2)) + s := struct{ A, B, C int }{A: 1, B: 2, C: 3} + if got := calc.AggregateSize(s); got != math.MaxUint32 { + t.Errorf("AggregateSize for native struct traversal > 2 got %d, want MaxUint32", got) + } + }) + + t.Run("maxTraversal native map limit", func(t *testing.T) { + calc := NewSizeCalculator(SizeCalculatorMaxTraversal(3)) + m := map[string]int{"a": 1, "b": 2} + if got := calc.AggregateSize(m); got != math.MaxUint32 { + t.Errorf("AggregateSize for native map traversal > 3 got %d, want MaxUint32", got) + } + }) + + t.Run("maxTraversal native slice limit", func(t *testing.T) { + calc := NewSizeCalculator(SizeCalculatorMaxTraversal(3)) + slice := []string{"a", "b", "c"} + if got := calc.AggregateSize(slice); got != math.MaxUint32 { + t.Errorf("AggregateSize for native slice traversal > 3 got %d, want MaxUint32", got) + } + }) + + t.Run("zero depth limit saturation", func(t *testing.T) { + calc := NewSizeCalculator(SizeCalculatorMaxDepth(0)) + if got := calc.AggregateSize(Int(42)); got != math.MaxUint32 { + t.Errorf("AggregateSize with depth 0 got %d, want MaxUint32", got) + } + }) + + t.Run("zero traversal limit saturation", func(t *testing.T) { + calc := NewSizeCalculator(SizeCalculatorMaxTraversal(0)) + if got := calc.AggregateSize(Int(42)); got != math.MaxUint32 { + t.Errorf("AggregateSize with traversal 0 got %d, want MaxUint32", got) + } + }) } type nestedNative struct { @@ -570,6 +687,14 @@ func BenchmarkCalculateSizeScaled(b *testing.B) { _ = calc.AggregateSize(customMap) } }) + b.Run(fmt.Sprintf("Amortized/custom_pure_map/N=%d", size), func(b *testing.B) { + b.ReportAllocs() + calc := NewSizeCalculator() + pureMap := customPureMapper{Mapper: builtinMap} + for i := 0; i < b.N; i++ { + _ = calc.AggregateSize(pureMap) + } + }) // First-time / Uncached calculation b.Run(fmt.Sprintf("FirstTime/builtin_list/N=%d", size), func(b *testing.B) { @@ -604,6 +729,14 @@ func BenchmarkCalculateSizeScaled(b *testing.B) { _ = calc.AggregateSize(m) } }) + b.Run(fmt.Sprintf("FirstTime/custom_pure_map/N=%d", size), func(b *testing.B) { + b.ReportAllocs() + calc := NewSizeCalculator() + for i := 0; i < b.N; i++ { + m := customPureMapper{Mapper: NewRefValMap(adapter, mapEntries)} + _ = calc.AggregateSize(m) + } + }) } // Benchmark nested tree complexity (Depth x Width)