feat: add SquashSelected

pkg/mutate/layer_selector.go

@@ -0,0 +1,19 @@
+// Copyright 2024 Sylabs Inc. All rights reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+package mutate
+
+import (
+	v1 "github.com/google/go-containerregistry/pkg/v1"
+)
+
+type LayerSelector interface {
+	// LayerSelected returns true if layer l is selected.
+	Selected(l v1.Layer) bool
+}
+
+type allLayerSelector struct{}
+
+func allLayersSelector() LayerSelector              { return allLayerSelector{} }
+func (s allLayerSelector) Selected(_ v1.Layer) bool { return true }

pkg/mutate/mutate.go

@@ -28,15 +28,35 @@ func SetLayer(i int, l v1.Layer) Mutation {
 	}
 }
 
-// ReplaceLayers replaces all layers in the image with l. The layer is annotated with the specified
-// values.
+// ReplaceLayers replaces all layers in the image with l.
 func ReplaceLayers(l v1.Layer) Mutation {
 	return func(img *image) error {
 		img.overrides = []v1.Layer{l}
 		return nil
 	}
 }
 
+// ReplaceSelectedLayers replaces selected layers in the image with rl.
+func ReplaceSelectedLayers(s LayerSelector, rl v1.Layer) Mutation {
+	return func(img *image) error {
+		var found bool
+		var overrides []v1.Layer
+
+		// Iterate over the current layers, replacing matching layers with rl.
+		for _, l := range img.overrides {
+			if !s.Selected(l) {
+				overrides = append(overrides, l)
+			} else if !found {
+				overrides = append(overrides, rl)
+				found = true
+			}
+		}
+
+		img.overrides = overrides
+		return nil
+	}
+}
+
 // SetHistory replaces the history in an image with the specified entry.
 func SetHistory(history v1.History) Mutation {
 	return func(img *image) error {

pkg/mutate/squash.go

@@ -11,6 +11,7 @@ import (
 	"fmt"
 	"io"
 	"path/filepath"
+	"slices"
 	"strings"
 
 	v1 "github.com/google/go-containerregistry/pkg/v1"
@@ -226,12 +227,17 @@ func (s *imageState) writeHardlinksFor(target string, root entry) (entry, error)
 }
 
 // squash writes a single, squashed TAR layer built from img to w.
-func squash(img v1.Image, w io.Writer) error {
+func squash(img v1.Image, s LayerSelector, w io.Writer) error {
 	ls, err := img.Layers()
 	if err != nil {
 		return fmt.Errorf("retrieving layers: %w", err)
 	}
 
+	// Remove layers that are not selected.
+	ls = slices.DeleteFunc(ls, func(l v1.Layer) bool {
+		return !s.Selected(l)
+	})
+
 	tw := tar.NewWriter(w)
 	defer tw.Close()
 
@@ -272,13 +278,18 @@ func squash(img v1.Image, w io.Writer) error {
 	return nil
 }
 
-// Squash replaces the layers in the base image with a single, squashed layer.
+// Squash replaces all layers in the base image with a single, squashed layer.
 func Squash(base v1.Image) (v1.Image, error) {
+	return SquashSelected(base, allLayersSelector())
+}
+
+// Squash replaces the layers selected by s in the base image with a single, squashed layer.
+func SquashSelected(base v1.Image, s LayerSelector) (v1.Image, error) {
 	opener := func() (io.ReadCloser, error) {
 		pr, pw := io.Pipe()
 
 		go func() {
-			pw.CloseWithError(squash(base, pw))
+			pw.CloseWithError(squash(base, s, pw))
 		}()
 
 		return pr, nil
@@ -289,5 +300,5 @@ func Squash(base v1.Image) (v1.Image, error) {
 		return nil, err
 	}
 
-	return Apply(base, ReplaceLayers(l))
+	return Apply(base, ReplaceSelectedLayers(s, l))
 }

pkg/mutate/squash_test.go

@@ -74,7 +74,7 @@ func TestSquash(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			var b bytes.Buffer
 
-			if err := squash(tt.base, &b); err != nil {
+			if err := squash(tt.base, allLayersSelector(), &b); err != nil {
 				t.Fatal(err)
 			}