From 0059690aad6c239cc9518bb86b5288e7c17cf13e Mon Sep 17 00:00:00 2001
From: Yulei Sui <7608399+yuleisui@users.noreply.github.com>
Date: Tue, 17 Sep 2024 10:33:10 +1000
Subject: [PATCH] add CallCFGEdge::getCallSite method

---
 svf/include/Graphs/ICFGEdge.h | 9 +++++++++
 svf/lib/Graphs/ICFG.cpp       | 7 +++++++
 2 files changed, 16 insertions(+)

diff --git a/svf/include/Graphs/ICFGEdge.h b/svf/include/Graphs/ICFGEdge.h
index d8c55a87f..2530e76d5 100644
--- a/svf/include/Graphs/ICFGEdge.h
+++ b/svf/include/Graphs/ICFGEdge.h
@@ -192,6 +192,12 @@ class CallCFGEdge : public ICFGEdge
     {
         callPEs.push_back(callPE);
     }
+    /// Return call ICFGNode at the callsite
+    inline const CallICFGNode* getCallSite() const
+    {
+        assert(SVFUtil::isa<CallICFGNode>(getSrcNode()) && "not a CallICFGNode?");
+        return SVFUtil::cast<CallICFGNode>(getSrcNode());
+    }
     /// Add get parameter edge to this CallCFGEdge
     inline const std::vector<const CallPE*>& getCallPEs() const
     {
@@ -243,6 +249,9 @@ class RetCFGEdge : public ICFGEdge
     {
         return retPE;
     }
+    /// Return call ICFGNode at the callsite
+    const CallICFGNode* getCallSite() const;
+
     /// Methods for support type inquiry through isa, cast, and dyn_cast:
     //@{
     static inline bool classof(const RetCFGEdge*)
diff --git a/svf/lib/Graphs/ICFG.cpp b/svf/lib/Graphs/ICFG.cpp
index 004b42277..d7c3284f1 100644
--- a/svf/lib/Graphs/ICFG.cpp
+++ b/svf/lib/Graphs/ICFG.cpp
@@ -192,6 +192,13 @@ const std::string RetCFGEdge::toString() const
     return rawstr.str();
 }
 
+/// Return call ICFGNode at the callsite
+const CallICFGNode* RetCFGEdge::getCallSite() const
+{
+    assert(SVFUtil::isa<RetICFGNode>(getDstNode()) && "not a RetICFGNode?");
+    return SVFUtil::cast<RetICFGNode>(getDstNode())->getCallICFGNode();
+}
+
 /*!
  * Constructor
  *  * Build ICFG