Merge pull request #2778 from advikkabra/set-args

Add support for in place sets and dictionaries in function calls
lcompilers · Aug 16, 2024 · f3389e6 · f3389e6
2 parents 81e2b1e + 4d91f3c
commit f3389e6
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 1 deletion.
diff --git a/integration_tests/CMakeLists.txt b/integration_tests/CMakeLists.txt
@@ -569,6 +569,7 @@ RUN(NAME test_tuple_04       LABELS cpython llvm llvm_jit c)
 RUN(NAME test_tuple_concat   LABELS cpython llvm llvm_jit)
 RUN(NAME test_tuple_nested   LABELS cpython llvm llvm_jit)
 RUN(NAME test_const_dict     LABELS cpython llvm llvm_jit)
+RUN(NAME test_params         LABELS cpython llvm llvm_jit NOFAST)
 RUN(NAME test_dict_01        LABELS cpython llvm llvm_jit c)
 RUN(NAME test_dict_02        LABELS cpython llvm llvm_jit c NOFAST)
 RUN(NAME test_dict_03        LABELS cpython llvm llvm_jit NOFAST)

diff --git a/integration_tests/test_params.py b/integration_tests/test_params.py
@@ -0,0 +1,14 @@
+from lpython import i32
+
+def takes_set(a: set[i32]) -> set[i32]:
+    return {1, 2, 3}
+
+def takes_dict(a: dict[i32, i32]) -> dict[i32, i32]:
+    return {1:1, 2:2}
+
+s: set[i32] = takes_set({1, 2})
+
+assert len(s) == 3
+
+w: dict[i32, i32] = takes_dict({1:1, 2:2})
+assert len(w)  == 2
diff --git a/src/libasr/codegen/asr_to_llvm.cpp b/src/libasr/codegen/asr_to_llvm.cpp
@@ -8859,6 +8859,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.");
                 }
@@ -8913,7 +8921,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);