Merge pull request #955 from Smit-create/list_ret

Support returning list from functions in LLVM backend

Author: czgdp1807

integration_tests/test_list_01.py

@@ -1,12 +1,21 @@
 from ltypes import f64, i32
 
+def fill_list_i32(size: i32) -> list[i32]:
+    aarg: list[i32] = [0, 1, 2, 3, 4]
+    i: i32
+    for i in range(10):
+        aarg.append(i + 5)
+    return aarg
+
+
 def test_list_01():
-    a: list[i32] = [0, 1, 2, 3, 4]
+    a: list[i32] = []
     f: list[f64] = [1.0, 2.0, 3.0, 4.0, 5.0]
     i: i32
 
+    a = fill_list_i32(10)
+
     for i in range(10):
-        a.append(i + 5)
         f.append(float(i + 6))
 
     for i in range(15):

integration_tests/test_list_02.py

@@ -1,14 +1,22 @@
 from ltypes import i32
 
-def test_list_01():
+def fill_list_str(size: i32) -> list[str]:
     a: list[str] = ["0_str", "1_str"]
+    i: i32
+    for i in range(size):
+        a.append(str(i + 2) + "_str")
+    return a
+
+def test_list_01():
+    a: list[str] = []
     b: list[str]
     string: str = "string_"
     b = [string, string]
     i: i32
 
+    a = fill_list_str(10)
+
     for i in range(10):
-        a.append(str(i + 2) + "_str")
         b.append(string + str(i + 2))
 
 

integration_tests/test_list_03.py

@@ -10,9 +10,7 @@ def test_list_01(n: i32) -> i32:
         sum += a[i]
     return sum
 
-def test_list_02(n: i32) -> i32:
-    x: list[i32] = [50, 1]
-
+def test_list_insert_02(x: list[i32], n: i32) -> list[i32]:
     i: i32
     imod: i32
     for i in range(n):
