ydb/contrib/tools/python3/Parser/action_helpers.c

#include <Python.h>

#include "pegen.h"
#include "pycore_runtime.h" // _PyRuntime
#include "string_parser.h"
#include "tokenizer.h"
#include "pycore_pystate.h"       // _PyInterpreterState_GET()

void *_PyPegen_dummy_name(Parser *p, ...) {
  return &_PyRuntime.parser.dummy_name;
}

/* Creates a single-element asdl_seq* that contains a */
asdl_seq *_PyPegen_singleton_seq(Parser *p, void *a) {
  assert(a != NULL);
  asdl_seq *seq = (asdl_seq *)_Py_asdl_generic_seq_new(1, p->arena);
  if (!seq) {
    return NULL;
  }
  asdl_seq_SET_UNTYPED(seq, 0, a);
  return seq;
}

/* Creates a copy of seq and prepends a to it */
asdl_seq *_PyPegen_seq_insert_in_front(Parser *p, void *a, asdl_seq *seq) {
  assert(a != NULL);
  if (!seq) {
    return _PyPegen_singleton_seq(p, a);
  }

  asdl_seq *new_seq =
      (asdl_seq *)_Py_asdl_generic_seq_new(asdl_seq_LEN(seq) + 1, p->arena);
  if (!new_seq) {
    return NULL;
  }

  asdl_seq_SET_UNTYPED(new_seq, 0, a);
  for (Py_ssize_t i = 1, l = asdl_seq_LEN(new_seq); i < l; i++) {
    asdl_seq_SET_UNTYPED(new_seq, i, asdl_seq_GET_UNTYPED(seq, i - 1));
  }
  return new_seq;
}

/* Creates a copy of seq and appends a to it */
asdl_seq *_PyPegen_seq_append_to_end(Parser *p, asdl_seq *seq, void *a) {
  assert(a != NULL);
  if (!seq) {
    return _PyPegen_singleton_seq(p, a);
  }

  asdl_seq *new_seq =
      (asdl_seq *)_Py_asdl_generic_seq_new(asdl_seq_LEN(seq) + 1, p->arena);
  if (!new_seq) {
    return NULL;
  }

  for (Py_ssize_t i = 0, l = asdl_seq_LEN(new_seq); i + 1 < l; i++) {
    asdl_seq_SET_UNTYPED(new_seq, i, asdl_seq_GET_UNTYPED(seq, i));
  }
  asdl_seq_SET_UNTYPED(new_seq, asdl_seq_LEN(new_seq) - 1, a);
  return new_seq;
}

static Py_ssize_t _get_flattened_seq_size(asdl_seq *seqs) {
  Py_ssize_t size = 0;
  for (Py_ssize_t i = 0, l = asdl_seq_LEN(seqs); i < l; i++) {
    asdl_seq *inner_seq = asdl_seq_GET_UNTYPED(seqs, i);
    size += asdl_seq_LEN(inner_seq);
  }
  return size;
}

/* Flattens an asdl_seq* of asdl_seq*s */
asdl_seq *_PyPegen_seq_flatten(Parser *p, asdl_seq *seqs) {
  Py_ssize_t flattened_seq_size = _get_flattened_seq_size(seqs);
  assert(flattened_seq_size > 0);

  asdl_seq *flattened_seq =
      (asdl_seq *)_Py_asdl_generic_seq_new(flattened_seq_size, p->arena);
  if (!flattened_seq) {
    return NULL;
  }

  int flattened_seq_idx = 0;
  for (Py_ssize_t i = 0, l = asdl_seq_LEN(seqs); i < l; i++) {
    asdl_seq *inner_seq = asdl_seq_GET_UNTYPED(seqs, i);
    for (Py_ssize_t j = 0, li = asdl_seq_LEN(inner_seq); j < li; j++) {
      asdl_seq_SET_UNTYPED(flattened_seq, flattened_seq_idx++,
                           asdl_seq_GET_UNTYPED(inner_seq, j));
    }
  }
  assert(flattened_seq_idx == flattened_seq_size);

  return flattened_seq;
}

void *_PyPegen_seq_last_item(asdl_seq *seq) {
  Py_ssize_t len = asdl_seq_LEN(seq);
  return asdl_seq_GET_UNTYPED(seq, len - 1);
}

void *_PyPegen_seq_first_item(asdl_seq *seq) {
  return asdl_seq_GET_UNTYPED(seq, 0);
}

/* Creates a new name of the form <first_name>.<second_name> */
expr_ty _PyPegen_join_names_with_dot(Parser *p, expr_ty first_name,
                                     expr_ty second_name) {
  assert(first_name != NULL && second_name != NULL);
  PyObject *first_identifier = first_name->v.Name.id;
  PyObject *second_identifier = second_name->v.Name.id;

  if (PyUnicode_READY(first_identifier) == -1) {
    return NULL;
  }
  if (PyUnicode_READY(second_identifier) == -1) {
    return NULL;
  }
  const char *first_str = PyUnicode_AsUTF8(first_identifier);
  if (!first_str) {
    return NULL;
  }
  const char *second_str = PyUnicode_AsUTF8(second_identifier);
  if (!second_str) {
    return NULL;
  }
  Py_ssize_t len = strlen(first_str) + strlen(second_str) + 1; // +1 for the dot

  PyObject *str = PyBytes_FromStringAndSize(NULL, len);
  if (!str) {
    return NULL;
  }

  char *s = PyBytes_AS_STRING(str);
  if (!s) {
    return NULL;
  }

  strcpy(s, first_str);
  s += strlen(first_str);
  *s++ = '.';
  strcpy(s, second_str);
  s += strlen(second_str);
  *s = '\0';

  PyObject *uni =
      PyUnicode_DecodeUTF8(PyBytes_AS_STRING(str), PyBytes_GET_SIZE(str), NULL);
  Py_DECREF(str);
  if (!uni) {
    return NULL;
  }
  PyInterpreterState *interp = _PyInterpreterState_GET();
  _PyUnicode_InternMortal(interp, &uni);
  if (_PyArena_AddPyObject(p->arena, uni) < 0) {
    Py_DECREF(uni);
    return NULL;
  }

  return _PyAST_Name(uni, Load, EXTRA_EXPR(first_name, second_name));
}

/* Counts the total number of dots in seq's tokens */
int _PyPegen_seq_count_dots(asdl_seq *seq) {
  int number_of_dots = 0;
  for (Py_ssize_t i = 0, l = asdl_seq_LEN(seq); i < l; i++) {
    Token *current_expr = asdl_seq_GET_UNTYPED(seq, i);
    switch (current_expr->type) {
    case ELLIPSIS:
      number_of_dots += 3;
      break;
    case DOT:
      number_of_dots += 1;
      break;
    default:
      Py_UNREACHABLE();
    }
  }

  return number_of_dots;
}

/* Creates an alias with '*' as the identifier name */
alias_ty _PyPegen_alias_for_star(Parser *p, int lineno, int col_offset,
                                 int end_lineno, int end_col_offset,
                                 PyArena *arena) {
  PyObject *str = PyUnicode_InternFromString("*");
  if (!str) {
    return NULL;
  }
  if (_PyArena_AddPyObject(p->arena, str) < 0) {
    Py_DECREF(str);
    return NULL;
  }
  return _PyAST_alias(str, NULL, lineno, col_offset, end_lineno, end_col_offset,
                      arena);
}

