From 05b95b7fa63d713dd2e7d696ae7f7e898a277594 Mon Sep 17 00:00:00 2001
From: Yan Da <dyanab@connect.ust.hk>
Date: Wed, 9 Dec 2020 12:35:24 +0800
Subject: [PATCH] [LANG] Add support for PREFIX_INC and PREFIX_DEC.

---
 include/triton/lang/code_gen.h |  1 +
 lib/lang/code_gen.cc           | 37 ++++++++++++++++++++++++++++++----
 2 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/include/triton/lang/code_gen.h b/include/triton/lang/code_gen.h
index 5347421bb..531da19c5 100644
--- a/include/triton/lang/code_gen.h
+++ b/include/triton/lang/code_gen.h
@@ -48,6 +48,7 @@ protected:
 
   void set_ret(ir::value* value);
   ir::value *GenUnaryMinus(ir::value* arg);
+  ir::value *GenUnaryInc(UnaryOp* arg, bool is_postfix, bool is_inc);
 
 public:
   Generator(Parser* parser) : parser_(parser) {}
diff --git a/lib/lang/code_gen.cc b/lib/lang/code_gen.cc
index 2d72e1794..5ac843b86 100644
--- a/lib/lang/code_gen.cc
+++ b/lib/lang/code_gen.cc
@@ -182,6 +182,35 @@ ir::value* Generator::GenUnaryMinus(ir::value* arg) {
     return bld_->create_sub(_0, arg);
 }
 
+ir::value* Generator::GenUnaryInc(UnaryOp* expr, bool is_postfix,
+                                  bool is_inc) {
+  Visit(expr->operand_);
+  ir::value* arg = ret_;
+
+  ir::value *_1 = nullptr;
+  if (arg->get_type()->is_half_ty() ||
+      arg->get_type()->is_float_ty() ||
+      arg->get_type()->is_double_ty())
+    _1 = ir::constant_fp::get(arg->get_type(), 1.0);
+  else if (arg->get_type()->is_integer_ty())
+    _1 = ir::constant_int::get(arg->get_type(), 1);
+  else 
+    error_not_implemented("data type not supported for unary inc");
+  
+  if (is_postfix) {
+    error_not_implemented("postfix inc not implemented");
+    return nullptr;
+  } else {
+    ir::value *instr = nullptr;
+    if (is_inc)
+      instr = bld_->create_add(arg, _1);
+    else 
+      instr = bld_->create_sub(arg, _1);
+    mod_->set_value(arg->get_name(), instr);
+    return instr;
+  }
+}
+
 void Generator::VisitUnaryOp(UnaryOp* unary) {
   // recursion
   Visit(unary->operand_);
@@ -190,10 +219,10 @@ void Generator::VisitUnaryOp(UnaryOp* unary) {
   ir::type *arg_scal_ty = arg_ty->get_scalar_ty();
   // return
   switch  (unary->op_) {
-    case Token::PREFIX_INC:  return error_not_implemented("prefix increment not implemented");
-    case Token::PREFIX_DEC:  return error_not_implemented("prefix decrement not implemented");
-    case Token::POSTFIX_INC: return error_not_implemented("postfix increment not implemented");
-    case Token::POSTFIX_DEC: return error_not_implemented("postfix decrement not implemented");
+    case Token::PREFIX_INC:  return set_ret(GenUnaryInc(unary, false, true));
+    case Token::PREFIX_DEC:  return set_ret(GenUnaryInc(unary, false, false));
+    case Token::POSTFIX_INC: return set_ret(GenUnaryInc(unary, true, true));
+    case Token::POSTFIX_DEC: return set_ret(GenUnaryInc(unary, true, false));
     case Token::ADDR:        return error_not_implemented("unary & not implemented");
     case Token::DEREF:       return set_ret(bld_->create_load(arg));
     case Token::PLUS:        return error_not_implemented("unary + not implemented");