package accrue.analysis.interprocanalysis;

import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import polyglot.ast.Branch;
import polyglot.types.ClassType;
import polyglot.types.Type;
import polyglot.util.InternalCompilerError;

/**
 * Map giving the possible exit paths and the lattice elements along those paths
 * 
 * @param <T>
 *            type of context on each exit path
 */
public class ExitMap<T extends Ordered<T>> {

    /**
     * Map of exit key to context
     */
    protected final Map<Key, T> map;

    /**
     * Type of exceptions that can be thrown on exception exit
     */
    protected final Set<ClassType> exceptionTypes;

    /**
     * Returns a copy so that it is ok to modify the map.
     */
    protected ExitMap<T> copy() {
        return new ExitMap<T>(new LinkedHashMap<Key, T>(this.map),
                              new LinkedHashSet<ClassType>(this.exceptionTypes));
    }

    /**
     * Create a new map from the constituent parts
     * 
     * @param m
     *            map of exit type to context
     * @param exTypes
     *            Type of exceptions that can be thrown on exception exit
     * @return {@link ExitMap} class
     */
    protected ExitMap<T> create(Map<Key, T> m, Set<ClassType> exTypes) {
        return new ExitMap<T>(m, exTypes);
    }

    /**
     * Create a new map from the constituent parts
     * 
     * @param m
     *            map of exit type to context
     * @param exTypes
     *            Type of exceptions that can be thrown on exception exit
     */
    protected ExitMap(Map<Key, T> m, Set<ClassType> exTypes) {
        this.map = m;
        this.exceptionTypes = exTypes;
    }

    /**
     * @return True if the given exception type can be thrown
     */
    public boolean hasExceptionType(ClassType exType) {
        return this.exceptionTypes.contains(exType);
    }

    /**
     * @return All exceptions types that can be thrown
     */
    public Set<ClassType> exceptionTypes() {
        return exceptionTypes;
    }

    /**
     * @return the context for normal termination
     */
    public T normalTermination() {
        return map.get(KEY_NORM_TERM);
    }

    /**
     * Set the context for normal termination
     * 
     * @param c
     *            the normal termination context
     * @return An exit map with the new normal termination element
     */
    public ExitMap<T> normalTermination(T c) {
        return addKey(KEY_NORM_TERM, c);
    }

    /**
     * Get the context for the true branch if any, otherwise return the
     * context for normal termination
     */
    public T trueResult() {
        if (hasTrueFalseResult()) {
            return map.get(KEY_TRUE_RESULT);
        }
        else {
            return map.get(KEY_NORM_TERM);
        }
    }

    /**
     * Get the context for the false branch if any, otherwise return the
     * context for normal termination
     */
    public T falseResult() {
        if (hasTrueFalseResult()) {
            return map.get(KEY_FALSE_RESULT);
        }
        else {
            return map.get(KEY_NORM_TERM);
        }
    }

    /**
     * Set the contexts for the true and the false branches
     * 
     * @param trueC new context for the true branch
     * @param falseC new context for the false branch
     * @return Exit map with the branch contexts added
     */
    public ExitMap<T> trueFalseResult(T trueC, T falseC) {
        return addKey(KEY_TRUE_RESULT, trueC).addKey(KEY_FALSE_RESULT, falseC);
    }

    /**
     * Set the context for the true branch
     * 
     * @param trueC new context for the true branch
     * @return Exit map with the branch contexts added
     */
    public ExitMap<T> trueResult(T trueC) {
        return addKey(KEY_TRUE_RESULT, trueC);
    }

    /**
     * Set the context for the false branch
     * 
     * @param falseC new context for the false branch
     * @return Exit map with the branch contexts added
     */
    public ExitMap<T> falseResult(T falseC) {
        return addKey(KEY_FALSE_RESULT, falseC);
    }

    /**
     * True if the exit map branches on exit
     */
    public boolean hasTrueFalseResult() {
        return map.containsKey(KEY_TRUE_RESULT)
                || map.containsKey(KEY_FALSE_RESULT);
    }

    /*
     * True if the exit map has a false edge
     */
    public boolean hasTrueResult() {
        return map.containsKey(KEY_TRUE_RESULT);
    }

    /*
     * True if the exit map has a false edge
     */
    public boolean hasFalseResult() {
        return map.containsKey(KEY_FALSE_RESULT);
    }

    /** 
     * @return the context for a return statement
     */
    public T returnItem() {
        return map.get(KEY_RETURN);
    }