/* Creates a new asdl_seq* with the identifiers of all the names in seq */
asdl_identifier_seq *_PyPegen_map_names_to_ids(Parser *p, asdl_expr_seq *seq) {
  Py_ssize_t len = asdl_seq_LEN(seq);
  assert(len > 0);

  asdl_identifier_seq *new_seq = _Py_asdl_identifier_seq_new(len, p->arena);
  if (!new_seq) {
    return NULL;
  }
  for (Py_ssize_t i = 0; i < len; i++) {
    expr_ty e = asdl_seq_GET(seq, i);
    asdl_seq_SET(new_seq, i, e->v.Name.id);
  }
  return new_seq;
}

/* Constructs a CmpopExprPair */
CmpopExprPair *_PyPegen_cmpop_expr_pair(Parser *p, cmpop_ty cmpop,
                                        expr_ty expr) {
  assert(expr != NULL);
  CmpopExprPair *a = _PyArena_Malloc(p->arena, sizeof(CmpopExprPair));
  if (!a) {
    return NULL;
  }
  a->cmpop = cmpop;
  a->expr = expr;
  return a;
}

asdl_int_seq *_PyPegen_get_cmpops(Parser *p, asdl_seq *seq) {
  Py_ssize_t len = asdl_seq_LEN(seq);
  assert(len > 0);

  asdl_int_seq *new_seq = _Py_asdl_int_seq_new(len, p->arena);
  if (!new_seq) {
    return NULL;
  }
  for (Py_ssize_t i = 0; i < len; i++) {
    CmpopExprPair *pair = asdl_seq_GET_UNTYPED(seq, i);
    asdl_seq_SET(new_seq, i, pair->cmpop);
  }
  return new_seq;
}

asdl_expr_seq *_PyPegen_get_exprs(Parser *p, asdl_seq *seq) {
  Py_ssize_t len = asdl_seq_LEN(seq);
  assert(len > 0);

  asdl_expr_seq *new_seq = _Py_asdl_expr_seq_new(len, p->arena);
  if (!new_seq) {
    return NULL;
  }
  for (Py_ssize_t i = 0; i < len; i++) {
    CmpopExprPair *pair = asdl_seq_GET_UNTYPED(seq, i);
    asdl_seq_SET(new_seq, i, pair->expr);
  }
  return new_seq;
}

/* Creates an asdl_seq* where all the elements have been changed to have ctx as
 * context */
static asdl_expr_seq *_set_seq_context(Parser *p, asdl_expr_seq *seq,
                                       expr_context_ty ctx) {
  Py_ssize_t len = asdl_seq_LEN(seq);
  if (len == 0) {
    return NULL;
  }

  asdl_expr_seq *new_seq = _Py_asdl_expr_seq_new(len, p->arena);
  if (!new_seq) {
    return NULL;
  }
  for (Py_ssize_t i = 0; i < len; i++) {
    expr_ty e = asdl_seq_GET(seq, i);
    asdl_seq_SET(new_seq, i, _PyPegen_set_expr_context(p, e, ctx));
  }
  return new_seq;
}

static expr_ty _set_name_context(Parser *p, expr_ty e, expr_context_ty ctx) {
  return _PyAST_Name(e->v.Name.id, ctx, EXTRA_EXPR(e, e));
}

static expr_ty _set_tuple_context(Parser *p, expr_ty e, expr_context_ty ctx) {
  return _PyAST_Tuple(_set_seq_context(p, e->v.Tuple.elts, ctx), ctx,
                      EXTRA_EXPR(e, e));
}

static expr_ty _set_list_context(Parser *p, expr_ty e, expr_context_ty ctx) {
  return _PyAST_List(_set_seq_context(p, e->v.List.elts, ctx), ctx,
                     EXTRA_EXPR(e, e));
}

static expr_ty _set_subscript_context(Parser *p, expr_ty e,
                                      expr_context_ty ctx) {
  return _PyAST_Subscript(e->v.Subscript.value, e->v.Subscript.slice, ctx,
                          EXTRA_EXPR(e, e));
}

static expr_ty _set_attribute_context(Parser *p, expr_ty e,
                                      expr_context_ty ctx) {
  return _PyAST_Attribute(e->v.Attribute.value, e->v.Attribute.attr, ctx,
                          EXTRA_EXPR(e, e));
}

static expr_ty _set_starred_context(Parser *p, expr_ty e, expr_context_ty ctx) {
  return _PyAST_Starred(_PyPegen_set_expr_context(p, e->v.Starred.value, ctx),
                        ctx, EXTRA_EXPR(e, e));
}

/* Creates an `expr_ty` equivalent to `expr` but with `ctx` as context */
expr_ty _PyPegen_set_expr_context(Parser *p, expr_ty expr,
                                  expr_context_ty ctx) {
  assert(expr != NULL);

  expr_ty new = NULL;
  switch (expr->kind) {
  case Name_kind:
    new = _set_name_context(p, expr, ctx);
    break;
  case Tuple_kind:
    new = _set_tuple_context(p, expr, ctx);
    break;
  case List_kind:
    new = _set_list_context(p, expr, ctx);
    break;
  case Subscript_kind:
    new = _set_subscript_context(p, expr, ctx);
    break;
  case Attribute_kind:
    new = _set_attribute_context(p, expr, ctx);
    break;
  case Starred_kind:
    new = _set_starred_context(p, expr, ctx);
    break;
  default:
    new = expr;
  }
  return new;
}

/* Constructs a KeyValuePair that is used when parsing a dict's key value pairs
 */
KeyValuePair *_PyPegen_key_value_pair(Parser *p, expr_ty key, expr_ty value) {
  KeyValuePair *a = _PyArena_Malloc(p->arena, sizeof(KeyValuePair));
  if (!a) {
    return NULL;
  }
  a->key = key;
  a->value = value;
  return a;
}

/* Extracts all keys from an asdl_seq* of KeyValuePair*'s */
asdl_expr_seq *_PyPegen_get_keys(Parser *p, asdl_seq *seq) {
  Py_ssize_t len = asdl_seq_LEN(seq);
  asdl_expr_seq *new_seq = _Py_asdl_expr_seq_new(len, p->arena);
  if (!new_seq) {
    return NULL;
  }
  for (Py_ssize_t i = 0; i < len; i++) {
    KeyValuePair *pair = asdl_seq_GET_UNTYPED(seq, i);
    asdl_seq_SET(new_seq, i, pair->key);
  }
  return new_seq;
}

/* Extracts all values from an asdl_seq* of KeyValuePair*'s */
asdl_expr_seq *_PyPegen_get_values(Parser *p, asdl_seq *seq) {
  Py_ssize_t len = asdl_seq_LEN(seq);
  asdl_expr_seq *new_seq = _Py_asdl_expr_seq_new(len, p->arena);
  if (!new_seq) {
    return NULL;
  }
  for (Py_ssize_t i = 0; i < len; i++) {
    KeyValuePair *pair = asdl_seq_GET_UNTYPED(seq, i);
    asdl_seq_SET(new_seq, i, pair->value);
  }
  return new_seq;
}

/* Constructs a KeyPatternPair that is used when parsing mapping & class
 * patterns */
KeyPatternPair *_PyPegen_key_pattern_pair(Parser *p, expr_ty key,
                                          pattern_ty pattern) {
  KeyPatternPair *a = _PyArena_Malloc(p->arena, sizeof(KeyPatternPair));
  if (!a) {
    return NULL;
  }
  a->key = key;
  a->pattern = pattern;
  return a;
}

