asdl_c.py

#! /usr/bin/env python
"""Generate C code from an ASDL description."""

import sys
import textwrap
import types

from argparse import ArgumentParser
from contextlib import contextmanager
from pathlib import Path

import asdl

TABSIZE = 4
MAX_COL = 80
AUTOGEN_MESSAGE = "// File automatically generated by {}.\n\n"

builtin_type_to_c_type = {
    "identifier": "PyUnicode_Type",
    "string": "PyUnicode_Type",
    "int": "PyLong_Type",
    "constant": "PyBaseObject_Type",
}

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.
    """
    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 reflow_c_string(s, depth):
    return '"%s"' % s.replace('\n', '\\n"\n%s"' % (' ' * depth * TABSIZE))

def is_simple(sum_type):
    """Return True if a sum is a simple.

    A sum is simple if its types have no fields and itself
    doesn't have any attributes. Instances of these types are
    cached at C level, and they act like singletons when propagating
    parser generated nodes into Python level, e.g.
    unaryop = Invert | Not | UAdd | USub
    """

    return not (
        sum_type.attributes or
        any(constructor.fields for constructor in sum_type.types)
    )

def asdl_of(name, obj):
    if isinstance(obj, asdl.Product) or isinstance(obj, asdl.Constructor):
        fields = ", ".join(map(str, obj.fields))
        if fields:
            fields = "({})".format(fields)
        return "{}{}".format(name, fields)
    else:
        if is_simple(obj):
            types = " | ".join(type.name for type in obj.types)
        else:
            sep = "\n{}| ".format(" " * (len(name) + 1))
            types = sep.join(
                asdl_of(type.name, type) for type in obj.types
            )
        return "{} = {}".format(name, types)

class EmitVisitor(asdl.VisitorBase):
    """Visit that emits lines"""

    def __init__(self, file, metadata = None):
        self.file = file
        self._metadata = metadata
        super(EmitVisitor, self).__init__()

    def emit(self, s, depth, reflow=True):
        # XXX reflow long lines?
        if reflow:
            lines = reflow_lines(s, depth)
        else:
            lines = [s]
        for line in lines:
            if line:
                line = (" " * TABSIZE * depth) + line
            self.file.write(line + "\n")

    @property
    def metadata(self):
        if self._metadata is None:
            raise ValueError(
                "%s was expecting to be annotated with metadata"
                % type(self).__name__
            )
        return self._metadata

    @metadata.setter
    def metadata(self, value):
        self._metadata = value

class MetadataVisitor(asdl.VisitorBase):
    ROOT_TYPE = "AST"

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        # Metadata:
        #    - simple_sums: Tracks the list of compound type
        #                   names where all the constructors
        #                   belonging to that type lack of any
        #                   fields.
        #    - identifiers: All identifiers used in the AST declarations
        #    - singletons:  List of all constructors that originates from
        #                   simple sums.
        #    - types:       List of all top level type names
        #
        self.metadata = types.SimpleNamespace(
            simple_sums=set(),
            identifiers=set(),
            singletons=set(),
            types={self.ROOT_TYPE},
        )

    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):
        self.metadata.types.add(name)

        simple_sum = is_simple(sum)
        if simple_sum:
            self.metadata.simple_sums.add(name)

        for constructor in sum.types:
            if simple_sum:
                self.metadata.singletons.add(constructor.name)
            self.visitConstructor(constructor)
        self.visitFields(sum.attributes)

    def visitConstructor(self, constructor):
        self.metadata.types.add(constructor.name)
        self.visitFields(constructor.fields)

    def visitProduct(self, product, name):
        self.metadata.types.add(name)
        self.visitFields(product.attributes)
        self.visitFields(product.fields)

    def visitFields(self, fields):
        for field in fields:
            self.visitField(field)

    def visitField(self, field):
        self.metadata.identifiers.add(field.name)


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 SequenceDefVisitor(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):
            return
        self.emit_sequence_constructor(name, depth)

    def emit_sequence_constructor(self, name,depth):
        ctype = get_c_type(name)
        self.emit("""\
typedef struct {
    _ASDL_SEQ_HEAD
    %(ctype)s typed_elements[1];
} asdl_%(name)s_seq;""" % locals(), reflow=False, depth=depth)
        self.emit("", depth)
        self.emit("asdl_%(name)s_seq *_Py_asdl_%(name)s_seq_new(Py_ssize_t size, PyArena *arena);" % locals(), depth)
        self.emit("", depth)

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

class StructVisitor(EmitVisitor):
    """Visitor to generate typedefs 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)

    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 in self.metadata.simple_sums:
                self.emit("asdl_int_seq *%(name)s;" % locals(), depth)
            else:
                _type = field.type
                self.emit("asdl_%(_type)s_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)
        for field in product.attributes:
            # rudimentary attribute handling
            type = str(field.type)
            assert type in asdl.builtin_types, type
            self.emit("%s %s;" % (type, field.name), depth + 1);
        self.emit("};", depth)
        self.emit("", depth)


def ast_func_name(name):
    return f"_PyAST_{name}"


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 info, 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 in self.metadata.simple_sums:
                    ctype = "asdl_int_seq *"
                else:
                    ctype = f"asdl_{f.type}_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=True):
        args = args + attrs
        if args:
            argstr = ", ".join(["%s %s" % (atype, aname)
                                for atype, aname, opt in args])
            argstr += ", PyArena *arena"
        else:
            argstr = "PyArena *arena"
        self.emit("%s %s(%s);" % (ctype, ast_func_name(name), argstr), False)

    def visitProduct(self, prod, name):
        self.emit_function(name, get_c_type(name),
                           self.get_args(prod.fields),
                           self.get_args(prod.attributes),
                           union=False)


class FunctionVisitor(PrototypeVisitor):
    """Visitor to generate constructor functions for AST."""

    def emit_function(self, name, ctype, args, attrs, union=True):
        def emit(s, depth=0, reflow=True):
            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)" % (ast_func_name(name), argstr))
        emit("{")
        emit("%s p;" % ctype, 1)
        for argtype, argname, opt in args:
            if not opt and 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=False)
                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=True):
            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=True):
            self.emit(s, depth, reflow)
        for argtype, argname, opt in args:
            emit("p->%s = %s;" % (argname, argname), 1)
        for argtype, argname, opt in attrs:
            emit("p->%s = %s;" % (argname, argname), 1)


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(struct ast_state *state, PyObject* obj, %s* out, const char* field, PyArena* arena);"
        self.emit(code % (name, get_c_type(name)), 0)

    visitSum = visitProduct


