From bb9e63aed4502a6ad3a2e52b7d476613552e8f99 Mon Sep 17 00:00:00 2001
From: Gabriele Cimato <gabriele.cimato+github@gmail.com>
Date: Wed, 26 Jun 2024 14:44:56 -0500
Subject: [PATCH] BAAS-32412: add nil check in mem usage for key and item in
 builtin_set (#124)

---
 builtin_map_test.go |  2 +-
 builtin_set.go      | 24 +++++++++++++-----------
 builtin_set_test.go | 32 ++++++++++++++++++++++++++++++++
 3 files changed, 46 insertions(+), 12 deletions(-)

diff --git a/builtin_map_test.go b/builtin_map_test.go
index 76284a46..6c63bd8e 100644
--- a/builtin_map_test.go
+++ b/builtin_map_test.go
@@ -260,7 +260,7 @@ func createOrderedMap(vm *Runtime, size int) *orderedMap {
 			value: value,
 		}
 		// These iter items are necessary for testing the mem usage
-		// estimation since that's who we iterate through the map.
+		// estimation since that's how we iterate through the map.
 		if i > 0 {
 			ht[uint64(i)].iterPrev = ht[uint64(i-1)]
 			ht[uint64(i-1)].iterNext = ht[uint64(i)]
diff --git a/builtin_set.go b/builtin_set.go
index 65400145..c393ae36 100644
--- a/builtin_set.go
+++ b/builtin_set.go
@@ -364,18 +364,20 @@ func (so *setObject) estimateMemUsage(ctx *MemUsageContext) (estimate uint64, er
 		}
 		samplesVisited += 1
 
-		// We still want to account for both key and value if we return a non-zero value on error.
-		// This could otherwise skew the estimate when in reality key/value pairs contribute to
-		// mem usage together.
-		inc, incErr := item.key.MemUsage(ctx)
-		memUsage += inc
-		inc, valErr := item.value.MemUsage(ctx)
-		memUsage += inc
-		if valErr != nil {
-			return computeMemUsageEstimate(memUsage, samplesVisited, totalItems), valErr
+		if item.key != nil {
+			inc, err := item.key.MemUsage(ctx)
+			memUsage += inc
+			if err != nil {
+				return computeMemUsageEstimate(memUsage, samplesVisited, totalItems), err
+			}
 		}
-		if incErr != nil {
-			return computeMemUsageEstimate(memUsage, samplesVisited, totalItems), incErr
+
+		if item.value != nil {
+			inc, err := item.value.MemUsage(ctx)
+			memUsage += inc
+			if err != nil {
+				return computeMemUsageEstimate(memUsage, samplesVisited, totalItems), err
+			}
 		}
 	}
 
diff --git a/builtin_set_test.go b/builtin_set_test.go
index de9c6789..e00ae265 100644
--- a/builtin_set_test.go
+++ b/builtin_set_test.go
@@ -193,6 +193,28 @@ func TestSetHasFloatVsInt(t *testing.T) {
 	testScript(SCRIPT, valueTrue, t)
 }
 
+func createOrderedMapWithNilValues(size int) *orderedMap {
+	ht := make(map[uint64]*mapEntry, 0)
+	for i := 0; i < size; i += 1 {
+		ht[uint64(i)] = &mapEntry{
+			key:   nil,
+			value: nil,
+		}
+		// These iter items are necessary for testing the mem usage
+		// estimation since that's how we iterate through the map.
+		if i > 0 {
+			ht[uint64(i)].iterPrev = ht[uint64(i-1)]
+			ht[uint64(i-1)].iterNext = ht[uint64(i)]
+		}
+	}
+	return &orderedMap{
+		size:      size,
+		iterFirst: ht[uint64(0)],
+		iterLast:  ht[uint64(size-1)],
+		hashTable: ht,
+	}
+}
+
 func TestSetObjectMemUsage(t *testing.T) {
 	vm := New()
 
@@ -274,6 +296,16 @@ func TestSetObjectMemUsage(t *testing.T) {
 				(5+SizeString)*20,
 			errExpected: nil,
 		},
+		{
+			name: "mem above estimate threshold and within memory limit and nil values returns correct mem usage",
+			mu:   NewMemUsageContext(vm, 88, 100, 50, 1, 0.1, TestNativeMemUsageChecker{}),
+			so: &setObject{
+				m: createOrderedMapWithNilValues(3),
+			},
+			// baseObject
+			expectedMem: SizeEmptyStruct,
+			errExpected: nil,
+		},
 		{
 			name:        "mem is SizeEmptyStruct given a nil orderedMap object",
 			mu:          NewMemUsageContext(vm, 88, 5000, 50, 50, 0.1, TestNativeMemUsageChecker{}),