/* Extracts all keys from an asdl_seq* of KeyPatternPair*'s */
asdl_expr_seq *_PyPegen_get_pattern_keys(Parser *p, asdl_seq *seq) {
  Py_ssize_t len = asdl_seq_LEN(seq);
  asdl_expr_seq *new_seq = _Py_asdl_expr_seq_new(len, p->arena);
  if (!new_seq) {
    return NULL;
  }
  for (Py_ssize_t i = 0; i < len; i++) {
    KeyPatternPair *pair = asdl_seq_GET_UNTYPED(seq, i);
    asdl_seq_SET(new_seq, i, pair->key);
  }
  return new_seq;
}

/* Extracts all patterns from an asdl_seq* of KeyPatternPair*'s */
asdl_pattern_seq *_PyPegen_get_patterns(Parser *p, asdl_seq *seq) {
  Py_ssize_t len = asdl_seq_LEN(seq);
  asdl_pattern_seq *new_seq = _Py_asdl_pattern_seq_new(len, p->arena);
  if (!new_seq) {
    return NULL;
  }
  for (Py_ssize_t i = 0; i < len; i++) {
    KeyPatternPair *pair = asdl_seq_GET_UNTYPED(seq, i);
    asdl_seq_SET(new_seq, i, pair->pattern);
  }
  return new_seq;
}

/* Constructs a NameDefaultPair */
NameDefaultPair *_PyPegen_name_default_pair(Parser *p, arg_ty arg,
                                            expr_ty value, Token *tc) {
  NameDefaultPair *a = _PyArena_Malloc(p->arena, sizeof(NameDefaultPair));
  if (!a) {
    return NULL;
  }
  a->arg = _PyPegen_add_type_comment_to_arg(p, arg, tc);
  a->value = value;
  return a;
}

/* Constructs a SlashWithDefault */
SlashWithDefault *_PyPegen_slash_with_default(Parser *p,
                                              asdl_arg_seq *plain_names,
                                              asdl_seq *names_with_defaults) {
  SlashWithDefault *a = _PyArena_Malloc(p->arena, sizeof(SlashWithDefault));
  if (!a) {
    return NULL;
  }
  a->plain_names = plain_names;
  a->names_with_defaults = names_with_defaults;
  return a;
}

/* Constructs a StarEtc */
StarEtc *_PyPegen_star_etc(Parser *p, arg_ty vararg, asdl_seq *kwonlyargs,
                           arg_ty kwarg) {
  StarEtc *a = _PyArena_Malloc(p->arena, sizeof(StarEtc));
  if (!a) {
    return NULL;
  }
  a->vararg = vararg;
  a->kwonlyargs = kwonlyargs;
  a->kwarg = kwarg;
  return a;
}

asdl_seq *_PyPegen_join_sequences(Parser *p, asdl_seq *a, asdl_seq *b) {
  Py_ssize_t first_len = asdl_seq_LEN(a);
  Py_ssize_t second_len = asdl_seq_LEN(b);
  asdl_seq *new_seq =
      (asdl_seq *)_Py_asdl_generic_seq_new(first_len + second_len, p->arena);
  if (!new_seq) {
    return NULL;
  }

  int k = 0;
  for (Py_ssize_t i = 0; i < first_len; i++) {
    asdl_seq_SET_UNTYPED(new_seq, k++, asdl_seq_GET_UNTYPED(a, i));
  }
  for (Py_ssize_t i = 0; i < second_len; i++) {
    asdl_seq_SET_UNTYPED(new_seq, k++, asdl_seq_GET_UNTYPED(b, i));
  }

  return new_seq;
}

static asdl_arg_seq *_get_names(Parser *p, asdl_seq *names_with_defaults) {
  Py_ssize_t len = asdl_seq_LEN(names_with_defaults);
  asdl_arg_seq *seq = _Py_asdl_arg_seq_new(len, p->arena);
  if (!seq) {
    return NULL;
  }
  for (Py_ssize_t i = 0; i < len; i++) {
    NameDefaultPair *pair = asdl_seq_GET_UNTYPED(names_with_defaults, i);
    asdl_seq_SET(seq, i, pair->arg);
  }
  return seq;
}

static asdl_expr_seq *_get_defaults(Parser *p, asdl_seq *names_with_defaults) {
  Py_ssize_t len = asdl_seq_LEN(names_with_defaults);
  asdl_expr_seq *seq = _Py_asdl_expr_seq_new(len, p->arena);
  if (!seq) {
    return NULL;
  }
  for (Py_ssize_t i = 0; i < len; i++) {
    NameDefaultPair *pair = asdl_seq_GET_UNTYPED(names_with_defaults, i);
    asdl_seq_SET(seq, i, pair->value);
  }
  return seq;
}

static int _make_posonlyargs(Parser *p, asdl_arg_seq *slash_without_default,
                             SlashWithDefault *slash_with_default,
                             asdl_arg_seq **posonlyargs) {
  if (slash_without_default != NULL) {
    *posonlyargs = slash_without_default;
  } else if (slash_with_default != NULL) {
    asdl_arg_seq *slash_with_default_names =
        _get_names(p, slash_with_default->names_with_defaults);
    if (!slash_with_default_names) {
      return -1;
    }
    *posonlyargs = (asdl_arg_seq *)_PyPegen_join_sequences(
        p, (asdl_seq *)slash_with_default->plain_names,
        (asdl_seq *)slash_with_default_names);
  } else {
    *posonlyargs = _Py_asdl_arg_seq_new(0, p->arena);
  }
  return *posonlyargs == NULL ? -1 : 0;
}

static int _make_posargs(Parser *p, asdl_arg_seq *plain_names,
                         asdl_seq *names_with_default, asdl_arg_seq **posargs) {
  if (plain_names != NULL && names_with_default != NULL) {
    asdl_arg_seq *names_with_default_names = _get_names(p, names_with_default);
    if (!names_with_default_names) {
      return -1;
    }
    *posargs = (asdl_arg_seq *)_PyPegen_join_sequences(
        p, (asdl_seq *)plain_names, (asdl_seq *)names_with_default_names);
  } else if (plain_names == NULL && names_with_default != NULL) {
    *posargs = _get_names(p, names_with_default);
  } else if (plain_names != NULL && names_with_default == NULL) {
    *posargs = plain_names;
  } else {
    *posargs = _Py_asdl_arg_seq_new(0, p->arena);
  }
  return *posargs == NULL ? -1 : 0;
}

static int _make_posdefaults(Parser *p, SlashWithDefault *slash_with_default,
                             asdl_seq *names_with_default,
                             asdl_expr_seq **posdefaults) {
  if (slash_with_default != NULL && names_with_default != NULL) {
    asdl_expr_seq *slash_with_default_values =
        _get_defaults(p, slash_with_default->names_with_defaults);
    if (!slash_with_default_values) {
      return -1;
    }
    asdl_expr_seq *names_with_default_values =
        _get_defaults(p, names_with_default);
    if (!names_with_default_values) {
      return -1;
    }
    *posdefaults = (asdl_expr_seq *)_PyPegen_join_sequences(
        p, (asdl_seq *)slash_with_default_values,
        (asdl_seq *)names_with_default_values);
  } else if (slash_with_default == NULL && names_with_default != NULL) {
    *posdefaults = _get_defaults(p, names_with_default);
  } else if (slash_with_default != NULL && names_with_default == NULL) {
    *posdefaults = _get_defaults(p, slash_with_default->names_with_defaults);
  } else {
    *posdefaults = _Py_asdl_expr_seq_new(0, p->arena);
  }
  return *posdefaults == NULL ? -1 : 0;
}

