/*
 * Copyright (c) 2002 - 2006 IBM Corporation.
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the Eclipse Public License v1.0
 * which accompanies this distribution, and is available at
 * http://www.eclipse.org/legal/epl-v10.html
 *
 * Contributors:
 *     IBM Corporation - initial API and implementation
 */
package com.ibm.wala.cast.tree.rewrite;

import com.ibm.wala.cast.tree.CAst;
import com.ibm.wala.cast.tree.CAstControlFlowMap;
import com.ibm.wala.cast.tree.CAstEntity;
import com.ibm.wala.cast.tree.CAstNode;
import com.ibm.wala.cast.tree.CAstNodeTypeMap;
import com.ibm.wala.cast.tree.CAstSourcePositionMap;
import com.ibm.wala.cast.tree.impl.CAstControlFlowRecorder;
import com.ibm.wala.cast.tree.impl.CAstNodeTypeMapRecorder;
import com.ibm.wala.cast.tree.impl.CAstSourcePositionRecorder;
import com.ibm.wala.cast.tree.impl.DelegatingEntity;
import com.ibm.wala.cast.util.CAstPrinter;
import com.ibm.wala.util.collections.EmptyIterator;
import com.ibm.wala.util.collections.HashMapFactory;
import com.ibm.wala.util.collections.HashSetFactory;
import com.ibm.wala.util.collections.Pair;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;

/**
 * Abstract superclass for types performing a rewrite operation on a CAst. The CAst is not mutated;
 * instead, a new CAst is created which delegates to the original CAst where no transformation was
 * performed.
 *
 * @param <C> type of the RewriteContext used when traversing the original CAst during the rewrite
 *     operation
 * @param <K> a key used to ease cloning of partial ASTs. When rewriting an AST, sub-classes
 *     maintain a mapping from (original node, key) pairs (where key is of type K) to new nodes; see
 *     {@link #copyNodes}
 */
