-- Lua profiler

-- Relies on 'clock()' to tell it CPU seconds used since startup

-- All profiler related stuff is stored in the top level table '_profiler'

_profiler = {}

-- Calling this will return a new profiler object which can be managed
function newProfiler()
   local newprof = {}
   for k,v in _profiler do
      newprof[k] = v;
   end
   return newprof
end

-- This function starts the profiler.
-- It will do nothing if this (or any other) profiler is already running
function _profiler.start(self)
   if( _profiler.running ) then
      return
   end
   -- Start the profiler. This begins by setting up internal profiler state
   _profiler.running = self;
   self.rawstats = {}
   self.callstack = {}
   setcallhook( _profiler_hook_wrapper )
end

-- This function stops the profiler
-- It will do nothing if a profiler is not running, and if it isn't the profiler running
function _profiler.stop(self)
   if( _profiler.running ~= self ) then
      return
   end
   -- Stop the profiler.
   setcallhook( nil )
   _profiler.running = nil
end

-- Simple wrapper to handle the hook
function _profiler_hook_wrapper(action)
   if( _profiler.running == nil ) then
      setcallhook( nil );
   end
   _profiler.running:_internal_profile(action)
end

-- This is the main internal function of the profiler and should not be called except by the hook wrapper
function _profiler._internal_profile(self,action)
   -- First entry on the stack will be the hook wrapper, ensure this...
   if( getinfo( 2, "n" ).name ~= "_profiler_hook_wrapper" ) then
      print("Profiler's internal profiler function called by: ",getinfo( 2, "n" ).name )
      exit(1);
   end
   -- Since we can obtain the 'function' for the item we've had call us, we can use that...
   local caller_info = getinfo( 3 );
   if( caller_info == nil ) then
      print "No caller_info";
      return
   end

   -- Retrieve the most recent activation record...
   local latest_ar = nil;
   if( getn(self.callstack) > 0 ) then
      latest_ar = self.callstack[getn(self.callstack)];
   end

   -- Are we allowed to profile this function?
   local should_not_profile = 0
   for k,v in self.prevented_functions do
      if( k == caller_info.func ) then
	 should_not_profile = v
      end
   end
   -- Also check the top activation record...
   if( latest_ar ) then
      if( latest_ar.should_not_profile == 2 ) then
	 should_not_profile = 2
      end
   end

   -- Now then, are we in 'call' or 'return' ?
   -- print("Profile:", caller_info.name, "SNP:", should_not_profile, "Action:", action );

   if( action == "call" ) then
      -- Making a call...
      local this_ar = {}
      this_ar.should_not_profile = should_not_profile;
      this_ar.parent_ar = latest_ar
      this_ar.anon_child = 0
      this_ar.name_child = 0
      this_ar.children = {}
      this_ar.children_time = {}
      this_ar.clock_start = clock()
      -- Last thing to do on a call is to insert this onto the ar stack...
      tinsert( self.callstack, this_ar );
   else
      local this_ar = latest_ar;
      if( this_ar == nil ) then
	 return; -- No point in doing anything if there's no upper activation record
      end

      -- Right, calculate the time in this function...
      this_ar.clock_end = clock()
      this_ar.this_time = this_ar.clock_end - this_ar.clock_start

      -- Now, if we have a parent, update its call info...
      if( this_ar.parent_ar ) then
	 this_ar.parent_ar.children[caller_info.func] = (this_ar.parent_ar.children[caller_info.func] or 0) + 1;
	 this_ar.parent_ar.children_time[caller_info.func] = (this_ar.parent_ar.children_time[caller_info.func] or 0 ) + this_ar.this_time;
	 if( caller_info.name == nil ) then
	    this_ar.parent_ar.anon_child = this_ar.parent_ar.anon_child + this_ar.this_time
	 else
	    this_ar.parent_ar.name_child = this_ar.parent_ar.name_child + this_ar.this_time
	 end
      end
      -- Now if we're meant to record information about ourselves, do so...
      if( this_ar.should_not_profile == 0 ) then
	 local inforec = self:_get_func_rec(caller_info.func,1);
	 inforec.count = inforec.count + 1;
	 inforec.time = inforec.time + this_ar.this_time
	 inforec.anon_child_time = inforec.anon_child_time + this_ar.anon_child
	 inforec.name_child_time = inforec.name_child_time + this_ar.name_child
	 for k,v in this_ar.children do
	    inforec.children[k] = (inforec.children[k] or 0) + v;
	    inforec.children_time[k] = (inforec.children_time[k] or 0) + this_ar.children_time[k];
	 end
      end

      -- Last thing to do on return is to drop the last activation record...
      tremove( self.callstack, getn( self.callstack ) );
   end

