2 files changed, 85 insertions, 52 deletions

diff --git a/clang/lib/Analysis/Consumed.cpp b/clang/lib/Analysis/Consumed.cpp
index 78153a0a684..b540d75898d 100644
--- a/clang/lib/Analysis/Consumed.cpp
+++ b/clang/lib/Analysis/Consumed.cpp
@@ -89,9 +89,25 @@ static bool isTestingFunction(const FunctionDecl *FunDecl) {
   return FunDecl->hasAttr<TestsUnconsumedAttr>();
 }
 
+static ConsumedState mapConsumableAttrState(const QualType QT) {
+  assert(isConsumableType(QT));
+
+  const ConsumableAttr *CAttr =
+      QT->getAsCXXRecordDecl()->getAttr<ConsumableAttr>();
+
+  switch (CAttr->getDefaultState()) {
+  case ConsumableAttr::Unknown:
+    return CS_Unknown;
+  case ConsumableAttr::Unconsumed:
+    return CS_Unconsumed;
+  case ConsumableAttr::Consumed:
+    return CS_Consumed;
+  }
+  llvm_unreachable("invalid enum");
+}
+
 static ConsumedState
 mapReturnTypestateAttrState(const ReturnTypestateAttr *RTSAttr) {
-
   switch (RTSAttr->getState()) {
   case ReturnTypestateAttr::Unknown:
     return CS_Unknown;
@@ -402,7 +418,7 @@ void ConsumedStmtVisitor::propagateReturnType(const Stmt *Call,
       ReturnState = mapReturnTypestateAttrState(
         Fun->getAttr<ReturnTypestateAttr>());
     else
-      ReturnState = CS_Unknown;
+      ReturnState = mapConsumableAttrState(ReturnType);
     
     PropagationMap.insert(PairType(Call,
       PropagationInfo(ReturnState)));
@@ -709,8 +725,18 @@ void ConsumedStmtVisitor::VisitMemberExpr(const MemberExpr *MExpr) {
 
 
 void ConsumedStmtVisitor::VisitParmVarDecl(const ParmVarDecl *Param) {
-  if (isConsumableType(Param->getType()))
-    StateMap->setState(Param, consumed::CS_Unknown);
+  QualType ParamType = Param->getType();
+  ConsumedState ParamState = consumed::CS_None;
+
+  if (!(ParamType->isPointerType() || ParamType->isReferenceType()) &&
+      isConsumableType(ParamType))
+    ParamState = mapConsumableAttrState(ParamType);
+  else if (ParamType->isReferenceType() &&
+           isConsumableType(ParamType->getPointeeType()))
+    ParamState = consumed::CS_Unknown;
+
+  if (ParamState)
+    StateMap->setState(Param, ParamState);
 }
 
 void ConsumedStmtVisitor::VisitReturnStmt(const ReturnStmt *Ret) {
@@ -952,6 +978,35 @@ void ConsumedStateMap::remove(const VarDecl *Var) {
   Map.erase(Var);
 }
 
+void ConsumedAnalyzer::determineExpectedReturnState(AnalysisDeclContext &AC,
+                                                    const FunctionDecl *D) {
+  QualType ReturnType;
+  if (const CXXConstructorDecl *Constructor = dyn_cast<CXXConstructorDecl>(D)) {
+    ASTContext &CurrContext = AC.getASTContext();
+    ReturnType = Constructor->getThisType(CurrContext)->getPointeeType();
+  } else
+    ReturnType = D->getCallResultType();
+
+  if (D->hasAttr<ReturnTypestateAttr>()) {
+    const ReturnTypestateAttr *RTSAttr = D->getAttr<ReturnTypestateAttr>();
+
+    const CXXRecordDecl *RD = ReturnType->getAsCXXRecordDecl();
+    if (!RD || !RD->hasAttr<ConsumableAttr>()) {
+      // FIXME: This should be removed when template instantiation propagates
+      //        attributes at template specialization definition, not
+      //        declaration. When it is removed the test needs to be enabled
+      //        in SemaDeclAttr.cpp.
+      WarningsHandler.warnReturnTypestateForUnconsumableType(
+          RTSAttr->getLocation(), ReturnType.getAsString());
+      ExpectedReturnState = CS_None;
+    } else
+      ExpectedReturnState = mapReturnTypestateAttrState(RTSAttr);
+  } else if (isConsumableType(ReturnType))
+    ExpectedReturnState = mapConsumableAttrState(ReturnType);
+  else
+    ExpectedReturnState = CS_None;
+}
+
 bool ConsumedAnalyzer::splitState(const CFGBlock *CurrBlock,
                                   const ConsumedStmtVisitor &Visitor) {
   
@@ -1051,52 +1106,7 @@ void ConsumedAnalyzer::run(AnalysisDeclContext &AC) {
   
   if (!D) return;
   
-  // FIXME: This should be removed when template instantiation propagates
-  //        attributes at template specialization definition, not declaration.
-  //        When it is removed the test needs to be enabled in SemaDeclAttr.cpp.
-  QualType ReturnType;
-  if (const CXXConstructorDecl *Constructor = dyn_cast<CXXConstructorDecl>(D)) {
-    ASTContext &CurrContext = AC.getASTContext();
-    ReturnType = Constructor->getThisType(CurrContext)->getPointeeType();
-    
-  } else {
-    ReturnType = D->getCallResultType();
-  }
-  
-  // Determine the expected return value.
-  if (D->hasAttr<ReturnTypestateAttr>()) {
-    
-    ReturnTypestateAttr *RTSAttr = D->getAttr<ReturnTypestateAttr>();
-    
-    const CXXRecordDecl *RD = ReturnType->getAsCXXRecordDecl();
-    if (!RD || !RD->hasAttr<ConsumableAttr>()) {
-        // FIXME: This branch can be removed with the code above.
-        WarningsHandler.warnReturnTypestateForUnconsumableType(
-          RTSAttr->getLocation(), ReturnType.getAsString());
-        ExpectedReturnState = CS_None;
-        
-    } else {
-      switch (RTSAttr->getState()) {
-      case ReturnTypestateAttr::Unknown:
-        ExpectedReturnState = CS_Unknown;
-        break;
-        
-      case ReturnTypestateAttr::Unconsumed:
-        ExpectedReturnState = CS_Unconsumed;
-        break;
-        
-      case ReturnTypestateAttr::Consumed:
-        ExpectedReturnState = CS_Consumed;
-        break;
-      }
-    }
-    
-  } else if (isConsumableType(ReturnType)) {
-    ExpectedReturnState = CS_Unknown;
-      
-  } else {
-    ExpectedReturnState = CS_None;
-  }
+  determineExpectedReturnState(AC, D);
   
   BlockInfo = ConsumedBlockInfo(AC.getCFG());
   
diff --git a/clang/lib/Sema/SemaDeclAttr.cpp b/clang/lib/Sema/SemaDeclAttr.cpp
index 97d12d57dcb..39347fa3949 100644
--- a/clang/lib/Sema/SemaDeclAttr.cpp
+++ b/clang/lib/Sema/SemaDeclAttr.cpp
@@ -968,7 +968,30 @@ static void handleLocksExcludedAttr(Sema &S, Decl *D,
 }
 
 static void handleConsumableAttr(Sema &S, Decl *D, const AttributeList &Attr) {
-  if (!checkAttributeNumArgs(S, Attr, 0)) return;
+  if (!checkAttributeNumArgs(S, Attr, 1))
+    return;
+
+  ConsumableAttr::ConsumedState DefaultState;
+
+  if (Attr.isArgIdent(0)) {
+    StringRef Param = Attr.getArgAsIdent(0)->Ident->getName();
+
+    if (Param == "unknown")
+      DefaultState = ConsumableAttr::Unknown;
+    else if (Param == "consumed")
+      DefaultState = ConsumableAttr::Consumed;
+    else if (Param == "unconsumed")
+      DefaultState = ConsumableAttr::Unconsumed;
+    else {
+      S.Diag(Attr.getLoc(), diag::warn_unknown_consumed_state) << Param;
+      return;
+    }
+
+  } else {
+    S.Diag(Attr.getLoc(), diag::err_attribute_argument_type)
+        << Attr.getName() << AANT_ArgumentIdentifier;
+    return;
+  }
 
   if (!isa<CXXRecordDecl>(D)) {
     S.Diag(Attr.getLoc(), diag::warn_attribute_wrong_decl_type) <<
@@ -977,7 +1000,7 @@ static void handleConsumableAttr(Sema &S, Decl *D, const AttributeList &Attr) {
   }
   
   D->addAttr(::new (S.Context)
-             ConsumableAttr(Attr.getRange(), S.Context,
+             ConsumableAttr(Attr.getRange(), S.Context, DefaultState,
                             Attr.getAttributeSpellingListIndex()));
 }