#-------------------------------------------------------------------------------
# visit.py
# 
# Basic classes for binding tree analysis
#
# Copyright (C) 2013, Shinya Takamaeda-Yamazaki
# License: Apache 2.0
#-------------------------------------------------------------------------------

import sys
import os
import re
import copy
import collections

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) )

import pyverilog
import pyverilog.utils.util as util
import pyverilog.utils.verror as verror
from pyverilog.utils.scope import ScopeLabel, ScopeChain
from pyverilog.vparser.ast import *

map_key   = lambda f,d: collections.OrderedDict([ (f(k),v) for k,v in d.items() ])
map_value = lambda f,d: collections.OrderedDict([ (k,f(v)) for k,v in d.items() ])

################################################################################
# Primitive list
################################################################################
primitives = {
    'and' : Uand,
    'nand': Unand,
    'or' : Uor,
    'nor' : Unor,
    'xor' : Uxor,
    'xnor' : Uxnor,
    'not' : Unot,
    'buf' : None
}

################################################################################
# Base Visitor
################################################################################
class NodeVisitor(object):
    def visit(self, node):
        method = 'visit_' + node.__class__.__name__
        visitor = getattr(self, method, self.generic_visit)
        return visitor(node)
    def generic_visit(self, node):
        for c in node.children():
            self.visit(c)

################################################################################
# Signal/Object Management Classes
################################################################################
class AlwaysInfo(object):
    def __init__(self, clock_name='', clock_edge=None, clock_bit=0,
                 reset_name='', reset_edge=None, reset_bit=0, senslist=()):
        self.clock_name = clock_name
        self.clock_edge = clock_edge
        self.clock_bit = clock_bit
        self.reset_name = reset_name
        self.reset_edge = reset_edge
        self.reset_bit = reset_bit
        self.senslist = senslist
    def getClockName(self):
        return self.clock_name
    def getClockEdge(self):
        return self.clock_edge
    def getClockBit(self):
        return self.clock_bit
    def getResetName(self):
        return self.reset_name
    def getResetEdge(self):
        return self.reset_edge
    def getResetBit(self):
        return self.reset_bit
    def isClockEdge(self):
        if self.clock_name != '' and self.clock_edge == 'posedge': return True
        if self.clock_name != '' and self.clock_edge == 'negedge': return True
        return False
    def isCombination(self):
        if self.clock_name is None: return True
        return False
    def isResetEdge(self):
        if self.reset_name != '' and self.reset_edge == 'posedge': return True
        if self.reset_name != '' and self.reset_edge == 'negedge': return True
        return False

################################################################################
class Label(object):
    def __init__(self, name):
        self.name = name
        self.cnt = 0
    def get(self):
        ret = self.name + str(self.cnt)
        self.inc()
        return ret
    def inc(self):
        self.cnt += 1

class Labels(object):
    def __init__(self):
        self.labels = {}
    def get(self, name):
        if not name in self.labels: self.labels[name] = Label(name)
        return self.labels[name].get()
    def inc(self, name):
        if not name in self.labels: self.labels[name] = Label(name)
        self.labels[name].inc()

################################################################################
class VariableTable(object):
    def __init__(self):
        self.dict = collections.OrderedDict()
    def add(self, name, var):
        if name in self.dict:
            self.dict[name] = self.dict[name] + (var, )
        else:
            self.dict[name] = (var, )
    def get(self, name):
        return self.dict[name]
    def has(self, name):
        return name in self.dict
    def update(self, table):
        self.dict.update(table)
    def getDict(self):
        return self.dict
        
class SignalTable(VariableTable): pass
class ConstTable(VariableTable): pass
class GenvarTable(VariableTable): pass

class Variables(object):
    def __init__(self):
        self.signal = SignalTable()
        self.const = ConstTable()
        self.genvar = GenvarTable()
    def addSignal(self, name, var):
        if self.const.has(name): return
        self.signal.add(name, var)
    def addConst(self, name, var):
        if self.const.has(name): return
        self.const.add(name, var)
    def addGenvar(self, name, var):
        if self.const.has(name): return
        self.genvar.add(name, var)
    def getSignal(self, name):
        return self.signal.get(name)
    def getConst(self, name):
        return self.const.get(name)
    def getGenvar(self, name):
        return self.genvar.get(name)
    def hasSignal(self, name):
        return self.signal.has(name)
    def hasConst(self, name):
        return self.const.has(name)
    def updateSignal(self, var):
        self.signal.update(var)
    def updateConst(self, var):
        self.const.update(var)
    def getSignals(self):
        return self.signal.getDict()
    def getConsts(self):
        return self.const.getDict()

