Avoid losing location when vertices duplicate

Signed-off-by: Kohei Tokunaga <[email protected]>
gmh5225 · May 10, 2022 · 959d63f · 959d63f
1 parent 9b45d45
commit 959d63f
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 2 deletions.
diff --git a/client/llb/sourcemap.go b/client/llb/sourcemap.go
@@ -61,7 +61,7 @@ func (smc *sourceMapCollector) Add(dgst digest.Digest, ls []*SourceLocation) {
 		}
 		smc.index[l.SourceMap] = idx
 	}
-	smc.locations[dgst] = ls
+	smc.locations[dgst] = append(smc.locations[dgst], ls...)
 }
 
 func (smc *sourceMapCollector) Marshal(ctx context.Context, co ...ConstraintsOpt) (*pb.Source, error) {

diff --git a/client/llb/state.go b/client/llb/state.go
@@ -199,10 +199,10 @@ func marshal(ctx context.Context, v Vertex, def *Definition, s *sourceMapCollect
 	if opMeta != nil {
 		def.Metadata[dgst] = mergeMetadata(def.Metadata[dgst], *opMeta)
 	}
+	s.Add(dgst, sls)
 	if _, ok := cache[dgst]; ok {
 		return def, nil
 	}
-	s.Add(dgst, sls)
 	def.Def = append(def.Def, dt)
 	cache[dgst] = struct{}{}
 	return def, nil

diff --git a/client/llb/state_test.go b/client/llb/state_test.go
@@ -99,6 +99,44 @@ func TestStateSourceMapMarshal(t *testing.T) {
 	require.Equal(t, int32(0), def.Source.Locations[dgst.String()].Locations[2].SourceIndex)
 	require.Equal(t, 1, len(def.Source.Locations[dgst.String()].Locations[2].Ranges))
 	require.Equal(t, int32(9), def.Source.Locations[dgst.String()].Locations[2].Ranges[0].Start.Line)
+
+	s = Merge([]State{s, Image("myimage",
+		sm1.Location([]*pb.Range{{Start: pb.Position{Line: 10}}}),
+	)})
+	def, err = s.Marshal(context.TODO())
+	require.NoError(t, err)
+	require.Equal(t, 3, len(def.Def))
+	dgst = digest.FromBytes(def.Def[0])
+
+	require.Equal(t, 2, len(def.Source.Infos))
+	require.Equal(t, 2, len(def.Source.Locations))
+
+	require.Equal(t, "foo", def.Source.Infos[0].Filename)
+	require.Equal(t, []byte("data1"), def.Source.Infos[0].Data)
+	require.Nil(t, def.Source.Infos[0].Definition)
+
+	require.Equal(t, "bar", def.Source.Infos[1].Filename)
+	require.Equal(t, []byte("data2"), def.Source.Infos[1].Data)
+	require.Nil(t, def.Source.Infos[1].Definition)
+
+	require.NotNil(t, def.Source.Locations[dgst.String()])
+	require.Equal(t, 4, len(def.Source.Locations[dgst.String()].Locations))
+
+	require.Equal(t, int32(0), def.Source.Locations[dgst.String()].Locations[0].SourceIndex)
+	require.Equal(t, 1, len(def.Source.Locations[dgst.String()].Locations[0].Ranges))
+	require.Equal(t, int32(7), def.Source.Locations[dgst.String()].Locations[0].Ranges[0].Start.Line)
+
+	require.Equal(t, int32(1), def.Source.Locations[dgst.String()].Locations[1].SourceIndex)
+	require.Equal(t, 1, len(def.Source.Locations[dgst.String()].Locations[1].Ranges))
+	require.Equal(t, int32(8), def.Source.Locations[dgst.String()].Locations[1].Ranges[0].Start.Line)
+
+	require.Equal(t, int32(0), def.Source.Locations[dgst.String()].Locations[2].SourceIndex)
+	require.Equal(t, 1, len(def.Source.Locations[dgst.String()].Locations[2].Ranges))
+	require.Equal(t, int32(9), def.Source.Locations[dgst.String()].Locations[2].Ranges[0].Start.Line)
+
+	require.Equal(t, int32(0), def.Source.Locations[dgst.String()].Locations[3].SourceIndex)
+	require.Equal(t, 1, len(def.Source.Locations[dgst.String()].Locations[3].Ranges))
+	require.Equal(t, int32(10), def.Source.Locations[dgst.String()].Locations[3].Ranges[0].Start.Line)
 }
 
 func TestPlatformFromImage(t *testing.T) {