From 40d5847034f62ac31cb38f0ca6f9e1647d7f6ac4 Mon Sep 17 00:00:00 2001
From: Brice Videau <bvideau@anl.gov>
Date: Fri, 8 Dec 2023 22:12:01 +0000
Subject: [PATCH] Fix hipGraphGetEdges behavior wrt numEdges

---
 src/CHIPBindings.cc | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/src/CHIPBindings.cc b/src/CHIPBindings.cc
index 4d56befe8..5c5f817d9 100644
--- a/src/CHIPBindings.cc
+++ b/src/CHIPBindings.cc
@@ -351,19 +351,26 @@ hipError_t hipGraphGetEdges(hipGraph_t graph, hipGraphNode_t *from,
                             hipGraphNode_t *to, size_t *numEdges) {
   CHIP_TRY
   CHIPInitialize();
+  NULLCHECK(numEdges);
   auto Edges = GRAPH(graph)->getEdges();
   if (!to && !from) {
     *numEdges = Edges.size();
     RETURN(hipSuccess);
   }
 
-  for (int i = 0; i < Edges.size(); i++) {
+  int to_return = *numEdges < Edges.size() ? *numEdges : Edges.size();
+  for (int i = 0; i < to_return; i++) {
     auto Edge = Edges[i];
     auto FromNode = Edge.first;
     auto ToNode = Edge.second;
     from[i] = HIPGRAPHNODE(FromNode);
     to[i] = HIPGRAPHNODE(ToNode);
   }
+  for (int i = to_return; i < *numEdges; i++) {
+    from[i] = nullptr;
+    to[i] = nullptr;
+  }
+  *numEdges = to_return;
   RETURN(hipSuccess);
   CHIP_CATCH
 }