Merge pull request lcompilers#59 from namannimmo10/str_repeat

Compute the result of string repetition
BrentMat · Jan 23, 2022 · ff4c13f · ff4c13f
2 parents 2f461bf + 66c662d
commit ff4c13f
Show file tree

Hide file tree

Showing 6 changed files with 54 additions and 10 deletions.
diff --git a/src/lpython/semantics/python_ast_to_asr.cpp b/src/lpython/semantics/python_ast_to_asr.cpp
@@ -8,6 +8,7 @@
 #include <bitset>
 #include <complex>
 #include <sstream>
+#include <iterator>
 
 #include <lpython/python_ast.h>
 #include <libasr/asr.h>
@@ -252,7 +253,6 @@ class SymbolTableVisitor : public CommonVisitor<SymbolTableVisitor> {
         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<SymbolTableVisitor> {
                 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<BodyVisitor> {
             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<ASR::Character_t>(left_type);
+                LFORTRAN_ASSERT(left_type2->n_dims == 0);
+                right_int = ASR::down_cast<ASR::ConstantInteger_t>(
+                                                   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::ttype_t>(
+                        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<ASR::Character_t>(right_type);
+                LFORTRAN_ASSERT(right_type2->n_dims == 0);
+                left_int = ASR::down_cast<ASR::ConstantInteger_t>(
+                                                   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::ttype_t>(
+                        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<ASR::ConstantString_t>(
+                                                ASRUtils::expr_value(left))->m_s :
+                                                ASR::down_cast<ASR::ConstantString_t>(
+                                                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<std::string>(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::expr_t>(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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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" ())) ())] [] () ())] [])