#!/usr/bin/env python

"""

Annotate program node structures. The code in this module operates upon nodes

which are produced when simplifying AST node trees originating from the compiler

module.

Copyright (C) 2006 Paul Boddie <paul@boddie.org.uk>

This software is free software; you can redistribute it and/or

modify it under the terms of the GNU General Public License as

published by the Free Software Foundation; either version 2 of

the License, or (at your option) any later version.

This software is distributed in the hope that it will be useful,

but WITHOUT ANY WARRANTY; without even the implied warranty of

MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

GNU General Public License for more details.

You should have received a copy of the GNU General Public

License along with this library; see the file LICENCE.txt

If not, write to the Free Software Foundation, Inc.,

51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA

--------

To use this module, the easiest approach is to use the annotate function:

annotate(module, builtins)

The more complicated approach involves obtaining an Annotator:

annotator = Annotator()

Then, processing an existing module with it:

annotator.process(module)

If a module containing built-in classes and functions has already been

annotated, such a module should be passed in as an additional argument:

annotator.process(module, builtins)

"""

from simplified import *

import compiler

class System:

    """

    A class maintaining the state of the annotation system. When the system

    counter can no longer be incremented by any annotation operation, the

    system may be considered stable and fully annotated.

    """

    def __init__(self):

        self.count = 0

    def init(self, node):

        if not hasattr(node, "types"):

            node.types = []

    def annotate(self, node, types):

        self.init(node)

        for type in types:

            if type not in node.types:

                node.types.append(type)

                self.count += 1

system = System()

# Exceptions.

class FailureError(Exception):

    def __init__(self, exc, node, *args):

        Exception.__init__(self, *args)

        self.nodes = [node]

        self.exc = exc

    def add(self, node):

        self.nodes.append(node)

    def __str__(self):

        return "%s, %s" % (self.exc, self.nodes)

# Annotation.

