[mypyc] Generate smaller code for casts (#12839)

Merge a cast op followed by a branch that does an error check and adds a
traceback entry. Since casts are very common, this reduces the size of
the generated code a fair amount.

Old code generated for a cast:
```
    if (likely(PyUnicode_Check(cpy_r_x)))
        cpy_r_r0 = cpy_r_x;
    else {
        CPy_TypeError("str", cpy_r_x);
        cpy_r_r0 = NULL;
    }
    if (unlikely(cpy_r_r0 == NULL)) {
        CPy_AddTraceback("t/t.py", "foo", 2, CPyStatic_globals);
        goto CPyL2;
    }
```

New code:
```
    if (likely(PyUnicode_Check(cpy_r_x)))
        cpy_r_r0 = cpy_r_x;
    else {
        CPy_TypeErrorTraceback("t/t.py", "foo", 2, CPyStatic_globals, "str", cpy_r_x);
        goto CPyL2;
    }
```

mypyc/codegen/emit.py

@@ -2,6 +2,8 @@
 
 from mypy.backports import OrderedDict
 from typing import List, Set, Dict, Optional, Callable, Union, Tuple
+from typing_extensions import Final
+
 import sys
 
 from mypyc.common import (
@@ -23,6 +25,10 @@
 from mypyc.sametype import is_same_type
 from mypyc.codegen.literals import Literals
 
+# Whether to insert debug asserts for all error handling, to quickly
+# catch errors propagating without exceptions set.
+DEBUG_ERRORS: Final = False
+
 
 class HeaderDeclaration:
  """A representation of a declaration in C.
@@ -104,6 +110,20 @@ def __init__(self, label: str) -> None:
  self.label = label
 
 
+class TracebackAndGotoHandler(ErrorHandler):
+ """Add traceback item and goto label on error."""
+
+ def __init__(self,
+ label: str,
+ source_path: str,
+ module_name: str,
+ traceback_entry: Tuple[str, int]) -> None:
+ self.label = label
+ self.source_path = source_path
+ self.module_name = module_name
+ self.traceback_entry = traceback_entry
+
+
 class ReturnHandler(ErrorHandler):
  """Return a constant value on error."""
 
@@ -439,18 +459,6 @@ def emit_cast(self,
  likely: If the cast is likely to succeed (can be False for unions)
  """
  error = error or AssignHandler()
- if isinstance(error, AssignHandler):
- handle_error = '%s = NULL;' % dest
- elif isinstance(error, GotoHandler):
- handle_error = 'goto %s;' % error.label
- else:
- assert isinstance(error, ReturnHandler)
- handle_error = 'return %s;' % error.value
- if raise_exception:
- raise_exc = f'CPy_TypeError("{self.pretty_name(typ)}", {src}); '
- err = raise_exc + handle_error
- else:
- err = handle_error
 
  # Special case casting *from* optional
  if src_type and is_optional_type(src_type) and not is_object_rprimitive(typ):
@@ -465,9 +473,9 @@ def emit_cast(self,
  self.emit_arg_check(src, dest, typ, check.format(src), optional)
  self.emit_lines(
  f' {dest} = {src};',
- 'else {',
-  err,
-  '}')
+ 'else {')
+ self.emit_cast_error_handler(error, src, dest, typ, raise_exception)
+ self.emit_line('}')
  return
 
  # TODO: Verify refcount handling.
@@ -500,9 +508,9 @@ def emit_cast(self,
  self.emit_arg_check(src, dest, typ, check.format(prefix, src), optional)
  self.emit_lines(
  f' {dest} = {src};',
- 'else {',
-  err,
-  '}')
+ 'else {')
+ self.emit_cast_error_handler(error, src, dest, typ, raise_exception)
+ self.emit_line('}')
  elif is_bytes_rprimitive(typ):
  if declare_dest:
  self.emit_line(f'PyObject *{dest};')
@@ -512,9 +520,9 @@ def emit_cast(self,
  self.emit_arg_check(src, dest, typ, check.format(src, src), optional)
  self.emit_lines(
  f' {dest} = {src};',
- 'else {',
-  err,
-  '}')
+ 'else {')
+ self.emit_cast_error_handler(error, src, dest, typ, raise_exception)
+ self.emit_line('}')
  elif is_tuple_rprimitive(typ):
  if declare_dest:
  self.emit_line(f'{self.ctype(typ)} {dest};')
@@ -525,9 +533,9 @@ def emit_cast(self,
  check.format(src), optional)
  self.emit_lines(
  f' {dest} = {src};',
- 'else {',
-  err,
-  '}')
+ 'else {')
+ self.emit_cast_error_handler(error, src, dest, typ, raise_exception)
+ self.emit_line('}')
  elif isinstance(typ, RInstance):
  if declare_dest:
  self.emit_line(f'PyObject *{dest};')
@@ -551,10 +559,10 @@ def emit_cast(self,
  check = f'(likely{check})'
  self.emit_arg_check(src, dest, typ, check, optional)
  self.emit_lines(
- f' {dest} = {src};',
- 'else {',
-  err,
-  '}')
+ f' {dest} = {src};'.format(dest, src),
+ 'else {')
+ self.emit_cast_error_handler(error, src, dest, typ, raise_exception)
+ self.emit_line('}')
  elif is_none_rprimitive(typ):
  if declare_dest:
  self.emit_line(f'PyObject *{dest};')
@@ -565,9 +573,9 @@ def emit_cast(self,
  check.format(src), optional)
  self.emit_lines(
  f' {dest} = {src};',
- 'else {',
-  err,
-  '}')
+ 'else {')
+ self.emit_cast_error_handler(error, src, dest, typ, raise_exception)
+ self.emit_line('}')
  elif is_object_rprimitive(typ):
  if declare_dest:
  self.emit_line(f'PyObject *{dest};')
