synclock_map.go

package mapsutil

import (
	"sync"
	"sync/atomic"

	errorutil "github.com/projectdiscovery/utils/errors"
)

var (
	ErrReadOnly = errorutil.New("map is currently in read-only mode").WithTag("syncLockMap")
)

// SyncLock adds sync and lock capabilities to generic map
type SyncLockMap[K, V comparable] struct {
	ReadOnly atomic.Bool
	mu       sync.RWMutex
	Map      Map[K, V]
}

type SyncLockMapOption[K, V comparable] func(slm *SyncLockMap[K, V])

func WithMap[K, V comparable](m Map[K, V]) SyncLockMapOption[K, V] {
	return func(slm *SyncLockMap[K, V]) {
		slm.Map = m
	}
}

// NewSyncLockMap creates a new SyncLockMap.
// If an existing map is provided, it is used; otherwise, a new map is created.
func NewSyncLockMap[K, V comparable](options ...SyncLockMapOption[K, V]) *SyncLockMap[K, V] {
	slm := &SyncLockMap[K, V]{}

	for _, option := range options {
		option(slm)
	}

	if slm.Map == nil {
		slm.Map = make(Map[K, V])
	}

	return slm
}

// Lock the current map to read-only mode
func (s *SyncLockMap[K, V]) Lock() {
	s.ReadOnly.Store(true)
}

// Unlock the current map
func (s *SyncLockMap[K, V]) Unlock() {
	s.ReadOnly.Store(false)
}

// Set an item with syncronous access
func (s *SyncLockMap[K, V]) Set(k K, v V) error {
	if s.ReadOnly.Load() {
		return ErrReadOnly
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	s.Map[k] = v

	return nil
}

// Get an item with syncronous access
func (s *SyncLockMap[K, V]) Get(k K) (V, bool) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	v, ok := s.Map[k]

	return v, ok
}

// Get an item with syncronous access
func (s *SyncLockMap[K, V]) Delete(k K) {
	s.mu.Lock()
	defer s.mu.Unlock()

	delete(s.Map, k)
}

// Iterate with a callback function synchronously
func (s *SyncLockMap[K, V]) Iterate(f func(k K, v V) error) error {
	s.mu.RLock()
	defer s.mu.RUnlock()

	for k, v := range s.Map {
		if err := f(k, v); err != nil {
			return err
		}
	}
	return nil
}

// Clone creates a new SyncLockMap with the same values
func (s *SyncLockMap[K, V]) Clone() *SyncLockMap[K, V] {
	s.mu.Lock()
	defer s.mu.Unlock()

	smap := &SyncLockMap[K, V]{
		ReadOnly: atomic.Bool{},
		mu:       sync.RWMutex{},
		Map:      s.Map.Clone(),
	}
	smap.ReadOnly.Store(s.ReadOnly.Load())
	return smap
}

// Has checks if the current map has the provided key
func (s *SyncLockMap[K, V]) Has(key K) bool {
	s.mu.RLock()
	defer s.mu.RUnlock()

	return s.Map.Has(key)
}

// IsEmpty checks if the current map is empty
func (s *SyncLockMap[K, V]) IsEmpty() bool {
	s.mu.RLock()
	defer s.mu.RUnlock()

	return s.Map.IsEmpty()
}

// IsEmpty checks if the current map is empty
func (s *SyncLockMap[K, V]) Clear() bool {
	s.mu.Lock()
	defer s.mu.Unlock()

	return s.Map.Clear()
}

// GetKeywithValue returns the first key having value
func (s *SyncLockMap[K, V]) GetKeyWithValue(value V) (K, bool) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	return s.Map.GetKeyWithValue(value)
}

// Merge the current map with the provided one
func (s *SyncLockMap[K, V]) Merge(n map[K]V) error {
	if s.ReadOnly.Load() {
		return ErrReadOnly
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	s.Map.Merge(n)
	return nil
}

// GetAll returns Copy of the current map
func (s *SyncLockMap[K, V]) GetAll() Map[K, V] {
	s.mu.RLock()
	defer s.mu.RUnlock()

	return s.Map.Clone()
}