Source code

001/*
002 * Licensed to the Apache Software Foundation (ASF) under one
003 * or more contributor license agreements.  See the NOTICE file
004 * distributed with this work for additional information
005 * regarding copyright ownership.  The ASF licenses this file
006 * to you under the Apache License, Version 2.0 (the
007 * "License"); you may not use this file except in compliance
008 * with the License.  You may obtain a copy of the License at
009 *
010 *   http://www.apache.org/licenses/LICENSE-2.0
011 *
012 * Unless required by applicable law or agreed to in writing,
013 * software distributed under the License is distributed on an
014 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
015 * KIND, either express or implied.  See the License for the
016 * specific language governing permissions and limitations
017 * under the License.
018 */
019package org.apache.reef.io.network.group.impl.utils;
020
021import org.apache.reef.io.network.group.impl.driver.MsgKey;
022
023import java.util.*;
024
025/**
026 * Map from K to {@code Set<V>}.
027 */
028public class SetMap<K, V> {
029  private final Map<K, Set<V>> map = new HashMap<>();
030
031  public boolean containsKey(final K key) {
032    return map.containsKey(key);
033  }
034
035  public boolean contains(final K key, final V value) {
036    if (!containsKey(key)) {
037      return false;
038    }
039    return map.get(key).contains(value);
040  }
041
042  public Set<V> get(final K key) {
043    if (map.containsKey(key)) {
044      return map.get(key);
045    } else {
046      return Collections.emptySet();
047    }
048  }
049
050  public void add(final K key, final V value) {
051    final Set<V> values;
052    if (!map.containsKey(key)) {
053      values = new HashSet<>();
054      map.put(key, values);
055    } else {
056      values = map.get(key);
057    }
058    values.add(value);
059  }
060
061  public boolean remove(final K key, final V value) {
062    if (!map.containsKey(key)) {
063      return false;
064    }
065    final Set<V> set = map.get(key);
066    final boolean retVal = set.remove(value);
067    if (set.isEmpty()) {
068      map.remove(key);
069    }
070    return retVal;
071  }
072
073  /**
074   * @param key
075   * @return
076   */
077  public int count(final K key) {
078    if (!containsKey(key)) {
079      return 0;
080    } else {
081      return map.get(key).size();
082    }
083  }
084
085  /**
086   * @param key
087   */
088  public Set<V> remove(final MsgKey key) {
089    return map.remove(key);
090  }
091
092  public Set<K> keySet() {
093    return map.keySet();
094  }
095}