Handle statements in the condition of while

NEWS

@@ -50,6 +50,7 @@ Changes from 0.13.0
    * Multiple expressions are now allowed in the else clause of
      a for loop
    * `else` clauses in `for` and `while` are recognized more reliably
+   * Statements in the condition of a `while` loop are repeated properly
    * Argument destructuring no longer interferes with function docstrings.
 
    [ Misc. Improvements ]

hy/compiler.py

@@ -1861,12 +1861,33 @@ class HyASTCompiler(object):
     @checkargs(min=2)
     def compile_while_expression(self, expr):
         expr.pop(0)  # "while"
-        ret = self.compile(expr.pop(0))
+        cond = expr.pop(0)
+        cond_compiled = self.compile(cond)
 
-        orel = Result()
+        else_expr = None
-        # (while cond body (else …))
         if ends_with_else(expr):
             else_expr = expr.pop()
+
+        if cond_compiled.stmts:
+            # We need to ensure the statements for the condition are
+            # executed on every iteration. Rewrite the loop to use a
+            # single anonymous variable as the condition.
+            def e(*x): return HyExpression(x)
+            s = HySymbol
+            cond_var = s(self.get_anon_var())
+            return self.compile(e(
+                s('do'),
+                e(s('setv'), cond_var, 1),
+                e(s('while'), cond_var,
+                  # Cast the condition to a bool in case it's mutable and
+                  # changes its truth value, but use (not (not ...)) instead of
+                  # `bool` in case `bool` has been redefined.
+                  e(s('setv'), cond_var, e(s('not'), e(s('not'), cond))),
+                  e(s('if*'), cond_var, e(s('do'), *expr)),
+                  *([else_expr] if else_expr is not None else []))).replace(expr))  # noqa
+
+        orel = Result()
+        if else_expr is not None:
             for else_body in else_expr[1:]:
                 orel += self.compile(else_body)
                 orel += orel.expr_as_stmt()
@@ -1874,11 +1895,9 @@ class HyASTCompiler(object):
         body = self._compile_branch(expr)
         body += body.expr_as_stmt()
 
-        ret += asty.While(expr,
+        ret = cond_compiled + asty.While(
-                          test=ret.force_expr,
+            expr, test=cond_compiled.force_expr,
-                          body=body.stmts,
+            body=body.stmts, orelse=orel.stmts)
-                          orelse=orel.stmts)
-
         ret.contains_yield = body.contains_yield
 
         return ret

tests/native_tests/language.hy

@@ -346,6 +346,68 @@
   (assert (= a [2 "e"])))
 
 
+(defn test-while-multistatement-condition []
+
+  ; The condition should be executed every iteration, before the body.
+  ; `else` should be executed last.
+  (setv s "")
+  (setv x 2)
+  (while (do (+= s "a") x)
+    (+= s "b")
+    (-= x 1)
+    (else
+      (+= s "z")))
+  (assert (= s "ababaz"))
+
+  ; `else` should still be skipped after `break`.
+  (setv s "")
+  (setv x 2)
+  (while (do (+= s "a") x)
+    (+= s "b")
+    (-= x 1)
+    (when (= x 0)
+      (break))
+    (else
+      (+= s "z")))
+  (assert (= s "abab"))
+
+  ; `continue` should jump to the condition.
+  (setv s "")
+  (setv x 2)
+  (setv continued? False)
+  (while (do (+= s "a") x)
+    (+= s "b")
+    (when (and (= x 1) (not continued?))
+      (+= s "c")
+      (setv continued? True)
+      (continue))
+    (-= x 1)
+    (else
+      (+= s "z")))
+  (assert (= s "ababcabaz"))
+
+  ; `break` in a condition applies to the `while`, not an outer loop.
+  (setv s "")
+  (for [x "123"]
+    (+= s x)
+    (setv y 0)
+    (while (do (when (and (= x "2") (= y 1)) (break)) (< y 3))
+      (+= s "y")
+      (+= y 1)))
+  (assert (= s "1yyy2y3yyy"))
+
+  ; The condition is still tested appropriately if its last variable
+  ; is set to a false value in the loop body.
+  (setv out [])
+  (setv x 0)
+  (setv a [1 1])
+  (while (do (.append out 2) (setv x (and a (.pop a))) x)
+    (setv x 0)
+    (.append out x))
+  (assert (= out [2 0 2 0 2]))
+  (assert (is x a)))
+
+
 (defn test-branching []
   "NATIVE: test if branching"
   (if True