package memo

import "sync"

type Func func(key interface{}) interface{}

type result struct {
	value interface{}
}

type Memo struct {
	f     Func
	cache map[interface{}]result
	mu    sync.RWMutex // Allows concurrent reads.
}

func New(f Func) *Memo {
	return &Memo{
		f:     f,
		cache: make(map[interface{}]result),
	}
}

func (memo *Memo) Get(key interface{}) interface{} {
	// First, try to read the cache using a read lock.
	memo.mu.RLock()
	res, ok := memo.cache[key]
	memo.mu.RUnlock()
	if ok {
		return res.value
	}

	// Compute the result without holding the lock.
	computed := memo.f(key)

	// Now acquire a write lock to update the cache.
	memo.mu.Lock()
	// Double-check: another goroutine may have stored the result in the meantime.
	res, ok = memo.cache[key]
	if !ok {
		res.value = computed
		memo.cache[key] = res
	}
	memo.mu.Unlock()

	return res.value
}