From c9adbb039c3b0150210356054527d634830cefa2 Mon Sep 17 00:00:00 2001
From: ayaanoncrypto <1.06945712e+08+ayaanoncrypto@users.noreply.github.com>
Date: Thu, 27 Aug 2026 02:02:06 +0000
Subject: [PATCH] Support uniform and zipfian gas distribution in workload
 modeler

---
 config/gas.go      | 59 +++++++++++++++++++++++++++++++++++++++-------
 config/gas_test.go | 44 ++++++++++++++++++++++++++++++++++
 2 files changed, 95 insertions(+), 8 deletions(-)

diff --git a/config/gas.go b/config/gas.go
index 506e204..4d039dc 100644
--- a/config/gas.go
+++ b/config/gas.go
@@ -10,6 +10,8 @@ var (
 	_ gasGenerator = (*GasPicker)(nil)
 	_ gasGenerator = (*FixedGasGenerator)(nil)
 	_ gasGenerator = (*RandomGasGenerator)(nil)
+	_ gasGenerator = (*UniformGasGenerator)(nil)
+	_ gasGenerator = (*ZipfianGasGenerator)(nil)
 )
 
 type gasGenerator interface {
@@ -48,12 +50,22 @@ func (g *GasPicker) UnmarshalJSON(data []byte) error {
 		}
 		g.delegate = &fixed
 		return nil
-	case "random":
-		var random RandomGasGenerator
-		if err := json.Unmarshal(data, &random); err != nil {
+	case "random", "uniform":
+		var uniform UniformGasGenerator
+		if err := json.Unmarshal(data, &uniform); err != nil {
 			return err
 		}
-		g.delegate = &random
+		g.delegate = &uniform
+		return nil
+	case "zipfian":
+		var zipfian ZipfianGasGenerator
+		if err := json.Unmarshal(data, &zipfian); err != nil {
+			return err
+		}
+		if err := zipfian.dist.validate(); err != nil {
+			return err
+		}
+		g.delegate = &zipfian
 		return nil
 	default:
 		return fmt.Errorf("unknown gas generator name: %s", g.name)
@@ -74,9 +86,40 @@ type RandomGasGenerator struct {
 }
 
 func (r *RandomGasGenerator) GenerateGas(rng *mrand.Rand) (uint64, error) {
-	if r.Min >= r.Max {
-		return 0, fmt.Errorf("invalid random gas range: min %d must be less than max %d", r.Min, r.Max)
+	u := UniformGasGenerator(*r)
+	return u.GenerateGas(rng)
+}
+
+type UniformGasGenerator struct {
+	Min uint64 `json:"Min"`
+	Max uint64 `json:"Max"`
+}
+
+func (u *UniformGasGenerator) GenerateGas(rng *mrand.Rand) (uint64, error) {
+	if u.Min >= u.Max {
+		return 0, fmt.Errorf("invalid uniform gas range: min %d must be less than max %d", u.Min, u.Max)
+	}
+	span := u.Max - u.Min + 1
+	return u.Min + rng.Uint64N(span), nil
+}
+
+type ZipfianGasGenerator struct {
+	Min   uint64  `json:"Min"`
+	Max   uint64  `json:"Max"`
+	Theta float64 `json:"theta"`
+
+	dist ZipfianDistribution
+}
+
+func (z *ZipfianGasGenerator) GenerateGas(rng *mrand.Rand) (uint64, error) {
+	if z.Min >= z.Max {
+		return 0, fmt.Errorf("invalid zipfian gas range: min %d must be less than max %d", z.Min, z.Max)
+	}
+	z.dist.Theta = z.Theta
+	n := z.Max - z.Min + 1
+	idx, err := z.dist.SampleIndex(rng, n)
+	if err != nil {
+		return 0, err
 	}
-	span := r.Max - r.Min + 1
-	return r.Min + rng.Uint64N(span), nil
+	return z.Min + idx, nil
 }
diff --git a/config/gas_test.go b/config/gas_test.go
index d57eeb1..2ce5fbe 100644
--- a/config/gas_test.go
+++ b/config/gas_test.go
@@ -37,6 +37,22 @@ func TestGasPicker(t *testing.T) {
 		require.GreaterOrEqual(t, gas, uint64(20000))
 		require.LessOrEqual(t, gas, uint64(30000))
 	})
+	t.Run("uniform", func(t *testing.T) {
+		var subject config.GasPicker
+		require.NoError(t, subject.UnmarshalJSON([]byte(`{"Name":"uniform","Min":20000,"Max":30000}`)))
+		gas, err := subject.GenerateGas(newTestRng(1))
+		require.NoError(t, err)
+		require.GreaterOrEqual(t, gas, uint64(20000))
+		require.LessOrEqual(t, gas, uint64(30000))
+	})
+	t.Run("zipfian", func(t *testing.T) {
+		var subject config.GasPicker
+		require.NoError(t, subject.UnmarshalJSON([]byte(`{"Name":"zipfian","Min":20000,"Max":30000,"theta":0.9}`)))
+		gas, err := subject.GenerateGas(newTestRng(1))
+		require.NoError(t, err)
+		require.GreaterOrEqual(t, gas, uint64(20000))
+		require.LessOrEqual(t, gas, uint64(30000))
+	})
 	t.Run("unknown", func(t *testing.T) {
 		var subject config.GasPicker
 		require.Error(t, subject.UnmarshalJSON([]byte(`{"Name":"unknown"}`)))
@@ -87,3 +103,31 @@ func TestGenerateGasForFixedAndEmpty(t *testing.T) {
 	require.NoError(t, err)
 	require.Zero(t, gas)
 }
+
+// zipfianPicker unmarshals a fresh zipfian gas picker over [min,max] with theta.
+func zipfianPicker(t *testing.T, min, max uint64, theta float64) *config.GasPicker {
+	t.Helper()
+	var gp config.GasPicker
+	require.NoError(t, gp.UnmarshalJSON(fmt.Appendf(nil, `{"Name":"zipfian","Min":%d,"Max":%d,"theta":%f}`, min, max, theta)))
+	return &gp
+}
+
+func TestZipfianGasPickerStreamSeeds(t *testing.T) {
+	const seed, n = 42, 64
+
+	seededA := drawN(t, zipfianPicker(t, 20000, 30000, 0.9), newTestRng(seed), n)
+	seededB := drawN(t, zipfianPicker(t, 20000, 30000, 0.9), newTestRng(seed), n)
+	require.Equal(t, seededA, seededB, "same seed must reproduce the zipfian draw sequence")
+}
+
+func TestZipfianGasPickerRange(t *testing.T) {
+	const min, max = 20000, 30000
+	gp := zipfianPicker(t, min, max, 0.99)
+	rng := newTestRng(1)
+	for i := 0; i < 1000; i++ {
+		gas, err := gp.GenerateGas(rng)
+		require.NoError(t, err)
+		require.GreaterOrEqual(t, gas, uint64(min))
+		require.LessOrEqual(t, gas, uint64(max))
+	}
+}