AST.py

import pprint

class ASTNode():
    # Base class for a node in the AST
    # Default fields:
    #   children : a list of child nodes, these nodes might
    #       be stored in other fields of the object, we double store the pointers
    #       for easier recursion
    #   parseTree: stores the parse tree that corresponds to this AST node
    #        This is a redundancy that can be cleaned up after the AST construction,
    #        but we will keep it for easier debugging, since effeciency is not a concern here
    def __init__(self, parseTree):
        self.parseTree = parseTree
        self.children = []
        self.myType = "" # either empty string or a TypeStruct

        # reachability: None = not set, True = maybe, False = no
        self.inMaybe = None # either None or True/False
        self.outMaybe = None # either None or True/False

    # Do certains actions on every node of the AST tree
    #   call the same method in each class and its children recursively
    #   the methods that represent an action would return arguments to be used in
    #   the child nodes' method if neccessary
    def recurseAction(self, actionName, args=None):
        func = getattr(self, actionName)
        result = None
        if func and args:
            result = func(args)
        elif func:
            result = func()
        for c in self.children:
            if hasattr(c, 'recurseAction'):
                c.recurseAction(actionName, result)

    # This is a function to recursively build environment
    # Modified from the function recurseAction above, to handle the proper linking of local variable environments
    def recurseBuildEnv(self, parentEnv):
        result = self.buildEnv(parentEnv)
        preVarDcl = None
        for c in self.children:
            if c and hasattr(c, 'recurseBuildEnv'):
                if preVarDcl:
                    c.recurseBuildEnv(preVarDcl.env)
                else:
                    c.recurseBuildEnv(result)
                if c.__class__.__name__ == 'VarDclNode':
                    preVarDcl = c

    def buildEnv(self, parentEnv):
        self.env = parentEnv
        return parentEnv

    def linkType(self):
        pass

    def checkHierarchy(self):
        pass

    def disambigName(self):
        for c in self.children:
            if c and hasattr(c, 'disambigName'):
                c.disambigName()

    def checkType(self):
        # self is type correct if all its children are type correct (no exception raised)
        for c in self.children:
            if c and hasattr(c, 'checkType'):
                c.checkType()

    def staticAnalysis(self):
        for c in self.children:
            if c and hasattr(c, 'staticAnalysis'):
                c.staticAnalysis()

    # return outMaybe
    def reachCheck(self, inMaybe = True):
        if inMaybe == False:
            # error if in[s] = no for any s
            # I don't think it should get to here... but maybe?
            raise Exception("in[s] = no for a certain {}".format(type(self)))
    
        self.inMaybe = inMaybe
        self.outMaybe = self.inMaybe
        lastOut = self.inMaybe

        for c in self.children:
            if c and hasattr(c, 'reachCheck'):
                lastOut = c.reachCheck(lastOut)
                self.outMaybe = self.outMaybe and lastOut

        return self.outMaybe

    def printNodePretty(self, prefix=0):
        pp = pprint.PrettyPrinter(indent=prefix)
        pp.pprint(self.__class__.__name__)
        pp.pprint(vars(self))
        pp.pprint("-----children-----")
        prefix += 1
        return prefix

    def printTree(self):
        self.recurseAction('printNodePretty')

    def printEnv(self, prefix=0):
        pp = pprint.PrettyPrinter(indent=prefix)
        pp.pprint(self.__class__.__name__)
        if self.env:
            pp.pprint(vars(self.env))
        pp.pprint("-----children-----")
        prefix += 1
        return prefix




# Utils ######################################################

#   given a parseTree and a list of names, traverse the tree
#        to return a list of tree nodes(on the same level) that
#        has one of those names. A termination list can also be supplied
#        to stop the recursive search at the specified nodes
def getParseTreeNodes(names, tree, terminateList = []):
    result = []
    if tree.name in names:
        result.append(tree)
        return result
    if not tree.children:
        return []
    for n in tree.children:
        if n.name in names:
            result.append(n)
        elif n.name in terminateList:
            continue
        else:
            result.extend(getParseTreeNodes(names, n, terminateList))
    return result

def getASTNode(names, AST):
    result = []
    if not AST:
        return result
    if AST.__class__.__name__ in names:
        result.append(AST)
        return result
    if not AST.children:
        return []
    for n in AST.children:
        if not n:
            continue
        if n.__class__.__name__ in names:
            result.append(n)
        else:
            result.extend(getASTNode(names, n))
    return result