import cPickle as pickle
import stackless
import sys
import signal

"""
pickled-sieve.py
Andrew Francis
October 30th, 2009

pickled-sieve.py is an implementation of the Sieve of Eratosthenes
with a twist: when Control-C is pressed, all the tasklets (one for each
prime) and associated channels are pickled to a file sieve.dat.

in the future, if pickle-sieve.py is executed with a file name: 

$python pickle-sieve.py sieve.dat

provided the file format is correct, all the tasklets and their associated
channels are deserialised and the programme resumes where it left off. This
is a good way of computing thousands of primes over multiple sessions and
on different machines.

<song> Almost Crimes - Broken Social Scene</song>
"""

flag = False

"""
we want to do the serialisation in the main tasklet's loop
"""
def signalHandler(signal, frame):
    global flag
    flag = True


"""
image is used to  stores tasklets and their associated channels. The 
channels are serialised separately because when a blocked tasklet is 
serialised, the channel is not included. 
"""
class Image(object):
    def __init__(self):
        self.channels = []
        self.tasklets = []


    # record tasklets 
    def makeTasklet(self, f, *args, **kwargs):
        t = stackless.tasklet(f)(*args, **kwargs)
        self.tasklets.append(t)
        return t


    # record channels
    def makeChannel(self):
        channel = stackless.channel()
        self.channels.append(channel)
        return channel


"""
the sieve proper
"""
def filter(prime, listen, send):
    while (True):
       i = listen.receive()
       if (i % prime):
          send.send(i)

       
def counter(c):
    i = 2
    while (True):
        c.send(i)
        i += 1


def sieve(prime, c):
    global image
    while (True):
        p = c.receive()
        prime.send(p)
        newc = image.makeChannel()
        image.makeTasklet(filter, p, c, newc) 
        c = newc

if __name__ == "__main__":

   signal.signal(signal.SIGINT, signalHandler)

   if len(sys.argv) == 2:
      fd = open('sieve.dat', 'r')
      image = pickle.load(fd)
      for t in image.tasklets:
          """
          we have to make sure that we are not inserting tasklets that
          were blocked.
          """
          if not t.blocked:
             t.insert()
      stackless.schedule()
   else:
      image = Image()
      image.makeChannel()
      image.makeChannel()
      image.makeTasklet(counter, image.channels[1])   
      image.makeTasklet(sieve, image.channels[0], image.channels[1])

   while(True):
      if (flag):
         fd = open('sieve.dat','w')
         pickle.dump(image, fd)
         fd.close()
         print "pickled ", len(image.tasklets),"and ", len(image.channels)
         break
      else:
         print image.channels[0].receive()