C_CPP: Support visit_ListRepeat()

RobertoMalatesta · Jul 30, 2023 · 64f75fa · 64f75fa
1 parent 7481e00
commit 64f75fa
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 0 deletions.
diff --git a/src/libasr/codegen/asr_to_c_cpp.h b/src/libasr/codegen/asr_to_c_cpp.h
@@ -1676,6 +1676,20 @@ PyMODINIT_FUNC PyInit_lpython_module_)" + fn_name + R"((void) {
         src += indent + list_remove_func + "(&" + list_var + ", " + element + ");\n";
     }
 
+    void visit_ListRepeat(const ASR::ListRepeat_t& x) {
+        CHECK_FAST_C_CPP(compiler_options, x)
+        ASR::List_t* t = ASR::down_cast<ASR::List_t>(x.m_type);
+        std::string list_repeat_func = c_ds_api->get_list_repeat_func(t);
+        bracket_open++;
+        self().visit_expr(*x.m_left);
+        std::string list_var = std::move(src);
+        self().visit_expr(*x.m_right);
+        std::string freq = std::move(src);
+        bracket_open--;
+        tmp_buffer_src.push_back(check_tmp_buffer());
+        src = "(*" + list_repeat_func + "(&" + list_var + ", " + freq + "))";
+    }
+
     void visit_ListLen(const ASR::ListLen_t& x) {
         CHECK_FAST_C_CPP(compiler_options, x)
         self().visit_expr(*x.m_arg);

diff --git a/src/libasr/codegen/c_utils.h b/src/libasr/codegen/c_utils.h
@@ -594,6 +594,7 @@ class CCPPDSUtils {
             list_remove(list_struct_type, list_type_code, list_element_type, list_type->m_type);
             list_clear(list_struct_type, list_type_code, list_element_type);
             list_concat(list_struct_type, list_type_code, list_element_type, list_type->m_type);
+            list_repeat(list_struct_type, list_type_code, list_element_type, list_type->m_type);
             list_section(list_struct_type, list_type_code);
             return list_struct_type;
         }
@@ -652,6 +653,11 @@ class CCPPDSUtils {
             return typecodeToDSfuncs[list_type_code]["list_concat"];
         }
 
+        std::string get_list_repeat_func(ASR::List_t* list_type) {
+            std::string list_type_code = ASRUtils::get_type_code(list_type->m_type, true);
+            return typecodeToDSfuncs[list_type_code]["list_repeat"];
+        }
+
         std::string get_list_find_item_position_function(std::string list_type_code) {
             return typecodeToDSfuncs[list_type_code]["list_find_item"];
         }
@@ -934,6 +940,40 @@ class CCPPDSUtils {
             generated_code += indent + "}\n\n";
         }
 
+        void list_repeat(std::string list_struct_type,
+            std::string list_type_code,
+            std::string list_element_type, ASR::ttype_t *m_type) {
+            std::string indent(indentation_level * indentation_spaces, ' ');
+            std::string tab(indentation_spaces, ' ');
+            std::string list_con_func = global_scope->get_unique_name("list_repeat_" + list_type_code);
+            typecodeToDSfuncs[list_type_code]["list_repeat"] = list_con_func;
+            std::string init_func = typecodeToDSfuncs[list_type_code]["list_init"];
+            std::string signature = list_struct_type + "* " + list_con_func + "("
+                                + list_struct_type + "* x, "
+                                + "int32_t freq)";
+            func_decls += "inline " + signature + ";\n";
+            generated_code += indent + signature + " {\n";
+            generated_code += indent + tab + list_struct_type + " *result = (" + list_struct_type + "*)malloc(sizeof(" +
+                                list_struct_type + "));\n";
+            generated_code += indent + tab + init_func + "(result, x->current_end_point * freq);\n";
+            generated_code += indent + tab + "for (int i=0; i<freq; i++) {\n";
+
+            if (ASR::is_a<ASR::List_t>(*m_type)) {
+                ASR::ttype_t *tt = ASR::down_cast<ASR::List_t>(m_type)->m_type;
+                std::string deep_copy_func = typecodeToDSfuncs[ASRUtils::get_type_code(tt, true)]["list_deepcopy"];
+                LCOMPILERS_ASSERT(deep_copy_func.size() > 0);
+                generated_code += indent + tab + tab + "for(int j=0; j<x->current_end_point; j++)\n";
+                generated_code += indent + tab + tab + tab + deep_copy_func + "(&x->data[j], &result->data[i*x->current_end_point+j]);\n";
+            } else {
+                generated_code += indent + tab + tab + "memcpy(&result->data[i*x->current_end_point], x->data, x->current_end_point * sizeof(" + list_element_type + "));\n";
+            }
+
+            generated_code += indent + tab + "}\n";
+            generated_code += indent + tab + "result->current_end_point = x->current_end_point * freq;\n";
+            generated_code += indent + tab + "return result;\n";
+            generated_code += indent + "}\n\n";
+        }
+
         void resize_if_needed(std::string list_struct_type,
             std::string list_type_code,
             std::string list_element_type) {