Add new `knock' protocol.

[tripe] / py / rmcr.py

### -*-python-*-
###
### Rich man's coroutines
###
### (c) 2006 Straylight/Edgeware
###

###----- Licensing notice ---------------------------------------------------
###
### This file is part of Trivial IP Encryption (TrIPE).
###
### TrIPE is free software: you can redistribute it and/or modify it under
### the terms of the GNU General Public License as published by the Free
### Software Foundation; either version 3 of the License, or (at your
### option) any later version.
###
### TrIPE is distributed in the hope that it will be useful, but WITHOUT
### ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
### FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
### for more details.
###
### You should have received a copy of the GNU General Public License
### along with TrIPE.  If not, see <https://www.gnu.org/licenses/>.

__pychecker__ = 'self=me'

###--------------------------------------------------------------------------
### External dependencies.

import thread as T
from sys import exc_info, excepthook

###--------------------------------------------------------------------------
### What's going on?
###
### A coroutine is just a thread.  Each coroutine has a recursive lock
### associated with it.  The lock is almost always held.  In order to switch
### to a coroutine, one releases its lock and then (re)locks one's own.

###--------------------------------------------------------------------------
### Low-level machinery.

__debug = False
def _debug(msg):
  if __debug:
    print '+++ %s: %s' % (T.get_ident(), msg)

def _switchto(cr, arg = None, exc = None):
  """Switch to coroutine CR."""
  global active
  _debug('> _switchto(%s, %s, %s)' % (cr, arg, exc))
  if not cr.livep:
    raise ValueError('coroutine is dead')
  cr._arg = arg
  cr._exc = exc
  if cr is active:
    _debug('  _switchto: self-switch to %s' % cr)
  else:
    _debug('  _switchto: switch from %s to %s' % (active, cr))
    active = cr
    _debug('  _switchto: %s releasing' % cr)
    assert cr._lk.locked()
    cr._lk.release()
  _debug('< _switchto')

###--------------------------------------------------------------------------
### Coroutine object.

def findvictim(cr):
  """Find an appropriate victim coroutine for something, starting from CR."""
  while not cr:
    if cr is None:
      return main
    cr = cr.parent
  return cr

class Coroutine (object):
  """Heard of lightweight threads?  Well, this is a heavyweight coroutine."""

  def __init__(me, func = None, name = None, parent = None, __tid = None):
    """
    Create a new coroutine object.

    The new coroutine is immediately activated.  FUNC is the function to run,
    and defaults to the coroutine object's `run' method, so subclassing is a
    reasonable thing to do; NAME is a friendly name for the coroutine, and
    shows up in debug traces.  The __TID argument is used internally for
    special effects, and shouldn't be used by external callers.
    """
    global active
    _debug('> __init__(%s, func = %s, tid = %s)' % (name, func, __tid))
    me.name = name
    me._lk = T.allocate_lock()
    _debug('  __init__: %s locking' % me)
    me._lk.acquire()
    me.livep = True
    me._exc = None
    me._arg = None
    me.parent = parent or active
    me._onexit = [None, None]
    if __tid is not None:
      me._tid = __tid
      _debug('  __init__: manufacture cr %s with existing thread %d' %
             (me, __tid))
      me._startp = False
    else:
      me._func = func or me.run
      me._tid = T.start_new_thread(me._start, ())
      me._startp = True
      assert me._lk.locked()
      _debug('  __init__: create %s with new thread %d' % (me, me._tid))
    _debug('< __init__(%s)' % me)

  def __str__(me):
    """Stringify a coroutine using its name if possible."""
    if me.name is not None:
      return '<Coroutine %s>' % me.name
    else:
      return repr(me)

  def _start(me):
    """
    Start up the coroutine.

    Wait for this coroutine to become active, and then run the user's
    function.  When (if) that returns, mark the coroutine as dead.
    """
    _debug('> _start(%s)' % (me))
    args, kwargs = me._wait()
    _debug('  start(%s): args = %s, kwargs = %s' % (me, args, kwargs))
    me._startp = False
    try:
      try:
        _debug('  _start(%s): call user (args = %s, kwargs = %s)' %
               (me, args, kwargs))
        me._onexit = [me._func(*args, **kwargs), None]
      except:
        exc = exc_info()
        _debug('  _start(%s): caught exception (%s)' % (me, exc))
        me._onexit = [None, exc]
    finally:
      _debug('  _start(%s): finally' % me)
      _debug('  _start(%s): _onexit = %s' % (me, me._onexit))
      me.livep = False
      _switchto(findvictim(me.parent), *me._onexit)
    _debug('< _start(%s)' % me)

  def _wait(me):
    """Wait for this coroutine to become active."""
    global active
    _debug('> _wait(%s)' % me)
    me._lk.acquire()
    while me is not active:
      _debug('  _wait(%s): locking' % me)
      me._lk.acquire()
    _debug('  _wait(%s): active, arg = %s, exc = %s' %
           (me, me._arg, me._exc))
    if me._exc:
      raise me._exc[0], me._exc[1], me._exc[2]
    _debug('< _wait(%s): %s' % (me, me._arg))
    return me._arg

  def switch(me, *args, **kwargs):
    """Switch to this coroutine, passing it the object OBJ."""
    global active
    _debug('> switch(%s, args = %s, kwargs = %s)' % (me, args, kwargs))
    if me._startp:
      obj = args, kwargs
    else:
      obj, = args or (None,)
      assert not kwargs
    old = active
    _switchto(me, obj)
    _debug('< switch')
    return old._wait()

  def getcurrent():
    return active
  getcurrent = staticmethod(getcurrent)

  def __nonzero__(me):
    return me.livep

  def throw(me, exc, arg = None, tb = None):
    """
    Switch to this coroutine, throwing it an exception.

    The exception is given by EXC, ARG and TB, which form the exception,
    argument, traceback triple.
    """
    global active
    _debug('> throw(%s, %s, args = %s)' % (me, exc, arg))
    old = active
    me._exc = [exc, arg, tb]
    _switchto(me, None)
    _debug('< throw')
    return old._wait()

  def run(me, *args, **kw):
    raise Exception('Coroutine has no body to run')

###--------------------------------------------------------------------------
### Setup stuff.

active = None
main = Coroutine(_Coroutine__tid = T.get_ident(), name = '_main')
active = main

###----- That's all, folks --------------------------------------------------