static int _make_kwargs(Parser *p, StarEtc *star_etc, asdl_arg_seq **kwonlyargs,
                        asdl_expr_seq **kwdefaults) {
  if (star_etc != NULL && star_etc->kwonlyargs != NULL) {
    *kwonlyargs = _get_names(p, star_etc->kwonlyargs);
  } else {
    *kwonlyargs = _Py_asdl_arg_seq_new(0, p->arena);
  }

  if (*kwonlyargs == NULL) {
    return -1;
  }

  if (star_etc != NULL && star_etc->kwonlyargs != NULL) {
    *kwdefaults = _get_defaults(p, star_etc->kwonlyargs);
  } else {
    *kwdefaults = _Py_asdl_expr_seq_new(0, p->arena);
  }

  if (*kwdefaults == NULL) {
    return -1;
  }

  return 0;
}

/* Constructs an arguments_ty object out of all the parsed constructs in the
 * parameters rule */
arguments_ty _PyPegen_make_arguments(Parser *p,
                                     asdl_arg_seq *slash_without_default,
                                     SlashWithDefault *slash_with_default,
                                     asdl_arg_seq *plain_names,
                                     asdl_seq *names_with_default,
                                     StarEtc *star_etc) {
  asdl_arg_seq *posonlyargs;
  if (_make_posonlyargs(p, slash_without_default, slash_with_default,
                        &posonlyargs) == -1) {
    return NULL;
  }

  asdl_arg_seq *posargs;
  if (_make_posargs(p, plain_names, names_with_default, &posargs) == -1) {
    return NULL;
  }

  asdl_expr_seq *posdefaults;
  if (_make_posdefaults(p, slash_with_default, names_with_default,
                        &posdefaults) == -1) {
    return NULL;
  }

  arg_ty vararg = NULL;
  if (star_etc != NULL && star_etc->vararg != NULL) {
    vararg = star_etc->vararg;
  }

  asdl_arg_seq *kwonlyargs;
  asdl_expr_seq *kwdefaults;
  if (_make_kwargs(p, star_etc, &kwonlyargs, &kwdefaults) == -1) {
    return NULL;
  }

  arg_ty kwarg = NULL;
  if (star_etc != NULL && star_etc->kwarg != NULL) {
    kwarg = star_etc->kwarg;
  }

  return _PyAST_arguments(posonlyargs, posargs, vararg, kwonlyargs, kwdefaults,
                          kwarg, posdefaults, p->arena);
}

/* Constructs an empty arguments_ty object, that gets used when a function
 * accepts no arguments. */
arguments_ty _PyPegen_empty_arguments(Parser *p) {
  asdl_arg_seq *posonlyargs = _Py_asdl_arg_seq_new(0, p->arena);
  if (!posonlyargs) {
    return NULL;
  }
  asdl_arg_seq *posargs = _Py_asdl_arg_seq_new(0, p->arena);
  if (!posargs) {
    return NULL;
  }
  asdl_expr_seq *posdefaults = _Py_asdl_expr_seq_new(0, p->arena);
  if (!posdefaults) {
    return NULL;
  }
  asdl_arg_seq *kwonlyargs = _Py_asdl_arg_seq_new(0, p->arena);
  if (!kwonlyargs) {
    return NULL;
  }
  asdl_expr_seq *kwdefaults = _Py_asdl_expr_seq_new(0, p->arena);
  if (!kwdefaults) {
    return NULL;
  }

  return _PyAST_arguments(posonlyargs, posargs, NULL, kwonlyargs, kwdefaults,
                          NULL, posdefaults, p->arena);
}

/* Encapsulates the value of an operator_ty into an AugOperator struct */
AugOperator *_PyPegen_augoperator(Parser *p, operator_ty kind) {
  AugOperator *a = _PyArena_Malloc(p->arena, sizeof(AugOperator));
  if (!a) {
    return NULL;
  }
  a->kind = kind;
  return a;
}

/* Construct a FunctionDef equivalent to function_def, but with decorators */
stmt_ty _PyPegen_function_def_decorators(Parser *p, asdl_expr_seq *decorators,
                                         stmt_ty function_def) {
  assert(function_def != NULL);
  if (function_def->kind == AsyncFunctionDef_kind) {
    return _PyAST_AsyncFunctionDef(
        function_def->v.AsyncFunctionDef.name,
        function_def->v.AsyncFunctionDef.args,
        function_def->v.AsyncFunctionDef.body, decorators,
        function_def->v.AsyncFunctionDef.returns,
        function_def->v.AsyncFunctionDef.type_comment,
        function_def->v.AsyncFunctionDef.type_params, function_def->lineno,
        function_def->col_offset, function_def->end_lineno,
        function_def->end_col_offset, p->arena);
  }

  return _PyAST_FunctionDef(
      function_def->v.FunctionDef.name, function_def->v.FunctionDef.args,
      function_def->v.FunctionDef.body, decorators,
      function_def->v.FunctionDef.returns,
      function_def->v.FunctionDef.type_comment,
      function_def->v.FunctionDef.type_params, function_def->lineno,
      function_def->col_offset, function_def->end_lineno,
      function_def->end_col_offset, p->arena);
}

/* Construct a ClassDef equivalent to class_def, but with decorators */
stmt_ty _PyPegen_class_def_decorators(Parser *p, asdl_expr_seq *decorators,
                                      stmt_ty class_def) {
  assert(class_def != NULL);
  return _PyAST_ClassDef(
      class_def->v.ClassDef.name, class_def->v.ClassDef.bases,
      class_def->v.ClassDef.keywords, class_def->v.ClassDef.body, decorators,
      class_def->v.ClassDef.type_params, class_def->lineno,
      class_def->col_offset, class_def->end_lineno, class_def->end_col_offset,
      p->arena);
}

/* Construct a KeywordOrStarred */
KeywordOrStarred *_PyPegen_keyword_or_starred(Parser *p, void *element,
                                              int is_keyword) {
  KeywordOrStarred *a = _PyArena_Malloc(p->arena, sizeof(KeywordOrStarred));
  if (!a) {
    return NULL;
  }
  a->element = element;
  a->is_keyword = is_keyword;
  return a;
}

/* Get the number of starred expressions in an asdl_seq* of KeywordOrStarred*s
 */
static int _seq_number_of_starred_exprs(asdl_seq *seq) {
  int n = 0;
  for (Py_ssize_t i = 0, l = asdl_seq_LEN(seq); i < l; i++) {
    KeywordOrStarred *k = asdl_seq_GET_UNTYPED(seq, i);
    if (!k->is_keyword) {
      n++;
    }
  }
  return n;
}

/* Extract the starred expressions of an asdl_seq* of KeywordOrStarred*s */
asdl_expr_seq *_PyPegen_seq_extract_starred_exprs(Parser *p, asdl_seq *kwargs) {
  int new_len = _seq_number_of_starred_exprs(kwargs);
  if (new_len == 0) {
    return NULL;
  }
  asdl_expr_seq *new_seq = _Py_asdl_expr_seq_new(new_len, p->arena);
  if (!new_seq) {
    return NULL;
  }

  int idx = 0;
  for (Py_ssize_t i = 0, len = asdl_seq_LEN(kwargs); i < len; i++) {
    KeywordOrStarred *k = asdl_seq_GET_UNTYPED(kwargs, i);
    if (!k->is_keyword) {
      asdl_seq_SET(new_seq, idx++, k->element);
    }
  }
  return new_seq;
}

