SymbolDefTree.h

Go to the documentation of this file.

//===-- SymbolDefTree.h -----------------------------------------*- C++ -*-===//
//
// Part of the LLZK Project, under the Apache License v2.0.
// See LICENSE.txt for license information.
// Copyright 2025 Veridise Inc.
// SPDX-License-Identifier: Apache-2.0
//
//===----------------------------------------------------------------------===//
 
#pragma once
 
#include <mlir/IR/SymbolTable.h>
 
#include <llvm/ADT/GraphTraits.h>
#include <llvm/ADT/MapVector.h>
#include <llvm/ADT/SetVector.h>
#include <llvm/Support/DOTGraphTraits.h>
 
namespace llzk {
 
class SymbolDefTreeNode {
  // The Symbol operation referenced by this node.
  mlir::SymbolOpInterface symbolDef;
 
  /* Tree structure. The SymbolDefTree owns the nodes so just pointers here. */
  SymbolDefTreeNode *parent;
  mlir::SetVector<SymbolDefTreeNode *> children;
 
  SymbolDefTreeNode(mlir::SymbolOpInterface symbolDefOp) : symbolDef(symbolDefOp), parent(nullptr) {
    assert(symbolDef && "must have a SymbolOpInterface node");
  }
 
  // Used only for creating the root node in the tree.
  SymbolDefTreeNode() : symbolDef(nullptr), parent(nullptr) {}

  void addChild(SymbolDefTreeNode *node);
 
  // Provide access to private members.
  friend class SymbolDefTree;
 
public:
  mlir::SymbolOpInterface getOp() const { return symbolDef; }

  const SymbolDefTreeNode *getParent() const { return parent; }

  bool hasChildren() const { return !children.empty(); }
  size_t numChildren() const { return children.size(); }

  using child_iterator = mlir::SetVector<SymbolDefTreeNode *>::const_iterator;
  child_iterator begin() const { return children.begin(); }
  child_iterator end() const { return children.end(); }

  inline llvm::iterator_range<child_iterator> childIter() const {
    return llvm::make_range(begin(), end());
  }

  std::string toString() const;
  void print(llvm::raw_ostream &os) const;
};

class SymbolDefTree {
  using NodeMapT = llvm::MapVector<mlir::SymbolOpInterface, std::unique_ptr<SymbolDefTreeNode>>;

  NodeMapT nodes;

  SymbolDefTreeNode root;

  class NodeIterator final
      : public llvm::mapped_iterator<
            NodeMapT::const_iterator, SymbolDefTreeNode *(*)(const NodeMapT::value_type &)> {
    static SymbolDefTreeNode *unwrap(const NodeMapT::value_type &value) {
      return value.second.get();
    }
 
  public:
    NodeIterator(NodeMapT::const_iterator it)
        : llvm::mapped_iterator<
              NodeMapT::const_iterator, SymbolDefTreeNode *(*)(const NodeMapT::value_type &)>(
              it, &unwrap
          ) {}
  };

  SymbolDefTreeNode *getOrAddNode(mlir::SymbolOpInterface symbolDef, SymbolDefTreeNode *parentNode);
 
  void buildTree(mlir::SymbolOpInterface symbolOp, SymbolDefTreeNode *parentNode);
 
public:
  SymbolDefTree(mlir::SymbolOpInterface root);

  const SymbolDefTreeNode *lookupNode(mlir::SymbolOpInterface symbolOp) const;

  const SymbolDefTreeNode *getRoot() const { return &root; }

  size_t size() const { return nodes.size(); }

  using iterator = NodeIterator;
  iterator begin() const { return nodes.begin(); }
  iterator end() const { return nodes.end(); }
 
  inline llvm::iterator_range<iterator> nodeIter() const {
    return llvm::make_range(begin(), end());
  }

  inline void dump() const { print(llvm::errs()); }
  void print(llvm::raw_ostream &os) const;

  void dumpToDotFile(std::string filename = "") const;
};
 
} // namespace llzk
 
namespace llvm {
// Provide graph traits for traversing SymbolDefTree using standard graph traversals.
 
template <> struct GraphTraits<const llzk::SymbolDefTreeNode *> {
  using NodeRef = const llzk::SymbolDefTreeNode *;
  static NodeRef getEntryNode(NodeRef node) { return node; }

  using ChildIteratorType = llzk::SymbolDefTreeNode::child_iterator;
  static ChildIteratorType child_begin(NodeRef node) { return node->begin(); }
  static ChildIteratorType child_end(NodeRef node) { return node->end(); }
};
 
template <>
struct GraphTraits<const llzk::SymbolDefTree *>
    : public GraphTraits<const llzk::SymbolDefTreeNode *> {
  using GraphType = const llzk::SymbolDefTree *;

  static NodeRef getEntryNode(GraphType g) { return g->getRoot(); }

  using nodes_iterator = llzk::SymbolDefTree::iterator;
  static nodes_iterator nodes_begin(GraphType g) { return g->begin(); }
  static nodes_iterator nodes_end(GraphType g) { return g->end(); }

  static unsigned size(GraphType g) { return g->size(); }
};
 
// Provide graph traits for printing SymbolDefTree using dot graph printer.
template <> struct DOTGraphTraits<const llzk::SymbolDefTreeNode *> : public DefaultDOTGraphTraits {
  using NodeRef = const llzk::SymbolDefTreeNode *;
  using GraphType = const llzk::SymbolDefTree *;
 
  DOTGraphTraits(bool isSimple = false) : DefaultDOTGraphTraits(isSimple) {}
 
  std::string getNodeLabel(NodeRef n, GraphType) { return n->toString(); }
};
 
template <>
struct DOTGraphTraits<const llzk::SymbolDefTree *>
    : public DOTGraphTraits<const llzk::SymbolDefTreeNode *> {
 
  DOTGraphTraits(bool isSimple = false) : DOTGraphTraits<NodeRef>(isSimple) {}
 
  static std::string getGraphName(GraphType) { return "Symbol Def Tree"; }
 
  std::string getNodeLabel(NodeRef n, GraphType g) {
    return DOTGraphTraits<NodeRef>::getNodeLabel(n, g);
  }
};
 
} // namespace llvm