Track contentDOM in descriptors

src/components/NodeWrapper.tsx

@@ -14,25 +14,42 @@ import {
   NodeViewDescriptorsContext,
 } from "../contexts/NodeViewPositionsContext.js";
 
+export function findChildDesc(
+  pos: number,
+  posToDesc: Map<number, NodeViewDescriptor>
+) {
+  const positions = Array.from(posToDesc.keys()).sort((a, b) => b - a);
+
+  let parentPos = null;
+  for (const nodePos of positions) {
+    if (nodePos < pos) break;
+
+    parentPos = nodePos;
+  }
+
+  return parentPos === null ? null : posToDesc.get(parentPos);
+}
+
 type NodeWrapperProps = {
   children: ReactNode;
   pos: number;
 };
 export function NodeWrapper({ children, pos }: NodeWrapperProps) {
-  const { posToDesc: posToDOM, domToDesc: domToPos } = useContext(
-    NodeViewDescriptorsContext
-  );
+  const { posToDesc, domToDesc } = useContext(NodeViewDescriptorsContext);
   const ref = useRef<Element | null>(null);
 
   useLayoutEffect(() => {
     if (!ref.current) return;
+
+    const childDesc = findChildDesc(pos, posToDesc);
+
     const desc: NodeViewDescriptor = {
       pos,
       dom: ref.current,
-      contentDOM: null,
+      contentDOM: childDesc?.dom.parentNode ?? null,
     };
-    posToDOM.set(pos, desc);
-    domToPos.set(ref.current, desc);
+    posToDesc.set(pos, desc);
+    domToDesc.set(ref.current, desc);
   });
 
   const child = Children.only(children);

src/components/TextNodeWrapper.tsx

@@ -18,16 +18,16 @@ export class TextNodeWrapper extends Component<TextNodeWrapperProps> {
     const textNode = findDOMNode(this);
     if (!textNode) return;
 
-    const { posToDesc: posToDOM, domToDesc: domToPos } = this
+    const { posToDesc, domToDesc } = this
       .context as NodeViewDescriptorsContextValue;
 
     const desc: NodeViewDescriptor = {
       pos: this.props.pos,
       dom: textNode,
       contentDOM: null,
     };
-    posToDOM.set(this.props.pos, desc);
-    domToPos.set(textNode, desc);
+    posToDesc.set(this.props.pos, desc);
+    domToDesc.set(textNode, desc);
   }
 
   componentDidUpdate(): void {
@@ -36,16 +36,16 @@ export class TextNodeWrapper extends Component<TextNodeWrapperProps> {
     const textNode = findDOMNode(this);
     if (!textNode) return;
 
-    const { posToDesc: posToDOM, domToDesc: domToPos } = this
+    const { posToDesc, domToDesc } = this
       .context as NodeViewDescriptorsContextValue;
 
     const desc: NodeViewDescriptor = {
       pos: this.props.pos,
       dom: textNode,
       contentDOM: null,
     };
-    posToDOM.set(this.props.pos, desc);
-    domToPos.set(textNode, desc);
+    posToDesc.set(this.props.pos, desc);
+    domToDesc.set(textNode, desc);
   }
 
   render() {