Fixed compatibility of sep, end keyword with list/tuple containers (#1512)

anutosh491 · web-flow · commit c1563c7eb3b5 · 2023-02-26T10:16:48.000-07:00
diff --git a/integration_tests/CMakeLists.txt b/integration_tests/CMakeLists.txt
@@ -228,6 +228,7 @@ RUN(NAME exit_02c    FAIL    LABELS cpython llvm c)
 RUN(NAME print_01            LABELS cpython llvm c wasm) # wasm not yet supports sep and end keywords
 RUN(NAME print_03            LABELS x86 c wasm wasm_x86 wasm_x64) # simple test case specifically for x86, wasm_x86 and wasm_x64
 RUN(NAME print_04            LABELS cpython llvm c)
+RUN(NAME print_06            LABELS cpython llvm c)
 RUN(NAME print_05            LABELS cpython llvm wasm wasm_x64)
 RUN(NAME print_float         LABELS cpython llvm wasm wasm_x64)
 RUN(NAME print_list_tuple    LABELS cpython c) # TODO: llvm doesn't completely support printing list
diff --git a/integration_tests/print_06.py b/integration_tests/print_06.py
@@ -0,0 +1,26 @@
+def main0():
+    a: i32 = 1
+    b: str = "abc"
+    c: list[i32] = [1, 2]
+    d: list[i32] = [3, 4]
+    e: tuple[i32, i32, i32] = (1, 2, 3)
+    f: tuple[i32, i32] = (4, 5)
+    print(a, b, c)
+    print(a, c, b)
+    print(c, a, b)
+    print(a, b, c, sep = "pqr")
+    print(a, c, b, sep = "pqr")
+    print(c, a, b, sep = "pqr")
+    print(a, b, c, end = "xyz\n")
+    print(a, c, b, end = "xyz\n")
+    print(c, a, b, end = "xyz\n")
+    print(a, b, c, sep = "pqr", end = "xyz\n")
+    print(a, c, b, sep = "pqr", end = "xyz\n")
+    print(c, a, b, sep = "pqr", end = "xyz\n")
+    # Tring out few cases with Lists and Tuples together
+    print(c, e)
+    print(c, d, a, sep = "pqr")
+    print(c, e, d, end = "xyz\n")
+    print(c, e, d, f, sep = "pqr", end = "xyz\n")
+
+main0()
diff --git a/src/libasr/pass/print_arr.cpp b/src/libasr/pass/print_arr.cpp
@@ -60,6 +60,10 @@ class PrintArrVisitor : public PassUtils::PassVisitor<PrintArrVisitor>
         ASR::stmt_t* doloop = nullptr;
         ASR::stmt_t* empty_print_endl = ASRUtils::STMT(ASR::make_Print_t(al, loc,
                                             nullptr, nullptr, 0, nullptr, nullptr));
+        ASR::ttype_t *str_type_len_1 = ASRUtils::TYPE(ASR::make_Character_t(
+        al, loc, 1, 1, nullptr, nullptr, 0));
+        ASR::expr_t *space = ASRUtils::EXPR(ASR::make_StringConstant_t(
+        al, loc, s2c(al, " "), str_type_len_1));
         for( int i = n_dims - 1; i >= 0; i-- ) {
             ASR::do_loop_head_t head;
             head.m_v = idx_vars[i];
@@ -75,7 +79,7 @@ class PrintArrVisitor : public PassUtils::PassVisitor<PrintArrVisitor>
                 print_args.reserve(al, 1);
                 print_args.push_back(al, ref);
                 ASR::stmt_t* print_stmt = ASRUtils::STMT(ASR::make_Print_t(al, loc, nullptr,
-                                                                print_args.p, print_args.size(), nullptr, nullptr));
+                                                                print_args.p, print_args.size(), nullptr, space));
                 doloop_body.push_back(al, print_stmt);
             } else {
                 doloop_body.push_back(al, doloop);
@@ -88,9 +92,16 @@ class PrintArrVisitor : public PassUtils::PassVisitor<PrintArrVisitor>
 
     void visit_Print(const ASR::Print_t& x) {
         std::vector<ASR::expr_t*> print_body;
-        ASR::stmt_t* empty_print_endl = ASRUtils::STMT(ASR::make_Print_t(al, x.base.base.loc,
-                                            nullptr, nullptr, 0, nullptr, nullptr));
+        ASR::stmt_t* empty_print_endl;
         ASR::stmt_t* print_stmt;
+        ASR::ttype_t *str_type_len_1 = ASRUtils::TYPE(ASR::make_Character_t(
+        al, x.base.base.loc, 1, 1, nullptr, nullptr, 0));
+        ASR::expr_t *space = ASRUtils::EXPR(ASR::make_StringConstant_t(
+        al, x.base.base.loc, s2c(al, " "), str_type_len_1));
+        ASR::expr_t *backspace = ASRUtils::EXPR(ASR::make_StringConstant_t(
+        al, x.base.base.loc, s2c(al, "\b"), str_type_len_1));
+        ASR::stmt_t* back = ASRUtils::STMT(ASR::make_Print_t(al, x.base.base.loc,
+                                            nullptr, nullptr, 0, nullptr, backspace));
         for (size_t i=0; i<x.n_values; i++) {
             // TODO: This will disallow printing array pointer in Fortran
             // Pointers are treated the same as normal variables in Fortran
@@ -107,15 +118,38 @@ class PrintArrVisitor : public PassUtils::PassVisitor<PrintArrVisitor>
                     for (size_t j=0; j<print_body.size(); j++) {
                         body.push_back(al, print_body[j]);
                     }
-                    print_stmt = ASRUtils::STMT(ASR::make_Print_t(
-                        al, x.base.base.loc, nullptr, body.p, body.size(),
-                        nullptr, nullptr));
+                    if (x.m_separator) {
+                        print_stmt = ASRUtils::STMT(ASR::make_Print_t(
+                            al, x.base.base.loc, nullptr, body.p, body.size(),
+                            x.m_separator, x.m_separator));
+                    } else {
+                        print_stmt = ASRUtils::STMT(ASR::make_Print_t(
+                            al, x.base.base.loc, nullptr, body.p, body.size(),
+                            nullptr, space));
+                    }
                     pass_result.push_back(al, print_stmt);
-                    pass_result.push_back(al, empty_print_endl);
                     print_body.clear();
                 }
                 print_stmt = print_array_using_doloop(x.m_values[i], x.base.base.loc);
                 pass_result.push_back(al, print_stmt);
+                pass_result.push_back(al, back);
+                if (x.m_separator) {
+                    if (i == x.n_values - 1) {
+                        empty_print_endl = ASRUtils::STMT(ASR::make_Print_t(al, x.base.base.loc,
+                                                nullptr, nullptr, 0, nullptr, x.m_end));
+                    } else {
+                        empty_print_endl = ASRUtils::STMT(ASR::make_Print_t(al, x.base.base.loc,
+                                                nullptr, nullptr, 0, nullptr, x.m_separator));
+                    }
+                } else {
+                    if (i == x.n_values - 1) {
+                        empty_print_endl = ASRUtils::STMT(ASR::make_Print_t(al, x.base.base.loc,
+                                                nullptr, nullptr, 0, nullptr, x.m_end));
+                    } else {
+                        empty_print_endl = ASRUtils::STMT(ASR::make_Print_t(al, x.base.base.loc,
+                                                nullptr, nullptr, 0, nullptr, nullptr));
+                    }
+                }
                 pass_result.push_back(al, empty_print_endl);
             } else {
                 print_body.push_back(x.m_values[i]);
@@ -129,7 +163,7 @@ class PrintArrVisitor : public PassUtils::PassVisitor<PrintArrVisitor>
             }
             print_stmt = ASRUtils::STMT(ASR::make_Print_t(
                 al, x.base.base.loc, nullptr, body.p, body.size(),
-                nullptr, nullptr));
+                x.m_separator, x.m_end));
             pass_result.push_back(al, print_stmt);
             print_body.clear();
         }