class Annotator(Visitor):

    """

    The type annotator which traverses the program nodes, typically depth-first,

    and maintains a record of the current set of types applying to the currently

    considered operation. Such types are also recorded on the nodes, and a

    special "system" record is maintained to monitor the level of annotation

    activity with a view to recognising when no more annotations are possible.

    Throughout the annotation activity, type information consists of lists of

    Attribute objects where such objects retain information about the context of

    the type (since a value in the program may be associated with an object or

    class) and the actual type of the value being manipulated. Upon accessing

    attribute information on namespaces, additional accessor information is also

    exchanged - this provides a means of distinguishing between the different

    types possible when the means of constructing the namespace may depend on

    run-time behaviour.

    """

    def __init__(self):

        "Initialise the visitor."

        Visitor.__init__(self)

        self.system = system

        # Satisfy visitor issues.

        self.visitor = self

    def process(self, module, builtins=None):

        """

        Process the given 'module', using the optional 'builtins' to access

        built-in classes and functions.

        """

        self.subprograms = []

        self.current_subprograms = []

        self.current_namespaces = []

        # Give constants their own namespace.

        for value, constant in module.simplifier.constants.items():

            constant.namespace = Namespace()

        # Process the module, supplying builtins if possible.

        self.global_namespace = Namespace()

        if builtins is not None:

            self.builtins_namespace = builtins.namespace

        else:

            self.builtins_namespace = self.global_namespace

        return self.process_node(module)

    def process_node(self, node, locals=None):

        """

        Process a subprogram or module 'node', indicating any initial 'locals'.

        Return an annotated subprogram or module. Note that this method may

        mutate nodes in the original program.

        """

        if locals:

            self.namespace = locals

        else:

            self.namespace = self.global_namespace

        # Record the current subprogram and namespace.

        self.current_subprograms.append(node)

        self.current_namespaces.append(self.namespace)

        # Add namespace details to any structure involved.

        if getattr(node, "structure", None) is not None:

            node.structure.namespace = Namespace()

            # Initialise bases where appropriate.

            if hasattr(node.structure, "bases"):

                base_refs = []

                for base in node.structure.bases:

                    self.dispatch(base)

                    base_refs.append(self.namespace.types)

                node.structure.base_refs = base_refs

        # Dispatch to the code itself.

        node.namespace = self.namespace

        result = self.dispatch(node)

        result.namespace = self.namespace

        # Obtain the return values.

        self.last_returns = self.namespace.returns

        self.returned_locals = self.namespace.return_locals

        # Restore the previous subprogram and namespace.

        self.current_namespaces.pop()

        if self.current_namespaces:

            self.namespace = self.current_namespaces[-1]

        self.current_subprograms.pop()

        return result

    def annotate(self, node):

        "Annotate the given 'node' in the system."

        self.system.annotate(node, self.namespace.types)

    # Visitor methods.

    def default(self, node):

        """

        Process the given 'node', given that it does not have a specific

        handler.

        """

        for attr in ("expr", "lvalue", "test", "handler"):

            value = getattr(node, attr, None)

            if value is not None:

                setattr(node, attr, self.dispatch(value))

        for attr in ("body", "else_", "finally_", "code"):

            value = getattr(node, attr, None)

            if value is not None:

                setattr(node, attr, self.dispatches(value))

        return node

    def dispatch(self, node, *args):

        try:

            return Visitor.dispatch(self, node, *args)

        except FailureError, exc:

            exc.add(node)

            raise

        except Exception, exc:

            raise FailureError(exc, node)

    def visitLoadRef(self, loadref):

        self.namespace.set_types([Attribute(None, loadref.ref)])

        self.annotate(loadref)

        return loadref

    def visitLoadName(self, loadname):

        self.namespace.set_types(self.namespace.load(loadname.name))

        result = loadname

        self.annotate(result)

        return result

    def visitStoreName(self, storename):

        storename.expr = self.dispatch(storename.expr)

        self.namespace.store(storename.name, self.namespace.types)

        return storename

    def visitLoadTemp(self, loadtemp):

        index = getattr(loadtemp, "index", None)

        self.namespace.set_types(self.namespace.temp[index][-1])

        self.annotate(loadtemp)

        return loadtemp

    def visitStoreTemp(self, storetemp):

        storetemp.expr = self.dispatch(storetemp.expr)

        index = getattr(storetemp, "index", None)

        if not self.namespace.temp.has_key(index):

            self.namespace.temp[index] = []

        self.namespace.temp[index].append(self.namespace.types)

        return storetemp

    def visitReleaseTemp(self, releasetemp):

        index = getattr(releasetemp, "index", None)

        self.namespace.temp[index].pop()

        return releasetemp

    def visitLoadAttr(self, loadattr):

        loadattr.expr = self.dispatch(loadattr.expr)

        types = []

        accesses = {}

        for attr in self.namespace.types:

            if not accesses.has_key(attr.type):

                accesses[attr.type] = []

            for attribute, accessor in get_attributes(attr.type, loadattr.name):

                if attribute is not None:

                    types.append(attribute)

                else:

                    print "Empty attribute via accessor", accessor

                accesses[attr.type].append((attribute, accessor))

        self.namespace.set_types(types)

        loadattr.accesses = accesses

        self.annotate(loadattr)

        return loadattr

    def visitStoreAttr(self, storeattr):

        storeattr.expr = self.dispatch(storeattr.expr)

        expr = self.namespace.types

        storeattr.lvalue = self.dispatch(storeattr.lvalue)

        writes = {}

        for attr in self.namespace.types:

            if attr is None:

                print "Empty attribute storage attempt"

                continue

            attr.type.namespace.store(storeattr.name, expr)

            writes[attr.type] = attr.type.namespace.load(storeattr.name)

        storeattr.writes = writes

        return storeattr

    def visitConditional(self, conditional):

        # Conditionals keep local namespace changes isolated.

        # With Return nodes inside the body/else sections, the changes are

        # communicated to the caller.

        conditional.test = self.dispatch(conditional.test)

        saved_namespace = self.namespace

        self.namespace = Namespace()

        self.namespace.merge_namespace(saved_namespace)

        conditional.body = self.dispatches(conditional.body)

        body_namespace = self.namespace

        self.namespace = Namespace()

        self.namespace.merge_namespace(saved_namespace)

        conditional.else_ = self.dispatches(conditional.else_)

        else_namespace = self.namespace

        self.namespace = Namespace()

        self.namespace.merge_namespace(body_namespace)

        self.namespace.merge_namespace(else_namespace)

        return conditional

    def visitReturn(self, return_):

        if hasattr(return_, "expr"):

            return_.expr = self.dispatch(return_.expr)

            self.namespace.returns += self.namespace.types

        self.namespace.snapshot()

        return return_

    def visitInvoke(self, invoke):

        invoke.expr = self.dispatch(invoke.expr)

        invocation_types = self.namespace.types

        self.process_args(invoke)

        # Now locate and invoke the subprogram. This can be complicated because

        # the target may be a class or object, and there may be many different

        # related subprograms.

        invocations = {}

        # Visit each callable in turn, finding subprograms.

        for attr in invocation_types:

            # Deal with class invocations by providing instance objects.

            # Here, each class is queried for the __init__ method, which may

            # exist for some combinations of classes in a hierarchy but not for

            # others.

            if isinstance(attr.type, Class):

                attributes = get_attributes(attr.type, "__init__")

            # Deal with object invocations by using __call__ methods.

            elif isinstance(attr.type, Instance):

                attributes = get_attributes(attr.type, "__call__")

            # Normal functions or methods are more straightforward.

            # Here, we model them using an attribute with no context and with

            # no associated accessor.

            else:

                attributes = [(attr, None)]

            # Inspect each attribute and extract the subprogram.

            for attribute, accessor in attributes:

                # If a class is involved, presume that it must create a new

                # object.

                if isinstance(attr.type, Class):

                    # Instantiate the class.

                    # NOTE: Should probably only allocate a single instance.

                    instance = Instance()

                    instance.namespace = Namespace()

                    instance.namespace.store("__class__", [attr.type])

                    # For instantiations, switch the context.

                    if attribute is not None:

                        attribute = Attribute(instance, attribute.type)

                # Skip cases where no callable is found.

                if attribute is not None:

                    # If a subprogram is defined, invoke it.

                    self.invoke_subprogram(invoke, attribute)

                    invocations[callable] = attribute.type

                else:

                    print "Invocation type is None"

                if isinstance(attr.type, Class):

                    # Associate the instance with the result of this invocation.

                    self.namespace.set_types([Attribute(None, instance)])

                    self.annotate(invoke)

        invoke.invocations = invocations

        return invoke

    # Utility methods.

    def invoke_subprogram(self, invoke, subprogram):

        """

        Invoke using the given 'invoke' node the given 'subprogram'.

        """

        # Test to see if anything has changed.

        if hasattr(invoke, "syscount") and invoke.syscount == self.system.count:

            return

        # Test for context information, making it into a real attribute.

        if subprogram.context is not None:

            context = Attribute(None, subprogram.context)

            target = subprogram.type

        else:

            context = None

            target = subprogram.type

        # Provide the correct namespace for the invocation.

        if isinstance(invoke, InvokeBlock):

            namespace = Namespace()

            namespace.merge_namespace(self.namespace)

        else:

            items = self.make_items(invoke, target, context)

            namespace = self.make_namespace(items)

        # Process the subprogram.

        self.process_node(target, namespace)

        # NOTE: Improve and verify this.

        if getattr(target, "returns_value", 0):

            self.namespace.set_types(self.last_returns)

            self.annotate(invoke)

        if isinstance(invoke, InvokeBlock):

            for locals in self.returned_locals:

                self.namespace.merge_namespace(locals)

        # Remember the state of the system.

        invoke.syscount = self.system.count

    def process_args(self, invoke):

        # NOTE: Consider initialiser invocation for classes.

        types = []

        args = []

        # Get type information for regular arguments.

        for arg in invoke.args:

            args.append(self.dispatch(arg))

            types.append(self.namespace.types)

        # Get type information for star and dstar arguments.

        if invoke.star is not None:

            param, default = invoke.star

            default = self.dispatch(default)

            invoke.star = param, default

            types.append(default.types)

        if invoke.dstar is not None:

            param, default = invoke.dstar

            default = self.dispatch(default)

            invoke.dstar = param, default

            types.append(default.types)

        invoke.args = args

    def make_items(self, invocation, subprogram, context):

        """

        Make an items mapping for the 'invocation' of the 'subprogram' using the

        given 'context' (which may be None).

        """

        if context is not None:

            args = [Self(context)] + invocation.args

        else:

            args = invocation.args

        params = subprogram.params

        items = []

        keywords = {}

        # Process the specified arguments.

        for arg in args:

            if isinstance(arg, Keyword):

                keywords[arg.name] = arg.expr

                continue

            elif params:

                param, default = params[0]

                if arg is None:

                    arg = default

            else:

                raise TypeError, "Invocation has too many arguments."

            items.append((param, arg.types))

            params = params[1:]

        # Collect the remaining defaults.

        while params:

            param, default = params[0]

            if keywords.has_key(param):

                arg = keywords[param]

            else:

                arg = self.dispatch(default) # NOTE: Review reprocessing.

            items.append((param, arg.types))

            params = params[1:]

        # Add star and dstar.

        if invocation.star is not None:

            if subprogram.star is not None:

                param, default = subprogram.star

                items.append((param, invocation.star.types))

            else:

                raise TypeError, "Invocation provides unwanted *args."

        elif subprogram.star is not None:

            param, default = subprogram.star

            arg = self.dispatch(default) # NOTE: Review reprocessing.

            items.append((param, arg.types))

        if invocation.dstar is not None:

            if subprogram.dstar is not None:

                param, default = subprogram.dstar

                items.append((param, invocation.dstar.types))

            else:

                raise TypeError, "Invocation provides unwanted **args."

        elif subprogram.dstar is not None:

            param, default = subprogram.dstar

            arg = self.dispatch(default) # NOTE: Review reprocessing.

            items.append((param, arg.types))

        return items

    def make_namespace(self, items):

        namespace = Namespace()

        namespace.merge_items(items)

        return namespace