################################################################################
class DefinitionInfo(object):
    def __init__(self, name, definition):
        self.name = name
        self.definition = definition
        self.variables = Variables()
        self.ioports = []
        self.params = []
    def addSignal(self, name, var):
        self.variables.addSignal(name, var)
    def addConst(self, name, var):
        self.variables.addConst(name, var)
    def addParamName(self, name):
        self.params.append(name)
    def addPort(self, port):
        self.ioports.append(port)
    def addPorts(self, ports):
        for p in ports:
            if isinstance(p, Ioport):
                self.ioports.append(p.first.name)
            else:
                self.ioports.append(p.name)
    def getSignals(self):
        return self.variables.getSignals()
    def getConsts(self):
        return self.variables.getConsts()
    def getDefinition(self):
        return self.definition
    def getIOPorts(self):
        return tuple(self.ioports)
    def getParamNames(self):
        return tuple(self.params)

class DefinitionInfoTable(object):
    def __init__(self):
        self.dict = collections.OrderedDict()
        self.current = None
    def addDefinition(self, name, definition):
        if name in self.dict:
            raise verror.DefinitionError('Already defined: %s' % name)
        self.dict[name] = DefinitionInfo(name, definition)
        self.current = name
    def addPorts(self, ports):
        self.dict[self.current].addPorts(ports)
    def addPort(self, port):
        self.dict[self.current].addPort(port)
    def addSignal(self, name, var):
        self.dict[self.current].addSignal(name, var)
    def addConst(self, name, var):
        self.dict[self.current].addConst(name, var)
    def addParamName(self, name):
        self.dict[self.current].addParamName(name)

    def getSignals(self, name):
        if name not in self.dict: raise verror.DefinitionError('No such module: %s' % name)
        return self.dict[name].getSignals()
    def getConsts(self, name):
        if name not in self.dict: raise verror.DefinitionError('No such module: %s' % name)
        return self.dict[name].getConsts()
    def getDefinition(self, name):
        if name not in self.dict: raise verror.DefinitionError('No such module: %s' % name)
        return self.dict[name].getDefinition()
    def getDefinitions(self):
        return self.dict
    def getIOPorts(self, name):
        if name not in self.dict: raise verror.DefinitionError('No such module: %s' % name)
        return self.dict[name].getIOPorts()
    def getParamNames(self, name):
        if name not in self.dict: raise verror.DefinitionError('No such module: %s' % name)
        return self.dict[name].getParamNames()
    def get_names(self):
        ret = []
        for dk, dv in self.dict.items():
            ret.append(dk)
        return ret

    def overwriteDefinition(self, name, definition):
        self.dict[name] = definition

    def copyDefinition(self, f, t):
        self.dict[t] = copy.deepcopy(self.dict[f])
        self.dict[t].definition.name = t
        self.dict[t].name = t

class ModuleInfo(DefinitionInfo): pass
class ModuleInfoTable(DefinitionInfoTable): pass
class FunctionInfo(DefinitionInfo): pass
class FunctionInfoTable(DefinitionInfoTable): pass
class TaskInfo(DefinitionInfo): pass
class TaskInfoTable(DefinitionInfoTable): pass

################################################################################
frametype_list = ('ifthen', 'ifelse', 'case', 'for', 'while', 'none')

