Fix code block highlighting not working reliably with many code blocks

src/components/views/messages/TextualBody.tsx

@@ -52,8 +52,6 @@ export default class TextualBody extends React.Component<IBodyProps, IState> {
     private tooltips = new ReactRootManager();
     private reactRoots = new ReactRootManager();
 
-    private ref = createRef<HTMLDivElement>();
-
     public static contextType = RoomContext;
     declare public context: React.ContextType<typeof RoomContext>;
 
@@ -86,7 +84,7 @@ export default class TextualBody extends React.Component<IBodyProps, IState> {
 
         if (this.props.mxEvent.getContent().format === "org.matrix.custom.html") {
             // Handle expansion and add buttons
-            const pres = this.ref.current?.getElementsByTagName("pre");
+            const pres = [...content.getElementsByTagName("pre")];
             if (pres && pres.length > 0) {
                 for (let i = 0; i < pres.length; i++) {
                     // If there already is a div wrapping the codeblock we want to skip this.
@@ -115,13 +113,14 @@ export default class TextualBody extends React.Component<IBodyProps, IState> {
         root.className = "mx_EventTile_pre_container";
 
         // Insert containing div in place of <pre> block
-        pre.parentNode?.replaceChild(root, pre);
+        pre.replaceWith(root);
 
         this.reactRoots.render(
             <StrictMode>
                 <CodeBlock onHeightChanged={this.props.onHeightChanged}>{pre}</CodeBlock>
             </StrictMode>,
             root,
+            pre,
         );
     }
 
@@ -196,10 +195,9 @@ export default class TextualBody extends React.Component<IBodyProps, IState> {
                     </StrictMode>
                 );
 
-                this.reactRoots.render(spoiler, spoilerContainer);
-
-                node.parentNode?.replaceChild(spoilerContainer, node);
+                this.reactRoots.render(spoiler, spoilerContainer, node);
 
+                node.replaceWith(spoilerContainer);
                 node = spoilerContainer;
             }
 
@@ -479,12 +477,7 @@ export default class TextualBody extends React.Component<IBodyProps, IState> {
 
         if (isEmote) {
             return (
-                <div
-                    className="mx_MEmoteBody mx_EventTile_content"
-                    onClick={this.onBodyLinkClick}
-                    dir="auto"
-                    ref={this.ref}
-                >
+                <div className="mx_MEmoteBody mx_EventTile_content" onClick={this.onBodyLinkClick} dir="auto">
                     *&nbsp;
                     <span className="mx_MEmoteBody_sender" onClick={this.onEmoteSenderClick}>
                         {mxEvent.sender ? mxEvent.sender.name : mxEvent.getSender()}
@@ -497,22 +490,22 @@ export default class TextualBody extends React.Component<IBodyProps, IState> {
         }
         if (isNotice) {
             return (
-                <div className="mx_MNoticeBody mx_EventTile_content" onClick={this.onBodyLinkClick} ref={this.ref}>
+                <div className="mx_MNoticeBody mx_EventTile_content" onClick={this.onBodyLinkClick}>
                     {body}
                     {widgets}
                 </div>
             );
         }
         if (isCaption) {
             return (
-                <div className="mx_MTextBody mx_EventTile_caption" onClick={this.onBodyLinkClick} ref={this.ref}>
+                <div className="mx_MTextBody mx_EventTile_caption" onClick={this.onBodyLinkClick}>
                     {body}
                     {widgets}
                 </div>
             );
         }
         return (
-            <div className="mx_MTextBody mx_EventTile_content" onClick={this.onBodyLinkClick} ref={this.ref}>
+            <div className="mx_MTextBody mx_EventTile_content" onClick={this.onBodyLinkClick}>
                 {body}
                 {widgets}
             </div>

src/utils/react.tsx

@@ -15,23 +15,38 @@ import { createRoot, Root } from "react-dom/client";
 export class ReactRootManager {
     private roots: Root[] = [];
     private rootElements: Element[] = [];
+    private revertElements: Array<null | Element> = [];
 
     public get elements(): Element[] {
         return this.rootElements;
     }
 
-    public render(children: ReactNode, element: Element): void {
-        const root = createRoot(element);
+    /**
+     * Render a React component into a new root based on the given root element
+     * @param children the React component to render
+     * @param rootElement the root element to render the component into
+     * @param revertElement the element to replace the root element with when unmounting
+     */
+    public render(children: ReactNode, rootElement: Element, revertElement?: Element): void {
+        const root = createRoot(rootElement);
         this.roots.push(root);
-        this.rootElements.push(element);
+        this.rootElements.push(rootElement);
+        this.revertElements.push(revertElement ?? null);
         root.render(children);
     }
 
+    /**
+     * Unmount all roots and revert the elements they were rendered into
+     */
     public unmount(): void {
         while (this.roots.length) {
             const root = this.roots.pop()!;
-            this.rootElements.pop();
+            const rootElement = this.rootElements.pop();
+            const revertElement = this.revertElements.pop();
             root.unmount();
+            if (revertElement) {
+                rootElement?.replaceWith(revertElement);
+            }
         }
     }
 }