/* Return a new asdl_seq* with only the keywords in kwargs */
asdl_keyword_seq *_PyPegen_seq_delete_starred_exprs(Parser *p,
                                                    asdl_seq *kwargs) {
  Py_ssize_t len = asdl_seq_LEN(kwargs);
  Py_ssize_t new_len = len - _seq_number_of_starred_exprs(kwargs);
  if (new_len == 0) {
    return NULL;
  }
  asdl_keyword_seq *new_seq = _Py_asdl_keyword_seq_new(new_len, p->arena);
  if (!new_seq) {
    return NULL;
  }

  int idx = 0;
  for (Py_ssize_t i = 0; i < len; i++) {
    KeywordOrStarred *k = asdl_seq_GET_UNTYPED(kwargs, i);
    if (k->is_keyword) {
      asdl_seq_SET(new_seq, idx++, k->element);
    }
  }
  return new_seq;
}

expr_ty _PyPegen_ensure_imaginary(Parser *p, expr_ty exp) {
  if (exp->kind != Constant_kind ||
      !PyComplex_CheckExact(exp->v.Constant.value)) {
    RAISE_SYNTAX_ERROR_KNOWN_LOCATION(
        exp, "imaginary number required in complex literal");
    return NULL;
  }
  return exp;
}

expr_ty _PyPegen_ensure_real(Parser *p, expr_ty exp) {
  if (exp->kind != Constant_kind ||
      PyComplex_CheckExact(exp->v.Constant.value)) {
    RAISE_SYNTAX_ERROR_KNOWN_LOCATION(
        exp, "real number required in complex literal");
    return NULL;
  }
  return exp;
}

mod_ty _PyPegen_make_module(Parser *p, asdl_stmt_seq *a) {
  asdl_type_ignore_seq *type_ignores = NULL;
  Py_ssize_t num = p->type_ignore_comments.num_items;
  if (num > 0) {
    // Turn the raw (comment, lineno) pairs into TypeIgnore objects in the arena
    type_ignores = _Py_asdl_type_ignore_seq_new(num, p->arena);
    if (type_ignores == NULL) {
      return NULL;
    }
    for (int i = 0; i < num; i++) {
      PyObject *tag = _PyPegen_new_type_comment(
          p, p->type_ignore_comments.items[i].comment);
      if (tag == NULL) {
        return NULL;
      }
      type_ignore_ty ti = _PyAST_TypeIgnore(
          p->type_ignore_comments.items[i].lineno, tag, p->arena);
      if (ti == NULL) {
        return NULL;
      }
      asdl_seq_SET(type_ignores, i, ti);
    }
  }
  return _PyAST_Module(a, type_ignores, p->arena);
}

PyObject *_PyPegen_new_type_comment(Parser *p, const char *s) {
  PyObject *res = PyUnicode_DecodeUTF8(s, strlen(s), NULL);
  if (res == NULL) {
    return NULL;
  }
  if (_PyArena_AddPyObject(p->arena, res) < 0) {
    Py_DECREF(res);
    return NULL;
  }
  return res;
}

arg_ty _PyPegen_add_type_comment_to_arg(Parser *p, arg_ty a, Token *tc) {
  if (tc == NULL) {
    return a;
  }
  const char *bytes = PyBytes_AsString(tc->bytes);
  if (bytes == NULL) {
    return NULL;
  }
  PyObject *tco = _PyPegen_new_type_comment(p, bytes);
  if (tco == NULL) {
    return NULL;
  }
  return _PyAST_arg(a->arg, a->annotation, tco, a->lineno, a->col_offset,
                    a->end_lineno, a->end_col_offset, p->arena);
}

/* Checks if the NOTEQUAL token is valid given the current parser flags
0 indicates success and nonzero indicates failure (an exception may be set) */
int _PyPegen_check_barry_as_flufl(Parser *p, Token *t) {
  assert(t->bytes != NULL);
  assert(t->type == NOTEQUAL);

  const char *tok_str = PyBytes_AS_STRING(t->bytes);
  if (p->flags & PyPARSE_BARRY_AS_BDFL && strcmp(tok_str, "<>") != 0) {
    RAISE_SYNTAX_ERROR("with Barry as BDFL, use '<>' instead of '!='");
    return -1;
  }
  if (!(p->flags & PyPARSE_BARRY_AS_BDFL)) {
    return strcmp(tok_str, "!=");
  }
  return 0;
}

int _PyPegen_check_legacy_stmt(Parser *p, expr_ty name) {
  if (name->kind != Name_kind) {
    return 0;
  }
  const char *candidates[2] = {"print", "exec"};
  for (int i = 0; i < 2; i++) {
    if (PyUnicode_CompareWithASCIIString(name->v.Name.id, candidates[i]) == 0) {
      return 1;
    }
  }
  return 0;
}

static ResultTokenWithMetadata *
result_token_with_metadata(Parser *p, void *result, PyObject *metadata) {
  ResultTokenWithMetadata *res =
      _PyArena_Malloc(p->arena, sizeof(ResultTokenWithMetadata));
  if (res == NULL) {
    return NULL;
  }
  res->metadata = metadata;
  res->result = result;
  return res;
}

ResultTokenWithMetadata *
_PyPegen_check_fstring_conversion(Parser *p, Token *conv_token, expr_ty conv) {
  if (conv_token->lineno != conv->lineno ||
      conv_token->end_col_offset != conv->col_offset) {
    return RAISE_SYNTAX_ERROR_KNOWN_RANGE(conv_token, conv,
                                          "f-string: conversion type must come "
                                          "right after the exclamanation mark");
  }
  return result_token_with_metadata(p, conv, conv_token->metadata);
}

static asdl_expr_seq *
unpack_top_level_joined_strs(Parser *p, asdl_expr_seq *raw_expressions);
ResultTokenWithMetadata *
_PyPegen_setup_full_format_spec(Parser *p, Token *colon, asdl_expr_seq *spec,
                                int lineno, int col_offset, int end_lineno,
                                int end_col_offset, PyArena *arena) {
  if (!spec) {
    return NULL;
  }

  // This is needed to keep compatibility with 3.11, where an empty format spec
  // is parsed as an *empty* JoinedStr node, instead of having an empty constant
  // in it.
  if (asdl_seq_LEN(spec) == 1) {
    expr_ty e = asdl_seq_GET(spec, 0);
    if (e->kind == Constant_kind && PyUnicode_Check(e->v.Constant.value) &&
        PyUnicode_GetLength(e->v.Constant.value) == 0) {
      spec = _Py_asdl_expr_seq_new(0, arena);
    }
  }
  expr_ty res;
  Py_ssize_t n = asdl_seq_LEN(spec);
  if (n == 0 || (n == 1 && asdl_seq_GET(spec, 0)->kind == Constant_kind)) {
    res = _PyAST_JoinedStr(spec, lineno, col_offset, end_lineno, end_col_offset,
                           p->arena);
  } else {
    res = _PyPegen_concatenate_strings(p, spec, lineno, col_offset, end_lineno,
                                       end_col_offset, arena);
  }
  if (!res) {
    return NULL;
  }
  return result_token_with_metadata(p, res, colon->metadata);
}

