Skip to content

Commit 44f9d26

Browse files
advikkabraadvikkabra
authored
Add looping over dictionaries and sets (#2710)
* Add looping over dictionaries and sets * Add looping over dictionaries and sets * Add support for subscripted items * Update * Update references * Fix scoping issue
1 parent 546777e commit 44f9d26

File tree

226 files changed

+540
-24
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

226 files changed

+540
-24
lines changed

‎integration_tests/loop_12.py

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
2+
def test_for_dict_int():
3+
dict_int: dict[i32, i32] = {1:2, 2:3, 3:4}
4+
key: i32
5+
s1: i32 = 0
6+
s2: i32 = 0
7+
8+
for key in dict_int:
9+
print(key)
10+
s1 += key
11+
s2 += dict_int[key]
12+
13+
assert s1 == 6
14+
assert s2 == 9
15+
16+
def test_for_dict_str():
17+
dict_str: dict[str, str] = {"a":"b", "c":"d"}
18+
key: str
19+
s1: str = ""
20+
s2: str = ""
21+
22+
for key in dict_str:
23+
print(key)
24+
s1 += key
25+
s2 += dict_str[key]
26+
27+
assert (s1 == "ac" or s1 == "ca")
28+
assert ((s1 == "ac" and s2 == "bd") or (s1 == "ca" and s2 == "db"))
29+
30+
def test_for_set_int():
31+
set_int: set[i32] = {1, 2, 3}
32+
el: i32
33+
s: i32 = 0
34+
35+
for el in set_int:
36+
print(el)
37+
s += el
38+
39+
assert s == 6
40+
41+
def test_for_set_str():
42+
set_str: set[str] = {'a', 'b'}
43+
el: str
44+
s: str = ""
45+
46+
for el in set_str:
47+
print(el)
48+
s += el
49+
50+
assert (s == "ab" or s == "ba")
51+
52+
def test_nested():
53+
graph: dict[i32, set[i32]] = {1: {2, 3}}
54+
el: i32
55+
s: i32 = 0
56+
for el in graph[1]:
57+
print(el)
58+
s += el
59+
60+
assert s == 5
61+
62+
63+
test_for_dict_int()
64+
test_for_set_int()
65+
test_for_dict_str()
66+
test_for_set_str()
67+
test_nested()

‎src/libasr/ASR.asdl

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ stmt
4040
| ErrorStop(expr? code)
4141
| Exit(identifier? stmt_name)
4242
| ForAllSingle(do_loop_head head, stmt assign_stmt)
43+
| ForEach(expr var, expr container, stmt* body)
4344
| GoTo(int target_id, identifier name)
4445
| GoToTarget(int id, identifier name)
4546
| If(expr test, stmt* body, stmt* orelse)

‎src/libasr/codegen/asr_to_llvm.cpp

Lines changed: 183 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5828,6 +5828,189 @@ class ASRToLLVMVisitor : public ASR::BaseVisitor<ASRToLLVMVisitor>
58285828
strings_to_be_deallocated.p = strings_to_be_deallocated_copy;
58295829
}
58305830

