Handle Op-Assign (<<=, +=, etc.) codegen in PyXASM

also, add the test for bit reverse in qjit kernel
Signed-off-by: Thien Nguyen <nguyentm@ornl.gov>

handlers/token_collector/pyxasm/pyxasm_visitor.hpp

@@ -457,7 +457,17 @@ class pyxasm_visitor : public pyxasmBaseVisitor {
         return 0;
       }
     } else {
-      return visitChildren(ctx);
+      // Visit child node:
+      auto child_result = visitChildren(ctx);
+      const auto translated_src = sub_node_translation.str();
+      sub_node_translation.str(std::string());
+      // If no child nodes, perform the codegen (result.first is not set)
+      // but just appending the incremental translation collector;
+      // return the collected C++ statement.
+      if (result.first.empty() && !translated_src.empty()) {
+        result.first = translated_src + ";\n";
+      }
+      return child_result;
     }
   }
 
@@ -509,6 +519,31 @@ class pyxasm_visitor : public pyxasmBaseVisitor {
     return 0;
   }
 
+  virtual antlrcpp::Any visitTestlist_star_expr(
+      pyxasmParser::Testlist_star_exprContext *context) override {
+    // std::cout << "Testlist_star_exprContext:" << context->getText() << "\n";
+    const auto var_name = context->getText();
+    if (xacc::container::contains(declared_var_names, var_name)) {
+      sub_node_translation << var_name << " ";
+      return 0;
+    }
+    return visitChildren(context);
+  }
+
+  virtual antlrcpp::Any
+  visitAugassign(pyxasmParser::AugassignContext *context) override {
+    // std::cout << "Augassign:" << context->getText() << "\n";
+    sub_node_translation << context->getText() << " ";
+    return 0;
+  }
+
+  virtual antlrcpp::Any
+  visitTestlist(pyxasmParser::TestlistContext *context) override {
+    // std::cout << "visitTestlist:" << context->getText() << "\n";
+    sub_node_translation << context->getText() << " ";
+    return 0;
+  }
+
  private:
   // Replaces common Python constants, e.g. 'math.pi' or 'numpy.pi'.
   // Note: the library names have been resolved to their original names.

python/tests/test_jit_simple.py

@@ -231,6 +231,35 @@ class TestKernelJIT(unittest.TestCase):
         # entangled
         self.assertEqual(q.counts()["0"], r.counts()["0"])
         self.assertEqual(q.counts()["1"], r.counts()["1"])
+    
+    def test_while_loop(self):
+        @qjit 
+        def while_loop_kernel(q: qubit, x: int, exp_inv: int):
+            rev = 0
+            while x:
+                rev <<= 1
+                rev += x & 1
+                x >>= 1
+
+            if rev == exp_inv:
+                print("Success")
+                X(q)
+
+        # Reference: pure python to check
+        def reverse_bit(num):
+            result = 0
+            while num:
+                result = (result << 1) + (num & 1)
+                num >>= 1
+            return result
+        
+        q = qalloc(1)
+        test_val = 123
+        expected = reverse_bit(test_val)
+        comp0 = while_loop_kernel.extract_composite(q[0], test_val, expected)    
+        print("Comp:", comp0)
+        # Has X applied.
+        self.assertEqual(comp0.nInstructions(), 1) 
 
 if __name__ == '__main__':
   unittest.main()
\ No newline at end of file