const char *_PyPegen_get_expr_name(expr_ty e) {
  assert(e != NULL);
  switch (e->kind) {
  case Attribute_kind:
    return "attribute";
  case Subscript_kind:
    return "subscript";
  case Starred_kind:
    return "starred";
  case Name_kind:
    return "name";
  case List_kind:
    return "list";
  case Tuple_kind:
    return "tuple";
  case Lambda_kind:
    return "lambda";
  case Call_kind:
    return "function call";
  case BoolOp_kind:
  case BinOp_kind:
  case UnaryOp_kind:
    return "expression";
  case GeneratorExp_kind:
    return "generator expression";
  case Yield_kind:
  case YieldFrom_kind:
    return "yield expression";
  case Await_kind:
    return "await expression";
  case ListComp_kind:
    return "list comprehension";
  case SetComp_kind:
    return "set comprehension";
  case DictComp_kind:
    return "dict comprehension";
  case Dict_kind:
    return "dict literal";
  case Set_kind:
    return "set display";
  case JoinedStr_kind:
  case FormattedValue_kind:
    return "f-string expression";
  case Constant_kind: {
    PyObject *value = e->v.Constant.value;
    if (value == Py_None) {
      return "None";
    }
    if (value == Py_False) {
      return "False";
    }
    if (value == Py_True) {
      return "True";
    }
    if (value == Py_Ellipsis) {
      return "ellipsis";
    }
    return "literal";
  }
  case Compare_kind:
    return "comparison";
  case IfExp_kind:
    return "conditional expression";
  case NamedExpr_kind:
    return "named expression";
  default:
    PyErr_Format(PyExc_SystemError,
                 "unexpected expression in assignment %d (line %d)", e->kind,
                 e->lineno);
    return NULL;
  }
}

expr_ty _PyPegen_get_last_comprehension_item(comprehension_ty comprehension) {
  if (comprehension->ifs == NULL || asdl_seq_LEN(comprehension->ifs) == 0) {
    return comprehension->iter;
  }
  return PyPegen_last_item(comprehension->ifs, expr_ty);
}

expr_ty _PyPegen_collect_call_seqs(Parser *p, asdl_expr_seq *a, asdl_seq *b,
                                   int lineno, int col_offset, int end_lineno,
                                   int end_col_offset, PyArena *arena) {
  Py_ssize_t args_len = asdl_seq_LEN(a);
  Py_ssize_t total_len = args_len;

  if (b == NULL) {
    return _PyAST_Call(_PyPegen_dummy_name(p), a, NULL, lineno, col_offset,
                       end_lineno, end_col_offset, arena);
  }

  asdl_expr_seq *starreds = _PyPegen_seq_extract_starred_exprs(p, b);
  asdl_keyword_seq *keywords = _PyPegen_seq_delete_starred_exprs(p, b);

  if (starreds) {
    total_len += asdl_seq_LEN(starreds);
  }

  asdl_expr_seq *args = _Py_asdl_expr_seq_new(total_len, arena);
  if (args == NULL) {
    return NULL;
  }

  Py_ssize_t i = 0;
  for (i = 0; i < args_len; i++) {
    asdl_seq_SET(args, i, asdl_seq_GET(a, i));
  }
  for (; i < total_len; i++) {
    asdl_seq_SET(args, i, asdl_seq_GET(starreds, i - args_len));
  }

  return _PyAST_Call(_PyPegen_dummy_name(p), args, keywords, lineno, col_offset,
                     end_lineno, end_col_offset, arena);
}

// AST Error reporting helpers

expr_ty _PyPegen_get_invalid_target(expr_ty e, TARGETS_TYPE targets_type) {
  if (e == NULL) {
    return NULL;
  }

#define VISIT_CONTAINER(CONTAINER, TYPE)                                       \
  do {                                                                         \
    Py_ssize_t len = asdl_seq_LEN((CONTAINER)->v.TYPE.elts);                   \
    for (Py_ssize_t i = 0; i < len; i++) {                                     \
      expr_ty other = asdl_seq_GET((CONTAINER)->v.TYPE.elts, i);               \
      expr_ty child = _PyPegen_get_invalid_target(other, targets_type);        \
      if (child != NULL) {                                                     \
        return child;                                                          \
      }                                                                        \
    }                                                                          \
  } while (0)

  // We only need to visit List and Tuple nodes recursively as those
  // are the only ones that can contain valid names in targets when
  // they are parsed as expressions. Any other kind of expression
  // that is a container (like Sets or Dicts) is directly invalid and
  // we don't need to visit it recursively.

  switch (e->kind) {
  case List_kind:
    VISIT_CONTAINER(e, List);
    return NULL;
  case Tuple_kind:
    VISIT_CONTAINER(e, Tuple);
    return NULL;
  case Starred_kind:
    if (targets_type == DEL_TARGETS) {
      return e;
    }
    return _PyPegen_get_invalid_target(e->v.Starred.value, targets_type);
  case Compare_kind:
    // This is needed, because the `a in b` in `for a in b` gets parsed
    // as a comparison, and so we need to search the left side of the comparison
    // for invalid targets.
    if (targets_type == FOR_TARGETS) {
      cmpop_ty cmpop = (cmpop_ty)asdl_seq_GET(e->v.Compare.ops, 0);
      if (cmpop == In) {
        return _PyPegen_get_invalid_target(e->v.Compare.left, targets_type);
      }
      return NULL;
    }
    return e;
  case Name_kind:
  case Subscript_kind:
  case Attribute_kind:
    return NULL;
  default:
    return e;
  }
}

void *_PyPegen_arguments_parsing_error(Parser *p, expr_ty e) {
  int kwarg_unpacking = 0;
  for (Py_ssize_t i = 0, l = asdl_seq_LEN(e->v.Call.keywords); i < l; i++) {
    keyword_ty keyword = asdl_seq_GET(e->v.Call.keywords, i);
    if (!keyword->arg) {
      kwarg_unpacking = 1;
    }
  }

  const char *msg = NULL;
  if (kwarg_unpacking) {
    msg = "positional argument follows keyword argument unpacking";
  } else {
    msg = "positional argument follows keyword argument";
  }

  return RAISE_SYNTAX_ERROR(msg);
}

void *_PyPegen_nonparen_genexp_in_call(Parser *p, expr_ty args,
                                       asdl_comprehension_seq *comprehensions) {
  /* The rule that calls this function is 'args for_if_clauses'.
     For the input f(L, x for x in y), L and x are in args and
     the for is parsed as a for_if_clause. We have to check if
     len <= 1, so that input like dict((a, b) for a, b in x)
     gets successfully parsed and then we pass the last
     argument (x in the above example) as the location of the
     error */
  Py_ssize_t len = asdl_seq_LEN(args->v.Call.args);
  if (len <= 1) {
    return NULL;
  }

  comprehension_ty last_comprehension =
      PyPegen_last_item(comprehensions, comprehension_ty);

  return RAISE_SYNTAX_ERROR_KNOWN_RANGE(
      (expr_ty)asdl_seq_GET(args->v.Call.args, len - 1),
      _PyPegen_get_last_comprehension_item(last_comprehension),
      "Generator expression must be parenthesized");
}

// Fstring stuff

static expr_ty _PyPegen_decode_fstring_part(Parser *p, int is_raw,
                                            expr_ty constant, Token *token) {
  assert(PyUnicode_CheckExact(constant->v.Constant.value));

  const char *bstr = PyUnicode_AsUTF8(constant->v.Constant.value);
  if (bstr == NULL) {
    return NULL;
  }

  size_t len;
  if (strcmp(bstr, "{{") == 0 || strcmp(bstr, "}}") == 0) {
    len = 1;
  } else {
    len = strlen(bstr);
  }

  is_raw = is_raw || strchr(bstr, '\\') == NULL;
  PyObject *str = _PyPegen_decode_string(p, is_raw, bstr, len, token);
  if (str == NULL) {
    _Pypegen_raise_decode_error(p);
    return NULL;
  }
  if (_PyArena_AddPyObject(p->arena, str) < 0) {
    Py_DECREF(str);
    return NULL;
  }
  return _PyAST_Constant(str, NULL, constant->lineno, constant->col_offset,
                         constant->end_lineno, constant->end_col_offset,
                         p->arena);
}

