#! /usr/bin/env python

"""Generate C code from an ASDL description."""

# TO DO

# handle fields that have a type but no name

import os, sys

import asdl

TABSIZE = 8

MAX_COL = 80

def get_c_type(name):

    """Return a string for the C name of the type.

    This function special cases the default types provided by asdl:

    identifier, string, int, bool.

    """

    # XXX ack!  need to figure out where Id is useful and where string

    if isinstance(name, asdl.Id):

        name = name.value

    if name in asdl.builtin_types:

        return name

    else:

        return "%s_ty" % name

def reflow_lines(s, depth):

    """Reflow the line s indented depth tabs.

    Return a sequence of lines where no line extends beyond MAX_COL

    when properly indented.  The first line is properly indented based

    exclusively on depth * TABSIZE.  All following lines -- these are

    the reflowed lines generated by this function -- start at the same

    column as the first character beyond the opening { in the first

    line.

    """

    size = MAX_COL - depth * TABSIZE

    if len(s) < size:

        return [s]

    lines = []

    cur = s

    padding = ""

    while len(cur) > size:

        i = cur.rfind(' ', 0, size)

        # XXX this should be fixed for real

        if i == -1 and 'GeneratorExp' in cur:

            i = size + 3

        assert i != -1, "Impossible line %d to reflow: %r" % (size, s)

        lines.append(padding + cur[:i])

        if len(lines) == 1:

            # find new size based on brace

            j = cur.find('{', 0, i)

            if j >= 0:

                j += 2 # account for the brace and the space after it

                size -= j

                padding = " " * j

            else:

                j = cur.find('(', 0, i)

                if j >= 0:

                    j += 1 # account for the paren (no space after it)

                    size -= j

                    padding = " " * j

        cur = cur[i+1:]

    else:

        lines.append(padding + cur)

    return lines

def is_simple(sum):

    """Return True if a sum is a simple.

    A sum is simple if its types have no fields, e.g.

    unaryop = Invert | Not | UAdd | USub

    """

    for t in sum.types:

        if t.fields:

            return False

    return True

class EmitVisitor(asdl.VisitorBase):

    """Visit that emits lines"""

    def __init__(self, file):

        self.file = file

        super(EmitVisitor, self).__init__()

    def emit(self, s, depth, reflow=1):

        # XXX reflow long lines?

        if reflow:

            lines = reflow_lines(s, depth)

        else:

            lines = [s]

        for line in lines:

            line = (" " * TABSIZE * depth) + line + "\n"

            self.file.write(line)

class TypeDefVisitor(EmitVisitor):

    def visitModule(self, mod):

        for dfn in mod.dfns:

            self.visit(dfn)

    def visitType(self, type, depth=0):

        self.visit(type.value, type.name, depth)

    def visitSum(self, sum, name, depth):

        if is_simple(sum):

            self.simple_sum(sum, name, depth)

        else:

            self.sum_with_constructors(sum, name, depth)

    def simple_sum(self, sum, name, depth):

        enum = []

        for i in range(len(sum.types)):

            type = sum.types[i]

            enum.append("%s=%d" % (type.name, i + 1))

        enums = ", ".join(enum)

        ctype = get_c_type(name)

        s = "typedef enum _%s { %s } %s;" % (name, enums, ctype)

        self.emit(s, depth)

        self.emit("", depth)

    def sum_with_constructors(self, sum, name, depth):

        ctype = get_c_type(name)

        s = "typedef struct _%(name)s *%(ctype)s;" % locals()

        self.emit(s, depth)

        self.emit("", depth)

    def visitProduct(self, product, name, depth):

        ctype = get_c_type(name)

        s = "typedef struct _%(name)s *%(ctype)s;" % locals()

        self.emit(s, depth)

        self.emit("", depth)

