Add looping over dictionaries and sets

integration_tests/loop_12.py

@@ -0,0 +1,67 @@
+
+def test_for_dict_int():
+    dict_int: dict[i32, i32] = {1:2, 2:3, 3:4}
+    key: i32
+    s1: i32 = 0
+    s2: i32 = 0
+
+    for key in dict_int:
+        print(key)
+        s1 += key
+        s2 += dict_int[key]
+
+    assert s1 == 6
+    assert s2 == 9
+
+def test_for_dict_str():
+    dict_str: dict[str, str] = {"a":"b", "c":"d"}
+    key: str
+    s1: str = ""
+    s2: str = ""
+
+    for key in dict_str:
+        print(key)
+        s1 += key
+        s2 += dict_str[key]
+
+    assert (s1 == "ac" or s1 == "ca")
+    assert ((s1 == "ac" and s2 == "bd") or (s1 == "ca" and s2 == "db"))
+
+def test_for_set_int():
+    set_int: set[i32] = {1, 2, 3}
+    el: i32
+    s: i32 = 0
+
+    for el in set_int:
+        print(el)
+        s += el
+
+    assert s == 6
+
+def test_for_set_str():
+    set_str: set[str] = {'a', 'b'}
+    el: str
+    s: str = ""
+
+    for el in set_str:
+        print(el)
+        s += el
+
+    assert (s == "ab" or s == "ba")
+
+def test_nested():
+    graph: dict[i32, set[i32]] = {1: {2, 3}}
+    el: i32
+    s: i32 = 0
+    for el in graph[1]:
+        print(el)
+        s += el
+
+    assert s == 5
+
+
+test_for_dict_int()
+test_for_set_int()
+test_for_dict_str()
+test_for_set_str()
+test_nested()

src/libasr/ASR.asdl

@@ -40,6 +40,7 @@ stmt
     | ErrorStop(expr? code)
     | Exit(identifier? stmt_name)
     | ForAllSingle(do_loop_head head, stmt assign_stmt)
+    | ForEach(expr var, expr container, stmt* body)
     | GoTo(int target_id, identifier name)
     | GoToTarget(int id, identifier name)
     | If(expr test, stmt* body, stmt* orelse)

src/libasr/codegen/asr_to_llvm.cpp

@@ -5783,6 +5783,189 @@ class ASRToLLVMVisitor : public ASR::BaseVisitor<ASRToLLVMVisitor>
         strings_to_be_deallocated.p = strings_to_be_deallocated_copy;
     }
 
