sampling_strategy.go

package dasmon

import (
	"math/rand"
	"time"

	"github.com/ethp2p/dasmon/store"
	"github.com/libp2p/go-libp2p/core/peer"
)

// Sampler selects slots and columns to test based on coverage history
type Sampler struct {
	config       SamplingConfig
	slotSelector SlotSelector
}

// NewSampler creates a new coverage-aware sampler
func NewSampler(config SamplingConfig) (*Sampler, error) {
	// Parse slot bias into a SlotSelector
	slotSelector, err := ParseSlotSelector(config.SlotBias)
	if err != nil {
		return nil, err
	}

	return &Sampler{
		config:       config,
		slotSelector: slotSelector,
	}, nil
}

// SamplingContext contains all information needed to select samples
type SamplingContext struct {
	PeerID          peer.ID
	CustodyRange    Range[uint64]
	CustodyColumns  []uint64
	MaxSlots        uint64
	MaxColumns      int
	BlockFilter     Range[uint64]
	CommitmentRange Range[uint64]
	History         *store.PeerHistory
}

// SamplingPlan specifies which slots and columns to test
type SamplingPlan struct {
	StartSlot uint64
	EndSlot   uint64
	Columns   []uint64
}

// SelectSamples chooses slots and columns to test for a peer
func (s *Sampler) SelectSamples(ctx SamplingContext) *SamplingPlan {
	// Find overlap of all constraints
	overlap, ok := ctx.CustodyRange.Intersect(ctx.BlockFilter)
	if !ok || !overlap.Valid() {
		return nil
	}

	overlap, ok = overlap.Intersect(ctx.CommitmentRange)
	if !ok || !overlap.Valid() {
		return nil
	}

	// Get history window for filtering
	historyWindow := time.Duration(s.config.HistoryWindow)

	// Select slot range
	startSlot, numSlots := s.selectSlots(ctx, overlap, historyWindow)
	if numSlots == 0 {
		return nil
	}

	// Select columns for the chosen slot range
	columns := s.selectColumns(ctx, startSlot, historyWindow)
	if len(columns) == 0 {
		return nil
	}

	return &SamplingPlan{
		StartSlot: startSlot,
		EndSlot:   startSlot + numSlots,
		Columns:   columns,
	}
}

// selectSlots chooses a slot range to test, preferring untested slots
func (s *Sampler) selectSlots(ctx SamplingContext, validRange Range[uint64], historyWindow time.Duration) (startSlot uint64, numSlots uint64) {
	if !ctx.History.IsLoaded() {
		// History not loaded yet, use slot selector on full range
		startSlot = s.slotSelector.SelectSlot(validRange)
		numSlots = s.determineNumSlots(ctx.MaxSlots, validRange, startSlot)
		return
	}

	if s.config.PreferUntested {
		// Try to find completely untested slots first
		untestedSlots := ctx.History.GetUntestedSlots(validRange.Start, validRange.End, historyWindow)

		if len(untestedSlots) > 0 {
			// Pick a random untested slot
			idx := randInt(len(untestedSlots))
			startSlot = untestedSlots[idx]
			numSlots = s.determineNumSlots(ctx.MaxSlots, validRange, startSlot)
			return
		}
	}

	// Fallback: use slot selector strategy (may select already-tested slots)
	startSlot = s.slotSelector.SelectSlot(validRange)
	numSlots = s.determineNumSlots(ctx.MaxSlots, validRange, startSlot)
	return
}

// selectColumns chooses which columns to test, preferring untested ones
func (s *Sampler) selectColumns(ctx SamplingContext, slot uint64, historyWindow time.Duration) []uint64 {
	if !ctx.History.IsLoaded() {
		// History not loaded yet, random selection
		return randomSample(ctx.CustodyColumns, ctx.MaxColumns)
	}

	// Get untested columns for this slot (considering history window)
	untestedCols := ctx.History.GetUntestedColumns(slot, ctx.CustodyColumns, historyWindow)

	if len(untestedCols) > 0 && s.config.PreferUntested {
		// Prefer untested columns
		return randomSample(untestedCols, ctx.MaxColumns)
	}

	// Fallback: random selection from all custody columns
	return randomSample(ctx.CustodyColumns, ctx.MaxColumns)
}

// determineNumSlots calculates how many slots to request
func (s *Sampler) determineNumSlots(maxSlots uint64, validRange Range[uint64], startSlot uint64) uint64 {
	if maxSlots == 0 {
		return 1
	}

	// Don't exceed the valid range
	slotsAvailable := validRange.End - startSlot
	if slotsAvailable == 0 {
		return 1
	}

	actualMax := min(maxSlots, slotsAvailable)
	if actualMax == 1 {
		return 1
	}

	// Random number of slots between 1 and actualMax
	return 1 + uint64(randInt(int(actualMax)))
}

// randomSample returns a random sample of up to n items from the input slice
func randomSample(items []uint64, n int) []uint64 {
	if n >= len(items) {
		// Return all items, shuffled
		result := make([]uint64, len(items))
		copy(result, items)
		shuffleUint64(result)
		return result
	}

	// Shuffle and take first n
	shuffled := make([]uint64, len(items))
	copy(shuffled, items)
	shuffleUint64(shuffled)
	return shuffled[:n]
}

// shuffleUint64 performs an in-place Fisher-Yates shuffle
func shuffleUint64(slice []uint64) {
	for i := len(slice) - 1; i > 0; i-- {
		j := randInt(i + 1)
		slice[i], slice[j] = slice[j], slice[i]
	}
}

// randInt returns a random integer in [0, n)
func randInt(n int) int {
	if n <= 0 {
		return 0
	}
	return int(rand.Uint64() % uint64(n))
}