static asdl_expr_seq *
unpack_top_level_joined_strs(Parser *p, asdl_expr_seq *raw_expressions) {
  /* The parser might put multiple f-string values into an individual
   * JoinedStr node at the top level due to stuff like f-string debugging
   * expressions. This function flattens those and promotes them to the
   * upper level. Only simplifies AST, but the compiler already takes care
   * of the regular output, so this is not necessary if you are not going
   * to expose the output AST to Python level. */

  Py_ssize_t i, req_size, raw_size;

  req_size = raw_size = asdl_seq_LEN(raw_expressions);
  expr_ty expr;
  for (i = 0; i < raw_size; i++) {
    expr = asdl_seq_GET(raw_expressions, i);
    if (expr->kind == JoinedStr_kind) {
      req_size += asdl_seq_LEN(expr->v.JoinedStr.values) - 1;
    }
  }

  asdl_expr_seq *expressions = _Py_asdl_expr_seq_new(req_size, p->arena);
  if (expressions == NULL) {
    return NULL;
  }

  Py_ssize_t raw_index, req_index = 0;
  for (raw_index = 0; raw_index < raw_size; raw_index++) {
    expr = asdl_seq_GET(raw_expressions, raw_index);
    if (expr->kind == JoinedStr_kind) {
      asdl_expr_seq *values = expr->v.JoinedStr.values;
      for (Py_ssize_t n = 0; n < asdl_seq_LEN(values); n++) {
        asdl_seq_SET(expressions, req_index, asdl_seq_GET(values, n));
        req_index++;
      }
    } else {
      asdl_seq_SET(expressions, req_index, expr);
      req_index++;
    }
  }
  return expressions;
}

expr_ty _PyPegen_joined_str(Parser *p, Token *a, asdl_expr_seq *raw_expressions,
                            Token *b) {

  asdl_expr_seq *expr = unpack_top_level_joined_strs(p, raw_expressions);
  Py_ssize_t n_items = asdl_seq_LEN(expr);

  const char *quote_str = PyBytes_AsString(a->bytes);
  if (quote_str == NULL) {
    return NULL;
  }
  int is_raw = strpbrk(quote_str, "rR") != NULL;

  asdl_expr_seq *seq = _Py_asdl_expr_seq_new(n_items, p->arena);
  if (seq == NULL) {
    return NULL;
  }

  Py_ssize_t index = 0;
  for (Py_ssize_t i = 0; i < n_items; i++) {
    expr_ty item = asdl_seq_GET(expr, i);
    if (item->kind == Constant_kind) {
      item = _PyPegen_decode_fstring_part(p, is_raw, item, b);
      if (item == NULL) {
        return NULL;
      }

      /* Tokenizer emits string parts even when the underlying string
      might become an empty value (e.g. FSTRING_MIDDLE with the value \\n)
      so we need to check for them and simplify it here. */
      if (PyUnicode_CheckExact(item->v.Constant.value) &&
          PyUnicode_GET_LENGTH(item->v.Constant.value) == 0) {
        continue;
      }
    }
    asdl_seq_SET(seq, index++, item);
  }

  asdl_expr_seq *resized_exprs;
  if (index != n_items) {
    resized_exprs = _Py_asdl_expr_seq_new(index, p->arena);
    if (resized_exprs == NULL) {
      return NULL;
    }
    for (Py_ssize_t i = 0; i < index; i++) {
      asdl_seq_SET(resized_exprs, i, asdl_seq_GET(seq, i));
    }
  } else {
    resized_exprs = seq;
  }

  return _PyAST_JoinedStr(resized_exprs, a->lineno, a->col_offset,
                          b->end_lineno, b->end_col_offset, p->arena);
}

expr_ty _PyPegen_decoded_constant_from_token(Parser *p, Token *tok) {
  Py_ssize_t bsize;
  char *bstr;
  if (PyBytes_AsStringAndSize(tok->bytes, &bstr, &bsize) == -1) {
    return NULL;
  }
  PyObject *str = _PyPegen_decode_string(p, 0, bstr, bsize, tok);
  if (str == NULL) {
    return NULL;
  }
  if (_PyArena_AddPyObject(p->arena, str) < 0) {
    Py_DECREF(str);
    return NULL;
  }
  return _PyAST_Constant(str, NULL, tok->lineno, tok->col_offset,
                         tok->end_lineno, tok->end_col_offset, p->arena);
}

expr_ty _PyPegen_constant_from_token(Parser *p, Token *tok) {
  char *bstr = PyBytes_AsString(tok->bytes);
  if (bstr == NULL) {
    return NULL;
  }
  PyObject *str = PyUnicode_FromString(bstr);
  if (str == NULL) {
    return NULL;
  }
  if (_PyArena_AddPyObject(p->arena, str) < 0) {
    Py_DECREF(str);
    return NULL;
  }
  return _PyAST_Constant(str, NULL, tok->lineno, tok->col_offset,
                         tok->end_lineno, tok->end_col_offset, p->arena);
}

expr_ty _PyPegen_constant_from_string(Parser *p, Token *tok) {
  char *the_str = PyBytes_AsString(tok->bytes);
  if (the_str == NULL) {
    return NULL;
  }
  PyObject *s = _PyPegen_parse_string(p, tok);
  if (s == NULL) {
    _Pypegen_raise_decode_error(p);
    return NULL;
  }
  if (_PyArena_AddPyObject(p->arena, s) < 0) {
    Py_DECREF(s);
    return NULL;
  }
  PyObject *kind = NULL;
  if (the_str && the_str[0] == 'u') {
    kind = _PyPegen_new_identifier(p, "u");
    if (kind == NULL) {
      return NULL;
    }
  }
  return _PyAST_Constant(s, kind, tok->lineno, tok->col_offset, tok->end_lineno,
                         tok->end_col_offset, p->arena);
}

expr_ty _PyPegen_formatted_value(Parser *p, expr_ty expression, Token *debug,
                                 ResultTokenWithMetadata *conversion,
                                 ResultTokenWithMetadata *format,
                                 Token *closing_brace, int lineno,
                                 int col_offset, int end_lineno,
                                 int end_col_offset, PyArena *arena) {
  int conversion_val = -1;
  if (conversion != NULL) {
    expr_ty conversion_expr = (expr_ty)conversion->result;
    assert(conversion_expr->kind == Name_kind);
    Py_UCS4 first = PyUnicode_READ_CHAR(conversion_expr->v.Name.id, 0);

    if (PyUnicode_GET_LENGTH(conversion_expr->v.Name.id) > 1 ||
        !(first == 's' || first == 'r' || first == 'a')) {
      RAISE_SYNTAX_ERROR_KNOWN_LOCATION(
          conversion_expr,
          "f-string: invalid conversion character %R: expected 's', 'r', or "
          "'a'",
          conversion_expr->v.Name.id);
      return NULL;
    }

    conversion_val = Py_SAFE_DOWNCAST(first, Py_UCS4, int);
  } else if (debug && !format) {
    /* If no conversion is specified, use !r for debug expressions */
    conversion_val = (int)'r';
  }

  expr_ty formatted_value = _PyAST_FormattedValue(
      expression, conversion_val, format ? (expr_ty)format->result : NULL,
      lineno, col_offset, end_lineno, end_col_offset, arena);

  if (debug) {
    /* Find the non whitespace token after the "=" */
    int debug_end_line, debug_end_offset;
    PyObject *debug_metadata;

    if (conversion) {
      debug_end_line = ((expr_ty)conversion->result)->lineno;
      debug_end_offset = ((expr_ty)conversion->result)->col_offset;
      debug_metadata = conversion->metadata;
    } else if (format) {
      debug_end_line = ((expr_ty)format->result)->lineno;
      debug_end_offset = ((expr_ty)format->result)->col_offset + 1;
      debug_metadata = format->metadata;
    } else {
      debug_end_line = end_lineno;
      debug_end_offset = end_col_offset;
      debug_metadata = closing_brace->metadata;
    }
    expr_ty debug_text =
        _PyAST_Constant(debug_metadata, NULL, lineno, col_offset + 1,
                        debug_end_line, debug_end_offset - 1, p->arena);
    if (!debug_text) {
      return NULL;
    }

    asdl_expr_seq *values = _Py_asdl_expr_seq_new(2, arena);
    if (values == NULL) {
      return NULL;
    }
    asdl_seq_SET(values, 0, debug_text);
    asdl_seq_SET(values, 1, formatted_value);
    return _PyAST_JoinedStr(values, lineno, col_offset, debug_end_line,
                            debug_end_offset, p->arena);
  } else {
    return formatted_value;
  }
}

