001    package org.jetbrains.jet.lang.resolve.calls.model;
002    
003    import com.google.common.collect.Maps;
004    import org.jetbrains.annotations.NotNull;
005    import org.jetbrains.jet.lang.psi.Call;
006    import org.jetbrains.jet.lang.psi.ValueArgument;
007    import org.jetbrains.jet.lang.resolve.calls.autocasts.DataFlowInfo;
008    
009    import java.util.Iterator;
010    import java.util.List;
011    import java.util.Map;
012    
013    public class DataFlowInfoForArgumentsImpl implements MutableDataFlowInfoForArguments {
014        private final Call call; //for better debug messages only
015        private final Map<ValueArgument, DataFlowInfo>  infoMap = Maps.newHashMap();
016        private final Map<ValueArgument, ValueArgument> nextArgument = Maps.newHashMap();
017        private DataFlowInfo initialInfo;
018        private DataFlowInfo resultInfo;
019    
020        public DataFlowInfoForArgumentsImpl(@NotNull Call call) {
021            this.call = call;
022            initNextArgMap(call.getValueArguments());
023        }
024    
025        private void initNextArgMap(@NotNull List<? extends ValueArgument> valueArguments) {
026            Iterator<? extends ValueArgument> iterator = valueArguments.iterator();
027            ValueArgument prev = null;
028            while (iterator.hasNext()) {
029                ValueArgument argument = iterator.next();
030                if (prev != null) {
031                    nextArgument.put(prev, argument);
032                }
033                prev = argument;
034            }
035        }
036    
037        @Override
038        public void setInitialDataFlowInfo(@NotNull DataFlowInfo dataFlowInfo) {
039            //TODO assert initialInfo == null
040            initialInfo = dataFlowInfo;
041        }
042    
043        @NotNull
044        @Override
045        public DataFlowInfo getInfo(@NotNull ValueArgument valueArgument) {
046            DataFlowInfo infoForArgument = infoMap.get(valueArgument);
047            if (infoForArgument == null) {
048                return initialInfo;
049            }
050            return initialInfo.and(infoForArgument);
051        }
052    
053        @Override
054        public void updateInfo(@NotNull ValueArgument valueArgument, @NotNull DataFlowInfo dataFlowInfo) {
055            ValueArgument next = nextArgument.get(valueArgument);
056            if (next != null) {
057                infoMap.put(next, dataFlowInfo);
058                return;
059            }
060            //TODO assert resultInfo == null
061            resultInfo = dataFlowInfo;
062        }
063    
064        @NotNull
065        @Override
066        public DataFlowInfo getResultInfo() {
067            if (resultInfo == null) return initialInfo;
068            assert initialInfo != null : "Initial data flow info was not set for call: " + call;
069            return initialInfo.and(resultInfo);
070        }
071    }