Add support for in place sets and dictionaries in function calls

lcompilers · Jul 16, 2024 · 426bd80 · 426bd80
1 parent 9e8da7d
commit 426bd80
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 1 deletion.
diff --git a/integration_tests/test_dict_14.py b/integration_tests/test_dict_14.py
@@ -1,5 +1,8 @@
 from lpython import i32
 
+def takes_dict(a: dict[i32, i32]) -> dict[i32, i32]:
+    return {1:1, 2:2}
+
 def test_dict():
     d_i32: dict[i32, i32] = {5: 1, 5: 2}
     d_str: dict[str, i32] = {'a': 1, 'a': 2}
@@ -62,4 +65,7 @@ def test_dict():
         l_i32_2.append(i)
     assert l_i32_2 == [30]
 
+    w: dict[i32, i32] = takes_dict({1:1, 2:2})
+    assert len(w)  == 2
+
 test_dict()
diff --git a/integration_tests/test_global_set.py b/integration_tests/test_global_set.py
@@ -1,9 +1,15 @@
 from lpython import i32
 
+def takes_set(a: set[i32]) -> set[i32]:
+    return {1, 2, 3}
+
 s1: set[str] = {"a", "b", "c", "a"}
 s2: set[i32] = {1, 2, 3, 1}
 s3: set[tuple[i32, i32]] = {(1, 2), (2, 3), (4, 5)}
 
+s4: set[i32] = takes_set({1, 2})
+
 assert len(s1) == 3
 assert len(s2) == 3
 assert len(s3) == 3
+assert len(s4) == 3
diff --git a/src/libasr/codegen/asr_to_llvm.cpp b/src/libasr/codegen/asr_to_llvm.cpp
@@ -8826,6 +8826,14 @@ class ASRToLLVMVisitor : public ASR::BaseVisitor<ASRToLLVMVisitor>
                         target_type = llvm_utils->get_type_from_ttype_t_util(arg_type_, module.get());
                         break;
                     }
+                    case (ASR::ttypeType::Dict): {
+                        target_type = llvm_utils->get_type_from_ttype_t_util(arg_type_, module.get());
+                        break;
+                    }
+                    case (ASR::ttypeType::Set): {
+                        target_type = llvm_utils->get_type_from_ttype_t_util(arg_type_, module.get());
+                        break;
+                    }
                     default :
                         throw CodeGenError("Type " + ASRUtils::type_to_str(arg_type) + " not implemented yet.");
                 }
@@ -8880,7 +8888,9 @@ class ASRToLLVMVisitor : public ASR::BaseVisitor<ASRToLLVMVisitor>
                                     llvm::AllocaInst *target = builder0.CreateAlloca(
                                         target_type, nullptr, "call_arg_value");
                                     if( ASR::is_a<ASR::Tuple_t>(*arg_type) ||
-                                        ASR::is_a<ASR::List_t>(*arg_type) ) {
+                                        ASR::is_a<ASR::List_t>(*arg_type)  ||
+                                        ASR::is_a<ASR::Set_t>(*arg_type)   ||
+                                        ASR::is_a<ASR::Dict_t>(*arg_type)) {
                                         llvm_utils->deepcopy(value, target, arg_type, module.get(), name2memidx);
                                     } else {
                                         builder->CreateStore(value, target);