# Namespace-related abstractions.

class Namespace:

    """

    A local namespace which may either relate to a genuine set of function

    locals or the initialisation of a structure or module.

    """

    def __init__(self):

        """

        Initialise the namespace with a mapping of local names to possible

        types, a list of return values and of possible returned local

        namespaces. The namespace also tracks the "current" types and a mapping

        of temporary value names to types.

        """

        self.names = {}

        self.returns = []

        self.return_locals = []

        self.temp = {}

        self.types = []

    def set_types(self, types):

        self.types = types

    def store(self, name, types):

        self.names[name] = types

    def load(self, name):

        return self.names[name]

    def merge(self, name, types):

        if not self.names.has_key(name):

            self.names[name] = types[:]

        else:

            existing = self.names[name]

            for type in types:

                if type not in existing:

                    existing.append(type)

    def merge_namespace(self, namespace):

        self.merge_items(namespace.names.items())

    def merge_items(self, items):

        for name, types in items:

            self.merge(name, types)

    def snapshot(self):

        "Make a snapshot of the locals and remember them."

        namespace = Namespace()

        namespace.merge_namespace(self)

        self.return_locals.append(namespace)

    def __repr__(self):

        return repr(self.names)

class Attribute:

    """

    An attribute abstraction, indicating the type of the attribute along with

    its context or origin.

    """

    def __init__(self, context, type):

        self.context = context

        self.type = type

    def __eq__(self, other):

        return hasattr(other, "type") and other.type == self.type or other == self.type

    def __repr__(self):

        return "Attribute(%s, %s)" % (repr(self.context), repr(self.type))

