modelverse/kernel/modelverse_kernel/request_handler.py

import sys
import modelverse_kernel.primitives as primitive_functions
import modelverse_kernel.jit as jit
from collections import defaultdict

class RequestHandler(object):
    def __init__(self):
        self.generator_stack = []
        self.handlers = {"CALL": lambda i: i[0],
                         "CALL_ARGS": lambda i: i[0](*(i[1])),
                         "CALL_KWARGS": lambda i: i[0](**(i[1])),
                         "RETURN": self.execute_return,
                         "SLEEP": self.execute_sleep}

    def push_generator(self, gen):
        self.generator_stack.append(gen)

    def handle_request(self, reply):
        self.reply = reply
        while self.generator_stack:
            try:
                requests = self.generator_stack[-1].send(self.reply)

                # Generated new request, so process
                if requests and requests[0][0] in self.handlers:
                    self.generator_stack.append(None)
                    # This next command potentially raises a finished message already, meaning that we should stop already
                    # We avoid an extra try/except block by putting the None on the stack already
                    self.reply = None
                    self.generator_stack[-1] = self.handlers[requests[0][0]](requests[0][1])
                else:
                    # MvS request, so forward that instead
                    return requests

            except StopIteration:
                # Exception, so finished execution of this generator, passing on None to the caller
                del self.generator_stack[-1]
                self.reply = [None]
                
            except primitive_functions.SleepKernel:
                # Processing sleep, so pop its generator and reraise
                del self.generator_stack[-1]
                raise

            except primitive_functions.PrimitiveFinished as ex:
                # Exception, so finished execution of this generator, passing on ex.result to the caller
                del self.generator_stack[-1]
                self.reply = [ex.result]

    def execute_sleep(self, request_args):
        raise primitive_functions.SleepKernel(request_args[0], request_args[1])

    def execute_return(self, request_args):
        del self.generator_stack[-1]
        del self.generator_stack[-1]
        self.reply = [request_args[0]]
        return self.generator_stack[-1]