/*
 * Copyright (C) 2015 Google 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.google.cloud.dataflow.sdk.util.state;

import com.google.cloud.dataflow.sdk.util.state.StateTag.StateBinder;
import com.google.common.base.Supplier;
import com.google.common.collect.Table;
import com.google.common.collect.Tables;

import java.util.HashMap;
import java.util.Map;
import java.util.Set;

/**
 * Table mapping {@code StateNamespace} and {@code StateTag<?>} to a {@code State} instance.
 */
public abstract class StateTable<K> {

  private final Table<StateNamespace, StateTag<? super K, ?>, State> stateTable =
      Tables.newCustomTable(new HashMap<StateNamespace, Map<StateTag<? super K, ?>, State>>(),
          new Supplier<Map<StateTag<? super K, ?>, State>>() {
        @Override
        public Map<StateTag<? super K, ?>, State> get() {
          return new HashMap<>();
        }
      });

  /**
   * Gets the {@link State} in the specified {@link StateNamespace} with the specified {@link
   * StateTag}, binding it using the {@link #binderForNamespace} if it is not
   * already present in this {@link StateTable}.
   */
  public <StateT extends State> StateT get(
      StateNamespace namespace, StateTag<? super K, StateT> tag, StateContext<?> c) {
    State storage = stateTable.get(namespace, tag);
    if (storage != null) {
      @SuppressWarnings("unchecked")
      StateT typedStorage = (StateT) storage;
      return typedStorage;
    }

    StateT typedStorage = tag.bind(binderForNamespace(namespace, c));
    stateTable.put(namespace, tag, typedStorage);
    return typedStorage;
  }

  public void clearNamespace(StateNamespace namespace) {
    stateTable.rowKeySet().remove(namespace);
  }

  public void clear() {
    stateTable.clear();
  }

  public Iterable<State> values() {
    return stateTable.values();
  }

  public boolean isNamespaceInUse(StateNamespace namespace) {
    return stateTable.containsRow(namespace);
  }

  public Map<StateTag<? super K, ?>, State> getTagsInUse(StateNamespace namespace) {
    return stateTable.row(namespace);
  }

  public Set<StateNamespace> getNamespacesInUse() {
    return stateTable.rowKeySet();
  }

  /**
   * Provide the {@code StateBinder} to use for creating {@code Storage} instances
   * in the specified {@code namespace}.
   */
  protected abstract StateBinder<K> binderForNamespace(StateNamespace namespace, StateContext<?> c);
}

@@ -738,21 +738,30 @@
 mon.
-base.Supplier
+collect.HashBasedTable
 ;%0Aim
@@ -802,75 +802,8 @@
 le;%0A
-import com.google.common.collect.Tables;%0A%0Aimport java.util.HashMap;
 %0Aimp
@@ -1076,287 +1076,30 @@
     
-Tables.newCustomTable(new HashMap%3CStateNamespace, Map%3CStateTag%3C? super K, ?%3E, State%3E%3E(),%0A          new Supplier%3CMap%3CStateTag%3C? super K, ?%3E, State%3E%3E() %7B%0A        @Override%0A        public Map%3CStateTag%3C? super K, ?%3E, State%3E get() %7B%0A          return new HashMap%3C%3E();%0A        %7D%0A      %7D
+HashBasedTable.create(
 );%0A%0A

Use HashBasedTable.create() instead of Tables.newCustomTable() (#336)