// Copyright (C) 2023  Luke Shumaker <lukeshu@lukeshu.com>
//
// SPDX-License-Identifier: GPL-2.0-or-later

package syncutil

import (
	"context"

	"git.lukeshu.com/go/containers/typedsync"
)

// A MapOnce wraps a Map (typically a
// [git.lukeshu.com/go/containers/typedsync.Map], but possibly other
// backends), in order to provide a LoadOrDo method that allows
// missing values to be constructed without duplicating
// work.
//
// It is similar to a [golang.org/x/sync/singleflight.Group], but
// values are persistent.
//
// It is similar to a map of [sync.Once] values, but without concerns
// about initialization.
//
// [git.lukeshu.com/go/containers/typedsync.Map]: https://pkg.go.dev/git.lukeshu.com/go/containers/typedsync#Map
// [golang.org/x/sync/singleflight.Group]: https://pkg.go.dev/golang.org/x/sync/singleflight#Group
// [sync.Once]: https://pkg.go.dev/sync#Once
type MapOnce[K mapkey, V any, M Map[K, *MapOnceVal[V]]] struct {
	// The techniques used by MapOnce are similar to the
	// techniques used by encoding/json's internal type cache.

	Inner M

	// Because LoadOrDo needs a "new" empty *MapOnceVal[V] that
	// will be immediately discarded for "Load" operations, it is
	// worth having a pool so that should-be-fast "Load"s don't
	// trigger slow allocations and create GC pressure.
	pool typedsync.Pool[*MapOnceVal[V]]
}

// A Map describes the parallel-safe "map" storage required by a
// MapOnce.
//
// The canonical Map implementation is
// git.lukeshu.com/go/containers/typedsync.Map.
type Map[K mapkey, V any] interface {
	Delete(K)
	LoadOrStore(K, V) (actual V, loaded bool)
}

// A MapOnceVal is a values that MapOnce stores in to its underlying
// Map.
type MapOnceVal[V any] struct {
	V V
	c chan struct{}
}

// Delete removes the value for a key.  If the value for that key is
// actively being constructed by LoadOrDo, this immediately removes
// the partial value from the underlying map, but outstanding LoadOrDo
// calls will still behave as if Delete had not been called.
func (m *MapOnce[K, V, M]) Delete(key K) {
	m.Inner.Delete(key)
}

// LoadOrDo returns the existing value stored for a key, if present.
// If not present, it calls the "do" function to construct the value,
// then stores and returns that value.  The "loaded" result is true if
// the value was loaded, false if constructed.  If a prior call to
// LoadOrDo is still constructing the value for that key, a latter
// call blocks until the constructing is complete, and then returns
// the initial call's value.
func (m *MapOnce[K, V, M]) LoadOrDo(key K, do func(K) V) (actual V, loaded bool) {
	_value, _ := m.pool.Get()
	if _value == nil {
		_value = &MapOnceVal[V]{
			c: make(chan struct{}),
		}
	}
	_actual, loaded := m.Inner.LoadOrStore(key, _value)
	if loaded {
		m.pool.Put(_value)
		<-_actual.c
	} else {
		_actual.V = do(key)
		close(_actual.c)
	}
	return _actual.V, loaded
}

// TryLoadOrDo is like LoadOrDo, but obeys context cancellation.  If a
// call is cancelled, the call to "do" continues running in a separate
// goroutine, in case other LoadOrDo calls are waiting on it.  If a
// call is cancelled, the error from ctx.Err() is returned, otherwise
// err is nil.
func (m *MapOnce[K, V, M]) TryLoadOrDo(ctx context.Context, key K, do func(K) V) (actual V, loaded bool, err error) {
	_value, _ := m.pool.Get()
	if _value == nil {
		_value = &MapOnceVal[V]{
			c: make(chan struct{}),
		}
	}
	_actual, loaded := m.Inner.LoadOrStore(key, _value)
	if loaded {
		m.pool.Put(_value)
	} else {
		go func() {
			_actual.V = do(key)
			close(_actual.c)
		}()
	}
	select {
	case <-ctx.Done():
		var zero V
		return zero, false, ctx.Err() //nolint:wrapcheck // We're too low level for that to be useful.
	case <-_actual.c:
		return _actual.V, loaded, nil
	}
}