package dataext

import "sync"

type MutexSet[T comparable] struct {
	master sync.RWMutex
	locks  map[T]*sync.RWMutex
}

func NewMutexSet[T comparable]() *MutexSet[T] {
	return &MutexSet[T]{
		master: sync.RWMutex{},
		locks:  make(map[T]*sync.RWMutex),
	}
}

func (ms *MutexSet[T]) get(key T) *sync.RWMutex {
	ms.master.RLock()
	if v, ok := ms.locks[key]; ok {
		ms.master.RUnlock()
		return v
	}
	ms.master.RUnlock()

	// ---------

	ms.master.Lock()
	defer ms.master.Unlock()

	if v, ok := ms.locks[key]; ok {
		return v
	}

	m := &sync.RWMutex{}
	ms.locks[key] = m
	return m
}

func (ms *MutexSet[T]) Lock(key T) {
	ms.get(key).Lock()
}

func (ms *MutexSet[T]) Unlock(key T) {
	ms.get(key).Unlock()
}

func (ms *MutexSet[T]) RLock(key T) {
	ms.get(key).RLock()
}

func (ms *MutexSet[T]) RUnlock(key T) {
	ms.get(key).RUnlock()
}