// Copyright (c) 2017 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef SOURCE_OPT_DOMINATOR_TREE_H_
#define SOURCE_OPT_DOMINATOR_TREE_H_

#include <algorithm>
#include <cstdint>
#include <map>
#include <utility>
#include <vector>

#include "source/opt/cfg.h"
#include "source/opt/tree_iterator.h"

namespace spvtools {
namespace opt {
// This helper struct forms the nodes in the tree, with each node containing its
// children. It also contains two values, for the pre and post indexes in the
// tree which are used to compare two nodes.
struct DominatorTreeNode {
  explicit DominatorTreeNode(BasicBlock* bb)
      : bb_(bb),
        parent_(nullptr),
        children_({}),
        dfs_num_pre_(-1),
        dfs_num_post_(-1) {}

  using iterator = std::vector<DominatorTreeNode*>::iterator;
  using const_iterator = std::vector<DominatorTreeNode*>::const_iterator;

  // depth first preorder iterator.
  using df_iterator = TreeDFIterator<DominatorTreeNode>;
  using const_df_iterator = TreeDFIterator<const DominatorTreeNode>;
  // depth first postorder iterator.
  using post_iterator = PostOrderTreeDFIterator<DominatorTreeNode>;
  using const_post_iterator = PostOrderTreeDFIterator<const DominatorTreeNode>;

  iterator begin() { return children_.begin(); }
  iterator end() { return children_.end(); }
  const_iterator begin() const { return cbegin(); }
  const_iterator end() const { return cend(); }
  const_iterator cbegin() const { return children_.begin(); }
  const_iterator cend() const { return children_.end(); }

  // Depth first preorder iterator using this node as root.
  df_iterator df_begin() { return df_iterator(this); }
  df_iterator df_end() { return df_iterator(); }
  const_df_iterator df_begin() const { return df_cbegin(); }
  const_df_iterator df_end() const { return df_cend(); }
  const_df_iterator df_cbegin() const { return const_df_iterator(this); }
  const_df_iterator df_cend() const { return const_df_iterator(); }

  // Depth first postorder iterator using this node as root.
  post_iterator post_begin() { return post_iterator::begin(this); }
  post_iterator post_end() { return post_iterator::end(nullptr); }
  const_post_iterator post_begin() const { return post_cbegin(); }
  const_post_iterator post_end() const { return post_cend(); }
  const_post_iterator post_cbegin() const {
    return const_post_iterator::begin(this);
  }
  const_post_iterator post_cend() const {
    return const_post_iterator::end(nullptr);
  }

  inline uint32_t id() const { return bb_->id(); }

  BasicBlock* bb_;
  DominatorTreeNode* parent_;
  std::vector<DominatorTreeNode*> children_;

  // These indexes are used to compare two given nodes. A node is a child or
  // grandchild of another node if its preorder index is greater than the
  // first nodes preorder index AND if its postorder index is less than the
  // first nodes postorder index.
  int dfs_num_pre_;
  int dfs_num_post_;
};

// A class representing a tree of BasicBlocks in a given function, where each
// node is dominated by its parent.
class DominatorTree {
 public:
  // Map OpLabel ids to dominator tree nodes
  using DominatorTreeNodeMap = std::map<uint32_t, DominatorTreeNode>;
  using iterator = TreeDFIterator<DominatorTreeNode>;
  using const_iterator = TreeDFIterator<const DominatorTreeNode>;
  using post_iterator = PostOrderTreeDFIterator<DominatorTreeNode>;
  using const_post_iterator = PostOrderTreeDFIterator<const DominatorTreeNode>;

  // List of DominatorTreeNode to define the list of roots
  using DominatorTreeNodeList = std::vector<DominatorTreeNode*>;
  using roots_iterator = DominatorTreeNodeList::iterator;
  using roots_const_iterator = DominatorTreeNodeList::const_iterator;

