From 5ca9a4dfd4c37a56f0dd886596494be4599f01c8 Mon Sep 17 00:00:00 2001
From: Quinn Klassen <klassenq@gmail.com>
Date: Tue, 9 Jan 2024 12:05:22 -0800
Subject: [PATCH] Add deterministic range over map helper (#1340)

Add deterministic range over map helper
---
 internal/workflow.go      | 27 +++++++++++++++
 internal/workflow_test.go | 72 +++++++++++++++++++++++++++++++++++++++
 workflow/workflow.go      | 15 ++++++++
 3 files changed, 114 insertions(+)

diff --git a/internal/workflow.go b/internal/workflow.go
index 0ef1c2635..55a2fbc62 100644
--- a/internal/workflow.go
+++ b/internal/workflow.go
@@ -31,6 +31,9 @@ import (
 	"strings"
 	"time"
 
+	"golang.org/x/exp/constraints"
+	"golang.org/x/exp/slices"
+
 	"google.golang.org/protobuf/types/known/durationpb"
 
 	commonpb "go.temporal.io/api/common/v1"
@@ -2043,3 +2046,27 @@ func convertFromPBRetryPolicy(retryPolicy *commonpb.RetryPolicy) *RetryPolicy {
 func GetLastCompletionResultFromWorkflowInfo(info *WorkflowInfo) *commonpb.Payloads {
 	return info.lastCompletionResult
 }
+
+// DeterministicKeys returns the keys of a map in deterministic (sorted) order. To be used in for
+// loops in workflows for deterministic iteration.
+func DeterministicKeys[K constraints.Ordered, V any](m map[K]V) []K {
+	r := make([]K, 0, len(m))
+	for k := range m {
+		r = append(r, k)
+	}
+	slices.Sort(r)
+	return r
+}
+
+// DeterministicKeysFunc returns the keys of a map in a deterministic (sorted) order.
+// cmp(a, b) should return a negative number when a < b, a positive number when
+// a > b and zero when a == b. Keys are sorted by cmp.
+// To be used in for loops in workflows for deterministic iteration.
+func DeterministicKeysFunc[K comparable, V any](m map[K]V, cmp func(a K, b K) int) []K {
+	r := make([]K, 0, len(m))
+	for k := range m {
+		r = append(r, k)
+	}
+	slices.SortStableFunc(r, cmp)
+	return r
+}
diff --git a/internal/workflow_test.go b/internal/workflow_test.go
index 610f2cd05..1abfa8209 100644
--- a/internal/workflow_test.go
+++ b/internal/workflow_test.go
@@ -189,3 +189,75 @@ func _assertNonZero(t *testing.T, i interface{}, prefix string) {
 		}
 	}
 }
+
+func TestDeterministicKeys(t *testing.T) {
+	t.Parallel()
+
+	var tests = []struct {
+		unsorted map[int]int
+		sorted   []int
+	}{
+		{
+			map[int]int{1: 1, 2: 2, 3: 3},
+			[]int{1, 2, 3},
+		},
+		{
+			map[int]int{},
+			[]int{},
+		},
+		{
+			map[int]int{1: 1, 5: 5, 3: 3},
+			[]int{1, 3, 5},
+		},
+		{
+			map[int]int{3: 3, 2: 2, 1: 1},
+			[]int{1, 2, 3},
+		},
+	}
+
+	for _, tt := range tests {
+		testname := fmt.Sprintf("%d,%d", tt.unsorted, tt.sorted)
+		t.Run(testname, func(t *testing.T) {
+			assert.Equal(t, tt.sorted, DeterministicKeys(tt.unsorted))
+		})
+	}
+}
+
+func TestDeterministicKeysFunc(t *testing.T) {
+	t.Parallel()
+
+	type keyStruct struct {
+		i int
+	}
+
+	var tests = []struct {
+		unsorted map[keyStruct]int
+		sorted   []keyStruct
+	}{
+		{
+			map[keyStruct]int{{1}: 1, {2}: 2, {3}: 3},
+			[]keyStruct{{1}, {2}, {3}},
+		},
+		{
+			map[keyStruct]int{},
+			[]keyStruct{},
+		},
+		{
+			map[keyStruct]int{{1}: 1, {5}: 5, {3}: 3},
+			[]keyStruct{{1}, {3}, {5}},
+		},
+		{
+			map[keyStruct]int{{3}: 3, {2}: 2, {1}: 1},
+			[]keyStruct{{1}, {2}, {3}},
+		},
+	}
+
+	for _, tt := range tests {
+		testname := fmt.Sprintf("%d,%d", tt.unsorted, tt.sorted)
+		t.Run(testname, func(t *testing.T) {
+			assert.Equal(t, tt.sorted, DeterministicKeysFunc(tt.unsorted, func(a, b keyStruct) int {
+				return a.i - b.i
+			}))
+		})
+	}
+}
diff --git a/workflow/workflow.go b/workflow/workflow.go
index 42b16aa2a..3811a3dfa 100644
--- a/workflow/workflow.go
+++ b/workflow/workflow.go
@@ -31,6 +31,7 @@ import (
 	"go.temporal.io/sdk/internal"
 	"go.temporal.io/sdk/internal/common/metrics"
 	"go.temporal.io/sdk/log"
+	"golang.org/x/exp/constraints"
 )
 
 type (
@@ -635,3 +636,17 @@ func IsContinueAsNewError(err error) bool {
 func DataConverterWithoutDeadlockDetection(c converter.DataConverter) converter.DataConverter {
 	return internal.DataConverterWithoutDeadlockDetection(c)
 }
+
+// DeterministicKeys returns the keys of a map in deterministic (sorted) order. To be used in for
+// loops in workflows for deterministic iteration.
+func DeterministicKeys[K constraints.Ordered, V any](m map[K]V) []K {
+	return internal.DeterministicKeys(m)
+}
+
+// DeterministicKeysFunc returns the keys of a map in a deterministic (sorted) order.
+// cmp(a, b) should return a negative number when a < b, a positive number when
+// a > b and zero when a == b. Keys are sorted by cmp.
+// To be used in for loops in workflows for deterministic iteration.
+func DeterministicKeysFunc[K comparable, V any](m map[K]V, cmp func(K, K) int) []K {
+	return internal.DeterministicKeysFunc(m, cmp)
+}