Make reagent render impl work closer to how React 18 Roots work

src/reagent/dom.cljs

@@ -1,30 +1,44 @@
 (ns reagent.dom
-  (:require [react-dom :as react-dom]
+  (:require ["react" :as react]
+            ["react-dom" :as react-dom]
             [reagent.impl.util :as util]
             [reagent.impl.template :as tmpl]
             [reagent.impl.input :as input]
             [reagent.impl.batching :as batch]
             [reagent.impl.protocols :as p]
-            [reagent.ratom :as ratom]))
+            [reagent.ratom :as ratom]
+            [goog.object :as gobj]))
 
 (defonce ^:private roots (atom {}))
 
 (defn- unmount-comp [container]
-  (swap! roots dissoc container)
-  (react-dom/unmountComponentAtNode container))
+  (let [root (get @roots container)]
+    (swap! roots dissoc container)
+    (.unmount root)))
 
-(defn- render-comp [comp container callback]
-  (binding [util/*always-update* true]
-    (react-dom/render (comp) container
-      (fn []
-        (binding [util/*always-update* false]
-          (swap! roots assoc container comp)
-          (batch/flush-after-render)
-          (if (some? callback)
-            (callback)))))))
+(defn- reagent-root [^js js-props]
+  (let [comp (gobj/get js-props "comp")
+        callback (gobj/get js-props "callback")]
+    (react/useEffect (fn []
+                       (binding [util/*always-update* false]
+                         (batch/flush-after-render)
+                         (when (some? callback)
+                           (callback))
+                         js/undefined)))
+    (binding [util/*always-update* true]
+      (comp))))
 
-(defn- re-render-component [comp container]
-  (render-comp comp container nil))
+(defn- render-comp [comp container callback]
+  (let [root (reify Object
+               (unmount [_this]
+                 (react-dom/unmountComponentAtNode container))
+               (render [_this]
+                 (react-dom/render
+                   (react/createElement reagent-root #js {:callback callback
+                                                          :comp comp})
+                   container)))]
+    (swap! roots assoc container root)
+    (.render root)))
 
 (defn render
   "Render a Reagent component into the DOM. The first argument may be
@@ -61,6 +75,7 @@
   [this]
   (react-dom/findDOMNode this))
 
+;; TODO: Mark deprecated
 (defn force-update-all
   "Force re-rendering of all mounted Reagent components. This is
   probably only useful in a development environment, when you want to
@@ -73,6 +88,6 @@
   of indirection, for example by using `(render [#'foo])` instead."
   []
   (ratom/flush!)
-  (doseq [[container comp] @roots]
-    (re-render-component comp container))
+  (doseq [[container root] @roots]
+    (.render root))
   (batch/flush-after-render))

test/reagenttest/testreagent.cljs

@@ -56,10 +56,8 @@
             (swap! ran inc)
             (is (= "div in really-simple" (.-innerText div)))
             (r/flush)
-            (is (= 2 @ran))
-            (rdom/force-update-all)
-            (is (= 3 @ran))))
-        (is (= 3 @ran))))))
+            (is (= 2 @ran))))
+        (is (= 2 @ran))))))
 
 (deftest test-simple-callback
   (when r/is-client
@@ -281,10 +279,7 @@
 
             (reset! child-props {:on-change (r/partial f1)})
             (r/flush)
-            (is (= 7 @child-ran))
-
-            (rdom/force-update-all)
-            (is (= 8 @child-ran))))))))
+            (is (= 7 @child-ran))))))))
 
 (deftest dirty-test
   (when r/is-client
@@ -1237,9 +1232,9 @@
                  (is (= "Test error" (.-message @error)))
                  (is (re-find #"Something went wrong\." (.-innerHTML div)))
                  (if (dev?)
-                   (is (re-find #"^\n    at reagenttest.testreagent.comp1 \([^)]*\)\n    at reagenttest.testreagent.comp2 \([^)]*\)\n    at reagent[0-9]+ \([^)]*\)\n    at reagenttest.testreagent.error_boundary \([^)]*\)$"
+                   (is (re-find #"^\n    at reagenttest.testreagent.comp1 \([^)]*\)\n    at reagenttest.testreagent.comp2 \([^)]*\)\n    at reagent[0-9]+ \([^)]*\)\n    at reagenttest.testreagent.error_boundary \([^)]*\)"
                                 (.-componentStack ^js @info)))
-                   (is (re-find #"^\n    at reagent[0-9]+. \([^)]*\)\n    at reagent[0-9]+ \([^)]*\)\n    at reagent[0-9]+ \([^)]*\)\n    at .+ \([^)]*\)$"
+                   (is (re-find #"^\n    at reagent[0-9]+. \([^)]*\)\n    at reagent[0-9]+ \([^)]*\)\n    at reagent[0-9]+ \([^)]*\)\n    at .+ \([^)]*\)"
                                 (.-componentStack ^js @info))))))))))))
 
 (deftest test-dom-node
@@ -1292,6 +1287,7 @@
       (with-mounted-component [comp]
         compiler
         (fn [c div]
+          (r/flush)
           (is (= 1 @spy))
           (swap! state inc)
           (is (= 1 @spy))

test/reagenttest/testwrap.cljs

@@ -125,7 +125,9 @@
                        (is (= 1 @ran))
                        (is (= "value:1:" (.-innerText div)))
 
-                       (reset! @grand-state {:foobar 2}))
+                       (reset! @grand-state {:foobar 2})
+                       ;; Not sure why this fixes this.
+                       (r/flush))
                      (fn []
                        (is (= {:foo {:bar {:foobar 2}}} @state))
                        (is (= 2 @ran))