  DominatorTree() : postdominator_(false) {}
  explicit DominatorTree(bool post) : postdominator_(post) {}

  // Depth first iterators.
  // Traverse the dominator tree in a depth first pre-order.
  // The pseudo-block is ignored.
  iterator begin() { return ++iterator(GetRoot()); }
  iterator end() { return iterator(); }
  const_iterator begin() const { return cbegin(); }
  const_iterator end() const { return cend(); }
  const_iterator cbegin() const { return ++const_iterator(GetRoot()); }
  const_iterator cend() const { return const_iterator(); }

  // Traverse the dominator tree in a depth first post-order.
  // The pseudo-block is ignored.
  post_iterator post_begin() { return post_iterator::begin(GetRoot()); }
  post_iterator post_end() { return post_iterator::end(GetRoot()); }
  const_post_iterator post_begin() const { return post_cbegin(); }
  const_post_iterator post_end() const { return post_cend(); }
  const_post_iterator post_cbegin() const {
    return const_post_iterator::begin(GetRoot());
  }
  const_post_iterator post_cend() const {
    return const_post_iterator::end(GetRoot());
  }

  roots_iterator roots_begin() { return roots_.begin(); }
  roots_iterator roots_end() { return roots_.end(); }
  roots_const_iterator roots_begin() const { return roots_cbegin(); }
  roots_const_iterator roots_end() const { return roots_cend(); }
  roots_const_iterator roots_cbegin() const { return roots_.begin(); }
  roots_const_iterator roots_cend() const { return roots_.end(); }

  // Get the unique root of the tree.
  // It is guaranteed to work on a dominator tree.
  // post-dominator might have a list.
  DominatorTreeNode* GetRoot() {
    assert(roots_.size() == 1);
    return *roots_.begin();
  }

  const DominatorTreeNode* GetRoot() const {
    assert(roots_.size() == 1);
    return *roots_.begin();
  }

  const DominatorTreeNodeList& Roots() const { return roots_; }

  // Dumps the tree in the graphvis dot format into the |out_stream|.
  void DumpTreeAsDot(std::ostream& out_stream) const;

  // Build the (post-)dominator tree for the given control flow graph
  // |cfg| and the function |f|. |f| must exist in the |cfg|. Any
  // existing data in the dominator tree will be overwritten
  void InitializeTree(const CFG& cfg, const Function* f);

  // Check if the basic block |a| dominates the basic block |b|.
  bool Dominates(const BasicBlock* a, const BasicBlock* b) const;

  // Check if the basic block id |a| dominates the basic block id |b|.
  bool Dominates(uint32_t a, uint32_t b) const;

  // Check if the dominator tree node |a| dominates the dominator tree node |b|.
  bool Dominates(const DominatorTreeNode* a, const DominatorTreeNode* b) const;

  // Check if the basic block |a| strictly dominates the basic block |b|.
  bool StrictlyDominates(const BasicBlock* a, const BasicBlock* b) const;

  // Check if the basic block id |a| strictly dominates the basic block id |b|.
  bool StrictlyDominates(uint32_t a, uint32_t b) const;

  // Check if the dominator tree node |a| strictly dominates the dominator tree
  // node |b|.
  bool StrictlyDominates(const DominatorTreeNode* a,
                         const DominatorTreeNode* b) const;

  // Returns the immediate dominator of basic block |a|.
  BasicBlock* ImmediateDominator(const BasicBlock* A) const;

  // Returns the immediate dominator of basic block id |a|.
  BasicBlock* ImmediateDominator(uint32_t a) const;

  // Returns true if the basic block |a| is reachable by this tree. A node would
  // be unreachable if it cannot be reached by traversal from the start node or
  // for a postdominator tree, cannot be reached from the exit nodes.
  inline bool ReachableFromRoots(const BasicBlock* a) const {
    if (!a) return false;
    return ReachableFromRoots(a->id());
  }

