py/emitnative: Implement viper unary ops positive, negative and invert.

Signed-off-by: Damien George <damien@micropython.org>

5 files changed, 74 insertions, 11 deletions

diff --git a/py/emitnative.c b/py/emitnative.c
index f80461dd4..0b84a2ec8 100644
--- a/py/emitnative.c
+++ b/py/emitnative.c
@@ -2259,15 +2259,38 @@ static void emit_native_pop_except_jump(emit_t *emit, mp_uint_t label, bool with
 }
 
 static void emit_native_unary_op(emit_t *emit, mp_unary_op_t op) {
-    vtype_kind_t vtype;
-    emit_pre_pop_reg(emit, &vtype, REG_ARG_2);
-    if (vtype == VTYPE_PYOBJ) {
+    vtype_kind_t vtype = peek_vtype(emit, 0);
+    if (vtype == VTYPE_INT || vtype == VTYPE_UINT) {
+        if (op == MP_UNARY_OP_POSITIVE) {
+            // No-operation, just leave the argument on the stack.
+        } else if (op == MP_UNARY_OP_NEGATIVE) {
+            int reg = REG_RET;
+            emit_pre_pop_reg_flexible(emit, &vtype, &reg, reg, reg);
+            ASM_NEG_REG(emit->as, reg);
+            emit_post_push_reg(emit, vtype, reg);
+        } else if (op == MP_UNARY_OP_INVERT) {
+            #ifdef ASM_NOT_REG
+            int reg = REG_RET;
+            emit_pre_pop_reg_flexible(emit, &vtype, &reg, reg, reg);
+            ASM_NOT_REG(emit->as, reg);
+            #else
+            int reg = REG_RET;
+            emit_pre_pop_reg_flexible(emit, &vtype, &reg, REG_ARG_1, reg);
+            ASM_MOV_REG_IMM(emit->as, REG_ARG_1, -1);
+            ASM_XOR_REG_REG(emit->as, reg, REG_ARG_1);
+            #endif
+            emit_post_push_reg(emit, vtype, reg);
+        } else {
+            EMIT_NATIVE_VIPER_TYPE_ERROR(emit,
+                MP_ERROR_TEXT("'not' not implemented"), mp_binary_op_method_name[op]);
+        }
+    } else if (vtype == VTYPE_PYOBJ) {
+        emit_pre_pop_reg(emit, &vtype, REG_ARG_2);
         emit_call_with_imm_arg(emit, MP_F_UNARY_OP, op, REG_ARG_1);
         emit_post_push_reg(emit, VTYPE_PYOBJ, REG_RET);
     } else {
-        adjust_stack(emit, 1);
         EMIT_NATIVE_VIPER_TYPE_ERROR(emit,
-            MP_ERROR_TEXT("unary op %q not implemented"), mp_unary_op_method_name[op]);
+            MP_ERROR_TEXT("can't do unary op of '%q'"), vtype_to_qstr(vtype));
     }
 }
 
diff --git a/tests/micropython/viper_error.py b/tests/micropython/viper_error.py
index 80617af0c..6c5c3ba20 100644
--- a/tests/micropython/viper_error.py
+++ b/tests/micropython/viper_error.py
@@ -50,6 +50,9 @@ def f():
 # incorrect return type
 test("@micropython.viper\ndef f() -> int: return []")
 
+# can't do unary op of incompatible type
+test("@micropython.viper\ndef f(x:ptr): -x")
+
 # can't do binary op between incompatible types
 test("@micropython.viper\ndef f(): 1 + []")
 test("@micropython.viper\ndef f(x:int, y:uint): x < y")
@@ -69,9 +72,7 @@ test("@micropython.viper\ndef f(x:ptr32): x[x] = None")
 test("@micropython.viper\ndef f(): raise 1")
 
 # unary ops not implemented
-test("@micropython.viper\ndef f(x:int): +x")
-test("@micropython.viper\ndef f(x:int): -x")
-test("@micropython.viper\ndef f(x:int): ~x")
+test("@micropython.viper\ndef f(x:int): not x")
 
 # binary op not implemented
 test("@micropython.viper\ndef f(x:uint, y:uint): res = x // y")
diff --git a/tests/micropython/viper_error.py.exp b/tests/micropython/viper_error.py.exp
index 31c85b1d8..51cbd6c70 100644
--- a/tests/micropython/viper_error.py.exp
+++ b/tests/micropython/viper_error.py.exp
@@ -5,6 +5,7 @@ ViperTypeError("local 'x' used before type known",)
 ViperTypeError("local 'x' has type 'int' but source is 'object'",)
 ViperTypeError("can't implicitly convert 'ptr' to 'bool'",)
 ViperTypeError("return expected 'int' but got 'object'",)
+ViperTypeError("can't do unary op of 'ptr'",)
 ViperTypeError("can't do binary op between 'int' and 'object'",)
 ViperTypeError('comparison of int and uint',)
 ViperTypeError("can't load from 'int'",)
@@ -15,9 +16,7 @@ ViperTypeError("can't store to 'int'",)
 ViperTypeError("can't store 'None'",)
 ViperTypeError("can't store 'None'",)
 ViperTypeError('must raise an object',)
-ViperTypeError('unary op __pos__ not implemented',)
-ViperTypeError('unary op __neg__ not implemented',)
-ViperTypeError('unary op __invert__ not implemented',)
+ViperTypeError("'not' not implemented",)
 ViperTypeError('div/mod not implemented for uint',)
 ViperTypeError('div/mod not implemented for uint',)
 ViperTypeError('binary op  not implemented',)
diff --git a/tests/micropython/viper_unop.py b/tests/micropython/viper_unop.py
new file mode 100644
index 000000000..61cbd5125
--- /dev/null
+++ b/tests/micropython/viper_unop.py
@@ -0,0 +1,31 @@
+# test unary operators
+
+
+@micropython.viper
+def pos(x: int) -> int:
+    return +x
+
+
+print(pos(0))
+print(pos(1))
+print(pos(-2))
+
+
+@micropython.viper
+def neg(x: int) -> int:
+    return -x
+
+
+print(neg(0))
+print(neg(1))
+print(neg(-2))
+
+
+@micropython.viper
+def inv(x: int) -> int:
+    return ~x
+
+
+print(inv(0))
+print(inv(1))
+print(inv(-2))
diff --git a/tests/micropython/viper_unop.py.exp b/tests/micropython/viper_unop.py.exp
new file mode 100644
index 000000000..6d93312ca
--- /dev/null
+++ b/tests/micropython/viper_unop.py.exp
@@ -0,0 +1,9 @@
+0
+1
+-2
+0
+-1
+2
+-1
+-2
+1