work

2024/go/utils/memo/memo.go

@@ -0,0 +1,47 @@
+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
+}