end


-- This returns a (possibly empty) function record for the specified function
function _profiler._get_func_rec(self,func,force)
   -- Find the function ref for the 'func' (if force and not present, create one)
   local ret = self.rawstats[func];
   if( ret == nil and force ~= 1 ) then
      return nil
   end
   if( ret == nil ) then
      -- Build a new function statistics table
      ret = {}
      ret.func = func;
      ret.count = 0;
      ret.time = 0;
      ret.anon_child_time = 0;
      ret.name_child_time = 0;
      ret.children = {}      
      ret.children_time = {}
      self.rawstats[func] = ret;
   end
   return ret
end

-- This returns a report in a string for you to do with what you will
function _profiler.report(self,outfile)
   write( outfile, [[Lua Profile output created by profiler.lua. Copyright Pepperfish 2002+

]] )
   local total_time = 0;
   local ordering = {}
   for func,record in self.rawstats do
      tinsert(ordering, func);
   end

   sort( ordering, function(a,b)
		      local arec,brec = %self.rawstats[a],%self.rawstats[b];
		      local atime = arec.time - ( arec.anon_child_time + arec.name_child_time )
		      local btime = brec.time - ( brec.anon_child_time + brec.name_child_time )
		      return atime > btime
		   end );

   for i=1,getn(ordering) do
      local func = ordering[i]
      local record = self.rawstats[func]
      local thisfuncname = " " .. self:_pretty_name(func) .. " ";
      if( strlen( thisfuncname ) < 42 ) then
	 thisfuncname = strrep( "-", (42 - strlen(thisfuncname))/2 ) .. thisfuncname;
	 thisfuncname = thisfuncname .. strrep( "-", 42 - strlen(thisfuncname) );
      end

      total_time = total_time + ( record.time - ( record.anon_child_time + record.name_child_time ) )
      write( outfile, strrep( "-", 19 ) .. thisfuncname .. strrep( "-", 19 ) .. "\n" )
      write( outfile, "                 Call count: " .. format( "%4d", record.count ) .. "\n" )
      write( outfile, "                 Total time: " .. format( "%4.3f", record.time ) .. "s\n" )
      write( outfile, "     Time spent in children: " .. format( "%4.3f", record.anon_child_time + record.name_child_time ) .. "s\n" )
      write( outfile, "         Time spent in self: " .. format( "%4.3f", record.time - ( record.anon_child_time + record.name_child_time ) ) .. "s\n" )
      write( outfile, "        Time spent per call: " .. format( "%4.5f", record.time/record.count ) .. "s/call\n" )
      write( outfile, "Time spent in self per call: " .. format( "%4.5f", (record.time - ( record.anon_child_time + record.name_child_time ))/record.count ) .. "s/call\n" )

      -- Report on each child in the form
      -- Child  <funcname> called n times and took a.bs
      local added_blank = 0
      for k,v in record.children do
	 if( self.prevented_functions[k] == nil or self.prevented_functions[k] == 0 ) then
	    if( added_blank == 0 ) then
	       write( outfile, "\n" ); -- extra separation line
	       added_blank = 1;
	    end
	    write( outfile, "Child " .. self:_pretty_name(k) .. strrep( " ", 41-strlen(self:_pretty_name(k)) ) .. " called " .. format("%6d", v) )
	    write( outfile, " times. Took " .. format("%4.2f", record.children_time[k] ) .. "s\n" )
	 end
      end

      write( outfile, "\n" ); -- extra separation line
      flush( outfile )
   end
   write( outfile, "\n\n" )
   write( outfile, "Total time spent in profiled functions: " .. format("%5.3g",total_time) .. "s\n" )
   write( outfile, [[

END
]] )
   flush( outfile )
end

