chore(tvix/nar-bridge): move to Export from storev1pb r/6773

This removes the Export method in nar-bridge, and updates all users to
the version now in storev1pb.

It moves the roundtrip test to the importer crate, and some of the
utility functions into a separate util_test.go file.

Change-Id: I81d9e0b35dfd78ef1042bed307281eecd2aaa2a8
Reviewed-on: https://cl.tvl.fyi/c/depot/+/9603
Reviewed-by: Brian McGee <brian@bmcgee.ie>
Tested-by: BuildkiteCI

8 files changed, 49 insertions, 335 deletions

diff --git a/tvix/nar-bridge/default.nix b/tvix/nar-bridge/default.nix
index 95d7cfd6be27..d2984c5bad53 100644
--- a/tvix/nar-bridge/default.nix
+++ b/tvix/nar-bridge/default.nix
@@ -6,5 +6,5 @@ pkgs.buildGoModule {
   name = "nar-bridge";
   src = depot.third_party.gitignoreSource ./.;
 
-  vendorHash = "sha256-UxKTfy1NNtZhor8Hj9LZja72vqW7OYDRn8/cUETPzoU=";
+  vendorHash = "sha256-Lb0MOziF86JrnrA9SibHYQAaS7H054Nuf3l8tm/9Sf8=";
 }
