package com.atlassian.vcache.internal.core.service;
   
   import java.util.Map;
   import java.util.Objects;
   import java.util.Optional;
   import java.util.Set;
   import java.util.function.Function;
   import java.util.stream.Collectors;
   import java.util.stream.StreamSupport;
  
  import static java.util.Objects.requireNonNull;
  
  /**
   * Provides common functionality for local cache operations.
   *
   * @since 1.6.2
   */
  public class LocalCacheUtils {
      /**
       * Arguments for a put operation
       *
       * @param <K> the key type
       * @param <V> the value type
       */
      public static class PutArgs<K, V> {
          public final K key;
          public final V value;
  
          private PutArgs(K key, V value) {
              this.key = requireNonNull(key);
              this.value = requireNonNull(value);
          }
      }
  
      /**
       * Performs a bulk get.
       *
       * @param factory       for creating missing values
       * @param keys          the keys to return
       * @param getFn         function for getting an existing value
       * @param putIfAbsentFn function for putting a new value
       * @param lock          the lock for the cache
       * @param <K>           the key type
       * @param <V>           the value type
       * @return the map of all key/value pairs.
       */
      public static <K, V> Map<K, V> getBulk(
              Function<Set<K>, Map<K, V>> factory,
              Iterable<K> keys,
              Function<K, Optional<V>> getFn,
              Function<PutArgs<K, V>, Optional<V>> putIfAbsentFn,
              VCacheLock lock) {
          // Get the existing values
          final Map<K, Optional<V>> existingValues =
                  lock.withLock(() ->
                          StreamSupport.stream(keys.spliterator(), false)
                                  .distinct()
                                  .collect(Collectors.toMap(
                                          Objects::requireNonNull,
                                          getFn)));
  
          // Add known values to the grand result
          @SuppressWarnings("OptionalGetWithoutIsPresent")
          final Map<K, V> grandResult = existingValues.entrySet().stream()
                  .filter(e -> e.getValue().isPresent())
                  .collect(Collectors.toMap(
                          Map.Entry::getKey,
                          e -> e.getValue().get()));
  
          // Bail out if we have all the values
          if (grandResult.size() == existingValues.size()) {
              return grandResult;
          }
  
          // Sadly we now need to call the factory to create the missing values and then merge into the grand result.
          final Set<K> missingKeys = existingValues.entrySet().stream()
                  .filter(e -> !e.getValue().isPresent())
                  .map(Map.Entry::getKey)
                  .collect(Collectors.toSet());
  
          final Map<K, V> missingValues = factory.apply(missingKeys);
          FactoryUtils.verifyFactoryResult(missingValues, missingKeys);
  
          lock.withLock(() ->
                  missingValues.entrySet().forEach(e -> {
                      // Handle that another thread may have beaten us to the punch.
                      final Optional<V> existing = putIfAbsentFn.apply(new PutArgs<>(e.getKey(), e.getValue()));
                      grandResult.put(e.getKey(), existing.orElse(e.getValue()));
                  }));
  
          return grandResult;
      }
  }