1 files changed, 79 insertions, 13 deletions

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()),
 		})