[FunctionAttr] Infer nonnull attributes on returns

[reames]

Summary

Teach FunctionAttr to infer the nonnull attribute on return values of functions which never return a potentially null value. This is done both via a conservative local analysis for the function itself and a optimistic per-SCC analysis. If no function in the SCC returns anything which could be null (other than values from other functions in the SCC), we can conclude no function returned a null pointer. Even if some function within the SCC returns a null pointer, we may be able to locally conclude that some don't.

This is inspired by discussion on http://reviews.llvm.org/D9132, but is mostly orthogonal to that patch except that they both relate to nonnul attributes.

[pete]

Hi Philip

I made a few comments, but mostly just about using foreach and some comments. Nothing major.

Otherwise, LGTM.

Cheers,
Pete

lib/Transforms/IPO/FunctionAttrs.cpp

73

I think you need \brief here

77

And here, with \\\ and without listing the function name.

849

All function types are pointer types. I think you need F->getReturnType()->isPointerTy() as you have below.

854

I think this can just be a foreach loop.

858

I think this would be better as a worklist you pop values off until its empty.

866

I think this can be a dyn_cast so that you can fold it with the cast<> below.

887

I don't know if there's a foreach on the phi values. If there is then good to use it, otherwise I guess someone else can add it later.

[nlewycky]

lib/Transforms/IPO/FunctionAttrs.cpp

858

Be careful not to break the case of resolving a self recursive phi. I think the normal way to handle this is with a set of "visited" elements, plus a worklist "queue" like Pete suggests.

920

"Speculative assume" -> "Speculatively assume" or possibly just "Speculate"

test/Transforms/FunctionAttrs/nonnull.ll

61

Extra newline.

[reames]

Responding to comments. Any comment not explicitly addressed with be fixed before submission.

lib/Transforms/IPO/FunctionAttrs.cpp

849

Thanks. Good catch.

854

Yep, I copied this from code above. If you don't mind, I'll leave this as is, then fix both in one patch.

858

As Nick pointed out, the use of the SetVector is important here. Given this is how other algorithms in the same file are structured, I'd prefer not to change this.

887

I copied this from code above. If you don't mind, I'll leave this as is, then fix both in one patch.

[pete]

SGTM. I have no problems with a separate patch. Thanks!

lib/Transforms/IPO/FunctionAttrs.cpp

 #include "llvm/ADT/SCCIterator.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/CallGraph.h"
 #include "llvm/Analysis/CallGraphSCCPass.h"
 #include "llvm/Analysis/CaptureTracking.h"
+#include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/LLVMContext.h"
+#include "llvm/Support/Debug.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 using namespace llvm;
 
 #define DEBUG_TYPE "functionattrs"
 
 STATISTIC(NumReadNone, "Number of functions marked readnone");
 STATISTIC(NumReadOnly, "Number of functions marked readonly");
 STATISTIC(NumNoCapture, "Number of arguments marked nocapture");
 STATISTIC(NumReadNoneArg, "Number of arguments marked readnone");
 STATISTIC(NumReadOnlyArg, "Number of arguments marked readonly");
 STATISTIC(NumNoAlias, "Number of function returns marked noalias");