    /**
     * @param c
     *            the context for a return statement
     * @return Exit map with the return context set
     */
    public ExitMap<T> returnItem(T c) {
        return addKey(KEY_RETURN, c);
    }

    /**
     * Get the context for a <code>break</code> or a <code>continue</code>
     * 
     * @param kind
     *            <code>break</code> or <code>continue</code>
     * @return context for the given branch
     */
    public T branch(Branch.Kind kind) {
        return map.get(branchKey(kind));
    }

    /**
     * Get the context for a <code>break</code> or a <code>continue</code>
     * 
     * @param kind
     *            <code>break</code> or <code>continue</code>
     * @param loopLabel
     *            label of the loop
     * @return context for the given branch
     */
    public T branch(Branch.Kind kind, String loopLabel) {
        return map.get(branchKey(kind, loopLabel));
    }

    /**
     * Set the context for a <code>break</code> or <code>continue</code>
     * 
     * @param kind
     *            <code>break</code> or <code>continue</code>
     * @param c
     *            context for the kind
     * @return map with the added context
     */
    public ExitMap<T> branch(Branch.Kind kind, T c) {
        return addKey(branchKey(kind), c);
    }

    /**
     * Add the context, <code>c</code>, to the map with a diverge key
     * 
     * @param c
     *            context upon diverging
     * 
     * @return map with the added context
     */
    public ExitMap<T> diverge(T c) {
        return addKey(KEY_DIVERGE, c);
    }

    /**
     * Add the exception type with the context, <code>c</code>
     * 
     * @param exType
     *            type of exception
     * @param c
     *            context when exiting with an exception of <code>exType</code>
     * @return map with the added context
     */
    public ExitMap<T> exception(ClassType exType, T c) {
        return addKey(ExceptionKey.memo(exType), c);
    }

    /**
     * Get the context for the given exception type
     * 
     * @param exType
     *            type of exception
     * @return context when exiting with an exception of <code>exType</code>
     */
    public T exception(ClassType exType) {
        return map.get(ExceptionKey.memo(exType));
    }

    /**
     * Get the key for the given exception type
     * 
     * @param exType
     *            type of exception
     * @return exit key for the given exception type
     */
    public static ExceptionKey exceptionKey(ClassType exType) {
        return ExceptionKey.memo(exType);
    }

    /**
     * Functional update of an {@link ExitMap}. Returns a new {@link ExitMap} 
     * with the new value replacing the existing value in the map at the given 
     * key similar to {@link Map#put(Object, Object)}.  If <code>c</code> is 
     * null then the value at the given key will be removed from the returned 
     * map.
     * 
     * @param k
     *            key for the map
     * @param c
     *            context to add
     * @return map with the added context
     */
    public ExitMap<T> addKey(Key k, T c) {
        if (c == null) {
            return removeKey(k);
        }

        if (c.equals(this.map.get(k))) return this;

        ExitMap<T> f = this.copy();
        f.map.put(k, c);
        if (k instanceof ExceptionKey) {
            f.exceptionTypes.add(((ExceptionKey) k).exType);
        }
        return f;
    }

    /**
     * Functionally remove the key from the map
     * 
     * @param k
     *            key to remove
     * @return New map with the key removed
     */
    private ExitMap<T> removeKey(Key k) {
        if (this.map.containsKey(k)) {
            ExitMap<T> f = this.copy();
            f.map.remove(k);
            if (k instanceof ExceptionKey) {
                f.exceptionTypes.remove(((ExceptionKey) k).exType);
            }
            return f;
        }
        return this;
    }

    /**
     * Functionally remove the keys from the map
     * 
     * @param ks
     *            keys to remove
     * @return New map with the keys removed
     */
    private ExitMap<T> removeKeys(Key... ks) {
        ExitMap<T> f = this;
        for (Key k : ks) {
            f = f.removeKey(k);
        }
        return f;
    }

    /**
     * Functionally remove the normal termination, true, and false results from
     * the map
     */
    public ExitMap<T> removeNormalTermination() {
        // also remove the true and false results.
        return removeKeys(KEY_NORM_TERM, KEY_TRUE_RESULT, KEY_FALSE_RESULT);
    }

    /**
     * Functionally remove the true, and false results from the map
     */
    public ExitMap<T> removeTrueFalseResults() {
        return removeKeys(KEY_TRUE_RESULT, KEY_FALSE_RESULT);
    }

    /**
     * Functionally remove the return result from the map
     */
    public ExitMap<T> removeReturnItem() {
        return removeKey(KEY_RETURN);
    }

