Gih's Blog

只言片语

Another way to make memoization function in Go

2014-09-03 by gihnius, tagged as go

Cache any function calls in Go without using reflect, very simple idea, use a normal memory cache.

Here is what I do:

package memoize

import (
    "sync"
    "time"
)

type memo struct {
    Timeout time.Time
    Result interface{}
}

type MemoPool struct {
    Pool map[string]*memo
    mutex *sync.RWMutex
}

var mp *MemoPool

func init() {
    var m = map[string]*memo{}
    mp = &MemoPool{Pool: m, mutex: new(sync.RWMutex)}
}

// memorize result return from caller() block, timeout in N seconds
func Memoize(key string, caller func() interface{}, timeout uint) interface{} {
    if timeout == 0 {
        // do not memoize
        return caller()
    }
    mp.mutex.RLock()
    memoized := mp.Pool[key]
    mp.mutex.RUnlock()
    // reached timeout or not memoized
    if memoized == nil || memoized.Timeout.Before(time.Now()) {
        result := caller()
        if result != nil {
            duration := time.Duration(timeout) * time.Second
            mp.mutex.Lock()
            mp.Pool[key] = &memo{
                Timeout: time.Now().Add(duration),
                Result: result,
            }
            mp.mutex.Unlock()
        }
        return result
    }
    return memoized.Result
}

func UnMemoize(key string) {
    delete(mp.Pool, key)
}

func UnMemoizeAll() {
    for key, _ := range mp.Pool {
        delete(mp.Pool, key)
    }
}

Check out the code from github: https://github.com/gihnius/gomemoize

usage

import "memoize"
// cache a function call for TimeOut (>= 0) seconds
memoize.Memoize("Cache_Key", func() interface{} {
  cache_return_value, err := your_function(args)
  if err != nil {
    // handle err
    // won't cache if return nil
    return nil
  }
  return cache_return_value
}, TimeOut)

// save result
result := memoize.Memoize("Cache_Key", func() interface{} {
  // call your function(s)
  cache_return_value, err := your_function(args)
  if err != nil {
    // handle err
    // won't cache if return nil
    return nil
  }
  return cache_return_value
}, TimeOut).(your_function_return_type)

// cache a long time calc subroutine
func calc_sub(i int, s string) (string, err) {
  // ...
}
result := Memoize("calc_sub", func() interface{} {
  res, err := calc_sub(1000, "string to process")
  if err != nil {
    return nil // not memoized
  }
  return res
}, 60) // cache for 1 minutes

// re-fetch the memoized(cached) result with the same call.

// remove a cached call
memoize.UnMemoize(key_name)
// remove all cached calls
memoize.UnMemoizeAll()

test

git clone https://github.com/gihnius/gomemoize
cd gomemoize
./test.sh
# or
VERBOSE="-v" ./test.sh

It works well but not so functional.

More about  Memoization from Wikipedia.