Removed extra method that intercepted desired result

Author: TristonianJones

common/types/memory_tracker.go

@@ -129,7 +129,7 @@ func (t *MemoryTracker) Track(val ref.Val) uint32 {
 		return sz
 	}
 
-	est := t.calc.EstimateAggregateSize(val)
+	est := t.calc.ApproximateAggregateSize(val)
 	if est.LimitExceeded {
 		t.calcLimitExceeded = true
 	}

common/types/size_calc.go

@@ -145,30 +145,21 @@ func cacheableAggregateSize(sizer AggregateSizer) bool {
 	return ok && !status.aggregateSizeLimitExceeded()
 }
 
-// AggregateSizeEstimate captures the outcome of an aggregate size computation.
+// ApproximateAggregateSize captures the outcome of an aggregate size computation.
 //
 // The Size saturates at math.MaxUint32 when the accumulated element count overflows uint32.
 // LimitExceeded reports the computation was aborted because the value was too expensive to
 // traverse (too deep, or too many nodes visited); in that case Size is also math.MaxUint32,
 // but the value's true size may be smaller — the two conditions are distinguishable by the flag.
-type AggregateSizeEstimate struct {
+type ApproximateAggregateSize struct {
 	Size          uint32
 	LimitExceeded bool
 }
 
-// AggregateSize returns the size of the input value, if known.
-// Otherwise, a unit size of 1 is returned.
-//
-// When the calculator's depth or traversal limits are exceeded, the size saturates to
-// math.MaxUint32. Use EstimateAggregateSize to distinguish limit-exceeded results from
-// genuine uint32 saturation.
-func (s *SizeCalculator) AggregateSize(val any) uint32 {
-	return s.EstimateAggregateSize(val).Size
-}
 
-// EstimateAggregateSize returns the aggregate size of the input value along with an indication
+// ApproximateAggregateSize returns the aggregate size of the input value along with an indication
 // of whether the computation was aborted due to the calculator's depth or traversal limits.
-func (s *SizeCalculator) EstimateAggregateSize(val any) AggregateSizeEstimate {
+func (s *SizeCalculator) ApproximateAggregateSize(val any) ApproximateAggregateSize {
 	ctx := sizeContextPool.Get().(*sizeContext)
 	ctx.calc = s
 	est := ctx.estimateAggregateSize(val)
@@ -178,6 +169,11 @@ func (s *SizeCalculator) EstimateAggregateSize(val any) AggregateSizeEstimate {
 	return est
 }
 
+// AggregateSize implements the AggregateSizer interface by returning the computed size.
+func (s *SizeCalculator) AggregateSize(val any) uint32 {
+	return s.ApproximateAggregateSize(val).Size
+}
+
 // stringSize converts a character or byte length to an element count where stringUnitLength
 // characters count as a single element, rounding up with a minimum size of 1.
 func (s *SizeCalculator) stringSize(length int) uint32 {
@@ -265,13 +261,13 @@ func (c *sizeContext) aggregateSize(val any) uint32 {
 }
 
 // estimateAggregateSize evaluates the aggregate size using this context and resets state.
-func (c *sizeContext) estimateAggregateSize(val any) AggregateSizeEstimate {
+func (c *sizeContext) estimateAggregateSize(val any) ApproximateAggregateSize {
 	c.depth = 0
 	c.traversalCount = 0
 	c.limitExceeded = false
 
 	size := c.AggregateSize(val)
-	return AggregateSizeEstimate{Size: size, LimitExceeded: c.limitExceeded}
+	return ApproximateAggregateSize{Size: size, LimitExceeded: c.limitExceeded}
 }
 
 func getProtoValueAggregateSize(c *sizeContext, v protoreflect.Value) uint32 {
@@ -385,24 +381,19 @@ func getReflectValueAggregateSize(c *sizeContext, fieldVal reflect.Value) uint32
 		}
 		total := safeAddUint32(1, safeUint32FromInt(fieldVal.Len()))
 		switch elemType.Kind() {
-		case reflect.String:
+		case reflect.String, reflect.Struct, reflect.Pointer, reflect.Slice, reflect.Array, reflect.Map, reflect.Interface:
 			total = 1
 			for i := 0; i < fieldVal.Len(); i++ {
-				total = safeAddUint32(total, c.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, c.AggregateSize(fieldVal.Index(i).Interface()))
+				total = safeAddUint32(total, c.AggregateSize(fieldVal.Index(i)))
 			}
 		}
 		return total
 	case reflect.Map:
 		total := uint32(1)
 		iter := fieldVal.MapRange()
 		for iter.Next() {
-			total = safeAddUint32(total, c.AggregateSize(iter.Key().Interface()))
-			total = safeAddUint32(total, c.AggregateSize(iter.Value().Interface()))
+			total = safeAddUint32(total, c.AggregateSize(iter.Key()))
+			total = safeAddUint32(total, c.AggregateSize(iter.Value()))
 		}
 		return total
 	case reflect.Pointer, reflect.Interface:
@@ -431,25 +422,45 @@ func getReflectValueAggregateSize(c *sizeContext, fieldVal reflect.Value) uint32
 			if !fVal.IsValid() || fVal.IsZero() {
 				continue
 			}
-			total = safeAddUint32(total, c.AggregateSize(fVal.Interface()))
+			total = safeAddUint32(total, c.AggregateSize(fVal))
 		}
 		return total
 	default:
 		return 1
 	}
 }
 
+var (
+	aggregateSizeVisitorType = reflect.TypeFor[AggregateSizeVisitor]()
+	sizerType                = reflect.TypeFor[traits.Sizer]()
+	protoMessageType         = reflect.TypeFor[proto.Message]()
+)
+
 func checkCustomSizer(c *sizeContext, fieldVal reflect.Value) (uint32, bool) {
-	if !fieldVal.CanInterface() {
+	if !fieldVal.IsValid() || !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
+	t := fieldVal.Type()
+	if t.Implements(aggregateSizeVisitorType) {
+		if sizer, ok := fieldVal.Interface().(AggregateSizeVisitor); ok {
+			return sizer.AggregateSize(c), true
+		}
+	}
+	if t.Implements(sizerType) {
+		if sizer, ok := fieldVal.Interface().(traits.Sizer); ok {
+			return safeUint32FromBoxedInt(sizer.Size().(Int)), true
+		}
+	}
+	if t.Implements(protoMessageType) {
+		if fieldVal.Kind() == reflect.Pointer && fieldVal.IsNil() {
+			return 0, true
+		}
+		if sizer, ok := fieldVal.Interface().(proto.Message); ok {
+			if sizer == nil {
+				return 0, true
+			}
+			return getProtoMessageAggregateSize(c, sizer.ProtoReflect()), true
+		}
 	}
+	return 0, false
 }