class StructVisitor(EmitVisitor):

    """Visitor to generate typdefs for AST."""

    def visitModule(self, mod):

        for dfn in mod.dfns:

            self.visit(dfn)

    def visitType(self, type, depth=0):

        self.visit(type.value, type.name, depth)

    def visitSum(self, sum, name, depth):

        if not is_simple(sum):

            self.sum_with_constructors(sum, name, depth)

    def sum_with_constructors(self, sum, name, depth):

        def emit(s, depth=depth):

            self.emit(s % sys._getframe(1).f_locals, depth)

        enum = []

        for i in range(len(sum.types)):

            type = sum.types[i]

            enum.append("%s_kind=%d" % (type.name, i + 1))

        emit("enum _%(name)s_kind {" + ", ".join(enum) + "};")

        emit("struct _%(name)s {")

        emit("enum _%(name)s_kind kind;", depth + 1)

        emit("union {", depth + 1)

        for t in sum.types:

            self.visit(t, depth + 2)

        emit("} v;", depth + 1)

        for field in sum.attributes:

            # rudimentary attribute handling

            type = str(field.type)

            assert type in asdl.builtin_types, type

            emit("%s %s;" % (type, field.name), depth + 1);

        emit("};")

        emit("")

    def visitConstructor(self, cons, depth):

        if cons.fields:

            self.emit("struct {", depth)

            for f in cons.fields:

                self.visit(f, depth + 1)

            self.emit("} %s;" % cons.name, depth)

            self.emit("", depth)

        else:

            # XXX not sure what I want here, nothing is probably fine

            pass

    def visitField(self, field, depth):

        # XXX need to lookup field.type, because it might be something

        # like a builtin...

        ctype = get_c_type(field.type)

        name = field.name

        if field.seq:

            if field.type.value in ('cmpop',):

                self.emit("asdl_int_seq *%(name)s;" % locals(), depth)

            else:

                self.emit("asdl_seq *%(name)s;" % locals(), depth)

        else:

            self.emit("%(ctype)s %(name)s;" % locals(), depth)

    def visitProduct(self, product, name, depth):

        self.emit("struct _%(name)s {" % locals(), depth)

        for f in product.fields:

            self.visit(f, depth + 1)

        self.emit("};", depth)

        self.emit("", depth)

class PrototypeVisitor(EmitVisitor):

    """Generate function prototypes for the .h file"""

    def visitModule(self, mod):

        for dfn in mod.dfns:

            self.visit(dfn)

    def visitType(self, type):

        self.visit(type.value, type.name)

    def visitSum(self, sum, name):

        if is_simple(sum):

            pass # XXX

        else:

            for t in sum.types:

                self.visit(t, name, sum.attributes)

    def get_args(self, fields):

        """Return list of C argument into, one for each field.

        Argument info is 3-tuple of a C type, variable name, and flag

        that is true if type can be NULL.

        """

        args = []

        unnamed = {}

        for f in fields:

            if f.name is None:

                name = f.type

                c = unnamed[name] = unnamed.get(name, 0) + 1

                if c > 1:

                    name = "name%d" % (c - 1)

            else:

                name = f.name

            # XXX should extend get_c_type() to handle this

            if f.seq:

                if f.type.value in ('cmpop',):

                    ctype = "asdl_int_seq *"

                else:

                    ctype = "asdl_seq *"

            else:

                ctype = get_c_type(f.type)

            args.append((ctype, name, f.opt or f.seq))

        return args

    def visitConstructor(self, cons, type, attrs):

        args = self.get_args(cons.fields)

        attrs = self.get_args(attrs)

        ctype = get_c_type(type)

        self.emit_function(cons.name, ctype, args, attrs)

    def emit_function(self, name, ctype, args, attrs, union=1):

        args = args + attrs

        if args:

            argstr = ", ".join(["%s %s" % (atype, aname)

                                for atype, aname, opt in args])

            argstr += ", PyArena *arena"

        else:

            argstr = "PyArena *arena"

        margs = "a0"

        for i in range(1, len(args)+1):

            margs += ", a%d" % i

        self.emit("#define %s(%s) _Py_%s(%s)" % (name, margs, name, margs), 0,

                reflow = 0)

        self.emit("%s _Py_%s(%s);" % (ctype, name, argstr), 0)

    def visitProduct(self, prod, name):

        self.emit_function(name, get_c_type(name),

                           self.get_args(prod.fields), [], union=0)

class FunctionVisitor(PrototypeVisitor):

    """Visitor to generate constructor functions for AST."""

    def emit_function(self, name, ctype, args, attrs, union=1):

        def emit(s, depth=0, reflow=1):

            self.emit(s, depth, reflow)

        argstr = ", ".join(["%s %s" % (atype, aname)

                            for atype, aname, opt in args + attrs])

        if argstr:

            argstr += ", PyArena *arena"

        else:

            argstr = "PyArena *arena"

        self.emit("%s" % ctype, 0)

        emit("%s(%s)" % (name, argstr))

        emit("{")

        emit("%s p;" % ctype, 1)

        for argtype, argname, opt in args:

            # XXX hack alert: false is allowed for a bool

            if not opt and not (argtype == "bool" or argtype == "int"):

                emit("if (!%s) {" % argname, 1)

                emit("PyErr_SetString(PyExc_ValueError,", 2)

                msg = "field %s is required for %s" % (argname, name)

                emit('                "%s");' % msg,

                     2, reflow=0)

                emit('return NULL;', 2)

                emit('}', 1)

        emit("p = (%s)PyArena_Malloc(arena, sizeof(*p));" % ctype, 1);

        emit("if (!p)", 1)

        emit("return NULL;", 2)

        if union:

            self.emit_body_union(name, args, attrs)

        else:

            self.emit_body_struct(name, args, attrs)

        emit("return p;", 1)

        emit("}")

        emit("")

    def emit_body_union(self, name, args, attrs):

        def emit(s, depth=0, reflow=1):

            self.emit(s, depth, reflow)

        emit("p->kind = %s_kind;" % name, 1)

        for argtype, argname, opt in args:

            emit("p->v.%s.%s = %s;" % (name, argname, argname), 1)

        for argtype, argname, opt in attrs:

            emit("p->%s = %s;" % (argname, argname), 1)

    def emit_body_struct(self, name, args, attrs):

        def emit(s, depth=0, reflow=1):

            self.emit(s, depth, reflow)

        for argtype, argname, opt in args:

            emit("p->%s = %s;" % (argname, argname), 1)

        assert not attrs