diff --git a/src/libasr/pass/print_list.cpp b/src/libasr/pass/print_list.cpp
@@ -18,8 +18,7 @@ comma_space, brackets and newline. The function
     print(a, b, l, sep="pqr", end="xyz") # l is a list (but not a & b)
 
 to:
-
-    print(a, b, sep="pqr")
+    print(a, b, sep="pqr", end="")
     print("[", end="")
     for i in range(len(l)):
         print(l[i], end="")
@@ -300,26 +299,61 @@ class PrintListVisitor
 
     void visit_Print(const ASR::Print_t &x) {
         std::vector<ASR::expr_t*> print_tmp;
+        ASR::ttype_t *str_type_len_1 = ASRUtils::TYPE(ASR::make_Character_t(
+        al, x.base.base.loc, 1, 1, nullptr, nullptr, 0));
+        ASR::expr_t *space = ASRUtils::EXPR(ASR::make_StringConstant_t(
+        al, x.base.base.loc, s2c(al, " "), str_type_len_1));
         for (size_t i=0; i<x.n_values; i++) {
             if (ASR::is_a<ASR::List_t>(*ASRUtils::expr_type(x.m_values[i])) ||
                 ASR::is_a<ASR::Tuple_t>(*ASRUtils::expr_type(x.m_values[i]))) {
                 if (!print_tmp.empty()) {
                     Vec<ASR::expr_t*> tmp_vec;
+                    ASR::stmt_t *print_stmt;
                     tmp_vec.reserve(al, print_tmp.size());
                     for (auto &e: print_tmp) {
                         tmp_vec.push_back(al, e);
                     }
-                    ASR::stmt_t *print_stmt = ASRUtils::STMT(
-                        ASR::make_Print_t(al, x.base.base.loc, nullptr, tmp_vec.p, tmp_vec.size(),
-                                    x.m_separator, nullptr));
+                    if (x.m_separator) {
+                        print_stmt = ASRUtils::STMT(ASR::make_Print_t(al,
+                            x.base.base.loc, nullptr, tmp_vec.p, tmp_vec.size(),
+                            x.m_separator, x.m_separator));
+                    } else {
+                        print_stmt = ASRUtils::STMT(ASR::make_Print_t(al,
+                            x.base.base.loc, nullptr, tmp_vec.p, tmp_vec.size(),
+                            x.m_separator, space));
+                    }
                     print_tmp.clear();
                     pass_result.push_back(al, print_stmt);
-
                 }
-                if (ASR::is_a<ASR::List_t>(*ASRUtils::expr_type(x.m_values[i])))
-                    print_list_helper(x.m_values[i], x.m_separator, nullptr, x.base.base.loc);
-                else
-                    print_tuple_helper(x.m_values[i], x.m_separator, nullptr, x.base.base.loc);
+                if (ASR::is_a<ASR::List_t>(*ASRUtils::expr_type(x.m_values[i]))){
+                    if (x.m_separator) {
+                        if (i == x.n_values - 1) {
+                            print_list_helper(x.m_values[i], x.m_separator, x.m_end, x.base.base.loc);
+                        } else {
+                            print_list_helper(x.m_values[i], x.m_separator, x.m_separator, x.base.base.loc);
+                        }
+                    } else {
+                        if (i == x.n_values - 1) {
+                            print_list_helper(x.m_values[i], x.m_separator, x.m_end, x.base.base.loc);
+                        } else {
+                            print_list_helper(x.m_values[i], x.m_separator, space, x.base.base.loc);
+                        }
+                    }
+                } else {
+                    if (x.m_separator){
+                        if (i == x.n_values - 1) {
+                            print_tuple_helper(x.m_values[i], x.m_separator, x.m_end, x.base.base.loc);
+                        } else {
+                            print_tuple_helper(x.m_values[i], x.m_separator, x.m_separator, x.base.base.loc);
+                        }
+                    } else {
+                        if (i == x.n_values - 1) {
+                            print_tuple_helper(x.m_values[i], x.m_separator, x.m_end, x.base.base.loc);
+                        } else {
+                            print_tuple_helper(x.m_values[i], x.m_separator, space, x.base.base.loc);
+                        }
+                    }
+                }
                 for (size_t j=0; j<print_pass_result_tmp.n; j++)
                     pass_result.push_back(al, print_pass_result_tmp[j]);
                 print_pass_result_tmp.n = 0;
diff --git a/tests/reference/asr-test_end_sep_keywords-2226a67.json b/tests/reference/asr-test_end_sep_keywords-2226a67.json
@@ -0,0 +1,13 @@
+{
+    "basename": "asr-test_end_sep_keywords-2226a67",
+    "cmd": "lpython --show-asr --no-color {infile} -o {outfile}",
+    "infile": "tests/test_end_sep_keywords.py",
+    "infile_hash": "5ea30711228d4ebb64266988c1a706a3d64f196457b939ed3bf15ecf",
+    "outfile": null,
+    "outfile_hash": null,
+    "stdout": "asr-test_end_sep_keywords-2226a67.stdout",
+    "stdout_hash": "00212987e8f6240db98f86440a28ed351609019cfd55ede681b86b4a",
+    "stderr": null,
+    "stderr_hash": null,
+    "returncode": 0
+}
diff --git a/tests/reference/asr-test_end_sep_keywords-2226a67.stdout b/tests/reference/asr-test_end_sep_keywords-2226a67.stdout
@@ -0,0 +1,3 @@
+(TranslationUnit (SymbolTable 1 {_lpython_main_program: (Function (SymbolTable 4 {}) _lpython_main_program (FunctionType [] () Source Implementation () .false. .false. .false. .false. .false. [] [] .false.) [testEndSepKeywords] [] [(SubroutineCall 1 testEndSepKeywords () [] ())] () Public .false. .false.), main_program: (Program (SymbolTable 3 {}) main_program [] [(SubroutineCall 1 _lpython_main_program () [] ())]), testEndSepKeywords: (Function (SymbolTable 2 {}) testEndSepKeywords (FunctionType [] () Source Implementation () .false. .false. .false. .false. .false. [] [] .false.) [] [] [(Print () [(StringConstant "abc" (Character 1 3 () [])) (StringConstant "lmn" (Character 1 3 () [])) (StringConstant "pqr" (Character 1 3 () []))] () ()) (Print () [(StringConstant "abc" (Character 1 3 () [])) (StringConstant "lmn" (Character 1 3 () [])) (StringConstant "pqr" (Character 1 3 () []))] (StringConstant "+" (Character 1 1 () [])) ()) (Print () [(StringConstant "abc" (Character 1 3 () [])) (StringConstant "lmn" (Character 1 3 () [])) (StringConstant "pqr" (Character 1 3 () []))] () (StringConstant "xyz
+" (Character 1 4 () []))) (Print () [(StringConstant "abc" (Character 1 3 () [])) (StringConstant "lmn" (Character 1 3 () [])) (StringConstant "pqr" (Character 1 3 () []))] (StringConstant "+" (Character 1 1 () [])) (StringConstant "xyz
+" (Character 1 4 () [])))] () Public .false. .false.)}) [])
diff --git a/tests/reference/pass_print_list-print_02-d2853f6.json b/tests/reference/pass_print_list-print_02-d2853f6.json
@@ -6,7 +6,7 @@
     "outfile": null,
     "outfile_hash": null,
     "stdout": "pass_print_list-print_02-d2853f6.stdout",
-    "stdout_hash": "ec6264230ae2d462976ddb1e59b7d1dd14d5d8ee649e51c9255190f4",
+    "stdout_hash": "98056f8958731d7403712531de96e396cf9a3643c640e9d91376ab88",
     "stderr": null,
     "stderr_hash": null,
     "returncode": 0
diff --git a/tests/reference/pass_print_list-print_02-d2853f6.stdout b/tests/reference/pass_print_list-print_02-d2853f6.stdout
diff --git a/tests/test_end_sep_keywords.py b/tests/test_end_sep_keywords.py
@@ -0,0 +1,7 @@
+def testEndSepKeywords():
+    print("abc", "lmn", "pqr")
+    print("abc", "lmn", "pqr", sep = "+")
+    print("abc", "lmn", "pqr", end = "xyz\n")
+    print("abc", "lmn", "pqr", sep = "+", end = "xyz\n")
+
+testEndSepKeywords()
diff --git a/tests/tests.toml b/tests/tests.toml
@@ -279,6 +279,10 @@ filename = "c_interop1.py"
 asr = true
 c = true
 
+[[test]]
+filename = "test_end_sep_keywords.py"
+asr = true
+
 # integration_tests
 
 [[test]]

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	+(TranslationUnit (SymbolTable 1 {_lpython_main_program: (Function (SymbolTable 4 {}) _lpython_main_program (FunctionType [] () Source Implementation () .false. .false. .false. .false. .false. [] [] .false.) [testEndSepKeywords] [] [(SubroutineCall 1 testEndSepKeywords () [] ())] () Public .false. .false.), main_program: (Program (SymbolTable 3 {}) main_program [] [(SubroutineCall 1 _lpython_main_program () [] ())]), testEndSepKeywords: (Function (SymbolTable 2 {}) testEndSepKeywords (FunctionType [] () Source Implementation () .false. .false. .false. .false. .false. [] [] .false.) [] [] [(Print () [(StringConstant "abc" (Character 1 3 () [])) (StringConstant "lmn" (Character 1 3 () [])) (StringConstant "pqr" (Character 1 3 () []))] () ()) (Print () [(StringConstant "abc" (Character 1 3 () [])) (StringConstant "lmn" (Character 1 3 () [])) (StringConstant "pqr" (Character 1 3 () []))] (StringConstant "+" (Character 1 1 () [])) ()) (Print () [(StringConstant "abc" (Character 1 3 () [])) (StringConstant "lmn" (Character 1 3 () [])) (StringConstant "pqr" (Character 1 3 () []))] () (StringConstant "xyz
	`2`	`+" (Character 1 4 () []))) (Print () [(StringConstant "abc" (Character 1 3 () [])) (StringConstant "lmn" (Character 1 3 () [])) (StringConstant "pqr" (Character 1 3 () []))] (StringConstant "+" (Character 1 1 () [])) (StringConstant "xyz`
	`3`	`+" (Character 1 4 () [])))] () Public .false. .false.)}) [])`