class Frame(object):
    def __init__(self, name, previous, frametype='none',
                 alwaysinfo=None, condition=None,
                 module=False, functioncall=False, taskcall=False,
                 generate=False, always=False, initial=False, loop=None, loop_iter=None,
                 modulename=None):
        self.name = name
        self.previous = previous
        self.next = []
        self.frametype = frametype

        self.alwaysinfo = alwaysinfo
        self.condition = condition

        self.module = module
        self.functioncall = functioncall
        self.taskcall = taskcall
        self.generate = generate
        self.always = always
        self.initial = initial
        self.loop = loop
        self.loop_iter = loop_iter

        self.variables = Variables()
        self.functions = FunctionInfoTable()
        self.tasks = TaskInfoTable()
        self.blockingassign = collections.OrderedDict()
        self.nonblockingassign = collections.OrderedDict()

        self.modulename = modulename

    def getName(self):
        return self.name
    def getPrevious(self):
        return self.previous
    def getNext(self):
        return self.next
    def getFrametype(self):
        return self.frametype

    def isIfelse(self):
        return self.frametype == 'ifelse'

    def getAlwaysInfo(self):
        return self.alwaysinfo
    def getCondition(self):
        return self.condition

    def isModule(self):
        return self.module
    def isFunctioncall(self):
        return self.functioncall
    def isTaskcall(self):
        return self.taskcall
    def isGenerate(self):
        return self.generate
    def isAlways(self):
        return self.always
    def isInitial(self):
        return self.initial

    def setNext(self, nextframe):
        self.next.append(nextframe)

    def setAlwaysInfo(self, clock_name, clock_edge, clock_bit,
                      reset_name, reset_edge, reset_bit, senslist):
        self.alwaysinfo = AlwaysInfo(clock_name, clock_edge, clock_bit,
                                     reset_name, reset_edge, reset_bit, senslist)

    def addSignal(self, node):
        self.variables.addSignal(node.name, node)
    def addConst(self, node):
        self.variables.addConst(node.name, node)
    def addGenvar(self, node):
        self.variables.addGenvar(node.name, 0)

    def addFunction(self, node):
        self.functions.addDefinition(node.name, node)
    def addTask(self, node):
        self.tasks.addDefinition(node.name, node)

    def addFunctionPort(self, node):
        self.functions.addPort(node)
    def addTaskPort(self, node):
        self.tasks.addPort(node)

    def updateSignal(self, signal):
        self.variables.updateSignal(signal)
    def updateConst(self, const):
        self.variables.updateConst(const)

    def getConstant(self, name):
        return self.variables.getConst(name)
    def hasConstant(self, name):
        return self.variables.hasConst(name)

    def getSignal(self, name):
        return self.variables.getSignal(name)
    def hasSignal(self, name):
        return self.variables.hasSignal(name)

    def getSignals(self):
        return self.variables.getSignals()
    def getConsts(self):
        return self.variables.getConsts()
    def getFunctions(self):
        return self.functions.getDefinitions()
    def getTasks(self):
        return self.tasks.getDefinitions()

    def getModuleName(self):
        return self.modulename

    ############################################################################
    def setBlockingAssign(self, dst, bind):
        if not dst in self.blockingassign:
            self.blockingassign[dst] = (bind,)
            return
        current = self.blockingassign[dst]
        for c_i, c in enumerate(current):
            if c.msb == bind.msb and c.msb == bind.msb and c.ptr == bind.ptr:
                self.blockingassign[dst][c_i].tree = bind.tree
                return
        self.blockingassign[dst] = current + (bind,)
            
    def getBlockingAssign(self, dst):
        if dst in self.blockingassign: return self.blockingassign[dst]
        return ()

    def getBlockingAssigns(self):
        return self.blockingassign

    ############################################################################
    def addNonblockingAssign(self, dst, bind):
        if not dst in self.nonblockingassign:
            self.nonblockingassign[dst] = (copy.deepcopy(bind),)
            return
        current = self.nonblockingassign[dst]
        for c_i, c in enumerate(current):
            if c.msb == bind.msb and c.msb == bind.msb and c.ptr == bind.ptr:
                self.nonblockingassign[dst][c_i].tree = copy.deepcopy(bind.tree)
                return
        self.nonblockingassign[dst] = current + (copy.deepcopy(bind),)

    def getNonblockingAssigns(self):
        return self.nonblockingassign

