Passes in inter stage props to range API

Summary:
This diff fixes a bug where the range lifecycle methods do not set inter stage props because the interstage props container is never passed to the API. So the AP was setting the container to null. This would throw an NPE when invoking a range lifecycle with inter stage props.

This fix
* The interstage props are held in the RangeTuple.
* Passed to the callbacks as method args.
* The AP and Component APIs have been updated to support this use case.

Differential Revision: D43246229

fbshipit-source-id: a840c1608a3fd06b64e360467bc811cbe2a8d22a

litho-core/src/main/java/com/facebook/litho/LayoutState.java

@@ -926,8 +926,19 @@ private static void collectResults(
       }
 
       for (WorkingRangeContainer.Registration registration : registrations) {
-        layoutState.mWorkingRangeContainer.registerWorkingRange(
-            registration.mName, registration.mWorkingRange, registration.mScopedComponentInfo);
+        if (component instanceof SpecGeneratedComponent) {
+          layoutState.mWorkingRangeContainer.registerWorkingRange(
+              registration.mName,
+              registration.mWorkingRange,
+              registration.mScopedComponentInfo,
+              (InterStagePropsContainer) result.getLayoutData());
+        } else {
+          layoutState.mWorkingRangeContainer.registerWorkingRange(
+              registration.mName,
+              registration.mWorkingRange,
+              registration.mScopedComponentInfo,
+              null);
+        }
       }
     }
 

litho-core/src/main/java/com/facebook/litho/SpecGeneratedComponent.java

@@ -276,11 +276,17 @@ StateContainer createInitialStateContainer(ComponentContext c) {
     return null;
   }
 
