From 47c250a779eddc985b0084e69e72e1ea837b5ea4 Mon Sep 17 00:00:00 2001
From: Tmk <i@tmk.im>
Date: Sun, 3 Apr 2022 22:53:37 +0800
Subject: [PATCH 1/2] fix: restore portal to dom

---
 src/Portal.tsx | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/src/Portal.tsx b/src/Portal.tsx
index 9aaf1c66..35c53a75 100644
--- a/src/Portal.tsx
+++ b/src/Portal.tsx
@@ -1,4 +1,4 @@
-import * as React from 'react';
+import type * as React from 'react';
 import { useRef, useEffect, forwardRef, useImperativeHandle } from 'react';
 import ReactDOM from 'react-dom';
 import canUseDom from './Dom/canUseDom';
@@ -14,6 +14,7 @@ export interface PortalProps {
 const Portal = forwardRef<PortalRef, PortalProps>((props, ref) => {
   const { didUpdate, getContainer, children } = props;
 
+  const parentRef = useRef<ParentNode>();
   const containerRef = useRef<HTMLElement>();
 
   // Ref return nothing, only for wrapper check exist
@@ -23,6 +24,7 @@ const Portal = forwardRef<PortalRef, PortalProps>((props, ref) => {
   const initRef = useRef(false);
   if (!initRef.current && canUseDom()) {
     containerRef.current = getContainer();
+    parentRef.current = containerRef.current.parentNode;
     initRef.current = true;
   }
 
@@ -32,6 +34,13 @@ const Portal = forwardRef<PortalRef, PortalProps>((props, ref) => {
   });
 
   useEffect(() => {
+    // Restore container to original place
+    if (
+      containerRef.current.parentNode === null &&
+      parentRef.current !== null
+    ) {
+      parentRef.current.appendChild(containerRef.current);
+    }
     return () => {
       // [Legacy] This should not be handle by Portal but parent PortalWrapper instead.
       // Since some component use `Portal` directly, we have to keep the logic here.

From 0e70c46806e49a540d23e09b90d5d4fc835c9071 Mon Sep 17 00:00:00 2001
From: Tmk <i@tmk.im>
Date: Wed, 6 Apr 2022 18:25:02 +0800
Subject: [PATCH 2/2] test: add StrictMode tests

---
 tests/Portal.test.tsx | 28 +++++++++++++++++++++++++++-
 1 file changed, 27 insertions(+), 1 deletion(-)

diff --git a/tests/Portal.test.tsx b/tests/Portal.test.tsx
index de964aa6..a5aa73a2 100644
--- a/tests/Portal.test.tsx
+++ b/tests/Portal.test.tsx
@@ -1,4 +1,4 @@
-import React from 'react';
+import React, { StrictMode, useEffect } from 'react';
 import { render } from '@testing-library/react';
 import { act } from 'react-dom/test-utils';
 import PortalWrapper, { getOpenCount } from '../src/PortalWrapper';
@@ -189,4 +189,30 @@ describe('Portal', () => {
     );
     expect(document.body.querySelector('.light')).toBeTruthy();
   });
+
+  it('should restore to original place in StrictMode', () => {
+    const parentContainer = document.createElement('div');
+    const domContainer = document.createElement('div');
+    parentContainer.appendChild(domContainer);
+    let mountCount = 0;
+    let unmountCount = 0;
+
+    const Demo = () => {
+      useEffect(() => {
+        mountCount += 1;
+        return () => {
+          unmountCount += 1;
+        };
+      }, []);
+
+      return <Portal getContainer={() => domContainer}>Contents</Portal>;
+    };
+
+    render(<Demo />, { wrapper: StrictMode });
+
+    expect(mountCount).toBe(2);
+    expect(unmountCount).toBe(1);
+    // portal should be attached to parent node
+    expect(parentContainer.textContent).toBe('Contents');
+  });
 });