public abstract class CAstRewriter<
    C extends CAstRewriter.RewriteContext<K>, K extends CAstRewriter.CopyKey<K>> {

  protected static final boolean DEBUG = false;

  /** interface to be implemented by keys used for cloning sub-trees during the rewrite */
  public interface CopyKey<Self extends CopyKey<Self>> {

    @Override
    int hashCode();

    @Override
    boolean equals(Object o);

    /**
     * keys have parent pointers, useful for when nesting cloning must occur (e.g., unrolling of
     * nested loops)
     */
    Self parent();
  }

  /** interface to be implemented by contexts used while traversing the AST */
  public interface RewriteContext<K extends CopyKey<K>> {

    /** get the cloning key for this context */
    K key();
  }

  /** represents a rewritten CAst */
  public interface Rewrite {

    CAstNode newRoot();

    CAstControlFlowMap newCfg();

    CAstSourcePositionMap newPos();

    CAstNodeTypeMap newTypes();

    Map<CAstNode, Collection<CAstEntity>> newChildren();

    CAstNode[] newDefaults();
  }

  protected final CAst Ast;

  /**
   * for CAstEntity nodes r s.t. r.getAst() == null, should the scoped entities of r be rewritten?
   */
  protected final boolean recursive;

  protected final C rootContext;

  public CAstRewriter(CAst Ast, boolean recursive, C rootContext) {
    this.Ast = Ast;
    this.recursive = recursive;
    this.rootContext = rootContext;
  }

  /**
   * rewrite the CAst rooted at root under some context, returning the node at the root of the
   * rewritten tree. mutate nodeMap in the process, indicating how (original node, copy key) pairs
   * are mapped to nodes in the rewritten tree.
   */
  protected abstract CAstNode copyNodes(
      CAstNode root,
      final CAstControlFlowMap cfg,
      C context,
      Map<Pair<CAstNode, K>, CAstNode> nodeMap);

  protected CAstNode copySubtreesIntoNewNode(
      CAstNode n, CAstControlFlowMap cfg, C c, Map<Pair<CAstNode, K>, CAstNode> nodeMap) {
    return copySubtreesIntoNewNode(n, cfg, c, nodeMap, Pair.make(n, c.key()));
  }

  protected CAstNode copySubtreesIntoNewNode(
      CAstNode n,
      CAstControlFlowMap cfg,
      C c,
      Map<Pair<CAstNode, K>, CAstNode> nodeMap,
      Pair<CAstNode, K> pairKey) {
    final List<CAstNode> newChildren = copyChildrenArray(n, cfg, c, nodeMap);
    CAstNode newN = Ast.makeNode(n.getKind(), newChildren);
    assert !nodeMap.containsKey(pairKey);
    nodeMap.put(pairKey, newN);
    return newN;
  }

  protected List<CAstNode> copyChildrenArray(
      CAstNode n, CAstControlFlowMap cfg, C context, Map<Pair<CAstNode, K>, CAstNode> nodeMap) {
    List<CAstNode> newChildren = new ArrayList<>(n.getChildCount());
    for (CAstNode child : n.getChildren()) {
      newChildren.add(copyNodes(child, cfg, context, nodeMap));
    }
    return newChildren;
  }

  protected List<CAstNode> copyChildrenArrayAndTargets(
      CAstNode n, CAstControlFlowMap cfg, C context, Map<Pair<CAstNode, K>, CAstNode> nodeMap) {
    final List<CAstNode> children = copyChildrenArray(n, cfg, context, nodeMap);
    if (cfg != null) {
      final Collection<Object> targetLabels = cfg.getTargetLabels(n);
      if (targetLabels != null)
        for (Object label : targetLabels)
          if (label instanceof CAstNode) copyNodes((CAstNode) label, cfg, context, nodeMap);
    }
    return children;
  }

  /**
   * in {@link #copyFlow(Map, CAstControlFlowMap, CAstSourcePositionMap)}, if the source of some
   * original CFG edge is replicated, but we find no replica for the target, what node should be the
   * target of the CFG edge in the rewritten AST? By default, just uses the original target.
   */
  @SuppressWarnings("unused")
  protected CAstNode flowOutTo(
      Map<Pair<CAstNode, K>, CAstNode> nodeMap,
      CAstNode oldSource,
      Object label,
      CAstNode oldTarget,
      CAstControlFlowMap orig,
      CAstSourcePositionMap src) {
    return oldTarget;
  }

  /**
   * create a control-flow map for the rewritten tree, given the mapping from (original node, copy
   * key) pairs ot new nodes and the original control-flow map.
   */
  protected CAstControlFlowMap copyFlow(
      Map<Pair<CAstNode, K>, CAstNode> nodeMap,
      CAstControlFlowMap orig,
      CAstSourcePositionMap newSrc) {

    // the new control-flow map
    final CAstControlFlowRecorder newMap = new CAstControlFlowRecorder(newSrc);
    // tracks which CAstNodes not present in nodeMap's key set (under any copy
    // key) are added as targets of CFG edges
    // via a call to flowOutTo() (see below); used to ensure these nodes are
    // only mapped to themselves once in newMap
    final Set<CAstNode> mappedOutsideNodes = HashSetFactory.make(1);
    // all edge targets in new control-flow map; must all be mapped to
    // themselves
    Set<CAstNode> allNewTargetNodes = HashSetFactory.make(1);
    Collection<CAstNode> oldSources = orig.getMappedNodes();

    for (Entry<Pair<CAstNode, K>, CAstNode> entry : nodeMap.entrySet()) {
      Pair<CAstNode, K> N = entry.getKey();
      CAstNode oldSource = N.fst;
      K key = N.snd;

      CAstNode newSource = entry.getValue();
      assert newSource != null;

      newMap.map(newSource, newSource);

      if (DEBUG) {
        System.err.println(("\n\nlooking at " + key + ':' + CAstPrinter.print(oldSource)));
      }

      if (oldSources.contains(oldSource)) {
        // if (orig.getTarget(oldSource, null) != null) {
        //  LS = IteratorPlusOne.make(LS, null);
        // }

        for (Object origLabel : orig.getTargetLabels(oldSource)) {
          CAstNode oldTarget = orig.getTarget(oldSource, origLabel);
          assert oldTarget != null;

          if (DEBUG) {
            System.err.println(("old: " + origLabel + " --> " + CAstPrinter.print(oldTarget)));
          }

          // try to find a k in key's parent chain such that (oldTarget, k) is
          // in nodeMap's key set
          Pair<CAstNode, CopyKey<K>> targetKey;
          CopyKey<K> k = key;
          do {
            targetKey = Pair.make(oldTarget, k);
            if (k != null) {
              k = k.parent();
            } else {
              break;
            }
          } while (!nodeMap.containsKey(targetKey));

          Object newLabel;
          if (nodeMap.containsKey(Pair.make(origLabel, targetKey.snd))) { // label
            // is
            // mapped
            // too
            newLabel = nodeMap.get(Pair.make(origLabel, targetKey.snd));
          } else {
            newLabel = origLabel;
          }

          CAstNode newTarget;
          if (nodeMap.containsKey(targetKey)) {
            newTarget = nodeMap.get(targetKey);
            newMap.add(newSource, newTarget, newLabel);
            allNewTargetNodes.add(newTarget);

          } else {
            // could not discover target of CFG edge in nodeMap under any key related to the current
            // source key.
            // the edge might have been deleted, or it may end at a node above the root where we
            // were
            // rewriting
            // ask flowOutTo() to just choose a target
            newTarget = flowOutTo(nodeMap, oldSource, origLabel, oldTarget, orig, newSrc);
            allNewTargetNodes.add(newTarget);
            newMap.add(newSource, newTarget, newLabel);
            if (newTarget != CAstControlFlowMap.EXCEPTION_TO_EXIT
                && !mappedOutsideNodes.contains(newTarget)) {
              mappedOutsideNodes.add(newTarget);
              newMap.map(newTarget, newTarget);
            }
          }

          if (DEBUG) {
            System.err.println(
                ("mapping:old: "
                    + CAstPrinter.print(oldSource)
                    + "-- "
                    + origLabel
                    + " --> "
                    + CAstPrinter.print(oldTarget)));
            System.err.println(
                ("mapping:new: "
                    + CAstPrinter.print(newSource)
                    + "-- "
                    + newLabel
                    + " --> "
                    + CAstPrinter.print(newTarget)));
          }
        }
      }
    }

    allNewTargetNodes.removeAll(newMap.getMappedNodes());
    for (CAstNode newTarget : allNewTargetNodes) {
      if (newTarget != CAstControlFlowMap.EXCEPTION_TO_EXIT) {
        newMap.map(newTarget, newTarget);
      }
    }

    assert !oldNodesInNewMap(nodeMap, newMap);

    return newMap;
  }

  // check whether newMap contains any CFG edges involving nodes in the domain of nodeMap
  private boolean oldNodesInNewMap(
      Map<Pair<CAstNode, K>, CAstNode> nodeMap, final CAstControlFlowRecorder newMap) {
    HashSet<CAstNode> oldNodes = HashSetFactory.make();
    for (Entry<Pair<CAstNode, K>, CAstNode> e : nodeMap.entrySet()) oldNodes.add(e.getKey().fst);
    for (CAstNode mappedNode : newMap.getMappedNodes()) {
      if (oldNodes.contains(mappedNode)) return true;
      for (Object lbl : newMap.getTargetLabels(mappedNode))
        if (oldNodes.contains(newMap.getTarget(mappedNode, lbl))) return true;
    }
    return false;
  }

  protected CAstSourcePositionMap copySource(
      Map<Pair<CAstNode, K>, CAstNode> nodeMap, CAstSourcePositionMap orig) {
    CAstSourcePositionRecorder newMap = new CAstSourcePositionRecorder();
    for (Entry<Pair<CAstNode, K>, CAstNode> entry : nodeMap.entrySet()) {
      Pair<CAstNode, K> N = entry.getKey();
      CAstNode oldNode = N.fst;

      CAstNode newNode = entry.getValue();

      if (orig.getPosition(oldNode) != null) {
        newMap.setPosition(newNode, orig.getPosition(oldNode));
      }
    }

    return newMap;
  }

  protected CAstNodeTypeMap copyTypes(
      Map<Pair<CAstNode, K>, CAstNode> nodeMap, CAstNodeTypeMap orig) {
    if (orig != null) {
      CAstNodeTypeMapRecorder newMap = new CAstNodeTypeMapRecorder();
      for (Entry<Pair<CAstNode, K>, CAstNode> entry : nodeMap.entrySet()) {
        Pair<CAstNode, K> N = entry.getKey();
        CAstNode oldNode = N.fst;

        CAstNode newNode = entry.getValue();

        if (orig.getNodeType(oldNode) != null) {
          newMap.add(newNode, orig.getNodeType(oldNode));
        }
      }

      return newMap;
    } else {
      return null;
    }
  }

  protected Map<CAstNode, Collection<CAstEntity>> copyChildren(
      @SuppressWarnings("unused") CAstNode root,
      Map<Pair<CAstNode, K>, CAstNode> nodeMap,
      Map<CAstNode, Collection<CAstEntity>> children) {
    final Map<CAstNode, Collection<CAstEntity>> newChildren = new LinkedHashMap<>();

    for (Entry<Pair<CAstNode, K>, CAstNode> entry : nodeMap.entrySet()) {
      Pair<CAstNode, K> N = entry.getKey();
      CAstNode oldNode = N.fst;

      CAstNode newNode = entry.getValue();

      if (children.containsKey(oldNode)) {
        Set<CAstEntity> newEntities = new LinkedHashSet<>();
        newChildren.put(newNode, newEntities);
        for (CAstEntity cAstEntity : children.get(oldNode)) {
          newEntities.add(rewrite(cAstEntity));
        }
      }
    }

    for (Entry<CAstNode, Collection<CAstEntity>> entry : children.entrySet()) {
      CAstNode key = entry.getKey();
      if (key == null) {
        Set<CAstEntity> newEntities = new LinkedHashSet<>();
        newChildren.put(key, newEntities);
        for (CAstEntity oldEntity : entry.getValue()) {
          newEntities.add(rewrite(oldEntity));
        }
      }
    }

    return newChildren;
  }

  /** rewrite the CAst sub-tree rooted at root */
  public Rewrite rewrite(
      final CAstNode root,
      final CAstControlFlowMap cfg,
      final CAstSourcePositionMap pos,
      final CAstNodeTypeMap types,
      final Map<CAstNode, Collection<CAstEntity>> children,
      final CAstNode[] defaults) {
    final Map<Pair<CAstNode, K>, CAstNode> nodes = HashMapFactory.make();
    final CAstNode newRoot = copyNodes(root, cfg, rootContext, nodes);
    final CAstNode newDefaults[] = new CAstNode[defaults == null ? 0 : defaults.length];
    for (int i = 0; i < newDefaults.length; i++) {
      newDefaults[i] = copyNodes(defaults[i], cfg, rootContext, nodes);
    }

    return new Rewrite() {
      private CAstControlFlowMap theCfg = null;

      private CAstSourcePositionMap theSource = null;

      private CAstNodeTypeMap theTypes = null;

      private Map<CAstNode, Collection<CAstEntity>> theChildren = null;

      @Override
      public CAstNode[] newDefaults() {
        return newDefaults;
      }

      @Override
      public CAstNode newRoot() {
        return newRoot;
      }

      @Override
      public CAstControlFlowMap newCfg() {
        if (theCfg == null && cfg != null) theCfg = copyFlow(nodes, cfg, newPos());
        return theCfg;
      }

      @Override
      public CAstSourcePositionMap newPos() {
        if (theSource == null && pos != null) theSource = copySource(nodes, pos);
        return theSource;
      }

      @Override
      public CAstNodeTypeMap newTypes() {
        if (theTypes == null && types != null) theTypes = copyTypes(nodes, types);
        return theTypes;
      }

      @Override
      public Map<CAstNode, Collection<CAstEntity>> newChildren() {
        if (theChildren == null) theChildren = copyChildren(root, nodes, children);
        return theChildren;
      }
    };
  }

  /**
   * perform the rewrite on a {@link CAstEntity}, returning the new {@link CAstEntity} as the result
   */
  public CAstEntity rewrite(final CAstEntity root) {

    if (root.getAST() != null) {
      final Rewrite rewrite =
          rewrite(
              root.getAST(),
              root.getControlFlow(),
              root.getSourceMap(),
              root.getNodeTypeMap(),
              root.getAllScopedEntities(),
              root.getArgumentDefaults());

      return new DelegatingEntity(root) {
        @Override
        public String toString() {
          return root.toString() + " (clone)";
        }

        @Override
        public Iterator<CAstEntity> getScopedEntities(CAstNode construct) {
          Map<CAstNode, Collection<CAstEntity>> newChildren = getAllScopedEntities();
          if (newChildren.containsKey(construct)) {
            return newChildren.get(construct).iterator();
          } else {
            return EmptyIterator.instance();
          }
        }

        @Override
        public Map<CAstNode, Collection<CAstEntity>> getAllScopedEntities() {
          return rewrite.newChildren();
        }

        @Override
        public CAstNode getAST() {
          return rewrite.newRoot();
        }

        @Override
        public CAstNodeTypeMap getNodeTypeMap() {
          return rewrite.newTypes();
        }

        @Override
        public CAstSourcePositionMap getSourceMap() {
          return rewrite.newPos();
        }

        @Override
        public CAstControlFlowMap getControlFlow() {
          return rewrite.newCfg();
        }

        @Override
        public CAstNode[] getArgumentDefaults() {
          return rewrite.newDefaults();
        }
      };

    } else if (recursive) {

      Map<CAstNode, Collection<CAstEntity>> children = root.getAllScopedEntities();
      final Map<CAstNode, Collection<CAstEntity>> newChildren = new LinkedHashMap<>();
      for (Entry<CAstNode, Collection<CAstEntity>> entry : children.entrySet()) {
        CAstNode key = entry.getKey();
        Set<CAstEntity> newValues = new LinkedHashSet<>();
        newChildren.put(key, newValues);
        for (CAstEntity entity : entry.getValue()) {
          newValues.add(rewrite(entity));
        }
      }

      return new DelegatingEntity(root) {
        @Override
        public String toString() {
          return root.toString() + " (clone)";
        }

        @Override
        public Iterator<CAstEntity> getScopedEntities(CAstNode construct) {
          if (newChildren.containsKey(construct)) {
            return newChildren.get(construct).iterator();
          } else {
            return EmptyIterator.instance();
          }
        }

        @Override
        public Map<CAstNode, Collection<CAstEntity>> getAllScopedEntities() {
          return newChildren;
        }
      };

    } else {
      return root;
    }
  }
}
