Support passing dict as return value of functions (#1659)

Author: harshsingh-24

integration_tests/CMakeLists.txt

@@ -288,6 +288,8 @@ RUN(NAME test_dict_02        LABELS cpython llvm c)
 RUN(NAME test_dict_03        LABELS cpython llvm)
 RUN(NAME test_dict_04        LABELS cpython llvm)
 RUN(NAME test_dict_05        LABELS cpython llvm)
+RUN(NAME test_dict_06        LABELS cpython llvm c)
+RUN(NAME test_dict_07        LABELS cpython llvm)
 RUN(NAME test_for_loop       LABELS cpython llvm c)
 RUN(NAME modules_01          LABELS cpython llvm c wasm wasm_x86 wasm_x64)
 RUN(NAME modules_02          LABELS cpython llvm c wasm wasm_x86 wasm_x64)

integration_tests/test_dict_06.py

@@ -0,0 +1,24 @@
+from lpython import f64, i32
+
+def fill_rollnumber2cpi(size: i32) -> dict[i32, f64]:
+    i : i32
+    rollnumber2cpi: dict[i32, f64] = {}
+    
+    rollnumber2cpi[0] = 1.1
+    for i in range(1000, 1000 + size):
+        rollnumber2cpi[i] = float(i/100.0 + 5.0)
+
+    return rollnumber2cpi
+
+def test_dict():
+    i: i32
+    size: i32 = 200
+    rollnumber2cpi: dict[i32, f64] = fill_rollnumber2cpi(size)
+
+    for i in range(1000 + size - 1, 1001, -1):
+        assert abs(rollnumber2cpi[i] - i/100.0 - 5.0) <= 1e-12
+
+    assert abs(rollnumber2cpi[0] - 1.1) <= 1e-12
+    assert len(rollnumber2cpi) == 201
+
+test_dict()

integration_tests/test_dict_07.py

@@ -0,0 +1,18 @@
+def fill_smalltocapital() -> dict[str, str]:
+    return {'a': 'A', 'b': 'B', 'c': 'C', 'd': 'D','e': 'E', 
+            'f': 'F', 'g': 'G', 'h': 'H', 'i': 'I','j': 'J', 
+            'k': 'K', 'l': 'L', 'm': 'M', 'n': 'N','o': 'O', 
+            'p': 'P', 'q': 'Q', 'r': 'R', 's': 'S','t': 'T', 
+            'u': 'U', 'v': 'V', 'w': 'W', 'x': 'X','y': 'Y', 
+            'z': 'Z'}
+
+def test_dict():
+    i : i32
+    smalltocaps: dict[str, str]
+    smalltocaps = fill_smalltocapital()
+
+    assert len(smalltocaps) == 26
+    for i in range(97, 97 + 26):
+        assert smalltocaps[chr(i)] == chr(i - 32)
+
+test_dict()

src/libasr/codegen/asr_to_c_cpp.h

@@ -442,6 +442,9 @@ R"(#include <stdio.h>
             } else if (ASR::is_a<ASR::TypeParameter_t>(*return_var->m_type)) {
                 has_typevar = true;
                 return "";
+            } else if (ASR::is_a<ASR::Dict_t>(*return_var->m_type)) {
+                ASR::Dict_t* dict_type = ASR::down_cast<ASR::Dict_t>(return_var->m_type);
+                sub = c_ds_api->get_dict_type(dict_type) + " ";
             } else {
                 throw CodeGenError("Return type not supported in function '" +
                     std::string(x.m_name) +
@@ -661,6 +664,16 @@ R"(#include <stdio.h>
                                 const_name + " = " + src + ";\n");
             src = const_name;
             return;
+        } else if( ASR::is_a<ASR::Dict_t>(*x.m_type) ) {
+            ASR::Dict_t* dict_type = ASR::down_cast<ASR::Dict_t>(x.m_type);
+            const_name += std::to_string(const_vars_count);
+            const_vars_count += 1;
+            const_name = current_scope->get_unique_name(const_name);
+            std::string indent(indentation_level*indentation_spaces, ' ');
+            tmp_buffer_src.push_back(check_tmp_buffer() + indent + c_ds_api->get_dict_type(dict_type) + 
+                                " " + const_name + " = " + src + ";\n");
+            src = const_name;
+            return;
         }
         src = check_tmp_buffer() + src;
     }
@@ -797,6 +810,9 @@ R"(#include <stdio.h>
             }
             src = check_tmp_buffer() + src_tmp;
             return;
+        } else if (ASR::is_a<ASR::DictItem_t>(*x.m_target)) {
+            self().visit_DictItem(*ASR::down_cast<ASR::DictItem_t>(x.m_target));
+            target = src;
         } else {
             LCOMPILERS_ASSERT(false)
         }

src/libasr/codegen/asr_to_llvm.cpp

@@ -4084,6 +4084,27 @@ class ASRToLLVMVisitor : public ASR::BaseVisitor<ASRToLLVMVisitor>
                     return_type = list_api->get_list_type(el_llvm_type, el_type_code, type_size);
                     break;
                 }
+                case (ASR::ttypeType::Dict) : {
+                    ASR::Dict_t* asr_dict = ASR::down_cast<ASR::Dict_t>(return_var_type0);
+                    std::string key_type_code = ASRUtils::get_type_code(asr_dict->m_key_type);
+                    std::string value_type_code = ASRUtils::get_type_code(asr_dict->m_value_type);
+                    
+                    bool is_local_array_type = false, is_local_malloc_array_type = false;
+                    bool is_local_list = false;
+                    ASR::dimension_t* local_m_dims = nullptr;
+                    ASR::storage_typeType local_m_storage = ASR::storage_typeType::Default;
+                    int local_n_dims = 0, local_a_kind = -1;
+
+                    llvm::Type* key_llvm_type = get_type_from_ttype_t(asr_dict->m_key_type, local_m_storage,is_local_array_type, is_local_malloc_array_type,is_local_list, local_m_dims, local_n_dims,local_a_kind);
+                    llvm::Type* value_llvm_type = get_type_from_ttype_t(asr_dict->m_value_type, local_m_storage,is_local_array_type, is_local_malloc_array_type,is_local_list, local_m_dims, local_n_dims,local_a_kind);
+                    int32_t key_type_size = get_type_size(asr_dict->m_key_type, key_llvm_type, local_a_kind);
+                    int32_t value_type_size = get_type_size(asr_dict->m_value_type, value_llvm_type, local_a_kind);
+
+                    set_dict_api(asr_dict);
+
+                    return_type = llvm_utils->dict_api->get_dict_type(key_type_code, value_type_code, key_type_size,value_type_size, key_llvm_type, value_llvm_type);
+                    break;
+                }
                 default :
                     throw CodeGenError("Type not implemented " + std::to_string(return_var_type));
             }

src/libasr/codegen/llvm_utils.h

@@ -95,7 +95,8 @@ namespace LCompilers {
             return ASR::is_a<ASR::Tuple_t>(*asr_type) ||
                    ASR::is_a<ASR::List_t>(*asr_type) ||
                    ASR::is_a<ASR::Struct_t>(*asr_type) ||
-                   ASR::is_a<ASR::Class_t>(*asr_type);
+                   ASR::is_a<ASR::Class_t>(*asr_type)|| 
+                   ASR::is_a<ASR::Dict_t>(*asr_type);
         }
     }