refactor(nixery): Extract layering logic into separate package r/4105

This will be required for making a standalone, Nixery-style image
builder function usable from Nix.

Change-Id: I5e36348bd4c32d249d56f6628cd046916691319f
Reviewed-on: https://cl.tvl.fyi/c/depot/+/5601
Tested-by: BuildkiteCI
Reviewed-by: sterni <sternenseemann@systemli.org>

4 files changed, 25 insertions, 21 deletions

diff --git a/tools/nixery/builder/archive.go b/tools/nixery/builder/archive.go
index 3bc02ab4d5b8..8763e4cb8566 100644
--- a/tools/nixery/builder/archive.go
+++ b/tools/nixery/builder/archive.go
@@ -16,6 +16,8 @@ import (
 	"io"
 	"os"
 	"path/filepath"
+
+	"github.com/google/nixery/layers"
 )
 
 // Create a new compressed tarball from each of the paths in the list
@@ -23,7 +25,7 @@ import (
 //
 // The uncompressed tarball is hashed because image manifests must
 // contain both the hashes of compressed and uncompressed layers.
-func packStorePaths(l *layer, w io.Writer) (string, error) {
+func packStorePaths(l *layers.Layer, w io.Writer) (string, error) {
 	shasum := sha256.New()
 	gz := gzip.NewWriter(w)
 	multi := io.MultiWriter(shasum, gz)
diff --git a/tools/nixery/builder/builder.go b/tools/nixery/builder/builder.go
index 37c9b9fcb763..901373b57e99 100644
--- a/tools/nixery/builder/builder.go
+++ b/tools/nixery/builder/builder.go
@@ -23,6 +23,7 @@ import (
 	"strings"
 
 	"github.com/google/nixery/config"
+	"github.com/google/nixery/layers"
 	"github.com/google/nixery/manifest"
 	"github.com/google/nixery/storage"
 	log "github.com/sirupsen/logrus"
@@ -39,7 +40,7 @@ type State struct {
 	Storage storage.Backend
 	Cache   *LocalCache
 	Cfg     config.Config
-	Pop     Popularity
+	Pop     layers.Popularity
 }
 
 // Architecture represents the possible CPU architectures for which
@@ -117,7 +118,7 @@ type ImageResult struct {
 	Pkgs  []string `json:"pkgs"`
 
 	// These fields are populated in case of success
-	Graph        runtimeGraph `json:"runtimeGraph"`
+	Graph        layers.RuntimeGraph `json:"runtimeGraph"`
 	SymlinkLayer struct {
 		Size    int    `json:"size"`
 		TarHash string `json:"tarHash"`
@@ -281,7 +282,7 @@ func prepareImage(s *State, image *Image) (*ImageResult, error) {
 // added only after successful uploads, which guarantees that entries
 // retrieved from the cache are present in the bucket.
 func prepareLayers(ctx context.Context, s *State, image *Image, result *ImageResult) ([]manifest.Entry, error) {
-	grouped := groupLayers(&result.Graph, &s.Pop, LayerBudget)
+	grouped := layers.GroupLayers(&result.Graph, &s.Pop, LayerBudget)
 
 	var entries []manifest.Entry
 
@@ -318,7 +319,7 @@ func prepareLayers(ctx context.Context, s *State, image *Image, result *ImageRes
 
 			var pkgs []string
 			for _, p := range l.Contents {
-				pkgs = append(pkgs, packageFromPath(p))
+				pkgs = append(pkgs, layers.PackageFromPath(p))
 			}
 
 			log.WithFields(log.Fields{
diff --git a/tools/nixery/builder/layers.go b/tools/nixery/layers/layers.go
index 5e37e626810f..131a0cdbed77 100644
--- a/tools/nixery/builder/layers.go
+++ b/tools/nixery/layers/layers.go
@@ -103,7 +103,7 @@
 //
 // Layer budget: 10
 // Layers: { E }, { D, F }, { A }, { B }, { C }
-package builder
+package layers
 
 import (
 	"crypto/sha1"
@@ -121,7 +121,7 @@ import (
 // dependencies of a derivation.
 //
 // This is generated in Nix by using the exportReferencesGraph feature.
-type runtimeGraph struct {
+type RuntimeGraph struct {
 	References struct {
 		Graph []string `json:"graph"`
 	} `json:"exportReferencesGraph"`
@@ -142,19 +142,19 @@ type Popularity = map[string]int
 
 // Layer represents the data returned for each layer that Nix should
 // build for the container image.
-type layer struct {
+type Layer struct {
 	Contents    []string `json:"contents"`
 	MergeRating uint64
 }
 
 // Hash the contents of a layer to create a deterministic identifier that can be
 // used for caching.
-func (l *layer) Hash() string {
+func (l *Layer) Hash() string {
 	sum := sha1.Sum([]byte(strings.Join(l.Contents, ":")))
 	return fmt.Sprintf("%x", sum)
 }
 
-func (a layer) merge(b layer) layer {
+func (a Layer) merge(b Layer) Layer {
 	a.Contents = append(a.Contents, b.Contents...)
 	a.MergeRating += b.MergeRating
 	return a
@@ -177,7 +177,7 @@ var nixRegexp = regexp.MustCompile(`^/nix/store/[a-z0-9]+-`)
 
 // PackageFromPath returns the name of a Nix package based on its
 // output store path.
-func packageFromPath(path string) string {
+func PackageFromPath(path string) string {
 	return nixRegexp.ReplaceAllString(path, "")
 }
 
@@ -185,7 +185,7 @@ func packageFromPath(path string) string {
 // the dot format used by GraphViz, into which the dependency graph
 // can be rendered.
 func (c *closure) DOTID() string {
-	return packageFromPath(c.Path)
+	return PackageFromPath(c.Path)
 }
 
 // bigOrPopular checks whether this closure should be considered for
@@ -228,7 +228,7 @@ func insertEdges(graph *simple.DirectedGraph, cmap *map[string]*closure, node *c
 }
 
 // Create a graph structure from the references supplied by Nix.
-func buildGraph(refs *runtimeGraph, pop *Popularity) *simple.DirectedGraph {
+func buildGraph(refs *RuntimeGraph, pop *Popularity) *simple.DirectedGraph {
 	cmap := make(map[string]*closure)
 	graph := simple.NewDirectedGraph()
 
@@ -288,7 +288,7 @@ func buildGraph(refs *runtimeGraph, pop *Popularity) *simple.DirectedGraph {
 // Extracts a subgraph starting at the specified root from the
 // dominator tree. The subgraph is converted into a flat list of
 // layers, each containing the store paths and merge rating.
-func groupLayer(dt *flow.DominatorTree, root *closure) layer {
+func groupLayer(dt *flow.DominatorTree, root *closure) Layer {
 	size := root.Size
 	contents := []string{root.Path}
 	children := dt.DominatedBy(root.ID())
@@ -305,7 +305,7 @@ func groupLayer(dt *flow.DominatorTree, root *closure) layer {
 	// Contents are sorted to ensure that hashing is consistent
 	sort.Strings(contents)
 
-	return layer{
+	return Layer{
 		Contents:    contents,
 		MergeRating: uint64(root.Popularity) * size,
 	}
@@ -316,10 +316,10 @@ func groupLayer(dt *flow.DominatorTree, root *closure) layer {
 //
 // Layers are merged together until they fit into the layer budget,
 // based on their merge rating.
-func dominate(budget int, graph *simple.DirectedGraph) []layer {
+func dominate(budget int, graph *simple.DirectedGraph) []Layer {
 	dt := flow.Dominators(graph.Node(0), graph)
 
-	var layers []layer
+	var layers []Layer
 	for _, n := range dt.DominatedBy(dt.Root().ID()) {
 		layers = append(layers, groupLayer(&dt, n.(*closure)))
 	}
@@ -347,7 +347,7 @@ func dominate(budget int, graph *simple.DirectedGraph) []layer {
 // groupLayers applies the algorithm described above the its input and returns a
 // list of layers, each consisting of a list of Nix store paths that it should
 // contain.
-func groupLayers(refs *runtimeGraph, pop *Popularity, budget int) []layer {
+func GroupLayers(refs *RuntimeGraph, pop *Popularity, budget int) []Layer {
 	graph := buildGraph(refs, pop)
 	return dominate(budget, graph)
 }
diff --git a/tools/nixery/main.go b/tools/nixery/main.go
index 2e633e0898cd..8fe1679cfad8 100644
--- a/tools/nixery/main.go
+++ b/tools/nixery/main.go
@@ -26,6 +26,7 @@ import (
 
 	"github.com/google/nixery/builder"
 	"github.com/google/nixery/config"
+	"github.com/google/nixery/layers"
 	"github.com/google/nixery/logs"
 	mf "github.com/google/nixery/manifest"
 	"github.com/google/nixery/storage"
@@ -52,7 +53,7 @@ var (
 
 // Downloads the popularity information for the package set from the
 // URL specified in Nixery's configuration.
-func downloadPopularity(url string) (builder.Popularity, error) {
+func downloadPopularity(url string) (layers.Popularity, error) {
 	resp, err := http.Get(url)
 	if err != nil {
 		return nil, err
@@ -67,7 +68,7 @@ func downloadPopularity(url string) (builder.Popularity, error) {
 		return nil, err
 	}
 
-	var pop builder.Popularity
+	var pop layers.Popularity
 	err = json.Unmarshal(j, &pop)
 	if err != nil {
 		return nil, err
@@ -246,7 +247,7 @@ func main() {
 		log.WithError(err).Fatal("failed to instantiate build cache")
 	}
 
-	var pop builder.Popularity
+	var pop layers.Popularity
 	if cfg.PopUrl != "" {
 		pop, err = downloadPopularity(cfg.PopUrl)
 		if err != nil {