+STATISTIC(NumNonNullReturn, "Number of function returns marked nonnull");
 STATISTIC(NumAnnotated, "Number of attributes added to library functions");
 
 namespace {
   struct FunctionAttrs : public CallGraphSCCPass {
     static char ID; // Pass identification, replacement for typeid
     FunctionAttrs() : CallGraphSCCPass(ID), AA(nullptr) {
       initializeFunctionAttrsPass(*PassRegistry::getPassRegistry());
     }
struct FunctionAttrs : public CallGraphSCCPass {
 
     // IsFunctionMallocLike - Does this function allocate new memory?
     bool IsFunctionMallocLike(Function *F,
                               SmallPtrSet<Function*, 8> &) const;
 
     // AddNoAliasAttrs - Deduce noalias attributes for the SCC.
     bool AddNoAliasAttrs(const CallGraphSCC &SCC);
 
+    /// Does this function return null?  
+    bool ReturnsNonNull(Function *F, SmallPtrSet<Function*, 8> &,
+                        bool &Speculative) const;
+
+    // AddNonNullAttrs - Deduce nonnull attributes for the SCC.
+    bool AddNonNullAttrs(const CallGraphSCC &SCC);
+
     // Utility methods used by inferPrototypeAttributes to add attributes
     // and maintain annotation statistics.
 
     void setDoesNotAccessMemory(Function &F) {
       if (!F.doesNotAccessMemory()) {
         F.setDoesNotAccessMemory();
         ++NumAnnotated;
       }
for (CallGraphSCC::iterator I = SCC.begin(), E = SCC.end(); I != E; ++I) {
     F->setDoesNotAlias(0);
     ++NumNoAlias;
     MadeChange = true;
   }
 
   return MadeChange;
 }
 
+bool FunctionAttrs::ReturnsNonNull(Function *F,
+                                   SmallPtrSet<Function*, 8> &SCCNodes,
+                                   bool &Speculative) const {
+  assert(F->getType()->isPointerTy() &&
+         "nonnull only meaningful on pointer types");
+  Speculative = false;
+  
+  SmallSetVector<Value *, 8> FlowsToReturn;
+  for (Function::iterator I = F->begin(), E = F->end(); I != E; ++I)
+    if (ReturnInst *Ret = dyn_cast<ReturnInst>(I->getTerminator()))
+      FlowsToReturn.insert(Ret->getReturnValue());
+
+  for (unsigned i = 0; i != FlowsToReturn.size(); ++i) {
+    Value *RetVal = FlowsToReturn[i];
+
+    // If this value is locally known to be non-null, we're good
+    if (isKnownNonNull(RetVal, TLI))
+      continue;
+
+    // It wasn't nonnull and we can't look through it..
+    if (!isa<Instruction>(RetVal))
+      return false;
+
+    // Otherwise, we need to look upwards since we can't make any local
+    // conclusions.  
+    Instruction *RVI = cast<Instruction>(RetVal);
+    switch (RVI->getOpcode()) {
+      // Extend the analysis by looking upwards.
+    case Instruction::BitCast:
+    case Instruction::GetElementPtr:
+    case Instruction::AddrSpaceCast:
+      FlowsToReturn.insert(RVI->getOperand(0));
+      continue;
+    case Instruction::Select: {
+      SelectInst *SI = cast<SelectInst>(RVI);
+      FlowsToReturn.insert(SI->getTrueValue());
+      FlowsToReturn.insert(SI->getFalseValue());
+      continue;
+    }
+    case Instruction::PHI: {
+      PHINode *PN = cast<PHINode>(RVI);
+      for (int i = 0, e = PN->getNumIncomingValues(); i != e; ++i)
+        FlowsToReturn.insert(PN->getIncomingValue(i));
+      continue;
+    }
+    case Instruction::Call:
+    case Instruction::Invoke: {
+      CallSite CS(RVI);
+      Function *Callee = CS.getCalledFunction();
+      // A call to a node within the SCC is assumed to return null until
+      // proven otherwise
+      if (Callee && SCCNodes.count(Callee)) {
+        Speculative = true;
+        continue;
+      }
+      return false;
+    }
+    default:
+      return false;  // Unknown source, may be null
+    };
+    llvm_unreachable("should have either continued or returned");
+  }
+
+  return true;
+}
+
+bool FunctionAttrs::AddNonNullAttrs(const CallGraphSCC &SCC) {
+  SmallPtrSet<Function*, 8> SCCNodes;
+
+  // Fill SCCNodes with the elements of the SCC.  Used for quickly
+  // looking up whether a given CallGraphNode is in this SCC.
+  for (CallGraphSCC::iterator I = SCC.begin(), E = SCC.end(); I != E; ++I)
+    SCCNodes.insert((*I)->getFunction());
+
+  // Speculative assume that all functions in the SCC return only nonnull
+  // pointers.  We may refute this as we analyze functions.
+  bool SCCReturnsNonNull = true;
+
+  bool MadeChange = false;
+
+  // Check each function in turn, determining which functions return nonnull
+  // pointers.
+  for (CallGraphSCC::iterator I = SCC.begin(), E = SCC.end(); I != E; ++I) {
+    Function *F = (*I)->getFunction();
+
+    if (!F || F->hasFnAttribute(Attribute::OptimizeNone))
+      // External node or node we don't want to optimize - skip it;
+      return false;
+
+    // Already nonnull.
+    if (F->getAttributes().hasAttribute(AttributeSet::ReturnIndex,
+                                        Attribute::NonNull))
+      continue;
+
+    // Definitions with weak linkage may be overridden at linktime, so
+    // treat them like declarations.
+    if (F->isDeclaration() || F->mayBeOverridden())
+      return false;
+
+    // We annotate nonnull return values, which are only applicable to 
+    // pointer types.
+    if (!F->getReturnType()->isPointerTy())
+      continue;
+
+    bool Speculative = false;
+    if (ReturnsNonNull(F, SCCNodes, Speculative)) {
+      if (!Speculative) {
+        // Mark the function eagerly since we may discover a function
+        // which prevents us from speculating about the entire SCC
+        DEBUG(dbgs() << "Eagerly marking " << F->getName() << " as nonnull\n");
+        MadeChange = true;
+        F->addAttribute(AttributeSet::ReturnIndex, Attribute::NonNull);
+      }
+      continue;
+    }
+    // At least one function returns something which could be null, can't
+    // speculate any more.
+    SCCReturnsNonNull = false;
+  }
+
+  if (SCCReturnsNonNull) {
+    for (CallGraphSCC::iterator I = SCC.begin(), E = SCC.end(); I != E; ++I) {
+      Function *F = (*I)->getFunction();
+      if (F->getAttributes().hasAttribute(AttributeSet::ReturnIndex,
+                                          Attribute::NonNull) ||
+          !F->getReturnType()->isPointerTy())
+        continue;
+
+      DEBUG(dbgs() << "SCC marking " << F->getName() << " as nonnull\n");
+      F->addAttribute(AttributeSet::ReturnIndex, Attribute::NonNull);
+      ++NumNonNullReturn;
+      MadeChange = true;
+    }
+  }
+
+  return MadeChange;
+}
+
 /// inferPrototypeAttributes - Analyze the name and prototype of the
 /// given function and set any applicable attributes.  Returns true
 /// if any attributes were set and false otherwise.
 bool FunctionAttrs::inferPrototypeAttributes(Function &F) {
   if (F.hasFnAttribute(Attribute::OptimizeNone))
     return false;
 
   FunctionType *FTy = F.getFunctionType();

 bool FunctionAttrs::runOnSCC(CallGraphSCC &SCC) {
   AA = &getAnalysis<AliasAnalysis>();
   TLI = &getAnalysis<TargetLibraryInfoWrapperPass>().getTLI();
 
   bool Changed = annotateLibraryCalls(SCC);
   Changed |= AddReadAttrs(SCC);
   Changed |= AddArgumentAttrs(SCC);
   Changed |= AddNoAliasAttrs(SCC);
+  Changed |= AddNonNullAttrs(SCC);
   return Changed;
 }

test/Transforms/FunctionAttrs/nonnull.ll

+; RUN: opt -S -functionattrs %s | FileCheck %s
+declare nonnull i8* @ret_nonnull()
+
+; Return a pointer trivially nonnull (call return attribute)
+define i8* @test1() {
+; CHECK: define nonnull i8* @test1
+  %ret = call i8* @ret_nonnull()
+  ret i8* %ret
+}
+
+; Return a pointer trivially nonnull (argument attribute)
+define i8* @test2(i8* nonnull %p) {
+; CHECK: define nonnull i8* @test2
+  ret i8* %p
+}
+
+; Given an SCC where one of the functions can not be marked nonnull,
+; can we still mark the other one which is trivially nonnull
+define i8* @scc_binder() {
+; CHECK: define i8* @scc_binder
+  call i8* @test3()
+  ret i8* null
+}
+
+define i8* @test3() {
+; CHECK: define nonnull i8* @test3
+  call i8* @scc_binder()
+  %ret = call i8* @ret_nonnull()
+  ret i8* %ret
+}
+
+; Given a mutual recursive set of functions, we can mark them
+; nonnull if neither can ever return null.  (In this case, they
+; just never return period.)
+define i8* @test4_helper() {
+; CHECK: define noalias nonnull i8* @test4_helper
+  %ret = call i8* @test4()
+  ret i8* %ret
+}
+
+define i8* @test4() {
+; CHECK: define noalias nonnull i8* @test4
+  %ret = call i8* @test4_helper()
+  ret i8* %ret
+}
+
+; Given a mutual recursive set of functions which *can* return null
+; make sure we haven't marked them as nonnull.
+define i8* @test5_helper() {
+; CHECK: define noalias i8* @test5_helper
+  %ret = call i8* @test5()
+  ret i8* null
+}
+
+define i8* @test5() {
+; CHECK: define noalias i8* @test5
+  %ret = call i8* @test5_helper()
+  ret i8* %ret
+}
+
+
+; Local analysis, but going through a self recursive phi
+define i8* @test6() {
+entry:
+; CHECK: define nonnull i8* @test6
+  %ret = call i8* @ret_nonnull()
+  br label %loop
+loop:
+  %phi = phi i8* [%ret, %entry], [%phi, %loop]
+  br i1 undef, label %loop, label %exit
+exit:
+  ret i8* %phi
+}
+
+