5 files changed, 45 insertions, 17 deletions
diff --git a/lib/Sema/Sema.cpp b/lib/Sema/Sema.cpp
index 1a9e5fb30e..13226c0b0e 100644
--- a/lib/Sema/Sema.cpp
+++ b/lib/Sema/Sema.cpp
@@ -201,9 +201,29 @@ const LangOptions &Sema::getLangOptions() const {
   return PP.getLangOptions();
 }
 
+/// getCurFunctionDecl - If inside of a function body, this returns a pointer
+/// to the function decl for the function being parsed.  If we're currently
+/// in a 'block', this returns the containing context.
+FunctionDecl *Sema::getCurFunctionDecl() {
+  DeclContext *DC = CurContext;
+  while (isa<BlockDecl>(DC))
+    DC = DC->getParent();
+  return dyn_cast<FunctionDecl>(DC);
+}
+
 ObjCMethodDecl *Sema::getCurMethodDecl() {
   DeclContext *DC = CurContext;
   while (isa<BlockDecl>(DC))
     DC = DC->getParent();
   return dyn_cast<ObjCMethodDecl>(DC);
 }
+
+NamedDecl *Sema::getCurFunctionOrMethodDecl() {
+  DeclContext *DC = CurContext;
+  while (isa<BlockDecl>(DC))
+    DC = DC->getParent();
+  if (isa<ObjCMethodDecl>(DC) || isa<FunctionDecl>(DC))
+    return cast<NamedDecl>(DC);
+  return 0;
+}
+
diff --git a/lib/Sema/Sema.h b/lib/Sema/Sema.h
index f82651bd28..ef15064638 100644
--- a/lib/Sema/Sema.h
+++ b/lib/Sema/Sema.h
@@ -329,16 +329,21 @@ public:
   void PushDeclContext(DeclContext *DC);
   void PopDeclContext();
   
-  /// CurFunctionDecl - If inside of a function body, this returns a pointer to
-  /// the function decl for the function being parsed.
-  FunctionDecl *getCurFunctionDecl() {
-    return dyn_cast<FunctionDecl>(CurContext);
-  }
-
-  /// CurMethodDecl - If inside of a method body, this returns a pointer to
-  /// the method decl for the method being parsed.
+  /// getCurFunctionDecl - If inside of a function body, this returns a pointer
+  /// to the function decl for the function being parsed.  If we're currently
+  /// in a 'block', this returns the containing context.
+  FunctionDecl *getCurFunctionDecl();
+  
+  /// getCurMethodDecl - If inside of a method body, this returns a pointer to
+  /// the method decl for the method being parsed.  If we're currently
+  /// in a 'block', this returns the containing context.
   ObjCMethodDecl *getCurMethodDecl();
 
+  /// getCurFunctionOrMethodDecl - Return the Decl for the current ObjC method
+  /// or C function we're in, otherwise return null.  If we're currently
+  /// in a 'block', this returns the containing context.
+  NamedDecl *getCurFunctionOrMethodDecl();
+
   /// Add this decl to the scope shadowed decl chains.
   void PushOnScopeChains(NamedDecl *D, Scope *S);
 
diff --git a/lib/Sema/SemaChecking.cpp b/lib/Sema/SemaChecking.cpp
index bbc50d6bac..0c7da0edb0 100644
--- a/lib/Sema/SemaChecking.cpp
+++ b/lib/Sema/SemaChecking.cpp
@@ -184,8 +184,8 @@ bool Sema::SemaBuiltinVAStart(CallExpr *TheCall) {
       // FIXME: This isn't correct for methods (results in bogus warning).
       // Get the last formal in the current function.
       const ParmVarDecl *LastArg;
-      if (getCurFunctionDecl())
-        LastArg = *(getCurFunctionDecl()->param_end()-1);
+      if (FunctionDecl *FD = getCurFunctionDecl())
+        LastArg = *(FD->param_end()-1);
       else
         LastArg = *(getCurMethodDecl()->param_end()-1);
       SecondArgIsLastNamedArgument = PV == LastArg;
diff --git a/lib/Sema/SemaExpr.cpp b/lib/Sema/SemaExpr.cpp
index 5aaa0dbc2f..ec86b9f2a6 100644
--- a/lib/Sema/SemaExpr.cpp
+++ b/lib/Sema/SemaExpr.cpp
@@ -505,14 +505,14 @@ Sema::ExprResult Sema::ActOnPredefinedExpr(SourceLocation Loc,
   }
 
   // Verify that this is in a function context.
-  if (getCurFunctionDecl() == 0 && getCurMethodDecl() == 0)
+  if (getCurFunctionOrMethodDecl() == 0)
     return Diag(Loc, diag::err_predef_outside_function);
   
   // Pre-defined identifiers are of type char[x], where x is the length of the
   // string.
   unsigned Length;
-  if (getCurFunctionDecl())
-    Length = getCurFunctionDecl()->getIdentifier()->getLength();
+  if (FunctionDecl *FD = getCurFunctionDecl())
+    Length = FD->getIdentifier()->getLength();
   else
     Length = getCurMethodDecl()->getSynthesizedMethodSize();
   
@@ -1438,7 +1438,7 @@ ActOnCompoundLiteral(SourceLocation LParenLoc, TypeTy *Ty,
                             DeclarationName()))
     return true;
 
-  bool isFileScope = !getCurFunctionDecl() && !getCurMethodDecl();
+  bool isFileScope = getCurFunctionOrMethodDecl() == 0;
   if (isFileScope) { // 6.5.2.5p3
     if (CheckForConstantInitializer(literalExpr, literalType))
       return true;
diff --git a/lib/Sema/SemaStmt.cpp b/lib/Sema/SemaStmt.cpp
index a469713121..0159f5f0a5 100644
--- a/lib/Sema/SemaStmt.cpp
+++ b/lib/Sema/SemaStmt.cpp
@@ -750,9 +750,12 @@ Sema::ActOnReturnStmt(SourceLocation ReturnLoc, ExprTy *rex) {
   Expr *RetValExp = static_cast<Expr *>(rex);
   if (CurBlock)
     return ActOnBlockReturnStmt(ReturnLoc, RetValExp);
-  QualType FnRetType =
-        getCurFunctionDecl() ? getCurFunctionDecl()->getResultType() : 
-                               getCurMethodDecl()->getResultType();
+  
+  QualType FnRetType;
+  if (FunctionDecl *FD = getCurFunctionDecl())
+    FnRetType = FD->getResultType();
+  else
+    FnRetType = getCurMethodDecl()->getResultType();
 
   if (FnRetType->isVoidType()) {
     if (RetValExp) {// C99 6.8.6.4p1 (ext_ since GCC warns)