Idiomatic React with heavy DOM manipulation (MathJax)

Question

I am using MathJax in a React application. MathJax brings a lot of complexity: it has its own system for managing concurrency, and makes changes to the DOM that React doesn't know about. This leads to a lot of DOM micromanaging which would generally be considered anti-patterns in React, and I'm wondering if my code can be made better.

In the code below, MJX is a component which takes a TeX string as input and feeds it into MathJax. RenderGroup is a convenience component which keeps track of when all its MJX descendants have finished typesetting.

/// <reference types="mathjax" />
import * as React from "react";

/* Promise that resolves once MathJax is loaded and ready to go */
export const MathJaxReady = new Promise<typeof MathJax>((resolve, reject) => {
  const script = $("#js-async-mathjax");
  if (!script) return;

  if (window.hasOwnProperty("MathJax")) {
    MathJax.Hub.Register.StartupHook("End", resolve);
  } else {
    script.addEventListener("load", () => MathJax.Hub.Register.StartupHook("End", resolve));
  }
});

interface Props extends React.HTMLAttributes<HTMLSpanElement> {
  display?: boolean;
}

export class MJX extends React.Component<Props, {}> {
  private resolveReady: () => void;
  domElement: HTMLSpanElement;
  jax: MathJax.ElementJax;

  // Promise that resolves after initial typeset
  ready: Promise<void>;

  static defaultProps = {
    display: false
  }

  constructor(props: Props) {
    super(props);

    this.ready = new Promise((resolve, reject) => this.resolveReady = resolve);

    this.Typeset = this.Typeset.bind(this);
  }

  async componentDidMount() {
    await MathJaxReady;

    this.Typeset()
    .then(() => this.jax = MathJax.Hub.getAllJax(this.domElement)[0])
    .then(this.resolveReady);
  }

  shouldComponentUpdate(nextProps, nextState) {
    /* original span has been eaten by MathJax, manage updates ourselves */
    const text = this.props.children instanceof Array ? this.props.children.join("") : this.props.children,
          nextText = nextProps.children instanceof Array ? nextProps.children.join("") : nextProps.children;

    // rerender?
    if (this.jax && text !== nextText) {
      this.jax.Text(nextProps.children);
    }

    // classes changed?
    if (this.props.className !== nextProps.className) {
      const classes = this.props.className ? this.props.className.split(" ") : [],
            newClasses = nextProps.className ? nextProps.className.split(" ") : [];

      const add = newClasses.filter(_ => !classes.includes(_)),
            remove = classes.filter(_ => !newClasses.includes(_));

      for (const _ of remove)
        this.domElement.classList.remove(_);
      for (const _ of add)
        this.domElement.classList.add(_);
    }

    // style attribute changed?
    if (JSON.stringify(this.props.style) !== JSON.stringify(nextProps.style)) {
      Object.keys(this.props.style || {})
      .filter(_ => !(nextProps.style || {}).hasOwnProperty(_))
      .forEach(_ => this.props.style[_] = null);
      Object.assign(this.domElement.style, nextProps.style);
    }

    return false;
  }

  Typeset(): Promise<void> {
    return new Promise((resolve, reject) => {
      MathJax.Hub.Queue(["Typeset", MathJax.Hub, this.domElement]);
      MathJax.Hub.Queue(resolve);
    });
  }

  render() {
    const {children, display, ...attrs} = this.props;

    const [open, close] = display ? ["\\[", "\\]"] : ["\\(", "\\)"];

    return (
      <span {...attrs} ref={node => this.domElement = node}>{open + children + close}</span>
    );
  }
}

// wait for a whole bunch of things to be rendered
export class RenderGroup extends React.Component {
  private promises: Promise<void>[];

  ready: Promise<void>;

  componentDidMount() {
    this.ready = Promise.all(this.promises).then(() => {});
  }

  render() {
    this.promises = [];

    return recursiveMap(this.props.children, node => {
      if (typeof node.type === "function" && node.type.prototype instanceof MJX) {
        const originalRef = node.ref;
        return React.cloneElement(node, {
          ref: (ref: MJX) => {
            if (!ref) return;
            this.promises.push(ref.ready);
            if (typeof originalRef === "function") {
              originalRef(ref);
            } else if (originalRef && typeof originalRef === "object") {
              originalRef.current = ref;
            }
          }
        });
      }

      return node;
    });
  }
}

// recursive React.Children.map
export function recursiveMap(
  children: React.ReactNode,
  fn: (child: React.ReactElement<any>) => React.ReactElement<any>
) {
  return React.Children.map(children, (child) => {
    if (!React.isValidElement<any>(child)) {
      return child;
    }

    if ("children" in child.props) {
      child = React.cloneElement(child, {
        children: recursiveMap(child.props.children, fn)
      });
    }

    return fn(child);
  });
}

