SVF-tools · yuleisui · Jan 17, 2024 · Jan 16, 2024 · Jan 16, 2024 · Jan 16, 2024
diff --git a/svf-llvm/include/SVF-LLVM/SVFIRBuilder.h b/svf-llvm/include/SVF-LLVM/SVFIRBuilder.h
@@ -298,6 +298,56 @@ class SVFIRBuilder: public llvm::InstVisitor<SVFIRBuilder>
         }
         return nullptr;
     }
+
+    /// Add Address edge from allocinst with arraysize like "%4 = alloca i8, i64 3"
+    inline AddrStmt* addAddrWithStackArraySz(NodeID src, NodeID dst, llvm::AllocaInst& inst) {
+        AddrStmt* edge = addAddrEdge(src, dst);
+        if (inst.getArraySize()) {
+            SVFValue* arrSz = LLVMModuleSet::getLLVMModuleSet()->getSVFValue(inst.getArraySize());
+            edge->addArrSize(arrSz);
+        }
+        return edge;
+    }
+
+    /// Add Address edge from ext call with args like "%5 = call i8* @malloc(i64 noundef 5)"
+    inline AddrStmt* addAddrWithHeapSz(NodeID src, NodeID dst, const CallBase* cs) {
+        // get name of called function
+        AddrStmt* edge = addAddrEdge(src, dst);
+
+        llvm::Function* calledFunc = cs->getCalledFunction();
+        std::string functionName;
+        if (calledFunc) {
+            functionName = calledFunc->getName().str();
+        }
+        else
+        {
+            SVFUtil::wrnMsg("not support indirect call to add AddrStmt.\n");
+        }
+        if (functionName == "malloc") {
+            if (cs->arg_size() > 0) {
+                const llvm::Value* val = cs->getArgOperand(0);
+                SVFValue* svfval = LLVMModuleSet::getLLVMModuleSet()->getSVFValue(val);
+                edge->addArrSize(svfval);
+            }
+        }
+        // Check if the function called is 'calloc' and process its arguments.
+        // e.g. "%5 = call i8* @calloc(1, 8)", edge should add two SVFValue (1 and 8)
+        else if (functionName == "calloc") {
+            if (cs->arg_size() > 1) {
+                edge->addArrSize(LLVMModuleSet::getLLVMModuleSet()->getSVFValue(cs->getArgOperand(0)));
+                edge->addArrSize(LLVMModuleSet::getLLVMModuleSet()->getSVFValue(cs->getArgOperand(1)));
+            }
+        }
+        else {
+            if (cs->arg_size() > 0) {
+                const llvm::Value* val = cs->getArgOperand(0);
+                SVFValue* svfval = LLVMModuleSet::getLLVMModuleSet()->getSVFValue(val);
+                edge->addArrSize(svfval);
+            }
+        }
+        return edge;
+    }
+
     /// Add Copy edge
     inline CopyStmt* addCopyEdge(NodeID src, NodeID dst)
     {

diff --git a/svf-llvm/lib/SVFIRBuilder.cpp b/svf-llvm/lib/SVFIRBuilder.cpp
@@ -636,7 +636,7 @@ void SVFIRBuilder::visitAllocaInst(AllocaInst &inst)
 
     NodeID src = getObjectNode(&inst);
 
-    addAddrEdge(src, dst);
+    addAddrWithStackArraySz(src, dst, inst);
 
 }
 

diff --git a/svf-llvm/lib/SVFIRExtAPI.cpp b/svf-llvm/lib/SVFIRExtAPI.cpp
@@ -133,7 +133,7 @@ void SVFIRBuilder::handleExtCall(const CallBase* cs, const SVFFunction* svfCalle
     {
         NodeID val = pag->getValueNode(svfInst);
         NodeID obj = pag->getObjectNode(svfInst);
-        addAddrEdge(obj, val);
+        addAddrWithHeapSz(obj, val, cs);
     }
     else if (isHeapAllocExtCallViaArg(svfCall))
     {
@@ -146,7 +146,7 @@ void SVFIRBuilder::handleExtCall(const CallBase* cs, const SVFFunction* svfCalle
             NodeID obj = pag->addDummyObjNode(arg->getType());
             if (vnArg && dummy && obj)
             {
-                addAddrEdge(obj, dummy);
+                addAddrWithHeapSz(obj, dummy, cs);
                 addStoreEdge(dummy, vnArg);
             }
         }

diff --git a/svf/include/SVFIR/SVFStatements.h b/svf/include/SVFIR/SVFStatements.h
@@ -321,6 +321,8 @@ class AddrStmt: public AssignStmt
     AddrStmt(const AddrStmt&);       ///< place holder
     void operator=(const AddrStmt&); ///< place holder
 
+    std::vector<SVFValue*> arrSize;	///< Array size of the allocated memory
+
 public:
     /// Methods for support type inquiry through isa, cast, and dyn_cast:
     //@{
@@ -343,6 +345,15 @@ class AddrStmt: public AssignStmt
 
     virtual const std::string toString() const override;
 
+    inline void addArrSize(SVFValue* size) { //TODO:addSizeVar
+        arrSize.push_back(size);
+    }
+
+    ///< get array size of the allocated memory
+    inline const std::vector<SVFValue*>& getArrSize() const { //TODO:getSizeVars
+        return arrSize;
+    }
+
 };
 
 /*!
-Original file line number
+Diff line change
@@ Expand Up / @@ -636,7 +636,7 @@ void SVFIRBuilder::visitAllocaInst(AllocaInst &inst) @@
         NodeID src = getObjectNode(&inst);
-        addAddrEdge(src, dst);
+        addAddrWithStackArraySz(src, dst, inst);
     }
@@ Expand Down @@