    /**
     * Functionally remove the <code>break</code> or <code>continue</code> from
     * the map
     */
    public ExitMap<T> removeBranch(Branch.Kind kind) {
        return removeKey(branchKey(kind));
    }

    /**
     * Functionally remove the <code>break</code> or <code>continue</code> with
     * the given label from the map
     */
    public ExitMap<T> removeBranch(Branch.Kind kind, String loopLabel) {
        return removeKey(branchKey(kind, loopLabel));
    }

    /**
     * Functionally remove the edge correspinding to the exception type
     */
    public ExitMap<T> removeException(ClassType exType) {
        return removeKey(ExceptionKey.memo(exType));
    }

    /**
     * Compute pairwise upper bound of <code>this</code> with <code>f</code>. If
     * a key only exists in one then it is added as-is
     * 
     * @param f
     *            map to upper bound with
     * @return new map that is the upper bound of <code>this</code> with
     *         <code>f</code>
     */
    public ExitMap<T> upperBound(ExitMap<T> f) {
        if (f == null) return this;

        Map<Key, T> newMap = new LinkedHashMap<Key, T>(this.map);
        Set<ClassType> exTypes =
                new LinkedHashSet<ClassType>(this.exceptionTypes);
        boolean changed = false;
        for (Key k : f.map.keySet()) {
            if (this.map.containsKey(k)) {
                T origItem = this.map.get(k);
                T newItem = origItem.upperBound(f.map.get(k));
                newMap.put(k, newItem);
                changed = changed || newItem != origItem;
            }
            else {
                newMap.put(k, f.map.get(k));
                if (k instanceof ExceptionKey) {
                    exTypes.add(((ExceptionKey) k).exType);
                }
                changed = true;
            }
        }

        if (!changed) return this;
        return this.create(newMap, exTypes);
    }

    /**
     * @return all the contexts from the map
     */
    public List<T> allItems() {
        return new ArrayList<T>(map.values());
    }

    /**
     * @return all the keys for the map
     */
    public Set<Key> allKeys() {
        return new LinkedHashSet<Key>(map.keySet());
    }

    /**
     * Upper bound <code>c</code> with the context for each key
     * 
     * @return New map with the upper bound <code>c</code> with the context for
     *         each key
     */
    public ExitMap<T> addToAll(T c) {
        Map<Key, T> newMap = new LinkedHashMap<Key, T>(this.map);
        boolean changed = false;
        for (Key k : this.map.keySet()) {
            T origItem = this.map.get(k);
            T newItem = origItem.upperBound(c);
            newMap.put(k, newItem);
            changed = changed || newItem != origItem;
        }

        if (!changed) return this;
        return new ExitMap<T>(newMap, this.exceptionTypes);
    }

    /**
     * Exit map key
     */
    public abstract static class Key {
        @Override
        public abstract boolean equals(Object o);

        @Override
        public abstract int hashCode();

        @Override
        public abstract String toString();
    }

    /**
     * Key for an exception exit path
     */
    public static class ExceptionKey extends Key {
        /**
         * Type of exception
         */
        public final ClassType exType;

        /**
         * Create a new key for the given type
         * @param exType exception type
         */
        private ExceptionKey(ClassType exType) {
            this.exType = exType;
        }

        /**
         * Memoized map of exception type to exit key
         */
        static Map<Type, ExceptionKey> memo =
                new LinkedHashMap<Type, ExceptionKey>();

        /**
         * Create and memoize the key for the given exception type
         * 
         * @param t
         *            exception type
         * @return key for the given exception type
         */
        static ExceptionKey memo(ClassType t) {
            if (!t.isSubtype(t.typeSystem().Throwable())) {
                throw new InternalCompilerError("Non-throwable type: " + t);
            }
            ExceptionKey ek = memo.get(t);
            if (ek == null) {
                ek = new ExceptionKey(t);
                memo.put(t, ek);
            }
            return ek;
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) return true;
            if (o instanceof ExceptionKey) {
                ExceptionKey that = (ExceptionKey) o;
                return this.exType.equals(that.exType);
            }
            return false;
        }

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

