lib: eagerly unmarshal items

5 files changed, 62 insertions, 14 deletions

diff --git a/pkg/btrfs/Makefile b/pkg/btrfs/Makefile
index fc6e9fd..b29e033 100644
--- a/pkg/btrfs/Makefile
+++ b/pkg/btrfs/Makefile
@@ -9,10 +9,17 @@ files += btrfsitem/items.txt
 btrfsitem/items_gen.go: btrfsitem/items.txt $(MAKEFILE_LIST)
 	{ \
 	  echo 'package $(@D)'; \
-	  echo 'import "lukeshu.com/btrfs-tools/pkg/btrfs/internal"'; \
+	  echo 'import ('; \
+	  echo   '"reflect"'; \
+	  echo; \
+	  echo   '"lukeshu.com/btrfs-tools/pkg/btrfs/internal"'; \
+	  echo ')'; \
 	  echo 'const ('; \
 	  sed -E 's,(.*)=(.*) (.*),\1_KEY=internal.\1_KEY,' $<; \
 	  echo ')'; \
+	  echo 'var keytype2gotype = map[Type]reflect.Type{'; \
+	  sed -E 's|(.*)=(.*) (.*)|\1_KEY: reflect.TypeOf(\3{}),|' $<; \
+	  echo '}'; \
 	  sed -E 's,(.*)=(.*) (.*),\3,' $< | LC_COLLATE=C sort -u | sed 's,.*,func (&) isItem() {},'; \
 	} | gofmt >$@
 files += btrfsitem/items_gen.go
diff --git a/pkg/btrfs/btrfsitem/items.go b/pkg/btrfs/btrfsitem/items.go
index 51e6344..91b7029 100644
--- a/pkg/btrfs/btrfsitem/items.go
+++ b/pkg/btrfs/btrfsitem/items.go
@@ -1,6 +1,10 @@
 package btrfsitem
 
 import (
+	"fmt"
+	"reflect"
+
+	"lukeshu.com/btrfs-tools/pkg/binstruct"
 	"lukeshu.com/btrfs-tools/pkg/btrfs/internal"
 )
 
@@ -9,3 +13,20 @@ type Type = internal.ItemType
 type Item interface {
 	isItem()
 }
+
+func UnmarshalItem(keytyp Type, dat []byte) (Item, error) {
+	gotyp, ok := keytype2gotype[keytyp]
+	if !ok {
+		return nil, fmt.Errorf("btrfsitem.UnmarshalItem: unknown item type: %v", keytyp)
+	}
+	retPtr := reflect.New(gotyp)
+	n, err := binstruct.Unmarshal(dat, retPtr.Interface())
+	if err != nil {
+		return nil, fmt.Errorf("btrfsitem.UnmarshalItem: %w", err)
+	}
+	if n < len(dat) {
+		return nil, fmt.Errorf("btrfsitem.UnmarshalItem: left over data: got %d bytes but only consumed %d",
+			len(dat), n)
+	}
+	return retPtr.Elem().Interface().(Item), nil
+}
diff --git a/pkg/btrfs/btrfsitem/items_gen.go b/pkg/btrfs/btrfsitem/items_gen.go
index 5234c96..40a0b53 100644
--- a/pkg/btrfs/btrfsitem/items_gen.go
+++ b/pkg/btrfs/btrfsitem/items_gen.go
@@ -1,6 +1,10 @@
 package btrfsitem
 
