proto: better cache codecs

proto/proto.go

@@ -3,12 +3,13 @@ package proto
 import (
 	"fmt"
 	"reflect"
-	"sync/atomic"
+	"sync"
 	"unsafe"
 )
 
 //go:generate go run ./gen/pointer
 //go:generate go run ./gen/required
+//go:generate syncmap -name structMap -pkg proto "map[unsafe.Pointer]*structInfo"
 
 func Size(v interface{}) int {
 	t, p := inspect(v)
@@ -97,20 +98,12 @@ type codec struct {
 	decode decodeFunc
 }
 
-var structInfoCache atomic.Value // map[unsafe.Pointer]*structInfo
-
-func loadCachedStruct(t reflect.Type) (*structInfo, map[unsafe.Pointer]*structInfo) {
-	cache, _ := structInfoCache.Load().(map[unsafe.Pointer]*structInfo)
-	return cache[pointer(t)], cache
-}
-
-func storeCachedStruct(newCache map[unsafe.Pointer]*structInfo) {
-	structInfoCache.Store(newCache)
-}
+var structInfoCache structMap // map[unsafe.Pointer]*structInfo
+var codecCache sync.Map       // map[reflect.Type]codec
 
 func cachedStructInfoOf(t reflect.Type) *structInfo {
-	c, oldCache := loadCachedStruct(t)
-	if c != nil {
+	c, ok := structInfoCache.Load(pointer(t))
+	if ok {
 		return c
 	}
 
@@ -119,18 +112,9 @@ func cachedStructInfoOf(t reflect.Type) *structInfo {
 		infos:  make(map[reflect.Type]*structInfo),
 	}
 
-	newCache := make(map[unsafe.Pointer]*structInfo, len(oldCache)+1)
-	for p, c := range oldCache {
-		newCache[p] = c
-	}
-
 	info := w.structInfo(t)
-	for t, info := range w.infos {
-		newCache[pointer(t)] = info
-	}
-
-	storeCachedStruct(newCache)
-	return info
+	actual, _ := structInfoCache.LoadOrStore(pointer(t), info)
+	return actual
 }
 
 // Bool stores v in a new bool value and returns a pointer to it.

proto/slice.go

@@ -2,51 +2,57 @@ package proto
 
 import (
 	"reflect"
+	"sync"
 	"unsafe"
 
 	. "github.com/RomiChan/protobuf/internal/runtime_reflect"
 )
 
-func sliceCodecOf(t reflect.Type, f structField, w *walker) *codec {
+var sliceMap sync.Map // map[*codec]*codec for slice
+
+func sliceCodecOf(t reflect.Type, c *codec, w *walker) *codec {
+	if loaded, ok := sliceMap.Load(c); ok {
+		return loaded.(*codec)
+	}
 	s := new(codec)
 	w.codecs[t] = s
 
-	s.size = sliceSizeFuncOf(t, &f)
-	s.encode = sliceEncodeFuncOf(t, &f)
-	s.decode = sliceDecodeFuncOf(t, &f)
-	return s
+	s.size = sliceSizeFuncOf(t, c)
+	s.encode = sliceEncodeFuncOf(t, c)
+	s.decode = sliceDecodeFuncOf(t, c)
+
+	actualCodec, _ := sliceMap.LoadOrStore(c, s)
+	return actualCodec.(*codec)
 }
 
-func sliceSizeFuncOf(t reflect.Type, f *structField) sizeFunc {
+func sliceSizeFuncOf(t reflect.Type, c *codec) sizeFunc {
 	elemSize := alignedSize(t.Elem())
-	codec := f.codec
 	return func(p unsafe.Pointer, sf *structField) int {
 		n := 0
 		if v := (*Slice)(p); v != nil {
 			for i := 0; i < v.Len(); i++ {
 				elem := v.Index(i, elemSize)
-				n += codec.size(elem, sf)
+				n += c.size(elem, sf)
 			}
 		}
 		return n
 	}
 }
 
-func sliceEncodeFuncOf(t reflect.Type, f *structField) encodeFunc {
+func sliceEncodeFuncOf(t reflect.Type, c *codec) encodeFunc {
 	elemSize := alignedSize(t.Elem())
-	codec := f.codec
 	return func(b []byte, p unsafe.Pointer, sf *structField) []byte {
 		if s := (*Slice)(p); s != nil {
 			for i := 0; i < s.Len(); i++ {
 				elem := s.Index(i, elemSize)
-				b = codec.encode(b, elem, sf)
+				b = c.encode(b, elem, sf)
 			}
 		}
 		return b
 	}
 }
 
-func sliceDecodeFuncOf(t reflect.Type, f *structField) decodeFunc {
+func sliceDecodeFuncOf(t reflect.Type, c *codec) decodeFunc {
 	elemType := t.Elem()
 	elemSize := alignedSize(elemType)
 	return func(b []byte, p unsafe.Pointer) (int, error) {
@@ -57,7 +63,7 @@ func sliceDecodeFuncOf(t reflect.Type, f *structField) decodeFunc {
 			*s = growSlice(elemType, s)
 		}
 
-		n, err := f.codec.decode(b, s.Index(i, elemSize))
+		n, err := c.decode(b, s.Index(i, elemSize))
 		if err == nil {
 			s.SetLen(i + 1)
 		}