diff --git a/src/lpython/semantics/python_ast_to_asr.cpp b/src/lpython/semantics/python_ast_to_asr.cpp index 3442908ac4..4c50ece55e 100644 --- a/src/lpython/semantics/python_ast_to_asr.cpp +++ b/src/lpython/semantics/python_ast_to_asr.cpp @@ -8,6 +8,7 @@ #include #include #include +#include #include #include @@ -252,7 +253,6 @@ class SymbolTableVisitor : public CommonVisitor { if (parent_scope->scope.find(sym_name) != parent_scope->scope.end()) { throw SemanticError("Subroutine already defined", tmp->loc); } - bool is_pure = false, is_module = false; ASR::accessType s_access = ASR::accessType::Public; ASR::deftypeType deftype = ASR::deftypeType::Implementation; char *bindc_name=nullptr; @@ -283,6 +283,7 @@ class SymbolTableVisitor : public CommonVisitor { current_procedure_abi_type, s_access, deftype, bindc_name); } else { + bool is_pure = false, is_module = false; tmp = ASR::make_Subroutine_t( al, x.base.base.loc, /* a_symtab */ current_scope, @@ -843,13 +844,53 @@ class BodyVisitor : public CommonVisitor { tmp = ASR::make_StrOp_t(al, loc, left, ops, right, dest_type, value); return; + } else if ((right_is_int || left_is_int) && op == ASR::binopType::Mul) { // string repeat - dest_type = right_is_int ? left_type : right_type; ASR::stropType ops = ASR::stropType::Repeat; - tmp = ASR::make_StrOp_t(al, loc, left, ops, right, dest_type, - value); + int64_t left_int = 0, right_int = 0, dest_len = 0; + if (right_is_int) { + ASR::Character_t *left_type2 = ASR::down_cast(left_type); + LFORTRAN_ASSERT(left_type2->n_dims == 0); + right_int = ASR::down_cast( + ASRUtils::expr_value(right))->m_n; + dest_len = left_type2->m_len * right_int; + if (dest_len < 0) dest_len = 0; + dest_type = ASR::down_cast( + ASR::make_Character_t(al, loc, left_type2->m_kind, + dest_len, nullptr, nullptr, 0)); + } else if (left_is_int) { + ASR::Character_t *right_type2 = ASR::down_cast(right_type); + LFORTRAN_ASSERT(right_type2->n_dims == 0); + left_int = ASR::down_cast( + ASRUtils::expr_value(left))->m_n; + dest_len = right_type2->m_len * left_int; + if (dest_len < 0) dest_len = 0; + dest_type = ASR::down_cast( + ASR::make_Character_t(al, loc, right_type2->m_kind, + dest_len, nullptr, nullptr, 0)); + } + + if (ASRUtils::expr_value(left) != nullptr && ASRUtils::expr_value(right) != nullptr) { + char* str = right_is_int ? ASR::down_cast( + ASRUtils::expr_value(left))->m_s : + ASR::down_cast( + ASRUtils::expr_value(right))->m_s; + int64_t repeat = right_is_int ? right_int : left_int; + char* result; + std::ostringstream os; + std::fill_n(std::ostream_iterator(os), repeat, std::string(str)); + std::string result_s = os.str(); + Str s; + s.from_str_view(result_s); + result = s.c_str(al); + LFORTRAN_ASSERT((int64_t)strlen(result) == dest_len) + value = ASR::down_cast(ASR::make_ConstantString_t( + al, loc, result, dest_type)); + } + tmp = ASR::make_StrOp_t(al, loc, left, ops, right, dest_type, value); return; + } else if (ASRUtils::is_complex(*left_type) && ASRUtils::is_complex(*right_type)) { dest_type = left_type; } else { diff --git a/tests/expr5.py b/tests/expr5.py index 6175d39c8b..f789934442 100644 --- a/tests/expr5.py +++ b/tests/expr5.py @@ -1,7 +1,10 @@ def test_strOP(): s: str s = 'a' * 2 + s = 'a' * -1 + s = 'test' * 5 s = 4 * 'bb' + s = -40 * 'bb' s = 3 * 'a' * 3 s = '3' + '4' s += 'test' diff --git a/tests/reference/asr-expr5-645ffcc.json b/tests/reference/asr-expr5-645ffcc.json index ecf80dcd30..d103645fff 100644 --- a/tests/reference/asr-expr5-645ffcc.json +++ b/tests/reference/asr-expr5-645ffcc.json @@ -2,11 +2,11 @@ "basename": "asr-expr5-645ffcc", "cmd": "lpython --show-asr --no-color {infile} -o {outfile}", "infile": "tests/expr5.py", - "infile_hash": "cdcbd274cd2d736afa60bebc7a966987f6972e2193ff66fc250f214d", + "infile_hash": "19625f22929d2c7e9a6999512d23e20451fd27947989654afa899948", "outfile": null, "outfile_hash": null, "stdout": "asr-expr5-645ffcc.stdout", - "stdout_hash": "a2cc572a0079b68291093ef39919a45cd3ff7789b436136aea065358", + "stdout_hash": "0b3e67ef332227347842abbf75ef365fde85d1776b2baf1fbe56a374", "stderr": null, "stderr_hash": null, "returncode": 0 diff --git a/tests/reference/asr-expr5-645ffcc.stdout b/tests/reference/asr-expr5-645ffcc.stdout index ba3c54f49e..91c43e3908 100644 --- a/tests/reference/asr-expr5-645ffcc.stdout +++ b/tests/reference/asr-expr5-645ffcc.stdout @@ -1 +1 @@ -(TranslationUnit (SymbolTable 1 {test_strOP: (Subroutine (SymbolTable 2 {s: (Variable 2 s Local () () Default (Character 1 -2 () []) Source Public Required .false.)}) test_strOP [] [(= (Var 2 s) (StrOp (ConstantString "a" (Character 1 1 () [])) Repeat (ConstantInteger 2 (Integer 4 [])) (Character 1 1 () []) ()) ()) (= (Var 2 s) (StrOp (ConstantInteger 4 (Integer 4 [])) Repeat (ConstantString "bb" (Character 1 2 () [])) (Character 1 2 () []) ()) ()) (= (Var 2 s) (StrOp (StrOp (ConstantInteger 3 (Integer 4 [])) Repeat (ConstantString "a" (Character 1 1 () [])) (Character 1 1 () []) ()) Repeat (ConstantInteger 3 (Integer 4 [])) (Character 1 1 () []) ()) ()) (= (Var 2 s) (StrOp (ConstantString "3" (Character 1 1 () [])) Concat (ConstantString "4" (Character 1 1 () [])) (Character 1 2 () []) (ConstantString "34" (Character 1 2 () []))) ()) (= (Var 2 s) (StrOp (Var 2 s) Concat (ConstantString "test" (Character 1 4 () [])) (Character 1 2 () []) ()) ()) (= (Var 2 s) (StrOp (StrOp (ConstantString "test" (Character 1 4 () [])) Concat (ConstantString "test" (Character 1 4 () [])) (Character 1 8 () []) (ConstantString "testtest" (Character 1 8 () []))) Concat (ConstantString "test" (Character 1 4 () [])) (Character 1 12 () []) (ConstantString "testtesttest" (Character 1 12 () []))) ())] Source Public Implementation () .false. .false.)}) []) +(TranslationUnit (SymbolTable 1 {test_strOP: (Subroutine (SymbolTable 2 {s: (Variable 2 s Local () () Default (Character 1 -2 () []) Source Public Required .false.)}) test_strOP [] [(= (Var 2 s) (StrOp (ConstantString "a" (Character 1 1 () [])) Repeat (ConstantInteger 2 (Integer 4 [])) (Character 1 2 () []) (ConstantString "aa" (Character 1 2 () []))) ()) (= (Var 2 s) (StrOp (ConstantString "a" (Character 1 1 () [])) Repeat (ConstantInteger -1 (Integer 4 [])) (Character 1 0 () []) (ConstantString "" (Character 1 0 () []))) ()) (= (Var 2 s) (StrOp (ConstantString "test" (Character 1 4 () [])) Repeat (ConstantInteger 5 (Integer 4 [])) (Character 1 20 () []) (ConstantString "testtesttesttesttest" (Character 1 20 () []))) ()) (= (Var 2 s) (StrOp (ConstantInteger 4 (Integer 4 [])) Repeat (ConstantString "bb" (Character 1 2 () [])) (Character 1 8 () []) (ConstantString "bbbbbbbb" (Character 1 8 () []))) ()) (= (Var 2 s) (StrOp (ConstantInteger -40 (Integer 4 [])) Repeat (ConstantString "bb" (Character 1 2 () [])) (Character 1 0 () []) (ConstantString "" (Character 1 0 () []))) ()) (= (Var 2 s) (StrOp (StrOp (ConstantInteger 3 (Integer 4 [])) Repeat (ConstantString "a" (Character 1 1 () [])) (Character 1 3 () []) (ConstantString "aaa" (Character 1 3 () []))) Repeat (ConstantInteger 3 (Integer 4 [])) (Character 1 9 () []) (ConstantString "aaaaaaaaa" (Character 1 9 () []))) ()) (= (Var 2 s) (StrOp (ConstantString "3" (Character 1 1 () [])) Concat (ConstantString "4" (Character 1 1 () [])) (Character 1 2 () []) (ConstantString "34" (Character 1 2 () []))) ()) (= (Var 2 s) (StrOp (Var 2 s) Concat (ConstantString "test" (Character 1 4 () [])) (Character 1 2 () []) ()) ()) (= (Var 2 s) (StrOp (StrOp (ConstantString "test" (Character 1 4 () [])) Concat (ConstantString "test" (Character 1 4 () [])) (Character 1 8 () []) (ConstantString "testtest" (Character 1 8 () []))) Concat (ConstantString "test" (Character 1 4 () [])) (Character 1 12 () []) (ConstantString "testtesttest" (Character 1 12 () []))) ())] Source Public Implementation () .false. .false.)}) []) diff --git a/tests/reference/ast-expr5-bbc6e71.json b/tests/reference/ast-expr5-bbc6e71.json index 0c4d2dc20a..01e6b61f3f 100644 --- a/tests/reference/ast-expr5-bbc6e71.json +++ b/tests/reference/ast-expr5-bbc6e71.json @@ -2,11 +2,11 @@ "basename": "ast-expr5-bbc6e71", "cmd": "lpython --show-ast --no-color {infile} -o {outfile}", "infile": "tests/expr5.py", - "infile_hash": "cdcbd274cd2d736afa60bebc7a966987f6972e2193ff66fc250f214d", + "infile_hash": "19625f22929d2c7e9a6999512d23e20451fd27947989654afa899948", "outfile": null, "outfile_hash": null, "stdout": "ast-expr5-bbc6e71.stdout", - "stdout_hash": "b9129bde458b2d81f5ae4b891369dfba8387254f8511f22393e7b9f7", + "stdout_hash": "f8a7e4775bda8d9a164bc267473569da8fdd3c670fd741d7e9d3151b", "stderr": null, "stderr_hash": null, "returncode": 0 diff --git a/tests/reference/ast-expr5-bbc6e71.stdout b/tests/reference/ast-expr5-bbc6e71.stdout index f9defa8130..8bf788e4b2 100644 --- a/tests/reference/ast-expr5-bbc6e71.stdout +++ b/tests/reference/ast-expr5-bbc6e71.stdout @@ -1 +1 @@ -(Module [(FunctionDef test_strOP ([] [] [] [] [] [] []) [(AnnAssign (Name s Store) (Name str Load) () 1) (Assign [(Name s Store)] (BinOp (ConstantStr "a" ()) Mult (ConstantInt 2 ())) ()) (Assign [(Name s Store)] (BinOp (ConstantInt 4 ()) Mult (ConstantStr "bb" ())) ()) (Assign [(Name s Store)] (BinOp (BinOp (ConstantInt 3 ()) Mult (ConstantStr "a" ())) Mult (ConstantInt 3 ())) ()) (Assign [(Name s Store)] (BinOp (ConstantStr "3" ()) Add (ConstantStr "4" ())) ()) (AugAssign (Name s Store) Add (ConstantStr "test" ())) (Assign [(Name s Store)] (BinOp (BinOp (ConstantStr "test" ()) Add (ConstantStr "test" ())) Add (ConstantStr "test" ())) ())] [] () ())] []) +(Module [(FunctionDef test_strOP ([] [] [] [] [] [] []) [(AnnAssign (Name s Store) (Name str Load) () 1) (Assign [(Name s Store)] (BinOp (ConstantStr "a" ()) Mult (ConstantInt 2 ())) ()) (Assign [(Name s Store)] (BinOp (ConstantStr "a" ()) Mult (UnaryOp USub (ConstantInt 1 ()))) ()) (Assign [(Name s Store)] (BinOp (ConstantStr "test" ()) Mult (ConstantInt 5 ())) ()) (Assign [(Name s Store)] (BinOp (ConstantInt 4 ()) Mult (ConstantStr "bb" ())) ()) (Assign [(Name s Store)] (BinOp (UnaryOp USub (ConstantInt 40 ())) Mult (ConstantStr "bb" ())) ()) (Assign [(Name s Store)] (BinOp (BinOp (ConstantInt 3 ()) Mult (ConstantStr "a" ())) Mult (ConstantInt 3 ())) ()) (Assign [(Name s Store)] (BinOp (ConstantStr "3" ()) Add (ConstantStr "4" ())) ()) (AugAssign (Name s Store) Add (ConstantStr "test" ())) (Assign [(Name s Store)] (BinOp (BinOp (ConstantStr "test" ()) Add (ConstantStr "test" ())) Add (ConstantStr "test" ())) ())] [] () ())] [])