-import "lukeshu.com/btrfs-tools/pkg/btrfs/internal"
+import (
+	"reflect"
+
+	"lukeshu.com/btrfs-tools/pkg/btrfs/internal"
+)
 
 const (
 	CHUNK_ITEM_KEY           = internal.CHUNK_ITEM_KEY
@@ -17,6 +21,21 @@ const (
 	UUID_RECEIVED_SUBVOL_KEY = internal.UUID_RECEIVED_SUBVOL_KEY
 )
 
+var keytype2gotype = map[Type]reflect.Type{
+	CHUNK_ITEM_KEY:           reflect.TypeOf(Chunk{}),
+	DEV_ITEM_KEY:             reflect.TypeOf(Dev{}),
+	DEV_EXTENT_KEY:           reflect.TypeOf(DevExtent{}),
+	UNTYPED_KEY:              reflect.TypeOf(Empty{}),
+	QGROUP_RELATION_KEY:      reflect.TypeOf(Empty{}),
+	INODE_ITEM_KEY:           reflect.TypeOf(Inode{}),
+	INODE_REF_KEY:            reflect.TypeOf(InodeRef{}),
+	ORPHAN_ITEM_KEY:          reflect.TypeOf(Orphan{}),
+	PERSISTENT_ITEM_KEY:      reflect.TypeOf(DevStats{}),
+	ROOT_ITEM_KEY:            reflect.TypeOf(Root{}),
+	UUID_SUBVOL_KEY:          reflect.TypeOf(UUIDMap{}),
+	UUID_RECEIVED_SUBVOL_KEY: reflect.TypeOf(UUIDMap{}),
+}
+
 func (Chunk) isItem()     {}
 func (Dev) isItem()       {}
 func (DevExtent) isItem() {}
diff --git a/pkg/btrfs/io2_fs.go b/pkg/btrfs/io2_fs.go
index bcf9777..b7752c7 100644
--- a/pkg/btrfs/io2_fs.go
+++ b/pkg/btrfs/io2_fs.go
@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"reflect"
 
-	"lukeshu.com/btrfs-tools/pkg/binstruct"
 	"lukeshu.com/btrfs-tools/pkg/btrfs/btrfsitem"
 	"lukeshu.com/btrfs-tools/pkg/util"
 )
@@ -123,17 +122,14 @@ func (fs *FS) Init() error {
 		for _, chunk := range syschunks {
 			fs.chunks = append(fs.chunks, chunk)
 		}
-		if err := fs.WalkTree(sb.Data.ChunkTree, func(key Key, dat []byte) error {
+		if err := fs.WalkTree(sb.Data.ChunkTree, func(key Key, body btrfsitem.Item) error {
 			if key.ItemType != btrfsitem.CHUNK_ITEM_KEY {
 				return nil
 			}
-			pair := SysChunk{
-				Key: key,
-			}
-			if _, err := binstruct.Unmarshal(dat, &pair.Chunk); err != nil {
-				return err
-			}
-			fs.chunks = append(fs.chunks, pair)
+			fs.chunks = append(fs.chunks, SysChunk{
+				Key:   key,
+				Chunk: body.(btrfsitem.Chunk),
+			})
 			return nil
 		}); err != nil {
 			fs.initErr = err
diff --git a/pkg/btrfs/types_btree.go b/pkg/btrfs/types_btree.go
index 9b06e5e..cac5f9a 100644
--- a/pkg/btrfs/types_btree.go
+++ b/pkg/btrfs/types_btree.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 
 	"lukeshu.com/btrfs-tools/pkg/binstruct"
+	"lukeshu.com/btrfs-tools/pkg/btrfs/btrfsitem"
 	"lukeshu.com/btrfs-tools/pkg/util"
 )
 
@@ -87,7 +88,7 @@ type ItemHeader struct {
 
 type Item struct {
 	Head ItemHeader
-	Body []byte
+	Body btrfsitem.Item
 }
 
 // MaxItems returns the maximum possible valid value of
@@ -135,8 +136,12 @@ func (node *Node) UnmarshalBinary(nodeBuf []byte) (int, error) {
 				return max(n, lastRead), fmt.Errorf("(leaf): item references byte %d, but node only has %d bytes",
 					dataOff+dataSize, len(nodeBuf))
 			}
+			dataBuf := nodeBuf[dataOff : dataOff+dataSize]
 			lastRead = max(lastRead, dataOff+dataSize)
-			item.Body = nodeBuf[dataOff : dataOff+dataSize]
+			item.Body, err = btrfsitem.UnmarshalItem(item.Head.Key.ItemType, dataBuf)
+			if err != nil {
+				return max(n, lastRead), fmt.Errorf("(leaf): item %d: %w", i, err)
+			}
 
 			node.BodyLeaf = append(node.BodyLeaf, item)
 		}
@@ -210,7 +215,7 @@ func (fs *FS) ReadNode(addr LogicalAddr) (util.Ref[LogicalAddr, Node], error) {
 	}, nil
 }
 
-func (fs *FS) WalkTree(nodeAddr LogicalAddr, fn func(Key, []byte) error) error {
+func (fs *FS) WalkTree(nodeAddr LogicalAddr, fn func(Key, btrfsitem.Item) error) error {
 	if nodeAddr == 0 {
 		return nil
 	}