refactor: add typed sync.Map wrapper to eliminate type assertions

Introduce util/bkmaps.SyncMap[K, V] as a generic typed
wrapper around sync.Map. Migrate usages in definition,
builder, solver, and contentutil packages, removing
unchecked type assertions at each call site.

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>

client/llb/definition.go

@@ -5,6 +5,7 @@ import (
 	"sync"
 
 	"github.com/moby/buildkit/solver/pb"
+	"github.com/moby/buildkit/util/bkmaps"
 	digest "github.com/opencontainers/go-digest"
 	ocispecs "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/pkg/errors"
@@ -24,7 +25,7 @@ type DefinitionOp struct {
 	platforms  map[digest.Digest]*ocispecs.Platform
 	dgst       digest.Digest
 	index      pb.OutputIndex
-	inputCache *sync.Map // shared and written among DefinitionOps so avoid race on this map using sync.Map
+	inputCache *bkmaps.SyncMap[string, []*DefinitionOp] // shared and written among DefinitionOps so avoid race on this map using sync.Map
 }
 
 // NewDefinitionOp returns a new operation from a marshalled definition.
@@ -106,7 +107,7 @@ func NewDefinitionOp(def *pb.Definition) (*DefinitionOp, error) {
 		platforms:  platforms,
 		dgst:       dgst,
 		index:      index,
-		inputCache: new(sync.Map),
+		inputCache: new(bkmaps.SyncMap[string, []*DefinitionOp]),
 	}, nil
 }
 
