kripkomat/astvisitors.py

from pycparser import c_ast
from itertools import pairwise

def path_to_str(p):
    return "->".join([str(n.__class__.__name__) for n in p])

def path_select_last(classname, p):
    for n in reversed(p):
        if n.__class__.__name__ == classname:
            return n
    return None

def path_contains(p, t):
    for n in p:
        if isinstance(n, t):
            return True
    return False

def path_select_parent(child, p):
    parent = None
    for n in p:
        if n == child:
            return parent
        parent = n
    return None

def path_slice_back(node, p):
    s = []
    for n in reversed(p):
        if n == node:
            break
        s.append(n)
    return [x for x in reversed(s)]


def path_filter(p, t):
    elems = []
    for n in p:
        if isinstance(n, t):
            elems.append(n)
    return elems

class ExprListSerializerVisitor(c_ast.NodeVisitor):
    def __init__(self):
        self.serial = []

    def visit_Constant(self, node):
        expr = node.value
        self.serial.append(expr)
    
    def visit_ID(self, node):
        expr = node.name
        self.serial.append(expr)

    #todo: expand

class CaseLabelExtractionVisitor(c_ast.NodeVisitor):
    def __init__(self):
        self.label = None

    def visit_ID(self, node):
        self.label = node.name


def expr_list_to_str(exprl):
    elsv = ExprListSerializerVisitor()
    elsv.visit(exprl)
    return ','.join(elsv.serial)

class NodeVisitorWithParent(object):
    def __init__(self):
        self.current_parent = None

    def visit(self, node):
        """ Visit a node.
        """
        method = 'visit_' + node.__class__.__name__
        visitor = getattr(self, method, self.generic_visit)
        return visitor(node)

    def generic_visit(self, node):
        """ Called if no explicit visitor function exists for a
            node. Implements preorder visiting of the node.
        """
        oldparent = self.current_parent
        self.current_parent = node
        for c in node.children():
            self.visit(c)
        self.current_parent = oldparent

class NodeVisitorFuncCallForward(object):
    def __init__(self, ast):
        self.current_parent = None
        self.ast = ast
        fdv = FuncDefVisitor()
        fdv.visit(ast)
        self.func_table = fdv.func_table
    
    def visit_FuncCall(self, node):
        print("Visiting FuncCall")
        print(node.show())
        print('---- parent ----')
        print(self.current_parent.show())

    def visit(self, node):
        """ Visit a node.
        """
        method = 'visit_' + node.__class__.__name__
        visitor = getattr(self, method, self.generic_visit)
        return visitor(node)

    def generic_visit(self, node):
        """ Called if no explicit visitor function exists for a
            node. Implements preorder visiting of the node.
        """
        oldparent = self.current_parent
        self.current_parent = node
        for c in node.children():
            self.visit(c)
        self.current_parent = oldparent

class FuncDefVisitor(c_ast.NodeVisitor):
    def __init__(self):
        self.func_table = {}

    def visit_FuncDef(self, node):
        self.func_table[node.decl.name] = node

class FuncCallVisitor(c_ast.NodeVisitor):
    def __init__(self):
        self.func_calls = []

    def visit_FuncCall(self, node):
        self.func_calls.append(node.children()[0][1].name)

class StateAssignmentVisitor(NodeVisitorFuncCallForward):
    def __init__(self, ast, state):
        super().__init__(ast)
        self._method_cache = {}
        self.state = state
        self.assignments = []

    def visit(self, node, path = None, invariants = None):
        """ Visit a node.
        """
        if path is None:
            path = []
        if invariants is None:
            invariants = []
        if self._method_cache is None:
            self._method_cache = {}

        visitor = self._method_cache.get(node.__class__.__name__, None)
        if visitor is None:
            method = 'visit_' + node.__class__.__name__
            visitor = getattr(self, method, self.generic_visit)
            self._method_cache[node.__class__.__name__] = visitor

        return visitor(node, path, invariants)
    
    def generic_visit(self, node, path, invariants):
        """ Called if no explicit visitor function exists for a
            node. Implements preorder visiting of the node.
        """
        path = path.copy()
        path.append(node)
        for c in node:
            self.visit(c, path, invariants)

    def visit_FuncCall(self, node, path, invariants):
        fcall = node.name.name
        #print("CAL path", path_to_str(path))
        cases = path_filter(path, c_ast.Case)
        #print(cases)
        new_invariants = [x.expr.name for x in cases]
        invariants = invariants.copy()
        for ni in new_invariants:
            if not(ni in invariants):
                invariants.append(ni)
        #print("invariants:", invariants)
        #print(f"Visiting FuncCall {fcall}")
        if fcall in self.func_table:
            #print("->deferring!")
            self.visit(self.func_table[fcall], path, invariants)
        #print('---- path ----')

    def visit_Assignment(self, n, path, invariants):
        # fallthrough detection
        case_node = path_select_last('Case', path)
        fallthrough_case_names = []
        if case_node is not None:
            case_parent = path_select_parent(case_node, path)
            siblings = [x for x in case_parent.block_items if not isinstance(x, c_ast.Default)]
            #for s in siblings:
            #    s.show()
            sibling_names = [(x.expr.expr.name if isinstance(x.expr, c_ast.Cast) else x.expr.name) for x in siblings]
            sibling_empty = [(len(x.stmts) == 0) for x in siblings]
                
            in_fallthrough = False
            slice_start = None
            slice_end = None
            for i,sibling in enumerate(zip(reversed(sibling_names),reversed(sibling_empty), reversed(siblings))):
                if sibling[0] == self.state:
                    slice_start = i + 1
                if (slice_start is not None) and sibling[1] and not(in_fallthrough) and (slice_start == i):
                    in_fallthrough = True
                if in_fallthrough:
                    if sibling[1]:
                        slice_end = i
                    else:
                        in_fallthrough = False
            if (slice_start is not None) and (slice_end is not None):
                slice_start_temp = slice_start
                slice_start = len(siblings) - slice_end
                slice_end = len(siblings) - slice_start_temp
                fallthrough_case_names = sibling_names[slice_start-1:slice_end]
        rval_str = ''
        
        # conditional assignment detection
        asm_if = path_select_last('If', path)
        subpath_case = path_slice_back(case_node, path)
        is_exhaustive_conditional = True if asm_if is None else (asm_if.iftrue is not None) and (asm_if.iffalse is not None)
        is_conditional = path_contains(subpath_case, c_ast.If) and not(is_exhaustive_conditional)

        if not(isinstance(n.rvalue, c_ast.Constant) or isinstance(n.rvalue, c_ast.BinaryOp) or isinstance(n.rvalue, c_ast.TernaryOp)):
            rval_str = n.rvalue.name
        #print(f">>>> {rval_str} == {self.state}")
        if rval_str == self.state:
            self.assignments.append((n,path,fallthrough_case_names,is_conditional,invariants))

