PLT-455: Distribution config primitive + discriminator

config/config.go

@@ -61,10 +61,12 @@ type AccountConfig struct {
 
 // Scenario represents each scenario in the load configuration.
 type Scenario struct {
-	Name            string         `json:"name,omitempty"`
-	Weight          int            `json:"weight,omitempty"`
-	Accounts        *AccountConfig `json:"accounts,omitempty"`
-	GasPicker       *GasPicker     `json:"gasPicker,omitempty"`
-	GasFeeCapPicker *GasPicker     `json:"gasFeeCapPicker,omitempty"`
-	GasTipCapPicker *GasPicker     `json:"gasTipCapPicker,omitempty"`
+	Name             string         `json:"name,omitempty"`
+	Weight           int            `json:"weight,omitempty"`
+	Accounts         *AccountConfig `json:"accounts,omitempty"`
+	GasPicker        *GasPicker     `json:"gasPicker,omitempty"`
+	GasFeeCapPicker  *GasPicker     `json:"gasFeeCapPicker,omitempty"`
+	GasTipCapPicker  *GasPicker     `json:"gasTipCapPicker,omitempty"`
+	KeyDistribution  *Distribution  `json:"keyDistribution,omitempty"`
+	SizeDistribution *Distribution  `json:"sizeDistribution,omitempty"`
 }

config/distribution.go

@@ -0,0 +1,102 @@
+package config
+
+import (
+	"encoding/json"
+	"fmt"
+)
+
+var (
+	_ indexSampler = (*Distribution)(nil)
+	_ indexSampler = (*UniformDistribution)(nil)
+	_ indexSampler = (*ZipfianDistribution)(nil)
+)
+
+// indexSampler draws an index in [0, n) from some keyspace distribution.
+type indexSampler interface {
+	SampleIndex(n uint64) (uint64, error)
+}
+
+// Distribution is a tagged wrapper over a keyspace index distribution, selected
+// by a "Name" discriminator on the JSON wire format. The discriminator strings
+// ("uniform", "zipfian") and the "theta" parameter name are a frozen
+// saved-workload contract; do not rename them.
+type Distribution struct {
+	name     string
+	delegate indexSampler
+}
+
+func (d *Distribution) Name() string { return d.name }
+
+// SampleIndex delegates to the selected distribution. A zero-value (no Name)
+// Distribution samples nothing and returns 0.
+func (d *Distribution) SampleIndex(n uint64) (uint64, error) {
+	if d.delegate == nil {
+		return 0, nil
+	}
+	return d.delegate.SampleIndex(n)
+}
+
+func (d *Distribution) UnmarshalJSON(data []byte) error {
+	var temp struct {
+		Name string `json:"Name"`
+	}
+	if err := json.Unmarshal(data, &temp); err != nil {
+		return err
+	}
+	d.name = temp.Name
+	switch d.name {
+	case "":
+		return nil
+	case "uniform":
+		var uniform UniformDistribution
+		if err := json.Unmarshal(data, &uniform); err != nil {
+			return err
+		}
+		d.delegate = &uniform
+		return nil
+	case "zipfian":
+		var zipfian ZipfianDistribution
+		if err := json.Unmarshal(data, &zipfian); err != nil {
+			return err
+		}
+		if err := zipfian.validate(); err != nil {
+			return err
+		}
+		d.delegate = &zipfian
+		return nil
+	default:
+		return fmt.Errorf("unknown distribution name: %s", d.name)
+	}
+}
+
+// UniformDistribution draws each index with equal probability.
+type UniformDistribution struct{}
+
+func (UniformDistribution) SampleIndex(n uint64) (uint64, error) {
+	// PLT-460: implement the seeded uniform draw. Out of scope for PLT-455
+	// (wire format + validation only).
+	return 0, nil
+}
+
+// ZipfianDistribution draws indices with a Zipf-distributed skew controlled by
+// theta. theta == 0 is uniform; larger theta concentrates draws on low indices.
+type ZipfianDistribution struct {
+	Theta float64 `json:"theta"`
+}
+
+// zipfianThetaMax bounds theta to the range over which the YCSB precomputed-zeta
+// generator (PLT-460) is numerically well-behaved.
+const zipfianThetaMax = 1.0
+
+func (z *ZipfianDistribution) validate() error {
+	if z.Theta < 0 || z.Theta >= zipfianThetaMax {
+		return fmt.Errorf("zipfian theta out of range: %v (want [0, %v))", z.Theta, zipfianThetaMax)
+	}
+	return nil
+}
+
+func (z *ZipfianDistribution) SampleIndex(n uint64) (uint64, error) {
+	// PLT-460: implement the YCSB precomputed-zeta zipfian draw with a seeded
+	// RNG. Out of scope for PLT-455 (wire format + validation only).
+	return 0, nil
+}

config/distribution_test.go

@@ -0,0 +1,68 @@
+package config_test
+
+import (
+	"encoding/json"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/sei-protocol/sei-load/config"
+	"github.com/stretchr/testify/require"
+)
+
+func TestDistribution(t *testing.T) {
+	t.Parallel()
+	t.Run("empty", func(t *testing.T) {
+		var subject config.Distribution
+		require.NoError(t, subject.UnmarshalJSON([]byte(`{}`)))
+		require.Empty(t, subject.Name())
+		idx, err := subject.SampleIndex(100)
+		require.NoError(t, err)
+		require.Zero(t, idx)
+	})
+	t.Run("uniform", func(t *testing.T) {
+		var subject config.Distribution
+		require.NoError(t, subject.UnmarshalJSON([]byte(`{"Name":"uniform"}`)))
+		require.Equal(t, "uniform", subject.Name())
+	})
+	t.Run("zipfian", func(t *testing.T) {
+		var subject config.Distribution
+		require.NoError(t, subject.UnmarshalJSON([]byte(`{"Name":"zipfian","theta":0.99}`)))
+		require.Equal(t, "zipfian", subject.Name())
+	})
+	t.Run("zipfian_theta_below_range", func(t *testing.T) {
+		var subject config.Distribution
+		require.Error(t, subject.UnmarshalJSON([]byte(`{"Name":"zipfian","theta":-0.1}`)))
+	})
+	t.Run("zipfian_theta_above_range", func(t *testing.T) {
+		var subject config.Distribution
+		require.Error(t, subject.UnmarshalJSON([]byte(`{"Name":"zipfian","theta":1.0}`)))
+	})
+	t.Run("unknown", func(t *testing.T) {
+		var subject config.Distribution
+		require.Error(t, subject.UnmarshalJSON([]byte(`{"Name":"weibull"}`)))
+	})
+}
+
+// TestScenarioDistributionAdditive proves the new fields are additive: a profile
+// carrying no distribution fields parses unchanged and round-trips without
+// introducing any distribution keys.
+func TestScenarioDistributionAdditive(t *testing.T) {
+	t.Parallel()
+	path := filepath.Join("..", "profiles", "conflict.json")
+	original, err := os.ReadFile(path)
+	require.NoError(t, err)
+
+	var cfg config.LoadConfig
+	require.NoError(t, json.Unmarshal(original, &cfg))
+
+	for _, s := range cfg.Scenarios {
+		require.Nil(t, s.KeyDistribution, "no distribution expected in baseline profile")
+		require.Nil(t, s.SizeDistribution, "no distribution expected in baseline profile")
+	}
+
+	remarshaled, err := json.Marshal(cfg)
+	require.NoError(t, err)
+	require.NotContains(t, string(remarshaled), "keyDistribution")
+	require.NotContains(t, string(remarshaled), "sizeDistribution")
+}