From ed9aa2854925fbce6970cfaa9c74a4794c3c19fc Mon Sep 17 00:00:00 2001
From: Aditya Sharat <adityasharat@fb.com>
Date: Thu, 20 Aug 2020 09:56:27 -0700
Subject: [PATCH] Fixes content pooling for Litho on RenderCore

Summary:
This diff fixes content pooling for RenderCore.

The pools should be keyed with `RenderUnit#getRenderContentType()` but were being keyed with `RenderUnit#getClass()`.

Reviewed By: mihaelao

Differential Revision: D23215571

fbshipit-source-id: f751f1c19473018abf10d03d2b22844345028f5c
---
 .../com/facebook/litho/LithoRenderUnit.java   |  5 ++
 .../com/facebook/litho/MountStateTest.java    | 81 +++++++++++++++++++
 .../facebook/rendercore/MountItemsPool.java   |  4 +-
 3 files changed, 88 insertions(+), 2 deletions(-)

diff --git a/litho-core/src/main/java/com/facebook/litho/LithoRenderUnit.java b/litho-core/src/main/java/com/facebook/litho/LithoRenderUnit.java
index 4b2f9b45201..e94af9bf765 100644
--- a/litho-core/src/main/java/com/facebook/litho/LithoRenderUnit.java
+++ b/litho-core/src/main/java/com/facebook/litho/LithoRenderUnit.java
@@ -55,6 +55,11 @@ public long getId() {
     return output.getId();
   }
 
+  @Override
+  public Object getRenderContentType() {
+    return output.getComponent().getClass();
+  }
+
   private boolean hasDefaultViewAttributeFlags() {
     return mDefaultViewAttributeFlags != -1;
   }
diff --git a/litho-it/src/test/java/com/facebook/litho/MountStateTest.java b/litho-it/src/test/java/com/facebook/litho/MountStateTest.java
index 8a8d46c154e..e0e87450701 100644
--- a/litho-it/src/test/java/com/facebook/litho/MountStateTest.java
+++ b/litho-it/src/test/java/com/facebook/litho/MountStateTest.java
@@ -21,11 +21,14 @@
 import static org.assertj.core.api.Java6Assertions.assertThat;
 
 import android.graphics.Color;
+import android.view.View;
 import com.facebook.litho.config.ComponentsConfiguration;
 import com.facebook.litho.testing.LithoViewRule;
 import com.facebook.litho.testing.testrunner.LithoTestRunner;
 import com.facebook.litho.widget.DynamicPropsComponentTester;
+import com.facebook.litho.widget.Progress;
 import com.facebook.litho.widget.SolidColor;
+import com.facebook.litho.widget.TextInput;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -123,4 +126,82 @@ public void onSetRootWithNoOutputsWithRenderCore_shouldSuccessfullyCompleteMount
     ComponentsConfiguration.useExtensionsWithMountDelegate = useExtensions;
     ComponentsConfiguration.useIncrementalMountExtension = incrementalMountExtension;
   }