+    void visit_ForEach(const ASR::ForEach_t &x) {
+        llvm::Value **strings_to_be_deallocated_copy = strings_to_be_deallocated.p;
+        size_t n = strings_to_be_deallocated.n;
+        strings_to_be_deallocated.reserve(al, 1);
+
+        int64_t ptr_loads_copy = ptr_loads;
+        ptr_loads = 0;
+        this->visit_expr(*x.m_container);
+        llvm::Value *pcontainer = tmp;
+        ptr_loads = 0;
+        this->visit_expr(*x.m_var);
+        llvm::Value *pvar = tmp;
+        ptr_loads = ptr_loads_copy;
+
+        if (ASR::is_a<ASR::Dict_t>(*ASRUtils::expr_type(x.m_container))) {
+            ASR::Dict_t *dict_type = ASR::down_cast<ASR::Dict_t>(
+                ASRUtils::expr_type(x.m_container));
+            ASR::ttype_t *key_type = dict_type->m_key_type;
+            llvm::Value *capacity = LLVM::CreateLoad(*builder,
+                llvm_utils->dict_api->get_pointer_to_capacity(pcontainer));
+            llvm::Value *key_mask = LLVM::CreateLoad(*builder, 
+                llvm_utils->dict_api->get_pointer_to_keymask(pcontainer));
+            llvm::Value *key_list = llvm_utils->dict_api->get_key_list(pcontainer);
+            llvm::AllocaInst *idx_ptr = builder->CreateAlloca(
+                llvm::Type::getInt32Ty(context), nullptr);
+            LLVM::CreateStore(*builder, llvm::ConstantInt::get(
+                llvm::Type::getInt32Ty(context), llvm::APInt(32, 0)), idx_ptr);
+
+            if (llvm_utils->dict_api == llvm_utils->dict_api_sc) {
+                llvm::Value *key_value_pairs = LLVM::CreateLoad(*builder, 
+                  llvm_utils->dict_api->get_pointer_to_key_value_pairs(pcontainer));
+                llvm::Type* kv_pair_type = 
+                    llvm_utils->dict_api->get_key_value_pair_type(key_type, dict_type->m_value_type);
+                llvm::AllocaInst *chain_itr = builder->CreateAlloca(
+                    llvm::Type::getInt8PtrTy(context), nullptr);
+
+                create_loop(nullptr, [=](){
+                    call_lcompilers_free_strings();
+                    return builder->CreateICmpSGT(capacity, LLVM::CreateLoad(*builder, idx_ptr));
+                }, [&](){
+                llvm::Value* idx = LLVM::CreateLoad(*builder, idx_ptr);
+                llvm::Value* key_mask_value = LLVM::CreateLoad(*builder,
+                    llvm_utils->create_ptr_gep(key_mask, idx));
+                llvm::Value* is_key_set = builder->CreateICmpEQ(key_mask_value,
+                    llvm::ConstantInt::get(llvm::Type::getInt8Ty(context), llvm::APInt(8, 1)));
+
+                llvm_utils->create_if_else(is_key_set, [&]() {
+                    llvm::Value* dict_i = llvm_utils->create_ptr_gep(key_value_pairs, idx);
+                    llvm::Value* kv_ll_i8 = builder->CreateBitCast(dict_i, llvm::Type::getInt8PtrTy(context));
+                    LLVM::CreateStore(*builder, kv_ll_i8, chain_itr);
+
+                    llvm::BasicBlock *loop2head = llvm::BasicBlock::Create(context, "loop2.head");
+                    llvm::BasicBlock *loop2body = llvm::BasicBlock::Create(context, "loop2.body");
+                    llvm::BasicBlock *loop2end = llvm::BasicBlock::Create(context, "loop2.end");
+
+                    // head
+                    llvm_utils->start_new_block(loop2head);
+                    {
+                        llvm::Value *cond = builder->CreateICmpNE(
+                            LLVM::CreateLoad(*builder, chain_itr),
+                            llvm::ConstantPointerNull::get(llvm::Type::getInt8PtrTy(context))
+                        );
+                        builder->CreateCondBr(cond, loop2body, loop2end);
+                    }
+
+                    // body
+                    llvm_utils->start_new_block(loop2body);
+                    {
+                        llvm::Value* kv_struct_i8 = LLVM::CreateLoad(*builder, chain_itr);
+                        llvm::Value* kv_struct = builder->CreateBitCast(kv_struct_i8, kv_pair_type->getPointerTo());
+                        llvm::Value* kv_el = llvm_utils->create_gep(kv_struct, 0);
+                        if( !LLVM::is_llvm_struct(key_type) ) {
+                            kv_el = LLVM::CreateLoad(*builder, kv_el);
+                        }
+                        LLVM::CreateStore(*builder, kv_el, pvar);
+                        for (size_t i=0; i<x.n_body; i++) {
+                            this->visit_stmt(*x.m_body[i]);
+                        }
+                        call_lcompilers_free_strings();
+                        llvm::Value* next_kv_struct = LLVM::CreateLoad(*builder, llvm_utils->create_gep(kv_struct, 2));
+                        LLVM::CreateStore(*builder, next_kv_struct, chain_itr);
+                    }
+
+                    builder->CreateBr(loop2head);
+
+                    // end
+                    llvm_utils->start_new_block(loop2end);
+                }, [=]() {
+                });
+                llvm::Value* tmp = builder->CreateAdd(idx,
+                            llvm::ConstantInt::get(context, llvm::APInt(32, 1)));
+                LLVM::CreateStore(*builder, tmp, idx_ptr);
+
+                });
+
+            } else {
+                create_loop(nullptr, [=](){
+                    call_lcompilers_free_strings();
+                    return builder->CreateICmpSGT(capacity, LLVM::CreateLoad(*builder, idx_ptr));
+                }, [&](){
+                llvm::Value *idx = LLVM::CreateLoad(*builder, idx_ptr);
+                llvm::Value *key_mask_value = LLVM::CreateLoad(*builder,
+                    llvm_utils->create_ptr_gep(key_mask, idx));
+                llvm::Value *is_key_skip = builder->CreateICmpEQ(key_mask_value,
+                    llvm::ConstantInt::get(llvm::Type::getInt8Ty(context),
+                    llvm::APInt(8, 3)));
+                llvm::Value *is_key_set = builder->CreateICmpNE(key_mask_value,
+                    llvm::ConstantInt::get(llvm::Type::getInt8Ty(context),
+                    llvm::APInt(8, 0)));
+
+                llvm::Value *el_exists = builder->CreateAnd(is_key_set,
+                    builder->CreateNot(is_key_skip));
+
+                llvm_utils->create_if_else(el_exists, [&]() {
+                    LLVM::CreateStore(*builder, llvm_utils->list_api->read_item(key_list, idx,
+                        false, *module, LLVM::is_llvm_struct(key_type)), pvar);
+
+                    for (size_t i=0; i<x.n_body; i++) {
+                        this->visit_stmt(*x.m_body[i]);
+                    }
+                    call_lcompilers_free_strings();
+                }, [=](){});
+
+                idx = builder->CreateAdd(idx, 
+                    llvm::ConstantInt::get(context, llvm::APInt(32, 1)));
+                LLVM::CreateStore(*builder, idx, idx_ptr);
+                });
+            }
+        } else if (ASR::is_a<ASR::Set_t>(*ASRUtils::expr_type(x.m_container))) {
+            ASR::Set_t *set_type = ASR::down_cast<ASR::Set_t>(
+                ASRUtils::expr_type(x.m_container));
+            ASR::ttype_t *el_type = set_type->m_type;
+
+            llvm::AllocaInst *idx_ptr = builder->CreateAlloca(
+                llvm::Type::getInt32Ty(context), nullptr);
+            LLVM::CreateStore(*builder, llvm::ConstantInt::get(
+                llvm::Type::getInt32Ty(context), llvm::APInt(32, 0)), idx_ptr);
+
+            llvm::Value *capacity = LLVM::CreateLoad(*builder, 
+                llvm_utils->set_api->get_pointer_to_capacity(pcontainer));
+            llvm::Value *el_list = llvm_utils->set_api->get_el_list(pcontainer);
+            llvm::Value *el_mask = LLVM::CreateLoad(*builder,
+                llvm_utils->set_api->get_pointer_to_mask(pcontainer));
+
+            create_loop(nullptr, [=](){
+                call_lcompilers_free_strings();
+                return builder->CreateICmpSGT(capacity, LLVM::CreateLoad(*builder, idx_ptr));
+            }, [&](){
+            llvm::Value *idx = LLVM::CreateLoad(*builder, idx_ptr);
+            llvm::Value *el_mask_value = LLVM::CreateLoad(*builder,
+                llvm_utils->create_ptr_gep(el_mask, idx));
+            llvm::Value *is_el_skip = builder->CreateICmpEQ(el_mask_value,
+                llvm::ConstantInt::get(llvm::Type::getInt8Ty(context),
+                llvm::APInt(8, 3)));
+            llvm::Value *is_el_set = builder->CreateICmpNE(el_mask_value,
+                llvm::ConstantInt::get(llvm::Type::getInt8Ty(context),
+                llvm::APInt(8, 0)));
+
+            llvm::Value *el_exists = builder->CreateAnd(is_el_set,
+                builder->CreateNot(is_el_skip));
+
+            llvm_utils->create_if_else(el_exists, [&]() {
+                LLVM::CreateStore(*builder, llvm_utils->list_api->read_item(el_list, idx,
+                    false, *module, LLVM::is_llvm_struct(el_type)), pvar);
+
+                for (size_t i=0; i<x.n_body; i++) {
+                    this->visit_stmt(*x.m_body[i]);
+                }
+                call_lcompilers_free_strings();
+            }, [=](){});
+
+            idx = builder->CreateAdd(idx, 
+                llvm::ConstantInt::get(context, llvm::APInt(32, 1)));
+            LLVM::CreateStore(*builder, idx, idx_ptr);
+            });
+        } else {
+            throw CodeGenError("Only sets and dictionaries are supported with this loop for now.");
+        }
+        strings_to_be_deallocated.reserve(al, n);
+        strings_to_be_deallocated.n = n;
+        strings_to_be_deallocated.p = strings_to_be_deallocated_copy;
+    }
+
     bool case_insensitive_string_compare(const std::string& str1, const std::string& str2) {
         if (str1.size() != str2.size()) {
             return false;

src/libasr/codegen/llvm_utils.cpp

@@ -2075,6 +2075,11 @@ namespace LCompilers {
         return get_key_value_pair_type(key_type_code, value_type_code);
     }
 
+    llvm::Type* LLVMDict::get_key_value_pair_type(
+        ASR::ttype_t* /*key_asr_type*/, ASR::ttype_t* /*value_asr_type*/) {
+        return nullptr;
+    }
+
     llvm::Type* LLVMDictSeparateChaining::get_dict_type(
         std::string key_type_code, std::string value_type_code,
         int32_t key_type_size, int32_t value_type_size,
@@ -2156,6 +2161,10 @@ namespace LCompilers {
         return llvm_utils->create_gep(dict, 1);
     }
 
+    llvm::Value* LLVMDict::get_pointer_to_key_value_pairs(llvm::Value* /*dict*/) {
+        return nullptr;
+    }
+
     llvm::Value* LLVMDictSeparateChaining::get_pointer_to_key_value_pairs(llvm::Value* dict) {
         return llvm_utils->create_gep(dict, 3);
     }

src/libasr/codegen/llvm_utils.h

@@ -567,6 +567,9 @@ namespace LCompilers {
             virtual
             llvm::Value* get_pointer_to_occupancy(llvm::Value* dict) = 0;
 
+            virtual
+            llvm::Value* get_pointer_to_keymask(llvm::Value* dict) = 0;
+
             virtual
             llvm::Value* get_pointer_to_capacity(llvm::Value* dict) = 0;
 
@@ -651,6 +654,13 @@ namespace LCompilers {
                 std::map<std::string, std::map<std::string, int>>& name2memidx,
                 bool key_or_value) = 0;
 
+            virtual
+            llvm::Type* get_key_value_pair_type(ASR::ttype_t* key_asr_type, ASR::ttype_t* value_pair_type) = 0;
+
+            virtual
+            llvm::Value* get_pointer_to_key_value_pairs(llvm::Value* dict) = 0;
+
+
             virtual ~LLVMDictInterface() = 0;
 
     };
@@ -744,6 +754,10 @@ namespace LCompilers {
                 std::map<std::string, std::map<std::string, int>>& name2memidx,
                 bool key_or_value);
 
+            llvm::Type* get_key_value_pair_type(ASR::ttype_t* key_asr_type, ASR::ttype_t* value_pair_type);
+
+            llvm::Value* get_pointer_to_key_value_pairs(llvm::Value* dict);
+
             virtual ~LLVMDict();
     };
 
@@ -791,8 +805,6 @@ namespace LCompilers {
 
             llvm::Value* get_pointer_to_number_of_filled_buckets(llvm::Value* dict);
 
-            llvm::Value* get_pointer_to_key_value_pairs(llvm::Value* dict);
-
             llvm::Value* get_pointer_to_rehash_flag(llvm::Value* dict);
 
             void deepcopy_key_value_pair_linked_list(llvm::Value* srci, llvm::Value* desti,
@@ -810,8 +822,6 @@ namespace LCompilers {
 
             llvm::Type* get_key_value_pair_type(std::string key_type_code, std::string value_type_code);
 
-            llvm::Type* get_key_value_pair_type(ASR::ttype_t* key_asr_type, ASR::ttype_t* value_pair_type);
-
             void dict_init_given_initial_capacity(std::string key_type_code, std::string value_type_code,
                 llvm::Value* dict, llvm::Module* module, llvm::Value* initial_capacity);
 
@@ -892,6 +902,10 @@ namespace LCompilers {
                 std::map<std::string, std::map<std::string, int>>& name2memidx,
                 bool key_or_value);
 
+            llvm::Type* get_key_value_pair_type(ASR::ttype_t* key_asr_type, ASR::ttype_t* value_pair_type);
+
+            llvm::Value* get_pointer_to_key_value_pairs(llvm::Value* dict);
+
             virtual ~LLVMDictSeparateChaining();
 
     };
@@ -939,6 +953,9 @@ namespace LCompilers {
             virtual
             llvm::Value* get_pointer_to_capacity(llvm::Value* set) = 0;
 
+            virtual
+            llvm::Value* get_pointer_to_mask(llvm::Value* set) = 0;
+
             llvm::Value* get_el_hash(llvm::Value* capacity, llvm::Value* el,
                 ASR::ttype_t* el_asr_type, llvm::Module& module);