[Inliner] Fix bug when propagating poison generating return attributes

Poison generating return attributes can't be propagated the same as
others, as they can change the behavior of other uses and/or create UB
where it otherwise wouldn't have occurred.

For example:
```
define nonnull ptr @foo() {
    %p = call ptr @bar()
    call void @use(ptr %p)
    ret ptr %p
}
```

If we inline `@foo` and propagate `nonnull` to `@bar`, it could change
the behavior of `@use` as instead of taking `null`, `@use` will
now be passed `poison`.

This can be even worth in a case like:
```
define nonnull ptr @foo() {
    %p = call noundef ptr @bar()
    ret ptr %p
}
```

Where propagating `nonnull` to `@bar` will cause UB on `null` return
of `@bar` (`noundef` + `poison`) where it previously wouldn't
have occurred.

To fix this, we only propagate poison generating return attributes if
either 1) The only use of the callsite to propagate too is return and
the callsite to propagate too doesn't have `noundef`. Or 2) the
callsite to be be inlined has `noundef`.

The former case ensures no new UB or `poison` values will be
added. The latter is UB anyways if the value is `poison` so we can go
ahead without worrying about behavior changes.

1 files changed, 66 insertions, 7 deletions

diff --git a/llvm/lib/Transforms/Utils/InlineFunction.cpp b/llvm/lib/Transforms/Utils/InlineFunction.cpp
index 56bc90a..548f949 100644
--- a/llvm/lib/Transforms/Utils/InlineFunction.cpp
+++ b/llvm/lib/Transforms/Utils/InlineFunction.cpp
@@ -1344,25 +1344,36 @@ static bool MayContainThrowingOrExitingCallAfterCB(CallBase *Begin,
       ++BeginIt, End->getIterator(), InlinerAttributeWindow + 1);
 }
 
-static AttrBuilder IdentifyValidAttributes(CallBase &CB) {
+// Only allow these white listed attributes to be propagated back to the
+// callee. This is because other attributes may only be valid on the call
+// itself, i.e. attributes such as signext and zeroext.
+
+// Attributes that are always okay to propagate as if they are violated its
+// immediate UB.
+static AttrBuilder IdentifyValidUBGeneratingAttributes(CallBase &CB) {
   AttrBuilder Valid(CB.getContext());
-  // Only allow these white listed attributes to be propagated back to the
-  // callee. This is because other attributes may only be valid on the call
-  // itself, i.e. attributes such as signext and zeroext.
   if (auto DerefBytes = CB.getRetDereferenceableBytes())
     Valid.addDereferenceableAttr(DerefBytes);
   if (auto DerefOrNullBytes = CB.getRetDereferenceableOrNullBytes())
     Valid.addDereferenceableOrNullAttr(DerefOrNullBytes);
   if (CB.hasRetAttr(Attribute::NoAlias))
     Valid.addAttribute(Attribute::NoAlias);
+  return Valid;
+}
+
+// Attributes that need additional checks as propagating them may change
+// behavior or cause new UB.
+static AttrBuilder IdentifyValidPoisonGeneratingAttributes(CallBase &CB) {
+  AttrBuilder Valid(CB.getContext());
   if (CB.hasRetAttr(Attribute::NonNull))
     Valid.addAttribute(Attribute::NonNull);
   return Valid;
 }
 
 static void AddReturnAttributes(CallBase &CB, ValueToValueMapTy &VMap) {
-  AttrBuilder Valid = IdentifyValidAttributes(CB);
-  if (!Valid.hasAttributes())
+  AttrBuilder ValidUB = IdentifyValidUBGeneratingAttributes(CB);
+  AttrBuilder ValidPG = IdentifyValidPoisonGeneratingAttributes(CB);
+  if (!ValidUB.hasAttributes() && !ValidPG.hasAttributes())
     return;
   auto *CalledFunction = CB.getCalledFunction();
   auto &Context = CalledFunction->getContext();
@@ -1406,7 +1417,55 @@ static void AddReturnAttributes(CallBase &CB, ValueToValueMapTy &VMap) {
     // existing attribute value (i.e. attributes such as dereferenceable,
     // dereferenceable_or_null etc). See AttrBuilder::merge for more details.
     AttributeList AL = NewRetVal->getAttributes();
-    AttributeList NewAL = AL.addRetAttributes(Context, Valid);
+    AttributeList NewAL = AL.addRetAttributes(Context, ValidUB);
+    // Attributes that may generate poison returns are a bit tricky. If we
+    // propagate them, other uses of the callsite might have their behavior
+    // change or cause UB (if they have noundef) b.c of the new potential
+    // poison.
+    // Take the following three cases:
+    //
+    // 1)
+    // define nonnull ptr @foo() {
+    //   %p = call ptr @bar()
+    //   call void @use(ptr %p) willreturn nounwind
+    //   ret ptr %p
+    // }
+    //
+    // 2)
+    // define noundef nonnull ptr @foo() {
+    //   %p = call ptr @bar()
+    //   call void @use(ptr %p) willreturn nounwind
+    //   ret ptr %p
+    // }
+    //
+    // 3)
+    // define nonnull ptr @foo() {
+    //   %p = call noundef ptr @bar()
+    //   ret ptr %p
+    // }
+    //
+    // In case 1, we can't propagate nonnull because poison value in @use may
+    // change behavior or trigger UB.
+    // In case 2, we don't need to be concerned about propagating nonnull, as
+    // any new poison at @use will trigger UB anyways.
+    // In case 3, we can never propagate nonnull because it may create UB due to
+    // the noundef on @bar.
+    if (ValidPG.hasAttributes()) {
+      // Three checks.
+      // If the callsite has `noundef`, then a poison due to violating the
+      // return attribute will create UB anyways so we can always propagate.
+      // Otherwise, if the return value (callee to be inlined) has `noundef`, we
+      // can't propagate as a new poison return will cause UB.
+      // Finally, check if the return value has no uses whose behavior may
+      // change/may cause UB if we potentially return poison. At the moment this
+      // is implemented overly conservatively with a single-use check.
+      // TODO: Update the single-use check to iterate through uses and only bail
+      // if we have a potentially dangerous use.
+
+      if (CB.hasRetAttr(Attribute::NoUndef) ||
+          (RetVal->hasOneUse() && !RetVal->hasRetAttr(Attribute::NoUndef)))
+        NewAL = NewAL.addRetAttributes(Context, ValidPG);
+    }
     NewRetVal->setAttributes(NewAL);
   }
 }