class Self:

    def __init__(self, attribute):

        self.types = [attribute]

def find_attributes(structure, name):

    """

    Find for the given 'structure' all attributes for the given 'name', visiting

    base classes where appropriate and returning the attributes in order of

    descending precedence for all possible base classes.

    The elements in the result list are 2-tuples which contain the attribute and

    the structure involved in accessing the attribute.

    """

    # First attempt to search the instance/class namespace.

    try:

        l = structure.namespace.load(name)

        attributes = []

        for attribute in l:

            attributes.append((attribute, structure))

    # If that does not work, attempt to investigate any class or base classes.

    except KeyError:

        attributes = []

        # Investigate any instance's implementing class.

        if isinstance(structure, Instance):

            for cls in structure.namespace.load("__class__"):

                l = get_attributes(cls, name)

                for attribute in l:

                    if attribute not in attributes:

                        attributes.append(attribute)

        # Investigate any class's base classes.

        elif isinstance(structure, Class):

            # If no base classes exist, return an indicator that no attribute

            # exists.

            if not structure.base_refs:

                return [(None, structure)]

            # Otherwise, find all possible base classes.

            for base_refs in structure.base_refs:

                base_attributes = []

                # For each base class, find attributes either in the base

                # class or its own base classes.

                for base_ref in base_refs:

                    l = get_attributes(base_ref, name)

                    for attribute in l:

                        if attribute not in base_attributes:

                            base_attributes.append(attribute)

                attributes += base_attributes

    return attributes

def get_attributes(structure, name):

    """

    Return all possible attributes for the given 'structure' having the given

    'name', wrapping each attribute in an Attribute object which includes

    context information for the attribute access.

    The elements in the result list are 2-tuples which contain the attribute and

    the structure involved in accessing the attribute.

    """

    if isinstance(structure, Attribute):

        structure = structure.type

    results = []

    for attribute, accessor in find_attributes(structure, name):

        if attribute is not None and isinstance(structure, Structure):

            results.append((Attribute(structure, attribute.type), accessor))

        else:

            results.append((attribute, accessor))

    return results

# Convenience functions.

def annotate(module, builtins=None):

    """

    Annotate the given 'module', also employing the optional 'builtins' module,

    if specified.

    """

    annotator = Annotator()

    if builtins is not None:

        annotator.process(module, builtins)

    else:

        annotator.process(module)

def annotate_all(modules, builtins):

    annotate(builtins)

    for module in modules:

        annotate(module, builtins)