###########################################################################

# 

#  Psyco profiler (Python part).

#   Copyright (C) 2001-2002  Armin Rigo et.al.

"""Psyco profiler (Python part).

The implementation of the non-time-critical parts of the profiler.

See profile() and full() in core.py for the easy interface.

"""

###########################################################################

import _psyco

from support import *

import math, time, types, atexit

now = time.time

try:

    import thread

except ImportError:

    import dummy_thread as thread

# current profiler instance

current = None

# enabled profilers, in order of priority

profilers = []

# logger module (when enabled by core.log())

logger = None

# a lock for a thread-safe go()

go_lock = thread.allocate_lock()

def go(stop=0):

    # run the highest-priority profiler in 'profilers'

    global current

    go_lock.acquire()

    try:

        prev = current

        if stop:

            del profilers[:]

        if prev:

            if profilers and profilers[0] is prev:

                return    # best profiler already running

            prev.stop()

            current = None

        for p in profilers[:]:

            if p.start():

                current = p

                if logger: # and p is not prev:

                    logger.write("%s: starting" % p.__class__.__name__, 5)

                return

    finally:

        go_lock.release()

    # no profiler is running now

    if stop:

        if logger:

            logger.writefinalstats()

    else:

        tag2bind()

atexit.register(go, 1)

def buildfncache(globals, cache):

    if hasattr(types.IntType, '__dict__'):

        clstypes = (types.ClassType, types.TypeType)

    else:

        clstypes = types.ClassType

    for x in globals.values():

        if isinstance(x, types.MethodType):

            x = x.im_func

        if isinstance(x, types.FunctionType):

            cache[x.func_code] = x, ''

        elif isinstance(x, clstypes):

            for y in x.__dict__.values():

                if isinstance(y, types.MethodType):

                    y = y.im_func

                if isinstance(y, types.FunctionType):

                    cache[y.func_code] = y, x.__name__

# code-to-function mapping (cache)

function_cache = {}

def trytobind(co, globals, log=1):

    try:

        f, clsname = function_cache[co]

    except KeyError:

        buildfncache(globals, function_cache)

        try:

            f, clsname = function_cache[co]

        except KeyError:

            if logger:

                logger.write('warning: cannot find function %s in %s' %

                             (co.co_name, globals.get('__name__', '?')), 3)

            return  # give up

    if logger and log:

        modulename = globals.get('__name__', '?')

        if clsname:

            modulename += '.' + clsname

        logger.write('bind function: %s.%s' % (modulename, co.co_name), 1)

    f.func_code = _psyco.proxycode(f)

# the list of code objects that have been tagged

tagged_codes = []

def tag(co, globals):

    if logger:

        try:

            f, clsname = function_cache[co]

        except KeyError:

            buildfncache(globals, function_cache)

            try:

                f, clsname = function_cache[co]

            except KeyError:

                clsname = ''  # give up

        modulename = globals.get('__name__', '?')

        if clsname:

            modulename += '.' + clsname

        logger.write('tag function: %s.%s' % (modulename, co.co_name), 1)

    tagged_codes.append((co, globals))

    _psyco.turbo_frame(co)

    _psyco.turbo_code(co)

def tag2bind():

    if tagged_codes:

        if logger:

            logger.write('profiling stopped, binding %d functions' %

                         len(tagged_codes), 2)

        for co, globals in tagged_codes:

            trytobind(co, globals, 0)

        function_cache.clear()

        del tagged_codes[:]

class Profiler:

    MemoryTimerResolution = 0.103

    def run(self, memory, time, memorymax, timemax):

        self.memory = memory

        self.memorymax = memorymax

        self.time = time

        if timemax is None:

            self.endtime = None

        else:

            self.endtime = now() + timemax

        self.alarms = []

        profilers.append(self)

        go()

    def start(self):

        curmem = _psyco.memory()

        memlimits = []

        if self.memorymax is not None:

            if curmem >= self.memorymax:

                if logger:

                    logger.writememory()

                return self.limitreached('memorymax')

            memlimits.append(self.memorymax)

        if self.memory is not None:

            if self.memory <= 0:

                if logger:

                    logger.writememory()

                return self.limitreached('memory')

            memlimits.append(curmem + self.memory)

            self.memory_at_start = curmem

        curtime = now()

        timelimits = []

        if self.endtime is not None:

            if curtime >= self.endtime:

                return self.limitreached('timemax')

            timelimits.append(self.endtime - curtime)

        if self.time is not None:

            if self.time <= 0.0:

                return self.limitreached('time')

            timelimits.append(self.time)

            self.time_at_start = curtime

        try:

            self.do_start()

        except error, e:

            if logger:

                logger.write('%s: disabled by psyco.error:' % (

                    self.__class__.__name__), 4)

                logger.write('    %s' % str(e), 3)

            return 0

        if memlimits:

            self.memlimits_args = (time.sleep, (self.MemoryTimerResolution,),

                                   self.check_memory, (min(memlimits),))

            self.alarms.append(_psyco.alarm(*self.memlimits_args))

        if timelimits:

            self.alarms.append(_psyco.alarm(time.sleep, (min(timelimits),),

                                            self.time_out))

        return 1

    def stop(self):

        for alarm in self.alarms:

            alarm.stop(0)

        for alarm in self.alarms:

            alarm.stop(1)   # wait for parallel threads to stop

        del self.alarms[:]

        if self.time is not None:

            self.time -= now() - self.time_at_start

        if self.memory is not None:

            self.memory -= _psyco.memory() - self.memory_at_start

        try:

            self.do_stop()

        except error:

            return 0

        return 1

    def check_memory(self, limit):

        if _psyco.memory() < limit:

            return self.memlimits_args

        go()

    def time_out(self):

        self.time = 0.0

        go()

    def limitreached(self, limitname):

        try:

            profilers.remove(self)

        except ValueError:

            pass

        if logger:

            logger.write('%s: disabled (%s limit reached)' % (

                self.__class__.__name__, limitname), 4)

        return 0

