Fix unmountComponentAtNode compat rax (#6424)

* fix: fixup unmountComponentAtNode compat

* fix(rax-compat): compat for rax driver-dom on nullish render container

* fix(rax-compat): compat for rax driver-dom on non-specified render container

.changeset/neat-apples-divide.md

@@ -0,0 +1,5 @@
+---
+'rax-compat': patch
+---
+
+fix: fix unmountComponentAtNode compat

packages/rax-compat/src/container-root-map.ts

@@ -0,0 +1,8 @@
+import type { Root } from 'react-dom/client';
+
+/**
+ * A WeakMap that keeps track of container mouted root.
+ * @key The container element passed to `render` function, which call `createRoot` and `root.render` internally.
+ * @value The root instance returned by `createRoot` function.
+ */
+export const containerRootMap = new WeakMap<Element | DocumentFragment, Root>();

packages/rax-compat/src/render.ts

@@ -3,7 +3,24 @@ import type { ReactNode } from 'react';
 import type { RootOptions } from 'react-dom/client';
 import { createRoot } from 'react-dom/client';
 import { isFunction } from './type.js';
+import { containerRootMap } from './container-root-map.js';
 
+export default function render(
+  element: RaxElement,
+  container: Element | DocumentFragment | null,
+  options?: RenderOption,
+): RaxElement;
+export default function render(
+  element: RaxElement,
+  container: Element | DocumentFragment | null,
+  callback?: Function,
+): RaxElement;
+export default function render(
+  element: RaxElement,
+  container: Element | DocumentFragment | null,
+  options?: RenderOption,
+  callback?: Function,
+): RaxElement;
 /**
  * Compat render for rax export.
  * https://github.com/alibaba/rax/blob/master/packages/rax/src/render.js#L5
@@ -15,17 +32,33 @@ import { isFunction } from './type.js';
  */
 export default function render(
   element: RaxElement,
-  container: Element | DocumentFragment | null,
-  options?: RenderOption,
+  // Doesnot provide overload like `render(<Container />, callback))`,
+  // should use `render(<Container />, null, callback))` instead,
+  // as rendering without specified container is NOT RECOMMENDED.
+  container?: Element | DocumentFragment | null,
+  optionsOrCallback?: RenderOption | Function,
   callback?: Function,
 ): RaxElement {
-  if (isFunction(options)) {
-    callback = options;
-    options = null;
+  if (isFunction(optionsOrCallback)) {
+    callback = optionsOrCallback;
+    optionsOrCallback = {} as RenderOption;
   }
 
-  const root = createRoot(container, options as RootOptions);
+  /**
+   * Compat for rax driver-dom behavior, which use body as container for non-specified container.
+   * ref: https://github.com/alibaba/rax/blob/13d80a491f18c74034aa9b05c36ac922ff8c3357/packages/rax/src/vdom/instance.js#L44
+   * ref: https://github.com/alibaba/rax/blob/13d80a491f18c74034aa9b05c36ac922ff8c3357/packages/driver-dom/src/index.js#L75
+   */
+  if (!container) {
+    container = document.body;
+  }
+
+  const root = createRoot(container, optionsOrCallback as RootOptions);
   root.render(element as ReactNode);
+
+  // Save container and root relation.
+  containerRootMap.set(container, root);
+
   if (isFunction(callback)) {
     callback.call(element);
   }

packages/rax-compat/src/unmount-component-at-node.ts

@@ -1,6 +1,25 @@
 import type { Container } from 'react-dom';
+import { containerRootMap } from './container-root-map.js';
 
+/**
+ * Unmount a component rendered to a DOM node.
+ * To unmount a component strictly, its state and event handlers should be cleaned up appropriately,
+ * Which means we should let `React` do the unmounting work.
+ *
+ * In React 18, this method will be replaced by [root.unmount](https://react.dev/reference/react-dom/client/createRoot#root-unmount),
+ * We preserve this method for Rax API compatibility.
+ *
+ * As we cannot approach the root instance from container element directly,
+ * we use a WeakMap to keep track of container mouted root,
+ * which saves the container-root relation in `render`(and `createRoot()` under the hood) execution.
+
+ * @param container The container element which to be unmounted.
+ */
 export default function unmountComponentAtNode(container: Container) {
-  // @TODO: unmount
-  container.textContent = '';
+  const relatedRoot = containerRootMap.get(container);
+  // Ensure the root exists, then unmount root and remove it from containerRootMap.
+  if (relatedRoot) {
+    relatedRoot.unmount();
+    containerRootMap.delete(container);
+  }
 }

packages/rax-compat/tsconfig.json

@@ -11,6 +11,7 @@
     "sourceMap": true,
     "allowJs": true,
     "noUnusedLocals": true,
+    "strictNullChecks": true,
     "noImplicitReturns": true,
     "noImplicitThis": true,
     "noImplicitAny": true,