5831+
void visit_ForEach(const ASR::ForEach_t &x) {
5832+
llvm::Value **strings_to_be_deallocated_copy = strings_to_be_deallocated.p;
5833+
size_t n = strings_to_be_deallocated.n;
5834+
strings_to_be_deallocated.reserve(al, 1);
5835+
5836+
int64_t ptr_loads_copy = ptr_loads;
5837+
ptr_loads = 0;
5838+
this->visit_expr(*x.m_container);
5839+
llvm::Value *pcontainer = tmp;
5840+
ptr_loads = 0;
5841+
this->visit_expr(*x.m_var);
5842+
llvm::Value *pvar = tmp;
5843+
ptr_loads = ptr_loads_copy;
5844+
5845+
if (ASR::is_a<ASR::Dict_t>(*ASRUtils::expr_type(x.m_container))) {
5846+
ASR::Dict_t *dict_type = ASR::down_cast<ASR::Dict_t>(
5847+
ASRUtils::expr_type(x.m_container));
5848+
ASR::ttype_t *key_type = dict_type->m_key_type;
5849+
llvm::Value *capacity = LLVM::CreateLoad(*builder,
5850+
llvm_utils->dict_api->get_pointer_to_capacity(pcontainer));
5851+
llvm::Value *key_mask = LLVM::CreateLoad(*builder,
5852+
llvm_utils->dict_api->get_pointer_to_keymask(pcontainer));
5853+
llvm::Value *key_list = llvm_utils->dict_api->get_key_list(pcontainer);
5854+
llvm::AllocaInst *idx_ptr = builder->CreateAlloca(
5855+
llvm::Type::getInt32Ty(context), nullptr);
5856+
LLVM::CreateStore(*builder, llvm::ConstantInt::get(
5857+
llvm::Type::getInt32Ty(context), llvm::APInt(32, 0)), idx_ptr);
5858+
5859+
if (llvm_utils->dict_api == llvm_utils->dict_api_sc) {
5860+
llvm::Value *key_value_pairs = LLVM::CreateLoad(*builder,
5861+
llvm_utils->dict_api->get_pointer_to_key_value_pairs(pcontainer));
5862+
llvm::Type* kv_pair_type =
5863+
llvm_utils->dict_api->get_key_value_pair_type(key_type, dict_type->m_value_type);
5864+
llvm::AllocaInst *chain_itr = builder->CreateAlloca(
5865+
llvm::Type::getInt8PtrTy(context), nullptr);
5866+
5867+
create_loop(nullptr, [=](){
5868+
call_lcompilers_free_strings();
5869+
return builder->CreateICmpSGT(capacity, LLVM::CreateLoad(*builder, idx_ptr));
5870+
}, [&](){
5871+
llvm::Value* idx = LLVM::CreateLoad(*builder, idx_ptr);
5872+
llvm::Value* key_mask_value = LLVM::CreateLoad(*builder,
5873+
llvm_utils->create_ptr_gep(key_mask, idx));
5874+
llvm::Value* is_key_set = builder->CreateICmpEQ(key_mask_value,
5875+
llvm::ConstantInt::get(llvm::Type::getInt8Ty(context), llvm::APInt(8, 1)));
5876+
5877+
llvm_utils->create_if_else(is_key_set, [&]() {
5878+
llvm::Value* dict_i = llvm_utils->create_ptr_gep(key_value_pairs, idx);
5879+
llvm::Value* kv_ll_i8 = builder->CreateBitCast(dict_i, llvm::Type::getInt8PtrTy(context));
5880+
LLVM::CreateStore(*builder, kv_ll_i8, chain_itr);
5881+
5882+
llvm::BasicBlock *loop2head = llvm::BasicBlock::Create(context, "loop2.head");
5883+
llvm::BasicBlock *loop2body = llvm::BasicBlock::Create(context, "loop2.body");
5884+
llvm::BasicBlock *loop2end = llvm::BasicBlock::Create(context, "loop2.end");
5885+
5886+
// head
5887+
llvm_utils->start_new_block(loop2head);
5888+
{
5889+
llvm::Value *cond = builder->CreateICmpNE(
5890+
LLVM::CreateLoad(*builder, chain_itr),
5891+
llvm::ConstantPointerNull::get(llvm::Type::getInt8PtrTy(context))
5892+
);
5893+
builder->CreateCondBr(cond, loop2body, loop2end);
5894+
}
5895+
5896+
// body
5897+
llvm_utils->start_new_block(loop2body);
5898+
{
5899+
llvm::Value* kv_struct_i8 = LLVM::CreateLoad(*builder, chain_itr);
5900+
llvm::Value* kv_struct = builder->CreateBitCast(kv_struct_i8, kv_pair_type->getPointerTo());
5901+
llvm::Value* kv_el = llvm_utils->create_gep(kv_struct, 0);
5902+
if( !LLVM::is_llvm_struct(key_type) ) {
5903+
kv_el = LLVM::CreateLoad(*builder, kv_el);
5904+
}
5905+
LLVM::CreateStore(*builder, kv_el, pvar);
5906+
for (size_t i=0; i<x.n_body; i++) {
5907+
this->visit_stmt(*x.m_body[i]);
5908+
}
5909+
call_lcompilers_free_strings();
5910+
llvm::Value* next_kv_struct = LLVM::CreateLoad(*builder, llvm_utils->create_gep(kv_struct, 2));
5911+
LLVM::CreateStore(*builder, next_kv_struct, chain_itr);
5912+
}
5913+
5914+
builder->CreateBr(loop2head);
5915+
5916+
// end
5917+
llvm_utils->start_new_block(loop2end);
5918+
}, [=]() {
5919+
});
5920+
llvm::Value* tmp = builder->CreateAdd(idx,
5921+
llvm::ConstantInt::get(context, llvm::APInt(32, 1)));
5922+
LLVM::CreateStore(*builder, tmp, idx_ptr);
5923+
5924+
});
5925+
5926+
} else {
5927+
create_loop(nullptr, [=](){
5928+
call_lcompilers_free_strings();
5929+
return builder->CreateICmpSGT(capacity, LLVM::CreateLoad(*builder, idx_ptr));
5930+
}, [&](){
5931+
llvm::Value *idx = LLVM::CreateLoad(*builder, idx_ptr);
5932+
llvm::Value *key_mask_value = LLVM::CreateLoad(*builder,
5933+
llvm_utils->create_ptr_gep(key_mask, idx));
5934+
llvm::Value *is_key_skip = builder->CreateICmpEQ(key_mask_value,
5935+
llvm::ConstantInt::get(llvm::Type::getInt8Ty(context),
5936+
llvm::APInt(8, 3)));
5937+
llvm::Value *is_key_set = builder->CreateICmpNE(key_mask_value,
5938+
llvm::ConstantInt::get(llvm::Type::getInt8Ty(context),
5939+
llvm::APInt(8, 0)));
5940+
5941+
llvm::Value *el_exists = builder->CreateAnd(is_key_set,
5942+
builder->CreateNot(is_key_skip));
5943+
5944+
llvm_utils->create_if_else(el_exists, [&]() {
5945+
LLVM::CreateStore(*builder, llvm_utils->list_api->read_item(key_list, idx,
5946+
false, *module, LLVM::is_llvm_struct(key_type)), pvar);
5947+
5948+
for (size_t i=0; i<x.n_body; i++) {
5949+
this->visit_stmt(*x.m_body[i]);
5950+
}
5951+
call_lcompilers_free_strings();
5952+
}, [=](){});
5953+
5954+
idx = builder->CreateAdd(idx,
5955+
llvm::ConstantInt::get(context, llvm::APInt(32, 1)));
5956+
LLVM::CreateStore(*builder, idx, idx_ptr);
5957+
});
5958+
}
5959+
} else if (ASR::is_a<ASR::Set_t>(*ASRUtils::expr_type(x.m_container))) {
5960+
ASR::Set_t *set_type = ASR::down_cast<ASR::Set_t>(
5961+
ASRUtils::expr_type(x.m_container));
5962+
ASR::ttype_t *el_type = set_type->m_type;
5963+
5964+
llvm::AllocaInst *idx_ptr = builder->CreateAlloca(
5965+
llvm::Type::getInt32Ty(context), nullptr);
5966+
LLVM::CreateStore(*builder, llvm::ConstantInt::get(
5967+
llvm::Type::getInt32Ty(context), llvm::APInt(32, 0)), idx_ptr);
5968+
5969+
llvm::Value *capacity = LLVM::CreateLoad(*builder,
5970+
llvm_utils->set_api->get_pointer_to_capacity(pcontainer));
5971+
llvm::Value *el_list = llvm_utils->set_api->get_el_list(pcontainer);
5972+
llvm::Value *el_mask = LLVM::CreateLoad(*builder,
5973+
llvm_utils->set_api->get_pointer_to_mask(pcontainer));
5974+
5975+
create_loop(nullptr, [=](){
5976+
call_lcompilers_free_strings();
5977+
return builder->CreateICmpSGT(capacity, LLVM::CreateLoad(*builder, idx_ptr));
5978+
}, [&](){
5979+
llvm::Value *idx = LLVM::CreateLoad(*builder, idx_ptr);
5980+
llvm::Value *el_mask_value = LLVM::CreateLoad(*builder,
5981+
llvm_utils->create_ptr_gep(el_mask, idx));
5982+
llvm::Value *is_el_skip = builder->CreateICmpEQ(el_mask_value,
5983+
llvm::ConstantInt::get(llvm::Type::getInt8Ty(context),
5984+
llvm::APInt(8, 3)));
5985+
llvm::Value *is_el_set = builder->CreateICmpNE(el_mask_value,
5986+
llvm::ConstantInt::get(llvm::Type::getInt8Ty(context),
5987+
llvm::APInt(8, 0)));
5988+
5989+
llvm::Value *el_exists = builder->CreateAnd(is_el_set,
5990+
builder->CreateNot(is_el_skip));
5991+
5992+
llvm_utils->create_if_else(el_exists, [&]() {
5993+
LLVM::CreateStore(*builder, llvm_utils->list_api->read_item(el_list, idx,
5994+
false, *module, LLVM::is_llvm_struct(el_type)), pvar);
5995+
5996+
for (size_t i=0; i<x.n_body; i++) {
5997+
this->visit_stmt(*x.m_body[i]);
5998+
}
5999+
call_lcompilers_free_strings();
6000+
}, [=](){});
6001+
6002+
idx = builder->CreateAdd(idx,
6003+
llvm::ConstantInt::get(context, llvm::APInt(32, 1)));
6004+
LLVM::CreateStore(*builder, idx, idx_ptr);
6005+
});
6006+
} else {
6007+
throw CodeGenError("Only sets and dictionaries are supported with this loop for now.");
6008+
}
6009+
strings_to_be_deallocated.reserve(al, n);
6010+
strings_to_be_deallocated.n = n;
6011+
strings_to_be_deallocated.p = strings_to_be_deallocated_copy;
6012+
}
6013+
58316014
bool case_insensitive_string_compare(const std::string& str1, const std::string& str2) {
58326015
if (str1.size() != str2.size()) {
58336016
return false;

‎src/libasr/codegen/llvm_utils.cpp

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2075,6 +2075,11 @@ namespace LCompilers {
20752075
return get_key_value_pair_type(key_type_code, value_type_code);
20762076
}
20772077

2078+
llvm::Type* LLVMDict::get_key_value_pair_type(
2079+
ASR::ttype_t* /*key_asr_type*/, ASR::ttype_t* /*value_asr_type*/) {
2080+
return nullptr;
2081+
}
2082+
20782083
llvm::Type* LLVMDictSeparateChaining::get_dict_type(
20792084
std::string key_type_code, std::string value_type_code,
20802085
int32_t key_type_size, int32_t value_type_size,
@@ -2156,6 +2161,10 @@ namespace LCompilers {
21562161
return llvm_utils->create_gep(dict, 1);
21572162
}
21582163

2164+
llvm::Value* LLVMDict::get_pointer_to_key_value_pairs(llvm::Value* /*dict*/) {
2165+
return nullptr;
2166+
}
2167+
21592168
llvm::Value* LLVMDictSeparateChaining::get_pointer_to_key_value_pairs(llvm::Value* dict) {
21602169
return llvm_utils->create_gep(dict, 3);
21612170
}

‎src/libasr/codegen/llvm_utils.h

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -567,6 +567,9 @@ namespace LCompilers {
567567
virtual
568568
llvm::Value* get_pointer_to_occupancy(llvm::Value* dict) = 0;
569569

570+
virtual
571+
llvm::Value* get_pointer_to_keymask(llvm::Value* dict) = 0;
572+
570573
virtual
571574
llvm::Value* get_pointer_to_capacity(llvm::Value* dict) = 0;
572575

@@ -652,6 +655,13 @@ namespace LCompilers {
652655
std::map<std::string, std::map<std::string, int>>& name2memidx,
653656
bool key_or_value) = 0;
654657

658+
virtual
659+
llvm::Type* get_key_value_pair_type(ASR::ttype_t* key_asr_type, ASR::ttype_t* value_pair_type) = 0;
660+
661+
virtual
662+
llvm::Value* get_pointer_to_key_value_pairs(llvm::Value* dict) = 0;
663+
664+
655665
virtual ~LLVMDictInterface() = 0;
656666

657667
};
@@ -746,6 +756,10 @@ namespace LCompilers {
746756
std::map<std::string, std::map<std::string, int>>& name2memidx,
747757
bool key_or_value);
748758

759+
llvm::Type* get_key_value_pair_type(ASR::ttype_t* key_asr_type, ASR::ttype_t* value_pair_type);
760+
761+
llvm::Value* get_pointer_to_key_value_pairs(llvm::Value* dict);
762+
749763
virtual ~LLVMDict();
750764
};
751765

@@ -794,8 +808,6 @@ namespace LCompilers {
794808

795809
llvm::Value* get_pointer_to_number_of_filled_buckets(llvm::Value* dict);
796810

797-
llvm::Value* get_pointer_to_key_value_pairs(llvm::Value* dict);
798-
799811
llvm::Value* get_pointer_to_rehash_flag(llvm::Value* dict);
800812

801813
void deepcopy_key_value_pair_linked_list(llvm::Value* srci, llvm::Value* desti,
@@ -813,8 +825,6 @@ namespace LCompilers {
813825

814826
llvm::Type* get_key_value_pair_type(std::string key_type_code, std::string value_type_code);
815827

816-
llvm::Type* get_key_value_pair_type(ASR::ttype_t* key_asr_type, ASR::ttype_t* value_pair_type);
817-
818828
void dict_init_given_initial_capacity(std::string key_type_code, std::string value_type_code,
819829
llvm::Value* dict, llvm::Module* module, llvm::Value* initial_capacity);
820830

@@ -896,6 +906,10 @@ namespace LCompilers {
896906
std::map<std::string, std::map<std::string, int>>& name2memidx,
897907
bool key_or_value);
898908

909+
llvm::Type* get_key_value_pair_type(ASR::ttype_t* key_asr_type, ASR::ttype_t* value_pair_type);
910+
911+
llvm::Value* get_pointer_to_key_value_pairs(llvm::Value* dict);
912+
899913
virtual ~LLVMDictSeparateChaining();
900914

901915
};
@@ -943,6 +957,9 @@ namespace LCompilers {
943957
virtual
944958
llvm::Value* get_pointer_to_capacity(llvm::Value* set) = 0;
945959

960+
virtual
961+
llvm::Value* get_pointer_to_mask(llvm::Value* set) = 0;
962+
946963
llvm::Value* get_el_hash(llvm::Value* capacity, llvm::Value* el,
947964
ASR::ttype_t* el_asr_type, llvm::Module& module);
948965

0 commit comments

Comments
 (0)