diff --git a/tvix/nar-bridge/go.mod b/tvix/nar-bridge/go.mod
index 4758b0f1793a..6e02277a4dbb 100644
--- a/tvix/nar-bridge/go.mod
+++ b/tvix/nar-bridge/go.mod
@@ -2,7 +2,7 @@ module code.tvl.fyi/tvix/nar-bridge
 
 require (
 	code.tvl.fyi/tvix/castore/protos v0.0.0-20231009220507-d6e0c5ab9bb7
-	code.tvl.fyi/tvix/store/protos v0.0.0-20231009220507-46652989e097
+	code.tvl.fyi/tvix/store/protos v0.0.0-20231010185549-e7ea67342035
 	github.com/alecthomas/kong v0.7.1
 	github.com/go-chi/chi v1.5.4
 	github.com/go-chi/chi/v5 v5.0.7
@@ -27,7 +27,7 @@ require (
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/spaolacci/murmur3 v1.1.0 // indirect
 	golang.org/x/crypto v0.14.0 // indirect
-	golang.org/x/net v0.16.0 // indirect
+	golang.org/x/net v0.17.0 // indirect
 	golang.org/x/sys v0.13.0 // indirect
 	golang.org/x/text v0.13.0 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20231009173412-8bfb1ae86b6c // indirect
diff --git a/tvix/nar-bridge/go.sum b/tvix/nar-bridge/go.sum
index 8e0218b473b1..469a6970bfdc 100644
--- a/tvix/nar-bridge/go.sum
+++ b/tvix/nar-bridge/go.sum
@@ -1,7 +1,7 @@
 code.tvl.fyi/tvix/castore/protos v0.0.0-20231009220507-d6e0c5ab9bb7 h1:gX2LWo/QHwGZK2QsDap9Lx1GrKLPX6mfgeNbGK3mwrU=
 code.tvl.fyi/tvix/castore/protos v0.0.0-20231009220507-d6e0c5ab9bb7/go.mod h1:hj0y8RPthqn1QPj8u2jFe2vzH7NouUoclrwo1/CSbuc=
-code.tvl.fyi/tvix/store/protos v0.0.0-20231009220507-46652989e097 h1:8BhyvAmI+yiFwpsttcCvAKDjzh+o2JL9wjbccurHfMc=
-code.tvl.fyi/tvix/store/protos v0.0.0-20231009220507-46652989e097/go.mod h1:WTugC5W8TV5mE0h38yBxCjQZsuenjmHtdvx2IvbYRqY=
+code.tvl.fyi/tvix/store/protos v0.0.0-20231010185549-e7ea67342035 h1:/CjyjG/4PiByWnO6q26bLGjFPP96oZCMMX63UA9wkdc=
+code.tvl.fyi/tvix/store/protos v0.0.0-20231010185549-e7ea67342035/go.mod h1:RmijF3bfElwtZpNkBtW66QEj/jldGNu+W2HlgZro7lw=
 github.com/alecthomas/assert/v2 v2.1.0 h1:tbredtNcQnoSd3QBhQWI7QZ3XHOVkw1Moklp2ojoH/0=
 github.com/alecthomas/kong v0.7.1 h1:azoTh0IOfwlAX3qN9sHWTxACE2oV8Bg2gAwBsMwDQY4=
 github.com/alecthomas/kong v0.7.1/go.mod h1:n1iCIO2xS46oE8ZfYCNDqdR0b0wZNrXAIAqro/2132U=
@@ -51,8 +51,8 @@ github.com/ulikunitz/xz v0.5.11 h1:kpFauv27b6ynzBNT/Xy+1k+fK4WswhN/6PN5WhFAGw8=
 github.com/ulikunitz/xz v0.5.11/go.mod h1:nbz6k7qbPmH4IRqmfOplQw/tblSgqTqBwxkY0oWt/14=
 golang.org/x/crypto v0.14.0 h1:wBqGXzWJW6m1XrIKlAH0Hs1JJ7+9KBwnIO8v66Q9cHc=
 golang.org/x/crypto v0.14.0/go.mod h1:MVFd36DqK4CsrnJYDkBA3VC4m2GkXAM0PvzMCn4JQf4=
-golang.org/x/net v0.16.0 h1:7eBu7KsSvFDtSXUIDbh3aqlK4DPsZ1rByC8PFfBThos=
-golang.org/x/net v0.16.0/go.mod h1:NxSsAGuq816PNPmqtQdLE42eU2Fs7NoRIZrHJAlaCOE=
+golang.org/x/net v0.17.0 h1:pVaXccu2ozPjCXewfr1S7xza/zcXTity9cCdXQYSjIM=
+golang.org/x/net v0.17.0/go.mod h1:NxSsAGuq816PNPmqtQdLE42eU2Fs7NoRIZrHJAlaCOE=
 golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.13.0 h1:Af8nKPmuFypiUBjVoU9V20FiaFXOcuZI21p0ycVYYGE=
diff --git a/tvix/nar-bridge/pkg/exporter/export.go b/tvix/nar-bridge/pkg/exporter/export.go
deleted file mode 100644
index 1f898ccad6dd..000000000000
--- a/tvix/nar-bridge/pkg/exporter/export.go
+++ /dev/null
@@ -1,276 +0,0 @@
-package exporter
-
-import (
-	"fmt"
-	"io"
-	"path"
-
-	castorev1pb "code.tvl.fyi/tvix/castore/protos"
-	storev1pb "code.tvl.fyi/tvix/store/protos"
-	"github.com/nix-community/go-nix/pkg/nar"
-)
-
-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.
-// It uses directoryLookupFn and blobLookupFn to resolve references.
-func Export(
-	w io.Writer,
-	pathInfo *storev1pb.PathInfo,
-	directoryLookupFn DirectoryLookupFn,
-	blobLookupFn BlobLookupFn,
-) error {
-	// initialize a NAR writer
-	narWriter, err := nar.NewWriter(w)
-	if err != nil {
-		return fmt.Errorf("unable to initialize nar writer: %w", err)
-	}
-	defer narWriter.Close()
-
-	// populate rootHeader
-	rootHeader := &nar.Header{
-		Path: "/",
-	}
-
-	// populate a stack
-	// we will push paths and directories to it when entering a directory,
-	// and emit individual elements to the NAR writer, draining the Directory object.
-	// once it's empty, we can pop it off the stack.
-	var stackPaths = []string{}
-	var stackDirectories = []*castorev1pb.Directory{}
-
-	// 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:
-		rootHeader.Type = nar.TypeRegular
-		rootHeader.Size = int64(v.File.GetSize())
-		rootHeader.Executable = v.File.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())
-		if err != nil {
-			return fmt.Errorf("unable to lookup blob: %w", err)
-		}
-		defer blobReader.Close()
-
-		_, err = io.Copy(narWriter, blobReader)
-		if err != nil {
-			return fmt.Errorf("unable to read from blobReader: %w", err)
-		}
-
-		err = blobReader.Close()
-		if err != nil {
-			return fmt.Errorf("unable to close content reader: %w", err)
-		}
-
-		err = narWriter.Close()
-		if err != nil {
-			return fmt.Errorf("unable to close nar reader: %w", err)
-		}
-
-		return nil
-
-	case *castorev1pb.Node_Symlink:
-		rootHeader.Type = nar.TypeSymlink
-		rootHeader.LinkTarget = string(v.Symlink.GetTarget())
-		err := narWriter.WriteHeader(rootHeader)
-		if err != nil {
-			return fmt.Errorf("unable to write root header: %w", err)
-		}
-
-		err = narWriter.Close()
-		if err != nil {
-			return fmt.Errorf("unable to close nar reader: %w", err)
-		}
-
-		return nil
-	case *castorev1pb.Node_Directory:
-		// We have a directory at the root, look it up and put in on the stack.
-		directory, err := directoryLookupFn(v.Directory.Digest)
-		if err != nil {
-			return fmt.Errorf("unable to lookup directory: %w", err)
-		}
-		stackDirectories = append(stackDirectories, directory)
-		stackPaths = append(stackPaths, "/")
-
-		err = narWriter.WriteHeader(&nar.Header{
-			Path: "/",
-			Type: nar.TypeDirectory,
-		})
-
-		if err != nil {
-			return fmt.Errorf("error writing header: %w", err)
-		}
-	}
-
-	// as long as the stack is not empty, we keep running.
-	for {
-		if len(stackDirectories) == 0 {
-			return nil
-		}
-
-		// Peek at the current top of the stack.
-		topOfStack := stackDirectories[len(stackDirectories)-1]
-		topOfStackPath := stackPaths[len(stackPaths)-1]
-
-		// get the next element that's lexicographically smallest, and drain it from
-		// the current directory on top of the stack.
-		nextNode := drainNextNode(topOfStack)
-
-		// If nextNode returns nil, there's nothing left in the directory node, so we
-		// can emit it from the stack.
-		// Contrary to the import case, we don't emit the node popping from the stack, but when pushing.
-		if nextNode == nil {
-			// pop off stack
-			stackDirectories = stackDirectories[:len(stackDirectories)-1]
-			stackPaths = stackPaths[:len(stackPaths)-1]
-
-			continue
-		}
-
-		switch n := (nextNode).(type) {
-		case *castorev1pb.DirectoryNode:
-			err := narWriter.WriteHeader(&nar.Header{
-				Path: path.Join(topOfStackPath, string(n.GetName())),
-				Type: nar.TypeDirectory,
-			})
-			if err != nil {
-				return fmt.Errorf("unable to write nar header: %w", err)
-			}
-
-			d, err := directoryLookupFn(n.GetDigest())
-			if err != nil {
-				return fmt.Errorf("unable to lookup directory: %w", err)
-			}
-
-			// add to stack
-			stackDirectories = append(stackDirectories, d)
-			stackPaths = append(stackPaths, path.Join(topOfStackPath, string(n.GetName())))
-		case *castorev1pb.FileNode:
-			err := narWriter.WriteHeader(&nar.Header{
-				Path:       path.Join(topOfStackPath, string(n.GetName())),
-				Type:       nar.TypeRegular,
-				Size:       int64(n.GetSize()),
-				Executable: n.GetExecutable(),
-			})
-			if err != nil {
-				return fmt.Errorf("unable to write nar header: %w", err)
-			}
-
-			// copy file contents
-			contentReader, err := blobLookupFn(n.GetDigest())
-			if err != nil {
-				return fmt.Errorf("unable to get blob: %w", err)
-			}
-			defer contentReader.Close()
-
-			_, err = io.Copy(narWriter, contentReader)
-			if err != nil {
-				return fmt.Errorf("unable to copy contents from contentReader: %w", err)
-			}
-
-			err = contentReader.Close()
-			if err != nil {
-				return fmt.Errorf("unable to close content reader: %w", err)
-			}
-		case *castorev1pb.SymlinkNode:
-			err := narWriter.WriteHeader(&nar.Header{
-				Path:       path.Join(topOfStackPath, string(n.GetName())),
-				Type:       nar.TypeSymlink,
-				LinkTarget: string(n.GetTarget()),
-			})
-			if err != nil {
-				return fmt.Errorf("unable to write nar header: %w", err)
-			}
-		}
-	}
-}
-
-// drainNextNode will drain a directory message with one of its child nodes,
-// whichever comes first alphabetically.
-func drainNextNode(d *castorev1pb.Directory) interface{} {
-	switch v := (smallestNode(d)).(type) {
-	case *castorev1pb.DirectoryNode:
-		d.Directories = d.Directories[1:]
-		return v
-	case *castorev1pb.FileNode:
-		d.Files = d.Files[1:]
-		return v
-	case *castorev1pb.SymlinkNode:
-		d.Symlinks = d.Symlinks[1:]
-		return v
-	case nil:
-		return nil
-	default:
-		panic("invalid type encountered")
-	}
-}
-
-// smallestNode will return the node from a directory message,
-// whichever comes first alphabetically.
-func smallestNode(d *castorev1pb.Directory) interface{} {
-	childDirectories := d.GetDirectories()
-	childFiles := d.GetFiles()
-	childSymlinks := d.GetSymlinks()
-
-	if len(childDirectories) > 0 {
-		if len(childFiles) > 0 {
-			if len(childSymlinks) > 0 {
-				// directories,files,symlinks
-				return smallerNode(smallerNode(childDirectories[0], childFiles[0]), childSymlinks[0])
-			} else {
-				// directories,files,!symlinks
-				return smallerNode(childDirectories[0], childFiles[0])
-			}
-		} else {
-			// directories,!files
-			if len(childSymlinks) > 0 {
-				// directories,!files,symlinks
-				return smallerNode(childDirectories[0], childSymlinks[0])
-			} else {
-				// directories,!files,!symlinks
-				return childDirectories[0]
-			}
-		}
-	} else {
-		// !directories
-		if len(childFiles) > 0 {
-			// !directories,files
-			if len(childSymlinks) > 0 {
-				// !directories,files,symlinks
-				return smallerNode(childFiles[0], childSymlinks[0])
-			} else {
-				// !directories,files,!symlinks
-				return childFiles[0]
-			}
-		} else {
-			//!directories,!files
-			if len(childSymlinks) > 0 {
-				//!directories,!files,symlinks
-				return childSymlinks[0]
-			} else {
-				//!directories,!files,!symlinks
-				return nil
-			}
-		}
-	}
-}
-
-// smallerNode compares two nodes by their name,
-// and returns the one with the smaller name.
-// both nodes may not be nil, we do check for these cases in smallestNode.
-func smallerNode(a interface{ GetName() []byte }, b interface{ GetName() []byte }) interface{ GetName() []byte } {
-	if string(a.GetName()) < string(b.GetName()) {
-		return a
-	} else {
-		return b
-	}
-}
diff --git a/tvix/nar-bridge/pkg/http/nar_get.go b/tvix/nar-bridge/pkg/http/nar_get.go
index 0c2b299e785b..85405d81f463 100644
--- a/tvix/nar-bridge/pkg/http/nar_get.go
+++ b/tvix/nar-bridge/pkg/http/nar_get.go
@@ -13,7 +13,6 @@ import (
 	"sync"
 
 	castorev1pb "code.tvl.fyi/tvix/castore/protos"
-	"code.tvl.fyi/tvix/nar-bridge/pkg/exporter"
 	storev1pb "code.tvl.fyi/tvix/store/protos"
 	"github.com/go-chi/chi/v5"
 	nixhash "github.com/nix-community/go-nix/pkg/hash"
@@ -94,7 +93,7 @@ func renderNar(
 	}
 
 	// render the NAR file
-	err := exporter.Export(
+	err := storev1pb.Export(
 		w,
 		pathInfo,
 		func(directoryDigest []byte) (*castorev1pb.Directory, error) {
diff --git a/tvix/nar-bridge/pkg/importer/importer_test.go b/tvix/nar-bridge/pkg/importer/importer_test.go
index de0548da9398..6754fd005551 100644
--- a/tvix/nar-bridge/pkg/importer/importer_test.go
+++ b/tvix/nar-bridge/pkg/importer/importer_test.go
@@ -11,35 +11,9 @@ import (
 	castorev1pb "code.tvl.fyi/tvix/castore/protos"
 	"code.tvl.fyi/tvix/nar-bridge/pkg/importer"
 	storev1pb "code.tvl.fyi/tvix/store/protos"
-	"github.com/google/go-cmp/cmp"
 	"github.com/stretchr/testify/require"
-	"google.golang.org/protobuf/testing/protocmp"
-	"lukechampine.com/blake3"
 )
 
-func requireProtoEq(t *testing.T, expected interface{}, actual interface{}) {
-	if diff := cmp.Diff(expected, actual, protocmp.Transform()); diff != "" {
-		t.Errorf("unexpected difference:\n%v", diff)
-	}
-}
-
-func mustDirectoryDigest(d *castorev1pb.Directory) []byte {
-	dgst, err := d.Digest()
-	if err != nil {
-		panic(err)
-	}
-	return dgst
-}
-
-func mustBlobDigest(r io.Reader) []byte {
-	hasher := blake3.New(32, nil)
-	_, err := io.Copy(hasher, r)
-	if err != nil {
-		panic(err)
-	}
-	return hasher.Sum([]byte{})
-}
-
 func TestSymlink(t *testing.T) {
 	f, err := os.Open("../../testdata/symlink.nar")
 	require.NoError(t, err)
diff --git a/tvix/nar-bridge/pkg/exporter/full_test.go b/tvix/nar-bridge/pkg/importer/roundtrip_test.go
index 4875c08e2133..89603cfcfdbe 100644
--- a/tvix/nar-bridge/pkg/exporter/full_test.go
+++ b/tvix/nar-bridge/pkg/importer/roundtrip_test.go
@@ -1,4 +1,4 @@
-package exporter_test
+package importer_test
 
 import (
 	"bytes"
@@ -10,15 +10,15 @@ import (
 	"testing"
 
 	castorev1pb "code.tvl.fyi/tvix/castore/protos"
-	"code.tvl.fyi/tvix/nar-bridge/pkg/exporter"
 	"code.tvl.fyi/tvix/nar-bridge/pkg/importer"
+	storev1pb "code.tvl.fyi/tvix/store/protos"
 	"github.com/stretchr/testify/require"
-	"lukechampine.com/blake3"
 )
 
-func TestFull(t *testing.T) {
-	// We pipe nar_1094wph9z4nwlgvsd53abfz8i117ykiv5dwnq9nnhz846s7xqd7d.nar to the exporter,
-	// and store all the file contents and directory objects received in two hashmaps.
+func TestRoundtrip(t *testing.T) {
+	// We pipe nar_1094wph9z4nwlgvsd53abfz8i117ykiv5dwnq9nnhz846s7xqd7d.nar to
+	// storev1pb.Export, and store all the file contents and directory objects
+	// received in two hashmaps.
 	// We then feed it to the writer, and test we come up with the same NAR file.
 
 	f, err := os.Open("../../testdata/nar_1094wph9z4nwlgvsd53abfz8i117ykiv5dwnq9nnhz846s7xqd7d.nar")
@@ -57,7 +57,7 @@ func TestFull(t *testing.T) {
 
 	// done populating everything, now actually test the export :-)
 	var buf bytes.Buffer
-	err = exporter.Export(
+	err = storev1pb.Export(
 		&buf,
 		pathInfo,
 		func(directoryDgst []byte) (*castorev1pb.Directory, error) {
@@ -79,20 +79,3 @@ func TestFull(t *testing.T) {
 	require.NoError(t, err, "exporter shouldn't fail")
 	require.Equal(t, narContents, buf.Bytes())
 }
-
-func mustDirectoryDigest(d *castorev1pb.Directory) []byte {
-	dgst, err := d.Digest()
-	if err != nil {
-		panic(err)
-	}
-	return dgst
-}
-
-func mustBlobDigest(r io.Reader) []byte {
-	hasher := blake3.New(32, nil)
-	_, err := io.Copy(hasher, r)
-	if err != nil {
-		panic(err)
-	}
-	return hasher.Sum([]byte{})
-}
diff --git a/tvix/nar-bridge/pkg/importer/util_test.go b/tvix/nar-bridge/pkg/importer/util_test.go
new file mode 100644
index 000000000000..623253ed1a35
--- /dev/null
+++ b/tvix/nar-bridge/pkg/importer/util_test.go
@@ -0,0 +1,34 @@
+package importer_test
+
+import (
+	"io"
+	"testing"
+
+	castorev1pb "code.tvl.fyi/tvix/castore/protos"
+	"github.com/google/go-cmp/cmp"
+	"google.golang.org/protobuf/testing/protocmp"
+	"lukechampine.com/blake3"
+)
+
+func requireProtoEq(t *testing.T, expected interface{}, actual interface{}) {
+	if diff := cmp.Diff(expected, actual, protocmp.Transform()); diff != "" {
+		t.Errorf("unexpected difference:\n%v", diff)
+	}
+}
+
+func mustDirectoryDigest(d *castorev1pb.Directory) []byte {
+	dgst, err := d.Digest()
+	if err != nil {
+		panic(err)
+	}
+	return dgst
+}
+
+func mustBlobDigest(r io.Reader) []byte {
+	hasher := blake3.New(32, nil)
+	_, err := io.Copy(hasher, r)
+	if err != nil {
+		panic(err)
+	}
+	return hasher.Sum([]byte{})
+}