class Obj2ModVisitor(PickleVisitor):

    attribute_special_defaults = {
        "end_lineno": "lineno",
        "end_col_offset": "col_offset",
    }

    @contextmanager
    def recursive_call(self, node, level):
        self.emit('if (_Py_EnterRecursiveCall(" while traversing \'%s\' node")) {' % node, level, reflow=False)
        self.emit('goto failed;', level + 1)
        self.emit('}', level)
        yield
        self.emit('_Py_LeaveRecursiveCall();', level)

    def funcHeader(self, name):
        ctype = get_c_type(name)
        self.emit("int", 0)
        self.emit("obj2ast_%s(struct ast_state *state, PyObject* obj, %s* out, const char* field, PyArena* arena)" % (name, ctype), 0)
        self.emit("{", 0)
        self.emit("int isinstance;", 1)
        self.emit("", 0)

    def sumTrailer(self, name, add_label=False):
        self.emit("", 0)
        # there's really nothing more we can do if this fails ...
        error = "expected some sort of %s, but got %%R" % name
        format = "PyErr_Format(PyExc_TypeError, \"%s\", obj);"
        self.emit(format % error, 1, reflow=False)
        if add_label:
            self.emit("failed:", 1)
            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:
            line = ("isinstance = PyObject_IsInstance(obj, "
                    "state->%s_type);")
            self.emit(line % (t.name,), 1)
            self.emit("if (isinstance == -1) {", 1)
            self.emit("return -1;", 2)
            self.emit("}", 1)
            self.emit("if (isinstance) {", 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 typeCheck(self, name):
        self.emit("tp = state->%s_type;" % name, 1)
        self.emit("isinstance = PyObject_IsInstance(obj, tp);", 1)
        self.emit("if (isinstance == -1) {", 1)
        self.emit("return 1;", 2)
        self.emit("}", 1)
        self.emit("if (!isinstance && field != NULL) {", 1)
        error = "field '%%s' was expecting node of type '%s', got '%%T'" % name
        self.emit("PyErr_Format(PyExc_TypeError, \"%s\", field, obj);" % error, 2, reflow=False)
        self.emit("return 1;", 2)
        self.emit("}", 1)

    def complexSum(self, sum, name):
        self.funcHeader(name)
        self.emit("PyObject *tmp = NULL;", 1)
        self.emit("PyObject *tp;", 1)
        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)
        self.typeCheck(name)
        for a in sum.attributes:
            self.visitField(a, name, sum=sum, depth=1)
        for t in sum.types:
            self.emit("tp = state->%s_type;" % (t.name,), 1)
            self.emit("isinstance = PyObject_IsInstance(obj, tp);", 1)
            self.emit("if (isinstance == -1) {", 1)
            self.emit("return -1;", 2)
            self.emit("}", 1)
            self.emit("if (isinstance) {", 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 for f in t.fields] + [a.name for a in sum.attributes]
            self.emit("*out = %s(%s);" % (ast_func_name(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, True)

    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(struct ast_state *state, PyObject* obj, %s* out, const char* field, 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)
        for a in prod.attributes:
            self.visitFieldDeclaration(a, name, prod=prod, depth=1)
        self.emit("", 0)
        for f in prod.fields:
            self.visitField(f, name, prod=prod, depth=1)
        for a in prod.attributes:
            self.visitField(a, name, prod=prod, depth=1)
        args = [f.name for f in prod.fields]
        args.extend([a.name for a in prod.attributes])
        self.emit("*out = %s(%s);" % (ast_func_name(name), self.buildArgs(args)), 1)
        self.emit("if (*out == NULL) goto failed;", 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:
                _type = field.type
                self.emit(f"asdl_{field.type}_seq* {field.name};", depth)
        else:
            ctype = get_c_type(field.type)
            self.emit("%s %s;" % (ctype, field.name), depth)

    def isNumeric(self, field):
        return get_c_type(field.type) in ("int", "bool")

    def isSimpleType(self, field):
        return field.type in self.metadata.simple_sums or self.isNumeric(field)

    def visitField(self, field, name, sum=None, prod=None, depth=0):
        ctype = get_c_type(field.type)
        line = "if (PyObject_GetOptionalAttr(obj, state->%s, &tmp) < 0) {"
        self.emit(line % field.name, depth)
        self.emit("return -1;", depth+1)
        self.emit("}", depth)
        if field.seq:
            self.emit("if (tmp == NULL) {", depth)
            self.emit("tmp = PyList_New(0);", depth+1)
            self.emit("if (tmp == NULL) {", depth+1)
            self.emit("return -1;", depth+2)
            self.emit("}", depth+1)
            self.emit("}", depth)
            self.emit("{", depth)
        else:
            if not field.opt:
                self.emit("if (tmp == NULL) {", depth)
                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:
                self.emit("if (tmp == NULL || tmp == Py_None) {", depth)
                self.emit("Py_CLEAR(tmp);", depth+1)
                if self.isNumeric(field):
                    if field.name in self.attribute_special_defaults:
                        self.emit(
                            "%s = %s;" % (field.name, self.attribute_special_defaults[field.name]),
                            depth+1,
                        )
                    else:
                        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)
            self.emit("else {", 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("if (!PyList_Check(tmp)) {", depth+1)
            self.emit("PyErr_Format(PyExc_TypeError, \"%s field \\\"%s\\\" must "
                      "be a list, not a %%T\", tmp);" %
                      (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 = _Py_asdl_int_seq_new(len, arena);" % field.name, depth+1)
            else:
                self.emit("%s = _Py_asdl_%s_seq_new(len, arena);" % (field.name, field.type), 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 val;" % ctype, depth+2)
            self.emit("PyObject *tmp2 = Py_NewRef(PyList_GET_ITEM(tmp, i));", depth+2)
            with self.recursive_call(name, depth+2):
                self.emit("res = obj2ast_%s(state, tmp2, &val, \"%s\", arena);" %
                          (field.type, field.name), depth+2, reflow=False)
            self.emit("Py_DECREF(tmp2);", depth+2)
            self.emit("if (res != 0) goto failed;", depth+2)
            self.emit("if (len != PyList_GET_SIZE(tmp)) {", depth+2)
            self.emit("PyErr_SetString(PyExc_RuntimeError, \"%s field \\\"%s\\\" "
                      "changed size during iteration\");" %
                      (name, field.name),
                      depth+3, reflow=False)
            self.emit("goto failed;", depth+3)
            self.emit("}", depth+2)
            self.emit("asdl_seq_SET(%s, i, val);" % field.name, depth+2)
            self.emit("}", depth+1)
        else:
            with self.recursive_call(name, depth+1):
                self.emit("res = obj2ast_%s(state, tmp, &%s, \"%s\", arena);" %
                          (field.type, field.name, field.name), depth+1)
            self.emit("if (res != 0) goto failed;", depth+1)

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


class SequenceConstructorVisitor(EmitVisitor):
    def visitModule(self, mod):
        for dfn in mod.dfns:
            self.visit(dfn)

    def visitType(self, type):
        self.visit(type.value, type.name)

    def visitProduct(self, prod, name):
        self.emit_sequence_constructor(name, get_c_type(name))

    def visitSum(self, sum, name):
        if not is_simple(sum):
            self.emit_sequence_constructor(name, get_c_type(name))

    def emit_sequence_constructor(self, name, type):
        self.emit(f"GENERATE_ASDL_SEQ_CONSTRUCTOR({name}, {type})", depth=0)

class PyTypesDeclareVisitor(PickleVisitor):

    def visitProduct(self, prod, name):
        self.emit("static PyObject* ast2obj_%s(struct ast_state *state, void*);" % name, 0)
        if prod.attributes:
            self.emit("static const char * const %s_attributes[] = {" % name, 0)
            for a in prod.attributes:
                self.emit('"%s",' % a.name, 1)
            self.emit("};", 0)
        if prod.fields:
            self.emit("static const char * const %s_fields[]={" % name,0)
            for f in prod.fields:
                self.emit('"%s",' % f.name, 1)
            self.emit("};", 0)

    def visitSum(self, sum, name):
        if sum.attributes:
            self.emit("static const char * const %s_attributes[] = {" % name, 0)
            for a in sum.attributes:
                self.emit('"%s",' % a.name, 1)
            self.emit("};", 0)
        ptype = "void*"
        if is_simple(sum):
            ptype = get_c_type(name)
        self.emit("static PyObject* ast2obj_%s(struct ast_state *state, %s);" % (name, ptype), 0)
        for t in sum.types:
            self.visitConstructor(t, name)

    def visitConstructor(self, cons, name):
        if cons.fields:
            self.emit("static const char * const %s_fields[]={" % cons.name, 0)
            for t in cons.fields:
                self.emit('"%s",' % t.name, 1)
            self.emit("};",0)


class AnnotationsVisitor(PickleVisitor):
    def visitModule(self, mod):
        self.nodes = []
        for dfn in mod.dfns:
            self.visit(dfn)
        builtins = list(dict.fromkeys(builtin_type_to_c_type.values()))
        self.file.write(textwrap.dedent('''
            static int
            add_ast_annotations(struct ast_state *state)
            {
                enum {
                    FIELD_OPTIONAL = 1,
                    FIELD_SEQUENCE = 2,
                    FIELD_BUILTIN = 4,
                };
                PyTypeObject *const builtin_types[] = {
        '''))
        for c_type in builtins:
            self.emit(f"&{c_type},", 2)
        self.file.write(textwrap.dedent('''\
            };
            // Offsets refer to this interpreter's AST state, not global types.
            static_assert(sizeof(struct ast_state) <= UINT16_MAX,
                          "ast_state offsets must fit in uint16_t");
            static const struct {
                uint16_t name_offset;
                uint16_t type_offset;  // An index into builtin_types for builtins.
                uint8_t flags;
            } fields[] = {
        #''').removesuffix('#'))  # Use d-string if it accepted.
        for name, fields in self.nodes:
            for field in fields:
                flags = []
                if field.opt:
                    flags.append("FIELD_OPTIONAL")
                elif field.seq:
                    flags.append("FIELD_SEQUENCE")
                if field.type in builtin_type_to_c_type:
                    c_type = builtin_type_to_c_type[field.type]
                    type_offset = str(builtins.index(c_type))
                    flags.append("FIELD_BUILTIN")
                else:
                    type_offset = f"offsetof(struct ast_state, {field.type}_type)"
                flags = " | ".join(flags) or "0"
                self.emit(f"{{offsetof(struct ast_state, {field.name}),", 2)
                self.emit(f" {type_offset}, {flags}}},", 2)
        self.file.write(textwrap.dedent('''\
            };
            static const struct {
                uint16_t type_offset;
                uint16_t first_field;
                uint16_t nfields;
            } nodes[] = {
        #''').removesuffix('#'))
        start = 0
        for name, fields in self.nodes:
            self.emit(f"{{offsetof(struct ast_state, {name}_type), "
                      f"{start}, {len(fields)}}},", 2)
            start += len(fields)
        self.file.write(textwrap.dedent('''\
                };
                char *base = (char *)state;
                PyObject *annotations = NULL;
                for (size_t i = 0; i < Py_ARRAY_LENGTH(nodes); i++) {
                    annotations = PyDict_New();
                    if (annotations == NULL) {
                        goto error;
                    }
                    size_t end = nodes[i].first_field + nodes[i].nfields;
                    for (size_t j = nodes[i].first_field; j < end; j++) {
                        PyObject *name = *(PyObject **)(base + fields[j].name_offset);
                        PyObject *type;
                        if (fields[j].flags & FIELD_BUILTIN) {
                            type = (PyObject *)builtin_types[fields[j].type_offset];
                        }
                        else {
                            type = *(PyObject **)(base + fields[j].type_offset);
                        }
                        if (fields[j].flags & FIELD_OPTIONAL) {
                            type = _Py_union_type_or(type, Py_None);
                        }
                        else if (fields[j].flags & FIELD_SEQUENCE) {
                            type = Py_GenericAlias((PyObject *)&PyList_Type, type);
                        }
                        else {
                            Py_INCREF(type);
                        }
                        if (type == NULL) {
                            goto error;
                        }
                        int res = PyDict_SetItem(annotations, name, type);
                        Py_DECREF(type);
                        if (res < 0) {
                            goto error;
                        }
                    }
                    PyObject *node = *(PyObject **)(base + nodes[i].type_offset);
                    if (PyObject_SetAttrString(node, "_field_types", annotations) < 0 ||
                        PyObject_SetAttrString(node, "__annotations__", annotations) < 0)
                    {
                        goto error;
                    }
                    Py_CLEAR(annotations);
                }
                return 1;
            error:
                Py_XDECREF(annotations);
                return 0;
            }
        '''))

    def visitProduct(self, prod, name):
        self.nodes.append((name, prod.fields))

    def visitSum(self, sum, name):
        for t in sum.types:
            self.visitConstructor(t, name)

    def visitConstructor(self, cons, name):
        self.nodes.append((cons.name, cons.fields))


class PyTypesVisitor(PickleVisitor):

    def visitModule(self, mod):
        self.emit("""

typedef struct {
    PyObject_HEAD
} AST_object;

static void
ast_dealloc(PyObject *op)
{
    AST_object *self = (AST_object*)op;
    /* bpo-31095: UnTrack is needed before calling any callbacks */
    PyTypeObject *tp = Py_TYPE(self);
    PyObject_GC_UnTrack(self);
    PyObject_ClearManagedDict(op);
    freefunc free_func = PyType_GetSlot(tp, Py_tp_free);
    assert(free_func != NULL);
    free_func(self);
    Py_DECREF(tp);
}

static int
ast_traverse(PyObject *op, visitproc visit, void *arg)
{
    AST_object *self = (AST_object*)op;
    Py_VISIT(Py_TYPE(self));
    return PyObject_VisitManagedDict(op, visit, arg);
}

static int
ast_clear(PyObject *op)
{
    PyObject_ClearManagedDict(op);
    return 0;
}

/*
 * Format the names in the set 'missing' into a natural language list,
 * sorted in the order in which they appear in 'fields'.
 *
 * Similar to format_missing() from 'Python/ceval.c'.
 *
 * Parameters
 *
 *      missing     Set of missing field names to render.
 *      fields      Sequence of AST node field names (self._fields).
 */
static PyObject *
format_missing(PyObject *missing, PyObject *fields)
{
    Py_ssize_t num_fields, num_total, num_left;
    num_fields = PySequence_Size(fields);
    if (num_fields == -1) {
        return NULL;
    }
    num_total = num_left = PySet_GET_SIZE(missing);
    PyUnicodeWriter *writer = PyUnicodeWriter_Create(0);
    if (writer == NULL) {
        goto error;
    }
    // Iterate all AST node fields in order so that the missing positional
    // arguments are rendered in the order in which __init__ expects them.
    for (Py_ssize_t i = 0; i < num_fields; i++) {
        PyObject *name = PySequence_GetItem(fields, i);
        if (name == NULL) {
            goto error;
        }
        int contains = PySet_Contains(missing, name);
        if (contains == -1) {
            Py_DECREF(name);
            goto error;
        }
        else if (contains == 1) {
            const char* fmt = NULL;
            if (num_left == 1) {
                fmt = "'%U'";
            }
            else if (num_total == 2) {
                fmt = "'%U' and ";
            }
            else if (num_left == 2) {
                fmt = "'%U', and ";
            }
            else {
                fmt = "'%U', ";
            }
            num_left--;
            if (PyUnicodeWriter_Format(writer, fmt, name) < 0) {
                Py_DECREF(name);
                goto error;
            }
        }
        Py_DECREF(name);
    }
    return PyUnicodeWriter_Finish(writer);
error:
    PyUnicodeWriter_Discard(writer);
    return NULL;
}

static int
ast_type_init(PyObject *self, PyObject *args, PyObject *kw)
{
    struct ast_state *state = get_ast_state();
    if (state == NULL) {
        return -1;
    }

    int contains = PySet_Contains(state->abstract_types, (PyObject *)Py_TYPE(self));
    if (contains == -1) {
        return -1;
    }
    else if (contains == 1) {
        if (PyErr_WarnFormat(
                PyExc_DeprecationWarning, 1,
                "Instantiating abstract AST node class %T is deprecated. "
                "This will become an error in Python 3.20", self) < 0) {
            return -1;
        }
    }

    Py_ssize_t i, numfields = 0;
    int res = -1;
    PyObject *key, *value, *fields, *attributes = NULL, *remaining_fields = NULL;

    fields = PyObject_GetAttr((PyObject*)Py_TYPE(self), state->_fields);
    if (fields == NULL) {
        goto cleanup;
    }

    numfields = PySequence_Size(fields);
    if (numfields == -1) {
        goto cleanup;
    }
    remaining_fields = PySet_New(fields);
    if (remaining_fields == NULL) {
        goto cleanup;
    }

    res = 0; /* if no error occurs, this stays 0 to the end */
    if (numfields < PyTuple_GET_SIZE(args)) {
        PyErr_Format(PyExc_TypeError, "%T constructor takes at most "
                     "%zd positional argument%s",
                     self, numfields, numfields == 1 ? "" : "s");
        res = -1;
        goto cleanup;
    }
    for (i = 0; i < PyTuple_GET_SIZE(args); i++) {
        /* cannot be reached when fields is NULL */
        PyObject *name = PySequence_GetItem(fields, i);
        if (!name) {
            res = -1;
            goto cleanup;
        }
        res = PyObject_SetAttr(self, name, PyTuple_GET_ITEM(args, i));
        if (PySet_Discard(remaining_fields, name) < 0) {
            res = -1;
            Py_DECREF(name);
            goto cleanup;
        }
        Py_DECREF(name);
        if (res < 0) {
            goto cleanup;
        }
    }
    if (kw) {
        i = 0;  /* needed by PyDict_Next */
        while (PyDict_Next(kw, &i, &key, &value)) {
            int contains = PySequence_Contains(fields, key);
            if (contains == -1) {
                res = -1;
                goto cleanup;
            }
            else if (contains == 1) {
                int p = PySet_Discard(remaining_fields, key);
                if (p == -1) {
                    res = -1;
                    goto cleanup;
                }
                if (p == 0) {
                    PyErr_Format(PyExc_TypeError,
                        "%T got multiple values for argument %R",
                        self, key);
                    res = -1;
                    goto cleanup;
                }
            }
            else {
                // Lazily initialize "attributes"
                if (attributes == NULL) {
                    attributes = PyObject_GetAttr((PyObject*)Py_TYPE(self), state->_attributes);
                    if (attributes == NULL) {
                        res = -1;
                        goto cleanup;
                    }
                }
                int contains = PySequence_Contains(attributes, key);
                if (contains == -1) {
                    res = -1;
                    goto cleanup;
                }
                else if (contains == 0) {
                    PyErr_Format(PyExc_TypeError,
                        "%T.__init__ got an unexpected keyword argument %R",
                        self, key);
                    res = -1;
                    goto cleanup;
                }
            }
            res = PyObject_SetAttr(self, key, value);
            if (res < 0) {
                goto cleanup;
            }
        }
    }
    Py_ssize_t size = PySet_Size(remaining_fields);
    PyObject *field_types = NULL, *remaining_list = NULL, *missing_names = NULL;
    if (size > 0) {
        if (PyObject_GetOptionalAttr((PyObject*)Py_TYPE(self), &_Py_ID(_field_types),
                                     &field_types) < 0) {
            res = -1;
            goto cleanup;
        }
        if (field_types == NULL) {
            // Probably a user-defined subclass of AST that lacks _field_types.
            // This will continue to work as it did before 3.13; i.e., attributes
            // that are not passed in simply do not exist on the instance.
            goto cleanup;
        }
        remaining_list = PySequence_List(remaining_fields);
        if (!remaining_list) {
            goto set_remaining_cleanup;
        }
        missing_names = PySet_New(NULL);
        if (!missing_names) {
            goto set_remaining_cleanup;
        }
        for (Py_ssize_t i = 0; i < size; i++) {
            PyObject *name = PyList_GET_ITEM(remaining_list, i);
            PyObject *type = PyDict_GetItemWithError(field_types, name);
            if (!type) {
                if (PyErr_Occurred()) {
                    goto set_remaining_cleanup;
                }
                else {
                    PyErr_Format(PyExc_TypeError,
                        "Field %R is missing from %T._field_types",
                        name, self);
                    goto set_remaining_cleanup;
                }
            }
            else if (_PyUnion_Check(type)) {
                // optional field
                // do nothing, we'll have set a None default on the class
            }
            else if (Py_IS_TYPE(type, &Py_GenericAliasType)) {
                // list field
                PyObject *empty = PyList_New(0);
                if (!empty) {
                    goto set_remaining_cleanup;
                }
                res = PyObject_SetAttr(self, name, empty);
                Py_DECREF(empty);
                if (res < 0) {
                    goto set_remaining_cleanup;
                }
            }
            else if (type == state->expr_context_type) {
                // special case for expr_context: default to Load()
                res = PyObject_SetAttr(self, name, state->Load_singleton);
                if (res < 0) {
                    goto set_remaining_cleanup;
                }
            }
            else {
                // simple field (e.g., identifier)
                res = PySet_Add(missing_names, name);
                if (res < 0) {
                    goto set_remaining_cleanup;
                }
            }
        }
        Py_ssize_t num_missing = PySet_GET_SIZE(missing_names);
        if (num_missing > 0) {
            PyObject *name_str = format_missing(missing_names, fields);
            if (!name_str) {
                goto set_remaining_cleanup;
            }
            PyErr_Format(PyExc_TypeError,
                "%T.__init__ missing %d required positional argument%s: %U",
                self, num_missing, num_missing == 1 ? "" : "s", name_str);
            Py_DECREF(name_str);
            goto set_remaining_cleanup;
        }
        Py_DECREF(missing_names);
        Py_DECREF(remaining_list);
        Py_DECREF(field_types);
    }
  cleanup:
    Py_XDECREF(attributes);
    Py_XDECREF(fields);
    Py_XDECREF(remaining_fields);
    return res;
  set_remaining_cleanup:
    Py_XDECREF(missing_names);
    Py_XDECREF(remaining_list);
    Py_XDECREF(field_types);
    res = -1;
    goto cleanup;
}

/* Pickling support */
static PyObject *
ast_type_reduce(PyObject *self, PyObject *unused)
{
    struct ast_state *state = get_ast_state();
    if (state == NULL) {
        return NULL;
    }

    PyObject *dict = NULL, *fields = NULL, *positional_args = NULL;
    if (PyObject_GetOptionalAttr(self, state->__dict__, &dict) < 0) {
        return NULL;
    }
    PyObject *result = NULL;
    if (dict) {
        // Unpickling (or copying) works as follows:
        // - Construct the object with only positional arguments
        // - Set the fields from the dict
        // We have two constraints:
        // - We must set all the required fields in the initial constructor call,
        //   or the unpickling or deepcopying of the object will trigger DeprecationWarnings.
        // - We must not include child nodes in the positional args, because
        //   that may trigger runaway recursion during copying (gh-120108).
        // To satisfy both constraints, we set all the fields to None in the
        // initial list of positional args, and then set the fields from the dict.
        if (PyObject_GetOptionalAttr((PyObject*)Py_TYPE(self), state->_fields, &fields) < 0) {
            goto cleanup;
        }
        if (fields) {
            Py_ssize_t numfields = PySequence_Size(fields);
            if (numfields == -1) {
                Py_DECREF(dict);
                goto cleanup;
            }
            positional_args = PyList_New(0);
            if (!positional_args) {
                goto cleanup;
            }
            for (Py_ssize_t i = 0; i < numfields; i++) {
                PyObject *name = PySequence_GetItem(fields, i);
                if (!name) {
                    goto cleanup;
                }
                PyObject *value;
                int rc = PyDict_GetItemRef(dict, name, &value);
                Py_DECREF(name);
                if (rc < 0) {
                    goto cleanup;
                }
                if (!value) {
                    break;
                }
                rc = PyList_Append(positional_args, Py_None);
                Py_DECREF(value);
                if (rc < 0) {
                    goto cleanup;
                }
            }
            PyObject *args_tuple = PyList_AsTuple(positional_args);
            if (!args_tuple) {
                goto cleanup;
            }
            result = Py_BuildValue("ONN", Py_TYPE(self), args_tuple, dict);
        }
        else {
            result = Py_BuildValue("O()N", Py_TYPE(self), dict);
        }
    }
    else {
        result = Py_BuildValue("O()", Py_TYPE(self));
    }
cleanup:
    Py_XDECREF(fields);
    Py_XDECREF(positional_args);
    return result;
}

/*
 * Python equivalent:
 *
 *   for key in keys:
 *       if hasattr(self, key):
 *           payload[key] = getattr(self, key)
 *
 * The 'keys' argument is a sequence corresponding to
 * the '_fields' or the '_attributes' of an AST node.
 *
 * This returns -1 if an error occurs and 0 otherwise.
 *
 * Parameters
 *
 *      payload   A dictionary to fill.
 *      keys      A sequence of keys or NULL for an empty sequence.
 *      dict      The AST node instance dictionary (must not be NULL).
 */
static inline int
ast_type_replace_update_payload(PyObject *payload,
                                PyObject *keys,
                                PyObject *dict)
{
    assert(dict != NULL);
    if (keys == NULL) {
        return 0;
    }
    Py_ssize_t n = PySequence_Size(keys);
    if (n == -1) {
        return -1;
    }
    for (Py_ssize_t i = 0; i < n; i++) {
        PyObject *key = PySequence_GetItem(keys, i);
        if (key == NULL) {
            return -1;
        }
        PyObject *value;
        if (PyDict_GetItemRef(dict, key, &value) < 0) {
            Py_DECREF(key);
            return -1;
        }
        if (value == NULL) {
            Py_DECREF(key);
            // If a field or attribute is not present at runtime, it should
            // be explicitly given in 'kwargs'. If not, the constructor will
            // issue a warning (which becomes an error in 3.15).
            continue;
        }
        int rc = PyDict_SetItem(payload, key, value);
        Py_DECREF(key);
        Py_DECREF(value);
        if (rc < 0) {
            return -1;
        }
    }
    return 0;
}

/* copy.replace() support (shallow copy) */
static PyObject *
ast_type_replace(PyObject *self, PyObject *args, PyObject *kwargs)
{
    if (!_PyArg_NoPositional("__replace__", args)) {
        return NULL;
    }

    struct ast_state *state = get_ast_state();
    if (state == NULL) {
        return NULL;
    }

    PyObject *result = NULL;
    // known AST class fields and attributes
    PyObject *fields = NULL, *attributes = NULL;
    // current instance dictionary
    PyObject *dict = NULL;
    // constructor positional and keyword arguments
    PyObject *empty_tuple = NULL, *payload = NULL;

    PyObject *type = (PyObject *)Py_TYPE(self);
    if (PyObject_GetOptionalAttr(type, state->_fields, &fields) < 0) {
        goto cleanup;
    }
    if (PyObject_GetOptionalAttr(type, state->_attributes, &attributes) < 0) {
        goto cleanup;
    }
    if (PyObject_GetOptionalAttr(self, state->__dict__, &dict) < 0) {
        goto cleanup;
    }
    empty_tuple = PyTuple_New(0);
    if (empty_tuple == NULL) {
        goto cleanup;
    }
    payload = PyDict_New();
    if (payload == NULL) {
        goto cleanup;
    }
    if (dict) { // in case __dict__ is missing (for some obscure reason)
        // copy the instance's fields (possibly NULL)
        if (ast_type_replace_update_payload(payload, fields, dict) < 0) {
            goto cleanup;
        }
        // copy the instance's attributes (possibly NULL)
        if (ast_type_replace_update_payload(payload, attributes, dict) < 0) {
            goto cleanup;
        }
    }
    if (kwargs && PyDict_Update(payload, kwargs) < 0) {
        goto cleanup;
    }
    result = PyObject_Call(type, empty_tuple, payload);
cleanup:
    Py_XDECREF(payload);
    Py_XDECREF(empty_tuple);
    Py_XDECREF(dict);
    Py_XDECREF(attributes);
    Py_XDECREF(fields);
    return result;
}

static PyMethodDef ast_type_methods[] = {
    {"__reduce__", ast_type_reduce, METH_NOARGS, NULL},
    {"__replace__", _PyCFunction_CAST(ast_type_replace), METH_VARARGS | METH_KEYWORDS,
     PyDoc_STR("__replace__($self, /, **fields)\\n--\\n\\n"
               "Return a copy of the AST node with new values "
               "for the specified fields.")},
    {NULL}
};

static PyGetSetDef ast_type_getsets[] = {
    {"__dict__", PyObject_GenericGetDict, PyObject_GenericSetDict},
    {NULL}
};

static PyObject *
ast_repr_max_depth(AST_object *self, int depth);

/* Format list and tuple properties of AST nodes.
   Note that, only the first and last elements are shown.
   Anything in between is represented with an ellipsis ('...').
   For example, the list [1, 2, 3] is formatted as
   'List(elts=[Constant(1), ..., Constant(3)])'. */
static PyObject *
ast_repr_list(PyObject *list, int depth)
{
    assert(PyList_Check(list) || PyTuple_Check(list));

    struct ast_state *state = get_ast_state();
    if (state == NULL) {
        return NULL;
    }

    Py_ssize_t length = PySequence_Size(list);
    if (length < 0) {
        return NULL;
    }
    else if (length == 0) {
        return PyObject_Repr(list);
    }

    PyObject *items[2] = {NULL, NULL};
    PyUnicodeWriter *writer = PyUnicodeWriter_Create(0);
    if (writer == NULL) {
        goto error;
    }

    items[0] = PySequence_GetItem(list, 0);
    if (!items[0]) {
        goto error;
    }
    if (length > 1) {
        items[1] = PySequence_GetItem(list, length - 1);
        if (!items[1]) {
            goto error;
        }
    }

    bool is_list = PyList_Check(list);
    if (PyUnicodeWriter_WriteChar(writer, is_list ? '[' : '(') < 0) {
        goto error;
    }

    for (Py_ssize_t i = 0; i < Py_MIN(length, 2); i++) {
        if (i > 0) {
            if (PyUnicodeWriter_WriteASCII(writer, ", ", 2) < 0) {
                goto error;
            }
        }

        PyObject *item = items[i];
        if (PyType_IsSubtype(Py_TYPE(item), (PyTypeObject *)state->AST_type)) {
            PyObject *item_repr;
            item_repr = ast_repr_max_depth((AST_object*)item, depth - 1);
            if (!item_repr) {
                goto error;
            }
            if (PyUnicodeWriter_WriteStr(writer, item_repr) < 0) {
                Py_DECREF(item_repr);
                goto error;
            }
            Py_DECREF(item_repr);
        } else {
            if (PyUnicodeWriter_WriteRepr(writer, item) < 0) {
                goto error;
            }
        }

        if (i == 0 && length > 2) {
            if (PyUnicodeWriter_WriteASCII(writer, ", ...", 5) < 0) {
                goto error;
            }
        }
    }

    if (PyUnicodeWriter_WriteChar(writer, is_list ? ']' : ')') < 0) {
        goto error;
    }

    Py_XDECREF(items[0]);
    Py_XDECREF(items[1]);
    return PyUnicodeWriter_Finish(writer);

error:
    Py_XDECREF(items[0]);
    Py_XDECREF(items[1]);
    PyUnicodeWriter_Discard(writer);
    return NULL;
}

static PyObject *
ast_repr_max_depth(AST_object *self, int depth)
{
    struct ast_state *state = get_ast_state();
    if (state == NULL) {
        return NULL;
    }

    PyObject *fields = PyObject_GetAttr((PyObject *)Py_TYPE(self), state->_fields);
    if (!fields) {
        return NULL;
    }

    Py_ssize_t numfields = PySequence_Size(fields);
    if (numfields < 0) {
        Py_DECREF(fields);
        return NULL;
    }

    if (numfields == 0) {
        Py_DECREF(fields);
        return PyUnicode_FromFormat("%s()", Py_TYPE(self)->tp_name);
    }

    if (depth <= 0) {
        Py_DECREF(fields);
        return PyUnicode_FromFormat("%s(...)", Py_TYPE(self)->tp_name);
    }

    int status = Py_ReprEnter((PyObject *)self);
    if (status != 0) {
        Py_DECREF(fields);
        if (status < 0) {
            return NULL;
        }
        return PyUnicode_FromFormat("%s(...)", Py_TYPE(self)->tp_name);
    }

    const char* tp_name = Py_TYPE(self)->tp_name;
    PyUnicodeWriter *writer = PyUnicodeWriter_Create(0);
    if (writer == NULL) {
        goto error;
    }

    if (PyUnicodeWriter_WriteUTF8(writer, tp_name, -1) < 0) {
        goto error;
    }
    if (PyUnicodeWriter_WriteChar(writer, '(') < 0) {
        goto error;
    }

    for (Py_ssize_t i = 0; i < numfields; i++) {
        PyObject *name = PySequence_GetItem(fields, i);
        if (!name) {
            goto error;
        }

        PyObject *value = PyObject_GetAttr((PyObject *)self, name);
        if (!value) {
            Py_DECREF(name);
            goto error;
        }

        PyObject *value_repr;
        if (PyList_Check(value) || PyTuple_Check(value)) {
            value_repr = ast_repr_list(value, depth);
        }
        else if (PyType_IsSubtype(Py_TYPE(value), (PyTypeObject *)state->AST_type)) {
            value_repr = ast_repr_max_depth((AST_object*)value, depth - 1);
        }
        else {
            value_repr = PyObject_Repr(value);
        }

        Py_DECREF(value);

        if (!value_repr) {
            Py_DECREF(name);
            goto error;
        }

        if (i > 0) {
            if (PyUnicodeWriter_WriteASCII(writer, ", ", 2) < 0) {
                Py_DECREF(name);
                Py_DECREF(value_repr);
                goto error;
            }
        }
        if (PyUnicodeWriter_WriteStr(writer, name) < 0) {
            Py_DECREF(name);
            Py_DECREF(value_repr);
            goto error;
        }

        Py_DECREF(name);

        if (PyUnicodeWriter_WriteChar(writer, '=') < 0) {
            Py_DECREF(value_repr);
            goto error;
        }
        if (PyUnicodeWriter_WriteStr(writer, value_repr) < 0) {
            Py_DECREF(value_repr);
            goto error;
        }

        Py_DECREF(value_repr);
    }

    if (PyUnicodeWriter_WriteChar(writer, ')') < 0) {
        goto error;
    }
    Py_ReprLeave((PyObject *)self);
    Py_DECREF(fields);
    return PyUnicodeWriter_Finish(writer);

error:
    Py_ReprLeave((PyObject *)self);
    Py_DECREF(fields);
    PyUnicodeWriter_Discard(writer);
    return NULL;
}

static PyObject *
ast_repr(PyObject *self)
{
    return ast_repr_max_depth((AST_object*)self, 3);
}

static PyType_Slot AST_type_slots[] = {
    {Py_tp_dealloc, ast_dealloc},
    {Py_tp_repr, ast_repr},
    {Py_tp_getattro, PyObject_GenericGetAttr},
    {Py_tp_setattro, PyObject_GenericSetAttr},
    {Py_tp_traverse, ast_traverse},
    {Py_tp_clear, ast_clear},
    {Py_tp_methods, ast_type_methods},
    {Py_tp_getset, ast_type_getsets},
    {Py_tp_init, ast_type_init},
    {Py_tp_alloc, PyType_GenericAlloc},
    {Py_tp_new, PyType_GenericNew},
    {Py_tp_free, PyObject_GC_Del},
    {0, 0},
};

static PyType_Spec AST_type_spec = {
    "ast.AST",
    sizeof(AST_object),
    0,
    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC |
    Py_TPFLAGS_MANAGED_DICT,
    AST_type_slots
};

static PyObject *
make_type(struct ast_state *state, const char *type, PyObject* base,
          const char* const* fields, int num_fields, const char *doc)
{
    PyObject *fnames, *result;
    int i;
    fnames = PyTuple_New(num_fields);
    if (!fnames) return NULL;
    for (i = 0; i < num_fields; i++) {
        PyObject *field = PyUnicode_InternFromString(fields[i]);
        if (!field) {
            Py_DECREF(fnames);
            return NULL;
        }
        PyTuple_SET_ITEM(fnames, i, field);
    }
    result = PyObject_CallFunction((PyObject*)&PyType_Type, "s(O){OOOOOOOs}",
                    type, base,
                    state->_fields, fnames,
                    state->__match_args__, fnames,
                    state->__module__,
                    state->ast,
                    state->__doc__, doc);
    Py_DECREF(fnames);
    return result;
}

static int
add_attributes(struct ast_state *state, PyObject *type, const char * const *attrs, int num_fields)
{
    int i, result;
    PyObject *s, *l = PyTuple_New(num_fields);
    if (!l)
        return -1;
    for (i = 0; i < num_fields; i++) {
        s = PyUnicode_InternFromString(attrs[i]);
        if (!s) {
            Py_DECREF(l);
            return -1;
        }
        PyTuple_SET_ITEM(l, i, s);
    }
    result = PyObject_SetAttr(type, state->_attributes, l);
    Py_DECREF(l);
    return result;
}

/* Conversion AST -> Python */

static PyObject* ast2obj_list(struct ast_state *state, asdl_seq *seq,
                              PyObject* (*func)(struct ast_state *state, void*))
{
    Py_ssize_t i, n = asdl_seq_LEN(seq);
    PyObject *result = PyList_New(n);
    PyObject *value;
    if (!result)
        return NULL;
    for (i = 0; i < n; i++) {
        value = func(state, asdl_seq_GET_UNTYPED(seq, i));
        if (!value) {
            Py_DECREF(result);
            return NULL;
        }
        PyList_SET_ITEM(result, i, value);
    }
    return result;
}

static PyObject* ast2obj_object(struct ast_state *Py_UNUSED(state), void *o)
{
    PyObject *op = (PyObject*)o;
    if (!op) {
        op = Py_None;
    }
    return Py_NewRef(op);
}
#define ast2obj_constant ast2obj_object
#define ast2obj_identifier ast2obj_object
#define ast2obj_string ast2obj_object

static PyObject* ast2obj_int(struct ast_state *Py_UNUSED(state), long b)
{
    return PyLong_FromLong(b);
}

/* Conversion Python -> AST */

static int obj2ast_object(struct ast_state *Py_UNUSED(state), PyObject* obj,
                          PyObject** out,
                          const char* Py_UNUSED(field), PyArena* arena)
{
    if (obj == Py_None)
        obj = NULL;
    if (obj) {
        if (_PyArena_AddPyObject(arena, obj) < 0) {
            *out = NULL;
            return -1;
        }
        *out = Py_NewRef(obj);
    }
    else {
        *out = NULL;
    }
    return 0;
}

static int obj2ast_constant(struct ast_state *Py_UNUSED(state), PyObject* obj,
                            PyObject** out,
                            const char* Py_UNUSED(field), PyArena* arena)
{
    if (_PyArena_AddPyObject(arena, obj) < 0) {
        *out = NULL;
        return -1;
    }
    *out = Py_NewRef(obj);
    return 0;
}

static int obj2ast_identifier(struct ast_state *state, PyObject* obj, PyObject** out, const char* field, PyArena* arena)
{
    if (!PyUnicode_CheckExact(obj) && obj != Py_None) {
        PyErr_Format(PyExc_TypeError, "field '%s' was expecting a string object, got %T", field, obj);
        return -1;
    }
    return obj2ast_object(state, obj, out, field, arena);
}

static int obj2ast_string(struct ast_state *state, PyObject* obj, PyObject** out, const char* field, PyArena* arena)
{
    if (!PyUnicode_CheckExact(obj) && !PyBytes_CheckExact(obj)) {
        PyErr_Format(PyExc_TypeError, "field '%s' was expecting a string or bytes object, got %T", field, obj);
        return -1;
    }
    return obj2ast_object(state, obj, out, field, arena);
}

static int obj2ast_int(struct ast_state* Py_UNUSED(state), PyObject* obj, int* out, const char* field, PyArena* arena)
{
    int i;
    if (!PyLong_Check(obj)) {
        PyErr_Format(PyExc_ValueError, "field \\"%s\\" got an invalid integer value: %R", field, obj);
        return -1;
    }

    i = PyLong_AsInt(obj);
    if (i == -1 && PyErr_Occurred())
        return -1;
    *out = i;
    return 0;
}

static int add_ast_fields(struct ast_state *state)
{
    PyObject *empty_tuple;
    empty_tuple = PyTuple_New(0);
    if (!empty_tuple ||
        PyObject_SetAttrString(state->AST_type, "_fields", empty_tuple) < 0 ||
        PyObject_SetAttrString(state->AST_type, "__match_args__", empty_tuple) < 0 ||
        PyObject_SetAttrString(state->AST_type, "_attributes", empty_tuple) < 0) {
        Py_XDECREF(empty_tuple);
        return -1;
    }
    Py_DECREF(empty_tuple);
    return 0;
}

""", 0, reflow=False)

        self.file.write(textwrap.dedent('''
            static int
            init_types(void *arg)
            {
                struct ast_state *state = arg;
                if (init_identifiers(state) < 0) {
                    return -1;
                }
                state->AST_type = PyType_FromSpec(&AST_type_spec);
                if (!state->AST_type) {
                    return -1;
                }
                state->abstract_types = PySet_New(NULL);
                if (!state->abstract_types) {
                    return -1;
                }
                if (PySet_Add(state->abstract_types, state->AST_type) < 0) {
                    return -1;
                }
                if (add_ast_fields(state) < 0) {
                    return -1;
                }
        '''))
        for dfn in mod.dfns:
            self.visit(dfn)
        self.file.write(textwrap.dedent('''
                if (!add_ast_annotations(state)) {
                    return -1;
                }
                return 0;
            }
        '''))

    def visitProduct(self, prod, name):
        if prod.fields:
            fields = name+"_fields"
        else:
            fields = "NULL"
        self.emit('state->%s_type = make_type(state, "%s", state->AST_type, %s, %d,' %
                        (name, name, fields, len(prod.fields)), 1)
        self.emit('%s);' % reflow_c_string(asdl_of(name, prod), 2), 2, reflow=False)
        self.emit("if (!state->%s_type) return -1;" % name, 1)
        if prod.attributes:
            self.emit("if (add_attributes(state, state->%s_type, %s_attributes, %d) < 0) return -1;" %
                            (name, name, len(prod.attributes)), 1)
        else:
            self.emit("if (add_attributes(state, state->%s_type, NULL, 0) < 0) return -1;" % name, 1)
        self.emit_defaults(name, prod.fields, 1)
        self.emit_defaults(name, prod.attributes, 1)

    def visitSum(self, sum, name):
        self.emit('state->%s_type = make_type(state, "%s", state->AST_type, NULL, 0,' %
                  (name, name), 1)
        self.emit('%s);' % reflow_c_string(asdl_of(name, sum), 2), 2, reflow=False)
        self.emit("if (!state->%s_type) return -1;" % name, 1)
        if sum.attributes:
            self.emit("if (add_attributes(state, state->%s_type, %s_attributes, %d) < 0) return -1;" %
                            (name, name, len(sum.attributes)), 1)
        else:
            self.emit("if (add_attributes(state, state->%s_type, NULL, 0) < 0) return -1;" % name, 1)
        self.emit("if (PySet_Add(state->abstract_types, state->%s_type) < 0) return -1;" % name, 1)
        self.emit_defaults(name, sum.attributes, 1)
        simple = is_simple(sum)
        for t in sum.types:
            self.visitConstructor(t, name, simple)

    def visitConstructor(self, cons, name, simple):
        if cons.fields:
            fields = cons.name+"_fields"
        else:
            fields = "NULL"
        self.emit('state->%s_type = make_type(state, "%s", state->%s_type, %s, %d,' %
                            (cons.name, cons.name, name, fields, len(cons.fields)), 1)
        self.emit('%s);' % reflow_c_string(asdl_of(cons.name, cons), 2), 2, reflow=False)
        self.emit("if (!state->%s_type) return -1;" % cons.name, 1)
        self.emit_defaults(cons.name, cons.fields, 1)
        if simple:
            self.emit("state->%s_singleton = PyType_GenericNew((PyTypeObject *)"
                      "state->%s_type, NULL, NULL);" %
                             (cons.name, cons.name), 1)
            self.emit("if (!state->%s_singleton) return -1;" % cons.name, 1)

    def emit_defaults(self, name, fields, depth):
        for field in fields:
            if field.opt:
                self.emit('if (PyObject_SetAttr(state->%s_type, state->%s, Py_None) == -1)' %
                            (name, field.name), depth)
                self.emit("return -1;", depth+1)


class ASTModuleVisitor(PickleVisitor):

    def visitModule(self, mod):
        self.emit("""
/* Helper for checking if a node class is abstract in the tests. */
static PyObject *
ast_is_abstract(PyObject *Py_UNUSED(module), PyObject *cls) {
    struct ast_state *state = get_ast_state();
    if (state == NULL) {
        return NULL;
    }
    int contains = PySet_Contains(state->abstract_types, cls);
    if (contains == -1) {
        return NULL;
    }
    else if (contains == 1) {
        Py_RETURN_TRUE;
    }
    Py_RETURN_FALSE;
}

static struct PyMethodDef astmodule_methods[] = {
    {"_is_abstract", ast_is_abstract, METH_O, NULL},
    {NULL}  /* Sentinel */
};
""".strip(), 0, reflow=False)
        self.emit("", 0)
        self.emit("static int", 0)
        self.emit("astmodule_exec(PyObject *m)", 0)
        self.emit("{", 0)
        self.emit('struct ast_state *state = get_ast_state();', 1)
        self.emit('if (state == NULL) {', 1)
        self.emit('return -1;', 2)
        self.emit('}', 1)
        self.emit('if (PyModule_AddObjectRef(m, "AST", state->AST_type) < 0) {', 1)
        self.emit('return -1;', 2)
        self.emit('}', 1)
        self.emit('if (PyModule_AddIntMacro(m, PyCF_ALLOW_TOP_LEVEL_AWAIT) < 0) {', 1)
        self.emit("return -1;", 2)
        self.emit('}', 1)
        self.emit('if (PyModule_AddIntMacro(m, PyCF_ONLY_AST) < 0) {', 1)
        self.emit("return -1;", 2)
        self.emit('}', 1)
        self.emit('if (PyModule_AddIntMacro(m, PyCF_TYPE_COMMENTS) < 0) {', 1)
        self.emit("return -1;", 2)
        self.emit('}', 1)
        self.emit('if (PyModule_AddIntMacro(m, PyCF_OPTIMIZED_AST) < 0) {', 1)
        self.emit("return -1;", 2)
        self.emit('}', 1)
        for dfn in mod.dfns:
            self.visit(dfn)
        self.emit("return 0;", 1)
        self.emit("}", 0)
        self.emit("", 0)
        self.emit("""
static PyModuleDef_Slot astmodule_slots[] = {
    _Py_ABI_SLOT,
    {Py_mod_exec, astmodule_exec},
    {Py_mod_multiple_interpreters, Py_MOD_PER_INTERPRETER_GIL_SUPPORTED},
    {Py_mod_gil, Py_MOD_GIL_NOT_USED},
    {0, NULL}
};

static struct PyModuleDef _astmodule = {
    PyModuleDef_HEAD_INIT,
    .m_name = "_ast",
    // The _ast module uses a per-interpreter state (PyInterpreterState.ast)
    .m_size = 0,
    .m_methods = astmodule_methods,
    .m_slots = astmodule_slots
};

PyMODINIT_FUNC
PyInit__ast(void)
{
    return PyModuleDef_Init(&_astmodule);
}
""".strip(), 0, reflow=False)

    def visitProduct(self, prod, name):
        self.addObj(name)

    def visitSum(self, sum, name):
        self.addObj(name)
        for t in sum.types:
            self.visitConstructor(t, name)

    def visitConstructor(self, cons, name):
        self.addObj(cons.name)

    def addObj(self, name):
        self.emit("if (PyModule_AddObjectRef(m, \"%s\", "
                  "state->%s_type) < 0) {" % (name, name), 1)
        self.emit("return -1;", 2)
        self.emit('}', 1)


class StaticVisitor(PickleVisitor):
    CODE = '''Very simple, always emit this static code.  Override CODE'''

    def visit(self, object):
        self.emit(self.CODE, 0, reflow=False)


class ObjVisitor(PickleVisitor):

    def func_begin(self, name):
        ctype = get_c_type(name)
        self.emit("PyObject*", 0)
        self.emit("ast2obj_%s(struct ast_state *state, void* _o)" % (name), 0)
        self.emit("{", 0)
        self.emit("%s o = (%s)_o;" % (ctype, ctype), 1)
        self.emit("PyObject *result = NULL, *value = NULL;", 1)
        self.emit("PyTypeObject *tp;", 1)
        self.emit('if (!o) {', 1)
        self.emit("Py_RETURN_NONE;", 2)
        self.emit("}", 1)
        self.emit('if (Py_EnterRecursiveCall("during  ast construction")) {', 1)
        self.emit("return NULL;", 2)
        self.emit("}", 1)

    def func_end(self):
        self.emit("Py_LeaveRecursiveCall();", 1)
        self.emit("return result;", 1)
        self.emit("failed:", 0)
        self.emit("Py_LeaveRecursiveCall();", 1)
        self.emit("Py_XDECREF(value);", 1)
        self.emit("Py_XDECREF(result);", 1)
        self.emit("return NULL;", 1)
        self.emit("}", 0)
        self.emit("", 0)

    def visitSum(self, sum, name):
        if is_simple(sum):
            self.simpleSum(sum, name)
            return
        self.func_begin(name)
        self.emit("switch (o->kind) {", 1)
        for i in range(len(sum.types)):
            t = sum.types[i]
            self.visitConstructor(t, i + 1, name)
        self.emit("}", 1)
        for a in sum.attributes:
            self.emit("value = ast2obj_%s(state, o->%s);" % (a.type, a.name), 1)
            self.emit("if (!value) goto failed;", 1)
            self.emit('if (PyObject_SetAttr(result, state->%s, value) < 0)' % a.name, 1)
            self.emit('goto failed;', 2)
            self.emit('Py_DECREF(value);', 1)
        self.func_end()

    def simpleSum(self, sum, name):
        self.emit("PyObject* ast2obj_%s(struct ast_state *state, %s_ty o)" % (name, name), 0)
        self.emit("{", 0)
        self.emit("switch(o) {", 1)
        for t in sum.types:
            self.emit("case %s:" % t.name, 2)
            self.emit("return Py_NewRef(state->%s_singleton);" % t.name, 3)
        self.emit("}", 1)
        self.emit("Py_UNREACHABLE();", 1);
        self.emit("}", 0)

    def visitProduct(self, prod, name):
        self.func_begin(name)
        self.emit("tp = (PyTypeObject *)state->%s_type;" % name, 1)
        self.emit("result = PyType_GenericNew(tp, NULL, NULL);", 1);
        self.emit("if (!result) return NULL;", 1)
        for field in prod.fields:
            self.visitField(field, name, 1, True)
        for a in prod.attributes:
            self.emit("value = ast2obj_%s(state, o->%s);" % (a.type, a.name), 1)
            self.emit("if (!value) goto failed;", 1)
            self.emit("if (PyObject_SetAttr(result, state->%s, value) < 0)" % a.name, 1)
            self.emit('goto failed;', 2)
            self.emit('Py_DECREF(value);', 1)
        self.func_end()

    def visitConstructor(self, cons, enum, name):
        self.emit("case %s_kind:" % cons.name, 1)
        self.emit("tp = (PyTypeObject *)state->%s_type;" % cons.name, 2)
        self.emit("result = PyType_GenericNew(tp, NULL, NULL);", 2);
        self.emit("if (!result) goto failed;", 2)
        for f in cons.fields:
            self.visitField(f, cons.name, 2, False)
        self.emit("break;", 2)

    def visitField(self, field, name, depth, product):
        def emit(s, d):
            self.emit(s, depth + d)
        if product:
            value = "o->%s" % field.name
        else:
            value = "o->v.%s.%s" % (name, field.name)
        self.set(field, value, depth)
        emit("if (!value) goto failed;", 0)
        emit("if (PyObject_SetAttr(result, state->%s, value) == -1)" % field.name, 0)
        emit("goto failed;", 1)
        emit("Py_DECREF(value);", 0)

    def set(self, field, value, depth):
        if field.seq:
            if field.type in self.metadata.simple_sums:
                # While the sequence elements are stored as void*,
                # simple sums expects an enum
                self.emit("{", depth)
                self.emit("Py_ssize_t i, n = asdl_seq_LEN(%s);" % value, depth+1)
                self.emit("value = PyList_New(n);", depth+1)
                self.emit("if (!value) goto failed;", depth+1)
                self.emit("for(i = 0; i < n; i++)", depth+1)
                # This cannot fail, so no need for error handling
                self.emit(
                    "PyList_SET_ITEM(value, i, ast2obj_{0}(state, ({0}_ty)asdl_seq_GET({1}, i)));".format(
                        field.type,
                        value
                    ),
                    depth + 2,
                    reflow=False,
                )
                self.emit("}", depth)
            else:
                self.emit("value = ast2obj_list(state, (asdl_seq*)%s, ast2obj_%s);" % (value, field.type), depth)
        else:
            self.emit("value = ast2obj_%s(state, %s);" % (field.type, value), depth, reflow=False)


class PartingShots(StaticVisitor):

    CODE = """
PyObject* PyAST_mod2obj(mod_ty t)
{
    struct ast_state *state = get_ast_state();
    if (state == NULL) {
        return NULL;
    }
    PyObject *result = ast2obj_mod(state, t);

    return result;
}

/* mode is 0 for "exec", 1 for "eval", 2 for "single" and 3 for "func_type"
   input */
int PyAst_CheckMode(PyObject *ast, int mode)
{
    const char * const req_name[] = {"Module", "Expression", "Interactive",
                                     "FunctionType"};

    struct ast_state *state = get_ast_state();
    if (state == NULL) {
        return -1;
    }

    PyObject *req_type[4];
    req_type[0] = state->Module_type;
    req_type[1] = state->Expression_type;
    req_type[2] = state->Interactive_type;
    req_type[3] = state->FunctionType_type;

    assert(0 <= mode && mode <= 3);
    int isinstance = PyObject_IsInstance(ast, req_type[mode]);
    if (isinstance == -1) {
        return -1;
    }
    if (!isinstance) {
        PyErr_Format(PyExc_TypeError, "expected %s node, got %T",
                     req_name[mode], ast);
        return -1;
    }
    return 0;
}

mod_ty PyAST_obj2mod(PyObject* ast, PyArena* arena, int mode)
{
    if (PySys_Audit("compile", "OO", ast, Py_None) < 0) {
        return NULL;
    }

    struct ast_state *state = get_ast_state();
    if (state == NULL) {
        return NULL;
    }

    if (PyAst_CheckMode(ast, mode) < 0) {
        return NULL;
    }

    mod_ty res = NULL;
    if (obj2ast_mod(state, ast, &res, NULL, arena) != 0)
        return NULL;
    else
        return res;
}

int PyAST_Check(PyObject* obj)
{
    struct ast_state *state = get_ast_state();
    if (state == NULL) {
        return -1;
    }
    return PyObject_IsInstance(obj, state->AST_type);
}
"""

class ChainOfVisitors:
    def __init__(self, *visitors, metadata = None):
        self.visitors = visitors
        self.metadata = metadata

    def visit(self, object):
        for v in self.visitors:
            v.metadata = self.metadata
            v.visit(object)
            v.emit("", 0)


def generate_ast_state(module_state, f):
    f.write('struct ast_state {\n')
    f.write('    _PyOnceFlag once;\n')
    f.write('    int finalized;\n')
    for s in module_state:
        f.write('    PyObject *' + s + ';\n')
    f.write('};')


def generate_ast_fini(module_state, f):
    f.write(textwrap.dedent("""
            void _PyAST_Fini(PyInterpreterState *interp)
            {
                struct ast_state *state = &interp->ast;

    """))
    f.write("    static const size_t offsets[] = {\n")
    for s in module_state:
        f.write("        offsetof(struct ast_state, " + s + "),\n")
    f.write("    };\n")
    f.write(textwrap.dedent("""
                for (size_t i = 0; i < Py_ARRAY_LENGTH(offsets); i++) {
                    PyObject **field = (PyObject **)((char *)state + offsets[i]);
                    Py_CLEAR(*field);
                }
                state->finalized = 1;
                state->once = (_PyOnceFlag){0};
            }

    """))


def generate_module_def(mod, metadata, f, internal_h):
    # Gather all the data needed for ModuleSpec
    state_strings = {
        "ast",
        "_fields",
        "__match_args__",
        "__doc__",
        "__dict__",
        "__module__",
        "_attributes",
        *metadata.identifiers
    }

    module_state = state_strings.copy()
    module_state.update(
        "%s_singleton" % singleton
        for singleton in metadata.singletons
    )
    module_state.update(
        "%s_type" % type
        for type in metadata.types
    )
    module_state.add("abstract_types")

    state_strings = sorted(state_strings)
    module_state = sorted(module_state)

    generate_ast_state(module_state, internal_h)

    print(textwrap.dedent("""
        #include "Python.h"
        #include "pycore_ast.h"
        #include "pycore_ast_state.h"     // struct ast_state
        #include "pycore_ceval.h"         // _Py_EnterRecursiveCall()
        #include "pycore_lock.h"          // _PyOnceFlag
        #include "pycore_modsupport.h"    // _PyArg_NoPositional()
        #include "pycore_pystate.h"       // _PyInterpreterState_GET()
        #include "pycore_runtime.h"       // _Py_ID()
        #include "pycore_setobject.h"     // _PySet_NextEntry()
        #include "pycore_unionobject.h"   // _Py_union_type_or

        #include <stddef.h>               // offsetof()


        // Forward declaration
        static int init_types(void *arg);

        static struct ast_state*
        get_ast_state(void)
        {
            PyInterpreterState *interp = _PyInterpreterState_GET();
            struct ast_state *state = &interp->ast;
            assert(!state->finalized);
            if (_PyOnceFlag_CallOnce(&state->once, (_Py_once_fn_t *)&init_types, state) < 0) {
                return NULL;
            }
            return state;
        }
    """).strip(), file=f)

    generate_ast_fini(module_state, f)

    f.write('static int init_identifiers(struct ast_state *state)\n')
    f.write('{\n')
    for identifier in state_strings:
        f.write('    if ((state->' + identifier)
        f.write(' = PyUnicode_InternFromString("')
        f.write(identifier + '")) == NULL) return -1;\n')
    f.write('    return 0;\n')
    f.write('};\n\n')

def write_header(mod, metadata, f):
    f.write(textwrap.dedent("""
        #ifndef Py_INTERNAL_AST_H
        #define Py_INTERNAL_AST_H
        #ifdef __cplusplus
        extern "C" {
        #endif

        #ifndef Py_BUILD_CORE
        #  error "this header requires Py_BUILD_CORE define"
        #endif

        #include "pycore_asdl.h"          // _ASDL_SEQ_HEAD

    """).lstrip())

    c = ChainOfVisitors(
        TypeDefVisitor(f),
        SequenceDefVisitor(f),
        StructVisitor(f),
        metadata=metadata
    )
    c.visit(mod)

    f.write("// Note: these macros affect function definitions, not only call sites.\n")
    prototype_visitor = PrototypeVisitor(f, metadata=metadata)
    prototype_visitor.visit(mod)

    f.write(textwrap.dedent("""

        PyObject* PyAST_mod2obj(mod_ty t);
        int PyAst_CheckMode(PyObject *ast, int mode);
        mod_ty PyAST_obj2mod(PyObject* ast, PyArena* arena, int mode);
        int PyAST_Check(PyObject* obj);

        extern int _PyAST_Validate(mod_ty);

        /* _PyAST_ExprAsUnicode is defined in ast_unparse.c */
        extern PyObject* _PyAST_ExprAsUnicode(expr_ty);

        /* Return the borrowed reference to the first literal string in the
           sequence of statements or NULL if it doesn't start from a literal string.
           Doesn't set exception. */
        extern PyObject* _PyAST_GetDocString(asdl_stmt_seq *);

        #ifdef __cplusplus
        }
        #endif
        #endif /* !Py_INTERNAL_AST_H */
    """))


def write_internal_h_header(mod, f):
    print(textwrap.dedent("""
        #ifndef Py_INTERNAL_AST_STATE_H
        #define Py_INTERNAL_AST_STATE_H

        #include "pycore_lock.h"          // _PyOnceFlag

        #ifdef __cplusplus
        extern "C" {
        #endif

        #ifndef Py_BUILD_CORE
        #  error "this header requires Py_BUILD_CORE define"
        #endif
    """).lstrip(), file=f)


def write_internal_h_footer(mod, f):
    print(textwrap.dedent("""

        #ifdef __cplusplus
        }
        #endif
        #endif /* !Py_INTERNAL_AST_STATE_H */
    """), file=f)

def write_source(mod, metadata, f, internal_h_file):
    generate_module_def(mod, metadata, f, internal_h_file)

    v = ChainOfVisitors(
        SequenceConstructorVisitor(f),
        PyTypesDeclareVisitor(f),
        AnnotationsVisitor(f),
        PyTypesVisitor(f),
        Obj2ModPrototypeVisitor(f),
        FunctionVisitor(f),
        ObjVisitor(f),
        Obj2ModVisitor(f),
        ASTModuleVisitor(f),
        PartingShots(f),
        metadata=metadata
    )
    v.visit(mod)

def main(input_filename, c_filename, h_filename, internal_h_filename, dump_module=False):
    auto_gen_msg = AUTOGEN_MESSAGE.format("/".join(Path(__file__).parts[-2:]))
    mod = asdl.parse(input_filename)
    if dump_module:
        print('Parsed Module:')
        print(mod)
    if not asdl.check(mod):
        sys.exit(1)

    metadata_visitor = MetadataVisitor()
    metadata_visitor.visit(mod)
    metadata = metadata_visitor.metadata

    with c_filename.open("w") as c_file, \
         h_filename.open("w") as h_file, \
         internal_h_filename.open("w") as internal_h_file:
        c_file.write(auto_gen_msg)
        h_file.write(auto_gen_msg)
        internal_h_file.write(auto_gen_msg)

        write_internal_h_header(mod, internal_h_file)
        write_source(mod, metadata, c_file, internal_h_file)
        write_header(mod, metadata, h_file)
        write_internal_h_footer(mod, internal_h_file)

    print(f"{c_filename}, {h_filename}, {internal_h_filename} regenerated.")

if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("input_file", type=Path)
    parser.add_argument("-C", "--c-file", type=Path, required=True)
    parser.add_argument("-H", "--h-file", type=Path, required=True)
    parser.add_argument("-I", "--internal-h-file", type=Path, required=True)
    parser.add_argument("-d", "--dump-module", action="store_true")

    args = parser.parse_args()
    main(args.input_file, args.c_file, args.h_file,
         args.internal_h_file, args.dump_module)