Paint by county with ID instead of features

app/src/app/components/Map.tsx

@@ -8,6 +8,7 @@ import {MAP_OPTIONS} from '../constants/configuration';
 import {mapEvents} from '../utils/events/mapEvents';
 import {INTERACTIVE_LAYERS} from '../constants/layers';
 import {useMapStore} from '../store/mapStore';
+import { parentIdCache } from '../store/idCache';
 
 export const MapComponent: React.FC = () => {
   const map: MutableRefObject<Map | null> = useRef(null);
@@ -45,6 +46,7 @@ export const MapComponent: React.FC = () => {
       zoom: MAP_OPTIONS.zoom,
       maxZoom: MAP_OPTIONS.maxZoom,
     });
+
     fitMapToBounds();
     map.current.scrollZoom.setWheelZoomRate(1 / 300);
     map.current.scrollZoom.setZoomRate(1 / 300);
@@ -86,3 +88,4 @@ export const MapComponent: React.FC = () => {
     />
   );
 };
+

app/src/app/components/sidebar/PaintByCounty.tsx

@@ -8,12 +8,15 @@ export default function PaintByCounty() {
   const mapRef = useMapStore(state => state.getMapRef());
   const addVisibleLayerIds = useMapStore(state => state.addVisibleLayerIds);
   const setPaintFunction = useMapStore(state => state.setPaintFunction);
-  const [checked, setChecked] = useState(false);
+  const paintByCounty = useMapStore(state => state.mapOptions.paintByCounty)
+  const setMapOptions = useMapStore(state => state.setMapOptions)
 
-  useEffect(() => {
+  const handleToggle = () => {
     if (!mapRef) return;
-
-    if (checked) {
+    setMapOptions({
+      paintByCounty: !paintByCounty
+    })
+    if (!paintByCounty) {
       COUNTY_LAYER_IDS.forEach(layerId => {
         mapRef.setLayoutProperty(layerId, 'visibility', 'visible');
       });
@@ -22,16 +25,16 @@ export default function PaintByCounty() {
     } else {
       setPaintFunction(getFeaturesInBbox);
     }
-  }, [checked, mapRef, addVisibleLayerIds]);
+  }
 
   return (
     <Box>
       <Text as="label" size="2">
         <Flex gap="2">
           <Checkbox
-            checked={checked}
+            checked={paintByCounty}
             defaultChecked={false}
-            onClick={() => setChecked(prevIsChecked => !prevIsChecked)}
+            onClick={handleToggle}
           />
           Paint by County
         </Flex>

app/src/app/constants/layers.ts

@@ -23,6 +23,7 @@ export const INTERACTIVE_LAYERS = [BLOCK_HOVER_LAYER_ID, BLOCK_HOVER_LAYER_ID_CH
 export const LINE_LAYERS = [BLOCK_LAYER_ID, BLOCK_LAYER_ID_CHILD] as const
 
 export const PARENT_LAYERS = [BLOCK_LAYER_ID, BLOCK_HOVER_LAYER_ID]; 
+export const COUNTY_LAYERS = ['counties_fill', 'counties_boundary','counties_label']
 
 export const CHILD_LAYERS = [
   BLOCK_LAYER_ID_CHILD,

app/src/app/store/idCache.ts

@@ -0,0 +1,25 @@
+class IdCache {
+  cachedTileIndices: Set<string> = new Set()
+  parentIds: Set<string> = new Set()
+
+  hasCached(index: string){
+    return this.cachedTileIndices.has(index)
+  }
+
+  add(index: string, ids: string[]){
+    this.cachedTileIndices.add(index)
+    ids.forEach(id => this.parentIds.add(id))
+  }
+
+  clear(){
+    this.parentIds.clear()
+    this.cachedTileIndices.clear()
+  }
+
+  getFilteredIds(id: string){
+    const regex = new RegExp(`^${id}`);
+    return Array.from(this.parentIds).filter(f => regex.test(f));
+  }
+}
+
+export const parentIdCache = new IdCache()

app/src/app/store/mapRenderSubs.ts

@@ -11,10 +11,13 @@ import {
   BLOCK_LAYER_ID_HIGHLIGHT,
   getHighlightLayerSpecification,
   BLOCK_LAYER_ID_HIGHLIGHT_CHILD,
+  COUNTY_LAYERS,
 } from '../constants/layers';
 import {
   ColorZoneAssignmentsState,
   colorZoneAssignments,
+  getFeaturesInBbox,
+  getFeaturesIntersectingCounties,
   shallowCompareArray,
 } from '../utils/helpers';
 import {useMapStore as _useMapStore, MapStore} from '@store/mapStore';
@@ -222,16 +225,21 @@ export const getRenderSubscriptions = (useMapStore: typeof _useMapStore) => {
     activeTool => {
       const mapRef = useMapStore.getState().getMapRef();
       if (!mapRef) return;
+      const mapOptions = useMapStore.getState().mapOptions
+      const defaultPaintFunction = mapOptions.paintByCounty ? getFeaturesIntersectingCounties : getFeaturesInBbox
       let cursor;
       switch (activeTool) {
         case 'pan':
           cursor = '';
+          useMapStore.getState().setPaintFunction(defaultPaintFunction);
           break;
         case 'brush':
           cursor = 'pointer';
+          useMapStore.getState().setPaintFunction(defaultPaintFunction);
           break;
         case 'eraser':
           cursor = 'pointer';
+          useMapStore.getState().setPaintFunction(defaultPaintFunction);
           break;
         case 'shatter':
           cursor = 'crosshair';
@@ -303,6 +311,18 @@ export const getRenderSubscriptions = (useMapStore: typeof _useMapStore) => {
       } 
     }
   );
+
+  const filterCountiesSub = useMapStore.subscribe<[string|undefined, MapStore['getMapRef']]>(state => [state.mapOptions.currentStateFp, state.getMapRef],
+    ([stateFp, getMapRef]) => {
+      const mapRef = getMapRef()
+      if (!mapRef) return
+      const filterExpression = (stateFp ? ["==", "STATEFP", stateFp] : true) as any
+      COUNTY_LAYERS.forEach(layer => {
+        mapRef.setFilter(layer,  ["any", filterExpression])
+      })
+    }
+  )
+
   return [
     addLayerSubMapDocument,
     _shatterMapSideEffectRender,
@@ -311,5 +331,6 @@ export const getRenderSubscriptions = (useMapStore: typeof _useMapStore) => {
     _updateMapCursor,
     _applyFocusFeatureState,
     highlightUnassignedSub,
+    filterCountiesSub
   ];
 };

app/src/app/store/mapStore.ts

@@ -37,6 +37,7 @@ import {BLOCK_SOURCE_ID} from '../constants/layers';
 import {DistrictrMapOptions} from './types';
 import {onlyUnique} from '../utils/arrays';
 import {queryClient} from '../utils/api/queryClient';
+import { parentIdCache } from './idCache';
 
 const combineSetValues = (setRecord: Record<string, Set<unknown>>, keys?: string[]) => {
   const combinedSet = new Set<unknown>(); // Create a new set to hold combined values
@@ -384,6 +385,7 @@ export const useMapStore = create(
           if (currentMapDocument?.document_id === mapDocument.document_id) {
             return;
           }
+          parentIdCache.clear()
           setFreshMap(true);
           resetZoneAssignments();
           upsertUserMap({mapDocument});

app/src/app/store/types.ts

@@ -5,4 +5,6 @@ export type DistrictrMapOptions = {
   higlightUnassigned?: boolean;
   lockPaintedAreas: boolean | Array<NullableZone>;
   mode: 'default' | 'break';
+  paintByCounty?: boolean;
+  currentStateFp?: string
 };

app/src/app/utils/events/mapEvents.ts

@@ -2,7 +2,7 @@
  Port over from map events declared at: https://github.com/uchicago-dsi/districtr-components/blob/2e8f9e5657b9f0fd2419b6f3258efd74ae310f32/src/Districtr/Districtr.tsx#L230
  */
 'use client';
-import type {Map as MapLibreMap, MapLayerMouseEvent, MapLayerTouchEvent} from 'maplibre-gl';
+import type {Map as MapLibreMap, MapLayerMouseEvent, MapLayerTouchEvent, MapDataEvent, MapSourceDataEvent} from 'maplibre-gl';
 import {useMapStore} from '@/app/store/mapStore';
 import {
   BLOCK_HOVER_LAYER_ID,
@@ -11,6 +11,7 @@ import {
 } from '@/app/constants/layers';
 import {ResetMapSelectState} from '@utils/events/handlers';
 import {ActiveTool} from '@/app/constants/types';
+import { parentIdCache } from '@/app/store/idCache';
 
 /*
 MapEvent handling; these functions are called by the event listeners in the MapComponent
@@ -227,6 +228,38 @@ export const handleMapContextMenu = (
   });
 };
 
+export const handleIdCache = (
+  _e: MapLayerMouseEvent | MapLayerTouchEvent,
+  map: MapLibreMap | null
+) => {
+  const e = _e as unknown as MapSourceDataEvent
+  const {tiles_s3_path, parent_layer} = useMapStore.getState().mapDocument || {}
+
+  if (
+    !tiles_s3_path || 
+    !parent_layer || 
+    e.dataType !== 'source' || 
+    !("url" in e.source) ||
+    !e.source.url?.includes(tiles_s3_path)
+  ) return
+
+  const tileData = e.tile.latestFeatureIndex;
+  if (!tileData) return
+
+  const index = `${tileData.x}-${tileData.y}-${tileData.z}`
+  if (parentIdCache.hasCached(index)) return
+  const vtLayers = tileData.loadVTLayers()
+
+  const parentLayerData = vtLayers[parent_layer]
+  const numFeatures = parentLayerData.length
+  const featureDataArray = parentLayerData._values
+  const idArray = featureDataArray.slice(-numFeatures,)
+  parentIdCache.add(index, idArray)
+  useMapStore.getState().setMapOptions({
+    currentStateFp: idArray[0].replace('vtd:','').slice(0,2)
+  })
+}
+
 export const mapEvents = [
   {action: 'click', handler: handleMapClick},
   {action: 'mouseup', handler: handleMapMouseUp},
@@ -246,4 +279,5 @@ export const mapEvents = [
   {action: 'moveend', handler: handleMapMoveEnd},
   {action: 'zoomend', handler: handleMapZoomEnd},
   {action: 'contextmenu', handler: handleMapContextMenu},
+  {action: 'data', handler: handleIdCache}
 ];

app/src/app/utils/helpers.ts

@@ -15,6 +15,7 @@ import {
 } from '@/app/constants/layers';
 import {MapStore, useMapStore} from '../store/mapStore';
 import {NullableZone} from '../constants/types';
+import {parentIdCache} from '../store/idCache';
 
 /**
  * PaintEventHandler
@@ -133,16 +134,25 @@ export const getFeaturesIntersectingCounties = (
 
   if (!countyFeatures?.length) return;
   const fips = countyFeatures[0].properties.STATEFP + countyFeatures[0].properties.COUNTYFP;
+  const {mapDocument, shatterIds} = useMapStore.getState();
+  const filterPrefix = mapDocument?.parent_layer.includes("vtd") ? "vtd:" : ""
+  const cachedParentFeatures = parentIdCache.getFilteredIds(`${filterPrefix}${fips}`).map(id => ({
+    id,
+    source: BLOCK_SOURCE_ID,
+    sourceLayer: mapDocument?.parent_layer,
+  }));
+
+  const childFeatures = shatterIds.children.size
+    ? (Array.from(shatterIds.children).map(id => ({
+        id,
+        source: BLOCK_SOURCE_ID,
+        sourceLayer: mapDocument?.child_layer,
+      })) as any)
+    : [];
 
-  const features = map.queryRenderedFeatures(undefined, {
-    layers,
-  });
-
-  return filterFeatures(
-    features,
-    true,
-    [(feature) => Boolean(feature?.id && feature.id.toString().match(/\d{5}/)?.[0] === fips)]
-  );
+  return filterFeatures([...cachedParentFeatures, ...childFeatures], true, [
+    feature => Boolean(feature?.id && feature.id.toString().match(/\d{5}/)?.[0] === fips),
+  ]);
 };
 
 /**