2 files changed, 34 insertions, 43 deletions

diff --git a/tvix/store/protos/export.go b/tvix/store/protos/export.go
index 8c4a27522977..889311efe70b 100644
--- a/tvix/store/protos/export.go
+++ b/tvix/store/protos/export.go
@@ -12,12 +12,12 @@ import (
 type DirectoryLookupFn func([]byte) (*castorev1pb.Directory, error)
 type BlobLookupFn func([]byte) (io.ReadCloser, error)
 
-// Export will traverse a given PathInfo structure, and write the contents
-// in NAR format to the passed Writer.
+// Export will traverse a given root node, and write the contents in NAR format
+// to the passed Writer.
 // It uses directoryLookupFn and blobLookupFn to resolve references.
 func Export(
 	w io.Writer,
-	pathInfo *PathInfo,
+	rootNode *castorev1pb.Node,
 	directoryLookupFn DirectoryLookupFn,
 	blobLookupFn BlobLookupFn,
 ) error {
@@ -43,18 +43,17 @@ func Export(
 	// peek at the pathInfo root and assemble the root node and write to writer
 	// in the case of a regular file, we retrieve and write the contents, close and exit
 	// in the case of a symlink, we write the symlink, close and exit
-	switch v := (pathInfo.GetNode().GetNode()).(type) {
-	case *castorev1pb.Node_File:
+	if fileNode := rootNode.GetFile(); fileNode != nil {
 		rootHeader.Type = nar.TypeRegular
-		rootHeader.Size = int64(v.File.GetSize())
-		rootHeader.Executable = v.File.GetExecutable()
+		rootHeader.Size = int64(fileNode.GetSize())
+		rootHeader.Executable = fileNode.GetExecutable()
 		err := narWriter.WriteHeader(rootHeader)
 		if err != nil {
 			return fmt.Errorf("unable to write root header: %w", err)
 		}
 
 		// if it's a regular file, retrieve and write the contents
-		blobReader, err := blobLookupFn(v.File.GetDigest())
+		blobReader, err := blobLookupFn(fileNode.GetDigest())
 		if err != nil {
 			return fmt.Errorf("unable to lookup blob: %w", err)
 		}
@@ -76,10 +75,9 @@ func Export(
 		}
 
 		return nil
-
-	case *castorev1pb.Node_Symlink:
+	} else if symlinkNode := rootNode.GetSymlink(); symlinkNode != nil {
 		rootHeader.Type = nar.TypeSymlink
-		rootHeader.LinkTarget = string(v.Symlink.GetTarget())
+		rootHeader.LinkTarget = string(symlinkNode.GetTarget())
 		err := narWriter.WriteHeader(rootHeader)
 		if err != nil {
 			return fmt.Errorf("unable to write root header: %w", err)
@@ -89,11 +87,9 @@ func Export(
 		if err != nil {
 			return fmt.Errorf("unable to close nar reader: %w", err)
 		}
-
-		return nil
-	case *castorev1pb.Node_Directory:
+	} else if directoryNode := rootNode.GetDirectory(); directoryNode != nil {
 		// We have a directory at the root, look it up and put in on the stack.
-		directory, err := directoryLookupFn(v.Directory.Digest)
+		directory, err := directoryLookupFn(directoryNode.GetDigest())
 		if err != nil {
 			return fmt.Errorf("unable to lookup directory: %w", err)
 		}
@@ -108,6 +104,8 @@ func Export(
 		if err != nil {
 			return fmt.Errorf("error writing header: %w", err)
 		}
+	} else {
+		panic("invalid type") // unreachable
 	}
 
 	// as long as the stack is not empty, we keep running.
diff --git a/tvix/store/protos/export_test.go b/tvix/store/protos/export_test.go
index d45dd1ee7cf3..9b1901758f83 100644
--- a/tvix/store/protos/export_test.go
+++ b/tvix/store/protos/export_test.go
@@ -30,21 +30,18 @@ func mustBlobDigest(r io.Reader) []byte {
 }
 
 func TestSymlink(t *testing.T) {
-	pathInfo := &storev1pb.PathInfo{
-
-		Node: &castorev1pb.Node{
-			Node: &castorev1pb.Node_Symlink{
-				Symlink: &castorev1pb.SymlinkNode{
-					Name:   []byte("doesntmatter"),
-					Target: []byte("/nix/store/somewhereelse"),
-				},
+	node := &castorev1pb.Node{
+		Node: &castorev1pb.Node_Symlink{
+			Symlink: &castorev1pb.SymlinkNode{
+				Name:   []byte("doesntmatter"),
+				Target: []byte("/nix/store/somewhereelse"),
 			},
 		},
 	}
 
 	var buf bytes.Buffer
 
-	err := storev1pb.Export(&buf, pathInfo, func([]byte) (*castorev1pb.Directory, error) {
+	err := storev1pb.Export(&buf, node, func([]byte) (*castorev1pb.Directory, error) {
 		panic("no directories expected")
 	}, func([]byte) (io.ReadCloser, error) {
 		panic("no files expected")
@@ -70,22 +67,20 @@ func TestRegular(t *testing.T) {
 		0x65, 0x2b,
 	}
 
-	pathInfo := &storev1pb.PathInfo{
-		Node: &castorev1pb.Node{
-			Node: &castorev1pb.Node_File{
-				File: &castorev1pb.FileNode{
-					Name:       []byte("doesntmatter"),
-					Digest:     BLAKE3_DIGEST_0X01,
-					Size:       1,
-					Executable: false,
-				},
+	node := &castorev1pb.Node{
+		Node: &castorev1pb.Node_File{
+			File: &castorev1pb.FileNode{
+				Name:       []byte("doesntmatter"),
+				Digest:     BLAKE3_DIGEST_0X01,
+				Size:       1,
+				Executable: false,
 			},
 		},
 	}
 
 	var buf bytes.Buffer
 
-	err := storev1pb.Export(&buf, pathInfo, func([]byte) (*castorev1pb.Directory, error) {
+	err := storev1pb.Export(&buf, node, func([]byte) (*castorev1pb.Directory, error) {
 		panic("no directories expected")
 	}, func(blobRef []byte) (io.ReadCloser, error) {
 		if !bytes.Equal(blobRef, BLAKE3_DIGEST_0X01) {
@@ -115,21 +110,19 @@ func TestEmptyDirectory(t *testing.T) {
 	}
 	emptyDirectoryDigest := mustDirectoryDigest(emptyDirectory)
 
-	pathInfo := &storev1pb.PathInfo{
-		Node: &castorev1pb.Node{
-			Node: &castorev1pb.Node_Directory{
-				Directory: &castorev1pb.DirectoryNode{
-					Name:   []byte("doesntmatter"),
-					Digest: emptyDirectoryDigest,
-					Size:   0,
-				},
+	node := &castorev1pb.Node{
+		Node: &castorev1pb.Node_Directory{
+			Directory: &castorev1pb.DirectoryNode{
+				Name:   []byte("doesntmatter"),
+				Digest: emptyDirectoryDigest,
+				Size:   0,
 			},
 		},
 	}
 
 	var buf bytes.Buffer
 
-	err := storev1pb.Export(&buf, pathInfo, func(directoryRef []byte) (*castorev1pb.Directory, error) {
+	err := storev1pb.Export(&buf, node, func(directoryRef []byte) (*castorev1pb.Directory, error) {
 		if !bytes.Equal(directoryRef, emptyDirectoryDigest) {
 			panic("unexpected directoryRef")
 		}