-- Dump the output as loadable Lua source
function _profiler.lua_report(self,outfile)
   -- Purpose: Write out the entire raw state in a cross-referenceable form.
   local ordering = {}
   local functonum = {}
   for func,record in self.rawstats do
      tinsert(ordering, func);
      functonum[func] = getn(ordering);
   end

   write( outfile, "-- Profile generated by profiler.lua Copyright Pepperfish 2002+\n\n" );
   write( outfile, "-- Function names\nfuncnames = {}\n" );
   for i=1,getn(ordering) do
      local thisfunc = ordering[i];
      write( outfile, "funcnames[" .. i .. "] = " .. format("%q", self:_pretty_name(thisfunc)) .. "\n" );
   end
   write( outfile, "\n" );
   write( outfile, "-- Function times\nfunctimes = {}\n" );
   for i=1,getn(ordering) do
      local thisfunc = ordering[i];
      local record = self.rawstats[thisfunc];
      write( outfile, "functimes[" .. i .. "] = { " );
      write( outfile, "tot=" .. record.time .. ", " );
      write( outfile, "achild=" .. record.anon_child_time .. ", " );
      write( outfile, "nchild=" .. record.name_child_time .. ", " );
      write( outfile, "count=" .. record.count .. " }\n" );
   end
   write( outfile, "\n" );
   write( outfile, "-- Child links\nchildren = {}\n" );
   for i=1,getn(ordering) do
      local thisfunc = ordering[i];
      local record = self.rawstats[thisfunc];
      write( outfile, "children[" .. i .. "] = { " );
      for k,v in record.children do
	 if( functonum[k] ) then -- non-recorded functions will be ignored now
	    write( outfile, functonum[k] .. ", " );
	 end
      end
      write( outfile, "}\n" );
   end
   write( outfile, "\n" );
   write( outfile, "-- Child call counts\nchildcounts = {}\n" );
   for i=1,getn(ordering) do
      local thisfunc = ordering[i];
      local record = self.rawstats[thisfunc];
      write( outfile, "children[" .. i .. "] = { " );
      for k,v in record.children do
	 if( functonum[k] ) then -- non-recorded functions will be ignored now
	    write( outfile, v .. ", " );
	 end
      end
      write( outfile, "}\n" );
   end
   write( outfile, "\n" );
   write( outfile, "-- Child call time\nchildtimes = {}\n" );
   for i=1,getn(ordering) do
      local thisfunc = ordering[i];
      local record = self.rawstats[thisfunc];
      write( outfile, "children[" .. i .. "] = { " );
      for k,v in record.children do
	 if( functonum[k] ) then -- non-recorded functions will be ignored now
	    write( outfile, record.children_time[k] .. ", " );
	 end
      end
      write( outfile, "}\n" );
   end
   write( outfile, "\n\n-- That is all.\n\n" );
   flush(outfile);
end

-- Internal function to calculate a pretty name for the profile output
function _profiler._pretty_name(self,func)
   local info = getinfo(func);
   local name = "";
   if( info.what == "Lua" ) then
      name = "L:"
   end
   if( info.what == "C" ) then
      name = "C:"
   end
   if( info.what == "main" ) then
      name = " :"
   end
   
   if( info.name == nil ) then
      name = name .. "<"..tostring(func) .. ">"
   else
      name = name .. info.name
   end

   if( strsub( info.source, 1, 1 ) == "@" ) then
      name = name .. info.source
   else
      if( info.what == "C" ) then
	 name = name .. "@?"
      else
	 name = name .. "@<string>"
      end
   end
   name = name .. ":"
   if( info.what == "C" ) then
      name = name .. "?"
   else
      name = name .. info.linedefined
   end

   return name
end

-- Finally, configure the prevented functions

function _profiler.prevent(self, func, level)
   self.prevented_functions[func] = (level or 1)
end

-- 1 == don't profile me
-- 2 == don't profile me or my kids

_profiler.prevented_functions = {
   [ _profiler.start             ] = 2,
   [ _profiler.stop              ] = 2,
   [ _profiler._internal_profile ] = 2,
   [ _profiler_hook_wrapper      ] = 2,
   [ _profiler.prevent           ] = 2,
   [ _profiler._get_func_rec     ] = 2,
   [ _profiler.report            ] = 2,
   [ _profiler.lua_report        ] = 2,
   [ _profiler._pretty_name      ] = 2
}