        @Override
        public String toString() {
            return exType.name();
        }
    }

    /**
     * Labeled <code>break</code> or <code>continue</code>
     */
    public static class LabeledBranchKey extends Key {
        /**
         * Label for the loop breaking or continuing
         */
        final String loopLabel;
        /**
         * <code>break</code> or <code>continue</code>
         */
        final Branch.Kind kind;

        /**
         * <code>break</code> or <code>continue</code> with a loop label
         * 
         * @param kind
         *            <code>break</code> or <code>continue</code>
         * @param loopLabel
         *            label of the loop we are breaking or continuing
         */
        private LabeledBranchKey(Branch.Kind kind, String loopLabel) {
            this.kind = kind;
            this.loopLabel = loopLabel;
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) return true;
            if (o instanceof LabeledBranchKey) {
                LabeledBranchKey that = (LabeledBranchKey) o;
                return this.kind == that.kind
                        && this.loopLabel.equals(that.loopLabel);
            }
            return false;
        }

        @Override
        public int hashCode() {
            return loopLabel.hashCode() ^ this.kind.hashCode();
        }

        @Override
        public String toString() {
            return this.kind.toString() + " " + loopLabel;
        }
    }

    /**
     * Key for normal termination
     */
    public static final Key KEY_NORM_TERM = new Key() {
        @Override
        public boolean equals(Object o) {
            return this == o;
        }

        @Override
        public int hashCode() {
            return 76952;
        }

        @Override
        public String toString() {
            return "KEY_NORM_TERM";
        }
    };

    /**
     * Key for exiting with a return statement
     */
    public static final Key KEY_RETURN = new Key() {
        @Override
        public boolean equals(Object o) {
            return this == o;
        }

        @Override
        public int hashCode() {
            return 3459827;
        }

        @Override
        public String toString() {
            return "KEY_RETURN";
        }
    };

    /**
     * Key for a loop break
     */
    public static final Key KEY_BREAK = new Key() {
        @Override
        public boolean equals(Object o) {
            return this == o;
        }

        @Override
        public int hashCode() {
            return 3457;
        }

        @Override
        public String toString() {
            return "KEY_BREAK";
        }
    };

    /**
     * Key for a loop continue
     */
    public static final Key KEY_CONTINUE = new Key() {
        @Override
        public boolean equals(Object o) {
            return this == o;
        }

        @Override
        public int hashCode() {
            return 3458;
        }

        @Override
        public String toString() {
            return "KEY_CONTINUE";
        }
    };

    /**
     * Key for divergence
     */
    public static final Key KEY_DIVERGE = new Key() {
        @Override
        public boolean equals(Object o) {
            return this == o;
        }

        @Override
        public int hashCode() {
            return 371;
        }

        @Override
        public String toString() {
            return "KEY_DIVERGE";
        }
    };
    /**
     * Key for a true branch
     */
    public static final Key KEY_TRUE_RESULT = new Key() {
        @Override
        public boolean equals(Object o) {
            return this == o;
        }

        @Override
        public int hashCode() {
            return 273;
        }

        @Override
        public String toString() {
            return "KEY_TRUE_RESULT";
        }
    };
    /**
     * Key for a false branch
     */
    public static final Key KEY_FALSE_RESULT = new Key() {
        @Override
        public boolean equals(Object o) {
            return this == o;
        }

        @Override
        public int hashCode() {
            return 274;
        }

        @Override
        public String toString() {
            return "KEY_FALSE_RESULT";
        }
    };

    @Override
    public String toString() {
        return "ItemMap [map=" + map + "]";
    }

    @Override
    public boolean equals(Object o) {
        if (o == this) return true;
        if (o instanceof ExitMap) {
            ExitMap<T> that = (ExitMap<T>) o;
            return this.map.equals(that.map)
                    && this.exceptionTypes.equals(that.exceptionTypes);
        }
        return false;
    }

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

    /**
     * Get the context for the given key
     */
    public T get(Key key) {
        return this.map.get(key);
    }

    /**
     * Get the key for continue or break
     */
    public static Key branchKey(Branch.Kind kind) {
        if (Branch.BREAK.equals(kind)) {
            return KEY_BREAK;
        }
        else if (Branch.CONTINUE.equals(kind)) {
            return KEY_CONTINUE;
        }
        throw new InternalCompilerError("Unexpected branch kind " + kind);
    }

    /**
     * Get the branch key (<code>break</code> or <code>continue</code>) for the
     * given kind and label
     * 
     * @param kind
     *            <code>break</code> or <code>continue</code>
     * @param loopLabel
     *            label for the loop we are breaking or continuing
     * @return the (possibly labeled) key
     */
    public static Key branchKey(Branch.Kind kind, String loopLabel) {
        if (loopLabel == null) {
            return branchKey(kind);
        }
        return new LabeledBranchKey(kind, loopLabel);
    }
}
