Fully split CustomNodeView into a separate, memoized component (#12)

Plain custom node views had been broken since we implemented SSR.
When they were pulled into their own component, it messed up some of
the layout effect ordering and triggering, such that we no longer
attached them to the DOM.

To resolve this, the NodeView component now either renders ReactNodeView
or CustomNodeView, each of which manages its own full lifecycle, decorations,
etc.

Author: smoores-dev

.yarn/versions/1cfdbe48.yml

@@ -0,0 +1,2 @@
+releases:
+  "@handlewithcare/react-prosemirror": patch

src/components/CustomNodeView.tsx

@@ -5,42 +5,139 @@ import {
   NodeViewConstructor,
   NodeView as NodeViewT,
 } from "prosemirror-view";
-import React, { MutableRefObject, createElement, useContext } from "react";
+import React, {
+  MutableRefObject,
+  cloneElement,
+  createElement,
+  memo,
+  useContext,
+  useLayoutEffect,
+  useMemo,
+  useRef,
+} from "react";
 import { createPortal } from "react-dom";
 
+import { ChildDescriptorsContext } from "../contexts/ChildDescriptorsContext.js";
 import { EditorContext } from "../contexts/EditorContext.js";
 import { useClientOnly } from "../hooks/useClientOnly.js";
+import { useNodeViewDescriptor } from "../hooks/useNodeViewDescriptor.js";
 
-import { ChildNodeViews } from "./ChildNodeViews.js";
+import { ChildNodeViews, wrapInDeco } from "./ChildNodeViews.js";
 
 interface Props {
-  customNodeViewRootRef: MutableRefObject<HTMLDivElement | null>;
-  customNodeViewRef: MutableRefObject<NodeViewT | null>;
-  contentDomRef: MutableRefObject<HTMLElement | null>;
   customNodeView: NodeViewConstructor;
-  initialNode: MutableRefObject<Node>;
   node: Node;
   getPos: MutableRefObject<() => number>;
-  initialOuterDeco: MutableRefObject<readonly Decoration[]>;
-  initialInnerDeco: MutableRefObject<DecorationSource>;
   innerDeco: DecorationSource;
+  outerDeco: readonly Decoration[];
 }
 
-export function CustomNodeView({
-  contentDomRef,
-  customNodeViewRef,
-  customNodeViewRootRef,
+export const CustomNodeView = memo(function CustomNodeView({
   customNodeView,
-  initialNode,
   node,
   getPos,
-  initialOuterDeco,
-  initialInnerDeco,
   innerDeco,
+  outerDeco,
 }: Props) {
   const { view } = useContext(EditorContext);
+  const domRef = useRef<HTMLElement | null>(null);
+  const nodeDomRef = useRef<HTMLElement | null>(null);
+  const contentDomRef = useRef<HTMLElement | null>(null);
+  const getPosFunc = useRef(() => getPos.current()).current;
+
+  // this is ill-conceived; should revisit
+  const initialNode = useRef(node);
+  const initialOuterDeco = useRef(outerDeco);
+  const initialInnerDeco = useRef(innerDeco);
+
+  const customNodeViewRootRef = useRef<HTMLDivElement | null>(null);
+  const customNodeViewRef = useRef<NodeViewT | null>(null);
 
   const shouldRender = useClientOnly();
+
+  useLayoutEffect(() => {
+    if (
+      !customNodeViewRef.current ||
+      !customNodeViewRootRef.current ||
+      !shouldRender
+    )
+      return;
+
+    const { dom } = customNodeViewRef.current;
+    nodeDomRef.current = customNodeViewRootRef.current;
+    customNodeViewRootRef.current.appendChild(dom);
+    return () => {
+      customNodeViewRef.current?.destroy?.();
+    };
+  }, [customNodeViewRef, customNodeViewRootRef, nodeDomRef, shouldRender]);
+
+  useLayoutEffect(() => {
+    if (!customNodeView || !customNodeViewRef.current || !shouldRender) return;
+
+    const { destroy, update } = customNodeViewRef.current;
+
+    const updated =
+      update?.call(customNodeViewRef.current, node, outerDeco, innerDeco) ??
+      true;
+    if (updated) return;
+
+    destroy?.call(customNodeViewRef.current);
+
+    if (!customNodeViewRootRef.current) return;
+
+    initialNode.current = node;
+    initialOuterDeco.current = outerDeco;
+    initialInnerDeco.current = innerDeco;
+
+    customNodeViewRef.current = customNodeView(
+      initialNode.current,
+      // customNodeView will only be set if view is set, and we can only reach
+      // this line if customNodeView is set
+      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+      view!,
+      getPosFunc,
+      initialOuterDeco.current,
+      initialInnerDeco.current
+    );
+    const { dom } = customNodeViewRef.current;
+    nodeDomRef.current = customNodeViewRootRef.current;
+    customNodeViewRootRef.current.appendChild(dom);
+  }, [
+    customNodeView,
+    view,
+    innerDeco,
+    node,
+    outerDeco,
+    getPos,
+    customNodeViewRef,
+    customNodeViewRootRef,
+    initialNode,
+    initialOuterDeco,
+    initialInnerDeco,
+    nodeDomRef,
+    shouldRender,
+    getPosFunc,
+  ]);
+
+  const { childDescriptors, nodeViewDescRef } = useNodeViewDescriptor(
+    node,
+    () => getPos.current(),
+    domRef,
+    nodeDomRef,
+    innerDeco,
+    outerDeco,
+    undefined,
+    contentDomRef
+  );
+
+  const childContextValue = useMemo(
+    () => ({
+      parentRef: nodeViewDescRef,
+      siblingsRef: childDescriptors,
+    }),
+    [childDescriptors, nodeViewDescRef]
+  );
+
   if (!shouldRender) return null;
 
   if (!customNodeViewRef.current) {
@@ -57,7 +154,7 @@ export function CustomNodeView({
   }
   const { contentDOM } = customNodeViewRef.current;
   contentDomRef.current = contentDOM ?? null;
-  return createElement(
+  const element = createElement(
     node.isInline ? "span" : "div",
     {
       ref: customNodeViewRootRef,
@@ -74,4 +171,21 @@ export function CustomNodeView({
         contentDOM
       )
   );
-}
+
+  const decoratedElement = cloneElement(
+    outerDeco.reduce(wrapInDeco, element),
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    outerDeco.some((d) => (d as any).type.attrs.nodeName)
+      ? { ref: domRef }
+      : // If all of the node decorations were attr-only, then
+        // we've already passed the domRef to the NodeView component
+        // as a prop
+        undefined
+  );
+
+  return (
+    <ChildDescriptorsContext.Provider value={childContextValue}>
+      {decoratedElement}
+    </ChildDescriptorsContext.Provider>
+  );
+});