@@ -186,11 +187,7 @@ func (d *DefinitionOp) Output() Output {
 }
 
 func (d *DefinitionOp) loadInputCache(dgst digest.Digest) ([]*DefinitionOp, bool) {
-	a, ok := d.inputCache.Load(dgst.String())
-	if ok {
-		return a.([]*DefinitionOp), true
-	}
-	return nil, false
+	return d.inputCache.Load(dgst.String())
 }
 
 func (d *DefinitionOp) storeInputCache(dgst digest.Digest, c []*DefinitionOp) {

frontend/dockerfile/builder/build.go

@@ -3,7 +3,6 @@ package builder
 import (
 	"context"
 	"strings"
-	"sync"
 
 	"github.com/containerd/platforms"
 	"github.com/moby/buildkit/client/llb"
@@ -23,6 +22,7 @@ import (
 	"github.com/moby/buildkit/solver/errdefs"
 	"github.com/moby/buildkit/solver/pb"
 	"github.com/moby/buildkit/solver/result"
+	"github.com/moby/buildkit/util/bkmaps"
 	dockerspec "github.com/moby/docker-image-spec/specs-go/v1"
 	ocispecs "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/pkg/errors"
@@ -128,7 +128,7 @@ func Build(ctx context.Context, c client.Client) (_ *client.Result, err error) {
 		}
 	}
 
-	scanTargets := sync.Map{}
+	scanTargets := bkmaps.SyncMap[string, *dockerfile2llb.SBOMTargets]{}
 
 	rb, err := bc.Build(ctx, func(ctx context.Context, platform *ocispecs.Platform, idx int) (client.Reference, *dockerspec.DockerOCIImage, *dockerspec.DockerOCIImage, error) {
 		opt := convertOpt
@@ -176,14 +176,10 @@ func Build(ctx context.Context, c client.Client) (_ *client.Result, err error) {
 
 	if scanner != nil {
 		if err := rb.EachPlatform(ctx, func(ctx context.Context, id string, p ocispecs.Platform) error {
-			v, ok := scanTargets.Load(id)
+			target, ok := scanTargets.Load(id)
 			if !ok {
 				return errors.Errorf("no scan targets for %s", id)
 			}
-			target, ok := v.(*dockerfile2llb.SBOMTargets)
-			if !ok {
-				return errors.Errorf("invalid scan targets for %T", v)
-			}
 
 			var opts []llb.ConstraintsOpt
 			if target.IgnoreCache {

solver/jobs.go

@@ -12,6 +12,7 @@ import (
 	"github.com/moby/buildkit/session"
 	"github.com/moby/buildkit/solver/errdefs"
 	"github.com/moby/buildkit/util/bklog"
+	"github.com/moby/buildkit/util/bkmaps"
 	"github.com/moby/buildkit/util/flightcontrol"
 	"github.com/moby/buildkit/util/progress"
 	"github.com/moby/buildkit/util/progress/controller"
@@ -339,7 +340,7 @@ type Job struct {
 	pr            *progress.MultiReader
 	pw            progress.Writer
 	span          trace.Span
-	values        sync.Map
+	values        bkmaps.SyncMap[string, any]
 	id            string
 	startedTime   time.Time
 	completedTime time.Time

solver/memorycachestorage.go

@@ -6,6 +6,7 @@ import (
 	"time"
 
 	"github.com/moby/buildkit/session"
+	"github.com/moby/buildkit/util/bkmaps"
 	"github.com/moby/buildkit/util/compression"
 	"github.com/pkg/errors"
 )
@@ -279,11 +280,11 @@ func (s *inMemoryStore) WalkBacklinks(id string, fn func(id string, link CacheIn
 }
 
 func NewInMemoryResultStorage() CacheResultStorage {
-	return &inMemoryResultStore{m: &sync.Map{}}
+	return &inMemoryResultStore{m: &bkmaps.SyncMap[string, Result]{}}
 }
 
 type inMemoryResultStore struct {
-	m *sync.Map
+	m *bkmaps.SyncMap[string, Result]
 }
 
 func (s *inMemoryResultStore) Save(r Result, createdAt time.Time) (CacheResult, error) {
@@ -296,7 +297,7 @@ func (s *inMemoryResultStore) Load(ctx context.Context, res CacheResult) (Result
 	if !ok {
 		return nil, errors.WithStack(ErrNotFound)
 	}
-	return v.(Result), nil
+	return v, nil
 }
 
 func (s *inMemoryResultStore) LoadRemotes(_ context.Context, _ CacheResult, _ *compression.Config, _ session.Group) ([]*Remote, error) {

util/bkmaps/syncmap.go

@@ -0,0 +1,42 @@
+package bkmaps
+
+import "sync"
+
+// SyncMap provides a typed wrapper around sync.Map.
+type SyncMap[K comparable, V any] struct {
+	m sync.Map
+}
+
+// Delete removes the value for a key.
+func (m *SyncMap[K, V]) Delete(key K) {
+	m.m.Delete(key)
+}
+
+// Load returns the value stored in the map for a key, if any.
+func (m *SyncMap[K, V]) Load(key K) (V, bool) {
+	v, ok := m.m.Load(key)
+	if !ok {
+		var zero V
+		return zero, false
+	}
+	return v.(V), true
+}
+
+// LoadOrStore returns the existing value for the key if present.
+// Otherwise it stores and returns the given value.
+func (m *SyncMap[K, V]) LoadOrStore(key K, value V) (V, bool) {
+	v, loaded := m.m.LoadOrStore(key, value)
+	return v.(V), loaded
+}
+
+// Range calls fn sequentially for each key and value present in the map.
+func (m *SyncMap[K, V]) Range(fn func(K, V) bool) {
+	m.m.Range(func(key, value any) bool {
+		return fn(key.(K), value.(V))
+	})
+}
+
+// Store sets the value for a key.
+func (m *SyncMap[K, V]) Store(key K, value V) {
+	m.m.Store(key, value)
+}

util/contentutil/copy.go

@@ -9,8 +9,10 @@ import (
 	"github.com/containerd/containerd/v2/core/content"
 	"github.com/containerd/containerd/v2/core/images"
 	cerrdefs "github.com/containerd/errdefs"
+	"github.com/moby/buildkit/util/bkmaps"
 	"github.com/moby/buildkit/util/resolver/limited"
 	"github.com/moby/buildkit/util/resolver/retryhandler"
+	digest "github.com/opencontainers/go-digest"
 	ocispecs "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/pkg/errors"
 )
@@ -74,10 +76,10 @@ func (r *rc) Seek(offset int64, whence int) (int64, error) {
 	return r.offset, nil
 }
 func CopyChain(ctx context.Context, ingester content.Ingester, provider content.Provider, desc ocispecs.Descriptor, opts ...CopyOption) error {
-	return copyChain(ctx, ingester, provider, desc, &sync.Map{}, opts...)
+	return copyChain(ctx, ingester, provider, desc, &bkmaps.SyncMap[digest.Digest, struct{}]{}, opts...)
 }
 
-func copyChain(ctx context.Context, ingester content.Ingester, provider content.Provider, desc ocispecs.Descriptor, visited *sync.Map, opts ...CopyOption) error {
+func copyChain(ctx context.Context, ingester content.Ingester, provider content.Provider, desc ocispecs.Descriptor, visited *bkmaps.SyncMap[digest.Digest, struct{}], opts ...CopyOption) error {
 	ci := &CopyInfo{}
 	for _, o := range opts {
 		if err := o(ci); err != nil {