minstrel/internal/recommendation/shuffle.go

package recommendation

import (
	"sort"
	"time"

	"github.com/jackc/pgx/v5/pgtype"

	"git.fabledsword.com/bvandeusen/minstrel/internal/db/dbq"
)

// Candidate pairs a track with the inputs needed to score it.
type Candidate struct {
	Track  dbq.Track
	Inputs ScoringInputs
}

// DiversityCaps bounds how much of one selection a single artist or album
// may occupy. A zero value means "no cap", which is what Shuffle did
// unconditionally before #3882.
type DiversityCaps struct {
	MaxPerArtist int
	MaxPerAlbum  int
}

// radioCapArtistPer25 / radioCapAlbumPer25 hold the caps at the same
// PROPORTION the system mixes already use — 3 artist / 2 album tracks in a
// 25-track mix — so a 20-track radio and a 200-track one feel alike rather
// than one of them being effectively uncapped.
//
// A fixed count cannot do that. Three-per-artist is a reasonable 12% of a
// 25-track mix and an absurd 1.5% of a 200-track radio, where it would put
// the selection permanently in the relaxation path below and quietly undo
// the cap it was meant to enforce.
const (
	radioCapArtistPer25 = 3
	radioCapAlbumPer25  = 2
)

// RadioDiversityCaps scales the diversity caps to the requested radio
// length. Floors of 2 and 1 keep a very short radio from being capped into
// a single track per artist, which would be its own kind of wrong.
func RadioDiversityCaps(limit int) DiversityCaps {
	artist := limit * radioCapArtistPer25 / 25
	if artist < 2 {
		artist = 2
	}
	album := limit * radioCapAlbumPer25 / 25
	if album < 1 {
		album = 1
	}
	return DiversityCaps{MaxPerArtist: artist, MaxPerAlbum: album}
}

// Shuffle scores each candidate, sorts descending by score, and returns the
// top `limit` candidates, preferring artist/album diversity. limit <= 0
// returns nil; nil input returns nil. Pure — no IO, no global state beyond
// the rng callback.
//
// THE COUNT IS NEVER REDUCED BY THE CAPS. Selection runs in two passes: the
// first takes candidates that fit under the caps, and the second fills any
// remaining slots from those the first pass skipped, still in score order.
// So the result holds min(limit, len(candidates)) either way — the caps
// change WHICH tracks are chosen, never HOW MANY.
//
// That two-pass shape is the whole design, and a hard cap would have been
// the easy mistake. Radio asks for 50 tracks by default and 200 at most; a
// pool concentrated on a few artists would return six tracks and call it a
// radio. Rule 131's principle — degrade, never vanish — applies past the
// system mixes it was written for.
//
// Before #3882 there was no cap here at all, while discover.go,
// you_might_like.go and home.go all had one. That asymmetry is what let a
// radio session come back entirely from a single artist: nothing between
// the pool and the output bounded any artist's share, so a pool dominated
// by one artist produced an output dominated by it too.
func Shuffle(
	candidates []Candidate,
	weights ScoringWeights,
	now time.Time,
	rng func() float64,
	limit int,
	caps DiversityCaps,
) []Candidate {
	if len(candidates) == 0 || limit <= 0 {
		return nil
	}
	scored := make([]struct {
		c     Candidate
		score float64
	}, len(candidates))
	for i, c := range candidates {
		scored[i].c = c
		scored[i].score = Score(c.Inputs, weights, now, rng)
	}
	sort.Slice(scored, func(i, j int) bool {
		return scored[i].score > scored[j].score
	})
	if limit > len(scored) {
		limit = len(scored)
	}

	out := make([]Candidate, 0, limit)
	deferred := make([]Candidate, 0, len(scored)-limit)
	artistCount := map[pgtype.UUID]int{}
	albumCount := map[pgtype.UUID]int{}
	// One song on several releases is played once (M498 #5438): after its
	// best-scoring copy, the others are dropped, not held back.
	songs := map[pgtype.UUID]bool{}
	repeat := func(c Candidate) bool { return c.Track.SongKey.Valid && songs[c.Track.SongKey] }
	take := func(c Candidate) {
		if c.Track.SongKey.Valid {
			songs[c.Track.SongKey] = true
		}
		out = append(out, c)
	}

	for _, s := range scored {
		if len(out) == limit {
			break
		}
		if repeat(s.c) {
			continue
		}
		overArtist := caps.MaxPerArtist > 0 && artistCount[s.c.Track.ArtistID] >= caps.MaxPerArtist
		overAlbum := caps.MaxPerAlbum > 0 && albumCount[s.c.Track.AlbumID] >= caps.MaxPerAlbum
		if overArtist || overAlbum {
			// Held back, not discarded — pass two may still need it.
			deferred = append(deferred, s.c)
			continue
		}
		artistCount[s.c.Track.ArtistID]++
		albumCount[s.c.Track.AlbumID]++
		take(s.c)
	}

	// Pass two: the caps could not fill the request, so relax them rather
	// than hand back a short radio. Still score order, so the best of the
	// held-back candidates go first.
	for _, c := range deferred {
		if len(out) == limit {
			break
		}
		if repeat(c) {
			continue
		}
		take(c)
	}
	return out
}