CallGraph.h

Go to the documentation of this file.

//===-- CallGraph.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/Analysis/CallGraph.h>
 
#include <llvm/ADT/GraphTraits.h>
#include <llvm/ADT/MapVector.h>
#include <llvm/ADT/PointerIntPair.h>
#include <llvm/ADT/SCCIterator.h>
#include <llvm/ADT/SetVector.h>
 
namespace mlir {
 
class Operation;
class CallOpInterface;
class SymbolTableCollection;
 
} // namespace mlir
 
namespace llzk {
 
namespace function {
class FuncDefOp;
}

class CallGraphNode {
public:
  class Edge {
    enum class Kind {
      // An 'Abstract' edge represents an opaque, non-operation, reference
      // between this node and the target. Edges of this type are only valid
      // from the external node (e.g., an external library call),
      // as there is no valid connection to an operation in the module.
      Abstract,
 
      // A 'Call' edge represents a direct reference to the target node via a
      // call-like operation within the callable region of this node.
      Call,
 
      // A 'Child' edge is used when the region of target node is defined inside
      // of the callable region of this node. This means that the region of this
      // node is an ancestor of the region for the target node. As such, this
      // edge cannot be used on the 'external' node.
      Child,
    };
 
  public:
    bool isAbstract() const { return targetAndKind.getInt() == Kind::Abstract; }

    bool isCall() const { return targetAndKind.getInt() == Kind::Call; }

    bool isChild() const { return targetAndKind.getInt() == Kind::Child; }

    CallGraphNode *getSource() const { return source; }

    CallGraphNode *getTarget() const { return targetAndKind.getPointer(); }
 
    bool operator==(const Edge &edge) const {
      return source == edge.source && targetAndKind == edge.targetAndKind;
    }
 
  private:
    Edge(CallGraphNode *src, CallGraphNode *target, Kind kind)
        : source(src), targetAndKind(target, kind) {}
    Edge(CallGraphNode *src, llvm::PointerIntPair<CallGraphNode *, 2, Kind> targetAndKind)
        : source(src), targetAndKind(targetAndKind) {}

    CallGraphNode *source;

    llvm::PointerIntPair<CallGraphNode *, 2, Kind> targetAndKind;
 
    // Provide access to the constructor and Kind.
    friend class CallGraphNode;
  };

  bool isExternal() const;

  mlir::Region *getCallableRegion() const;

  llzk::function::FuncDefOp getCalledFunction() const;

  void addAbstractEdge(CallGraphNode *node);

  void addCallEdge(CallGraphNode *node);

  void addChildEdge(CallGraphNode *child);

  using iterator = mlir::SmallVectorImpl<Edge>::const_iterator;
  iterator begin() const { return edges.begin(); }
  iterator end() const { return edges.end(); }
 
  llvm::iterator_range<iterator> edgesOut() const { return llvm::make_range(begin(), end()); }

  bool hasChildren() const;
 
private:
  struct EdgeKeyInfo {
    using SourceInfo = mlir::DenseMapInfo<CallGraphNode *>;
    using BaseInfo = mlir::DenseMapInfo<llvm::PointerIntPair<CallGraphNode *, 2, Edge::Kind>>;
 
    static Edge getEmptyKey() { return Edge(nullptr, BaseInfo::getEmptyKey()); }
    static Edge getTombstoneKey() { return Edge(nullptr, BaseInfo::getTombstoneKey()); }
    static unsigned getHashValue(const Edge &edge) {
      return SourceInfo::getHashValue(edge.source) ^ BaseInfo::getHashValue(edge.targetAndKind);
    }
    static bool isEqual(const Edge &lhs, const Edge &rhs) { return lhs == rhs; }
  };
 
  CallGraphNode(mlir::Region *callableRegion) : callableRegion(callableRegion) {}

  void addEdge(CallGraphNode *node, Edge::Kind kind);

  mlir::Region *callableRegion;

  mlir::SetVector<Edge, mlir::SmallVector<Edge, 4>, llvm::SmallDenseSet<Edge, 4, EdgeKeyInfo>>
      edges;
 
  // Provide access to private methods.
  friend class CallGraph;
};

class CallGraph {
  using NodeMapT = llvm::MapVector<mlir::Region *, std::unique_ptr<CallGraphNode>>;

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

  CallGraphNode *getOrAddNode(mlir::Region *region, CallGraphNode *parentNode);

  CallGraphNode *lookupNode(mlir::Region *region) const;

  CallGraphNode *getExternalCallerNode() const {
    return const_cast<CallGraphNode *>(&externalCallerNode);
  }

  CallGraphNode *getUnknownCalleeNode() const {
    return const_cast<CallGraphNode *>(&unknownCalleeNode);
  }

  CallGraphNode *
  resolveCallable(mlir::CallOpInterface call, mlir::SymbolTableCollection &symbolTable) const;

  void eraseNode(CallGraphNode *node);

  using iterator = NodeIterator;
  iterator begin() const { return nodes.begin(); }
  iterator end() const { return nodes.end(); }
 
  size_t size() const { return nodes.size(); }

  void dump() const;
  void print(llvm::raw_ostream &os) const;
 
private:
  NodeMapT nodes;

  CallGraphNode externalCallerNode;

  CallGraphNode unknownCalleeNode;
};
 
} // namespace llzk
 
namespace llvm {
// Provide graph traits for traversing call graphs using standard graph
// traversals.
template <> struct GraphTraits<const llzk::CallGraphNode *> {
  using NodeRef = const llzk::CallGraphNode *;
  static NodeRef getEntryNode(NodeRef node) { return node; }
 
  static NodeRef unwrap(const llzk::CallGraphNode::Edge &edge) { return edge.getTarget(); }
 
  // ChildIteratorType/begin/end - Allow iteration over all nodes in the graph.
  using ChildIteratorType = mapped_iterator<llzk::CallGraphNode::iterator, decltype(&unwrap)>;
  static ChildIteratorType child_begin(NodeRef node) { return {node->begin(), &unwrap}; }
  static ChildIteratorType child_end(NodeRef node) { return {node->end(), &unwrap}; }
};
 
template <>
struct GraphTraits<const llzk::CallGraph *> : public GraphTraits<const llzk::CallGraphNode *> {
  static NodeRef getEntryNode(const llzk::CallGraph *cg) { return cg->getExternalCallerNode(); }
 
  // nodes_iterator/begin/end - Allow iteration over all nodes in the graph
  using nodes_iterator = llzk::CallGraph::iterator;
  static nodes_iterator nodes_begin(llzk::CallGraph *cg) { return cg->begin(); }
  static nodes_iterator nodes_end(llzk::CallGraph *cg) { return cg->end(); }
};
} // namespace llvm