+
+  @Test
+  public void onSetRootWithSimilarComponent_MountContentShouldUsePools() {
+    final boolean delegateToRenderCoreMount = ComponentsConfiguration.delegateToRenderCoreMount;
+    final boolean useExtensions = ComponentsConfiguration.useExtensionsWithMountDelegate;
+    final boolean incrementalMountExtension = ComponentsConfiguration.useIncrementalMountExtension;
+
+    ComponentsConfiguration.delegateToRenderCoreMount = true;
+    ComponentsConfiguration.useExtensionsWithMountDelegate = true;
+    ComponentsConfiguration.useIncrementalMountExtension = true;
+
+    final Component root =
+        Column.create(mContext)
+            .child(TextInput.create(mContext).widthDip(100).heightDip(100))
+            .build();
+
+    mLithoViewRule
+        .setRoot(root)
+        .attachToWindow()
+        .setSizeSpecs(makeSizeSpec(1000, EXACTLY), makeSizeSpec(1000, EXACTLY))
+        .measure()
+        .layout();
+
+    View view = mLithoViewRule.getLithoView().getChildAt(0);
+
+    final Component newRoot =
+        Row.create(mContext)
+            .child(TextInput.create(mContext).initialText("testing").widthDip(120).heightDip(120))
+            .build();
+
+    mLithoViewRule
+        .setRoot(newRoot)
+        .setSizeSpecs(makeSizeSpec(1000, EXACTLY), makeSizeSpec(1000, EXACTLY));
+
+    View newView = mLithoViewRule.getLithoView().getChildAt(0);
+
+    assertThat(newView).isSameAs(view);
+
+    ComponentsConfiguration.delegateToRenderCoreMount = delegateToRenderCoreMount;
+    ComponentsConfiguration.useExtensionsWithMountDelegate = useExtensions;
+    ComponentsConfiguration.useIncrementalMountExtension = incrementalMountExtension;
+  }
+
+  @Test
+  public void onSetRootWithDifferentComponent_MountContentPoolsShouldNoCollide() {
+    final boolean delegateToRenderCoreMount = ComponentsConfiguration.delegateToRenderCoreMount;
+    final boolean useExtensions = ComponentsConfiguration.useExtensionsWithMountDelegate;
+    final boolean incrementalMountExtension = ComponentsConfiguration.useIncrementalMountExtension;
+
+    ComponentsConfiguration.delegateToRenderCoreMount = true;
+    ComponentsConfiguration.useExtensionsWithMountDelegate = true;
+    ComponentsConfiguration.useIncrementalMountExtension = true;
+
+    final Component root =
+        Column.create(mContext)
+            .child(TextInput.create(mContext).widthDip(100).heightDip(100))
+            .build();
+
+    mLithoViewRule
+        .setRoot(root)
+        .attachToWindow()
+        .setSizeSpecs(makeSizeSpec(1000, EXACTLY), makeSizeSpec(1000, EXACTLY))
+        .measure()
+        .layout();
+
+    final Component newRoot =
+        Column.create(mContext)
+            .child(Progress.create(mContext).widthDip(100).heightDip(100))
+            .build();
+
+    mLithoViewRule
+        .setRoot(newRoot)
+        .setSizeSpecs(makeSizeSpec(1000, EXACTLY), makeSizeSpec(1000, EXACTLY));
+
+    ComponentsConfiguration.delegateToRenderCoreMount = delegateToRenderCoreMount;
+    ComponentsConfiguration.useExtensionsWithMountDelegate = useExtensions;
+    ComponentsConfiguration.useIncrementalMountExtension = incrementalMountExtension;
+  }
 }
diff --git a/litho-rendercore/src/main/java/com/facebook/rendercore/MountItemsPool.java b/litho-rendercore/src/main/java/com/facebook/rendercore/MountItemsPool.java
index a5982a8ca44..f3ec873d709 100644
--- a/litho-rendercore/src/main/java/com/facebook/rendercore/MountItemsPool.java
+++ b/litho-rendercore/src/main/java/com/facebook/rendercore/MountItemsPool.java
@@ -75,7 +75,7 @@ static Object acquireMountContent(Context context, RenderUnit renderUnit) {
   }
 
   static void release(Context context, RenderUnit renderUnit, Object mountContent) {
-    final Pools.SimplePool pool = getMountContentPool(context, renderUnit.getClass());
+    final Pools.SimplePool pool = getMountContentPool(context, renderUnit.getRenderContentType());
     if (pool != null) {
       pool.release(mountContent);
     }
@@ -98,7 +98,7 @@ static void release(Context context, RenderUnit renderUnit, Object mountContent)
     Pools.SimplePool pool = poolsMap.get(lifecycle);
     if (pool == null) {
       pool = new Pools.SimplePool(DEFAULT_POOL_SIZE);
-      poolsMap.put(lifecycle.getClass(), pool);
+      poolsMap.put(lifecycle, pool);
     }
 
     return pool;