From b423649d7b7c8c79470d480ea1b657e50b787e9c Mon Sep 17 00:00:00 2001
From: ZhuangYumin <zhuangyumin@sjtu.edu.cn>
Date: Thu, 9 Nov 2023 18:44:24 +0800
Subject: [PATCH] upd: write loop

---
 src/Evalvisitor.cpp | 29 +++++++++++++++++++++++------
 1 file changed, 23 insertions(+), 6 deletions(-)

diff --git a/src/Evalvisitor.cpp b/src/Evalvisitor.cpp
index 7598edf..9fd6de4 100644
--- a/src/Evalvisitor.cpp
+++ b/src/Evalvisitor.cpp
@@ -187,22 +187,39 @@ std::any EvalVisitor::visitReturn_stmt(Python3Parser::Return_stmtContext *ctx) {
 
 std::any EvalVisitor::visitCompound_stmt(
     Python3Parser::Compound_stmtContext *ctx) {
-  // TODO
   return visitChildren(ctx);
 }
 
 std::any EvalVisitor::visitIf_stmt(Python3Parser::If_stmtContext *ctx) {
-  // TODO
-  return visitChildren(ctx);
+  auto test_list = ctx->test();
+  auto suite_list = ctx->suite();
+  if (test_list.size() - suite_list.size() != 1 &&
+      test_list.size() - suite_list.size() != 0)
+    throw InterpretException("test and suite doesn't match in If_stmt");
+  for (int i = 0; i < test_list.size(); i++)
+    if (Any2Bool(DeQuate(visitTest(test_list[i]), Variables)))
+      return visitSuite(suite_list[i]);
+  if (test_list.size() - suite_list.size() == 1)
+    return visitSuite(suite_list[suite_list.size() - 1]);
+  return nullptr;
 }
 
 std::any EvalVisitor::visitWhile_stmt(Python3Parser::While_stmtContext *ctx) {
-  // TODO
-  return visitChildren(ctx);
+  auto test = ctx->test();
+  auto suite = ctx->suite();
+  while (Any2Bool(DeQuate(visitTest(test), Variables))) {
+    auto res = visitSuite(suite);
+    FlowType *res_ptr = std::any_cast<FlowType>(&res);
+    if (res_ptr) {
+      if (res_ptr->Status == BREAK) break;
+      if (res_ptr->Status == CONTINUE) continue;
+      if (res_ptr->Status == RETURN) return *res_ptr;
+    }
+  }
+  return nullptr;
 }
 
 std::any EvalVisitor::visitSuite(Python3Parser::SuiteContext *ctx) {
-  // TODO
   return visitChildren(ctx);
 }
 std::any EvalVisitor::visitTestlist(Python3Parser::TestlistContext *ctx) {