internal/memoize/memoize.go

// Copyright 2019 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package memoize supports memoizing the return values of functions with
// idempotent results that are expensive to compute.
//
// The memoized result is returned again the next time the function is invoked.
// To prevent excessive memory use, the return values are only remembered
// for as long as they still have a user.
//
// To use this package, build a store and use it to acquire handles with the
// Bind method.
//
package memoize

import (
	"context"
	"runtime"
	"sync"
	"unsafe"

	"golang.org/x/tools/internal/xcontext"
)

// Store binds keys to functions, returning handles that can be used to access
// the functions results.
type Store struct {
	mu sync.Mutex
	// entries is the set of values stored.
	entries map[interface{}]uintptr
}

// Function is the type for functions that can be memoized.
// The result must be a pointer.
type Function func(ctx context.Context) interface{}

type state int

const (
	stateIdle = iota
	stateRunning
	stateCompleted
)

// Handle is returned from a store when a key is bound to a function.
// It is then used to access the results of that function.
//
// A Handle starts out in idle state, waiting for something to demand its
// evaluation. It then transitions into running state. While it's running,
// waiters tracks the number of Get calls waiting for a result, and the done
// channel is used to notify waiters of the next state transition. Once the
// evaluation finishes, value is set, state changes to completed, and done
// is closed, unblocking waiters. Alternatively, as Get calls are cancelled,
// they decrement waiters. If it drops to zero, the inner context is cancelled,
// computation is abandoned, and state resets to idle to start the process over
// again.
type Handle struct {
	store *Store
	key   interface{}

	mu    sync.Mutex
	state state
	// done is set in running state, and closed when exiting it.
	done chan struct{}
	// cancel is set in running state. It cancels computation.
	cancel context.CancelFunc
	// waiters is the number of Gets outstanding.
	waiters uint
	// the function that will be used to populate the value
	function Function
	// value is set in completed state.
	value interface{}
}

// Has returns true if they key is currently valid for this store.
func (s *Store) Has(key interface{}) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	_, found := s.entries[key]
	return found
}

// Bind returns a handle for the given key and function.
//
// Each call to bind will return the same handle if it is already bound.
// Bind will always return a valid handle, creating one if needed.
// Each key can only have one handle at any given time.
// The value will be held for as long as the handle is, once it has been
// generated.
// Bind does not cause the value to be generated.
func (s *Store) Bind(key interface{}, function Function) *Handle {
	// panic early if the function is nil
	// it would panic later anyway, but in a way that was much harder to debug
	if function == nil {
		panic("the function passed to bind must not be nil")
	}
	// check if we already have the key
	s.mu.Lock()
	defer s.mu.Unlock()
	h := s.get(key)
	if h != nil {
		// we have a handle already, just return it
		return h
	}
	// we have not seen this key before, add a new entry
	if s.entries == nil {
		s.entries = make(map[interface{}]uintptr)
	}
	h = &Handle{
		store:    s,
		key:      key,
		function: function,
	}
	// now add the weak reference to the handle into the map
	s.entries[key] = uintptr(unsafe.Pointer(h))
	// add the deletion the entry when the handle is garbage collected
	runtime.SetFinalizer(h, release)
	return h
}

// Find returns the handle associated with a key, if it is bound.
//
// It cannot cause a new handle to be generated, and thus may return nil.
func (s *Store) Find(key interface{}) *Handle {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.get(key)
}

// Cached returns the value associated with a key.
//
// It cannot cause the value to be generated.
// It will return the cached value, if present.
func (s *Store) Cached(key interface{}) interface{} {
	h := s.Find(key)
	if h == nil {
		return nil
	}
	return h.Cached()
}

//go:nocheckptr
// nocheckptr because: https://github.com/golang/go/issues/35125#issuecomment-545671062
func (s *Store) get(key interface{}) *Handle {
	// this must be called with the store mutex already held
	e, found := s.entries[key]
	if !found {
		return nil
	}
	return (*Handle)(unsafe.Pointer(e))
}

// Cached returns the value associated with a handle.
//
// It will never cause the value to be generated.
// It will return the cached value, if present.
func (h *Handle) Cached() interface{} {
	h.mu.Lock()
	defer h.mu.Unlock()
	if h.state == stateCompleted {
		return h.value
	}
	return nil
}

// Get returns the value associated with a handle.
//
// If the value is not yet ready, the underlying function will be invoked.
// This activates the handle, and it will remember the value for as long as it exists.
// If ctx is cancelled, Get returns nil.
func (h *Handle) Get(ctx context.Context) interface{} {
	if ctx.Err() != nil {
		return nil
	}
	h.mu.Lock()
	switch h.state {
	case stateIdle:
		return h.run(ctx)
	case stateRunning:
		return h.wait(ctx)
	case stateCompleted:
		defer h.mu.Unlock()
		return h.value
	default:
		panic("unknown state")
	}
}

// run starts h.function and returns the result. h.mu must be locked.
func (h *Handle) run(ctx context.Context) interface{} {
	childCtx, cancel := context.WithCancel(xcontext.Detach(ctx))
	h.cancel = cancel
	h.state = stateRunning
	h.done = make(chan struct{})
	go func() {
		v := h.function(childCtx)
		if childCtx.Err() != nil {
			return
		}

		h.mu.Lock()
		defer h.mu.Unlock()
		// It's theoretically possible that the handle has been cancelled out
		// of the run that started us, and then started running again since we
		// checked childCtx above. Even so, that should be harmless, since each
		// run should produce the same results.
		if h.state != stateRunning {
			return
		}
		h.value = v
		h.function = nil
		h.state = stateCompleted
		close(h.done)
	}()

	return h.wait(ctx)
}

// wait waits for the value to be computed, or ctx to be cancelled. h.mu must be locked.
func (h *Handle) wait(ctx context.Context) interface{} {
	h.waiters++
	done := h.done
	h.mu.Unlock()

	select {
	case <-done:
		h.mu.Lock()
		defer h.mu.Unlock()
		if h.state == stateCompleted {
			return h.value
		}
		return nil
	case <-ctx.Done():
		h.mu.Lock()
		defer h.mu.Unlock()
		h.waiters--
		if h.waiters == 0 && h.state == stateRunning {
			h.cancel()
			close(h.done)
			h.state = stateIdle
			h.done = nil
			h.cancel = nil
		}
		return nil
	}
}

func release(p interface{}) {
	h := p.(*Handle)
	h.store.mu.Lock()
	defer h.store.mu.Unlock()
	// there is a small gap between the garbage collector deciding that the handle
	// is liable for collection and the finalizer being called
	// if the handle is recovered during that time, you will end up with a valid
	// handle that no longer has an entry in the map, and that no longer has a
	// finalizer associated with it, but that is okay.
	delete(h.store.entries, h.key)
}