Merge branch 'lukeshu/fast'

1 files changed, 19 insertions, 18 deletions

diff --git a/lib/binstruct/structs.go b/lib/binstruct/structs.go
index 7eea600..91bfec7 100644
--- a/lib/binstruct/structs.go
+++ b/lib/binstruct/structs.go
@@ -10,6 +10,8 @@ import (
 	"strconv"
 	"strings"
 
+	"git.lukeshu.com/go/typedsync"
+
 	"git.lukeshu.com/btrfs-progs-ng/lib/binstruct/binutil"
 )
 
@@ -67,7 +69,8 @@ type structHandler struct {
 }
 
 type structField struct {
-	name string
+	name          string
+	isUnmarshaler bool
 	tag
 }
 
@@ -80,7 +83,7 @@ func (sh structHandler) Unmarshal(dat []byte, dst reflect.Value) (int, error) {
 		if field.skip {
 			continue
 		}
-		_n, err := Unmarshal(dat[n:], dst.Field(i).Addr().Interface())
+		_n, err := unmarshal(dat[n:], dst.Field(i), field.isUnmarshaler)
 		if err != nil {
 			if _n >= 0 {
 				n += _n
@@ -164,8 +167,9 @@ func genStructHandler(structInfo reflect.Type) (structHandler, error) {
 		curOffset += fieldTag.siz
 
 		ret.fields = append(ret.fields, structField{
-			name: fieldInfo.Name,
-			tag:  fieldTag,
+			name:          fieldInfo.Name,
+			isUnmarshaler: reflect.PtrTo(fieldInfo.Type).Implements(unmarshalerType),
+			tag:           fieldTag,
 		})
 	}
 	ret.Size = curOffset
@@ -178,21 +182,18 @@ func genStructHandler(structInfo reflect.Type) (structHandler, error) {
 	return ret, nil
 }
 
-var structCache = make(map[reflect.Type]structHandler)
+var structCache typedsync.CacheMap[reflect.Type, structHandler]
 
 func getStructHandler(typ reflect.Type) structHandler {
-	h, ok := structCache[typ]
-	if ok {
+	ret, _ := structCache.LoadOrCompute(typ, func(typ reflect.Type) structHandler {
+		h, err := genStructHandler(typ)
+		if err != nil {
+			panic(&InvalidTypeError{
+				Type: typ,
+				Err:  err,
+			})
+		}
 		return h
-	}
-
-	h, err := genStructHandler(typ)
-	if err != nil {
-		panic(&InvalidTypeError{
-			Type: typ,
-			Err:  err,
-		})
-	}
-	structCache[typ] = h
-	return h
+	})
+	return ret
 }