Add option for codecov and allocation tracking to be restricted by path (#44359)

IanButterworth · web-flow · commit 166b829a208a · 2022-03-03T08:16:02.000-05:00
diff --git a/base/options.jl b/base/options.jl
@@ -20,6 +20,7 @@ struct JLOptions
     compile_enabled::Int8
     code_coverage::Int8
     malloc_log::Int8
+    tracked_path::Ptr{UInt8}
     opt_level::Int8
     opt_level_min::Int8
     debug_level::Int8
diff --git a/base/util.jl b/base/util.jl
@@ -196,6 +196,8 @@ function julia_cmd(julia=joinpath(Sys.BINDIR, julia_exename()))
                 push!(addflags, "--code-coverage=user")
             elseif opts.code_coverage == 2
                 push!(addflags, "--code-coverage=all")
+            elseif opts.code_coverage == 3
+                push!(addflags, "--code-coverage=@$(unsafe_string(opts.tracked_path))")
             end
             isempty(coverage_file) || push!(addflags, "--code-coverage=$coverage_file")
         end
@@ -204,6 +206,8 @@ function julia_cmd(julia=joinpath(Sys.BINDIR, julia_exename()))
         push!(addflags, "--track-allocation=user")
     elseif opts.malloc_log == 2
         push!(addflags, "--track-allocation=all")
+    elseif opts.malloc_log == 3
+        push!(addflags, "--track-allocation=@$(unsafe_string(opts.tracked_path))")
     end
     if opts.color == 1
         push!(addflags, "--color=yes")
diff --git a/doc/man/julia.1 b/doc/man/julia.1
@@ -191,6 +191,11 @@ Disallow or enable unsafe floating point optimizations (overrides @fastmath decl
 --code-coverage[={none*|user|all}]
 Count executions of source lines (omitting setting is equivalent to `user`)
 
+.TP
+--code-coverage=@<path>
+Count executions of source lines in a file or files under a given directory. A `@` must
+be placed before the path to indicate this option. A `@` with no path will track the current directory.
+
 .TP
  --code-coverage=tracefile.info
  Append coverage information to the LCOV tracefile (filename supports format tokens)
@@ -199,6 +204,11 @@ Count executions of source lines (omitting setting is equivalent to `user`)
 --track-allocation[={none*|user|all}]
 Count bytes allocated by each source line (omitting setting is equivalent to `user`)
 
+.TP
+--track-allocation=@<path>
+Count bytes allocated by each source line in a file or files under a given directory. A `@`
+must be placed before the path to indicate this option. A `@` with no path will track the current directory.
+
 .TP
 --bug-report=KIND
 Launch a bug report session. It can be used to start a REPL, run a script, or evaluate
diff --git a/src/codegen.cpp b/src/codegen.cpp
@@ -6976,15 +6976,20 @@ static std::pair<std::unique_ptr<Module>, jl_llvm_functions_t>
         return (!jl_is_submodule(mod, jl_base_module) &&
                 !jl_is_submodule(mod, jl_core_module));
     };
+    auto in_tracked_path = [] (StringRef file) {
+        return jl_options.tracked_path != NULL && file.startswith(jl_options.tracked_path);
+    };
     bool mod_is_user_mod = in_user_mod(ctx.module);
+    bool mod_is_tracked = in_tracked_path(ctx.file);
     struct DebugLineTable {
         DebugLoc loc;
         StringRef file;
         ssize_t line;
         bool is_user_code;
+        bool is_tracked; // falls within an explicitly set file or directory
         unsigned inlined_at;
         bool operator ==(const DebugLineTable &other) const {
-            return other.loc == loc && other.file == file && other.line == line && other.is_user_code == is_user_code && other.inlined_at == inlined_at;
+            return other.loc == loc && other.file == file && other.line == line && other.is_user_code == is_user_code && other.is_tracked == is_tracked && other.inlined_at == inlined_at;
         }
     };
     std::vector<DebugLineTable> linetable;
@@ -6997,6 +7002,7 @@ static std::pair<std::unique_ptr<Module>, jl_llvm_functions_t>
         topinfo.file = ctx.file;
         topinfo.line = toplineno;
         topinfo.is_user_code = mod_is_user_mod;
+        topinfo.is_tracked = mod_is_tracked;
         topinfo.inlined_at = 0;
         topinfo.loc = topdebugloc;
         for (size_t i = 0; i < nlocs; i++) {
@@ -7010,13 +7016,14 @@ static std::pair<std::unique_ptr<Module>, jl_llvm_functions_t>
             info.line = jl_unbox_long(jl_fieldref(locinfo, 3));
             info.inlined_at = jl_unbox_long(jl_fieldref(locinfo, 4));
             assert(info.inlined_at <= i);
+            info.file = jl_symbol_name(filesym);
+            if (info.file.empty())
+                info.file = "<missing>";
             if (module == ctx.module)
                 info.is_user_code = mod_is_user_mod;
             else
                 info.is_user_code = in_user_mod(module);
-            info.file = jl_symbol_name(filesym);
-            if (info.file.empty())
-                info.file = "<missing>";
+            info.is_tracked = in_tracked_path(info.file);
             if (ctx.debug_enabled) {
                 StringRef fname;
                 if (jl_is_method_instance(method))
@@ -7130,13 +7137,15 @@ static std::pair<std::unique_ptr<Module>, jl_llvm_functions_t>
         cursor = -1;
     };
 
-    auto do_coverage = [&] (bool in_user_code) {
+    auto do_coverage = [&] (bool in_user_code, bool is_tracked) {
         return (coverage_mode == JL_LOG_ALL ||
-                (coverage_mode == JL_LOG_USER && in_user_code));
+                (in_user_code && coverage_mode == JL_LOG_USER) ||
+                (is_tracked && coverage_mode == JL_LOG_PATH));
     };
-    auto do_malloc_log = [&] (bool in_user_code) {
+    auto do_malloc_log = [&] (bool in_user_code, bool is_tracked) {
         return (malloc_log_mode == JL_LOG_ALL ||
-                (malloc_log_mode == JL_LOG_USER && in_user_code));
+                (in_user_code && malloc_log_mode == JL_LOG_USER) ||
+                (is_tracked && malloc_log_mode == JL_LOG_PATH));
     };
     std::vector<unsigned> current_lineinfo, new_lineinfo;
     auto coverageVisitStmt = [&] (size_t dbg) {
@@ -7155,15 +7164,15 @@ static std::pair<std::unique_ptr<Module>, jl_llvm_functions_t>
             if (newdbg != current_lineinfo[dbg]) {
                 current_lineinfo[dbg] = newdbg;
                 const auto &info = linetable.at(newdbg);
-                if (do_coverage(info.is_user_code))
+                if (do_coverage(info.is_user_code, info.is_tracked))
                     coverageVisitLine(ctx, info.file, info.line);
             }
         }
         new_lineinfo.clear();
     };
     auto mallocVisitStmt = [&] (unsigned dbg, Value *sync) {
-        if (!do_malloc_log(mod_is_user_mod) || dbg == 0) {
-            if (do_malloc_log(true) && sync)
+        if (!do_malloc_log(mod_is_user_mod, mod_is_tracked) || dbg == 0) {
+            if (do_malloc_log(true, mod_is_tracked) && sync)
                 ctx.builder.CreateCall(prepare_call(sync_gc_total_bytes_func), {sync});
             return;
         }
@@ -7174,7 +7183,7 @@ static std::pair<std::unique_ptr<Module>, jl_llvm_functions_t>
     if (coverage_mode != JL_LOG_NONE) {
         // record all lines that could be covered
         for (const auto &info : linetable)
-            if (do_coverage(info.is_user_code))
+            if (do_coverage(info.is_user_code, info.is_tracked))
                 jl_coverage_alloc_line(info.file, info.line);
     }
 
@@ -7229,15 +7238,15 @@ static std::pair<std::unique_ptr<Module>, jl_llvm_functions_t>
     }
 
     Value *sync_bytes = nullptr;
-    if (do_malloc_log(true))
+    if (do_malloc_log(true, mod_is_tracked))
         sync_bytes = ctx.builder.CreateCall(prepare_call(diff_gc_total_bytes_func), {});
     { // coverage for the function definition line number
         const auto &topinfo = linetable.at(0);
         if (linetable.size() > 1) {
             if (topinfo == linetable.at(1))
                 current_lineinfo.push_back(1);
         }
-        if (do_coverage(topinfo.is_user_code))
+        if (do_coverage(topinfo.is_user_code, topinfo.is_tracked))
             coverageVisitLine(ctx, topinfo.file, topinfo.line);
     }
 
diff --git a/src/init.c b/src/init.c
@@ -583,6 +583,8 @@ static void jl_resolve_sysimg_location(JL_IMAGE_SEARCH rel)
         jl_options.machine_file = abspath(jl_options.machine_file, 0);
     if (jl_options.output_code_coverage)
         jl_options.output_code_coverage = absformat(jl_options.output_code_coverage);
+    if (jl_options.tracked_path)
+        jl_options.tracked_path = absformat(jl_options.tracked_path);
 
     const char **cmdp = jl_options.cmds;
     if (cmdp) {
diff --git a/src/jloptions.c b/src/jloptions.c
@@ -49,6 +49,7 @@ JL_DLLEXPORT void jl_init_options(void)
                         JL_OPTIONS_COMPILE_DEFAULT, // compile_enabled
                         0,    // code_coverage
                         0,    // malloc_log
+                        NULL, // tracked_path
                         2,    // opt_level
                         0,    // opt_level_min
 #ifdef JL_DEBUG_BUILD
@@ -154,11 +155,20 @@ static const char opts[]  =
     // instrumentation options
     " --code-coverage[={none*|user|all}]\n"
     "                            Count executions of source lines (omitting setting is equivalent to `user`)\n"
+    " --code-coverage=@<path>\n"
+    "                            Count executions but only in files that fall under the given file path/directory.\n"
+    "                            The `@` prefix is required to select this option. A `@` with no path will track the\n"
+    "                            current directory.\n"
+
     " --code-coverage=tracefile.info\n"
     "                            Append coverage information to the LCOV tracefile (filename supports format tokens)\n"
 // TODO: These TOKENS are defined in `runtime_ccall.cpp`. A more verbose `--help` should include that list here.
     " --track-allocation[={none*|user|all}]\n"
     "                            Count bytes allocated by each source line (omitting setting is equivalent to `user`)\n"
+    " --track-allocation=@<path>\n"
+    "                            Count bytes but only in files that fall under the given file path/directory.\n"
+    "                            The `@` prefix is required to select this option. A `@` with no path will track the\n"
+    "                            current directory.\n"
     " --bug-report=KIND          Launch a bug report session. It can be used to start a REPL, run a script, or evaluate\n"
     "                            expressions. It first tries to use BugReporting.jl installed in current environment and\n"
     "                            fallbacks to the latest compatible BugReporting.jl if not. For more information, see\n"
@@ -520,6 +530,10 @@ JL_DLLEXPORT void jl_parse_opts(int *argcp, char ***argvp)
                         codecov = JL_LOG_ALL;
                     jl_options.output_code_coverage = optarg;
                 }
+                else if (!strncmp(optarg, "@", 1)) {
+                    codecov = JL_LOG_PATH;
+                    jl_options.tracked_path = optarg + 1; // skip `@`
+                }
                 else
                     jl_errorf("julia: invalid argument to --code-coverage (%s)", optarg);
                 break;
@@ -536,6 +550,10 @@ JL_DLLEXPORT void jl_parse_opts(int *argcp, char ***argvp)
                     malloclog = JL_LOG_ALL;
                 else if (!strcmp(optarg,"none"))
                     malloclog = JL_LOG_NONE;
+                else if (!strncmp(optarg, "@", 1)) {
+                    malloclog = JL_LOG_PATH;
+                    jl_options.tracked_path = optarg + 1; // skip `@`
+                }
                 else
                     jl_errorf("julia: invalid argument to --track-allocation (%s)", optarg);
                 break;
diff --git a/src/jloptions.h b/src/jloptions.h
@@ -24,6 +24,7 @@ typedef struct {
     int8_t compile_enabled;
     int8_t code_coverage;
     int8_t malloc_log;
+    const char *tracked_path;
     int8_t opt_level;
     int8_t opt_level_min;
     int8_t debug_level;
diff --git a/src/julia.h b/src/julia.h
@@ -2075,6 +2075,7 @@ JL_DLLEXPORT int jl_generating_output(void) JL_NOTSAFEPOINT;
 #define JL_LOG_NONE 0
 #define JL_LOG_USER 1
 #define JL_LOG_ALL  2
+#define JL_LOG_PATH 3
 
 #define JL_OPTIONS_CHECK_BOUNDS_DEFAULT 0
 #define JL_OPTIONS_CHECK_BOUNDS_ON 1
diff --git a/test/cmdlineargs.jl b/test/cmdlineargs.jl
@@ -326,6 +326,39 @@ let exename = `$(Base.julia_cmd()) --startup-file=no --color=no`
         rm(covfile)
         @test occursin(expected, got) || (expected, got)
         @test_broken occursin(expected_good, got)
+
+        # Ask for coverage in specific file
+        # TODO: Figure out why asking for a specific file/dir means some lines are under-counted
+        # NOTE that a different expected reference is loaded here
+        expected = replace(read(joinpath(helperdir, "coverage_file.info.bad2"), String),
+            "<FILENAME>" => realpath(inputfile))
+        tfile = realpath(inputfile)
+        @test readchomp(`$exename -E "(Base.JLOptions().code_coverage, unsafe_string(Base.JLOptions().tracked_path))" -L $inputfile
+            --code-coverage=$covfile --code-coverage=@$tfile`) == "(3, $(repr(tfile)))"
+        @test isfile(covfile)
+        got = read(covfile, String)
+        rm(covfile)
+        @test occursin(expected, got) || (expected, got)
+        @test_broken occursin(expected_good, got)
+
+        # Ask for coverage in directory
+        tdir = dirname(realpath(inputfile))
+        @test readchomp(`$exename -E "(Base.JLOptions().code_coverage, unsafe_string(Base.JLOptions().tracked_path))" -L $inputfile
+            --code-coverage=$covfile --code-coverage=@$tdir`) == "(3, $(repr(tdir)))"
+        @test isfile(covfile)
+        got = read(covfile, String)
+        rm(covfile)
+        @test occursin(expected, got) || (expected, got)
+        @test_broken occursin(expected_good, got)
+
+        # Ask for coverage in a different directory
+        tdir = mktempdir() # a dir that contains no code
+        @test readchomp(`$exename -E "(Base.JLOptions().code_coverage, unsafe_string(Base.JLOptions().tracked_path))" -L $inputfile
+            --code-coverage=$covfile --code-coverage=@$tdir`) == "(3, $(repr(tdir)))"
+        @test isfile(covfile)
+        got = read(covfile, String)
+        @test isempty(got)
+        rm(covfile)
     end
 
     # --track-allocation
diff --git a/test/testhelpers/coverage_file.info.bad2 b/test/testhelpers/coverage_file.info.bad2
@@ -0,0 +1,20 @@
+SF:<FILENAME>
+DA:3,1
+DA:4,1
+DA:5,0
+DA:7,1
+DA:8,1
+DA:9,3
+DA:10,5
+DA:11,0
+DA:12,1
+DA:14,0
+DA:17,1
+DA:18,0
+DA:19,0
+DA:20,0
+DA:22,1
+DA:1234,0
+LH:9
+LF:16
+end_of_record