  // Returns true if the basic block id |a| is reachable by this tree.
  bool ReachableFromRoots(uint32_t a) const {
    return GetTreeNode(a) != nullptr;
  }

  // Returns true if this tree is a post dominator tree.
  bool IsPostDominator() const { return postdominator_; }

  // Clean up the tree.
  void ClearTree() {
    nodes_.clear();
    roots_.clear();
  }

  // Applies the std::function |func| to all nodes in the dominator tree.
  // Tree nodes are visited in a depth first pre-order.
  bool Visit(std::function<bool(DominatorTreeNode*)> func) {
    for (auto n : *this) {
      if (!func(&n)) return false;
    }
    return true;
  }

  // Applies the std::function |func| to all nodes in the dominator tree.
  // Tree nodes are visited in a depth first pre-order.
  bool Visit(std::function<bool(const DominatorTreeNode*)> func) const {
    for (auto n : *this) {
      if (!func(&n)) return false;
    }
    return true;
  }

  // Applies the std::function |func| to all nodes in the dominator tree from
  // |node| downwards. The boolean return from |func| is used to determine
  // whether or not the children should also be traversed. Tree nodes are
  // visited in a depth first pre-order.
  void VisitChildrenIf(std::function<bool(DominatorTreeNode*)> func,
                       iterator node) {
    if (func(&*node)) {
      for (auto n : *node) {
        VisitChildrenIf(func, n->df_begin());
      }
    }
  }

  // Returns the DominatorTreeNode associated with the basic block |bb|.
  // If the |bb| is unknown to the dominator tree, it returns null.
  inline DominatorTreeNode* GetTreeNode(BasicBlock* bb) {
    return GetTreeNode(bb->id());
  }
  // Returns the DominatorTreeNode associated with the basic block |bb|.
  // If the |bb| is unknown to the dominator tree, it returns null.
  inline const DominatorTreeNode* GetTreeNode(BasicBlock* bb) const {
    return GetTreeNode(bb->id());
  }

  // Returns the DominatorTreeNode associated with the basic block id |id|.
  // If the id |id| is unknown to the dominator tree, it returns null.
  inline DominatorTreeNode* GetTreeNode(uint32_t id) {
    DominatorTreeNodeMap::iterator node_iter = nodes_.find(id);
    if (node_iter == nodes_.end()) {
      return nullptr;
    }
    return &node_iter->second;
  }
  // Returns the DominatorTreeNode associated with the basic block id |id|.
  // If the id |id| is unknown to the dominator tree, it returns null.
  inline const DominatorTreeNode* GetTreeNode(uint32_t id) const {
    DominatorTreeNodeMap::const_iterator node_iter = nodes_.find(id);
    if (node_iter == nodes_.end()) {
      return nullptr;
    }
    return &node_iter->second;
  }

  // Adds the basic block |bb| to the tree structure if it doesn't already
  // exist.
  DominatorTreeNode* GetOrInsertNode(BasicBlock* bb);

  // Recomputes the DF numbering of the tree.
  void ResetDFNumbering();

 private:
  // Wrapper function which gets the list of pairs of each BasicBlocks to its
  // immediately  dominating BasicBlock and stores the result in the the edges
  // parameter.
  //
  // The |edges| vector will contain the dominator tree as pairs of nodes.
  // The first node in the pair is a node in the graph. The second node in the
  // pair is its immediate dominator.
  // The root of the tree has themself as immediate dominator.
  void GetDominatorEdges(
      const Function* f, const BasicBlock* dummy_start_node,
      std::vector<std::pair<BasicBlock*, BasicBlock*>>* edges);

  // The roots of the tree.
  std::vector<DominatorTreeNode*> roots_;

  // Pairs each basic block id to the tree node containing that basic block.
  DominatorTreeNodeMap nodes_;

  // True if this is a post dominator tree.
  bool postdominator_;
};

}  // namespace opt
}  // namespace spvtools

#endif  // SOURCE_OPT_DOMINATOR_TREE_H_