Merge pull request lcompilers#775 from Shaikh-Ubaid/for_in_strings

Supporting For in looping over strings

integration_tests/CMakeLists.txt

@@ -140,6 +140,7 @@ RUN(NAME expr_09             LABELS cpython llvm)
 RUN(NAME expr_10             LABELS cpython llvm)
 RUN(NAME expr_11             LABELS cpython llvm c)
 RUN(NAME expr_12             LABELS llvm c)
+RUN(NAME loop_01             LABELS cpython llvm)
 RUN(NAME test_types_01       LABELS cpython llvm)
 RUN(NAME test_str_01         LABELS cpython llvm)
 RUN(NAME test_str_02         LABELS cpython llvm)

integration_tests/loop_01.py

@@ -0,0 +1,8 @@
+def main0():
+    s: str
+    s = 'aabbcc'
+    c: str
+    for c in s:
+        print(c)
+
+main0()

src/lpython/semantics/python_ast_to_asr.cpp

@@ -2684,8 +2684,7 @@ class BodyVisitor : public CommonVisitor<BodyVisitor> {
         this->visit_expr(*x.m_target);
         ASR::expr_t *target=ASRUtils::EXPR(tmp);
         Vec<ASR::stmt_t*> body;
-        body.reserve(al, x.n_body);
-        transform_stmts(body, x.n_body, x.m_body);
+        bool is_explicit_iterator_required = false;
         ASR::expr_t *loop_end = nullptr, *loop_start = nullptr, *inc = nullptr;
         if (AST::is_a<AST::Call_t>(*x.m_iter)) {
             AST::Call_t *c = AST::down_cast<AST::Call_t>(x.m_iter);
@@ -2723,6 +2722,44 @@ class BodyVisitor : public CommonVisitor<BodyVisitor> {
                     x.base.base.loc);
             }
 
+        } else if (AST::is_a<AST::Name_t>(*x.m_iter)) {
+            std::string loop_src_var_name = AST::down_cast<AST::Name_t>(x.m_iter)->m_id;
+            auto loop_src_var_symbol = current_scope->get_symbol(loop_src_var_name);
+            auto loop_src_var_ttype = ASRUtils::symbol_type(loop_src_var_symbol);
+            if (ASR::is_a<ASR::Character_t>(*loop_src_var_ttype)) {
+                auto int_type = ASR::make_Integer_t(al, x.base.base.loc, 4, nullptr, 0);
+
+                {
+                    // create a new variable called/named __explicit_iterator of type i32 and add it to symbol table
+                    std::string explicit_iter_name = current_scope->get_unique_name("__explicit_iterator");
+                    auto explicit_iter_variable = ASR::make_Variable_t(al, x.base.base.loc, current_scope,
+                        s2c(al, explicit_iter_name), ASR::intentType::Local, nullptr, nullptr, ASR::storage_typeType::Default,
+                        ASRUtils::TYPE(int_type), ASR::abiType::Source, ASR::accessType::Public, ASR::presenceType::Required, false
+                    );
+
+                    current_scope->add_symbol(explicit_iter_name, ASR::down_cast<ASR::symbol_t>(explicit_iter_variable));
+                }
+
+                {
+                    // make loop_end = len(loop_src_var), where loop_src_var is the variable over which
+                    // we are iterating the for in loop
+                    auto loop_src_var = ASR::make_Var_t(al, x.base.base.loc, loop_src_var_symbol);
+                    auto call_str_len = ASR::make_StringLen_t(al, x.base.base.loc, ASRUtils::EXPR(loop_src_var), ASRUtils::TYPE(int_type), nullptr);
+                    loop_end = ASRUtils::EXPR(call_str_len);
+                }
+
+                is_explicit_iterator_required = true;
+            } else if (ASR::is_a<ASR::List_t>(*loop_src_var_ttype)) {
+                throw SemanticError("Iterating on Lists using for in loop not yet supported as "
+                    "visit_Len() is not yet supported in the LLVM Backend", x.base.base.loc);
+            } else if (ASR::is_a<ASR::Set_t>(*loop_src_var_ttype)) {
+                throw SemanticError("Iterating on Set using for in loop not yet supported", x.base.base.loc);
+            } else if (ASR::is_a<ASR::Tuple_t>(*loop_src_var_ttype)) {
+                throw SemanticError("Iterating on Tuple using for in loop not yet supported", x.base.base.loc);
+            } else {
+                throw SemanticError("Only Strings, Lists, Sets and Tuples can be used with for in loop, not " +
+                    ASRUtils::type_to_str(loop_src_var_ttype), x.base.base.loc);
+            }
         } else {
             throw SemanticError("Only function call `range(..)` supported as for loop iteration for now",
                 x.base.base.loc);
@@ -2736,7 +2773,28 @@ class BodyVisitor : public CommonVisitor<BodyVisitor> {
                             x.base.base.loc, false);
         loop_end = ASRUtils::EXPR(tmp);
         ASR::do_loop_head_t head;
-        head.m_v = target;
+
+        if(is_explicit_iterator_required) {
+            body.reserve(al, x.n_body + 1);
+            // add an assignment instruction to body to assign value of loop_src_var at an index to the loop_target_var
+            auto explicit_iter_var = ASR::make_Var_t(al, x.base.base.loc, current_scope->get_symbol("__explicit_iterator"));
+            auto index_plus_one = ASR::make_IntegerBinOp_t(al, x.base.base.loc, ASRUtils::EXPR(explicit_iter_var),
+                ASR::binopType::Add, constant_one, a_type, nullptr);
+            std::string loop_src_var_name = AST::down_cast<AST::Name_t>(x.m_iter)->m_id;
+            auto loop_src_var = ASR::make_Var_t(al, x.base.base.loc, current_scope->get_symbol(loop_src_var_name));
+            auto loop_src_var_element = ASR::make_StringItem_t(al, x.base.base.loc, ASRUtils::EXPR(loop_src_var),
+                               ASRUtils::EXPR(index_plus_one), a_type, nullptr);
+            auto loop_target_assignment = ASR::make_Assignment_t(al, x.base.base.loc, target, ASRUtils::EXPR(loop_src_var_element), nullptr);
+            body.push_back(al, ASRUtils::STMT(loop_target_assignment));
+
+            head.m_v = ASRUtils::EXPR(explicit_iter_var);
+        } else {
+            body.reserve(al, x.n_body);
+            head.m_v = target;
+        }
+
+        transform_stmts(body, x.n_body, x.m_body);
+
         if (loop_start) {
             head.m_start = loop_start;
         } else {