1 files changed, 88 insertions, 7 deletions

diff --git a/src/jit/p.lua b/src/jit/p.lua
index 90b50bca..f3bec9f0 100644
--- a/src/jit/p.lua
+++ b/src/jit/p.lua
@@ -31,6 +31,8 @@
 --   v  Show VM states. Can be combined with stack dumps, e.g. vf or fv.
 --   z  Show zones. Can be combined with stack dumps, e.g. zf or fz.
 --   r  Show raw sample counts. Default: show percentages.
+--   a  Annotate excerpts from source code files.
+--   A  Annotate complete source code files.
 --   G  Produce output suitable for graphical tools (e.g. flame graphs).
 --   m<number> Minimum sample percentage to be shown. Default: 3.
 --   i<number> Sampling interval in milliseconds. Default: 10.
@@ -42,7 +44,8 @@ local jit = require("jit")
 assert(jit.version_num == 20100, "LuaJIT core/library version mismatch")
 local profile = require("jit.profile")
 local vmdef = require("jit.vmdef")
-local pairs, tonumber, floor, min = pairs, tonumber, math.floor, math.min
+local math = math
+local pairs, ipairs, tonumber, floor = pairs, ipairs, tonumber, math.floor
 local sort, format = table.sort, string.format
 local stdout = io.stdout
 local zone -- Load jit.zone module on demand.
@@ -54,7 +57,7 @@ local out
 
 local prof_ud
 local prof_states, prof_split, prof_min, prof_raw, prof_fmt, prof_depth
-local prof_count1, prof_count2, prof_samples
+local prof_ann, prof_count1, prof_count2, prof_samples
 
 local map_vmmode = {
   N = "Compiled",
@@ -120,15 +123,14 @@ local function prof_top(count1, count2, samples, indent)
     t[n] = k
   end
   sort(t, function(a, b) return count1[a] > count1[b] end)
-  local raw = prof_raw
   for i=1,n do
     local k = t[i]
     local v = count1[k]
     local pct = floor(v*100/samples + 0.5)
     if pct < prof_min then break end
-    if not raw then
+    if not prof_raw then
       out:write(format("%s%2d%%  %s\n", indent, pct, k))
-    elseif raw == "r" then
+    elseif prof_raw == "r" then
       out:write(format("%s%5d  %s\n", indent, v, k))
     else
       out:write(format("%s %d\n", k, v))
@@ -142,6 +144,75 @@ local function prof_top(count1, count2, samples, indent)
   end
 end
 
+-- Annotate source code
+local function prof_annotate(count1, samples)
+  local files = {}
+  local ms = 0
+  for k, v in pairs(count1) do
+    local pct = floor(v*100/samples + 0.5)
+    ms = math.max(ms, v)
+    if pct >= prof_min then
+      local file, line = k:match("^(.*):(%d+)$")
+      local fl = files[file]
+      if not fl then fl = {}; files[file] = fl; files[#files+1] = file end
+      line = tonumber(line)
+      fl[line] = prof_raw and v or pct
+    end
+  end
+  sort(files)
+  local fmtv, fmtn = " %3d%% | %s\n", "      | %s\n"
+  if prof_raw then
+    local n = math.max(5, math.ceil(math.log10(ms)))
+    fmtv = "%"..n.."d | %s\n"
+    fmtn = (" "):rep(n).." | %s\n"
+  end
+  local ann = prof_ann
+  for _, file in ipairs(files) do
+    local f0 = file:byte()
+    if f0 == 40 or f0 == 91 then
+      out:write(format("\n====== %s ======\n[Cannot annotate non-file]\n", file))
+      break
+    end
+    local fp, err = io.open(file)
+    if not fp then
+      out:write(format("====== ERROR: %s: %s\n", file, err))
+      break
+    end
+    out:write(format("\n====== %s ======\n", file))
+    local fl = files[file]
+    local n, show = 1, false
+    if ann ~= 0 then
+      for i=1,ann do
+        if fl[i] then show = true; out:write("@@ 1 @@\n"); break end
+      end
+    end
+    for line in fp:lines() do
+      if line:byte() == 27 then
+        out:write("[Cannot annotate bytecode file]\n")
+        break
+      end
+      local v = fl[n]
+      if ann ~= 0 then
+        if show then
+          if v then show = n elseif show+ann < n then show = false end
+        elseif fl[n+ann] then
+          show = n+ann
+          out:write(format("@@ %d @@\n", n))
+        end
+        if not show then goto next end
+      end
+      if v then
+        out:write(format(fmtv, v, line))
+      else
+        out:write(format(fmtn, line))
+      end
+    ::next::
+      n = n + 1
+    end
+    fp:close()
+  end
+end
+
 ------------------------------------------------------------------------------
 
 -- Finish profiling and dump result.
@@ -153,7 +224,11 @@ local function prof_finish()
       if prof_raw ~= true then out:write("[no samples collected]\n") end
       return
     end
-    prof_top(prof_count1, prof_count2, samples, "")
+    if prof_ann then
+      prof_annotate(prof_count1, samples)
+    else
+      prof_top(prof_count1, prof_count2, samples, "")
+    end
     prof_count1 = nil
     prof_count2 = nil
     prof_ud = nil
@@ -182,7 +257,13 @@ local function prof_start(mode)
   else
     prof_split = (scope == "" or mode:find("[zv].*[lfF]")) and 1 or 0
   end
-  if m.G and scope ~= "" then
+  prof_ann = m.A and 0 or (m.a and 3)
+  if prof_ann then
+    scope = "l"
+    prof_fmt = "pl"
+    prof_split = 0
+    prof_depth = 1
+  elseif m.G and scope ~= "" then
     prof_fmt = flags..scope.."Z;"
     prof_depth = -100
     prof_raw = true