Skip to content

Commit e05a22c

Browse files
authored
Avoid quadratic deduping in location expansion (#3636)
Does what it says on the tin. Usually the deduping was not necessary at all because the argument were already coming from a depset or an attribute (though kept for safety), but doing the same work for every arg was truly unnecessary Before: <img width="898" height="222" alt="image" src="https://github.com/user-attachments/assets/1ffce739-20eb-4c3b-8cc6-fed0ba702aa0" /> After: <img width="898" height="150" alt="image" src="https://github.com/user-attachments/assets/2497e70f-3434-4e9d-9dc8-4c9535cd2feb" />
1 parent 7dff955 commit e05a22c

File tree

5 files changed

+48
-40
lines changed

5 files changed

+48
-40
lines changed

cargo/private/cargo_build_script.bzl

Lines changed: 14 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ load(
1919
load(
2020
"//rust/private:utils.bzl",
2121
"dedent",
22+
"deduplicate",
2223
"expand_dict_value_locations",
2324
"find_cc_toolchain",
2425
"find_toolchain",
@@ -485,16 +486,19 @@ def _cargo_build_script_impl(ctx):
485486
variables = getattr(target[platform_common.TemplateVariableInfo], "variables", depset([]))
486487
known_variables.update(variables)
487488

488-
_merge_env_dict(env, expand_dict_value_locations(
489-
ctx,
490-
ctx.attr.build_script_env,
491-
getattr(ctx.attr, "data", []) +
492-
getattr(ctx.attr, "compile_data", []) +
493-
getattr(ctx.attr, "tools", []) +
494-
script_info.data +
495-
script_info.tools,
496-
known_variables,
497-
))
489+
if ctx.attr.build_script_env:
490+
_merge_env_dict(env, expand_dict_value_locations(
491+
ctx,
492+
ctx.attr.build_script_env,
493+
deduplicate(
494+
getattr(ctx.attr, "data", []) +
495+
getattr(ctx.attr, "compile_data", []) +
496+
getattr(ctx.attr, "tools", []) +
497+
script_info.data +
498+
script_info.tools,
499+
),
500+
known_variables,
501+
))
498502

499503
tools = depset(
500504
direct = [

rust/private/rust.bzl

Lines changed: 17 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ load(
3535
"compute_crate_name",
3636
"crate_root_src",
3737
"dedent",
38+
"deduplicate",
3839
"determine_lib_name",
3940
"determine_output_hash",
4041
"expand_dict_value_locations",
@@ -401,13 +402,13 @@ def _rust_test_impl(ctx):
401402

402403
# crate.rustc_env is already expanded upstream in rust_library rule implementation
403404
rustc_env = dict(crate.rustc_env)
404-
data_paths = depset(direct = getattr(ctx.attr, "data", [])).to_list()
405-
rustc_env.update(expand_dict_value_locations(
406-
ctx,
407-
ctx.attr.rustc_env,
408-
data_paths,
409-
{},
410-
))
405+
if ctx.attr.rustc_env:
406+
rustc_env.update(expand_dict_value_locations(
407+
ctx,
408+
ctx.attr.rustc_env,
409+
deduplicate(getattr(ctx.attr, "data", [])),
410+
{},
411+
))
411412
aliases = dict(crate.aliases)
412413
aliases.update(ctx.attr.aliases)
413414

@@ -467,13 +468,15 @@ def _rust_test_impl(ctx):
467468
)
468469
rustc_rmeta_output = generate_output_diagnostics(ctx, rust_metadata)
469470

470-
data_paths = depset(direct = getattr(ctx.attr, "data", [])).to_list()
471-
rustc_env = expand_dict_value_locations(
472-
ctx,
473-
ctx.attr.rustc_env,
474-
data_paths,
475-
{},
476-
)
471+
if ctx.attr.rustc_env:
472+
rustc_env = expand_dict_value_locations(
473+
ctx,
474+
ctx.attr.rustc_env,
475+
deduplicate(getattr(ctx.attr, "data", [])),
476+
{},
477+
)
478+
else:
479+
rustc_env = {}
477480

478481
# Target is a standalone crate. Build the test binary as its own crate.
479482
crate_info_dict = dict(

rust/private/rust_analyzer.bzl

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ load(
2828
"//rust/private:utils.bzl",
2929
"concat",
3030
"dedent",
31-
"dedup_expand_location",
31+
"deduplicate",
3232
"find_toolchain",
3333
)
3434

@@ -270,9 +270,9 @@ def _create_single_crate(ctx, attrs, info):
270270

271271
# TODO: The only imagined use case is an env var holding a filename in the workspace passed to a
272272
# macro like include_bytes!. Other use cases might exist that require more complex logic.
273-
expand_targets = concat([getattr(attrs, attr, []) for attr in ["data", "compile_data"]])
273+
expand_targets = deduplicate(concat([getattr(attrs, attr, []) for attr in ["data", "compile_data"]]))
274274

275-
crate["env"].update({k: dedup_expand_location(ctx, v, expand_targets) for k, v in info.env.items()})
275+
crate["env"].update({k: ctx.expand_location(v, expand_targets) for k, v in info.env.items()})
276276

277277
# Omit when a crate appears to depend on itself (e.g. foo_test crates).
278278
# It can happen a single source file is present in multiple crates - there can

rust/private/utils.bzl

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -250,20 +250,20 @@ def get_preferred_artifact(library_to_link, use_pic):
250250
# We do this to work around a potential crash, see
251251
# https://github.com/bazelbuild/bazel/issues/16664
252252
def dedup_expand_location(ctx, input, targets = []):
253-
return ctx.expand_location(input, _deduplicate(targets))
253+
return ctx.expand_location(input, deduplicate(targets))
254254

255-
def _deduplicate(xs):
255+
def deduplicate(xs):
256256
return {x: True for x in xs}.keys()
257257

258258
def concat(xss):
259259
return [x for xs in xss for x in xs]
260260

261-
def _expand_location_for_build_script_runner(ctx, env, data, known_variables):
261+
def _expand_location_for_build_script_runner(ctx, v, data, known_variables):
262262
"""A trivial helper for `expand_dict_value_locations` and `expand_list_element_locations`
263263
264264
Args:
265265
ctx (ctx): The rule's context object
266-
env (str): The value possibly containing location macros to expand.
266+
v (str): The value possibly containing location macros to expand.
267267
data (sequence of Targets): See one of the parent functions.
268268
known_variables (dict): Make variables (probably from toolchains) to substitute in when doing make variable expansion.
269269
@@ -272,16 +272,16 @@ def _expand_location_for_build_script_runner(ctx, env, data, known_variables):
272272
"""
273273

274274
# Fast-path - both location expansions and make vars have a `$` so we can short-circuit everything.
275-
if "$" not in env:
276-
return env
275+
if "$" not in v:
276+
return v
277277

278278
for directive in ("$(execpath ", "$(location "):
279-
if directive in env:
279+
if directive in v:
280280
# build script runner will expand pwd to execroot for us
281-
env = env.replace(directive, "$${pwd}/" + directive)
281+
v = v.replace(directive, "$${pwd}/" + directive)
282282
return ctx.expand_make_variables(
283-
env,
284-
dedup_expand_location(ctx, env, data),
283+
v,
284+
ctx.expand_location(v, data),
285285
known_variables,
286286
)
287287

@@ -315,7 +315,7 @@ def expand_dict_value_locations(ctx, env, data, known_variables):
315315
Returns:
316316
dict: A dict of environment variables with expanded location macros
317317
"""
318-
return dict([(k, _expand_location_for_build_script_runner(ctx, v, data, known_variables)) for (k, v) in env.items()])
318+
return {k: _expand_location_for_build_script_runner(ctx, v, data, known_variables) for (k, v) in env.items()}
319319

320320
def expand_list_element_locations(ctx, args, data, known_variables):
321321
"""Performs location-macro expansion on a list of string values.

rust/toolchain.bzl

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ load(
1818
load(
1919
"//rust/private:utils.bzl",
2020
"dedent",
21-
"dedup_expand_location",
21+
"deduplicate",
2222
"find_cc_toolchain",
2323
"is_exec_configuration",
2424
"is_std_dylib",
@@ -534,9 +534,10 @@ def _experimental_use_cc_common_link(ctx):
534534
return ctx.attr.experimental_use_cc_common_link[BuildSettingInfo].value
535535

536536
def _expand_flags(ctx, flags, targets):
537+
targets = deduplicate(targets)
537538
expanded_flags = []
538539
for flag in flags:
539-
expanded_flags.append(dedup_expand_location(ctx, flag, targets))
540+
expanded_flags.append(ctx.expand_location(flag, targets))
540541
return expanded_flags
541542

542543
def _rust_toolchain_impl(ctx):

0 commit comments

Comments
 (0)