Dedup ImportVar entries in merge_imports / merge_parsed_imports

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).

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(

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"),
     [

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)