Add stringSet for crane.Optimize (#903)

cmd/crane/cmd/optimize.go

@@ -19,7 +19,6 @@ import (
 
 	"github.com/google/go-containerregistry/pkg/crane"
 	"github.com/spf13/cobra"
-	"k8s.io/apimachinery/pkg/util/sets"
 )
 
 // NewCmdOptimize creates a new cobra.Command for the optimize subcommand.
@@ -34,7 +33,7 @@ func NewCmdOptimize(options *[]crane.Option) *cobra.Command {
 		Args:    cobra.ExactArgs(2),
 		Run: func(_ *cobra.Command, args []string) {
 			src, dst := args[0], args[1]
-			if err := crane.Optimize(src, dst, sets.NewString(files...), *options...); err != nil {
+			if err := crane.Optimize(src, dst, files, *options...); err != nil {
 				log.Fatal(err)
 			}
 		},

pkg/crane/optimize.go

@@ -27,12 +27,12 @@ import (
 	"github.com/google/go-containerregistry/pkg/v1/remote"
 	"github.com/google/go-containerregistry/pkg/v1/tarball"
 	"github.com/google/go-containerregistry/pkg/v1/types"
-	"k8s.io/apimachinery/pkg/util/sets"
 )
 
 // Optimize optimizes a remote image or index from src to dst.
 // THIS API IS EXPERIMENTAL AND SUBJECT TO CHANGE WITHOUT WARNING.
-func Optimize(src, dst string, prioritize sets.String, opt ...Option) error {
+func Optimize(src, dst string, prioritize []string, opt ...Option) error {
+	pset := newStringSet(prioritize)
 	o := makeOptions(opt...)
 	srcRef, err := name.ParseReference(src, o.name...)
 	if err != nil {
@@ -55,11 +55,11 @@ func Optimize(src, dst string, prioritize sets.String, opt ...Option) error {
 		// Handle indexes separately.
 		if o.platform != nil {
 			// If platform is explicitly set, don't optimize the whole index, just the appropriate image.
-			if err := optimizeAndPushImage(desc, dstRef, prioritize, o); err != nil {
+			if err := optimizeAndPushImage(desc, dstRef, pset, o); err != nil {
 				return fmt.Errorf("failed to optimize image: %v", err)
 			}
 		} else {
-			if err := optimizeAndPushIndex(desc, dstRef, prioritize, o); err != nil {
+			if err := optimizeAndPushIndex(desc, dstRef, pset, o); err != nil {
 				return fmt.Errorf("failed to optimize index: %v", err)
 			}
 		}
@@ -69,15 +69,15 @@ func Optimize(src, dst string, prioritize sets.String, opt ...Option) error {
 
 	default:
 		// Assume anything else is an image, since some registries don't set mediaTypes properly.
-		if err := optimizeAndPushImage(desc, dstRef, prioritize, o); err != nil {
+		if err := optimizeAndPushImage(desc, dstRef, pset, o); err != nil {
 			return fmt.Errorf("failed to optimize image: %v", err)
 		}
 	}
 
 	return nil
 }
 
-func optimizeAndPushImage(desc *remote.Descriptor, dstRef name.Reference, prioritize sets.String, o options) error {
+func optimizeAndPushImage(desc *remote.Descriptor, dstRef name.Reference, prioritize stringSet, o options) error {
 	img, err := desc.Image()
 	if err != nil {
 		return err
@@ -95,7 +95,7 @@ func optimizeAndPushImage(desc *remote.Descriptor, dstRef name.Reference, priori
 	return remote.Write(dstRef, oimg, o.remote...)
 }
 
-func optimizeImage(img v1.Image, prioritize sets.String) (sets.String, v1.Image, error) {
+func optimizeImage(img v1.Image, prioritize stringSet) (stringSet, v1.Image, error) {
 	cfg, err := img.ConfigFile()
 	if err != nil {
 		return nil, nil, err
@@ -114,7 +114,7 @@ func optimizeImage(img v1.Image, prioritize sets.String) (sets.String, v1.Image,
 		return nil, nil, err
 	}
 
-	missingFromImage := sets.NewString(prioritize.UnsortedList()...)
+	missingFromImage := newStringSet(prioritize.List())
 	olayers := make([]mutate.Addendum, 0, len(layers))
 	for _, layer := range layers {
 		missingFromLayer := []string{}
@@ -127,7 +127,7 @@ func optimizeImage(img v1.Image, prioritize sets.String) (sets.String, v1.Image,
 		if err != nil {
 			return nil, nil, err
 		}
-		missingFromImage = missingFromImage.Intersection(sets.NewString(missingFromLayer...))
+		missingFromImage = missingFromImage.Intersection(newStringSet(missingFromLayer))
 
 		olayers = append(olayers, mutate.Addendum{
 			Layer:     olayer,
@@ -142,7 +142,7 @@ func optimizeImage(img v1.Image, prioritize sets.String) (sets.String, v1.Image,
 	return missingFromImage, oimg, nil
 }
 
-func optimizeAndPushIndex(desc *remote.Descriptor, dstRef name.Reference, prioritize sets.String, o options) error {
+func optimizeAndPushIndex(desc *remote.Descriptor, dstRef name.Reference, prioritize stringSet, o options) error {
 	idx, err := desc.ImageIndex()
 	if err != nil {
 		return err
@@ -160,13 +160,13 @@ func optimizeAndPushIndex(desc *remote.Descriptor, dstRef name.Reference, priori
 	return remote.WriteIndex(dstRef, oidx, o.remote...)
 }
 
-func optimizeIndex(idx v1.ImageIndex, prioritize sets.String) (sets.String, v1.ImageIndex, error) {
+func optimizeIndex(idx v1.ImageIndex, prioritize stringSet) (stringSet, v1.ImageIndex, error) {
 	im, err := idx.IndexManifest()
 	if err != nil {
 		return nil, nil, err
 	}
 
-	missingFromIndex := sets.NewString(prioritize.UnsortedList()...)
+	missingFromIndex := newStringSet(prioritize.List())
 
 	// Build an image for each child from the base and append it to a new index to produce the result.
 	adds := make([]mutate.IndexAddendum, 0, len(im.Manifests))
@@ -199,3 +199,39 @@ func optimizeIndex(idx v1.ImageIndex, prioritize sets.String) (sets.String, v1.I
 
 	return missingFromIndex, mutate.IndexMediaType(mutate.AppendManifests(empty.Index, adds...), idxType), nil
 }
+
+type stringSet map[string]struct{}
+
+func newStringSet(in []string) stringSet {
+	ss := stringSet{}
+	for _, s := range in {
+		ss[s] = struct{}{}
+	}
+	return ss
+}
+
+func (s stringSet) List() []string {
+	result := make([]string, 0, len(s))
+	for k := range s {
+		result = append(result, k)
+	}
+	return result
+}
+
+func (s stringSet) Intersection(rhs stringSet) stringSet {
+	// To appease ST1016
+	lhs := s
+
+	// Make sure len(lhs) >= len(rhs)
+	if len(lhs) < len(rhs) {
+		return rhs.Intersection(lhs)
+	}
+
+	result := stringSet{}
+	for k := range lhs {
+		if _, ok := rhs[k]; ok {
+			result[k] = struct{}{}
+		}
+	}
+	return result
+}

pkg/crane/optimize_test.go

@@ -0,0 +1,68 @@
+// Copyright 2020 Google LLC All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package crane
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
+)
+
+func TestStringSet(t *testing.T) {
+	for _, tc := range []struct {
+		lhs    []string
+		rhs    []string
+		result []string
+	}{{
+		lhs:    []string{},
+		rhs:    []string{},
+		result: []string{},
+	}, {
+		lhs:    []string{"a"},
+		rhs:    []string{},
+		result: []string{},
+	}, {
+		lhs:    []string{},
+		rhs:    []string{"a"},
+		result: []string{},
+	}, {
+		lhs:    []string{"a", "b", "c"},
+		rhs:    []string{"a", "b", "c"},
+		result: []string{"a", "b", "c"},
+	}, {
+		lhs:    []string{"a", "b"},
+		rhs:    []string{"a"},
+		result: []string{"a"},
+	}, {
+		lhs:    []string{"a"},
+		rhs:    []string{"a", "b"},
+		result: []string{"a"},
+	}} {
+		got := newStringSet(tc.lhs).Intersection(newStringSet(tc.rhs))
+		want := newStringSet(tc.result)
+		if diff := cmp.Diff(want, got); diff != "" {
+			t.Errorf("%v.intersect(%v) (-want +got): %s", tc.lhs, tc.rhs, diff)
+		}
+
+		less := func(a, b string) bool {
+			return strings.Compare(a, b) <= -1
+		}
+		if diff := cmp.Diff(tc.result, got.List(), cmpopts.SortSlices(less)); diff != "" {
+			t.Errorf("%v.List() (-want +got): = %v", tc.result, diff)
+		}
+	}
+}