class PickleVisitor(EmitVisitor):

    def visitModule(self, mod):

        for dfn in mod.dfns:

            self.visit(dfn)

    def visitType(self, type):

        self.visit(type.value, type.name)

    def visitSum(self, sum, name):

        pass

    def visitProduct(self, sum, name):

        pass

    def visitConstructor(self, cons, name):

        pass

    def visitField(self, sum):

        pass

class Obj2ModPrototypeVisitor(PickleVisitor):

    def visitProduct(self, prod, name):

        code = "static int obj2ast_%s(PyObject* obj, %s* out, PyArena* arena);"

        self.emit(code % (name, get_c_type(name)), 0)

    visitSum = visitProduct

class Obj2ModVisitor(PickleVisitor):

    def funcHeader(self, name):

        ctype = get_c_type(name)

        self.emit("int", 0)

        self.emit("obj2ast_%s(PyObject* obj, %s* out, PyArena* arena)" % (name, ctype), 0)

        self.emit("{", 0)

        self.emit("PyObject* tmp = NULL;", 1)

        self.emit("", 0)

    def sumTrailer(self, name):

        self.emit("", 0)

        self.emit("tmp = PyObject_Repr(obj);", 1)

        # there's really nothing more we can do if this fails ...

        self.emit("if (tmp == NULL) goto failed;", 1)

        error = "expected some sort of %s, but got %%.400s" % name

        format = "PyErr_Format(PyExc_TypeError, \"%s\", PyString_AS_STRING(tmp));"

        self.emit(format % error, 1, reflow=False)

        self.emit("failed:", 0)

        self.emit("Py_XDECREF(tmp);", 1)

        self.emit("return 1;", 1)

        self.emit("}", 0)

        self.emit("", 0)

    def simpleSum(self, sum, name):

        self.funcHeader(name)

        for t in sum.types:

            self.emit("if (PyObject_IsInstance(obj, (PyObject*)%s_type)) {" % t.name, 1)

            self.emit("*out = %s;" % t.name, 2)

            self.emit("return 0;", 2)

            self.emit("}", 1)

        self.sumTrailer(name)

    def buildArgs(self, fields):

        return ", ".join(fields + ["arena"])

    def complexSum(self, sum, name):

        self.funcHeader(name)

        for a in sum.attributes:

            self.visitAttributeDeclaration(a, name, sum=sum)

        self.emit("", 0)

        # XXX: should we only do this for 'expr'?

        self.emit("if (obj == Py_None) {", 1)

        self.emit("*out = NULL;", 2)

        self.emit("return 0;", 2)

        self.emit("}", 1)

        for a in sum.attributes:

            self.visitField(a, name, sum=sum, depth=1)

        for t in sum.types:

            self.emit("if (PyObject_IsInstance(obj, (PyObject*)%s_type)) {" % t.name, 1)

            for f in t.fields:

                self.visitFieldDeclaration(f, t.name, sum=sum, depth=2)

            self.emit("", 0)

            for f in t.fields:

                self.visitField(f, t.name, sum=sum, depth=2)

            args = [f.name.value for f in t.fields] + [a.name.value for a in sum.attributes]

            self.emit("*out = %s(%s);" % (t.name, self.buildArgs(args)), 2)

            self.emit("if (*out == NULL) goto failed;", 2)

            self.emit("return 0;", 2)

            self.emit("}", 1)

        self.sumTrailer(name)

    def visitAttributeDeclaration(self, a, name, sum=sum):

        ctype = get_c_type(a.type)

        self.emit("%s %s;" % (ctype, a.name), 1)

    def visitSum(self, sum, name):

        if is_simple(sum):

            self.simpleSum(sum, name)

        else:

            self.complexSum(sum, name)

    def visitProduct(self, prod, name):

        ctype = get_c_type(name)

        self.emit("int", 0)

        self.emit("obj2ast_%s(PyObject* obj, %s* out, PyArena* arena)" % (name, ctype), 0)

        self.emit("{", 0)

        self.emit("PyObject* tmp = NULL;", 1)

        for f in prod.fields:

            self.visitFieldDeclaration(f, name, prod=prod, depth=1)

        self.emit("", 0)

        for f in prod.fields:

            self.visitField(f, name, prod=prod, depth=1)

        args = [f.name.value for f in prod.fields]

        self.emit("*out = %s(%s);" % (name, self.buildArgs(args)), 1)

        self.emit("return 0;", 1)

        self.emit("failed:", 0)

        self.emit("Py_XDECREF(tmp);", 1)

        self.emit("return 1;", 1)

        self.emit("}", 0)

        self.emit("", 0)

    def visitFieldDeclaration(self, field, name, sum=None, prod=None, depth=0):

        ctype = get_c_type(field.type)

        if field.seq:

            if self.isSimpleType(field):

                self.emit("asdl_int_seq* %s;" % field.name, depth)

            else:

                self.emit("asdl_seq* %s;" % field.name, depth)

        else:

            ctype = get_c_type(field.type)

            self.emit("%s %s;" % (ctype, field.name), depth)

    def isSimpleSum(self, field):

        # XXX can the members of this list be determined automatically?

        return field.type.value in ('expr_context', 'boolop', 'operator',

                                    'unaryop', 'cmpop')

    def isNumeric(self, field):

        return get_c_type(field.type) in ("int", "bool")

    def isSimpleType(self, field):

        return self.isSimpleSum(field) or self.isNumeric(field)

    def visitField(self, field, name, sum=None, prod=None, depth=0):

        ctype = get_c_type(field.type)

        self.emit("if (PyObject_HasAttrString(obj, \"%s\")) {" % field.name, depth)

        self.emit("int res;", depth+1)

        if field.seq:

            self.emit("Py_ssize_t len;", depth+1)

            self.emit("Py_ssize_t i;", depth+1)

        self.emit("tmp = PyObject_GetAttrString(obj, \"%s\");" % field.name, depth+1)

        self.emit("if (tmp == NULL) goto failed;", depth+1)

        if field.seq:

            self.emit("if (!PyList_Check(tmp)) {", depth+1)

            self.emit("PyErr_Format(PyExc_TypeError, \"%s field \\\"%s\\\" must "

                      "be a list, not a %%.200s\", tmp->ob_type->tp_name);" %

                      (name, field.name),

                      depth+2, reflow=False)

            self.emit("goto failed;", depth+2)

            self.emit("}", depth+1)

            self.emit("len = PyList_GET_SIZE(tmp);", depth+1)

            if self.isSimpleType(field):

                self.emit("%s = asdl_int_seq_new(len, arena);" % field.name, depth+1)

            else:

                self.emit("%s = asdl_seq_new(len, arena);" % field.name, depth+1)

            self.emit("if (%s == NULL) goto failed;" % field.name, depth+1)

            self.emit("for (i = 0; i < len; i++) {", depth+1)

            self.emit("%s value;" % ctype, depth+2)

            self.emit("res = obj2ast_%s(PyList_GET_ITEM(tmp, i), &value, arena);" %

                      field.type, depth+2, reflow=False)

            self.emit("if (res != 0) goto failed;", depth+2)

            self.emit("asdl_seq_SET(%s, i, value);" % field.name, depth+2)

            self.emit("}", depth+1)

        else:

            self.emit("res = obj2ast_%s(tmp, &%s, arena);" %

                      (field.type, field.name), depth+1)

            self.emit("if (res != 0) goto failed;", depth+1)

        self.emit("Py_XDECREF(tmp);", depth+1)

        self.emit("tmp = NULL;", depth+1)

        self.emit("} else {", depth)

        if not field.opt:

            message = "required field \\\"%s\\\" missing from %s" % (field.name, name)

            format = "PyErr_SetString(PyExc_TypeError, \"%s\");"

            self.emit(format % message, depth+1, reflow=False)

            self.emit("return 1;", depth+1)

        else:

            if self.isNumeric(field):

                self.emit("%s = 0;" % field.name, depth+1)

            elif not self.isSimpleType(field):

                self.emit("%s = NULL;" % field.name, depth+1)

            else:

                raise TypeError("could not determine the default value for %s" % field.name)

        self.emit("}", depth)

class MarshalPrototypeVisitor(PickleVisitor):