@@ -576,21 +584,51 @@ def emit_cast(self,
  if optional:
  self.emit_line('}')
  elif isinstance(typ, RUnion):
- self.emit_union_cast(src, dest, typ, declare_dest, err, optional, src_type)
+ self.emit_union_cast(src, dest, typ, declare_dest, error, optional, src_type,
+ raise_exception)
  elif isinstance(typ, RTuple):
  assert not optional
- self.emit_tuple_cast(src, dest, typ, declare_dest, err, src_type)
+ self.emit_tuple_cast(src, dest, typ, declare_dest, error, src_type)
  else:
  assert False, 'Cast not implemented: %s' % typ
 
+ def emit_cast_error_handler(self,
+ error: ErrorHandler,
+ src: str,
+ dest: str,
+ typ: RType,
+ raise_exception: bool) -> None:
+ if raise_exception:
+ if isinstance(error, TracebackAndGotoHandler):
+ # Merge raising and emitting traceback entry into a single call.
+ self.emit_type_error_traceback(
+ error.source_path, error.module_name, error.traceback_entry,
+ typ=typ,
+ src=src)
+ self.emit_line('goto %s;' % error.label)
+ return
+ self.emit_line('CPy_TypeError("{}", {}); '.format(self.pretty_name(typ), src))
+ if isinstance(error, AssignHandler):
+ self.emit_line('%s = NULL;' % dest)
+ elif isinstance(error, GotoHandler):
+ self.emit_line('goto %s;' % error.label)
+ elif isinstance(error, TracebackAndGotoHandler):
+ self.emit_line('%s = NULL;' % dest)
+ self.emit_traceback(error.source_path, error.module_name, error.traceback_entry)
+ self.emit_line('goto %s;' % error.label)
+ else:
+ assert isinstance(error, ReturnHandler)
+ self.emit_line('return %s;' % error.value)
+
  def emit_union_cast(self,
  src: str,
  dest: str,
  typ: RUnion,
  declare_dest: bool,
- err: str,
+ error: ErrorHandler,
  optional: bool,
- src_type: Optional[RType]) -> None:
+ src_type: Optional[RType],
+ raise_exception: bool) -> None:
  """Emit cast to a union type.
 
  The arguments are similar to emit_cast.
@@ -613,11 +651,11 @@ def emit_union_cast(self,
  likely=False)
  self.emit_line(f'if ({dest} != NULL) goto {good_label};')
  # Handle cast failure.
- self.emit_line(err)
+ self.emit_cast_error_handler(error, src, dest, typ, raise_exception)
  self.emit_label(good_label)
 
  def emit_tuple_cast(self, src: str, dest: str, typ: RTuple, declare_dest: bool,
- err: str, src_type: Optional[RType]) -> None:
+ error: ErrorHandler, src_type: Optional[RType]) -> None:
  """Emit cast to a tuple type.
 
  The arguments are similar to emit_cast.
