Start adding support for Optional[...]

JukkaL · JukkaL · commit 3e4a7202b2d2 · 2017-09-10T16:07:00.000+01:00
Also fix a fix bugs that were exposed.
diff --git a/doc/dev-intro.md b/doc/dev-intro.md
@@ -182,6 +182,8 @@ Here are some hints about how to add support for a new primitive type
   `supports_unbox`, `ctype` and various other properties work
   correctly for the new type.
 
+* Update `mypyc.genops.Mapper.type_to_rtype()`.
+
 * Update `emit_box` in `mypyc.emit`.
 
 * Update `emit_unbox` or `emit_cast` in `mypyc.emit`.
@@ -193,8 +195,6 @@ Here are some hints about how to add support for a new primitive type
 
 * Update `emit_error_check` in `mypyc.emit` for unboxed types.
 
-* Update `mypyc.genops.Mapper.type_to_rtype()`.
-
 The above may be enough to allow you to declare variables with the
 type and pass values around. You likely also want to add support for
 some primitive operations for the type (see Built-in Operation for an
diff --git a/mypyc/analysis.py b/mypyc/analysis.py
@@ -254,7 +254,7 @@ def analyze_undefined_regs(blocks: List[BasicBlock],
 
 class LivenessVisitor(BaseAnalysisVisitor):
     def visit_branch(self, op: Branch) -> GenAndKill:
-        return {op.left, op.right}, set()
+        return set(op.sources()), set()
 
     def visit_return(self, op: Return) -> GenAndKill:
         return {op.reg}, set()
diff --git a/mypyc/emit.py b/mypyc/emit.py
@@ -4,7 +4,8 @@
 
 from mypyc.common import REG_PREFIX
 from mypyc.ops import (
-    Environment, Label, Register, RType, ObjectRType, TupleRType, UserRType, type_struct_name
+    Environment, Label, Register, RType, ObjectRType, TupleRType, UserRType, OptionalRType,
+    type_struct_name
 )
 
 
@@ -162,7 +163,20 @@ def emit_cast(self, src: str, dest: str, typ: RType, failure: str,
         elif typ.name == 'None':
             if declare_dest:
                 self.emit_line('PyObject *{};'.format(dest))
-            self.emit_line('{} = {};'.format(dest, src))
+            self.emit_lines(
+                'if ({} == Py_None)'.format(src),
+                '    {} = {};'.format(dest, src),
+                'else',
+                failure)
+        elif isinstance(typ, OptionalRType):
+            if declare_dest:
+                self.emit_line('PyObject *{};'.format(dest))
+            self.emit_lines(
+                'if ({} == Py_None)'.format(src),
+                '    {} = {};'.format(dest, src),
+                'else {')
+            self.emit_cast(src, dest, typ.value_type, failure.lstrip())
+            self.emit_line('}')
         else:
             assert False, 'Cast not implemented: %s' % typ
 
@@ -263,6 +277,7 @@ def emit_box(self, src: str, dest: str, typ: RType, failure: str,
                                   declare_dest=True)
                     self.emit_line('PyTuple_SetItem({}, {}, {});'.format(dest, i, inner_name, i))
         else:
+            assert not typ.supports_unbox
             # Type is boxed -- trivially just assign.
             self.emit_line('{}{} = {};'.format(declaration, dest, src))
 
diff --git a/mypyc/emitfunc.py b/mypyc/emitfunc.py
@@ -73,6 +73,9 @@ def visit_branch(self, op: Branch) -> None:
         if op.op == Branch.BOOL_EXPR:
             expr_result = self.reg(op.left) # right isn't used
             self.emit_line('if ({}({}))'.format(neg, expr_result))
+        elif op.op == Branch.IS_NONE:
+            compare = '!=' if op.negated else '=='
+            self.emit_line('if ({} {} Py_None)'.format(self.reg(op.left), compare))
         else:
             left = self.reg(op.left)
             right = self.reg(op.right)
diff --git a/mypyc/genops.py b/mypyc/genops.py
@@ -23,7 +23,7 @@ def f(x: int) -> int:
     ContinueStmt, ConditionalExpr, OperatorAssignmentStmt, TupleExpr, ClassDef, TypeInfo,
     Import, ImportFrom, ImportAll, ARG_POS, MODULE_REF
 )
-from mypy.types import Type, Instance, CallableType, NoneTyp, TupleType
+from mypy.types import Type, Instance, CallableType, NoneTyp, TupleType, UnionType
 from mypy.visitor import NodeVisitor
 from mypy.subtypes import is_named_instance
 
@@ -32,7 +32,7 @@ def f(x: int) -> int:
     PrimitiveOp, Branch, Goto, RuntimeArg, Call, Box, Unbox, Cast, TupleRType,
     Unreachable, TupleGet, ClassIR, UserRType, ModuleIR, GetAttr, SetAttr, LoadStatic,
     PyGetAttr, PyCall, IntRType, BoolRType, ListRType, SequenceTupleRType, ObjectRType, NoneRType,
-    c_module_name, INVALID_REGISTER, INVALID_LABEL
+    OptionalRType, c_module_name, INVALID_REGISTER, INVALID_LABEL
 )
 
 
@@ -69,6 +69,13 @@ def type_to_rtype(self, typ: Type) -> RType:
             return ObjectRType()
         elif isinstance(typ, NoneTyp):
             return NoneRType()
+        elif isinstance(typ, UnionType):
+            assert len(typ.items) == 2 and any(isinstance(it, NoneTyp) for it in typ.items)
+            if isinstance(typ.items[0], NoneTyp):
+                value_type = typ.items[1]
+            else:
+                value_type = typ.items[0]
+            return OptionalRType(self.type_to_rtype(value_type))
         assert False, '%s unsupported' % type(typ)
 
 
@@ -788,9 +795,16 @@ def process_conditional(self, e: Node) -> List[Branch]:
                 right = self.accept(e.operands[1])
                 opcode = self.int_relative_ops[op]
                 branch = Branch(left, right, INVALID_LABEL, INVALID_LABEL, opcode)
-                self.add(branch)
-                return [branch]
-            assert False, "unsupported comparison epxression"
+            elif op in ['is', 'is not']:
+                left = self.accept(e.operands[0])
+                branch = Branch(left, INVALID_REGISTER, INVALID_LABEL, INVALID_LABEL,
+                                Branch.IS_NONE)
+                if op == 'is not':
+                    branch.negated = True
+            else:
+                assert False, "unsupported comparison epxression"
+            self.add(branch)
+            return [branch]
         elif isinstance(e, OpExpr) and e.op in ['and', 'or']:
             if e.op == 'and':
                 # Short circuit 'and' in a conditional context.
diff --git a/mypyc/ops.py b/mypyc/ops.py
@@ -29,12 +29,12 @@
 # Eventually we may want to separate expression visitors and statement-like visitors at
 # the type level but until then returning INVALID_REGISTER from a statement-like visitor
 # seems acceptable.
-INVALID_REGISTER = Register(-1)
+INVALID_REGISTER = Register(-99999)
 
 
 # Similarly this is used for placeholder labels which aren't assigned yet (but will
 # be eventually. Its kind of a hack.
-INVALID_LABEL = Label(-1)
+INVALID_LABEL = Label(-88888)
 
 
 def c_module_name(module_name: str):
@@ -255,6 +255,23 @@ def __repr__(self) -> str:
         return '<UserRType %s>' % self.name
 
 
+class OptionalRType(PyObjectRType):
+    """Optional[x]"""
+
+    def __init__(self, value_type: RType) -> None:
+        self.name = 'optional'
+        self.value_type = value_type
+
+    def __repr__(self) -> str:
+        return '<OptionalRType %s>' % self.value_type
+
+    def __eq__(self, other: object) -> bool:
+        return isinstance(other, OptionalRType) and other.value_type == self.value_type
+
+    def __hash__(self) -> int:
+        return hash(('optional', self.value_type))
+
+
 class Environment:
     """Keep track of names and types of registers."""
 
@@ -361,9 +378,10 @@ class Branch(Op):
     INT_GT = 14
     INT_GE = 15
 
-    # Unlike the above, a unary operation so it only uses the "left" register
+    # Unlike the above, these are unary operations so they only uses the "left" register
     # ("right" should be INVALID_REGISTER).
-    BOOL_EXPR = 16
+    BOOL_EXPR = 100
+    IS_NONE = 101
 
     op_names = {
         INT_EQ:  ('==', 'int'),
@@ -374,6 +392,11 @@ class Branch(Op):
         INT_GE:  ('>=', 'int'),
     }
 
+    unary_op_names = {
+        BOOL_EXPR: ('%r', 'bool'),
+        IS_NONE: ('%r is None', 'object'),
+    }
+
     def __init__(self, left: Register, right: Register, true_label: Label,
                  false_label: Label, op: int) -> None:
         self.left = left
@@ -384,24 +407,24 @@ def __init__(self, left: Register, right: Register, true_label: Label,
         self.negated = False
 
     def sources(self) -> List[Register]:
-        return [self.left, self.right]
+        if self.right != INVALID_REGISTER:
+            return [self.left, self.right]
+        else:
+            return [self.left]
 
     def to_str(self, env: Environment) -> str:
         # Right not used for BOOL_EXPR
-        if self.op != Branch.BOOL_EXPR:
+        if self.op in self.op_names:
             if self.negated:
                 fmt = 'not %r {} %r'
             else:
                 fmt = '%r {} %r'
             op, typ = self.op_names[self.op]
             fmt = fmt.format(op)
-
         else:
+            fmt, typ = self.unary_op_names[self.op]
             if self.negated:
-                fmt = 'not %r'
-            else:
-                fmt = '%r'
-            typ = 'bool'
+                fmt = 'not {}'.format(fmt)
 
         cond = env.format(fmt, self.left, self.right)
         fmt = 'if {} goto %l else goto %l :: {}'.format(cond, typ)
@@ -674,7 +697,7 @@ def sources(self) -> List[Register]:
 
     def to_str(self, env: Environment) -> str:
         params = {}  # type: Dict[str, Any]
-        if self.dest is not None:
+        if self.dest is not None and self.dest != INVALID_REGISTER:
             params['dest'] = env.format('%r', self.dest)
         args = [env.format('%r', arg) for arg in self.args]
         params['args'] = args
diff --git a/test-data/genops-basic.test b/test-data/genops-basic.test
@@ -578,3 +578,32 @@ L0:
     r5 = cast(None, r2)
     r6 = None
     return r6
+
+[case testIsNone]
+from typing import Optional
+
+class A: pass
+
+def f(x: Optional[A]) -> int:
+    if x is None:
+        return 1
+    if x is not None:
+        return 2
+    return 3
+[out]
+def f(x):
+    x :: optional
+    r0, r1, r2 :: int
+L0:
+    if x is None goto L1 else goto L2 :: object
+L1:
+    r0 = 1
+    return r0
+L2:
+    if not x is None goto L3 else goto L4 :: object
+L3:
+    r1 = 2
+    return r1
+L4:
+    r2 = 3
+    return r2
diff --git a/test-data/refcount.test b/test-data/refcount.test
@@ -544,3 +544,18 @@ L0:
     dec_ref d
     r3 = None
     return r3
+
+[case testUnaryBranchSpecialCase]
+def f(x: bool) -> int:
+    if x:
+        return 1
+    return 2
+[out]
+L0:
+    if x goto L1 else goto L2 :: bool
+L1:
+    r0 = 1
+    return r0
+L2:
+    r1 = 2
+    return r1
diff --git a/test-data/run.test b/test-data/run.test
@@ -257,3 +257,25 @@ from native import f
 f(5)
 [out]
 5
+
+[case testOptional]
+from typing import Optional
+
+class A: pass
+
+def f(x: Optional[A]) -> Optional[A]:
+    return x
+
+def g(x: Optional[A]) -> int:
+    if x is None:
+        return 1
+    if x is not None:
+        return 2
+    return 3
+[file driver.py]
+from native import f, g, A
+a = A()
+assert f(None) is None
+assert f(a) is a
+assert g(None) == 1
+assert g(a) == 2