-  protected void dispatchOnEnteredRange(ComponentContext c, String name) {
+  protected void dispatchOnEnteredRange(
+      final ComponentContext c,
+      final String name,
+      final @Nullable InterStagePropsContainer interStageProps) {
     // Do nothing by default
   }
 
-  protected void dispatchOnExitedRange(ComponentContext c, String name) {
+  protected void dispatchOnExitedRange(
+      final ComponentContext c,
+      final String name,
+      final @Nullable InterStagePropsContainer interStageProps) {
     // Do nothing by default
   }
 

litho-core/src/main/java/com/facebook/litho/WorkingRangeContainer.java

@@ -41,15 +41,17 @@ class WorkingRangeContainer {
   void registerWorkingRange(
       final String name,
       final WorkingRange workingRange,
-      final ScopedComponentInfo scopedComponentInfo) {
+      final ScopedComponentInfo scopedComponentInfo,
+      final @Nullable InterStagePropsContainer interStageProps) {
     if (mWorkingRanges == null) {
       mWorkingRanges = new LinkedHashMap<>();
     }
 
     final String key = name + "_" + workingRange.hashCode();
     final RangeTuple rangeTuple = mWorkingRanges.get(key);
     if (rangeTuple == null) {
-      mWorkingRanges.put(key, new RangeTuple(name, workingRange, scopedComponentInfo));
+      mWorkingRanges.put(
+          key, new RangeTuple(name, workingRange, scopedComponentInfo, interStageProps));
     } else {
       rangeTuple.addComponent(scopedComponentInfo);
     }
@@ -88,7 +90,8 @@ && isEnteringRange(
                 firstFullyVisibleIndex,
                 lastFullyVisibleIndex)) {
           try {
-            component.dispatchOnEnteredRange(scopedContext, rangeTuple.mName);
+            component.dispatchOnEnteredRange(
+                scopedContext, rangeTuple.mName, rangeTuple.mInterStagePropsContainer);
           } catch (Exception e) {
             ComponentUtils.handle(scopedContext, e);
           }
@@ -103,7 +106,8 @@ && isExitingRange(
                 firstFullyVisibleIndex,
                 lastFullyVisibleIndex)) {
           try {
-            component.dispatchOnExitedRange(scopedContext, rangeTuple.mName);
+            component.dispatchOnExitedRange(
+                scopedContext, rangeTuple.mName, rangeTuple.mInterStagePropsContainer);
           } catch (Exception e) {
             ComponentUtils.handle(scopedContext, e);
           }
@@ -134,7 +138,8 @@ void dispatchOnExitedRangeIfNeeded(WorkingRangeStatusHandler statusHandler) {
         String globalKey = scopedContext.getGlobalKey();
         if (statusHandler.isInRange(rangeTuple.mName, component, globalKey)) {
           try {
-            component.dispatchOnExitedRange(scopedContext, rangeTuple.mName);
+            component.dispatchOnExitedRange(
+                scopedContext, rangeTuple.mName, rangeTuple.mInterStagePropsContainer);
           } catch (Exception e) {
             ComponentUtils.handle(scopedContext, e);
           }
@@ -189,15 +194,18 @@ static class RangeTuple {
     final String mName;
     final WorkingRange mWorkingRange;
     final List<ScopedComponentInfo> mScopedComponentInfos;
+    final @Nullable InterStagePropsContainer mInterStagePropsContainer;
 
     RangeTuple(
         final String name,
         final WorkingRange workingRange,
-        final ScopedComponentInfo scopedComponentInfo) {
+        final ScopedComponentInfo scopedComponentInfo,
+        final @Nullable InterStagePropsContainer interStagePropsContainer) {
       mName = name;
       mWorkingRange = workingRange;
       mScopedComponentInfos = new ArrayList<>();
       mScopedComponentInfos.add(scopedComponentInfo);
+      mInterStagePropsContainer = interStagePropsContainer;
     }
 
     void addComponent(final ScopedComponentInfo scopedComponentInfo) {

litho-it/src/test/java/com/facebook/litho/WorkingRangeContainerTest.java

@@ -18,6 +18,7 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.ArgumentMatchers.isA;
+import static org.mockito.ArgumentMatchers.isNull;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
@@ -71,7 +72,7 @@ public void setup() {
 
   @Test
   public void testRegisterWorkingRange() {
-    mWorkingRangeContainer.registerWorkingRange(NAME, mWorkingRange, mScopedComponentInfo);
+    mWorkingRangeContainer.registerWorkingRange(NAME, mWorkingRange, mScopedComponentInfo, null);
 
     final Map<String, RangeTuple> workingRanges =
         mWorkingRangeContainer.getWorkingRangesForTestOnly();
@@ -88,7 +89,7 @@ public void testRegisterWorkingRange() {
 
   @Test
   public void testIsEnteredRange() {
-    RangeTuple rangeTuple = new RangeTuple(NAME, mWorkingRange, mScopedComponentInfo);
+    RangeTuple rangeTuple = new RangeTuple(NAME, mWorkingRange, mScopedComponentInfo, null);
     WorkingRange workingRange = rangeTuple.mWorkingRange;
 
     assertThat(WorkingRangeContainer.isEnteringRange(workingRange, 0, 0, 1, 0, 1)).isEqualTo(true);
@@ -97,7 +98,7 @@ public void testIsEnteredRange() {
 
   @Test
   public void testIsExitedRange() {
-    RangeTuple rangeTuple = new RangeTuple(NAME, mWorkingRange, mScopedComponentInfo);
+    RangeTuple rangeTuple = new RangeTuple(NAME, mWorkingRange, mScopedComponentInfo, null);
     WorkingRange workingRange = rangeTuple.mWorkingRange;
 
     assertThat(WorkingRangeContainer.isExitingRange(workingRange, 0, 0, 1, 0, 1)).isEqualTo(false);
@@ -107,28 +108,28 @@ public void testIsExitedRange() {
   @Test
   public void testDispatchOnExitedRangeIfNeeded() {
     TestWorkingRange workingRange = new TestWorkingRange();
-    mWorkingRangeContainer.registerWorkingRange(NAME, workingRange, mScopedComponentInfo);
+    mWorkingRangeContainer.registerWorkingRange(NAME, workingRange, mScopedComponentInfo, null);
 
     TestWorkingRange workingRange2 = new TestWorkingRange();
-    mWorkingRangeContainer.registerWorkingRange(NAME, workingRange2, mScopedComponentInfo2);
+    mWorkingRangeContainer.registerWorkingRange(NAME, workingRange2, mScopedComponentInfo2, null);
 
     final WorkingRangeStatusHandler statusHandler = new WorkingRangeStatusHandler();
     statusHandler.setStatus(
         NAME, mComponent, "component", WorkingRangeStatusHandler.STATUS_IN_RANGE);
     doNothing()
         .when(mComponent)
-        .dispatchOnExitedRange(isA(ComponentContext.class), isA(String.class));
+        .dispatchOnExitedRange(isA(ComponentContext.class), isA(String.class), isNull());
 
     statusHandler.setStatus(
         NAME, mComponent2, "component2", WorkingRangeStatusHandler.STATUS_OUT_OF_RANGE);
     doNothing()
         .when(mComponent2)
-        .dispatchOnExitedRange(isA(ComponentContext.class), isA(String.class));
+        .dispatchOnExitedRange(isA(ComponentContext.class), isA(String.class), isNull());
 
     mWorkingRangeContainer.dispatchOnExitedRangeIfNeeded(statusHandler);
 
-    verify(mComponent, times(1)).dispatchOnExitedRange(mComponentContext, NAME);
-    verify(mComponent2, times(0)).dispatchOnExitedRange(mComponentContext, NAME);
+    verify(mComponent, times(1)).dispatchOnExitedRange(mComponentContext, NAME, null);
+    verify(mComponent2, times(0)).dispatchOnExitedRange(mComponentContext, NAME, null);
   }
 
   private static class TestWorkingRange implements WorkingRange {

litho-it/src/test/java/com/facebook/litho/specmodels/generator/WorkingRangeGeneratorTest.java

@@ -83,14 +83,15 @@ public void testGenerateDispatchOnEnteredRange() {
     assertThat(WorkingRangeGenerator.generateDispatchOnEnteredRangeMethod(mSpecModel).toString())
         .isEqualTo(
             "@java.lang.Override\n"
-                + "public void dispatchOnEnteredRange(com.facebook.litho.ComponentContext c, java.lang.String name) {\n"
+                + "public void dispatchOnEnteredRange(com.facebook.litho.ComponentContext c, java.lang.String name,\n"
+                + "    com.facebook.litho.InterStagePropsContainer interStageProps) {\n"
                 + "  switch (name) {\n"
                 + "    case \"enter\": {\n"
-                + "      testEnteredRangeMethod(c);\n"
+                + "      testEnteredRangeMethod(c, interStageProps);\n"
                 + "      return;\n"
                 + "    }\n"
                 + "    case \"prefetch\": {\n"
-                + "      testEnteredPrefetchMethod(c);\n"
+                + "      testEnteredPrefetchMethod(c, interStageProps);\n"
                 + "      return;\n"
                 + "    }\n"
                 + "  }\n"
@@ -102,14 +103,15 @@ public void testGenerateDispatchOnExitedRange() {
     assertThat(WorkingRangeGenerator.generateDispatchOnExitedRangeMethod(mSpecModel).toString())
         .isEqualTo(
             "@java.lang.Override\n"
-                + "public void dispatchOnExitedRange(com.facebook.litho.ComponentContext c, java.lang.String name) {\n"
+                + "public void dispatchOnExitedRange(com.facebook.litho.ComponentContext c, java.lang.String name,\n"
+                + "    com.facebook.litho.InterStagePropsContainer interStageProps) {\n"
                 + "  switch (name) {\n"
                 + "    case \"exit\": {\n"
-                + "      testExitedRangeMethod(c);\n"
+                + "      testExitedRangeMethod(c, interStageProps);\n"
                 + "      return;\n"
                 + "    }\n"
                 + "    case \"prefetch\": {\n"
-                + "      testExitedPrefetchMethod(c);\n"
+                + "      testExitedPrefetchMethod(c, interStageProps);\n"
                 + "      return;\n"
                 + "    }\n"
                 + "  }\n"
@@ -125,7 +127,8 @@ public void testGenerateWorkingRangeMethodDelegates() {
 
     assertThat(dataHolder.getMethodSpecs().get(0).toString())
         .isEqualTo(
-            "private void testEnteredRangeMethod(com.facebook.litho.ComponentContext c) {\n"
+            "private void testEnteredRangeMethod(com.facebook.litho.ComponentContext c,\n"
+                + "    com.facebook.litho.InterStagePropsContainer interStageProps) {\n"
                 + "  TestSpec.testEnteredRangeMethod(\n"
                 + "    (com.facebook.litho.ComponentContext) c,\n"
                 + "    (boolean) arg0,\n"
@@ -134,7 +137,8 @@ public void testGenerateWorkingRangeMethodDelegates() {
 
     assertThat(dataHolder.getMethodSpecs().get(1).toString())
         .isEqualTo(
-            "private void testExitedRangeMethod(com.facebook.litho.ComponentContext c) {\n"
+            "private void testExitedRangeMethod(com.facebook.litho.ComponentContext c,\n"
+                + "    com.facebook.litho.InterStagePropsContainer interStageProps) {\n"
                 + "  TestSpec.testExitedRangeMethod(\n"
                 + "    (com.facebook.litho.ComponentContext) c,\n"
                 + "    (T) arg2,\n"
@@ -143,7 +147,8 @@ public void testGenerateWorkingRangeMethodDelegates() {
 
     assertThat(dataHolder.getMethodSpecs().get(2).toString())
         .isEqualTo(
-            "private void testEnteredPrefetchMethod(com.facebook.litho.ComponentContext c) {\n"
+            "private void testEnteredPrefetchMethod(com.facebook.litho.ComponentContext c,\n"
+                + "    com.facebook.litho.InterStagePropsContainer interStageProps) {\n"
                 + "  TestSpec.testEnteredPrefetchMethod(\n"
                 + "    (com.facebook.litho.ComponentContext) c,\n"
                 + "    (boolean) arg0,\n"
@@ -152,7 +157,8 @@ public void testGenerateWorkingRangeMethodDelegates() {
 
     assertThat(dataHolder.getMethodSpecs().get(3).toString())
         .isEqualTo(
-            "private void testExitedPrefetchMethod(com.facebook.litho.ComponentContext c) {\n"
+            "private void testExitedPrefetchMethod(com.facebook.litho.ComponentContext c,\n"
+                + "    com.facebook.litho.InterStagePropsContainer interStageProps) {\n"
                 + "  TestSpec.testExitedPrefetchMethod(\n"
                 + "    (com.facebook.litho.ComponentContext) c,\n"
                 + "    (T) arg2,\n"

litho-processor/src/main/java/com/facebook/litho/specmodels/generator/WorkingRangeGenerator.java

@@ -59,15 +59,16 @@ static MethodSpec generateDispatchOnEnteredRangeMethod(SpecModel specModel) {
             .addAnnotation(Override.class)
             .returns(TypeName.VOID)
             .addParameter(specModel.getContextClass(), "c")
-            .addParameter(ClassNames.STRING, "name");
+            .addParameter(ClassNames.STRING, "name")
+            .addParameter(ClassNames.INTER_STAGE_PROPS_CONTAINER, "interStageProps");
 
     methodBuilder.beginControlFlow("switch (name)");
 
     for (WorkingRangeMethodModel model : specModel.getWorkingRangeMethods()) {
       if (model.enteredRangeModel != null && model.enteredRangeModel.typeModel != null) {
         final String nameInAnnotation = model.enteredRangeModel.typeModel.name;
         methodBuilder.beginControlFlow("case \"$L\":", nameInAnnotation);
-        methodBuilder.addStatement("$L(c)", model.enteredRangeModel.name);
+        methodBuilder.addStatement("$L(c, interStageProps)", model.enteredRangeModel.name);
         methodBuilder.addStatement("return");
         methodBuilder.endControlFlow();
       }
@@ -84,15 +85,16 @@ static MethodSpec generateDispatchOnExitedRangeMethod(SpecModel specModel) {
             .addAnnotation(Override.class)
             .returns(TypeName.VOID)
             .addParameter(specModel.getContextClass(), "c")
-            .addParameter(ClassNames.STRING, "name");
+            .addParameter(ClassNames.STRING, "name")
+            .addParameter(ClassNames.INTER_STAGE_PROPS_CONTAINER, "interStageProps");
 
     methodBuilder.beginControlFlow("switch (name)");
 
     for (WorkingRangeMethodModel model : specModel.getWorkingRangeMethods()) {
       if (model.exitedRangeModel != null && model.exitedRangeModel.typeModel != null) {
         final String nameInAnnotation = model.exitedRangeModel.typeModel.name;
         methodBuilder.beginControlFlow("case \"$L\":", nameInAnnotation);
-        methodBuilder.addStatement("$L(c)", model.exitedRangeModel.name);
+        methodBuilder.addStatement("$L(c, interStageProps)", model.exitedRangeModel.name);
         methodBuilder.addStatement("return");
         methodBuilder.endControlFlow();
       }
@@ -124,7 +126,8 @@ private static MethodSpec generateWorkingRangeMethodDelegate(
         MethodSpec.methodBuilder(methodModel.name.toString())
             .addModifiers(Modifier.PRIVATE)
             .returns(TypeName.VOID)
-            .addParameter(ClassNames.COMPONENT_CONTEXT, "c");
+            .addParameter(ClassNames.COMPONENT_CONTEXT, "c")
+            .addParameter(ClassNames.INTER_STAGE_PROPS_CONTAINER, "interStageProps");
 
     final CodeBlock.Builder delegation = CodeBlock.builder();
 
@@ -134,7 +137,7 @@ private static MethodSpec generateWorkingRangeMethodDelegate(
           "$L $L = $L",
           ClassNames.INTER_STAGE_PROPS_CONTAINER,
           ComponentBodyGenerator.LOCAL_INTER_STAGE_PROPS_CONTAINER_NAME,
-          "null");
+          "interStageProps");
     }
 
     final String sourceDelegateAccessor = SpecModelUtils.getSpecAccessor(specModel);