@@ -740,7 +778,8 @@ def emit_unbox(self,
  self.emit_line('} else {')
 
  cast_temp = self.temp_name()
- self.emit_tuple_cast(src, cast_temp, typ, declare_dest=True, err='', src_type=None)
+ self.emit_tuple_cast(src, cast_temp, typ, declare_dest=True, error=error,
+ src_type=None)
  self.emit_line(f'if (unlikely({cast_temp} == NULL)) {{')
 
  # self.emit_arg_check(src, dest, typ,
@@ -886,3 +925,44 @@ def emit_gc_clear(self, target: str, rtype: RType) -> None:
  self.emit_line(f'Py_CLEAR({target});')
  else:
  assert False, 'emit_gc_clear() not implemented for %s' % repr(rtype)
+
+ def emit_traceback(self,
+ source_path: str,
+ module_name: str,
+ traceback_entry: Tuple[str, int]) -> None:
+ return self._emit_traceback('CPy_AddTraceback', source_path, module_name, traceback_entry)
+
+ def emit_type_error_traceback(
+ self,
+ source_path: str,
+ module_name: str,
+ traceback_entry: Tuple[str, int],
+ *,
+ typ: RType,
+ src: str) -> None:
+ func = 'CPy_TypeErrorTraceback'
+ type_str = f'"{self.pretty_name(typ)}"'
+ return self._emit_traceback(
+ func, source_path, module_name, traceback_entry, type_str=type_str, src=src)
+
+ def _emit_traceback(self,
+ func: str,
+ source_path: str,
+ module_name: str,
+ traceback_entry: Tuple[str, int],
+ type_str: str = '',
+ src: str = '') -> None:
+ globals_static = self.static_name('globals', module_name)
+ line = '%s("%s", "%s", %d, %s' % (
+ func,
+ source_path.replace("\\", "\\\\"),
+ traceback_entry[0],
+ traceback_entry[1],
+ globals_static)
+ if type_str:
+ assert src
+ line += f', {type_str}, {src}'
+ line += ');'
+ self.emit_line(line)
+ if DEBUG_ERRORS:
+ self.emit_line('assert(PyErr_Occurred() != NULL && "failure w/o err!");')

mypyc/codegen/emitfunc.py

@@ -6,7 +6,7 @@
 from mypyc.common import (
  REG_PREFIX, NATIVE_PREFIX, STATIC_PREFIX, TYPE_PREFIX, MODULE_PREFIX,
 )
-from mypyc.codegen.emit import Emitter
+from mypyc.codegen.emit import Emitter, TracebackAndGotoHandler, DEBUG_ERRORS
 from mypyc.ir.ops import (
  Op, OpVisitor, Goto, Branch, Return, Assign, Integer, LoadErrorValue, GetAttr, SetAttr,
  LoadStatic, InitStatic, TupleGet, TupleSet, Call, IncRef, DecRef, Box, Cast, Unbox,
@@ -23,10 +23,6 @@
 from mypyc.ir.pprint import generate_names_for_ir
 from mypyc.analysis.blockfreq import frequently_executed_blocks
 
-# Whether to insert debug asserts for all error handling, to quickly
-# catch errors propagating without exceptions set.
-DEBUG_ERRORS = False
-
 
 def native_function_type(fn: FuncIR, emitter: Emitter) -> str:
  args = ', '.join(emitter.ctype(arg.type) for arg in fn.args) or 'void'
@@ -322,7 +318,7 @@ def visit_get_attr(self, op: GetAttr) -> None:
  and branch.traceback_entry is not None
  and not branch.negated):
  # Generate code for the following branch here to avoid
- # redundant branches in the generate code.
+ # redundant branches in the generated code.
  self.emit_attribute_error(branch, cl.name, op.attr)
  self.emit_line('goto %s;' % self.label(branch.true))
  merged_branch = branch
@@ -485,8 +481,24 @@ def visit_box(self, op: Box) -> None:
  self.emitter.emit_box(self.reg(op.src), self.reg(op), op.src.type, can_borrow=True)
 
  def visit_cast(self, op: Cast) -> None:
+ branch = self.next_branch()
+ handler = None
+ if branch is not None:
+ if (branch.value is op
+ and branch.op == Branch.IS_ERROR
+ and branch.traceback_entry is not None
+ and not branch.negated
+ and branch.false is self.next_block):
+ # Generate code also for the following branch here to avoid
+ # redundant branches in the generated code.
+ handler = TracebackAndGotoHandler(self.label(branch.true),
+ self.source_path,
+ self.module_name,
+ branch.traceback_entry)
+ self.op_index += 1
+
  self.emitter.emit_cast(self.reg(op.src), self.reg(op), op.type,
- src_type=op.src.type)
+ src_type=op.src.type, error=handler)
 
  def visit_unbox(self, op: Unbox) -> None:
  self.emitter.emit_unbox(self.reg(op.src), self.reg(op), op.type)
