001    /*
002     * Copyright 2010-2015 JetBrains s.r.o.
003     *
004     * Licensed under the Apache License, Version 2.0 (the "License");
005     * you may not use this file except in compliance with the License.
006     * You may obtain a copy of the License at
007     *
008     * http://www.apache.org/licenses/LICENSE-2.0
009     *
010     * Unless required by applicable law or agreed to in writing, software
011     * distributed under the License is distributed on an "AS IS" BASIS,
012     * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013     * See the License for the specific language governing permissions and
014     * limitations under the License.
015     */
016    
017    package org.jetbrains.kotlin.util.slicedMap;
018    
019    import com.google.common.collect.ArrayListMultimap;
020    import com.google.common.collect.ImmutableMap;
021    import com.google.common.collect.Multimap;
022    import com.intellij.openapi.util.Key;
023    import com.intellij.openapi.util.UserDataHolder;
024    import gnu.trove.THashMap;
025    import kotlin.jvm.functions.Function3;
026    import org.jetbrains.annotations.NotNull;
027    
028    import java.util.Collection;
029    import java.util.Collections;
030    import java.util.Map;
031    
032    public class SlicedMapImpl implements MutableSlicedMap {
033    
034        public static SlicedMapImpl create() {
035            return new SlicedMapImpl();
036        }
037    
038        private final Map<Object, UserDataHolderImpl> map = new THashMap<Object, UserDataHolderImpl>(0);
039        private Multimap<WritableSlice<?, ?>, Object> collectiveSliceKeys = null;
040    
041        @Override
042        public <K, V> void put(WritableSlice<K, V> slice, K key, V value) {
043            if (!slice.check(key, value)) {
044                return;
045            }
046    
047            UserDataHolderImpl holder = map.get(key);
048            if (holder == null) {
049                holder = new UserDataHolderImpl();
050                map.put(key, holder);
051            }
052    
053            Key<V> sliceKey = slice.getKey();
054    
055            RewritePolicy rewritePolicy = slice.getRewritePolicy();
056            if (rewritePolicy.rewriteProcessingNeeded(key)) {
057                V oldValue = holder.getUserData(sliceKey);
058                if (oldValue != null) {
059                    //noinspection unchecked
060                    if (!rewritePolicy.processRewrite(slice, key, oldValue, value)) {
061                        return;
062                    }
063                }
064            }
065    
066            if (slice.isCollective()) {
067                if (collectiveSliceKeys == null) {
068                    collectiveSliceKeys = ArrayListMultimap.create();
069                }
070    
071                collectiveSliceKeys.put(slice, key);
072            }
073    
074            holder.putUserData(sliceKey, value);
075            slice.afterPut(this, key, value);
076        }
077    
078        @Override
079        public void clear() {
080            map.clear();
081            collectiveSliceKeys = null;
082        }
083    
084        @Override
085        public <K, V> V get(ReadOnlySlice<K, V> slice, K key) {
086            UserDataHolderImpl holder = map.get(key);
087    
088            V value = holder == null ? null : holder.getUserData(slice.getKey());
089    
090            return slice.computeValue(this, key, value, value == null);
091        }
092    
093        @Override
094        @SuppressWarnings("unchecked")
095        public <K, V> Collection<K> getKeys(WritableSlice<K, V> slice) {
096            assert slice.isCollective() : "Keys are not collected for slice " + slice;
097    
098            if (collectiveSliceKeys == null) return Collections.emptyList();
099            return (Collection<K>) collectiveSliceKeys.get(slice);
100        }
101    
102        @Override
103        public void forEach(@NotNull Function3<WritableSlice, Object, Object, Void> f) {
104            for (Map.Entry<Object, UserDataHolderImpl> entry : map.entrySet()) {
105                Object key = entry.getKey();
106                UserDataHolderImpl holder = entry.getValue();
107    
108                if (holder == null) continue;
109    
110                for (Key<?> sliceKey : holder.getKeys()) {
111                    Object value = holder.getUserData(sliceKey);
112    
113                    f.invoke(((AbstractWritableSlice) sliceKey).getSlice(), key, value);
114                }
115            }
116        }
117    
118        @NotNull
119        @Override
120        public <K, V> ImmutableMap<K, V> getSliceContents(@NotNull ReadOnlySlice<K, V> slice) {
121            ImmutableMap.Builder<K, V> builder = ImmutableMap.builder();
122    
123            for (Map.Entry<Object, UserDataHolderImpl> entry : map.entrySet()) {
124    
125                UserDataHolder holder = entry.getValue();
126    
127                V value = holder.getUserData(slice.getKey());
128    
129                if (value != null) {
130                    //noinspection unchecked
131                    builder.put((K) entry.getKey(), value);
132                }
133            }
134            return builder.build();
135        }
136    }