Make BlockingButton use a context instead of cloning children

2 files changed, 82 insertions, 79 deletions

diff --git a/gui/src/renderer/components/AppButton.tsx b/gui/src/renderer/components/AppButton.tsx
index 5b2c7fa62e..f35f7b2049 100644
--- a/gui/src/renderer/components/AppButton.tsx
+++ b/gui/src/renderer/components/AppButton.tsx
@@ -1,7 +1,8 @@
 import log from 'electron-log';
-import React, { useContext } from 'react';
+import React, { useCallback, useContext, useEffect, useMemo, useRef, useState } from 'react';
 import styled from 'styled-components';
 import { colors } from '../../config.json';
+import { useMounted } from '../lib/utilityHooks';
 import {
   StyledButton,
   StyledButtonContent,
@@ -10,9 +11,13 @@ import {
 } from './AppButtonStyles';
 import ImageView from './ImageView';
 
-const ButtonContext = React.createContext({
+interface IButtonContext {
+  textAdjustment: number;
+  textRef?: React.Ref<HTMLDivElement>;
+}
+
+const ButtonContext = React.createContext<IButtonContext>({
   textAdjustment: 0,
-  textRef: React.createRef<HTMLDivElement>(),
 });
 
 interface ILabelProps {
@@ -46,51 +51,19 @@ export interface IProps extends React.HTMLAttributes<HTMLButtonElement> {
   textOffset?: number;
 }
 
-interface IState {
-  textAdjustment: number;
-}
-
-class BaseButton extends React.Component<IProps, IState> {
-  public state: IState = {
-    textAdjustment: 0,
-  };
+const BaseButton = React.memo(function BaseButtonT(props: IProps) {
+  const { children, disabled, onClick, ...otherProps } = props;
 
-  private buttonRef = React.createRef<HTMLButtonElement>();
-  private textRef = React.createRef<HTMLDivElement>();
+  const blockingContext = useContext(BlockingContext);
+  const [textAdjustment, setTextAdjustment] = useState(0);
+  const buttonRef = useRef() as React.RefObject<HTMLButtonElement>;
+  const textRef = useRef() as React.RefObject<HTMLDivElement>;
 
-  public componentDidMount() {
-    this.updateTextAdjustment();
-  }
+  const contextValue = useMemo(() => ({ textAdjustment, textRef }), [textAdjustment, textRef]);
 
-  public componentDidUpdate() {
-    this.updateTextAdjustment();
-  }
-
-  public render() {
-    const { children, ...otherProps } = this.props;
-
-    return (
-      <ButtonContext.Provider
-        value={{
-          textAdjustment: this.state.textAdjustment,
-          textRef: this.textRef,
-        }}>
-        <StyledButton ref={this.buttonRef} {...otherProps}>
-          <StyledButtonContent>
-            {React.Children.map(children, (child) =>
-              typeof child === 'string' ? <Label>{child as string}</Label> : child,
-            )}
-          </StyledButtonContent>
-        </StyledButton>
-      </ButtonContext.Provider>
-    );
-  }
-
-  private updateTextAdjustment() {
-    const textOffset = this.props.textOffset ?? 0;
-
-    const buttonRect = this.buttonRef.current?.getBoundingClientRect();
-    const textRect = this.textRef.current?.getBoundingClientRect();
+  useEffect(() => {
+    const buttonRect = buttonRef.current?.getBoundingClientRect();
+    const textRect = textRef.current?.getBoundingClientRect();
 
     if (buttonRect && textRect) {
       const leftDiff = textRect.left - buttonRect.left;
@@ -99,55 +72,70 @@ class BaseButton extends React.Component<IProps, IState> {
       const trailingSpace = buttonRect.width - (leftDiff + textRect.width);
 
       // calculate text adjustment
+      const textOffset = props.textOffset ?? 0;
       const textAdjustment = leftDiff - trailingSpace - textOffset;
 
       // re-render the view with the new text adjustment if it changed
-      if (this.state.textAdjustment !== textAdjustment) {
-        this.setState({ textAdjustment });
-      }
+      setTextAdjustment(textAdjustment);
     }
-  }
-}
+  });
 
-interface IBlockingState {
-  isBlocked: boolean;
+  return (
+    <ButtonContext.Provider value={contextValue}>
+      <StyledButton
+        ref={buttonRef}
+        disabled={blockingContext.disabled || disabled}
+        onClick={blockingContext.onClick ?? onClick}
+        {...otherProps}>
+        <StyledButtonContent>
+          {React.Children.map(children, (child) =>
+            typeof child === 'string' ? <Label>{child as string}</Label> : child,
+          )}
+        </StyledButtonContent>
+      </StyledButton>
+    </ButtonContext.Provider>
+  );
+});
+
+interface IBlockingContext {
+  disabled?: boolean;
+  onClick?: () => Promise<void>;
 }
 
+const BlockingContext = React.createContext<IBlockingContext>({});
+
 interface IBlockingProps {
   children?: React.ReactNode;
   onClick: () => Promise<void>;
   disabled?: boolean;
 }
 
-export class BlockingButton extends React.Component<IBlockingProps, IBlockingState> {
-  public state = {
-    isBlocked: false,
-  };
+export function BlockingButton(props: IBlockingProps) {
+  const isMounted = useMounted();
+  const [isBlocked, setIsBlocked] = useState(false);
 
-  public render() {
-    return React.Children.map(this.props.children, (child) => {
-      if (React.isValidElement(child)) {
-        return React.cloneElement(child as React.ReactElement, {
-          ...child.props,
-          disabled: this.state.isBlocked || this.props.disabled,
-          onClick: this.onClick,
-        });
-      } else {
-        return child;
-      }
-    });
-  }
+  const onClick = useCallback(async () => {
+    setIsBlocked(true);
+    try {
+      await props.onClick();
+    } catch (error) {
+      log.error(`onClick() failed - ${error}`);
+    }
+
+    if (isMounted()) {
+      setIsBlocked(false);
+    }
+  }, [props.onClick]);
+
+  const contextValue = useMemo(
+    () => ({
+      disabled: isBlocked || props.disabled,
+      onClick,
+    }),
+    [isBlocked, props.disabled, onClick],
+  );
 
-  private onClick = () => {
-    this.setState({ isBlocked: true }, async () => {
-      try {
-        await this.props.onClick();
-      } catch (error) {
-        log.error(`onClick() failed - ${error}`);
-      }
-      this.setState({ isBlocked: false });
-    });
-  };
+  return <BlockingContext.Provider value={contextValue}>{props.children}</BlockingContext.Provider>;
 }
 
 export const RedButton = styled(BaseButton)({
diff --git a/gui/src/renderer/lib/utilityHooks.ts b/gui/src/renderer/lib/utilityHooks.ts
new file mode 100644
index 0000000000..cbb5575f9b
--- /dev/null
+++ b/gui/src/renderer/lib/utilityHooks.ts
@@ -0,0 +1,15 @@
+import { useCallback, useEffect, useRef } from 'react';
+
+export function useMounted() {
+  const mountedRef = useRef(false);
+  const isMounted = useCallback(() => mountedRef.current, []);
+
+  useEffect(() => {
+    mountedRef.current = true;
+    return () => {
+      mountedRef.current = false;
+    };
+  }, []);
+
+  return isMounted;
+}