Merge branch 'lukeshu/fast'

5 files changed, 132 insertions, 340 deletions

diff --git a/lib/binstruct/binint.go b/lib/binstruct/binint.go
deleted file mode 100644
index 302ab8d..0000000
--- a/lib/binstruct/binint.go
+++ /dev/null
@@ -1,39 +0,0 @@
-// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
-//
-// SPDX-License-Identifier: GPL-2.0-or-later
-
-package binstruct
-
-import (
-	"reflect"
-
-	"git.lukeshu.com/btrfs-progs-ng/lib/binstruct/binint"
-)
-
-type (
-	U8    = binint.U8
-	U16le = binint.U16le
-	U32le = binint.U32le
-	U64le = binint.U64le
-	U16be = binint.U16be
-	U32be = binint.U32be
-	U64be = binint.U64be
-	I8    = binint.I8
-	I16le = binint.I16le
-	I32le = binint.I32le
-	I64le = binint.I64le
-	I16be = binint.I16be
-	I32be = binint.I32be
-	I64be = binint.I64be
-)
-
-var intKind2Type = map[reflect.Kind]reflect.Type{
-	reflect.Uint8:  reflect.TypeOf(U8(0)),
-	reflect.Int8:   reflect.TypeOf(I8(0)),
-	reflect.Uint16: reflect.TypeOf(U16le(0)),
-	reflect.Int16:  reflect.TypeOf(I16le(0)),
-	reflect.Uint32: reflect.TypeOf(U32le(0)),
-	reflect.Int32:  reflect.TypeOf(I32le(0)),
-	reflect.Uint64: reflect.TypeOf(U64le(0)),
-	reflect.Int64:  reflect.TypeOf(I64le(0)),
-}
diff --git a/lib/binstruct/binint/builtins.go b/lib/binstruct/binint/builtins.go
deleted file mode 100644
index cfd0fc2..0000000
--- a/lib/binstruct/binint/builtins.go
+++ /dev/null
@@ -1,258 +0,0 @@
-// Copyright (C) 2022-2023  Luke Shumaker <lukeshu@lukeshu.com>
-//
-// SPDX-License-Identifier: GPL-2.0-or-later
-
-package binint
-
-import (
-	"encoding/binary"
-
-	"git.lukeshu.com/btrfs-progs-ng/lib/binstruct/binutil"
-)
-
-const (
-	sizeof8  = 1
-	sizeof16 = 2
-	sizeof32 = 4
-	sizeof64 = 8
-)
-
-// unsigned
-
-type U8 uint8
-
-func (U8) BinaryStaticSize() int            { return sizeof8 }
-func (x U8) MarshalBinary() ([]byte, error) { return []byte{byte(x)}, nil }
-func (x *U8) UnmarshalBinary(dat []byte) (int, error) {
-	if err := binutil.NeedNBytes(dat, sizeof8); err != nil {
-		return 0, err
-	}
-	*x = U8(dat[0])
-	return sizeof8, nil
-}
-
-// unsigned little endian
-
-type U16le uint16
-
-func (U16le) BinaryStaticSize() int { return sizeof16 }
-func (x U16le) MarshalBinary() ([]byte, error) {
-	var buf [sizeof16]byte
-	binary.LittleEndian.PutUint16(buf[:], uint16(x))
-	return buf[:], nil
-}
-
-func (x *U16le) UnmarshalBinary(dat []byte) (int, error) {
-	if err := binutil.NeedNBytes(dat, sizeof16); err != nil {
-		return 0, err
-	}
-	*x = U16le(binary.LittleEndian.Uint16(dat))
-	return sizeof16, nil
-}
-
-type U32le uint32
-
-func (U32le) BinaryStaticSize() int { return sizeof32 }
-func (x U32le) MarshalBinary() ([]byte, error) {
-	var buf [sizeof32]byte
-	binary.LittleEndian.PutUint32(buf[:], uint32(x))
-	return buf[:], nil
-}
-
-func (x *U32le) UnmarshalBinary(dat []byte) (int, error) {
-	if err := binutil.NeedNBytes(dat, sizeof32); err != nil {
-		return 0, err
-	}
-	*x = U32le(binary.LittleEndian.Uint32(dat))
-	return sizeof32, nil
-}
-
-type U64le uint64
-
-func (U64le) BinaryStaticSize() int { return sizeof64 }
-func (x U64le) MarshalBinary() ([]byte, error) {
-	var buf [sizeof64]byte
-	binary.LittleEndian.PutUint64(buf[:], uint64(x))
-	return buf[:], nil
-}
-
-func (x *U64le) UnmarshalBinary(dat []byte) (int, error) {
-	if err := binutil.NeedNBytes(dat, sizeof64); err != nil {
-		return 0, err
-	}
-	*x = U64le(binary.LittleEndian.Uint64(dat))
-	return sizeof64, nil
-}
-
-// unsigned big endian
-
-type U16be uint16
-
-func (U16be) BinaryStaticSize() int { return sizeof16 }
-func (x U16be) MarshalBinary() ([]byte, error) {
-	var buf [sizeof16]byte
-	binary.BigEndian.PutUint16(buf[:], uint16(x))
-	return buf[:], nil
-}
-
-func (x *U16be) UnmarshalBinary(dat []byte) (int, error) {
-	if err := binutil.NeedNBytes(dat, sizeof16); err != nil {
-		return 0, err
-	}
-	*x = U16be(binary.BigEndian.Uint16(dat))
-	return sizeof16, nil
-}
-
-type U32be uint32
-
-func (U32be) BinaryStaticSize() int { return sizeof32 }
-func (x U32be) MarshalBinary() ([]byte, error) {
-	var buf [sizeof32]byte
-	binary.BigEndian.PutUint32(buf[:], uint32(x))
-	return buf[:], nil
-}
-
-func (x *U32be) UnmarshalBinary(dat []byte) (int, error) {
-	if err := binutil.NeedNBytes(dat, sizeof32); err != nil {
-		return 0, err
-	}
-	*x = U32be(binary.BigEndian.Uint32(dat))
-	return sizeof32, nil
-}
-
-type U64be uint64
-
-func (U64be) BinaryStaticSize() int { return sizeof64 }
-func (x U64be) MarshalBinary() ([]byte, error) {
-	var buf [sizeof64]byte
-	binary.BigEndian.PutUint64(buf[:], uint64(x))
-	return buf[:], nil
-}
-
-func (x *U64be) UnmarshalBinary(dat []byte) (int, error) {
-	if err := binutil.NeedNBytes(dat, sizeof64); err != nil {
-		return 0, err
-	}
-	*x = U64be(binary.BigEndian.Uint64(dat))
-	return sizeof64, nil
-}
-
-// signed
-
-type I8 int8
-
-func (I8) BinaryStaticSize() int            { return sizeof8 }
-func (x I8) MarshalBinary() ([]byte, error) { return []byte{byte(x)}, nil }
-func (x *I8) UnmarshalBinary(dat []byte) (int, error) {
-	if err := binutil.NeedNBytes(dat, sizeof8); err != nil {
-		return 0, err
-	}
-	*x = I8(dat[0])
-	return sizeof8, nil
-}
-
-// signed little endian
-
-type I16le int16
-
-func (I16le) BinaryStaticSize() int { return sizeof16 }
-func (x I16le) MarshalBinary() ([]byte, error) {
-	var buf [sizeof16]byte
-	binary.LittleEndian.PutUint16(buf[:], uint16(x))
-	return buf[:], nil
-}
-
-func (x *I16le) UnmarshalBinary(dat []byte) (int, error) {
-	if err := binutil.NeedNBytes(dat, sizeof16); err != nil {
-		return 0, err
-	}
-	*x = I16le(binary.LittleEndian.Uint16(dat))
-	return sizeof16, nil
-}
-
-type I32le int32
-
-func (I32le) BinaryStaticSize() int { return sizeof32 }
-func (x I32le) MarshalBinary() ([]byte, error) {
-	var buf [sizeof32]byte
-	binary.LittleEndian.PutUint32(buf[:], uint32(x))
-	return buf[:], nil
-}
-
-func (x *I32le) UnmarshalBinary(dat []byte) (int, error) {
-	if err := binutil.NeedNBytes(dat, sizeof32); err != nil {
-		return 0, err
-	}
-	*x = I32le(binary.LittleEndian.Uint32(dat))
-	return sizeof32, nil
-}
-
-type I64le int64
-
-func (I64le) BinaryStaticSize() int { return sizeof64 }
-func (x I64le) MarshalBinary() ([]byte, error) {
-	var buf [sizeof64]byte
-	binary.LittleEndian.PutUint64(buf[:], uint64(x))
-	return buf[:], nil
-}
-
-func (x *I64le) UnmarshalBinary(dat []byte) (int, error) {
-	if err := binutil.NeedNBytes(dat, sizeof64); err != nil {
-		return 0, err
-	}
-	*x = I64le(binary.LittleEndian.Uint64(dat))
-	return sizeof64, nil
-}
-
-// signed big endian
-
-type I16be int16
-
-func (I16be) BinaryStaticSize() int { return sizeof16 }
-func (x I16be) MarshalBinary() ([]byte, error) {
-	var buf [sizeof16]byte
-	binary.BigEndian.PutUint16(buf[:], uint16(x))
-	return buf[:], nil
-}
-
-func (x *I16be) UnmarshalBinary(dat []byte) (int, error) {
-	if err := binutil.NeedNBytes(dat, sizeof16); err != nil {
-		return 0, err
-	}
-	*x = I16be(binary.BigEndian.Uint16(dat))
-	return sizeof16, nil
-}
-
-type I32be int32
-
-func (I32be) BinaryStaticSize() int { return sizeof32 }
-func (x I32be) MarshalBinary() ([]byte, error) {
-	var buf [sizeof32]byte
-	binary.BigEndian.PutUint32(buf[:], uint32(x))
-	return buf[:], nil
-}
-
-func (x *I32be) UnmarshalBinary(dat []byte) (int, error) {
-	if err := binutil.NeedNBytes(dat, sizeof32); err != nil {
-		return 0, err
-	}
-	*x = I32be(binary.BigEndian.Uint32(dat))
-	return sizeof32, nil
-}
-
-type I64be int64
-
-func (I64be) BinaryStaticSize() int { return sizeof64 }
-func (x I64be) MarshalBinary() ([]byte, error) {
-	var buf [sizeof64]byte
-	binary.BigEndian.PutUint64(buf[:], uint64(x))
-	return buf[:], nil
-}
-
-func (x *I64be) UnmarshalBinary(dat []byte) (int, error) {
-	if err := binutil.NeedNBytes(dat, sizeof64); err != nil {
-		return 0, err
-	}
-	*x = I64be(binary.BigEndian.Uint64(dat))
-	return sizeof64, nil
-}
diff --git a/lib/binstruct/marshal.go b/lib/binstruct/marshal.go
index 78a4bb5..48e2f1d 100644
--- a/lib/binstruct/marshal.go
+++ b/lib/binstruct/marshal.go
@@ -1,4 +1,4 @@
-// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
+// Copyright (C) 2022-2023  Luke Shumaker <lukeshu@lukeshu.com>
 //
 // SPDX-License-Identifier: GPL-2.0-or-later
 
