Best Pattern for Implementing a visitor that only handles certain nodes in the tree

[jacbro2021]

Hello!

I am trying to figure out the best way to implement visitors using pyslang to handle sets of specific AST node kinds. Below is an example of what I have at the moment

import pyslang as ps 

class DirectDispatchVisitor:

    def __init__(self):
        self.lookup_table = {
            ps.SymbolKind.InstanceBody: self.handle_instancebody_sym
        }
        
    def __call__(self, node):
        if not isinstance(node, ps.Symbol):
            return
        
        handler = self.lookup_table.get(node.kind) 
        if handler:
            handler(node)

    def handle_instancebody_sym(self, node: ps.Symbol):
        print(f"--- Entering Instance Body: {node.definition.name} ---")

        ast_ctx = ps.ASTContext(node, ps.LookupLocation.max)
        context = ps.EvalContext(ast_ctx)
        for param in node.parameters:
            param_name = param.name
            
            if hasattr(param, "initializer") and param.initializer:
                param_val = param.initializer.eval(context)
                print(f"  Parameter: {param_name} = {param_val.value}")

def main():
    path, relative = CVA6_FLIST_PATH_CFG 

    driver = ps.Driver()
    driver.addStandardArgs()
    driver.processCommandFiles(path, relative, False)
    driver.processOptions()
    driver.parseAllSources()
    c = driver.createCompilation()

    v = DirectDispatchVisitor()
    c.getRoot().visit(v)

In the example above I am only visiting instance body symbols and evaluating their parameters, but there are other cases in my project where I need to handle a larger set of node kinds, hence the lookup table pattern. I've found that this pattern is much cleaner for my codebase than a giant list of if statements testing for specific node kinds, but it is much slower. Surprisingly, I believe that the slowness of this implementation is not caused by my table lookup and dispatch, but rather by the fact that slang calls into this visitor for every node in the AST (and not just the instance body that I care about in this scenario).

I have also tried to achieve the same goal via a pattern that manually walks the tree like so:

import pyslang as ps

class ConditionalVisitor:
    def visit_instance_recursive(self, node):
        print(f"--- Entering Instance Body: {node.definition.name} ---")

        ast_ctx = ps.ASTContext(node, ps.LookupLocation.max)
        context = ps.EvalContext(ast_ctx)
        for param in node.parameters:
            param_name = param.name
            
            if hasattr(param, "initializer") and param.initializer:
                param_val = param.initializer.eval(context)
                print(f"  Parameter: {param_name} = {param_val.value}")
        
        for member in node.__iter__():
            if member.kind == ps.SymbolKind.Instance:
                self.visit_instance_recursive(member.body)

def main():
    path, relative = CVA6_FLIST_PATH_CFG 

    driver = ps.Driver()
    driver.addStandardArgs()
    driver.processCommandFiles(path, relative, False)
    driver.processOptions()
    driver.parseAllSources()
    c = driver.createCompilation()

    v3 = ConditionalVisitor()
    root_body = c.getRoot().topInstances[0].body
    v3.visit_instance_recursive(root_body)

This pattern is much faster than the first one listed, but I've found that it is difficult to ensure that all of the instance bodies actually get visited due to the structure of some symbols being different from others.

Would anyone recommend a better way to implement the visitor that is also scalable for times when I need to support a larger set of Symbol kinds? Thanks!

[CheeksTheGeek]

Basically, the slowness is from the py/cpp boundary crossing that's happening on every AST logic and not your dispatch code, and I'd say there's two ways to do this

1. you pass visitStatements=False, visitExpressions=False to visit()
   for when you only care about symbol-level nodes (like InstanceBody), you skip statement and expression traversal completely

def __call__(self, node):
    if not isinstance(node, ps.Symbol):
        return
    handler = self.lookup_table.get(node.kind)
    if handler:
        handler(node)
c.getRoot().visit(v, visitStatements=False, visitExpressions=False)

You'll notice this dramatically reduces the callbacks since statements and the expressions make up the vast majority of AST nodes:
2. another way to do this is a quite-robust form of manual walking
I'd say your ConditionalVisitor approach is maybe the right idea, since the trick would be to handle all container symbols that can hold nested instances and not just InstanceBody:

class ConditionalVisitor:
    def visit_instance_recursive(self, body):
        print(f"--- entering instance body: {body.definition.name} ---")
        # ... handle params ...

        for member in body:
            if member.kind == ps.SymbolKind.Instance:
                self.visit_instance_recursive(member.body)
            elif member.kind == ps.SymbolKind.GenerateBlock:
                self._walk_generate(member)

    def _walk_generate(self, block):
        for member in block:
            if member.kind == ps.SymbolKind.Instance:
                self.visit_instance_recursive(member.body)
            elif member.kind == ps.SymbolKind.GenerateBlock:
                self._walk_generate(member)

for inst in c.getRoot().topInstances:
    v.visit_instance_recursive(inst.body)

and this covers gen blocks and multi top-level instances, and it becomes obvious this manual approach will be the fastest since you can control which subtrees to enter.
This covers generate blocks and multiple top-level instances. The manual approach will always be fastest since you control exactly which subtrees to enter.

the lookup table pattern itself is super fine since dict.get() is essentially free compared to the cross-boundary overhead.

[MikePopoloski]

Maybe the Python visitor system could be reworked to use the dict-based solution instead. Like you pass it a dict of handlers at creation time and then the C++ would only call into Python for node types you care about.

[jacbro2021]

@CheeksTheGeek thank you so much for the input! It makes sense to me that the will always be the fastest if I manage to implement the walk correctly for all of the symbols/statements/expressions I need to support in a given visitor. It seems to me that it leaves a lot of room for me to write small bugs that aren't always easy to notice.

I think if I could control which node types trigger slang to cross the cpp/python boundary with the built in visitor then the performance would be good enough for many of my use cases and I could avoid the anxiety of worrying that I incorrectly walked the tree because the visiting logic is delegated out to slang.

Modifying the call to visit to take the dictionary like this would make things better in my use cases (and possibly others as well).

import pyslang as ps 

...

lookup_table = {
    ps.SymbolKind.InstanceBody: handle_instancebody_sym
}
        
...

c.getRoot().visit(v, lookup_table)

The visitor could then walk the tree and if it encounters a node with one of the specified kinds it crosses the cpp/python boundary, otherwise it continues walking. The dict here could default to empty in which case the visitor behaves like it normally does now. There may also be further room for improvement by analyzing the kinds of nodes when the dictionary is non empty to check if all statements / expressions could be skipped over since there is already support for that as noted by @CheeksTheGeek. @MikePopoloski would this be a change you'd be open to if I started to look into it?

[MikePopoloski]

Yes, that sounds good to me.