Skip to content

Commit 85a9fab

Browse files
EdSchoutenEdSchouten
committed
Add better "reference equality" for depsets
In our case it's a bit annoying to provide reference equality, for the reason that depsets get written to/read from storage. We may duplicate depsets, but may also collapse identical instances. As we should try to remain compatible with Bazel in this regard, give every depset an identifier. For empty depsets we keep it empty, as all empty depsets are equal to each other. For non-empty depsets we number them sequentially, using a cryptographic hash of the current function's key as the initial offset. This is all suboptimal, but probably the best we can do for now. In the future we should investigate whether there are better things we can do here. Maybe we should just .to_list() the depset when any attempt is made to hash/compare it?
1 parent a94bbfa commit 85a9fab

10 files changed

Lines changed: 356 additions & 132 deletions

File tree

‎pkg/model/analysis/compiled_bzl_file.go‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package analysis
33
import (
44
"bytes"
55
"context"
6+
"crypto/sha256"
67
"errors"
78
"fmt"
89
"strings"
@@ -14,10 +15,46 @@ import (
1415
model_starlark "github.com/buildbarn/bonanza/pkg/model/starlark"
1516
model_analysis_pb "github.com/buildbarn/bonanza/pkg/proto/model/analysis"
1617

18+
"google.golang.org/protobuf/proto"
19+
1720
"go.starlark.net/starlark"
1821
"go.starlark.net/syntax"
1922
)
2023

24+
// getReferenceEqualIdentifierGenerator creates an identifier generator
25+
// for Starlark values that provide reference equality, such as depsets.
26+
//
27+
// As these identifiers end up getting written to storage, we want to
28+
// use a deterministic process. We therefore hand out these identifiers
29+
// sequentially, using a cryptographic hash of the current key as the
30+
// initial offset.
31+
func (c *baseComputer[TReference, TMetadata]) getReferenceEqualIdentifierGenerator(key model_core.Message[proto.Message, TReference]) (model_starlark.ReferenceEqualIdentifierGenerator, error) {
32+
topLevelKey, _ := model_core.Patch(c.discardingObjectCapturer, key).SortAndSetReferences()
33+
anyKey, err := model_core.MarshalTopLevelAny(topLevelKey)
34+
if err != nil {
35+
return nil, fmt.Errorf("failed to marshal key for creating an identifier generator for reference equal values: %w", err)
36+
}
37+
marshaledKey, err := model_core.MarshalTopLevelMessage(anyKey)
38+
if err != nil {
39+
return nil, fmt.Errorf("failed to marshal key for creating an identifier generator for reference equal values: %w", err)
40+
}
41+
hashedKey := sha256.Sum256(marshaledKey)
42+
43+
var nextIdentifier [16]byte
44+
copy(nextIdentifier[:], hashedKey[:])
45+
46+
return func() []byte {
47+
identifier := append([]byte(nil), nextIdentifier[:]...)
48+
for i := 0; i < len(nextIdentifier); i++ {
49+
nextIdentifier[i]++
50+
if nextIdentifier[i] != 0 {
51+
break
52+
}
53+
}
54+
return identifier
55+
}, nil
56+
}
57+
2158
func (c *baseComputer[TReference, TMetadata]) ComputeCompiledBzlFileValue(ctx context.Context, key *model_analysis_pb.CompiledBzlFile_Key, e CompiledBzlFileEnvironment[TReference, TMetadata]) (PatchedCompiledBzlFileValue, error) {
2259
canonicalLabel, err := label.NewCanonicalLabel(key.Label)
2360
if err != nil {
@@ -80,6 +117,12 @@ func (c *baseComputer[TReference, TMetadata]) ComputeCompiledBzlFileValue(ctx co
80117
return PatchedCompiledBzlFileValue{}, err
81118
}
82119

120+
identifierGenerator, err := c.getReferenceEqualIdentifierGenerator(model_core.NewSimpleMessage[TReference](proto.Message(key)))
121+
if err != nil {
122+
return PatchedCompiledBzlFileValue{}, err
123+
}
124+
thread.SetLocal(model_starlark.ReferenceEqualIdentifierGeneratorKey, identifierGenerator)
125+
83126
globals, err := program.Init(thread, bzlFileBuiltins)
84127
if err != nil {
85128
if !errors.Is(err, evaluation.ErrMissingDependency) {
@@ -145,6 +188,12 @@ func (c *baseComputer[TReference, TMetadata]) ComputeCompiledBzlFileFunctionFact
145188
return nil, bzlFileBuiltinsErr
146189
}
147190

191+
identifierGenerator, err := c.getReferenceEqualIdentifierGenerator(model_core.NewSimpleMessage[TReference](proto.Message(key)))
192+
if err != nil {
193+
return nil, err
194+
}
195+
thread.SetLocal(model_starlark.ReferenceEqualIdentifierGeneratorKey, identifierGenerator)
196+
148197
program, err := starlark.CompiledProgram(bytes.NewBuffer(compiledBzlFile.Message.CompiledProgram.GetCode()))
149198
if err != nil {
150199
return nil, fmt.Errorf("failed to load previously compiled file %#v: %w", key.Label, err)

‎pkg/model/analysis/configured_target.go‎

Lines changed: 62 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -212,7 +212,10 @@ var emptyRunfilesValue = &model_starlark_pb.Value{
212212
},
213213
}
214214

215-
func getSingleFileConfiguredTargetValue[TMetadata model_core.WalkableReferenceMetadata](file model_core.PatchedMessage[*model_starlark_pb.File, TMetadata]) PatchedConfiguredTargetValue {
215+
func getSingleFileConfiguredTargetValue[TMetadata model_core.WalkableReferenceMetadata](
216+
file model_core.PatchedMessage[*model_starlark_pb.File, TMetadata],
217+
identifierGenerator model_starlark.ReferenceEqualIdentifierGenerator,
218+
) PatchedConfiguredTargetValue {
216219
fileValue := &model_starlark_pb.Value{
217220
Kind: &model_starlark_pb.Value_File{
218221
File: file.Message,
@@ -252,6 +255,7 @@ func getSingleFileConfiguredTargetValue[TMetadata model_core.WalkableReferenceMe
252255
Leaf: fileValue,
253256
},
254257
}},
258+
Identifier: identifierGenerator(),
255259
},
256260
},
257261
},
@@ -559,6 +563,11 @@ func (c *baseComputer[TReference, TMetadata]) ComputeConfiguredTargetValue(ctx c
559563
), nil
560564
case *model_starlark_pb.Target_Definition_PredeclaredOutputFileTarget:
561565
// Handcraft a DefaultInfo provider for this source file.
566+
identifierGenerator, err := c.getReferenceEqualIdentifierGenerator(model_core.Nested(key, proto.Message(key.Message)))
567+
if err != nil {
568+
return PatchedConfiguredTargetValue{}, err
569+
}
570+
562571
configurationReference := model_core.Patch(e, model_core.Nested(key, key.Message.ConfigurationReference))
563572
return getSingleFileConfiguredTargetValue(
564573
model_core.NewPatchedMessage(
@@ -572,6 +581,7 @@ func (c *baseComputer[TReference, TMetadata]) ComputeConfiguredTargetValue(ctx c
572581
},
573582
configurationReference.Patcher,
574583
),
584+
identifierGenerator,
575585
), nil
576586
case *model_starlark_pb.Target_Definition_RuleTarget:
577587
ruleTarget := targetKind.RuleTarget
@@ -1193,8 +1203,10 @@ func (c *baseComputer[TReference, TMetadata]) ComputeConfiguredTargetValue(ctx c
11931203
splitAttrValues[namedAttr.Name] = splitAttrValue
11941204
}
11951205

1196-
filesDepset := model_starlark.NewDepsetFromList[TReference, TMetadata](filesDepsetElements, model_starlark_pb.Depset_DEFAULT)
1197-
filesElements, err := filesDepset.ToList(thread)
1206+
filesElements, err := model_starlark.NewDepsetContentsFromList[TReference, TMetadata](
1207+
filesDepsetElements,
1208+
model_starlark_pb.Depset_DEFAULT,
1209+
).ToList(thread)
11981210
if err != nil {
11991211
return PatchedConfiguredTargetValue{}, fmt.Errorf("converting files depset to list: %w", err)
12001212
}
@@ -1329,6 +1341,12 @@ func (c *baseComputer[TReference, TMetadata]) ComputeConfiguredTargetValue(ctx c
13291341
)
13301342
})
13311343

1344+
identifierGenerator, err := c.getReferenceEqualIdentifierGenerator(model_core.Nested(key, proto.Message(key.Message)))
1345+
if err != nil {
1346+
return PatchedConfiguredTargetValue{}, err
1347+
}
1348+
thread.SetLocal(model_starlark.ReferenceEqualIdentifierGeneratorKey, identifierGenerator)
1349+
13321350
// Invoke the rule implementation function. Instead of
13331351
// calling it directly, we call the rule implementation
13341352
// wrapper function, having both the actual
@@ -1403,6 +1421,13 @@ func (c *baseComputer[TReference, TMetadata]) ComputeConfiguredTargetValue(ctx c
14031421
providerInstancesByIdentifier[providerIdentifier] = providerInstance
14041422
}
14051423

1424+
emptyDepset := model_starlark.NewDepset(
1425+
model_starlark.NewDepsetContentsFromList[TReference, TMetadata](
1426+
/* children = */ nil,
1427+
model_starlark_pb.Depset_DEFAULT,
1428+
),
1429+
/* identifierGenerator = */ nil,
1430+
)
14061431
defaultInfoProviderInstanceProperties := model_starlark.NewProviderInstanceProperties[TReference, TMetadata](&defaultInfoProviderIdentifier, false, nil)
14071432
if defaultInfo, ok := providerInstancesByIdentifier[defaultInfoProviderIdentifier]; ok {
14081433
// Rule returned DefaultInfo. Make sure that
@@ -1422,15 +1447,11 @@ func (c *baseComputer[TReference, TMetadata]) ComputeConfiguredTargetValue(ctx c
14221447
switch attrName {
14231448
case "data_runfiles", "default_runfiles":
14241449
if attrValue == starlark.None {
1425-
attrValue = model_starlark.NewRunfiles(
1426-
model_starlark.NewDepsetFromList[TReference, TMetadata](nil, model_starlark_pb.Depset_DEFAULT),
1427-
model_starlark.NewDepsetFromList[TReference, TMetadata](nil, model_starlark_pb.Depset_DEFAULT),
1428-
model_starlark.NewDepsetFromList[TReference, TMetadata](nil, model_starlark_pb.Depset_DEFAULT),
1429-
)
1450+
attrValue = model_starlark.NewRunfiles(emptyDepset, emptyDepset, emptyDepset)
14301451
}
14311452
case "files":
14321453
if attrValue == starlark.None {
1433-
attrValue = model_starlark.NewDepsetFromList[TReference, TMetadata](nil, model_starlark_pb.Depset_DEFAULT)
1454+
attrValue = emptyDepset
14341455
}
14351456
}
14361457
newAttrs[attrName] = attrValue
@@ -1442,17 +1463,9 @@ func (c *baseComputer[TReference, TMetadata]) ComputeConfiguredTargetValue(ctx c
14421463
providerInstancesByIdentifier[defaultInfoProviderIdentifier] = model_starlark.NewStructFromDict[TReference, TMetadata](
14431464
defaultInfoProviderInstanceProperties,
14441465
map[string]any{
1445-
"data_runfiles": model_starlark.NewRunfiles[TReference](
1446-
model_starlark.NewDepsetFromList[TReference, TMetadata](nil, model_starlark_pb.Depset_DEFAULT),
1447-
model_starlark.NewDepsetFromList[TReference, TMetadata](nil, model_starlark_pb.Depset_DEFAULT),
1448-
model_starlark.NewDepsetFromList[TReference, TMetadata](nil, model_starlark_pb.Depset_DEFAULT),
1449-
),
1450-
"default_runfiles": model_starlark.NewRunfiles[TReference](
1451-
model_starlark.NewDepsetFromList[TReference, TMetadata](nil, model_starlark_pb.Depset_DEFAULT),
1452-
model_starlark.NewDepsetFromList[TReference, TMetadata](nil, model_starlark_pb.Depset_DEFAULT),
1453-
model_starlark.NewDepsetFromList[TReference, TMetadata](nil, model_starlark_pb.Depset_DEFAULT),
1454-
),
1455-
"files": model_starlark.NewDepsetFromList[TReference, TMetadata](nil, model_starlark_pb.Depset_DEFAULT),
1466+
"data_runfiles": model_starlark.NewRunfiles[TReference](emptyDepset, emptyDepset, emptyDepset),
1467+
"default_runfiles": model_starlark.NewRunfiles[TReference](emptyDepset, emptyDepset, emptyDepset),
1468+
"files": emptyDepset,
14561469
"files_to_run": model_starlark.NewStructFromDict[TReference, TMetadata](
14571470
model_starlark.NewProviderInstanceProperties[TReference, TMetadata](&filesToRunProviderIdentifier, false, nil),
14581471
map[string]any{
@@ -1596,13 +1609,18 @@ func (c *baseComputer[TReference, TMetadata]) ComputeConfiguredTargetValue(ctx c
15961609
), nil
15971610
case *model_starlark_pb.Target_Definition_SourceFileTarget:
15981611
// Handcraft a DefaultInfo provider for this source file.
1612+
identifierGenerator, err := c.getReferenceEqualIdentifierGenerator(model_core.Nested(key, proto.Message(key.Message)))
1613+
if err != nil {
1614+
return PatchedConfiguredTargetValue{}, err
1615+
}
15991616
return getSingleFileConfiguredTargetValue(
16001617
model_core.NewSimplePatchedMessage[TMetadata](
16011618
&model_starlark_pb.File{
16021619
Label: targetLabel.String(),
16031620
Type: model_starlark_pb.File_FILE,
16041621
},
16051622
),
1623+
identifierGenerator,
16061624
), nil
16071625
default:
16081626
return PatchedConfiguredTargetValue{}, errors.New("only source file targets and rule targets can be configured")
@@ -1882,12 +1900,13 @@ func (rc *ruleContext[TReference, TMetadata]) AttrNames() []string {
18821900
return attrNames
18831901
}
18841902

1885-
func toSymlinkEntryDepset[TReference object.BasicReference, TMetadata BaseComputerReferenceMetadata](v any) *model_starlark.Depset[TReference, TMetadata] {
1903+
func toSymlinkEntryDepset[TReference object.BasicReference, TMetadata BaseComputerReferenceMetadata](v any, identifierGenerator model_starlark.ReferenceEqualIdentifierGenerator) *model_starlark.Depset[TReference, TMetadata] {
1904+
var entries []any
18861905
switch typedV := v.(type) {
18871906
case *model_starlark.Depset[TReference, TMetadata]:
18881907
return typedV
18891908
case map[string]string:
1890-
entries := make([]any, 0, len(typedV))
1909+
entries = make([]any, 0, len(typedV))
18911910
for _, path := range slices.Sorted(maps.Keys(typedV)) {
18921911
entries = append(entries, model_starlark.NewStructFromDict[TReference, TMetadata](
18931912
nil,
@@ -1897,12 +1916,15 @@ func toSymlinkEntryDepset[TReference object.BasicReference, TMetadata BaseComput
18971916
},
18981917
))
18991918
}
1900-
return model_starlark.NewDepsetFromList[TReference, TMetadata](entries, model_starlark_pb.Depset_DEFAULT)
19011919
case nil:
1902-
return model_starlark.NewDepsetFromList[TReference, TMetadata](nil, model_starlark_pb.Depset_DEFAULT)
1920+
// Return an empty depset.
19031921
default:
19041922
panic("unknown type")
19051923
}
1924+
return model_starlark.NewDepset(
1925+
model_starlark.NewDepsetContentsFromList[TReference, TMetadata](entries, model_starlark_pb.Depset_DEFAULT),
1926+
identifierGenerator,
1927+
)
19061928
}
19071929

19081930
func (ruleContext[TReference, TMetadata]) doRunfiles(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
@@ -1924,10 +1946,19 @@ func (ruleContext[TReference, TMetadata]) doRunfiles(thread *starlark.Thread, b
19241946
return nil, err
19251947
}
19261948

1949+
identifierGeneratorValue := thread.Local(model_starlark.ReferenceEqualIdentifierGeneratorKey)
1950+
if identifierGeneratorValue == nil {
1951+
return nil, errors.New("depsets cannot be created from within this context")
1952+
}
1953+
identifierGenerator := identifierGeneratorValue.(model_starlark.ReferenceEqualIdentifierGenerator)
1954+
19271955
if transitiveFiles == nil {
1928-
transitiveFiles = model_starlark.NewDepsetFromList[TReference, TMetadata](nil, model_starlark_pb.Depset_DEFAULT)
1956+
transitiveFiles = model_starlark.NewDepset(
1957+
model_starlark.NewDepsetContentsFromList[TReference, TMetadata](nil, model_starlark_pb.Depset_DEFAULT),
1958+
nil,
1959+
)
19291960
}
1930-
filesDepset, err := model_starlark.NewDepset(
1961+
filesDepsetContents, err := model_starlark.NewDepsetContents(
19311962
thread,
19321963
files,
19331964
[]*model_starlark.Depset[TReference, TMetadata]{transitiveFiles},
@@ -1938,9 +1969,9 @@ func (ruleContext[TReference, TMetadata]) doRunfiles(thread *starlark.Thread, b
19381969
}
19391970

19401971
return model_starlark.NewRunfiles(
1941-
filesDepset,
1942-
toSymlinkEntryDepset[TReference, TMetadata](rootSymlinks),
1943-
toSymlinkEntryDepset[TReference, TMetadata](symlinks),
1972+
model_starlark.NewDepset(filesDepsetContents, identifierGenerator),
1973+
toSymlinkEntryDepset[TReference, TMetadata](rootSymlinks, identifierGenerator),
1974+
toSymlinkEntryDepset[TReference, TMetadata](symlinks, identifierGenerator),
19441975
), nil
19451976
}
19461977

@@ -2535,7 +2566,7 @@ func (rca *ruleContextActions[TReference, TMetadata]) doRun(thread *starlark.Thr
25352566
}
25362567
}
25372568

2538-
mergedInputs, err := model_starlark.NewDepset(thread, inputsDirect, inputsTransitive, model_starlark_pb.Depset_DEFAULT)
2569+
mergedInputs, err := model_starlark.NewDepsetContents(thread, inputsDirect, inputsTransitive, model_starlark_pb.Depset_DEFAULT)
25392570
if err != nil {
25402571
return nil, err
25412572
}
@@ -2544,7 +2575,7 @@ func (rca *ruleContextActions[TReference, TMetadata]) doRun(thread *starlark.Thr
25442575
return nil, err
25452576
}
25462577

2547-
mergedTools, err := model_starlark.NewDepset[TReference, TMetadata](thread, toolsDirect, nil, model_starlark_pb.Depset_DEFAULT)
2578+
mergedTools, err := model_starlark.NewDepsetContents[TReference, TMetadata](thread, toolsDirect, nil, model_starlark_pb.Depset_DEFAULT)
25482579
if err != nil {
25492580
return nil, err
25502581
}

‎pkg/model/starlark/BUILD.bazel‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ go_library(
2222
"parse_repo_dot_bazel.go",
2323
"path.go",
2424
"provider.go",
25+
"reference_equal.go",
2526
"repo_registrar.go",
2627
"repository_rule.go",
2728
"rule.go",

‎pkg/model/starlark/builtins.go‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -850,7 +850,17 @@ func GetBuiltins[TReference object.BasicReference, TMetadata model_core.Cloneabl
850850
default:
851851
return nil, fmt.Errorf("unknown order %#v", order)
852852
}
853-
return NewDepset(thread, direct, transitive, orderValue)
853+
854+
identifierGenerator := thread.Local(ReferenceEqualIdentifierGeneratorKey)
855+
if identifierGenerator == nil {
856+
return nil, errors.New("depsets cannot be created from within this context")
857+
}
858+
859+
dc, err := NewDepsetContents(thread, direct, transitive, orderValue)
860+
if err != nil {
861+
return nil, err
862+
}
863+
return NewDepset(dc, identifierGenerator.(ReferenceEqualIdentifierGenerator)), nil
854864
},
855865
),
856866
"exec_group": starlark.NewBuiltin(

0 commit comments

Comments
 (0)