Fix while assignments tabs

progs/test.py

@@ -1,13 +1,8 @@
 n = int(input())
 x = int(input())
-y = int(input())
 m = 1
-
-if x == 10:
-    x = 1
-while n > 2 * x - 3 * y + 5:
+while n > 1:
     tmp = 2 * x
     m = m * tmp
     n = n - 1
-
 print(m + tmp)

src/ast/Python3VisitorImpl.java

@@ -338,6 +338,8 @@ public Node visitWhile_stmt(While_stmtContext ctx) {
             return whileStmt;
         }
 
+        rewriter.insertAfter(ctx.COLON(0).getSymbol().getTokenIndex(), "\n");
+
         int lineStart = ctx.getStart().getLine();
         int lineStop = ctx.getStop().getLine();
         int index = ctx.getStart().getTokenIndex();
@@ -429,6 +431,10 @@ private void optimizeBlock(BlockNode block, int lineStart, int lineStop, int ind
                     // int lastToken = ctx.expr().expr(counter).getStop().getTokenIndex();
                     // int firstToken = ctx.expr().expr(counter).getStart().getTokenIndex();
                     // rewriter.replace(firstToken, lastToken, newVar);
+                    System.out.println("1 " + assignment.toPrint(""));
+                } else {
+
+                    rewriter.insertBefore(assignment.getLhrIndex(), "\t");
                 }
             }
         }

test/2a.py

@@ -1,9 +1,9 @@
-n = 1
-x = 2
-y = 3
+n = int(input())
+x = int(input())
 m = 1
-while n < (2 * x - 3 * y + 5):
-    m = m + n
-    n = n + 1
-
-print(m)
+tmp = 0
+while n > 1:
+    m = m * tmp
+    tmp = 2 * x
+    n = n - 1
+print(m + tmp)

test/2b.py

@@ -0,0 +1,9 @@
+n = 1
+x = 2
+y = 3
+m = 1
+while n < 2 * x - 3 * y + 5:
+    m = m + n
+    n = n + 1
+
+print(m)