expr_ty _PyPegen_concatenate_strings(Parser *p, asdl_expr_seq *strings,
                                     int lineno, int col_offset, int end_lineno,
                                     int end_col_offset, PyArena *arena) {
  Py_ssize_t len = asdl_seq_LEN(strings);
  assert(len > 0);

  int f_string_found = 0;
  int unicode_string_found = 0;
  int bytes_found = 0;

  Py_ssize_t i = 0;
  Py_ssize_t n_flattened_elements = 0;
  for (i = 0; i < len; i++) {
    expr_ty elem = asdl_seq_GET(strings, i);
    switch (elem->kind) {
    case Constant_kind:
      if (PyBytes_CheckExact(elem->v.Constant.value)) {
        bytes_found = 1;
      } else {
        unicode_string_found = 1;
      }
      n_flattened_elements++;
      break;
    case JoinedStr_kind:
      n_flattened_elements += asdl_seq_LEN(elem->v.JoinedStr.values);
      f_string_found = 1;
      break;
    default:
      n_flattened_elements++;
      f_string_found = 1;
      break;
    }
  }

  if ((unicode_string_found || f_string_found) && bytes_found) {
    RAISE_SYNTAX_ERROR("cannot mix bytes and nonbytes literals");
    return NULL;
  }

  if (bytes_found) {
    PyObject *res = PyBytes_FromString("");

    /* Bytes literals never get a kind, but just for consistency
       since they are represented as Constant nodes, we'll mirror
       the same behavior as unicode strings for determining the
       kind. */
    PyObject *kind = asdl_seq_GET(strings, 0)->v.Constant.kind;
    for (i = 0; i < len; i++) {
      expr_ty elem = asdl_seq_GET(strings, i);
      PyBytes_Concat(&res, elem->v.Constant.value);
    }
    if (!res || _PyArena_AddPyObject(arena, res) < 0) {
      Py_XDECREF(res);
      return NULL;
    }
    return _PyAST_Constant(res, kind, lineno, col_offset, end_lineno,
                           end_col_offset, p->arena);
  }

  if (!f_string_found && len == 1) {
    return asdl_seq_GET(strings, 0);
  }

  asdl_expr_seq *flattened =
      _Py_asdl_expr_seq_new(n_flattened_elements, p->arena);
  if (flattened == NULL) {
    return NULL;
  }

  /* build flattened list */
  Py_ssize_t current_pos = 0;
  Py_ssize_t j = 0;
  for (i = 0; i < len; i++) {
    expr_ty elem = asdl_seq_GET(strings, i);
    switch (elem->kind) {
    case JoinedStr_kind:
      for (j = 0; j < asdl_seq_LEN(elem->v.JoinedStr.values); j++) {
        expr_ty subvalue = asdl_seq_GET(elem->v.JoinedStr.values, j);
        if (subvalue == NULL) {
          return NULL;
        }
        asdl_seq_SET(flattened, current_pos++, subvalue);
      }
      break;
    default:
      asdl_seq_SET(flattened, current_pos++, elem);
      break;
    }
  }

  /* calculate folded element count */
  Py_ssize_t n_elements = 0;
  int prev_is_constant = 0;
  for (i = 0; i < n_flattened_elements; i++) {
    expr_ty elem = asdl_seq_GET(flattened, i);

    /* The concatenation of a FormattedValue and an empty Contant should
       lead to the FormattedValue itself. Thus, we will not take any empty
       constants into account, just as in `_PyPegen_joined_str` */
    if (f_string_found && elem->kind == Constant_kind &&
        PyUnicode_CheckExact(elem->v.Constant.value) &&
        PyUnicode_GET_LENGTH(elem->v.Constant.value) == 0)
      continue;

    if (!prev_is_constant || elem->kind != Constant_kind) {
      n_elements++;
    }
    prev_is_constant = elem->kind == Constant_kind;
  }

  asdl_expr_seq *values = _Py_asdl_expr_seq_new(n_elements, p->arena);
  if (values == NULL) {
    return NULL;
  }

  /* build folded list */
  _PyUnicodeWriter writer;
  current_pos = 0;
  for (i = 0; i < n_flattened_elements; i++) {
    expr_ty elem = asdl_seq_GET(flattened, i);

    /* if the current elem and the following are constants,
       fold them and all consequent constants */
    if (elem->kind == Constant_kind) {
      if (i + 1 < n_flattened_elements &&
          asdl_seq_GET(flattened, i + 1)->kind == Constant_kind) {
        expr_ty first_elem = elem;

        /* When a string is getting concatenated, the kind of the string
           is determined by the first string in the concatenation
           sequence.

           u"abc" "def" -> u"abcdef"
           "abc" u"abc" ->  "abcabc" */
        PyObject *kind = elem->v.Constant.kind;

        _PyUnicodeWriter_Init(&writer);
        expr_ty last_elem = elem;
        for (j = i; j < n_flattened_elements; j++) {
          expr_ty current_elem = asdl_seq_GET(flattened, j);
          if (current_elem->kind == Constant_kind) {
            if (_PyUnicodeWriter_WriteStr(&writer,
                                          current_elem->v.Constant.value)) {
              _PyUnicodeWriter_Dealloc(&writer);
              return NULL;
            }
            last_elem = current_elem;
          } else {
            break;
          }
        }
        i = j - 1;

        PyObject *concat_str = _PyUnicodeWriter_Finish(&writer);
        if (concat_str == NULL) {
          _PyUnicodeWriter_Dealloc(&writer);
          return NULL;
        }
        if (_PyArena_AddPyObject(p->arena, concat_str) < 0) {
          Py_DECREF(concat_str);
          return NULL;
        }
        elem = _PyAST_Constant(concat_str, kind, first_elem->lineno,
                               first_elem->col_offset, last_elem->end_lineno,
                               last_elem->end_col_offset, p->arena);
        if (elem == NULL) {
          return NULL;
        }
      }

      /* Drop all empty contanst strings */
      if (f_string_found && PyUnicode_CheckExact(elem->v.Constant.value) &&
          PyUnicode_GET_LENGTH(elem->v.Constant.value) == 0) {
        continue;
      }
    }

    asdl_seq_SET(values, current_pos++, elem);
  }

  if (!f_string_found) {
    assert(n_elements == 1);
    expr_ty elem = asdl_seq_GET(values, 0);
    assert(elem->kind == Constant_kind);
    return elem;
  }

  assert(current_pos == n_elements);
  return _PyAST_JoinedStr(values, lineno, col_offset, end_lineno,
                          end_col_offset, p->arena);
}