/*
 * Copyright (c) 2017-present, Facebook, Inc.
 * All rights reserved.
 *
 * This source code is licensed under the BSD-style license found in the
 * LICENSE file in the root directory of this source tree. An additional grant
 * of patent rights can be found in the PATENTS file in the same directory.
 */

package com.facebook.litho;

import java.util.ArrayList;
import java.util.List;

import android.os.Looper;
import android.support.annotation.Nullable;

/**
 * A {@link ComponentTree} for testing purposes. Leverages test classes to create component layouts
 * and exposes additional information useful for testing.
 */
public class TestComponentTree extends ComponentTree {

  public static Builder create(ComponentContext context, Component<?> root) {
    return new Builder(context, root);
  }

  private TestComponentTree(ComponentTree.Builder builder) {
    super(builder);
  }

  public List<Component> getSubComponents() {
    final List<Component> subComponents = new ArrayList<>();
    extractSubComponents(getMainThreadLayoutState().getDiffTree(), subComponents);

    return subComponents;
  }

  @Override
  protected LayoutState calculateLayoutState(
      @Nullable Object lock,
      ComponentContext context,
      Component<?> root,
      int widthSpec,
      int heightSpec,
      boolean diffingEnabled,
      @Nullable DiffNode diffNode) {
    return LayoutState.calculate(
        new TestComponentContext(
            ComponentContext.withComponentTree(new TestComponentContext(context), this),
            new StateHandler()),
        root,
        mId,
        widthSpec,
        heightSpec,
        diffingEnabled,
        diffNode,
        false /* canPrefetchDisplayLists */,
        false /* canCacheDrawingDisplayLists */,
        true /* clipChildren */);
  }

  private static void extractSubComponents(DiffNode root, List<Component> output) {
    if (root == null) {
      return;
    }

    if (root.getChildCount() == 0) {
      if (root.getComponent() != null && root.getComponent() instanceof TestComponent) {
        TestComponent testSubcomponent = (TestComponent) root.getComponent();
        output.add(testSubcomponent.getWrappedComponent());
      }

      return;
    }

    for (DiffNode child : root.getChildren()) {
      extractSubComponents(child, output);
    }
  }

  public static class Builder extends ComponentTree.Builder {

    private Builder(ComponentContext context, Component<?> root) {
      super(context, root);
    }

    @Override
    public Builder incrementalMount(boolean isEnabled) {
      return (Builder) super.incrementalMount(isEnabled);
    }

    @Override
    public Builder layoutDiffing(boolean enabled) {
      return (Builder) super.layoutDiffing(enabled);
    }

    @Override
    public Builder layoutThreadLooper(Looper looper) {
      return (Builder) super.layoutThreadLooper(looper);
    }

    @Override
    public Builder layoutLock(Object layoutLock) {
      return (Builder) super.layoutLock(layoutLock);
    }

    @Override
    public TestComponentTree build() {

      return new TestComponentTree(this);
    }
  }
}

@@ -359,16 +359,46 @@
 ayList;%0A
+import java.util.Collections;%0A
 import j
@@ -968,68 +968,14 @@
     
-final List%3CComponent%3E subComponents = new ArrayList%3C%3E();%0A   
+return
  ext
@@ -1036,51 +1036,9 @@
 ee()
-, subComponents);%0A%0A    return subComponents
+)
 ;%0A  
@@ -1746,12 +1746,23 @@
 tic 
-void
+List%3CComponent%3E
  ext
@@ -1796,32 +1796,8 @@
 root
-, List%3CComponent%3E output
 ) %7B%0A
@@ -1828,31 +1828,110 @@
       return
-;%0A    %7D
+ Collections.emptyList();%0A    %7D%0A%0A    final List%3CComponent%3E output = new ArrayList%3C%3E();
 %0A%0A    if (ro
@@ -2205,16 +2205,23 @@
   return
+ output
 ;%0A    %7D%0A
@@ -2275,16 +2275,30 @@
 %7B%0A      
+output.addAll(
 extractS
@@ -2319,24 +2319,37 @@
 hild
-, output);%0A    %7D
+));%0A    %7D%0A%0A    return output;
 %0A  %7D

Simplify SubComponent extraction