Improve AST serialization grammar for For nodes

roastduck · Jan 15, 2024 · da83797 · da83797
1 parent 051ab9e
commit da83797
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 8 deletions.
diff --git a/grammar/ast_lexer.g b/grammar/ast_lexer.g
@@ -6,7 +6,10 @@ Comment:    '/*' .*? '*/' -> skip;
 IF:         'if';
 ELSE:       'else';
 FOR:        'for';
-IN:         'in';
+FROM:       'from';
+UNTIL:      'until';
+STEP:       'step';
+LENGTH:     'length';
 ASSERT_TOKEN:     'assert';
 ASSUME:     'assume';
 FUNC:       'func';

diff --git a/grammar/ast_parser.g b/grammar/ast_parser.g
@@ -347,7 +347,7 @@ forProperty returns [Ref<ForProperty> property]
 
 for returns [Stmt node]
     : forProperty
-        FOR var IN begin=expr ':' end=expr ':' step=expr ':' len=expr
+        FOR var FROM begin=expr UNTIL end=expr STEP step=expr LENGTH len=expr
         LBRACE stmts RBRACE
       {
           $node = makeFor($var.name, $begin.node, $end.node, $step.node, $len.node,

diff --git a/include/serialize/print_ast.h b/include/serialize/print_ast.h
@@ -15,7 +15,8 @@ class PrintVisitor : public CodeGen<CodeGenStream> {
          hexFloat_ = false, parenDespitePriority_ = false,
          printSourceLocation_ = false;
     const std::unordered_set<std::string> keywords = {
-        "if", "else", "for", "in", "assert", "assume", "func", "true", "false",
+        "if",     "else",   "for",    "from", "until", "step",
+        "length", "assert", "assume", "func", "true",  "false",
     };
 
     /**

diff --git a/src/serialize/print_ast.cc b/src/serialize/print_ast.cc
@@ -657,14 +657,14 @@ void PrintVisitor::visit(const For &op) {
         os() << "@!prefer_libs" << std::endl;
     }
     makeIndent();
-    os() << prettyKeyword("for ") << prettyIterName(op->iter_)
-         << prettyKeyword(" in ");
+    os() << prettyKeyword("for") << " " << prettyIterName(op->iter_) << " "
+         << prettyKeyword("from") << " ";
     recur(op->begin_);
-    os() << " : ";
+    os() << " " << prettyKeyword("until") << " ";
     recur(op->end_);
-    os() << " : ";
+    os() << " " << prettyKeyword("step") << " ";
     recur(op->step_);
-    os() << " : ";
+    os() << " " << prettyKeyword("length") << " ";
     recur(op->len_);
     os() << " ";
     beginBlock();