diff --git a/integration_tests/test_str_01.py b/integration_tests/test_str_01.py index a889a85d83..6be357aa3b 100644 --- a/integration_tests/test_str_01.py +++ b/integration_tests/test_str_01.py @@ -1,3 +1,5 @@ +from lpython import i32 + def f(): x: str x = "ok" @@ -58,9 +60,38 @@ def test_str_repeat(): assert a*3 == "XyzXyzXyz" assert a*2*3 == "XyzXyzXyzXyzXyzXyz" assert 3*a*3 == "XyzXyzXyzXyzXyzXyzXyzXyzXyz" - assert a*-1 == "" + b: str = a * -1 + assert b == "" assert len(a*(10**6)) == (3 * 10 ** 6) + # string repeat with a non-constant integer + s: str = "#" + n: i32 = 5 + + assert s * n == "#####" + assert n * s == "#####" + + assert "@" * n == "@@@@@" + assert "@#$%" * n == "@#$%@#$%@#$%@#$%@#$%" + + s = "@#$%" + assert n * s == "@#$%@#$%@#$%@#$%@#$%" + + n = 10 ** 6 + assert len(s * n) == (4 * 10 ** 6) + + s = "$" + m: i32 = 2 + n = 5 + t: str = s * m * n + assert t == "$$$$$$$$$$" + assert s * m * 2 == "$$$$" + assert 2 * (m + n) * s == "$$$$$$$$$$$$$$" + + t = 2 * (m + n) * "abc-" + assert t == "abc-abc-abc-abc-abc-abc-abc-abc-abc-abc-abc-abc-abc-abc-" + + def test_str_join(): a: str a = "," diff --git a/src/libasr/runtime/lfortran_intrinsics.c b/src/libasr/runtime/lfortran_intrinsics.c index ce6c6901a4..d6ea899619 100644 --- a/src/libasr/runtime/lfortran_intrinsics.c +++ b/src/libasr/runtime/lfortran_intrinsics.c @@ -2140,7 +2140,7 @@ LFORTRAN_API void _lfortran_strrepeat(char** s, int32_t n, char** dest) char* dest_char = (char*)malloc(f_len+trmn_size); if (s_len == 1) { - memset(dest_char, *(*s), n); + memset(dest_char, *(*s), f_len); } else { memcpy(dest_char, *s, s_len); int chars_copied = s_len; diff --git a/src/lpython/semantics/python_ast_to_asr.cpp b/src/lpython/semantics/python_ast_to_asr.cpp index 571883a84e..06da016050 100644 --- a/src/lpython/semantics/python_ast_to_asr.cpp +++ b/src/lpython/semantics/python_ast_to_asr.cpp @@ -2079,7 +2079,7 @@ class CommonVisitor : public AST::BaseVisitor { } else if ((right_is_int || left_is_int) && op == ASR::binopType::Mul) { // string repeat int64_t left_int = 0, right_int = 0, dest_len = 0; - if (right_is_int) { + if (right_is_int && ASRUtils::expr_value(right) != nullptr) { ASR::Character_t *left_type2 = ASR::down_cast( ASRUtils::type_get_past_array(left_type)); LCOMPILERS_ASSERT(ASRUtils::extract_n_dims_from_ttype(left_type) == 0); @@ -2090,7 +2090,7 @@ class CommonVisitor : public AST::BaseVisitor { dest_type = ASR::down_cast( ASR::make_Character_t(al, loc, left_type2->m_kind, dest_len, nullptr)); - } else if (left_is_int) { + } else if (left_is_int && ASRUtils::expr_value(left) != nullptr) { ASR::Character_t *right_type2 = ASR::down_cast( ASRUtils::type_get_past_array(right_type)); LCOMPILERS_ASSERT(ASRUtils::extract_n_dims_from_ttype(right_type) == 0); @@ -2101,6 +2101,9 @@ class CommonVisitor : public AST::BaseVisitor { dest_type = ASR::down_cast( ASR::make_Character_t(al, loc, right_type2->m_kind, dest_len, nullptr)); + } else { + dest_type = ASRUtils::TYPE(ASR::make_Character_t(al, + loc, 1, -1, nullptr)); } if (ASRUtils::expr_value(left) != nullptr && ASRUtils::expr_value(right) != nullptr) {