class FullCompiler(Profiler):

    def do_start(self):

        _psyco.profiling('f')

    def do_stop(self):

        _psyco.profiling('.')

class RunOnly(Profiler):

    def do_start(self):

        _psyco.profiling('n')

    def do_stop(self):

        _psyco.profiling('.')

class ChargeProfiler(Profiler):

    def __init__(self, watermark, parentframe):

        self.watermark = watermark

        self.parent2 = parentframe * 2.0

        self.lock = thread.allocate_lock()

    def init_charges(self):

        _psyco.statwrite(watermark = self.watermark,

                         parent2   = self.parent2)

    def do_stop(self):

        _psyco.profiling('.')

        _psyco.statwrite(callback = None)

class ActiveProfiler(ChargeProfiler):

    def active_start(self):

        _psyco.profiling('p')

    def do_start(self):

        self.init_charges()

        self.active_start()

        _psyco.statwrite(callback = self.charge_callback)

    def charge_callback(self, frame, charge):

        tag(frame.f_code, frame.f_globals)

class PassiveProfiler(ChargeProfiler):

    initial_charge_unit   = _psyco.statread('unit')

    reset_stats_after     = 120      # half-lives (maximum 200!)

    reset_limit           = initial_charge_unit * (2.0 ** reset_stats_after)

    def __init__(self, watermark, halflife, pollfreq, parentframe):

        ChargeProfiler.__init__(self, watermark, parentframe)

        self.pollfreq = pollfreq

        # self.progress is slightly more than 1.0, and computed so that

        # do_profile() will double the change_unit every 'halflife' seconds.

        self.progress = 2.0 ** (1.0 / (halflife * pollfreq))

    def reset(self):

        _psyco.statwrite(unit = self.initial_charge_unit, callback = None)

        _psyco.statreset()

        if logger:

            logger.write("%s: resetting stats" % self.__class__.__name__, 1)

    def passive_start(self):

        self.passivealarm_args = (time.sleep, (1.0 / self.pollfreq,),

                                  self.do_profile)

        self.alarms.append(_psyco.alarm(*self.passivealarm_args))

    def do_start(self):

        tag2bind()

        self.init_charges()

        self.passive_start()

    def do_profile(self):

        _psyco.statcollect()

        if logger:

            logger.dumpcharges()

        nunit = _psyco.statread('unit') * self.progress

        if nunit > self.reset_limit:

            self.reset()

        else:

            _psyco.statwrite(unit = nunit, callback = self.charge_callback)

        return self.passivealarm_args

    def charge_callback(self, frame, charge):

        trytobind(frame.f_code, frame.f_globals)

class ActivePassiveProfiler(PassiveProfiler, ActiveProfiler):

    def do_start(self):

        self.init_charges()

        self.active_start()

        self.passive_start()

    def charge_callback(self, frame, charge):

        tag(frame.f_code, frame.f_globals)

#

# we register our own version of sys.settrace(), sys.setprofile()

# and thread.start_new_thread().

#

def psyco_settrace(*args, **kw):

    "This is the Psyco-aware version of sys.settrace()."

    result = original_settrace(*args, **kw)

    go()

    return result

def psyco_setprofile(*args, **kw):

    "This is the Psyco-aware version of sys.setprofile()."

    result = original_setprofile(*args, **kw)

    go()

    return result

def psyco_thread_stub(callable, args, kw):

    _psyco.statcollect()

    if kw is None:

        return callable(*args)

    else:

        return callable(*args, **kw)

def psyco_start_new_thread(callable, args, kw=None):

    "This is the Psyco-aware version of thread.start_new_thread()."

    return original_start_new_thread(psyco_thread_stub, (callable, args, kw))

original_settrace         = sys.settrace

original_setprofile       = sys.setprofile

original_start_new_thread = thread.start_new_thread

sys.settrace            = psyco_settrace

sys.setprofile          = psyco_setprofile

thread.start_new_thread = psyco_start_new_thread

# hack to patch threading._start_new_thread if the module is

# already loaded

if (sys.modules.has_key('threading') and

    hasattr(sys.modules['threading'], '_start_new_thread')):

    sys.modules['threading']._start_new_thread = psyco_start_new_thread