Fix calculated properties when searching in relay list

7 files changed, 87 insertions, 38 deletions

diff --git a/gui/src/renderer/components/select-location/CustomLists.tsx b/gui/src/renderer/components/select-location/CustomLists.tsx
index 843be91055..b87de1fe46 100644
--- a/gui/src/renderer/components/select-location/CustomLists.tsx
+++ b/gui/src/renderer/components/select-location/CustomLists.tsx
@@ -82,7 +82,7 @@ export default function CustomLists(props: CustomListsProps) {
     hideAddList();
   }, []);
 
-  if (searchTerm !== '' && customLists.length === 0) {
+  if (searchTerm !== '' && !customLists.some((list) => list.visible)) {
     return null;
   }
 
diff --git a/gui/src/renderer/components/select-location/LocationRow.tsx b/gui/src/renderer/components/select-location/LocationRow.tsx
index 3d6c818e42..07c7d0003b 100644
--- a/gui/src/renderer/components/select-location/LocationRow.tsx
+++ b/gui/src/renderer/components/select-location/LocationRow.tsx
@@ -150,7 +150,7 @@ interface IProps<C extends LocationSpecification> {
 
 // Renders the rows and its children for countries, cities and relays
 function LocationRow<C extends LocationSpecification>(props: IProps<C>) {
-  const hasChildren = React.Children.count(props.children) > 0;
+  const hasChildren = getLocationChildren(props.source).some((child) => child.visible);
   const buttonRef = useRef<HTMLButtonElement>() as React.RefObject<HTMLButtonElement>;
   const userInvokedExpand = useRef(false);
 
@@ -226,6 +226,10 @@ function LocationRow<C extends LocationSpecification>(props: IProps<C>) {
     }
   }, [props.source.location.customList]);
 
+  if (!props.source.visible) {
+    return null;
+  }
+
   // The selectedRef should only be used if the element is selected
   const selectedRef = props.source.selected ? props.selectedElementRef : undefined;
   return (
@@ -353,6 +357,7 @@ function compareLocation(
   nextLocation: LocationSpecification,
 ): boolean {
   return (
+    oldLocation.visible === nextLocation.visible &&
     oldLocation.label === nextLocation.label &&
     oldLocation.active === nextLocation.active &&
     oldLocation.disabled === nextLocation.disabled &&
@@ -367,16 +372,16 @@ function compareChildren(
   oldLocation: LocationSpecification,
   nextLocation: LocationSpecification,
 ): boolean {
-  const oldChildren = getLocationChildren(oldLocation);
-  const nextChildren = getLocationChildren(nextLocation);
+  const oldVisibleChildren = getLocationChildren(oldLocation).filter((child) => child.visible);
+  const nextVisibleChildren = getLocationChildren(nextLocation).filter((child) => child.visible);
 
   // Children shouldn't be checked if the row is collapsed
   const nextExpanded = 'expanded' in nextLocation && nextLocation.expanded;
 
   return (
-    (!nextExpanded && oldChildren.length > 0 && nextChildren.length > 0) ||
-    (oldChildren.length === nextChildren.length &&
-      oldChildren.every((oldChild, i) => compareLocation(oldChild, nextChildren[i])))
+    (!nextExpanded && oldVisibleChildren.length > 0 && nextVisibleChildren.length > 0) ||
+    (oldVisibleChildren.length === nextVisibleChildren.length &&
+      oldVisibleChildren.every((oldChild, i) => compareLocation(oldChild, nextVisibleChildren[i])))
   );
 }
 
diff --git a/gui/src/renderer/components/select-location/RelayListContext.tsx b/gui/src/renderer/components/select-location/RelayListContext.tsx
index e81ba4986b..d0ee5d6e85 100644
--- a/gui/src/renderer/components/select-location/RelayListContext.tsx
+++ b/gui/src/renderer/components/select-location/RelayListContext.tsx
@@ -27,6 +27,7 @@ import {
   DisabledReason,
   GeographicalRelayList,
   LocationType,
+  RelayLocationCountryWithVisibility,
 } from './select-location-types';
 import { useSelectLocationContext } from './SelectLocationContainer';
 
@@ -127,7 +128,7 @@ export function RelayListContextProvider(props: RelayListContextProviderProps) {
 // Return the final filtered and formatted relay list. This should be the only place in the app
 // where processing of the relay list is performed.
 function useRelayList(
-  relayList: Array<IRelayLocationCountryRedux>,
+  relayList: Array<RelayLocationCountryWithVisibility>,
   expandedLocations?: Array<RelayLocation>,
 ): GeographicalRelayList {
   const locale = useSelector((state) => state.userInterface.locale);
diff --git a/gui/src/renderer/components/select-location/SelectLocation.tsx b/gui/src/renderer/components/select-location/SelectLocation.tsx
index 8cd3dfcde1..1ea7aeb5aa 100644
--- a/gui/src/renderer/components/select-location/SelectLocation.tsx
+++ b/gui/src/renderer/components/select-location/SelectLocation.tsx
@@ -355,7 +355,7 @@ function LocationList<T>(props: CombinedLocationListProps<T>) {
 
   if (
     searchTerm !== '' &&
-    props.relayLocations.length === 0 &&
+    !props.relayLocations.some((country) => country.visible) &&
     (props.specialLocations === undefined || props.specialLocations.length === 0)
   ) {
     return null;
@@ -381,8 +381,8 @@ function NoSearchResult(props: NoSearchResultProps) {
 
   if (
     searchTerm === '' ||
-    relayList.length > 0 ||
-    customLists.length > 0 ||
+    relayList.some((country) => country.visible) ||
+    customLists.some((list) => list.visible) ||
     props.specialLocationsLength > 0
   ) {
     return null;
@@ -393,7 +393,7 @@ function NoSearchResult(props: NoSearchResultProps) {
       <StyledNoResultText>
         {formatHtml(
           sprintf(messages.gettext('No result for <b>%(searchTerm)s</b>.'), {
-            searchTerm: searchTerm,
+            searchTerm,
           }),
         )}
       </StyledNoResultText>
diff --git a/gui/src/renderer/components/select-location/custom-list-helpers.ts b/gui/src/renderer/components/select-location/custom-list-helpers.ts
index 39c8aa0d60..163acb11b9 100644
--- a/gui/src/renderer/components/select-location/custom-list-helpers.ts
+++ b/gui/src/renderer/components/select-location/custom-list-helpers.ts
@@ -28,11 +28,16 @@ export function useCustomListsRelayList(
   // Populate all custom lists with the real location trees for the list locations.
   return useMemo(
     () =>
-      customLists
-        .map((list) =>
-          prepareCustomList(list, relayList, selectedLocation, disabledLocation, expandedLocations),
-        )
-        .filter((list) => searchMatch(searchTerm, list.label)),
+      customLists.map((list) =>
+        prepareCustomList(
+          list,
+          relayList,
+          searchTerm,
+          selectedLocation,
+          disabledLocation,
+          expandedLocations,
+        ),
+      ),
     [customLists, relayList, selectedLocation, disabledLocation, expandedLocations],
   );
 }
@@ -41,6 +46,7 @@ export function useCustomListsRelayList(
 function prepareCustomList(
   list: ICustomList,
   fullRelayList: GeographicalRelayList,
+  searchTerm: string,
   selectedLocation?: RelayLocation,
   disabledLocation?: { location: RelayLocation; reason: DisabledReason },
   expandedLocations?: Array<RelayLocation>,
@@ -59,6 +65,7 @@ function prepareCustomList(
     disabledReason,
     expanded: isExpanded(location, expandedLocations),
     selected: isSelected(location, selectedLocation),
+    visible: searchMatch(searchTerm, list.name),
     locations,
   };
 }
@@ -121,6 +128,7 @@ function updateCountry(
     location,
     expanded: isExpanded(location, expandedLocations),
     selected: false,
+    visible: true,
     cities: country.cities.map((city) =>
       updateCity(city, customList, locationCounter, expandedLocations),
     ),
@@ -147,6 +155,7 @@ function updateCity(
     location,
     expanded: isExpanded(location, expandedLocations),
     selected: false,
+    visible: true,
     relays: city.relays.map((relay) => updateRelay(relay, customList)),
   };
 }
@@ -158,6 +167,7 @@ function updateRelay(relay: RelaySpecification, customList: string): RelaySpecif
     ...relay,
     location: { ...relay.location, customList },
     selected: false,
+    visible: true,
   };
 }
 
diff --git a/gui/src/renderer/components/select-location/select-location-types.ts b/gui/src/renderer/components/select-location/select-location-types.ts
index b96fee5f06..29d897e0e5 100644
--- a/gui/src/renderer/components/select-location/select-location-types.ts
+++ b/gui/src/renderer/components/select-location/select-location-types.ts
@@ -28,6 +28,10 @@ export enum SpecialLocationIcon {
   geoLocation = 'icon-nearest',
 }
 
+export interface LocationVisibility {
+  visible: boolean;
+}
+
 interface CommonLocationSpecification {
   label: string;
   selected: boolean;
@@ -48,10 +52,25 @@ type GeographicalLocationSpecification =
 
 export type LocationSpecification = GeographicalLocationSpecification | CustomListSpecification;
 
-interface CommonNormalLocationSpecification extends CommonLocationSpecification {
+export interface RelayLocationCountryWithVisibility
+  extends IRelayLocationCountryRedux,
+    LocationVisibility {
+  cities: Array<RelayLocationCityWithVisibility>;
+}
+
+export interface RelayLocationCityWithVisibility
+  extends IRelayLocationCityRedux,
+    LocationVisibility {
+  relays: Array<RelayLocationRelayWithVisibility>;
+}
+
+export type RelayLocationRelayWithVisibility = IRelayLocationRelayRedux & LocationVisibility;
+
+interface CommonNormalLocationSpecification
+  extends CommonLocationSpecification,
+    LocationVisibility {
   location: RelayLocation;
   disabled: boolean;
-  selected: boolean;
   active: boolean;
 }
 
diff --git a/gui/src/renderer/lib/filter-locations.ts b/gui/src/renderer/lib/filter-locations.ts
index a4e2c12a03..63ca56cbf2 100644
--- a/gui/src/renderer/lib/filter-locations.ts
+++ b/gui/src/renderer/lib/filter-locations.ts
@@ -1,6 +1,11 @@
 import { Ownership, RelayEndpointType, RelayLocation } from '../../shared/daemon-rpc-types';
 import { relayLocations } from '../../shared/gettext';
-import { SpecialLocation } from '../components/select-location/select-location-types';
+import {
+  RelayLocationCityWithVisibility,
+  RelayLocationCountryWithVisibility,
+  RelayLocationRelayWithVisibility,
+  SpecialLocation,
+} from '../components/select-location/select-location-types';
 import {
   IRelayLocationCityRedux,
   IRelayLocationCountryRedux,
@@ -90,35 +95,44 @@ function filterLocationsImpl(
 export function searchForLocations(
   countries: Array<IRelayLocationCountryRedux>,
   searchTerm: string,
-): Array<IRelayLocationCountryRedux> {
-  if (searchTerm === '') {
-    return countries;
-  }
-
-  return countries.reduce((countries, country) => {
-    const matchingCities = searchCities(country.cities, searchTerm);
-    const expanded = matchingCities.length > 0;
+): Array<RelayLocationCountryWithVisibility> {
+  return countries.map((country) => {
     const match =
+      searchTerm === '' ||
       searchMatch(searchTerm, country.code) ||
       searchMatch(searchTerm, relayLocations.gettext(country.name));
-    const resultingCities = match ? country.cities : matchingCities;
-    return expanded || match ? [...countries, { ...country, cities: resultingCities }] : countries;
-  }, [] as Array<IRelayLocationCountryRedux>);
+    const cities = searchCities(country.cities, searchTerm, match);
+    const expanded = cities.some((city) => city.visible);
+    return { ...country, cities: cities, visible: expanded || match };
+  });
 }
 
 function searchCities(
   cities: Array<IRelayLocationCityRedux>,
   searchTerm: string,
-): Array<IRelayLocationCityRedux> {
-  return cities.reduce((cities, city) => {
-    const matchingRelays = city.relays.filter((relay) => searchMatch(searchTerm, relay.hostname));
-    const expanded = matchingRelays.length > 0;
+  countryMatch: boolean,
+): Array<RelayLocationCityWithVisibility> {
+  return cities.map((city) => {
     const match =
+      searchTerm === '' ||
+      countryMatch ||
       searchMatch(searchTerm, city.code) ||
       searchMatch(searchTerm, relayLocations.gettext(city.name));
-    const resultingRelays = match ? city.relays : matchingRelays;
-    return expanded || match ? [...cities, { ...city, relays: resultingRelays }] : cities;
-  }, [] as Array<IRelayLocationCityRedux>);
+    const relays = searchRelays(city.relays, searchTerm, match);
+    const expanded = match || relays.some((relay) => relay.visible);
+    return { ...city, relays: relays, visible: expanded };
+  });
+}
+
+function searchRelays(
+  relays: Array<IRelayLocationRelayRedux>,
+  searchTerm: string,
+  cityMatch: boolean,
+): Array<RelayLocationRelayWithVisibility> {
+  return relays.map((relay) => ({
+    ...relay,
+    visible: searchTerm === '' || cityMatch || searchMatch(searchTerm, relay.hostname),
+  }));
 }
 
 export function getLocationsExpandedBySearch(