Track components per location

Author: nikic

llvm/docs/LangRef.rst

@@ -1394,8 +1394,9 @@ Currently, only the following parameter attributes are defined:
     - ``read_provenance`` (subset of ``provenance``): The ability to access the
       pointer only for reads after the function returns.
 
-    Additionally, it is possible to specify that the pointer is captured via
-    the return value only, by using ``captures(ret: ...)``.
+    Additionally, it is possible to specify that some components are only
+    captured in certain locations. Currently only the return value (``ret``)
+    and other (default) locations are supported.
 
     The `pointer capture section <pointercapture>` discusses these semantics
     in more detail.
@@ -1410,6 +1411,9 @@ Currently, only the following parameter attributes are defined:
       captured, but only for read-only access.
     - ``captures(ret: address, provenance)``: Pointer captured through return
       value only.
+    - ``captures(address_is_null, ret: address, provenance)``: The whole pointer
+      is captured through the return value, and additionally whether the pointer
+      is null is captured in some other way.
 
 .. _nocapture:
 

llvm/include/llvm/Support/ModRef.h

@@ -311,55 +311,64 @@ inline bool capturesFullProvenance(CaptureComponents CC) {
 
 raw_ostream &operator<<(raw_ostream &OS, CaptureComponents CC);
 
-/// Represents which components of the pointer may be captured and whether
-/// the capture is via the return value only. This represents the captures(...)
-/// attribute in IR.
+/// Represents which components of the pointer may be captured in which
+/// location. This represents the captures(...) attribute in IR.
 ///
 /// For more information on the precise semantics see LangRef.
 class CaptureInfo {
-  CaptureComponents Components;
-  bool ReturnOnly;
+  CaptureComponents OtherComponents;
+  CaptureComponents RetComponents;
 
 public:
-  CaptureInfo(CaptureComponents Components, bool ReturnOnly = false)
-      : Components(Components),
-        ReturnOnly(capturesAnything(Components) && ReturnOnly) {}
+  CaptureInfo(CaptureComponents OtherComponents,
+              CaptureComponents RetComponents)
+      : OtherComponents(OtherComponents), RetComponents(RetComponents) {}
+
+  CaptureInfo(CaptureComponents Components)
+      : OtherComponents(Components), RetComponents(Components) {}
 
   /// Create CaptureInfo that may capture all components of the pointer.
   static CaptureInfo all() { return CaptureInfo(CaptureComponents::All); }
 
-  /// Get the potentially captured components of the pointer.
-  operator CaptureComponents() const { return Components; }
+  /// Get components potentially captured by the return value.
+  CaptureComponents getRetComponents() const { return RetComponents; }
+
+  /// Get components potentially captured through locations other than the
+  /// return value.
+  CaptureComponents getOtherComponents() const { return OtherComponents; }
 
-  /// Whether the pointer is captured through the return value only.
-  bool isReturnOnly() const { return ReturnOnly; }
+  /// Get the potentially captured components of the pointer (regardless of
+  /// location).
+  operator CaptureComponents() const { return OtherComponents | RetComponents; }
 
   bool operator==(CaptureInfo Other) const {
-    return Components == Other.Components && ReturnOnly == Other.ReturnOnly;
+    return OtherComponents == Other.OtherComponents &&
+           RetComponents == Other.RetComponents;
   }
 
   bool operator!=(CaptureInfo Other) const { return !(*this == Other); }
 
   /// Compute union of CaptureInfos.
   CaptureInfo operator|(CaptureInfo Other) const {
-    return CaptureInfo(Components | Other.Components,
-                       ReturnOnly && Other.ReturnOnly);
+    return CaptureInfo(OtherComponents | Other.OtherComponents,
+                       RetComponents | Other.RetComponents);
   }
 
   /// Compute intersection of CaptureInfos.
   CaptureInfo operator&(CaptureInfo Other) const {
-    return CaptureInfo(Components & Other.Components,
-                       ReturnOnly || Other.ReturnOnly);
+    return CaptureInfo(OtherComponents & Other.OtherComponents,
+                       RetComponents & Other.RetComponents);
   }
 
   static CaptureInfo createFromIntValue(uint32_t Data) {
-    return CaptureInfo(CaptureComponents(Data >> 1), Data & 1);
+    return CaptureInfo(CaptureComponents(Data >> 4),
+                       CaptureComponents(Data & 0xf));
   }
 
   /// Convert CaptureInfo into an encoded integer value (used by captures
   /// attribute).
   uint32_t toIntValue() const {
-    return (uint32_t(Components) << 1) | ReturnOnly;
+    return (uint32_t(OtherComponents) << 4) | uint32_t(RetComponents);
   }
 };
 

llvm/lib/AsmParser/LLParser.cpp

@@ -3168,8 +3168,8 @@ bool LLParser::parseInitializesAttr(AttrBuilder &B) {
 }
 
 bool LLParser::parseCapturesAttr(AttrBuilder &B) {
-  CaptureComponents CC = CaptureComponents::None;
-  bool ReturnOnly = false;
+  CaptureComponents Other = CaptureComponents::None;
+  std::optional<CaptureComponents> Ret;
 
   // We use syntax like captures(ret: address, provenance), so the colon
   // should not be interpreted as a label terminator.
@@ -3180,39 +3180,49 @@ bool LLParser::parseCapturesAttr(AttrBuilder &B) {
   if (parseToken(lltok::lparen, "expected '('"))
     return true;
 
-  if (EatIfPresent(lltok::kw_ret)) {
-    if (parseToken(lltok::colon, "expected ':'"))
-      return true;
+  CaptureComponents *Current = &Other;
+  bool SeenComponent = false;
+  while (true) {
+    if (EatIfPresent(lltok::kw_ret)) {
+      if (parseToken(lltok::colon, "expected ':'"))
+        return true;
+      if (Ret)
+        return tokError("duplicate 'ret' location");
+      Ret = CaptureComponents::None;
+      Current = &*Ret;
+      SeenComponent = false;
+    }
 
-    ReturnOnly = true;
-  }
+    if (EatIfPresent(lltok::kw_none)) {
+      if (SeenComponent)
+        return tokError("cannot use 'none' with other component");
+      *Current = CaptureComponents::None;
+    } else {
+      if (SeenComponent && capturesNothing(*Current))
+        return tokError("cannot use 'none' with other component");
 
-  if (EatIfPresent(lltok::kw_none)) {
-    if (parseToken(lltok::rparen, "expected ')'"))
-      return true;
-  } else {
-    while (true) {
       if (EatIfPresent(lltok::kw_address_is_null))
-        CC |= CaptureComponents::AddressIsNull;
+        *Current |= CaptureComponents::AddressIsNull;
       else if (EatIfPresent(lltok::kw_address))
-        CC |= CaptureComponents::Address;
+        *Current |= CaptureComponents::Address;
       else if (EatIfPresent(lltok::kw_provenance))
-        CC |= CaptureComponents::Provenance;
+        *Current |= CaptureComponents::Provenance;
       else if (EatIfPresent(lltok::kw_read_provenance))
-        CC |= CaptureComponents::ReadProvenance;
+        *Current |= CaptureComponents::ReadProvenance;
       else
-        return tokError("expected one of 'address', 'address_is_null', "
+        return tokError("expected one of 'none', 'address', 'address_is_null', "
                         "'provenance' or 'read_provenance'");
+    }
 
-      if (EatIfPresent(lltok::rparen))
-        break;
+    SeenComponent = true;
+    if (EatIfPresent(lltok::rparen))
+      break;
 
-      if (parseToken(lltok::comma, "expected ',' or ')'"))
-        return true;
-    }
+    if (parseToken(lltok::comma, "expected ',' or ')'"))
+      return true;
   }
 
-  B.addCapturesAttr(CaptureInfo(CC, ReturnOnly));
+  B.addCapturesAttr(CaptureInfo(Other, Ret.value_or(Other)));
   return false;
 }
 

llvm/lib/Support/ModRef.cpp

@@ -72,9 +72,15 @@ raw_ostream &llvm::operator<<(raw_ostream &OS, CaptureComponents CC) {
 }
 
 raw_ostream &llvm::operator<<(raw_ostream &OS, CaptureInfo CI) {
+  ListSeparator LS;
+  CaptureComponents Other = CI.getOtherComponents();
+  CaptureComponents Ret = CI.getRetComponents();
+
   OS << "captures(";
-  if (CI.isReturnOnly())
-    OS << "ret: ";
-  OS << CaptureComponents(CI) << ")";
+  if (!capturesNothing(Other) || Other == Ret)
+    OS << LS << Other;
+  if (Other != Ret)
+    OS << LS << "ret: " << Ret;
+  OS << ")";
   return OS;
 }

llvm/test/Assembler/captures-errors.ll

@@ -4,6 +4,9 @@
 ; RUN: not llvm-as < %t/missing-rparen-none.ll 2>&1 | FileCheck %s --check-prefix=CHECK-MISSING-RPAREN-NONE
 ; RUN: not llvm-as < %t/missing-colon.ll 2>&1 | FileCheck %s --check-prefix=CHECK-MISSING-COLON
 ; RUN: not llvm-as < %t/invalid-component.ll 2>&1 | FileCheck %s --check-prefix=CHECK-INVALID-COMPONENT
+; RUN: not llvm-as < %t/duplicate-ret.ll 2>&1 | FileCheck %s --check-prefix=CHECK-DUPLICATE-RET
+; RUN: not llvm-as < %t/none-after.ll 2>&1 | FileCheck %s --check-prefix=CHECK-NONE-AFTER
+; RUN: not llvm-as < %t/none-before.ll 2>&1 | FileCheck %s --check-prefix=CHECK-NONE-BEFORE
 
 ;--- missing-lparen.ll
 
@@ -21,7 +24,7 @@ define void @test(ptr captures(address %p) {
 
 ;--- missing-rparen-none.ll
 
-; CHECK-MISSING-RPAREN-NONE: <stdin>:[[@LINE+1]]:37: error: expected ')'
+; CHECK-MISSING-RPAREN-NONE: <stdin>:[[@LINE+1]]:37: error: expected ',' or ')'
 define void @test(ptr captures(none %p) {
   ret void
 }
@@ -35,7 +38,28 @@ define void @test(ptr captures(ret address) %p) {
 
 ;--- invalid-component.ll
 
-; CHECK-INVALID-COMPONENT: <stdin>:[[@LINE+1]]:32: error: expected one of 'address', 'address_is_null', 'provenance' or 'read_provenance'
+; CHECK-INVALID-COMPONENT: <stdin>:[[@LINE+1]]:32: error: expected one of 'none', 'address', 'address_is_null', 'provenance' or 'read_provenance'
 define void @test(ptr captures(foo) %p) {
   ret void
 }
+
+;--- duplicate-ret.ll
+
+; CHECK-DUPLICATE-RET: <stdin>:[[@LINE+1]]:51: error: duplicate 'ret' location
+define void @test(ptr captures(ret: address, ret: provenance) %p) {
+  ret void
+}
+
+;--- none-after.ll
+
+; CHECK-NONE-AFTER: <stdin>:[[@LINE+1]]:45: error: cannot use 'none' with other component
+define void @test(ptr captures(address, none) %p) {
+  ret void
+}
+
+;--- none-before.ll
+
+; CHECK-NONE-BEFORE: <stdin>:[[@LINE+1]]:38: error: cannot use 'none' with other component
+define void @test(ptr captures(none, address) %p) {
+  ret void
+}

llvm/test/Assembler/captures.ll

@@ -50,6 +50,22 @@ define void @test_ret(ptr captures(ret: address, provenance) %p) {
   ret void
 }
 
+define void @test_address_is_null_and_ret(ptr captures(address_is_null, ret: address, provenance) %p) {
+; CHECK-LABEL: define void @test_address_is_null_and_ret(
+; CHECK-SAME: ptr captures(address_is_null, ret: address, provenance) [[P:%.*]]) {
+; CHECK-NEXT:    ret void
+;
+  ret void
+}
+
+define void @test_address_and_ret_none(ptr captures(address, ret: none) %p) {
+; CHECK-LABEL: define void @test_address_and_ret_none(
+; CHECK-SAME: ptr captures(address, ret: none) [[P:%.*]]) {
+; CHECK-NEXT:    ret void
+;
+  ret void
+}
+
 ; Duplicates callpse into one.
 define void @test_duplicate(ptr captures(address, address) %p) {
 ; CHECK-LABEL: define void @test_duplicate(

llvm/unittests/IR/AttributesTest.cpp

@@ -438,9 +438,11 @@ TEST(Attributes, SetIntersect) {
       case Attribute::Range:
         break;
       case Attribute::Captures:
-        V0 = CaptureInfo(CaptureComponents::Address, /*ReturnOnly=*/false)
+        V0 = CaptureInfo(CaptureComponents::AddressIsNull,
+                         CaptureComponents::None)
                  .toIntValue();
-        V1 = CaptureInfo(CaptureComponents::ReadProvenance, /*ReturnOnly=*/true)
+        V1 = CaptureInfo(CaptureComponents::None,
+                         CaptureComponents::ReadProvenance)
                  .toIntValue();
         break;
       default:
@@ -524,9 +526,8 @@ TEST(Attributes, SetIntersect) {
         break;
       case Attribute::Captures:
         ASSERT_EQ(Res->getCaptureInfo(),
-                  CaptureInfo(CaptureComponents::Address |
-                                  CaptureComponents::ReadProvenance,
-                              /*ReturnOnly=*/false));
+                  CaptureInfo(CaptureComponents::AddressIsNull,
+                              CaptureComponents::ReadProvenance));
         break;
       default:
         ASSERT_FALSE(true);