Repository: storybookjs/storybook
File: code/addons/a11y/src/components/A11yContext.tsx (TypeScript)

Code before:
    const unsubscribe = experimental_getStatusStore('storybook/component-test').onAllStatusChange(
      (statuses, previousStatuses) => {
        const current = statuses[storyId]?.[STATUS_TYPE_ID_COMPONENT_TEST];
        const previous = previousStatuses[storyId]?.[STATUS_TYPE_ID_COMPONENT_TEST];
        if (current?.value === 'status-value:error' && previous?.value !== 'status-value:error') {
          setStatus('component-test-error');
        }
      }
    );
    return unsubscribe;
  }, [storyId]);

  const handleToggleHighlight = useCallback(() => {
    setHighlighted((prevHighlighted) => {
      const next = !prevHighlighted;
      setUiState((prev) => ({ ...prev, highlighted: next }));
      return next;
    });
  }, [setUiState]);
  const setTab = useCallback(
    (type: RuleType) => {
      setTabState(type);
      setUiState((prev) => ({ ...prev, tab: type }));
    },
    [setUiState]
  );

  const [selectedItems, setSelectedItems] = useState<Map<string, string>>(() => {
    const initialValue = new Map();
    // Check if the a11ySelection param is a valid format before parsing it
    // It should look like `violation.aria-hidden-body.1`
    if (a11ySelection && /^[a-z]+.[a-z-]+.[0-9]+$/.test(a11ySelection)) {
      const [type, id] = a11ySelection.split('.');
      initialValue.set(`${type}.${id}`, a11ySelection);
    }
    return initialValue;
  });

  // All items are expanded if something is selected from each result for the current tab
  const allExpanded = useMemo(() => {
    const currentResults = results?.[tab];
    return currentResults?.every((result) => selectedItems.has(`${tab}.${result.id}`)) ?? false;
  }, [results, selectedItems, tab]);

  const toggleOpen = useCallback(
    (event: React.SyntheticEvent<Element>, type: RuleType, item: EnhancedResult) => {
      event.stopPropagation();
      const key = `${type}.${item.id}`;
      setSelectedItems((prev) => new Map(prev.delete(key) ? prev : prev.set(key, `${key}.1`)));
    },
    []

Review comment (refactor):
I think we can simplify this a lot, by not having the local state anymore, but just the upstream uiState.

Code after:

  const handleToggleHighlight = useCallback(() => {
    setState((prev) => ({ ...prev, ui: { ...prev.ui, highlighted: !prev.ui.highlighted } }));
  }, [setState]);

  const [selectedItems, setSelectedItems] = useState<Map<string, string>>(() => {
    const initialValue = new Map();
    // Check if the a11ySelection param is a valid format before parsing it
    // It should look like `violation.aria-hidden-body.1`
    if (a11ySelection && /^[a-z]+.[a-z-]+.[0-9]+$/.test(a11ySelection)) {
      const [type, id] = a11ySelection.split('.');
      initialValue.set(`${type}.${id}`, a11ySelection);
    }
    return initialValue;
  });

  // All items are expanded if something is selected from each result for the current tab
  const allExpanded = useMemo(() => {
    const currentResults = results?.[ui.tab];
    return currentResults?.every((result) => selectedItems.has(`${ui.tab}.${result.id}`)) ?? false;
  }, [results, selectedItems, ui.tab]);

  const toggleOpen = useCallback(
    (event: React.SyntheticEvent<Element>, type: RuleType, item: EnhancedResult) => {
      event.stopPropagation();
      const key = `${type}.${item.id}`;
      setSelectedItems((prev) => new Map(prev.delete(key) ? prev : prev.set(key, `${key}.1`)));
    },
    []
  );

  const handleCollapseAll = useCallback(() => {
    setSelectedItems(new Map());
  }, []);

  const handleExpandAll = useCallback(() => {
    setSelectedItems(
      (prev) =>
        new Map(
          results?.[ui.tab]?.map((result) => {
            const key = `${ui.tab}.${result.id}`;
            return [key, prev.get(key) ?? `${key}.1`];
          }) ?? []
        )
    );
  }, [results, ui.tab]);

  const handleSelectionChange = useCallback((key: string) => {
    const [type, id] = key.split('.');
    setSelectedItems((prev) => new Map(prev.set(`${type}.${id}`, key)));
  }, []);