3 files changed, 23 insertions, 6 deletions
diff --git a/lib/CodeGen/CodeGenFunction.cpp b/lib/CodeGen/CodeGenFunction.cpp
index 75c60edbba..a3c39649b7 100644
--- a/lib/CodeGen/CodeGenFunction.cpp
+++ b/lib/CodeGen/CodeGenFunction.cpp
@@ -222,8 +222,18 @@ void CodeGenFunction::FinishFunction(SourceLocation EndLoc) {
   // Emit function epilog (to return).
   EmitReturnBlock();
 
-  if (ShouldInstrumentFunction())
-    EmitFunctionInstrumentation("__cyg_profile_func_exit");
+  if (ShouldInstrumentFunction()) {
+      // The size of the function isn't known during StartFunction, so in order
+      // to instrument selectively based on function size, we need to wait
+      // until the end and insert both the entry and exit instrumentation
+      llvm::BasicBlock *CurBlock = Builder.GetInsertBlock();
+      llvm::BasicBlock *BeginBlock = &CurFn->getEntryBlock();
+
+      Builder.SetInsertPoint(BeginBlock, BeginBlock->begin());
+      EmitFunctionInstrumentation("__cyg_profile_func_enter");
+      Builder.SetInsertPoint(CurBlock);
+      EmitFunctionInstrumentation("__cyg_profile_func_exit");
+  }
 
   // Emit debug descriptor for function end.
   if (CGDebugInfo *DI = getDebugInfo()) {
@@ -274,7 +284,10 @@ bool CodeGenFunction::ShouldInstrumentFunction() {
     return false;
   if (!CurFuncDecl || CurFuncDecl->hasAttr<NoInstrumentFunctionAttr>())
     return false;
-  return true;
+
+  // If not specified, defaults to 0.
+  int Size = CGM.getCodeGenOpts().InstrumentFunctionsSize;
+  return CurFn->getBasicBlockList().size() >= Size;
 }
 
 /// EmitFunctionInstrumentation - Emit LLVM code to call the specified
@@ -541,9 +554,6 @@ void CodeGenFunction::StartFunction(GlobalDecl GD,
     DI->EmitFunctionStart(GD, FnType, CurFn, Builder);
   }
 
-  if (ShouldInstrumentFunction())
-    EmitFunctionInstrumentation("__cyg_profile_func_enter");
-
   if (CGM.getCodeGenOpts().InstrumentForProfiling)
     EmitMCountInstrumentation();
 
diff --git a/lib/Driver/Tools.cpp b/lib/Driver/Tools.cpp
index aba1fe4d2d..c04fc13b23 100644
--- a/lib/Driver/Tools.cpp
+++ b/lib/Driver/Tools.cpp
@@ -2462,6 +2462,11 @@ void Clang::ConstructJob(Compilation &C, const JobAction &JA,
 
   Args.AddAllArgs(CmdArgs, options::OPT_finstrument_functions);
 
+  if (Arg *A = Args.getLastArg(options::OPT_finstrument_functions_size_EQ)) {
+    CmdArgs.push_back("-finstrument-functions-size");
+    CmdArgs.push_back(A->getValue());
+  }
+
   if (Args.hasArg(options::OPT_ftest_coverage) ||
       Args.hasArg(options::OPT_coverage))
     CmdArgs.push_back("-femit-coverage-notes");
diff --git a/lib/Frontend/CompilerInvocation.cpp b/lib/Frontend/CompilerInvocation.cpp
index 42ea96f0f2..06d8b12e6c 100644
--- a/lib/Frontend/CompilerInvocation.cpp
+++ b/lib/Frontend/CompilerInvocation.cpp
@@ -407,6 +407,8 @@ static bool ParseCodeGenArgs(CodeGenOptions &Opts, ArgList &Args, InputKind IK,
   }
 
   Opts.InstrumentFunctions = Args.hasArg(OPT_finstrument_functions);
+  Opts.InstrumentFunctionsSize =
+    Args.getLastArgIntValue(OPT_finstrument_functions_size, 0);
   Opts.InstrumentForProfiling = Args.hasArg(OPT_pg);
   Opts.EmitOpenCLArgMetadata = Args.hasArg(OPT_cl_kernel_arg_info);
   Opts.DebugCompilationDir = Args.getLastArgValue(OPT_fdebug_compilation_dir);