diff --git a/packages/reflex-base/news/6737.performance.md b/packages/reflex-base/news/6737.performance.md new file mode 100644 index 00000000000..016372c05bc --- /dev/null +++ b/packages/reflex-base/news/6737.performance.md @@ -0,0 +1 @@ +Dedup `ImportVar` entries in `merge_imports`, keeping merged `VarData` constant-sized for chained var operations (previously entries doubled per nesting level: depth 18 produced ~1.57M entries and a multi-hundred-millisecond merge). diff --git a/packages/reflex-base/src/reflex_base/utils/imports.py b/packages/reflex-base/src/reflex_base/utils/imports.py index 7e614f99524..405c07c8ca7 100644 --- a/packages/reflex-base/src/reflex_base/utils/imports.py +++ b/packages/reflex-base/src/reflex_base/utils/imports.py @@ -22,6 +22,12 @@ def merge_parsed_imports( ) -> ParsedImportDict: """Merge multiple parsed import dicts together. + Entries are concatenated without dedup: this is the component-tree fold + (``_get_all_imports``), where duplication is linear and the result is + collapsed downstream. Per-entry hashing here measurably slows the walk; + the compounding-duplication path (``VarData.merge``) goes through + :func:`merge_imports`, which does dedup. + Args: *imports: The list of import dicts to merge. @@ -40,29 +46,34 @@ def merge_imports( ) -> ParsedImportDict: """Merge multiple import dicts together. + Duplicate ``ImportVar`` entries are dropped (keeping first-seen order) so + repeated merges stay constant-sized instead of growing with every level + of nesting — e.g. chained var operations merge the same operand imports + via both ``_args`` and ``_return``, which doubled entries per level. + Args: *imports: The list of import dicts to merge. Returns: The merged import dicts. """ - all_imports: defaultdict[str, list[ImportVar]] = defaultdict(list) + all_imports: defaultdict[str, dict[ImportVar, None]] = defaultdict(dict) for import_dict in imports: for lib, fields in ( import_dict if isinstance(import_dict, tuple) else import_dict.items() ): # If the lib is an absolute path, we need to prefix it with a $ lib = "$" + lib if lib.startswith(ABSOLUTE_IMPORT_PREFIXES) else lib + bucket = all_imports[lib] if isinstance(fields, (list, tuple, set)): - all_imports[lib].extend( - ImportVar(field) if isinstance(field, str) else field - for field in fields - ) + for field in fields: + bucket[ImportVar(field) if isinstance(field, str) else field] = None else: - all_imports[lib].append( - ImportVar(fields) if isinstance(fields, str) else fields - ) - return all_imports + bucket[ImportVar(fields) if isinstance(fields, str) else fields] = None + merged: defaultdict[str, list[ImportVar]] = defaultdict(list) + for lib, bucket in all_imports.items(): + merged[lib] = list(bucket) + return merged def parse_imports( diff --git a/tests/units/utils/test_imports.py b/tests/units/utils/test_imports.py index 7b456d892c9..b372f7c3ac2 100644 --- a/tests/units/utils/test_imports.py +++ b/tests/units/utils/test_imports.py @@ -4,6 +4,7 @@ ImportVar, ParsedImportDict, merge_imports, + merge_parsed_imports, parse_imports, ) @@ -91,6 +92,43 @@ def test_merge_imports(input_1, input_2, output): assert set(res[key]) == set(output[key]) +def test_merge_imports_dedups_duplicates(): + """Duplicate ImportVars collapse to one entry, preserving first-seen order. + + Regression: without dedup, chained var operations (which merge the same + operand imports via both ``_args`` and ``_return``) doubled the entry + count per nesting level. + """ + merged = merge_imports( + {"react": ["useEffect", "useState", "useEffect"]}, + {"react": ["useState", "useMemo"]}, + ) + assert merged["react"] == [ + ImportVar(tag="useEffect"), + ImportVar(tag="useState"), + ImportVar(tag="useMemo"), + ] + + +def test_merge_imports_keeps_distinct_fields(): + """Entries differing in any field (e.g. alias) are not collapsed.""" + parsed = {"react": [ImportVar(tag="useEffect"), ImportVar(tag="useState")]} + aliased = {"react": [ImportVar(tag="useEffect", alias="effectAlias")]} + merged = merge_imports(parsed, aliased) + assert merged["react"] == [ + ImportVar(tag="useEffect"), + ImportVar(tag="useState"), + ImportVar(tag="useEffect", alias="effectAlias"), + ] + + +def test_merge_parsed_imports_concatenates(): + """The component-tree fold concatenates without dedup (collapsed later).""" + parsed = {"react": [ImportVar(tag="useEffect")]} + merged = merge_parsed_imports(parsed, parsed) + assert merged["react"] == [ImportVar(tag="useEffect"), ImportVar(tag="useEffect")] + + @pytest.mark.parametrize( ("input", "output"), [ diff --git a/tests/units/vars/test_base.py b/tests/units/vars/test_base.py index a8075357fa7..a814009d681 100644 --- a/tests/units/vars/test_base.py +++ b/tests/units/vars/test_base.py @@ -1,7 +1,8 @@ from collections.abc import Mapping, Sequence import pytest -from reflex_base.vars.base import computed_var, figure_out_type +from reflex_base.utils.imports import ImportVar +from reflex_base.vars.base import Var, VarData, computed_var, figure_out_type from reflex.state import State @@ -42,6 +43,28 @@ def test_figure_out_type(value, expected): assert figure_out_type(value) == expected +def test_chained_var_operations_var_data_stays_bounded() -> None: + """Deeply chained var operations keep merged VarData entries deduped. + + Regression: every ``@var_operation`` merges the same operand VarData via + both ``_args`` and ``_return``, so without dedup in ``merge_imports`` the + merged import entries doubled per nesting level (depth 18 produced ~1.57M + entries and a multi-hundred-millisecond merge). + """ + base = Var( + _js_expr="a", + _var_data=VarData(imports={"my-lib": [ImportVar(tag="thing")]}), + ).to(int) + + chained = base + for i in range(18): + chained = chained + i + + var_data = chained._get_all_var_data() + assert var_data is not None + assert dict(var_data.imports)["my-lib"] == (ImportVar(tag="thing"),) + + def test_computed_var_replace() -> None: class StateTest(State): @computed_var(cache=True)