Generic Memoization

Author: coreyog

.vscode/launch.json

@@ -0,0 +1,15 @@
+{
+  // Use IntelliSense to learn about possible attributes.
+  // Hover to view descriptions of existing attributes.
+  // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+  "version": "0.2.0",
+  "configurations": [
+    {
+      "name": "Run Tests",
+      "type": "go",
+      "request": "launch",
+      "mode": "test",
+      "program": "${workspaceFolder}"
+    }
+  ]
+}

README.md

@@ -1 +1,5 @@
-# memoize2
+# Memoize
+
+A generic memoization library.
+
+Needs testing.

errors.go

@@ -0,0 +1,9 @@
+package memoize
+
+import "github.com/pkg/errors"
+
+var (
+	ErrNotAFunc       = errors.New("not a function")
+	ErrMissingArgs    = errors.New("target function must accept at least 1 argument")
+	ErrMissingReturns = errors.New("target function must return at least 1 value")
+)

go.mod

@@ -0,0 +1,14 @@
+module github.com/coreyog/memoize
+
+go 1.18
+
+require (
+	github.com/pkg/errors v0.9.1
+	github.com/stretchr/testify v1.8.1
+)
+
+require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
+)

go.sum

@@ -0,0 +1,19 @@
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
+github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
+github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
+github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
+github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk=
+github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

memoize.go

@@ -0,0 +1,75 @@
+package memoize
+
+import "reflect"
+
+func Memo[T any](fn T) (m T, err error) {
+	fnt := reflect.TypeOf(fn)
+
+	if fnt.Kind() != reflect.Func {
+		return fn, ErrNotAFunc
+	}
+
+	if fnt.NumIn() == 0 {
+		return fn, ErrMissingArgs
+	}
+
+	if fnt.NumOut() == 0 {
+		return fn, ErrMissingReturns
+	}
+
+	fnv := reflect.ValueOf(fn)
+
+	cacheRoot := map[interface{}]interface{}{}
+
+	ret := reflect.MakeFunc(fnt, func(args []reflect.Value) (results []reflect.Value) {
+		cResults := fillAndCheck(cacheRoot, args)
+		if cResults == nil {
+			results = fnv.Call(args)
+			fillAndSet(cacheRoot, args, results)
+		} else {
+			results = cResults
+		}
+
+		return results
+	})
+
+	return ret.Interface().(T), nil
+}
+
+func fillAndCheck(cacheRoot map[interface{}]interface{}, args []reflect.Value) (results []reflect.Value) {
+	var m, next interface{}
+	var ok bool
+	m = cacheRoot
+
+	for _, arg := range args {
+		next, ok = m.(map[interface{}]interface{})[arg.Interface()]
+		if !ok {
+			next = map[interface{}]interface{}{}
+			m.(map[interface{}]interface{})[arg.Interface()] = next
+		}
+
+		m = next
+	}
+
+	results, ok = m.([]reflect.Value)
+
+	if !ok || results == nil {
+		return nil
+	}
+
+	return results
+}
+
+func fillAndSet(cacheRoot map[interface{}]interface{}, args []reflect.Value, results []reflect.Value) {
+	var m interface{}
+	var prev map[interface{}]interface{}
+
+	m = cacheRoot
+
+	for _, arg := range args {
+		prev = m.(map[interface{}]interface{})
+		m = prev[arg.Interface()]
+	}
+
+	prev[args[len(args)-1].Interface()] = results
+}

memoize_test.go

@@ -0,0 +1,76 @@
+package memoize
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestBadFuncErrors(t *testing.T) {
+	_, err := Memo("not a func")
+	assert.Equal(t, err, ErrNotAFunc)
+
+	noArgs := func() bool { return true }
+	_, err = Memo(noArgs)
+	assert.Equal(t, err, ErrMissingArgs)
+
+	noReturns := func(x int) {}
+	_, err = Memo(noReturns)
+	assert.Equal(t, err, ErrMissingReturns)
+}
+
+func TestMultiCall(t *testing.T) {
+	called := 0
+	square := func(x int) int {
+		called++
+		return x * x
+	}
+
+	m, err := Memo(square)
+	assert.NoError(t, err)
+
+	for i := 0; i < 100; i++ {
+		mResult := m(i)
+		m2Result := m(i)
+		squareResult := square(i)
+
+		assert.Equal(t, mResult, m2Result)
+		assert.Equal(t, mResult, squareResult)
+	}
+
+	assert.Equal(t, called, 200)
+}
+
+func TestSimple(t *testing.T) {
+	called := 0
+	work := func(x int) int {
+		called++
+		return x
+	}
+
+	m, err := Memo(work)
+	assert.NoError(t, err)
+
+	for i := 0; i < 1000; i++ {
+		m(0)
+	}
+
+	assert.Equal(t, 1, called)
+}
+
+func TestManyArgsManyRets(t *testing.T) {
+	called := 0
+	work := func(x int, y string, z float64) (int, string, float64) {
+		called++
+		return x, y, z
+	}
+
+	m, err := Memo(work)
+	assert.NoError(t, err)
+
+	for i := 0; i < 1000; i++ {
+		m(0, "x", 3.14)
+	}
+
+	assert.Equal(t, 1, called)
+}