@@ -24,7 +22,15 @@ def test_list_02(n: i32) -> i32:
         elif imod == 2:
             x.insert(len(x)//2, i + n + 2)
 
+    return x
+
+def test_list_02(n: i32) -> i32:
+    x: list[i32] = [50, 1]
     acc: i32 = 0
+    i: i32
+
+    x = test_list_insert_02(x, n)
+
     for i in range(n):
         acc += x[i]
     return acc
@@ -49,6 +55,7 @@ def test_list_02_string():
     for i in range(50):
         assert x[i] == y[i]
 
+
 def verify():
     assert test_list_01(11) == 55
     assert test_list_02(50) == 3628

integration_tests/test_list_05.py

@@ -18,17 +18,38 @@ def check_list_of_tuples(l: list[tuple[i32, f64, str]], sign: i32):
         assert t[2] == string
         assert l[i][2] == string
 
-def test_list_of_tuples():
+def fill_list_of_tuples(size: i32) -> list[tuple[i32, f64, str]]:
     l1: list[tuple[i32, f64, str]] = []
     t: tuple[i32, f64, str]
-    size: i32 = 20
     i: i32
-    string: str
 
     for i in range(size):
         t = (i, float(i), str(i) + "_str")
         l1.append(t)
 
+    return l1
+
+def insert_tuples_into_list(l: list[tuple[i32, f64, str]], size: i32) -> list[tuple[i32, f64, str]]:
+    i: i32
+    string: str
+    t: tuple[i32, f64, str]
+
+    for i in range(size//2, size):
+        string  = str(i) + "_str"
+        t = (i, float(i), string)
+        l.insert(i, t)
+
+    return l
+
+def test_list_of_tuples():
+    l1: list[tuple[i32, f64, str]] = []
+    t: tuple[i32, f64, str]
+    size: i32 = 20
+    i: i32
+    string: str
+
+    l1 = fill_list_of_tuples(size)
+
     check_list_of_tuples(l1, 1)
 
     for i in range(size//2):
@@ -39,10 +60,7 @@ def test_list_of_tuples():
     assert t[1] == size//2 - 1
     assert t[2] == str(size//2 - 1) + "_str"
 
-    for i in range(size//2, size):
-        string  = str(i) + "_str"
-        t = (i, float(i), string)
-        l1.insert(i, t)
+    l1 = insert_tuples_into_list(l1, size)
 
     check_list_of_tuples(l1, 1)
 

integration_tests/test_list_07.py

@@ -1,6 +1,28 @@
 from ltypes import c64, i32
 from copy import deepcopy
 
+def generate_complex_tensors(mat: list[list[c64]], vec: list[c64]) -> list[tuple[list[list[c64]], list[c64]]]:
+    tensor: tuple[list[list[c64]], list[c64]]
+    tensors: list[tuple[list[list[c64]], list[c64]]] = []
+    rows: i32 = len(mat)
+    cols: i32 = len(vec)
+    i: i32; j: i32; k: i32
+
+    tensor = (deepcopy(mat), deepcopy(vec))
+
+    for k in range(2 * rows):
+        tensors.append(deepcopy(tensor))
+        for i in range(rows):
+            for j in range(cols):
+                mat[i][j] += complex(1.0, 2.0)
+
+        for i in range(cols):
+            vec[i] += complex(1.0, 2.0)
+
+        tensor = (deepcopy(mat), deepcopy(vec))
+
+    return tensors
+
 def test_tuple_with_lists():
     mat: list[list[c64]] = []
     vec: list[c64] = []
@@ -42,18 +64,7 @@ def test_tuple_with_lists():
     for i in range(cols):
         assert tensor[1][i] - vec[i] == -complex(0, 2.0)
 
-    tensor = (deepcopy(mat), deepcopy(vec))
-
-    for k in range(2 * rows):
-        tensors.append(deepcopy(tensor))
-        for i in range(rows):
-            for j in range(cols):
-                mat[i][j] += complex(1.0, 2.0)
-
-        for i in range(cols):
-            vec[i] += complex(1.0, 2.0)
-
-        tensor = (deepcopy(mat), deepcopy(vec))
+    tensors = generate_complex_tensors(mat, vec)
 
     for k in range(2 * rows):
         for i in range(rows):

src/libasr/codegen/asr_to_llvm.cpp

@@ -1134,10 +1134,19 @@ class ASRToLLVMVisitor : public ASR::BaseVisitor<ASRToLLVMVisitor>
 
     void visit_ListConstant(const ASR::ListConstant_t& x) {
         ASR::List_t* list_type = ASR::down_cast<ASR::List_t>(x.m_type);
-        llvm::Type* llvm_el_type = get_el_type(list_type->m_type);
+        bool is_array_type_local = false, is_malloc_array_type_local = false;
+        bool is_list_local = false;
+        ASR::dimension_t* m_dims_local = nullptr;
+        int n_dims_local = -1, a_kind_local = -1;
+        llvm::Type* llvm_el_type = get_type_from_ttype_t(list_type->m_type,
+                                    ASR::storage_typeType::Default, is_array_type_local,
+                                    is_malloc_array_type_local, is_list_local, m_dims_local,
+                                    n_dims_local, a_kind_local);
         std::string type_code = ASRUtils::get_type_code(list_type->m_type);
         int32_t type_size = -1;
-        if( ASR::is_a<ASR::Character_t>(*list_type->m_type) ) {
+        if( ASR::is_a<ASR::Character_t>(*list_type->m_type) ||
+            LLVM::is_llvm_struct(list_type->m_type) ||
+            ASR::is_a<ASR::Complex_t>(*list_type->m_type) ) {
             llvm::DataLayout data_layout(module.get());
             type_size = data_layout.getTypeAllocSize(llvm_el_type);
         } else {
@@ -2720,6 +2729,31 @@ class ASRToLLVMVisitor : public ASR::BaseVisitor<ASRToLLVMVisitor>
                     return_type = tuple_api->get_tuple_type(type_code, llvm_el_types);
                     break;
                 }
+                case (ASR::ttypeType::List) : {
+                    bool is_array_type = false, is_malloc_array_type = false;
+                    bool is_list = true;
+                    ASR::dimension_t *m_dims = nullptr;
+                    ASR::storage_typeType m_storage = ASR::storage_typeType::Default;
+                    int n_dims = 0, a_kind = -1;
+                    ASR::List_t* asr_list = ASR::down_cast<ASR::List_t>(return_var_type0);
+                    llvm::Type* el_llvm_type = get_type_from_ttype_t(asr_list->m_type, m_storage,
+                                                                     is_array_type,
+                                                                     is_malloc_array_type,
+                                                                     is_list, m_dims, n_dims,
+                                                                     a_kind);
+                    int32_t type_size = -1;
+                    if( LLVM::is_llvm_struct(asr_list->m_type) ||
+                        ASR::is_a<ASR::Character_t>(*asr_list->m_type) ||
+                        ASR::is_a<ASR::Complex_t>(*asr_list->m_type) ) {
+                        llvm::DataLayout data_layout(module.get());
+                        type_size = data_layout.getTypeAllocSize(el_llvm_type);
+                    } else {
+                        type_size = a_kind;
+                    }
+                    std::string el_type_code = ASRUtils::get_type_code(asr_list->m_type);
+                    return_type = list_api->get_list_type(el_llvm_type, el_type_code, type_size);
+                    break;
+                }
                 default :
                     LFORTRAN_ASSERT(false);
                     throw CodeGenError("Type not implemented");
@@ -3122,15 +3156,11 @@ class ASRToLLVMVisitor : public ASR::BaseVisitor<ASRToLLVMVisitor>
                 this->visit_expr(*x.m_target);
                 llvm::Value* target_tuple = tmp;
                 ptr_loads = ptr_loads_copy;
-                if( ASR::is_a<ASR::FunctionCall_t>(*x.m_value) ) {
-                    builder->CreateStore(value_tuple, target_tuple);
-                } else {
-                    ASR::Tuple_t* value_tuple_type = ASR::down_cast<ASR::Tuple_t>(asr_value_type);
-                    std::string type_code = ASRUtils::get_type_code(value_tuple_type->m_type,
-                                                                    value_tuple_type->n_type);
-                    tuple_api->tuple_deepcopy(value_tuple, target_tuple,
-                                              value_tuple_type, *module);
-                }
+                ASR::Tuple_t* value_tuple_type = ASR::down_cast<ASR::Tuple_t>(asr_value_type);
+                std::string type_code = ASRUtils::get_type_code(value_tuple_type->m_type,
+                                                                value_tuple_type->n_type);
+                tuple_api->tuple_deepcopy(value_tuple, target_tuple,
+                                          value_tuple_type, *module);
             }
             return ;
         } else if( is_target_dict && is_value_dict ) {
@@ -5333,6 +5363,33 @@ class ASRToLLVMVisitor : public ASR::BaseVisitor<ASRToLLVMVisitor>
         tmp = builder->CreateOr(arg1, arg2);
     }
 
+    llvm::Value* CreatePointerToStructReturnValue(llvm::FunctionType* fnty,
+                                                  llvm::Value* return_value,
+                                                  ASR::ttype_t* asr_return_type) {
+        if( !LLVM::is_llvm_struct(asr_return_type) ) {
+            return return_value;
+        }
+
+        // Call to LLVM APIs not needed to fetch the return type of the function.
+        // We can use asr_return_type as well but anyways for compactness I did it here.
+        llvm::Value* pointer_to_struct = builder->CreateAlloca(fnty->getReturnType(), nullptr);
+        LLVM::CreateStore(*builder, return_value, pointer_to_struct);
+        return pointer_to_struct;
+    }
+
+    llvm::Value* CreateCallUtil(llvm::FunctionType* fnty, llvm::Function* fn,
+                                std::vector<llvm::Value*>& args,
+                                ASR::ttype_t* asr_return_type) {
+        llvm::Value* return_value = builder->CreateCall(fn, args);
+        return CreatePointerToStructReturnValue(fnty, return_value,
+                                                asr_return_type);
+    }
+
+    llvm::Value* CreateCallUtil(llvm::Function* fn, std::vector<llvm::Value*>& args,
+                                ASR::ttype_t* asr_return_type) {
+        return CreateCallUtil(fn->getFunctionType(), fn, args, asr_return_type);
+    }
+
     void visit_FunctionCall(const ASR::FunctionCall_t &x) {
         if( ASRUtils::is_intrinsic_optimization(x.m_name) ) {
             ASR::Function_t* routine = ASR::down_cast<ASR::Function_t>(
@@ -5426,8 +5483,8 @@ class ASRToLLVMVisitor : public ASR::BaseVisitor<ASRToLLVMVisitor>
             std::string m_name = std::string(((ASR::Function_t*)(&(x.m_name->base)))->m_name);
             std::vector<llvm::Value *> args2 = convert_call_args(x, m_name);
             args.insert(args.end(), args2.begin(), args2.end());
+            ASR::ttype_t *return_var_type0 = EXPR2VAR(s->m_return_var)->m_type;
             if (s->m_abi == ASR::abiType::BindC) {
-                ASR::ttype_t *return_var_type0 = EXPR2VAR(s->m_return_var)->m_type;
                 if (is_a<ASR::Complex_t>(*return_var_type0)) {
                     int a_kind = down_cast<ASR::Complex_t>(return_var_type0)->m_kind;
                     if (a_kind == 8) {
@@ -5447,7 +5504,7 @@ class ASRToLLVMVisitor : public ASR::BaseVisitor<ASRToLLVMVisitor>
                     tmp = builder->CreateCall(fn, args);
                 }
             } else {
-                tmp = builder->CreateCall(fn, args);
+                tmp = CreateCallUtil(fn, args, return_var_type0);
             }
         }
         if (s->m_abi == ASR::abiType::BindC) {