pascal.taie.util.collection.MapSetMultiMap Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of tai-e Show documentation
An easy-to-learn/use static analysis framework for Java
The newest version!
/*
 * Tai-e: A Static Analysis Framework for Java
 *
 * Copyright (C) 2022 Tian Tan 
 * Copyright (C) 2022 Yue Li 
 *
 * This file is part of Tai-e.
 *
 * Tai-e is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License
 * as published by the Free Software Foundation, either version 3
 * of the License, or (at your option) any later version.
 *
 * Tai-e is distributed in the hope that it will be useful,but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
 * or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General
 * Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with Tai-e. If not, see .
 */

package pascal.taie.util.collection;

import pascal.taie.util.function.SSupplier;

import javax.annotation.Nonnull;
import java.io.Serializable;
import java.util.Collection;
import java.util.Collections;
import java.util.Iterator;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Objects;
import java.util.Set;
import java.util.function.BiConsumer;

/**
 * An implementation of {@link MultiMap} that stores key-value pairs
 * as a map from key to the sets of its corresponding values.
 *
 * @param  type of the keys in this map
 * @param  type of the values in this map
 */
public class MapSetMultiMap extends AbstractMultiMap
        implements Serializable {

    /**
     * The backing map.
     */
    private final Map> map;

    /**
     * Factory function for creating new sets.
     */
    private final SSupplier> setFactory;

    private int size = 0;

    public MapSetMultiMap(Map> map, SSupplier> setFactory) {
        this.map = map;
        this.setFactory = setFactory;
    }

    @Override
    public boolean contains(K key, V value) {
        return get(key).contains(value);
    }

    @Override
    public boolean containsKey(K key) {
        return map.containsKey(key);
    }

    @Override
    public Set get(@Nonnull K key) {
        Objects.requireNonNull(key, NULL_KEY);
        Set values = map.get(key);
        return values == null ? Set.of() :
                Collections.unmodifiableSet(values);
    }

    @Override
    public boolean put(@Nonnull K key, @Nonnull V value) {
        Objects.requireNonNull(key, NULL_KEY);
        Objects.requireNonNull(value, NULL_VALUE);
        if (getOrCreateSet(key).add(value)) {
            ++size;
            return true;
        } else {
            return false;
        }
    }

    @Override
    public boolean putAll(@Nonnull K key, @Nonnull Collection values) {
        Objects.requireNonNull(key, NULL_KEY);
        Objects.requireNonNull(values);
        if (!values.isEmpty()) {
            Set set = getOrCreateSet(key);
            int diff = 0;
            for (V v : values) {
                if (set.add(Objects.requireNonNull(v, NULL_VALUE))) {
                    ++diff;
                }
            }
            if (diff > 0) {
                size += diff;
                return true;
            }
        }
        return false;
    }

    private Set getOrCreateSet(@Nonnull K key) {
        return map.computeIfAbsent(key, __ -> setFactory.get());
    }

    @Override
    public boolean putAll(@Nonnull MultiMap multiMap) {
        Objects.requireNonNull(multiMap);
        boolean[] changed = {false};
        multiMap.forEachSet((k, vs) -> changed[0] |= putAll(k, vs));
        return changed[0];
    }

    @Override
    public boolean remove(K key, V value) {
        Collection values = map.get(key);
        if (values != null && values.remove(value)) {
            if (values.isEmpty()) {
                map.remove(key);
            }
            --size;
            return true;
        } else {
            return false;
        }
    }

    @Override
    public boolean removeAll(K key) {
        Set pre = map.remove(key);
        if (pre != null) {
            size -= pre.size();
            return true;
        } else {
            return false;
        }
    }

    @Override
    public boolean removeAll(K key, Collection values) {
        Collection currValues = map.get(key);
        if (currValues != null) {
            int beforeSize = currValues.size();
            currValues.removeAll(values);
            int diff = beforeSize - currValues.size();
            if (diff > 0) {
                size -= diff;
                if (currValues.isEmpty()) {
                    map.remove(key);
                }
                return true;
            }
        }
        return false;
    }

    @Override
    public Set keySet() {
        return Collections.unmodifiableSet(map.keySet());
    }

    @Override
    protected Iterator> entryIterator() {
        return new EntryIterator();
    }

    private final class EntryIterator implements Iterator> {

        private final Iterator>> mapIt;

        private Iterator valueIt;

        private K currKey;

        private EntryIterator() {
            mapIt = map.entrySet().iterator();
            if (mapIt.hasNext()) {
                advanceKey();
            } else {
                valueIt = Collections.emptyIterator();
            }
        }

        @Override
        public boolean hasNext() {
            return valueIt.hasNext() || mapIt.hasNext();
        }

        @Override
        public Map.Entry next() {
            if (valueIt.hasNext()) {
                return new ImmutableMapEntry<>(currKey, valueIt.next());
            } else if (mapIt.hasNext()) {
                advanceKey();
                return new ImmutableMapEntry<>(currKey, valueIt.next());
            } else {
                throw new NoSuchElementException();
            }
        }

        private void advanceKey() {
            var entry = mapIt.next();
            currKey = entry.getKey();
            valueIt = entry.getValue().iterator();
        }
    }

    @Override
    public void forEachSet(@Nonnull BiConsumer> action) {
        map.forEach(action);
    }

    @Override
    public void clear() {
        map.clear();
        size = 0;
    }

    @Override
    public int size() {
        return size;
    }

    @Override
    public int hashCode() {
        return map.hashCode();
    }
}