/*
 * Copyright 2017 Google Inc. All rights reserved.
 *
 * 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 org.inferred.freebuilder.processor.util;

import static java.util.Objects.requireNonNull;

import com.google.common.collect.FluentIterable;
import com.google.common.collect.ImmutableSet;

import java.util.ConcurrentModificationException;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Set;
import java.util.function.Supplier;

public abstract class Scope {

  public enum Level {
    FILE, METHOD;
  }

  @SuppressWarnings("unused")
  public interface Key<V> {
    Level level();
  }

  private static final Object RECURSION_SENTINEL = new Object();

  private final Map<Key<?>, Object> entries = new LinkedHashMap<>();
  private final Scope parent;
  private final Level level;

  private Scope(Scope parent, Level level) {
    this.parent = parent;
    this.level = level;
  }

  public boolean contains(Key<?> key) {
    return get(key) != null;
  }

  public <V> V get(Key<V> key) {
    @SuppressWarnings("unchecked")
    V value = (V) entries.get(key);
    if (value == RECURSION_SENTINEL) {
      throw new ConcurrentModificationException(
          "Cannot access scope key " + key + " while computing its value");
    } else if (value != null) {
      return value;
    } else if (parent != null) {
      return parent.get(key);
    } else {
      return null;
    }
  }

  public <V> V computeIfAbsent(Key<V> key, Supplier<V> supplier) {
    V value = get(key);
    if (value != null) {
      return value;
    } else if (level == key.level()) {
      entries.put(key, RECURSION_SENTINEL);
      value = supplier.get();
      entries.put(key, value);
      return value;
    } else if (parent != null) {
      return parent.computeIfAbsent(key, supplier);
    } else {
      throw new IllegalStateException(
          "Not in " + key.level().toString().toLowerCase() + " scope");
    }
  }

  public <V> Set<V> keysOfType(Class<V> keyType) {
    ImmutableSet.Builder<V> keys = ImmutableSet.builder();
    if (parent != null) {
      keys.addAll(parent.keysOfType(keyType));
    }
    keys.addAll(FluentIterable.from(entries.keySet()).filter(keyType).toSet());
    return keys.build();
  }

  public <V> V putIfAbsent(Key<V> key, V value) {
    requireNonNull(key);
    requireNonNull(value);
    if (level == key.level()) {
      @SuppressWarnings("unchecked")
      V existingValue = (V) entries.get(key);
      if (value == RECURSION_SENTINEL) {
        throw new ConcurrentModificationException(
            "Cannot access scope key " + key + " while computing its value");
      } else if (existingValue == null) {
        entries.put(key, value);
      }
      return existingValue;
    } else if (parent != null) {
      return parent.putIfAbsent(key, value);
    }
    return null;
  }

  static class FileScope extends Scope {
    FileScope() {
      super(null, Level.FILE);
    }
  }

  static class MethodScope extends Scope {
    MethodScope(Scope parent) {
      super(parent, Level.METHOD);
    }
  }
}

@@ -964,16 +964,721 @@
 plier;%0A%0A
+/**%0A * An append-only, hierarchical map with key-specific value typing.%0A *%0A * %3Cp%3EScopes allow source generators to cooperate with each other&mdash;to avoid namespace clashes%0A * or reuse variables, for example&mdash;in a decoupled fashion using a shared key-value space.%0A * The key type dictates the type of the value that can be stored in the map, as well as at what%0A * level in the source the key should be unique at: %7B@code FILE%7D keys are unique within a%0A * compilation unit, while %7B@code METHOD%7D keys are scoped to the current method.%0A *%0A * %3Cp%3EWhile a Scope has to be mutable, it limits the potential for complex bugs by only permitting%0A * key-values pairs to be inserted, not modified or removed.%0A */%0A
 public a
@@ -2614,32 +2614,268 @@
 ull;%0A    %7D%0A  %7D%0A%0A
+  /**%0A   * If %7B@code key%7D is not already associated with a value, computes its value using%0A   * %7B@code supplier%7D and enters it into the scope.%0A   *%0A   * @return the current (existing or computed) value associated with %7B@code key%7D%0A   */%0A
   public %3CV%3E V c
@@ -3122,38 +3122,54 @@
 ntries.put(key, 
+requireNonNull(
 value)
+)
 ;%0A      return v
@@ -3689,24 +3689,217 @@
 ild();%0A  %7D%0A%0A
+  /**%0A   * If %7B@code key%7D is not already associated with a value, associates it with %7B@code value%7D.%0A   *%0A   * @return the original value, or %7B@code null%7D if there was no value associated%0A   */%0A
   public %3CV%3E

Clarify the Scope API