Parse case expr list with ranges

suhanigarg29 · Jul 8, 2021 · 4ff7839 · 4ff7839
1 parent 6536f30
commit 4ff7839
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 15 deletions.
diff --git a/grammar/AST.asdl b/grammar/AST.asdl
@@ -333,10 +333,13 @@ bind = Bind(expr* args, keyword* kwargs)
 array_index = ArrayIndex(expr? left, expr? right, expr? step)
 
 case_stmt
-    = CaseStmt(expr* test, stmt* body)
-    | CaseStmt_Range(expr? start, expr? end, stmt* body)
+    = CaseStmt(case_cond* test, stmt* body)
     | CaseStmt_Default(stmt* body)
 
+case_cond
+    = CaseCondExpr(expr cond)
+    | CaseCondRange(expr? start, expr? end)
+
 rank_stmt
     = RankExpr(expr value, stmt* body)
     | RankStar(stmt* body)

diff --git a/src/lfortran/parser/parser.yy b/src/lfortran/parser/parser.yy
@@ -4,7 +4,7 @@
 %param {LFortran::Parser &p}
 %locations
 %glr-parser
-%expect    599 // shift/reduce conflicts
+%expect    593 // shift/reduce conflicts
 %expect-rr 97  // reduce/reduce conflicts
 
 // Uncomment this to get verbose error messages
@@ -355,6 +355,8 @@ void yyerror(YYLTYPE *yyloc, LFortran::Parser &p, const std::string &msg)
 %type <ast> select_rank_case_stmt
 %type <vec_ast> case_statements
 %type <ast> case_statement
+%type <vec_ast> case_conditions
+%type <ast> case_condition
 %type <ast> while_statement
 %type <ast> critical_statement
 %type <ast> do_statement
@@ -1442,14 +1444,23 @@ case_statements
     ;
 
 case_statement
-    : KW_CASE "(" expr_list ")" sep statements { $$ = CASE_STMT($3, $6, @$); }
-    | KW_CASE "(" expr ":" ")" sep statements { $$ = CASE_STMT2($3, $7, @$); }
-    | KW_CASE "(" ":" expr ")" sep statements { $$ = CASE_STMT3($4, $7, @$); }
-    | KW_CASE "(" expr ":" expr ")" sep statements {
-            $$ = CASE_STMT4($3, $5, $8, @$); }
+    : KW_CASE "(" case_conditions ")" sep statements {
+            $$ = CASE_STMT($3, $6, @$); }
     | KW_CASE KW_DEFAULT sep statements { $$ = CASE_STMT_DEFAULT($4, @$); }
     ;
 
+case_conditions
+    : case_conditions "," case_condition { $$ = $1; LIST_ADD($$, $3); }
+    | case_condition { LIST_NEW($$); LIST_ADD($$, $1); }
+    ;
+
+case_condition
+    : expr { $$ = CASE_EXPR($1, @$); }
+    | expr ":" { $$ = CASE_RANGE1($1, @$); }
+    | ":" expr { $$ = CASE_RANGE2($2, @$); }
+    | expr ":" expr { $$ = CASE_RANGE3($1, $3, @$); }
+    ;
+
 select_rank_statement
     : KW_SELECT KW_RANK "(" expr ")" sep select_rank_case_stmts
         KW_END KW_SELECT { $$ = SELECT_RANK1($4, $7, @$); }

diff --git a/src/lfortran/parser/semantics.h b/src/lfortran/parser/semantics.h
@@ -1586,16 +1586,16 @@ ast_t* COARRAY(Allocator &al, const ast_t *id,
         CASE_STMTS(body), body.size())
 
 #define CASE_STMT(cond, body, l) make_CaseStmt_t(p.m_a, l, \
-        EXPRS(cond), cond.size(), STMTS(body), body.size())
-#define CASE_STMT2(cond, body, l) make_CaseStmt_Range_t(p.m_a, l, \
-        EXPR(cond), nullptr, STMTS(body), body.size())
-#define CASE_STMT3(cond, body, l) make_CaseStmt_Range_t(p.m_a, l, \
-        nullptr, EXPR(cond), STMTS(body), body.size())
-#define CASE_STMT4(cond1, cond2, body, l) make_CaseStmt_Range_t(p.m_a, l, \
-        EXPR(cond1), EXPR(cond2), STMTS(body), body.size())
+        VEC_CAST(cond, case_cond), cond.size(), STMTS(body), body.size())
 #define CASE_STMT_DEFAULT(body, l) make_CaseStmt_Default_t(p.m_a, l, \
         STMTS(body), body.size())
 
+#define CASE_EXPR(cond, l) make_CaseCondExpr_t(p.m_a, l, EXPR(cond))
+#define CASE_RANGE1(cond, l) make_CaseCondRange_t(p.m_a, l, EXPR(cond), nullptr)
+#define CASE_RANGE2(cond, l) make_CaseCondRange_t(p.m_a, l, nullptr, EXPR(cond))
+#define CASE_RANGE3(cond1, cond2, l) make_CaseCondRange_t(p.m_a, l, \
+        EXPR(cond1), EXPR(cond2))
+
 #define SELECT_RANK1(sel, body, l) make_SelectRank_t(p.m_a, l, 0, nullptr, \
         nullptr, EXPR(sel), RANK_STMTS(body), body.size())
 #define SELECT_RANK2(assoc, sel, body, l) make_SelectRank_t(p.m_a, l, \