/*
 * Copyright 2018-present Facebook, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may
 * not use this file except in compliance with the License. You may obtain
 * a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations
 * under the License.
 */

package com.facebook.buck.cxx;

import com.facebook.buck.core.model.BuildTarget;
import com.facebook.buck.core.rules.ActionGraphBuilder;
import com.facebook.buck.cxx.toolchain.CxxPlatform;
import com.facebook.buck.util.concurrent.Parallelizer;
import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSortedMap;
import com.google.common.util.concurrent.UncheckedExecutionException;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ExecutionException;
import javax.annotation.Nonnull;

/** Transitive C++ preprocessor input cache */
public class TransitiveCxxPreprocessorInputCache {
  private final Cache<CxxPlatform, ImmutableSortedMap<BuildTarget, CxxPreprocessorInput>> cache =
      CacheBuilder.newBuilder().build();
  private final CxxPreprocessorDep preprocessorDep;

  public TransitiveCxxPreprocessorInputCache(CxxPreprocessorDep preprocessorDep) {
    this.preprocessorDep = preprocessorDep;
  }

  /** Get a value from the cache */
  public ImmutableMap<BuildTarget, CxxPreprocessorInput> getUnchecked(
      CxxPlatform key, ActionGraphBuilder graphBuilder) {
    try {
      return cache.get(
          key,
          () ->
              computeTransitiveCxxToPreprocessorInputMap(
                  key, preprocessorDep, true, graphBuilder, graphBuilder.getParallelizer()));
    } catch (ExecutionException e) {
      throw new UncheckedExecutionException(e.getCause());
    }
  }

  public static ImmutableMap<BuildTarget, CxxPreprocessorInput>
      computeTransitiveCxxToPreprocessorInputMap(
          @Nonnull CxxPlatform key,
          CxxPreprocessorDep preprocessorDep,
          boolean includeDep,
          ActionGraphBuilder graphBuilder) {
    return computeTransitiveCxxToPreprocessorInputMap(
        key, preprocessorDep, includeDep, graphBuilder, graphBuilder.getParallelizer());
  }

  private static ImmutableSortedMap<BuildTarget, CxxPreprocessorInput>
      computeTransitiveCxxToPreprocessorInputMap(
          @Nonnull CxxPlatform key,
          CxxPreprocessorDep preprocessorDep,
          boolean includeDep,
          ActionGraphBuilder graphBuilder,
          Parallelizer parallelizer) {
    Map<BuildTarget, CxxPreprocessorInput> builder = new HashMap<>();
    if (includeDep) {
      builder.put(
          preprocessorDep.getBuildTarget(),
          preprocessorDep.getCxxPreprocessorInput(key, graphBuilder));
    }

    Collection<ImmutableMap<BuildTarget, CxxPreprocessorInput>> transitiveDepInputs =
        parallelizer.maybeParallelizeTransform(
            ImmutableList.copyOf(preprocessorDep.getCxxPreprocessorDeps(key, graphBuilder)),
            dep -> dep.getTransitiveCxxPreprocessorInput(key, graphBuilder));
    transitiveDepInputs.forEach(builder::putAll);

    // Using an ImmutableSortedMap here:
    //
    // 1. Memory efficiency. ImmutableSortedMap is implemented with 2 lists (an ImmutableSortedSet
    // of keys, and a ImmutableList of values). This is much more efficient than an ImmutableMap,
    // which creates an Entry instance for each entry.
    //
    // 2. Historically we seem to care that the result has some definite order.
    //
    // 3. We mostly iterate over these maps rather than do lookups, so ImmutableSortedMap
    // binary-search based lookup is not an issue.
    return ImmutableSortedMap.copyOf(builder);
  }
}

@@ -840,24 +840,89 @@
 rallelizer;%0A
+import com.facebook.buck.util.concurrent.WorkThreadTrackingTask;%0A
 import com.g
@@ -1144,78 +1144,8 @@
 ap;%0A
-import com.google.common.util.concurrent.UncheckedExecutionException;%0A
 impo
@@ -1238,37 +1238,95 @@
 til.
-concurrent.ExecutionException
+Objects;%0Aimport java.util.concurrent.ConcurrentHashMap;%0Aimport javax.annotation.Nonnull
 ;%0Aim
@@ -1348,22 +1348,23 @@
 tation.N
-onn
 ull
+able
 ;%0A%0A/** T
@@ -1646,16 +1646,210 @@
 ssorDep;
+%0A  private final ConcurrentHashMap%3C%0A          CxxPlatform,%0A          WorkThreadTrackingTask%3CImmutableSortedMap%3CBuildTarget, CxxPreprocessorInput%3E%3E%3E%0A      cacheLoader = new ConcurrentHashMap%3C%3E();
 %0A%0A  publ
@@ -2147,38 +2147,379 @@
     
-try %7B%0A      return cache.ge
+@Nullable%0A    ImmutableSortedMap%3CBuildTarget, CxxPreprocessorInput%3E result = cache.getIfPresent(key);%0A    if (result != null) %7B%0A      return result;%0A    %7D%0A%0A    // We schedule the compute for the given key so that it is computed only once.%0A    WorkThreadTrackingTask%3CImmutableSortedMap%3CBuildTarget, CxxPreprocessorInput%3E%3E loadTask =%0A        cacheLoader.computeIfAbsen
 t(%0A
+  
     
@@ -2543,14 +2543,102 @@
     
-() -%3E%0A
+  mapKey -%3E%0A                new WorkThreadTrackingTask%3C%3E(%0A                    () -%3E%0A          
     
@@ -2713,49 +2713,171 @@
     
+        
+  
 key,
- preprocessorDep, true, graphBuilder,
+%0A                            preprocessorDep,%0A                            true,%0A                            graphBuilder,%0A                           
  gra
@@ -2909,111 +2909,162 @@
 ()))
-;
+);%0A
 %0A    
-%7D catch (ExecutionException e) %7B%0A      throw new UncheckedExecutionException(e.getCause());%0A    %7D
+result = Objects.requireNonNull(loadTask.externalCompute());%0A%0A    // store to cache after computing%0A    cache.put(key, result);%0A    return result;
 %0A  %7D

Make TransitiveCxxPreprocessorInputCache work steal load