Here is an example which is close to real code. We use MathJax to create some <input>s inside 2D vectors. In my case this will integrate with a graphical display which is also interactive, so the values of the entries will be stored in the state of a parent component, and Example can both receive values from the parent and set those values. Since the <input>s do not exist until MathJax has finished typesetting, we have to manage them manually.

interface Props {
  setParentValue: (i: number, value: number) => void;
  values: number[];
}

class Example extends React.PureComponent<Props> {
  private div: HTMLDivElement;
  private inputs: HTMLInputElement[];
  private rg: RenderGroup;

  componentDidMount() {
    this.rg.ready.then(() => {
      this.inputs = this.div.querySelectorAll("input");
      for (let i = 0; i < this.inputs.length; ++i) {
        this.inputs[i].addEventListener("change", e => this.props.setParentValue(i, e.target.value));
      }
    });
  }

  shouldComponentUpdate(nextProps) {
    if (this.inputs) {
      for (let i = 0; i < nextProps.values.length; ++i) {
        if (this.props.values[i] !== nextProps.values[i])
          this.inputs[i].value = nextProps.values[i];
      }
    }
    return false;
  }

  render() {
    // render only runs once, using initial values
    return (
      <div ref={ref => this.div = ref}>
        <RenderGroup ref={ref => this.rg = ref}>
          <MJX>{String.raw`
            \begin{bmatrix}
              \FormInput[4][matrix-entry][${this.props.values[0]}]{input1}\\
              \FormInput[4][matrix-entry][${this.props.values[1]}]{input2}
            \end{bmatrix}
          `}</MJX>

          <MJX>+</MJX>

          <MJX>{String.raw`
            \begin{bmatrix}
              \FormInput[4][matrix-entry][${this.props.values[2]}]{input3}\\
              \FormInput[4][matrix-entry][${this.props.values[3]}]{input4}
            \end{bmatrix}
          `}</MJX>

          <MJX>=</MJX>

          <MJX>{String.raw`
            \begin{bmatrix}
              ${this.props.values[0]+this.props.values[2]}\\
              ${this.props.values[1]+this.props.values[3]}
            \end{bmatrix}
          `}</MJX>
        </RenderGroup>
      </div>
    );
  }
}

Here are my questions.

1. RenderGroup is brittle. For example I don't understand why I need to check if (!ref); but if I omit that line, then ref will (for reasons I don't understand) become null on subsequent updates and cause an error. Intercepting the ref to grab the ready Promise also seems sketchy.

2. I'm slowly trying to migrate my class components to Hooks; while this isn't strictly necessary, according to the React team it should be possible. The problem is that function components don't have instances, so I don't see how to expose .ready to parent components like Example. I see there is a useImperativeHandle hook for this scenario, but that seems to depend on ultimately having a ref to an HTML component. I guess in the case of MJX I could put a ref on the <span>, but this wouldn't work for RenderGroup.

3. Managing the inputs imperatively is painful and error-prone. Is there any way to get some of React's declarative goodness back?

4. Bonus: I haven't been able to figure out how to type recursiveMap properly; TypeScript gets angry at the fn(child) line. It would also be good to replace the any's with generics.

Answer 1

I haven't personally used MathJax, but in my experience, an "idiomatic React" way to handle the resolveReady stuff would probably be to pass a callback down through context that lets the children notify the parent when they're loading or ready. Example (with hooks!):

const LoadingContext = createContext(() => () => {});
const LoadingProvider = memo(LoadingContext.Provider);

function RenderGroup({ children }) {
  const [areChildrenReady, setAreChildrenReady] = useState(false);

  const nextChildIdRef = useRef(0);
  const unfinishedChildrenRef = useRef(new Set());
  const startLoading = useCallback(() => {
    const childId = nextChildIdRef.current++;
    unfinishedChildrenRef.current.add(childId);
    setAreChildrenReady(!!unfinishedChildrenRef.current.size);
    const finishLoading = () => {
      unfinishedChildrenRef.current.delete(childId);
      setAreChildrenReady(!!unfinishedChildrenRef.current.size);
    };
    return finishLoading;
  }, []);

  useEffect(() => {
    if (areChildrenReady) {
      // do whatever
    }
  }, [areChildrenReady]);

  return (
    <LoadingProvider value={startLoading}>
      {children}
    </LoadingProvider>
  );
}

function ChildComponent() {
  const startLoading = useContext(LoadingContext);
  useEffect(() => {
    const finishLoading = startLoading();
    MathJaxReady
      .then(anotherPromise)
      .then(finishLoading);
  }, [startLoading]);
  return (
    // elements
  );
}