modelverse/interface/HUTN/hutn_compiler/bootstrap_visitor.py

import string
from primitives_visitor import PrimitivesVisitor, Action

class BootstrapVisitor(PrimitivesVisitor):
    def __init__(self, args):
        PrimitivesVisitor.__init__(self, args)
        self.main = "--main" in args

        for f in args:
            if "--prepend:" in f:
                self.prepend_name = f.split(":", 1)[1]
                break
        else:
            self.prepend_name = "auto"

        self.object_symbols = {}

    def rename(self, name):
        if name == "initial_IP":
            return "%s_%s" % (self.prepend_name, name)
        else:
            return "%s_%s" % (self.prepend_name, name) if isinstance(name, int) else name

    def dump(self):
        link_id = 0

        if self.main:
            call = self.value(Action("call"))
            access = self.value(Action("access"))
            resolve = self.value(Action("resolve"))
            main = self.value("__main")
            self.dict(resolve, "var", main)
            self.dict(access, "var", resolve)
            self.dict(call, "func", access)
            self.dict(self.last_instruction, "next", call)
            ret = self.value(Action("return"))
            self.dict(call, "next", ret)

        output = []
        for t, data in self.output:
            if t == "N":
                output.append("Node %s_%s()\n" % (self.prepend_name, data))
            elif t == "V":
                name, value = data
                name = name if self.first != name else "initial_IP"
                output.append("Node %s_%s(%s)\n" % (self.prepend_name, name, value))
            elif t == "D":
                source, value, target = data
                source = source if self.first != source else "initial_IP"
                target = target if self.first != target else "initial_IP"

                source = self.rename(source)
                target = self.rename(target)
                linkname = "%s_%s_%s" % (source, link_id, target)
                link_id += 1
                output.append("Edge _%s_0(%s, %s)\n" % (linkname, source, target))
                output.append("Node _%s_2(%s)\n" % (linkname, value))
                output.append("Edge _%s_1(_%s_0, _%s_2)\n" % (linkname, linkname, linkname))
            elif t == "E":
                name, source, target = data
                source = source if self.first != source else "initial_IP"
                target = target if self.first != target else "initial_IP"
                name = self.rename(name)
                source = self.rename(source)
                target = self.rename(target)
                output.append("Edge _%s(_%s, _%s)\n" % (name, source, target))
        return ''.join(output)

    def visit_definition(self, tree):
        for a in tree.get_children("ID"):
            name = a.get_tail()[0]
            self.object_symbols[name] = True
            return PrimitivesVisitor.visit_definition(self, tree)

    def visit_vardecl(self, tree):
        if len(tree.get_tail()) > 2:
            for a in tree.get_children("ID"):
                name = a.get_tail()[0]
                self.object_symbols.setdefault(name, False)
                return PrimitivesVisitor.visit_vardecl(self, tree)
        else:
            return PrimitivesVisitor.visit_vardecl(self, tree)

    def visit_funcdecl(self, tree):
        for a in tree.get_children("func_name"):
            for b in a.get_children("ID"):
                name = b.get_tail()[0]
                if tree.get_children("func_body") or tree.get_children("ASSIGN"):
                    self.object_symbols[name] = True
                else:
                    self.object_symbols.setdefault(name, False)
                return PrimitivesVisitor.visit_funcdecl(self, tree)