################################################################################
class FrameTable(object):
    def __init__(self):
        self.dict = {}
        self.current = ScopeChain()

        self.function_def = False
        self.task_def = False
        self.for_pre = False
        self.for_post = False
        self.for_iter = None

    def toScopeChain(self, scopename):
        if scopename is None: return self.current
        return self.current + scopename

    def addFrame(self, scopename,
                 frametype='none',
                 alwaysinfo=None, condition=None,
                 module=False, functioncall=False, taskcall=False,
                 generate=False, always=False, initial=False, loop=None, loop_iter=None,
                 modulename=None):

        scopechain = self.toScopeChain(scopename)
        if scopechain in self.dict:
            raise verror.DefinitionError('Already Exists: %s' % str(scopechain))
        ret = self.current
        previous = self.current
        if len(previous) > 0:
            self.dict[previous].setNext(scopechain)
        self.dict[scopechain] = Frame(scopechain, previous, frametype=frametype,
                                      alwaysinfo=alwaysinfo, condition=condition,
                                      module=module, functioncall=functioncall,
                                      taskcall=taskcall, generate=generate,
                                      always=always, initial=initial, loop=loop, loop_iter=loop_iter,
                                      modulename=modulename)
        self.current = scopechain
        return ret

    def hasFrame(self, scope):
        return scope in self.dict

    def getFrametype(self):
        return self.dict[self.current].getFrametype()
    def getAlwaysInfo(self):
        return self.dict[self.current].getAlwaysInfo()
    def getCondition(self):
        return self.dict[self.current].getCondition()

    def isModule(self):
        return self.dict[self.current].isModule()
    def isFunctioncall(self):
        return self.dict[self.current].isFunctioncall()
    def isTaskcall(self):
        return self.dict[self.current].isTaskcall()
    def isGenerate(self):
        return self.dict[self.current].isGenerate()
    def isAlways(self):
        return self.dict[self.current].isAlways()
    def isInitial(self):
        return self.dict[self.current].isInitial()
    def getLoop(self):
        return self.dict[self.current].loop
    def getLoopIter(self):
        return self.dict[self.current].loop_iter

    def isFunctiondef(self):
        return self.function_def
    def isTaskdef(self):
        return self.task_def
    def isForpre(self):
        return self.for_pre
    def isForpost(self):
        return self.for_post

    def setFunctionDef(self):
        self.function_def = True
    def unsetFunctionDef(self):
        self.function_def = False
    def setTaskDef(self):
        self.task_def = True
    def unsetTaskDef(self):
        self.task_def = False
    def setForPre(self):
        self.for_pre = True
    def unsetForPre(self):
        self.for_pre = False
    def setForPost(self):
        self.for_post = True
    def unsetForPost(self):
        self.for_post = False

    def setForIter(self, iter):
        self.for_iter = iter
    def getForIter(self):
        return self.for_iter

    def setAlwaysInfo(self, clock_name, clock_edge, clock_bit,
                      reset_name, reset_edge, reset_bit, senslist):
        self.dict[self.current].setAlwaysInfo(clock_name, clock_edge, clock_bit,
                                              reset_name, reset_edge, reset_bit, senslist)

    def setCurrent(self, current):
        self.current = current

    def getCurrent(self):
        return self.current

    def getModuleName(self):
        return self.dict[self.current].getModuleName()

    def getLabelKey(self, name):
        ret = ''
        if self.isModule(): ret += 'md_'
        if self.isGenerate(): ret += 'ge_'
        if self.isAlways(): ret += 'al_'
        if self.isInitial(): ret += 'in_'
        if self.isFunctioncall(): ret += 'fc_'
        if self.isTaskcall(): ret += 'tc_'
        return ret + name

    ############################################################################
    def addSignal(self, var):
        if self.function_def and isinstance(var, Input):
            self.dict[self.current].addFunctionPort(var)
            return
        if self.task_def and isinstance(var, Input):
            self.dict[self.current].addTaskPort(var)
            return
        self.dict[self.current].addSignal(var)
            
    def addConst(self, var):
        self.dict[self.current].addConst(var)
    def addGenvar(self, var):
        self.dict[self.current].addGenvar(var)

    def addFunction(self, var):
        self.dict[self.current].addFunction(var)
    def addTask(self, var):
        self.dict[self.current].addTask(var)

    ############################################################################
    def updateSignal(self, signal):
        self.dict[self.current].updateSignal(signal)
    def updateConst(self, const):
        self.dict[self.current].updateConst(const)

    ############################################################################
    def getAllInstances(self):
        ret = []
        for dk in self.dict.keys():
            if dk[-1].scopetype == 'module':
                ret.append( (dk, self.dict[dk].getModuleName()) )
        return tuple(ret)
        
    def getAllSignals(self):
        ret = collections.OrderedDict()
        for dk, dv in self.dict.items():
            ret.update(
                map_key((lambda x: dk+ScopeLabel(x,'signal')), dv.getSignals()))
        return ret

    def getAllConsts(self):
        ret = collections.OrderedDict()
        for dk, dv in self.dict.items():
            ret.update(
                map_key((lambda x: dk+ScopeLabel(x,'signal')), dv.getConsts()))
        return ret

    def getAllFunctions(self):
        ret = collections.OrderedDict()
        for dk, dv in self.dict.items():
            ret.update(
                map_key((lambda x: dk+ScopeLabel(x,'function')), dv.getFunctions()))
        return ret

    def getAllTasks(self):
        ret = collections.OrderedDict()
        for dk, dv in self.dict.items():
            ret.update(
                map_key((lambda x: dk+ScopeLabel(x,'task')), dv.getTasks()))
        return ret

    ############################################################################
    def getSignals(self, scope):
        return map_key(
            (lambda x: scope+ScopeLabel(x,'signal')), self.dict[scope].getSignals())

    def getConsts(self, scope):
        return map_key(
            (lambda x: scope+ScopeLabel(x,'signal')), self.dict[scope].getConsts())

    def getFunctions(self, scope):
        return map_key(
            (lambda x: scope+ScopeLabel(x,'function')), self.dict[scope].getFunctions())

    def getTasks(self, scope):
        return map_key(
            (lambda x: scope+ScopeLabel(x,'task')), self.dict[scope].getTasks())

    ################################################################################
    def getGenerateConditions(self):
        ret = []
        ptr = self.current
        while True:
            frame = self.dict[ptr]
            if not frame.isGenerate(): break
            if frame.loop is not None: # generate for
                ret.append( (frame.loop_iter, frame.loop) )
            elif frame.condition is not None: # generate if
                if frame.isIfelse():
                    ret.append( (None, Ulnot(frame.condition)) )
                else:
                    ret.append( (None, frame.condition) )
            ptr = self.dict[ptr].previous
        return tuple(reversed(ret))

    def getAlwaysStatus(self):
        ptr = self.current
        frame = self.dict[ptr]
        alwaysinfo = frame.getAlwaysInfo()
        while True:
            if frame.isFunctioncall(): return None
            if frame.isTaskcall(): return None
            if not frame.isAlways(): return None
            if alwaysinfo is not None: return alwaysinfo
            ptr = self.dict[ptr].previous
            frame = self.dict[ptr]
            alwaysinfo = frame.getAlwaysInfo()
    
    def setBlockingAssign(self, dst, bind, scope):
        self.dict[scope].setBlockingAssign(dst, bind)

    def getBlockingAssign(self, dst, scope):
        p = scope
        while self.dict[p].isAlways():
            ret = self.dict[p].getBlockingAssign(dst)
            if ret: return ret
            p = self.dict[p].getPrevious()
        return ()

    def addNonblockingAssign(self, dst, bind):
        self.dict[self.current].addNonblockingAssign(dst, bind)

    def getBlockingAssigns(self):
        return self.dict[self.current].getBlockingAssigns()

    def getBlockingAssignsScope(self, scope):
        return self.dict[scope].getBlockingAssigns()

    def getNonblockingAssigns(self):
        return self.dict[self.current].getNonblockingAssigns()

    def getPreviousNonblockingAssign(self):
        previous = self.dict[self.current].previous
        if len(previous) == 0: return collections.OrderedDict()
        previous_frame = self.dict[previous]
        if not previous_frame.isAlways(): return collections.OrderedDict()
        return previous_frame.getNonblockingAssigns()

    def searchConstantDefinition(self, key, name):
        if self.dict[key].hasConstant(name):
            return key, self.dict[key].getConstant(name)
        if self.dict[key].isModule(): return None, None
        previous = self.dict[key].getPrevious()
        return self.searchConstantDefinition(previous, name)

    def getConstantDefinition(self, key, name):
        if self.dict[key].hasConstant(name):
            return self.dict[key].getConstant(name)
        return None

    def searchSignalDefinition(self, key, name):
        if self.dict[key].hasSignal(name):
            return key, self.dict[key].getSignal(name)
        if self.dict[key].isModule(): return None, None
        previous = self.dict[key].getPrevious()
        return self.searchSignalDefinition(previous, name)

    def searchMatchedScopeChain(self, currentchain, targetchain):
        skiplength = len(currentchain) - 1
        printable = currentchain[skiplength].isPrintable()

        if printable:
            if currentchain[skiplength] != targetchain[0]: # not match
                return None
            if len(targetchain) == 1:
                return ScopeChain( [targetchain[0]] )
            for nextframe in self.dict[currentchain].getNext():
                rslt = self.searchMatchedScopeChain(nextframe, targetchain[1:])
                if rslt is not None:
                    return ScopeChain( [currentchain[skiplength]] ) + rslt
            return None

        if (currentchain[skiplength].scopetype == 'for' and 
            targetchain[0].scopetype == 'for'):
            if currentchain[skiplength].scopeloop != targetchain[0].scopeloop:
                return None
            if len(targetchain) == 1:
                return ScopeChain( [targetchain[0]] )
            for nextframe in self.dict[currentchain].getNext():
                rslt = self.searchMatchedScopeChain(nextframe, targetchain[1:])
                if rslt is not None:
                    return ScopeChain( [currentchain[skiplength]] ) + rslt
                    
            return None

        for nextframe in self.dict[currentchain].getNext():
            rslt = self.searchMatchedScopeChain(nextframe, targetchain)
            if rslt is not None:
                return ScopeChain( [currentchain[skiplength]] ) + rslt
        return None

