Fix issues with custom list search

desktop/packages/mullvad-vpn/src/renderer/components/views/select-location/SelectLocationViewContext.tsx

@@ -1,9 +1,11 @@
 import React from 'react';
 
 import {
-  useCustomListLocations,
-  useFilteredCountryLocations,
+  useFilterCountryLocations,
+  useMapCustomListsToLocations,
+  useMapReduxCountriesToCountryLocations,
   useSearchCountryLocations,
+  useSearchCustomListLocations,
   useSelectedLocation,
 } from '../../../features/locations/hooks';
 import { LocationType } from '../../../features/locations/types';
@@ -17,9 +19,8 @@ type SelectLocationViewContextProps = Omit<SelectLocationViewProviderProps, 'chi
   setLocationType: (locationType: LocationType) => void;
   searchTerm: string;
   setSearchTerm: (value: string) => void;
-  filteredLocations: ReturnType<typeof useFilteredCountryLocations>;
-  searchedLocations: ReturnType<typeof useSearchCountryLocations>;
-  customListLocations: ReturnType<typeof useCustomListLocations>;
+  countryLocations: ReturnType<typeof useSearchCountryLocations>;
+  customListLocations: ReturnType<typeof useSearchCustomListLocations>;
 };
 
 const SelectLocationViewContext = React.createContext<SelectLocationViewContextProps | undefined>(
@@ -44,16 +45,23 @@ export function SelectLocationViewProvider({ children }: SelectLocationViewProvi
   const [searchTerm, setSearchTerm] = React.useState('');
   const relaySettings = useNormalRelaySettings();
   const selectedLocation = useSelectedLocation(locationTypeSelector);
-  const filteredLocations = useFilteredCountryLocations(locationTypeSelector);
-  const searchedLocations = useSearchCountryLocations(filteredLocations, searchTerm);
 
-  const activeSearch = searchTerm.length > 0;
+  const filteredCountries = useFilterCountryLocations(locationTypeSelector);
+  const filteredCountryLocations = useMapReduxCountriesToCountryLocations(
+    locationTypeSelector,
+    filteredCountries,
+  );
+  const searchedCountryLocations = useSearchCountryLocations(filteredCountryLocations, searchTerm);
 
-  const customListLocations = useCustomListLocations({
-    locations: activeSearch ? searchedLocations : filteredLocations,
+  const filteredCustomListLocations = useMapCustomListsToLocations(
+    filteredCountryLocations,
+    searchTerm,
     selectedLocation,
+  );
+  const searchedCustomListLocations = useSearchCustomListLocations(
+    filteredCustomListLocations,
     searchTerm,
-  });
+  );
 
   const locationType = React.useMemo(() => {
     const allowEntryLocations = relaySettings?.wireguard.useMultihop;
@@ -70,16 +78,14 @@ export function SelectLocationViewProvider({ children }: SelectLocationViewProvi
       setLocationType: setSelectLocationView,
       searchTerm,
       setSearchTerm,
-      filteredLocations,
-      searchedLocations,
-      customListLocations,
+      countryLocations: searchedCountryLocations,
+      customListLocations: searchedCustomListLocations,
     }),
     [
-      customListLocations,
-      filteredLocations,
+      searchedCustomListLocations,
+      searchedCountryLocations,
       locationType,
       searchTerm,
-      searchedLocations,
       setSearchTerm,
       setSelectLocationView,
     ],

desktop/packages/mullvad-vpn/src/renderer/components/views/select-location/components/country-locations/CountryLocations.tsx

@@ -9,7 +9,7 @@ import { CountryLocation } from '../country-location';
 import { useRelayCount } from './hooks';
 
 export function CountryLocations() {
-  const { searchedLocations } = useSelectLocationViewContext();
+  const { countryLocations } = useSelectLocationViewContext();
   const { visibleRelays, totalRelays } = useRelayCount();
 
   const showFilterText = visibleRelays !== totalRelays;
@@ -41,7 +41,7 @@ export function CountryLocations() {
         )}
       </SectionTitle>
       <FlexColumn>
-        {searchedLocations.map((location) => {
+        {countryLocations.map((location) => {
           const { key } = getLocationListItemMapProps(location, undefined);
           return <CountryLocation key={key} location={location} />;
         })}

desktop/packages/mullvad-vpn/src/renderer/components/views/select-location/components/country-locations/hooks/use-relay-count.ts

@@ -2,10 +2,10 @@ import { useRelayLocations } from '../../../../../../features/locations/hooks';
 import { useSelectLocationViewContext } from '../../../SelectLocationViewContext';
 
 export function useRelayCount() {
-  const { searchedLocations } = useSelectLocationViewContext();
+  const { countryLocations } = useSelectLocationViewContext();
   const { relayLocations } = useRelayLocations();
 
-  const visibleRelays = searchedLocations.reduce(
+  const visibleRelays = countryLocations.reduce(
     (countryAcc, country) =>
       countryAcc + country.cities.reduce((cityAcc, city) => cityAcc + city.relays.length, 0),
     0,

desktop/packages/mullvad-vpn/src/renderer/components/views/select-location/components/custom-list-location/CustomListLocation.tsx

@@ -35,6 +35,11 @@ export function CustomListLocationImpl({
     }
   }, [customList.locations.length, setExpanded]);
 
+  // If custom list state is updated from outside, update state accordingly
+  useEffect(() => {
+    setExpanded(customList.expanded);
+  }, [customList.expanded]);
+
   const handleClick = useCallback(() => {
     void handleSelect(customList);
   }, [customList, handleSelect]);

desktop/packages/mullvad-vpn/src/renderer/components/views/select-location/components/location-lists/hooks/use-has-visible-locations.ts

@@ -1,9 +1,9 @@
 import { useSelectLocationViewContext } from '../../../SelectLocationViewContext';
 
 export function useHasSearchedLocations() {
-  const { searchedLocations } = useSelectLocationViewContext();
+  const { countryLocations } = useSelectLocationViewContext();
 
-  const hasSearchedLocations = searchedLocations.length > 0;
+  const hasSearchedLocations = countryLocations.length > 0;
 
   return hasSearchedLocations;
 }

desktop/packages/mullvad-vpn/src/renderer/features/locations/hooks/index.ts

@@ -11,6 +11,7 @@ export * from './use-relay-locations';
 export * from './use-selected-location';
 export * from './use-active-filters';
 export * from './use-search-country-locations';
-export * from './use-filtered-country-locations';
+export * from './use-filter-country-locations';
 export * from './use-map-redux-countries-to-country-locations';
-export * from './use-custom-list-locations';
+export * from './use-map-custom-lists-to-locations';
+export * from './use-search-custom-list-locations';

desktop/packages/mullvad-vpn/src/renderer/features/locations/hooks/use-filtered-country-locations.ts → desktop/packages/mullvad-vpn/src/renderer/features/locations/hooks/use-filter-country-locations.ts

@@ -6,9 +6,8 @@ import { useIpVersion } from '../../tunnel/hooks';
 import { type LocationType } from '../types';
 import { filterLocations } from '../utils';
 import { useOwnership, useProviders } from '.';
-import { useMapReduxCountriesToCountryLocations } from './use-map-redux-countries-to-country-locations';
 
-export function useFilteredCountryLocations(locationType: LocationType) {
+export function useFilterCountryLocations(locationType: LocationType) {
   const locations = useSelector((state) => state.settings.relayLocations);
   const { activeOwnership } = useOwnership();
   const { providers } = useProviders();
@@ -18,7 +17,7 @@ export function useFilteredCountryLocations(locationType: LocationType) {
   const { multihop } = useMultihop();
   const { ipVersion } = useIpVersion();
 
-  const filteredRelayLocations = filterLocations({
+  return filterLocations({
     locations,
     ownership: activeOwnership,
     providers,
@@ -29,6 +28,4 @@ export function useFilteredCountryLocations(locationType: LocationType) {
     obfuscation,
     ipVersion,
   });
-
-  return useMapReduxCountriesToCountryLocations(filteredRelayLocations, locationType);
 }

desktop/packages/mullvad-vpn/src/renderer/features/locations/hooks/use-custom-list-locations.ts → desktop/packages/mullvad-vpn/src/renderer/features/locations/hooks/use-map-custom-lists-to-locations.ts

@@ -1,5 +1,3 @@
-import React from 'react';
-
 import type { RelayLocation } from '../../../../shared/daemon-rpc-types';
 import { useCustomLists } from '../../custom-lists/hooks';
 import {
@@ -15,21 +13,16 @@ import {
   searchMatchesLocation,
 } from '../utils';
 
-export function useCustomListLocations({
-  locations,
-  selectedLocation,
-  searchTerm,
-}: {
-  locations: CountryLocation[];
-  selectedLocation?: RelayLocation;
-  searchTerm: string;
-}): CustomListLocation[] {
+export function useMapCustomListsToLocations(
+  countryLocations: CountryLocation[],
+  searchTerm: string,
+  selectedLocation?: RelayLocation,
+): CustomListLocation[] {
   const { customLists } = useCustomLists();
 
-  const locationMap = React.useMemo(() => createLocationMap(locations), [locations]);
-
   const customListLocations: CustomListLocation[] = customLists.map((customList) => {
     const customListMatchesSearch = searchMatchesLocation(customList.name, searchTerm);
+    const locationMap = createLocationMap(countryLocations);
 
     // Get all ids of locations that are in the custom list
     const customListLocationIds = customList.locations.flatMap((location) => {
@@ -43,6 +36,7 @@ export function useCustomListLocations({
       return location.country;
     });
 
+    // Pick the locations from the map that are in the custom list, and add custom list details to them
     const customListGeographicalLocations: GeographicalLocation[] = [];
     for (const id of customListLocationIds) {
       const location = locationMap.get(id);

desktop/packages/mullvad-vpn/src/renderer/features/locations/hooks/use-map-redux-countries-to-country-locations.ts

@@ -8,8 +8,8 @@ import { useDisabledLocation } from './use-disabled-location';
 import { useSelectedLocation } from './use-selected-location';
 
 export function useMapReduxCountriesToCountryLocations(
-  relayList: Array<IRelayLocationCountryRedux>,
   locationType: LocationType,
+  relayList: Array<IRelayLocationCountryRedux>,
 ): CountryLocation[] {
   const locale = useSelector((state) => state.userInterface.locale);
   const selectedLocation = useSelectedLocation(locationType);

desktop/packages/mullvad-vpn/src/renderer/features/locations/hooks/use-search-country-locations.ts

@@ -1,49 +1,19 @@
 import React from 'react';
 
-import type { CityLocation, CountryLocation, RelayLocation } from '../types';
-import { searchMatchesLocation } from '../utils';
+import type { CountryLocation } from '../types';
+import { searchCountryAndCities } from '../utils';
 
 export function useSearchCountryLocations(
-  locations: CountryLocation[],
+  countryLocations: CountryLocation[],
   searchTerm: string,
 ): CountryLocation[] {
-  return React.useMemo(() => formatCountriesResult(locations, searchTerm), [locations, searchTerm]);
-}
-
-export function formatCountriesResult(countries: CountryLocation[], searchTerm: string) {
-  if (!searchTerm) {
-    return countries;
-  }
-  return countries
-    .map((country) => {
-      const citiesResult = formatCitiesResult(country, searchTerm);
-      if (citiesResult.length > 0) {
-        return { ...country, expanded: true, cities: citiesResult };
-      }
-      if (searchMatchesLocation(country.label, searchTerm)) {
-        return country;
-      }
-      return undefined;
-    })
-    .filter((country) => country !== undefined);
-}
-
-export function formatCitiesResult(country: CountryLocation, searchTerm: string): CityLocation[] {
-  return country.cities
-    .map((city) => {
-      const relaysResult = formatRelaysResult(city, searchTerm);
-      if (relaysResult.length > 0) {
-        return { ...city, expanded: true, relays: relaysResult };
-      }
-      if (searchMatchesLocation(city.label, searchTerm)) {
-        return city;
-      }
-
-      return undefined;
-    })
-    .filter((city) => city !== undefined);
-}
+  return React.useMemo(() => {
+    if (!searchTerm) {
+      return countryLocations;
+    }
 
-export function formatRelaysResult(city: CityLocation, searchTerm: string): RelayLocation[] {
-  return city.relays.filter((relay) => searchMatchesLocation(relay.label, searchTerm));
+    return countryLocations
+      .map((country) => searchCountryAndCities(country, searchTerm))
+      .filter((country) => country !== undefined);
+  }, [countryLocations, searchTerm]);
 }

desktop/packages/mullvad-vpn/src/renderer/features/locations/hooks/use-search-custom-list-locations.ts

@@ -0,0 +1,19 @@
+import React from 'react';
+
+import type { CustomListLocation } from '../types';
+import { searchCustomListAndLocations } from '../utils';
+
+export function useSearchCustomListLocations(
+  customListLocations: CustomListLocation[],
+  searchTerm: string,
+): CustomListLocation[] {
+  return React.useMemo(() => {
+    if (!searchTerm) {
+      return customListLocations;
+    }
+
+    return customListLocations
+      .map((customList) => searchCustomListAndLocations(customList, searchTerm))
+      .filter((customList) => customList !== undefined);
+  }, [customListLocations, searchTerm]);
+}

desktop/packages/mullvad-vpn/src/renderer/features/locations/utils/index.ts

@@ -21,3 +21,6 @@ export * from './is-location-selected';
 export * from './search-matches-location';
 export * from './create-location-map';
 export * from './create-location-label';
+export * from './search-city-and-relays';
+export * from './search-country-and-cities';
+export * from './search-custom-list-and-locations';

desktop/packages/mullvad-vpn/src/renderer/features/locations/utils/search-city-and-relays.ts

@@ -0,0 +1,19 @@
+import type { CityLocation } from '../types';
+import { searchMatchesLocation } from './search-matches-location';
+
+export function searchCityAndRelays(
+  city: CityLocation,
+  searchTerm: string,
+): CityLocation | undefined {
+  const relaysResult = city.relays.filter((relay) =>
+    searchMatchesLocation(relay.label, searchTerm),
+  );
+  if (relaysResult.length > 0) {
+    return { ...city, expanded: true, relays: relaysResult };
+  }
+  if (searchMatchesLocation(city.label, searchTerm)) {
+    return city;
+  }
+
+  return undefined;
+}

desktop/packages/mullvad-vpn/src/renderer/features/locations/utils/search-country-and-cities.ts

@@ -0,0 +1,19 @@
+import type { CountryLocation } from '../types';
+import { searchCityAndRelays } from './search-city-and-relays';
+import { searchMatchesLocation } from './search-matches-location';
+
+export function searchCountryAndCities(
+  country: CountryLocation,
+  searchTerm: string,
+): CountryLocation | undefined {
+  const citiesResult = country.cities
+    .map((city) => searchCityAndRelays(city, searchTerm))
+    .filter((city) => city !== undefined);
+  if (citiesResult.length > 0) {
+    return { ...country, expanded: true, cities: citiesResult };
+  }
+  if (searchMatchesLocation(country.label, searchTerm)) {
+    return country;
+  }
+  return undefined;
+}

desktop/packages/mullvad-vpn/src/renderer/features/locations/utils/search-custom-list-and-locations.ts

@@ -0,0 +1,32 @@
+import type { CustomListLocation } from '../types';
+import { searchCityAndRelays } from './search-city-and-relays';
+import { searchCountryAndCities } from './search-country-and-cities';
+import { searchMatchesLocation } from './search-matches-location';
+
+export function searchCustomListAndLocations(
+  customList: CustomListLocation,
+  searchTerm: string,
+): CustomListLocation | undefined {
+  const locationsResult = customList.locations.filter((location) => {
+    if (location.type === 'relay') {
+      return searchMatchesLocation(location.label, searchTerm);
+    } else if (location.type === 'city') {
+      const cityResult = searchCityAndRelays(location, searchTerm);
+      return cityResult !== undefined;
+    } else if (location.type === 'country') {
+      const countryResult = searchCountryAndCities(location, searchTerm);
+      return countryResult !== undefined;
+    }
+    return false;
+  });
+
+  if (locationsResult.length > 0) {
+    return { ...customList, expanded: true, locations: locationsResult };
+  }
+
+  const customListMatchesSearch = searchMatchesLocation(customList.label, searchTerm);
+  if (customListMatchesSearch) {
+    return customList;
+  }
+  return undefined;
+}