llb: carry platform from inputs for merge/diff

Signed-off-by: Tonis Tiigi <[email protected]>

Author: tonistiigi

client/llb/diff.go

@@ -104,5 +104,5 @@ func Diff(lower, upper State, opts ...ConstraintsOpt) State {
 	for _, o := range opts {
 		o.SetConstraintsOption(&c)
 	}
-	return NewState(NewDiff(lower, upper, c).Output())
+	return lower.WithOutput(NewDiff(lower, upper, c).Output())
 }

client/llb/merge.go

@@ -92,5 +92,5 @@ func Merge(inputs []State, opts ...ConstraintsOpt) State {
 		o.SetConstraintsOption(&c)
 	}
 	addCap(&c, pb.CapMergeOp)
-	return NewState(NewMerge(filteredInputs, c).Output())
+	return filteredInputs[0].WithOutput(NewMerge(filteredInputs, c).Output())
 }

client/llb/state_test.go

@@ -213,17 +213,56 @@ func TestPlatformFromImage(t *testing.T) {
 	require.Equal(t, "docker-image://docker.io/library/srcimage:latest", src.Source.Identifier)
 	require.Equal(t, "s390x", vtx.Platform.Architecture)
 }
+
+func TestPlatformFromImageWithMerge(t *testing.T) {
+	t.Parallel()
+
+	s := Image("srcimage", LinuxS390x)
+
+	s2 := Scratch().File(Mkdir("/foo", 0700).Mkfile("/bar", 0600, []byte("bar")))
+
+	dest := Merge([]State{s, s2}).Run(Args([]string{"aftermerge"}))
+
+	def, err := dest.Marshal(context.TODO(), LinuxPpc64le)
+	require.NoError(t, err)
+
+	m, arr := parseDef(t, def.Def)
+	_ = m
+	require.Equal(t, 5, len(arr))
+
+	dgst, idx := last(t, arr)
+	require.Equal(t, 0, idx)
+
+	vtx, ok := m[dgst]
+	require.Equal(t, true, ok)
+
 	_, ok = vtx.Op.(*pb.Op_Exec)
 	require.Equal(t, true, ok)
 	require.Equal(t, "s390x", vtx.Platform.Architecture)
 
 	vtx, ok = m[vtx.Inputs[0].Digest]
 	require.Equal(t, true, ok)
 
-	src, ok = vtx.Op.(*pb.Op_Source)
+	_, ok = vtx.Op.(*pb.Op_Merge)
+	require.Equal(t, true, ok)
+	require.Nil(t, vtx.Platform)
+
+	mainVtx := vtx
+	vtx, ok = m[vtx.Inputs[0].Digest]
+	require.Equal(t, true, ok)
+
+	src, ok := vtx.Op.(*pb.Op_Source)
 	require.Equal(t, true, ok)
 	require.Equal(t, "docker-image://docker.io/library/srcimage:latest", src.Source.Identifier)
 	require.Equal(t, "s390x", vtx.Platform.Architecture)
+
+	vtx, ok = m[mainVtx.Inputs[1].Digest]
+	require.Equal(t, true, ok)
+
+	f, ok := vtx.Op.(*pb.Op_File)
+	require.Equal(t, true, ok)
+	require.Equal(t, 2, len(f.File.Actions))
+	require.Nil(t, vtx.Platform)
 }
 
 func getEnvHelper(t *testing.T, s State, k string) (string, bool) {