@@ -647,14 +659,7 @@ def emit_declaration(self, line: str) -> None:
 
  def emit_traceback(self, op: Branch) -> None:
  if op.traceback_entry is not None:
- globals_static = self.emitter.static_name('globals', self.module_name)
- self.emit_line('CPy_AddTraceback("%s", "%s", %d, %s);' % (
- self.source_path.replace("\\", "\\\\"),
- op.traceback_entry[0],
- op.traceback_entry[1],
- globals_static))
- if DEBUG_ERRORS:
- self.emit_line('assert(PyErr_Occurred() != NULL && "failure w/o err!");')
+ self.emitter.emit_traceback(self.source_path, self.module_name, op.traceback_entry)
 
  def emit_attribute_error(self, op: Branch, class_name: str, attr: str) -> None:
  assert op.traceback_entry is not None

mypyc/lib-rt/CPy.h

@@ -500,6 +500,8 @@ void _CPy_GetExcInfo(PyObject **p_type, PyObject **p_value, PyObject **p_traceba
 void CPyError_OutOfMemory(void);
 void CPy_TypeError(const char *expected, PyObject *value);
 void CPy_AddTraceback(const char *filename, const char *funcname, int line, PyObject *globals);
+void CPy_TypeErrorTraceback(const char *filename, const char *funcname, int line,
+ PyObject *globals, const char *expected, PyObject *value);
 void CPy_AttributeError(const char *filename, const char *funcname, const char *classname,
  const char *attrname, int line, PyObject *globals);
 

mypyc/lib-rt/exc_ops.c

@@ -233,6 +233,13 @@ void CPy_AddTraceback(const char *filename, const char *funcname, int line, PyOb
  _PyErr_ChainExceptions(exc, val, tb);
 }
 
+CPy_NOINLINE
+void CPy_TypeErrorTraceback(const char *filename, const char *funcname, int line,
+ PyObject *globals, const char *expected, PyObject *value) {
+ CPy_TypeError(expected, value);
+ CPy_AddTraceback(filename, funcname, line, globals);
+}
+
 void CPy_AttributeError(const char *filename, const char *funcname, const char *classname,
  const char *attrname, int line, PyObject *globals) {
  char buf[500];

mypyc/test-data/run-functions.test

@@ -1220,3 +1220,18 @@ def sub(s: str, f: Callable[[str], str]) -> str: ...
 def sub(s: bytes, f: Callable[[bytes], bytes]) -> bytes: ...
 def sub(s, f):
  return f(s)
+
+[case testContextManagerSpecialCase]
+from typing import Generator, Callable, Iterator
+from contextlib import contextmanager
+
+@contextmanager
+def f() -> Iterator[None]:
+ yield
+
+def g() -> None:
+ a = ['']
+ with f():
+ a.pop()
+
+g()