class EnumDefVisitor(c_ast.NodeVisitor):
    def __init__(self, name):
        super().__init__()
        self._name = name
        self.enums = {}

    def visit_Enum(self, node):
        self.enums[self._name] = node

class EnumTypedefVisitor(c_ast.NodeVisitor):
    def __init__(self):
        self.enums = {}

    def visit_Typedef(self, node):
        ev = EnumDefVisitor(node.name)
        ev.visit(node)
        self.enums = {**self.enums, **ev.enums}


class EnumVisitor(c_ast.NodeVisitor):
    def __init__(self):
        super().__init__()
        self.enum_names = []

    def visit_Enumerator(self, node):
        self.enum_names.append(node.name)        

class SwitchCaseTermVisitor(c_ast.NodeVisitor):
    def __init__(self, asm_node):
        super().__init__()
        self._asm_node = asm_node
        self.hit = False

    def visit_Assignment(self, node):
        if node == self._asm_node:
            self.hit = True

class SwitchCaseTranVisitor(c_ast.NodeVisitor):
    def __init__(self, asm_node, path, states, fallthrough_states, is_conditional, invariants):
        super().__init__()
        self.states = states
        self.path = path
        self.invariants = invariants
        self.fallthrough_states = fallthrough_states
        self.is_conditional = is_conditional
        self._asm_node = asm_node
        self.tran_table = []

    def visit_Case(self, node):
        clev = CaseLabelExtractionVisitor()
        clev.visit(node.children()[0][1])
        state_from = clev.label
        # highly inefficient but it is what it is
        hit = False
        #for n in self.path:
        sctv = SwitchCaseTermVisitor(self._asm_node)
        #node.show()
        sctv.visit(node)
        hit = sctv.hit

        #print(state_from, "->", self._asm_node.rvalue.name, "? hit=", hit, "isstate=", (state_from in self.states), "invar=", self.invariants)
        if (hit or (state_from in self.invariants)) and (state_from in self.states):
            #if conditional, state remains in state sometimes
            if self.is_conditional:
                self.tran_table.append((state_from,state_from))
            # process state assignment
            self.tran_table.append((state_from, self._asm_node.rvalue.name))
            for ft in self.fallthrough_states:
                self.tran_table.append((ft, self._asm_node.rvalue.name))

class SwitchCasePropertyVisitor(c_ast.NodeVisitor):
    def __init__(self, state_asmts, func_blacklist = None):
        super().__init__()
        self._func_bl = func_blacklist
        self._sas = state_asmts
        self.properties = []

    def visit_FuncCall(self, node):
        name = node.name.name
        args = expr_list_to_str(node.args) if node.args is not None else ""
        fcall = f"{name}({args})"
        if not(name in self._func_bl):
            self.properties.append(fcall)

    def visit_Assignment(self, node):
        if not(node in self._sas):
            lvalue = None
            rvalue = None
            if isinstance(node.lvalue, c_ast.StructRef):
                lvalue = f"{node.lvalue.children()[0][1].name}->{node.lvalue.children()[1][1].name}";
            else:
                lvalue = node.lvalue.name
            
            if isinstance(node.rvalue, c_ast.Constant):
                rvalue = f"{node.rvalue.type }({node.rvalue.value})"
            else:
                rvalue = node.rvalue.name

            prop = f"{lvalue}<={rvalue}"
            self.properties.append(prop)


class SwitchCaseCodePropertyVisitor(c_ast.NodeVisitor):
    def __init__(self, case, state_asmts, func_blacklist):
        super().__init__()
        self._func_bl = func_blacklist
        self._case = case
        self._sas = state_asmts
        self.properties = []

    def visit_Case(self, node):
        label = node.children()[0][1]
        block = node
        clev = CaseLabelExtractionVisitor()
        clev.visit(label)
        if clev.label == self._case:
            scpv = SwitchCasePropertyVisitor(self._sas, self._func_bl)
            scpv.visit(block)
            self.properties += scpv.properties