package org.inferred.freebuilder.processor.util.feature;

import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;

import javax.annotation.processing.ProcessingEnvironment;

/**
 * A set of {@link Feature} instances, determined dynamically by calling
 * {@link FeatureType#forEnvironment}.
 */
public class EnvironmentFeatureSet implements FeatureSet {

  private class FeatureFromEnvironmentLoader
      extends CacheLoader<FeatureType<?>, Feature<?>> {
    private final ProcessingEnvironment env;

    /** <pre>featureType -> featureType.forEnvironment(env)</pre> */
    private FeatureFromEnvironmentLoader(ProcessingEnvironment env) {
      this.env = env;
    }

    @Override
    public Feature<?> load(FeatureType<?> featureType) {
      return featureType.forEnvironment(env, EnvironmentFeatureSet.this);
    }
  }

  private final LoadingCache<FeatureType<?>, Feature<?>> featuresByType;

  /** Constructs a feature set using the given processing environment. */
  public EnvironmentFeatureSet(ProcessingEnvironment env) {
    featuresByType = CacheBuilder.newBuilder()
        .concurrencyLevel(1)
        .build(new FeatureFromEnvironmentLoader(env));
  }

  @Override
  public <T extends Feature<T>> T get(FeatureType<T> featureType) {
    @SuppressWarnings("unchecked")
    T feature = (T) featuresByType.getUnchecked(featureType);
    return feature;
  }
}

@@ -62,133 +62,47 @@
 ort 
-com.google.common.cache.CacheBuilder;%0Aimport com.google.common.cache.CacheLoader;%0Aimport com.google.common.cache.LoadingCache
+java.util.HashMap;%0Aimport java.util.Map
 ;%0A%0Ai
@@ -343,111 +343,8 @@
  %7B%0A%0A
-  private class FeatureFromEnvironmentLoader%0A      extends CacheLoader%3CFeatureType%3C?%3E, Feature%3C?%3E%3E %7B%0A  
   pr
@@ -386,240 +386,63 @@
 nv;%0A
-%0A
   
-  /** %3Cpre%3EfeatureType -%3E featureType.forEnvironment(env)%3C/pre%3E */%0A    private FeatureFromEnvironmentLoader(ProcessingEnvironment env) %7B%0A      this.env = env;%0A    %7D%0A%0A    @Override%0A    public Feature%3C?%3E load(
+private final Map%3CFeatureType%3C?%3E, 
 Feature
-Type
 %3C?%3E
+%3E
  feature
 Type
@@ -441,172 +441,33 @@
 ture
+sBy
 Type
-) %7B%0A      return featureType.forEnvironment(env, EnvironmentFeatureSet.this);%0A    %7D%0A  %7D%0A%0A  private final LoadingCache%3CFeatureType%3C?%3E, Feature%3C?%3E%3E featuresByType
+ = new HashMap%3C%3E()
 ;%0A%0A 
@@ -607,133 +607,22 @@
     
-featuresByType = CacheBuilder.newBuilder()%0A        .concurrencyLevel(1)%0A        .build(new FeatureFromEnvironmentLoader(env))
+this.env = env
 ;%0A  
@@ -778,32 +778,69 @@
 ype.
-getUnchecked(featureType
+computeIfAbsent(featureType, $ -%3E $.forEnvironment(env, this)
 );%0A 

Use Map.computeIfAbsent rather than a LoadingCache