3 files changed, 54 insertions, 14 deletions

diff --git a/llvm/include/llvm/IR/Statepoint.h b/llvm/include/llvm/IR/Statepoint.h
index 3d4091c3bf2..a27c43b1e1e 100644
--- a/llvm/include/llvm/IR/Statepoint.h
+++ b/llvm/include/llvm/IR/Statepoint.h
@@ -8,8 +8,9 @@
 //===----------------------------------------------------------------------===//
 //
 // This file contains utility functions and a wrapper class analogous to
-// CallSite for accessing the fields of gc.statepoint, gc.relocate, and
-// gc.result intrinsics
+// CallSite for accessing the fields of gc.statepoint, gc.relocate,
+// gc.result intrinsics; and some general utilities helpful when dealing with
+// gc.statepoint.
 //
 //===----------------------------------------------------------------------===//
 
@@ -17,6 +18,7 @@
 #define LLVM_IR_STATEPOINT_H
 
 #include "llvm/ADT/iterator_range.h"
+#include "llvm/ADT/Optional.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/CallSite.h"
 #include "llvm/IR/Constants.h"
@@ -399,6 +401,23 @@ StatepointBase<FunTy, InstructionTy, ValueTy, CallSiteTy>::getRelocates()
   }
   return Result;
 }
+
+/// Call sites that get wrapped by a gc.statepoint (currently only in
+/// RewriteStatepointsForGC and potentially in other passes in the future) can
+/// have attributes that describe properties of gc.statepoint call they will be
+/// eventually be wrapped in.  This struct is used represent such directives.
+struct StatepointDirectives {
+  Optional<uint32_t> NumPatchBytes;
+  Optional<uint64_t> StatepointID;
+};
+
+/// Parse out statepoint directives from the function attributes present in \p
+/// AS.
+StatepointDirectives parseStatepointDirectivesFromAttrs(AttributeSet AS);
+
+/// Return \c true if the the \p Attr is an attribute that is a statepoint
+/// directive.
+bool isStatepointDirectiveAttr(Attribute Attr);
 }
 
 #endif
diff --git a/llvm/lib/IR/Statepoint.cpp b/llvm/lib/IR/Statepoint.cpp
index e56da6beaff..13e8c6312a6 100644
--- a/llvm/lib/IR/Statepoint.cpp
+++ b/llvm/lib/IR/Statepoint.cpp
@@ -53,3 +53,28 @@ bool llvm::isGCResult(ImmutableCallSite CS) {
 bool llvm::isGCResult(const Value *V) {
   return isGCResult(ImmutableCallSite(V));
 }
+
+bool llvm::isStatepointDirectiveAttr(Attribute Attr) {
+  return Attr.hasAttribute("statepoint-id") ||
+         Attr.hasAttribute("statepoint-num-patch-bytes");
+}
+
+StatepointDirectives llvm::parseStatepointDirectivesFromAttrs(AttributeSet AS) {
+  StatepointDirectives Result;
+
+  Attribute AttrID =
+      AS.getAttribute(AttributeSet::FunctionIndex, "statepoint-id");
+  uint64_t StatepointID;
+  if (AttrID.isStringAttribute())
+    if (!AttrID.getValueAsString().getAsInteger(10, StatepointID))
+      Result.StatepointID = StatepointID;
+
+  uint32_t NumPatchBytes;
+  Attribute AttrNumPatchBytes = AS.getAttribute(AttributeSet::FunctionIndex,
+                                                "statepoint-num-patch-bytes");
+  if (AttrNumPatchBytes.isStringAttribute())
+    if (!AttrNumPatchBytes.getValueAsString().getAsInteger(10, NumPatchBytes))
+      Result.NumPatchBytes = NumPatchBytes;
+
+  return Result;
+}
diff --git a/llvm/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp b/llvm/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp
index 3cdf84456ff..85aaf62fb3c 100644
--- a/llvm/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp
+++ b/llvm/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp
@@ -1187,8 +1187,7 @@ static AttributeSet legalizeCallAttributes(AttributeSet AS) {
         // These attributes control the generation of the gc.statepoint call /
         // invoke itself; and once the gc.statepoint is in place, they're of no
         // use.
-        if (Attr.hasAttribute("statepoint-num-patch-bytes") ||
-            Attr.hasAttribute("statepoint-id"))
+        if (isStatepointDirectiveAttr(Attr))
           continue;
 
         Ret = Ret.addAttributes(
@@ -1332,17 +1331,14 @@ makeStatepointExplicitImpl(const CallSite CS, /* to replace */
     TransitionArgs = TransitionBundle->Inputs;
   }
 
+  StatepointDirectives SD =
+      parseStatepointDirectivesFromAttrs(CS.getAttributes());
+  if (SD.NumPatchBytes)
+    NumPatchBytes = *SD.NumPatchBytes;
+  if (SD.StatepointID)
+    StatepointID = *SD.StatepointID;
+
   Value *CallTarget = CS.getCalledValue();
-  AttributeSet OriginalAttrs = CS.getAttributes();
-  Attribute AttrID = OriginalAttrs.getAttribute(AttributeSet::FunctionIndex,
-                                                "statepoint-id");
-  if (AttrID.isStringAttribute())
-    AttrID.getValueAsString().getAsInteger(10, StatepointID);
-
-  Attribute AttrNumPatchBytes = OriginalAttrs.getAttribute(
-    AttributeSet::FunctionIndex, "statepoint-num-patch-bytes");
-  if (AttrNumPatchBytes.isStringAttribute())
-    AttrNumPatchBytes.getValueAsString().getAsInteger(10, NumPatchBytes);
 
   // Create the statepoint given all the arguments
   Instruction *Token = nullptr;