@@ -6,6 +6,7 @@ package binstruct
 
 import (
 	"encoding"
+	"encoding/binary"
 	"fmt"
 	"reflect"
 )
@@ -30,17 +31,38 @@ func Marshal(obj any) ([]byte, error) {
 func MarshalWithoutInterface(obj any) ([]byte, error) {
 	val := reflect.ValueOf(obj)
 	switch val.Kind() {
-	case reflect.Uint8, reflect.Int8, reflect.Uint16, reflect.Int16, reflect.Uint32, reflect.Int32, reflect.Uint64, reflect.Int64:
-		typ := intKind2Type[val.Kind()]
-		dat, err := val.Convert(typ).Interface().(Marshaler).MarshalBinary()
-		if err != nil {
-			err = &UnmarshalError{
-				Type:   typ,
-				Method: "MarshalBinary",
-				Err:    err,
-			}
-		}
-		return dat, err
+	case reflect.Uint8:
+		var buf [sizeof8]byte
+		buf[0] = byte(val.Uint())
+		return buf[:], nil
+	case reflect.Int8:
+		var buf [sizeof8]byte
+		buf[0] = byte(val.Int())
+		return buf[:], nil
+	case reflect.Uint16:
+		var buf [sizeof16]byte
+		binary.LittleEndian.PutUint16(buf[:], uint16(val.Uint()))
+		return buf[:], nil
+	case reflect.Int16:
+		var buf [sizeof16]byte
+		binary.LittleEndian.PutUint16(buf[:], uint16(val.Int()))
+		return buf[:], nil
+	case reflect.Uint32:
+		var buf [sizeof32]byte
+		binary.LittleEndian.PutUint32(buf[:], uint32(val.Uint()))
+		return buf[:], nil
+	case reflect.Int32:
+		var buf [sizeof32]byte
+		binary.LittleEndian.PutUint32(buf[:], uint32(val.Int()))
+		return buf[:], nil
+	case reflect.Uint64:
+		var buf [sizeof64]byte
+		binary.LittleEndian.PutUint64(buf[:], val.Uint())
+		return buf[:], nil
+	case reflect.Int64:
+		var buf [sizeof64]byte
+		binary.LittleEndian.PutUint64(buf[:], uint64(val.Int()))
+		return buf[:], nil
 	case reflect.Ptr:
 		return Marshal(val.Elem().Interface())
 	case reflect.Array:
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
 }
diff --git a/lib/binstruct/unmarshal.go b/lib/binstruct/unmarshal.go
index 61c2a4a..eae4b84 100644
--- a/lib/binstruct/unmarshal.go
+++ b/lib/binstruct/unmarshal.go
@@ -1,13 +1,16 @@
-// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
+// Copyright (C) 2022-2023  Luke Shumaker <lukeshu@lukeshu.com>
 //
 // SPDX-License-Identifier: GPL-2.0-or-later
 
 package binstruct
 
 import (
+	"encoding/binary"
 	"errors"
 	"fmt"
 	"reflect"
+
+	"git.lukeshu.com/btrfs-progs-ng/lib/binstruct/binutil"
 )
 
 type Unmarshaler interface {
@@ -29,6 +32,23 @@ func Unmarshal(dat []byte, dstPtr any) (int, error) {
 	return UnmarshalWithoutInterface(dat, dstPtr)
 }
 
+// unmarshal is like Unmarshal, but for internal use to avoid some
+// slow round-tripping between `any` and `reflect.Value`.
+func unmarshal(dat []byte, dst reflect.Value, isUnmarshaler bool) (int, error) {
+	if isUnmarshaler {
+		n, err := dst.Addr().Interface().(Unmarshaler).UnmarshalBinary(dat)
+		if err != nil {
+			err = &UnmarshalError{
+				Type:   reflect.PtrTo(dst.Type()),
+				Method: "UnmarshalBinary",
+				Err:    err,
+			}
+		}
+		return n, err
+	}
+	return unmarshalWithoutInterface(dat, dst)
+}
+
 func UnmarshalWithoutInterface(dat []byte, dstPtr any) (int, error) {
 	_dstPtr := reflect.ValueOf(dstPtr)
 	if _dstPtr.Kind() != reflect.Ptr {
@@ -37,24 +57,70 @@ func UnmarshalWithoutInterface(dat []byte, dstPtr any) (int, error) {
 			Err:  errors.New("not a pointer"),
 		})
 	}
-	dst := _dstPtr.Elem()
+	return unmarshalWithoutInterface(dat, _dstPtr.Elem())
+}
 
+func unmarshalWithoutInterface(dat []byte, dst reflect.Value) (int, error) {
 	switch dst.Kind() {
-	case reflect.Uint8, reflect.Int8, reflect.Uint16, reflect.Int16, reflect.Uint32, reflect.Int32, reflect.Uint64, reflect.Int64:
-		typ := intKind2Type[dst.Kind()]
-		newDstPtr := reflect.New(typ)
-		n, err := Unmarshal(dat, newDstPtr.Interface())
-		dst.Set(newDstPtr.Elem().Convert(dst.Type()))
-		return n, err
+	case reflect.Uint8:
+		if err := binutil.NeedNBytes(dat, sizeof8); err != nil {
+			return 0, err
+		}
+		dst.SetUint(uint64(dat[0]))
+		return sizeof8, nil
+	case reflect.Int8:
+		if err := binutil.NeedNBytes(dat, sizeof8); err != nil {
+			return 0, err
+		}
+		dst.SetInt(int64(dat[0]))
+		return sizeof8, nil
+	case reflect.Uint16:
+		if err := binutil.NeedNBytes(dat, sizeof16); err != nil {
+			return 0, err
+		}
+		dst.SetUint(uint64(binary.LittleEndian.Uint16(dat[:sizeof16])))
+		return sizeof16, nil
+	case reflect.Int16:
+		if err := binutil.NeedNBytes(dat, sizeof16); err != nil {
+			return 0, err
+		}
+		dst.SetInt(int64(binary.LittleEndian.Uint16(dat[:sizeof16])))
+		return sizeof16, nil
+	case reflect.Uint32:
+		if err := binutil.NeedNBytes(dat, sizeof32); err != nil {
+			return 0, err
+		}
+		dst.SetUint(uint64(binary.LittleEndian.Uint32(dat[:sizeof32])))
+		return sizeof32, nil
+	case reflect.Int32:
+		if err := binutil.NeedNBytes(dat, sizeof32); err != nil {
+			return 0, err
+		}
+		dst.SetInt(int64(binary.LittleEndian.Uint32(dat[:sizeof32])))
+		return sizeof32, nil
+	case reflect.Uint64:
+		if err := binutil.NeedNBytes(dat, sizeof64); err != nil {
+			return 0, err
+		}
+		dst.SetUint(binary.LittleEndian.Uint64(dat[:sizeof64]))
+		return sizeof64, nil
+	case reflect.Int64:
+		if err := binutil.NeedNBytes(dat, sizeof64); err != nil {
+			return 0, err
+		}
+		dst.SetInt(int64(binary.LittleEndian.Uint64(dat[:sizeof64])))
+		return sizeof64, nil
 	case reflect.Ptr:
-		elemPtr := reflect.New(dst.Type().Elem())
-		n, err := Unmarshal(dat, elemPtr.Interface())
-		dst.Set(elemPtr.Convert(dst.Type()))
+		typ := dst.Type()
+		elemPtr := reflect.New(typ.Elem())
+		n, err := unmarshal(dat, elemPtr.Elem(), typ.Implements(unmarshalerType))
+		dst.SetPointer(elemPtr.UnsafePointer())
 		return n, err
 	case reflect.Array:
+		isUnmarshaler := dst.Type().Elem().Implements(unmarshalerType)
 		var n int
 		for i := 0; i < dst.Len(); i++ {
-			_n, err := Unmarshal(dat[n:], dst.Index(i).Addr().Interface())
+			_n, err := unmarshal(dat[n:], dst.Index(i), isUnmarshaler)
 			n += _n
 			if err != nil {
 				return n, err
@@ -65,7 +131,7 @@ func UnmarshalWithoutInterface(dat []byte, dstPtr any) (int, error) {
 		return getStructHandler(dst.Type()).Unmarshal(dat, dst)
 	default:
 		panic(&InvalidTypeError{
-			Type: _dstPtr.Type(),
+			Type: reflect.PtrTo(dst.Type()),
 			Err: fmt.Errorf("does not implement binfmt.Unmarshaler and kind=%v is not a supported statically-sized kind",
 				dst.Kind()),
 		})