Added the ._match primitive to the parser to factor out common code.

Also, getting started on parsing 'for'

chapter2.py

@@ -30,7 +30,7 @@ class Lexer(object):
         self.buf = buf
         self.pos = 0
         self.lastchar = self.buf[0]
-    
+
     def tokens(self):
         while self.lastchar:
             # Skip whitespace
@@ -137,7 +137,7 @@ class PrototypeAST(ASTNode):
     def dump(self, indent=0):
         return '{0}{1}[{2}]'.format(
             ' ' * indent, self.__class__.__name__, ', '.join(self.argnames))
-    
+
 
 class FunctionAST(ASTNode):
     def __init__(self, proto, body):
@@ -175,6 +175,18 @@ class Parser(object):
     def _get_next_token(self):
         self.cur_tok = next(self.token_generator)
 
+    def _match(self, expected_kind, expected_value=None):
+        """Consume the current token; verify that it's of the expected kind.
+
+        If expected_kind == TokenKind.OPERATOR, verify the operator's value.
+        """
+        if (expected_kind == TokenKind.OPERATOR and
+            not self._cur_tok_is_operator(expected_value)):
+            raise ParseError('Expected "{0}"'.format(expected_value))
+        elif expected_kind != self.cur_tok.kind:
+            raise ParseError('Expected "{0}"'.format(expected_kind))
+        self._get_next_token()
+
     _precedence_map = {'<': 10, '+': 20, '-': 20, '*': 40}
 
     def _cur_tok_precedence(self):
@@ -198,7 +210,7 @@ class Parser(object):
         # If followed by a '(' it's a call; otherwise, a simple variable ref.
         if not self._cur_tok_is_operator('('):
             return VariableExprAST(id_name)
-        
+
         self._get_next_token()
         args = []
         if not self._cur_tok_is_operator(')'):
@@ -206,9 +218,7 @@ class Parser(object):
                 args.append(self._parse_expression())
                 if self._cur_tok_is_operator(')'):
                     break
-                if not self._cur_tok_is_operator(','):
+                self._match(TokenKind.OPERATOR, ',')
-                    raise ParseError('Expected ")" or "," in argument list')
-                self._get_next_token()
 
         self._get_next_token()  # consume the ')'
         return CallExprAST(id_name, args)
@@ -223,9 +233,7 @@ class Parser(object):
     def _parse_paren_expr(self):
         self._get_next_token()  # consume the '('
         expr = self._parse_expression()
-        if not self._cur_tok_is_operator(')'):
+        self._match(TokenKind.OPERATOR, ')')
-            raise ParseError('Expected ")"')
-        self._get_next_token()  # consume the ')'
         return expr
 
     # primary
@@ -284,20 +292,14 @@ class Parser(object):
 
     # prototype ::= id '(' id* ')'
     def _parse_prototype(self):
-        if self.cur_tok.kind != TokenKind.IDENTIFIER:
-            raise ParseError('Expected function name in prototype')
         name = self.cur_tok.value
-        self._get_next_token()  # consume the name
+        self._match(TokenKind.IDENTIFIER)
-        if not self._cur_tok_is_operator('('):
+        self._match(TokenKind.OPERATOR, '(')
-            raise ParseError('Expected "(" in prototype')
-        self._get_next_token()  # consume '('
         argnames = []
         while self.cur_tok.kind == TokenKind.IDENTIFIER:
             argnames.append(self.cur_tok.value)
             self._get_next_token()
-        if not self._cur_tok_is_operator(')'):
+        self._match(TokenKind.OPERATOR, ')')
-            raise ParseError('Expected ")" in prototype')
-        self._get_next_token()  # consume ')'
         return PrototypeAST(name, argnames)
 
     # external ::= 'extern' prototype
@@ -317,7 +319,7 @@ class Parser(object):
         expr = self._parse_expression()
         # Anonymous function
         return FunctionAST(PrototypeAST('', []), expr)
-        
+
 
 #---- Some unit tests ----#
 
@@ -343,13 +345,13 @@ class TestLexer(unittest.TestCase):
     def test_token_kinds(self):
         l = Lexer('10.1 def der extern foo (')
         self._assert_toks(
-            list(l.tokens()), 
+            list(l.tokens()),
             ['NUMBER', 'DEF', 'IDENTIFIER', 'EXTERN', 'IDENTIFIER',
              'OPERATOR', 'EOF'])
 
         l = Lexer('+- 1 2 22 22.4 a b2 C3d')
         self._assert_toks(
-            list(l.tokens()), 
+            list(l.tokens()),
             ['OPERATOR', 'OPERATOR', 'NUMBER', 'NUMBER', 'NUMBER', 'NUMBER',
              'IDENTIFIER', 'IDENTIFIER', 'IDENTIFIER', 'EOF'])
 
@@ -360,7 +362,7 @@ class TestLexer(unittest.TestCase):
             \t\t\t10
             ''')
         self._assert_toks(
-            list(l.tokens()), 
+            list(l.tokens()),
             ['DEF', 'IDENTIFIER', 'NUMBER', 'EOF'])
 
 

chapter3and4.py

@@ -160,7 +160,7 @@ class FunctionAST(ASTNode):
 
     def is_anonymous(self):
         return self.proto.name.startswith('_anon')
-        
+
     def dump(self, indent=0):
         s = '{0}{1}[{2}]\n'.format(
             ' ' * indent, self.__class__.__name__, self.proto.dump())
@@ -192,6 +192,18 @@ class Parser(object):
     def _get_next_token(self):
         self.cur_tok = next(self.token_generator)
 
+    def _match(self, expected_kind, expected_value=None):
+        """Consume the current token; verify that it's of the expected kind.
+
+        If expected_kind == TokenKind.OPERATOR, verify the operator's value.
+        """
+        if (expected_kind == TokenKind.OPERATOR and
+            not self._cur_tok_is_operator(expected_value)):
+            raise ParseError('Expected "{0}"'.format(expected_value))
+        elif expected_kind != self.cur_tok.kind:
+            raise ParseError('Expected "{0}"'.format(expected_kind))
+        self._get_next_token()
+
     _precedence_map = {'<': 10, '+': 20, '-': 20, '*': 40}
 
     def _cur_tok_precedence(self):
@@ -223,9 +235,7 @@ class Parser(object):
                 args.append(self._parse_expression())
                 if self._cur_tok_is_operator(')'):
                     break
-                if not self._cur_tok_is_operator(','):
+                self._match(TokenKind.OPERATOR, ',')
-                    raise ParseError('Expected ")" or "," in argument list')
-                self._get_next_token()
 
         self._get_next_token()  # consume the ')'
         return CallExprAST(id_name, args)
@@ -240,9 +250,7 @@ class Parser(object):
     def _parse_paren_expr(self):
         self._get_next_token()  # consume the '('
         expr = self._parse_expression()
-        if not self._cur_tok_is_operator(')'):
+        self._match(TokenKind.OPERATOR, ')')
-            raise ParseError('Expected ")"')
-        self._get_next_token()  # consume the ')'
         return expr
 
     # primary
@@ -301,20 +309,14 @@ class Parser(object):
 
     # prototype ::= id '(' id* ')'
     def _parse_prototype(self):
-        if self.cur_tok.kind != TokenKind.IDENTIFIER:
-            raise ParseError('Expected function name in prototype')
         name = self.cur_tok.value
-        self._get_next_token()  # consume the name
+        self._match(TokenKind.IDENTIFIER)
-        if not self._cur_tok_is_operator('('):
+        self._match(TokenKind.OPERATOR, '(')
-            raise ParseError('Expected "(" in prototype')
-        self._get_next_token()  # consume '('
         argnames = []
         while self.cur_tok.kind == TokenKind.IDENTIFIER:
             argnames.append(self.cur_tok.value)
             self._get_next_token()
-        if not self._cur_tok_is_operator(')'):
+        self._match(TokenKind.OPERATOR, ')')
-            raise ParseError('Expected ")" in prototype')
-        self._get_next_token()  # consume ')'
         return PrototypeAST(name, argnames)
 
     # external ::= 'extern' prototype
@@ -332,6 +334,7 @@ class Parser(object):
     # toplevel ::= expression
     def _parse_toplevel_expression(self):
         expr = self._parse_expression()
+        # Anonymous function
         return FunctionAST.create_anonymous(expr)
 
 
@@ -402,7 +405,7 @@ class LLVMCodeGenerator(object):
             raise CodegenError('Call argument length mismatch', node.callee)
         call_args = [self._codegen(arg) for arg in node.args]
         return self.builder.call(callee_func, call_args, 'calltmp')
-        
+
     def _codegen_PrototypeAST(self, node):
         funcname = node.name
         # Create a function type
@@ -461,7 +464,7 @@ class KaleidoscopeEvaluator(object):
         self.codegen = LLVMCodeGenerator()
 
         self.target = llvm.Target.from_default_triple()
- 
+
     def evaluate(self, codestr, optimize=True, llvmdump=False):
         """Evaluate code in codestr.
 
@@ -471,7 +474,7 @@ class KaleidoscopeEvaluator(object):
         # Parse the given code and generate code from it
         ast = Parser(codestr).parse_toplevel()
         self.codegen.generate_code(ast)
-        
+
         if llvmdump:
             print('======== Unoptimized LLVM IR')
             print(str(self.codegen.module))

chapter5.py

@@ -19,6 +19,8 @@ class TokenKind(Enum):
     IF = -7
     THEN = -8
     ELSE = -9
+    FOR = -10
+    IN = -11
 
 
 Token = namedtuple('Token', 'kind value')
@@ -146,6 +148,27 @@ class IfExprAST(ExprAST):
         return s
 
 
+class ForExprAST(ExprAST):
+    def __init__(self, start_expr, end_expr, step_expr, body):
+        self.start_expr = start_expr
+        self.end_expr = end_expr
+        self.step_expr = step_expr
+        self.body = body
+
+    def dump(self, indent=0):
+        prefix = ' ' * indent
+        s = '{0}{1}\n'.format(prefix, self._class__.__name__)
+        s += '{0} Start:\n{1}\n'.format(
+            prefix. self.start_expr.dump(indent + 2))
+        s += '{0} End:\n{1}\n'.format(
+            prefix. self.end_expr.dump(indent + 2))
+        s += '{0} Step:\n{1}\n'.format(
+            prefix. self.step_expr.dump(indent + 2))
+        s += '{0} Body:\n{1}\n'.format(
+            prefix. self.body.dump(indent + 2))
+        return s
+
+
 class CallExprAST(ExprAST):
     def __init__(self, callee, args):
         self.callee = callee
@@ -219,6 +242,18 @@ class Parser(object):
     def _get_next_token(self):
         self.cur_tok = next(self.token_generator)
 
+    def _match(self, expected_kind, expected_value=None):
+        """Consume the current token; verify that it's of the expected kind.
+
+        If expected_kind == TokenKind.OPERATOR, verify the operator's value.
+        """
+        if (expected_kind == TokenKind.OPERATOR and
+            not self._cur_tok_is_operator(expected_value)):
+            raise ParseError('Expected "{0}"'.format(expected_value))
+        elif expected_kind != self.cur_tok.kind:
+            raise ParseError('Expected "{0}"'.format(expected_kind))
+        self._get_next_token()
+
     _precedence_map = {'<': 10, '+': 20, '-': 20, '*': 40}
 
     def _cur_tok_precedence(self):
@@ -250,9 +285,7 @@ class Parser(object):
                 args.append(self._parse_expression())
                 if self._cur_tok_is_operator(')'):
                     break
-                if not self._cur_tok_is_operator(','):
+                self._match(TokenKind.OPERATOR, ',')
-                    raise ParseError('Expected ")" or "," in argument list')
-                self._get_next_token()
 
         self._get_next_token()  # consume the ')'
         return CallExprAST(id_name, args)
@@ -267,9 +300,7 @@ class Parser(object):
     def _parse_paren_expr(self):
         self._get_next_token()  # consume the '('
         expr = self._parse_expression()
-        if not self._cur_tok_is_operator(')'):
+        self._match(TokenKind.OPERATOR, ')')
-            raise ParseError('Expected ")"')
-        self._get_next_token()  # consume the ')'
         return expr
 
     # primary
@@ -277,6 +308,7 @@ class Parser(object):
     #   ::= numberexpr
     #   ::= parenexpr
     #   ::= ifexpr
+    #   ::= forexpr
     def _parse_primary(self):
         if self.cur_tok.kind == TokenKind.IDENTIFIER:
             return self._parse_identifier_expr()
@@ -286,6 +318,8 @@ class Parser(object):
             return self._parse_paren_expr()
         elif self.cur_tok.kind == TokenKind.IF:
             return self._parse_if_expr()
+        elif self.cur_tok.kind == TokenKind.FOR:
+            return self._parse_for_expr()
         else:
             raise ParseError('Unknown token when expecting an expression')
 
@@ -293,16 +327,16 @@ class Parser(object):
     def _parse_if_expr(self):
         self._get_next_token()  # consume the 'if'
         cond_expr = self._parse_expression()
-        if self.cur_tok.kind != TokenKind.THEN:
+        self._match(TokenKind.THEN)
-            raise ParseError('Expected "then" in ifexpr')
-        self._get_next_token()  # consume the 'then'
         then_expr = self._parse_expression()
-        if self.cur_tok.kind != TokenKind.ELSE:
+        self._match(TokenKind.ELSE)
-            raise ParseError('Expected "else" in ifexpr')
-        self._get_next_token()  # consume the 'else'
         else_expr = self._parse_expression()
         return IfExprAST(cond_expr, then_expr, else_expr)
         
+    # forexpr ::= 'for' identifier '=' expr ',' expr (',' expr)? 'in' expr
+    #def _parse_for_expr(self):
+
+        
     # binoprhs ::= (<binop> primary)*
     def _parse_binop_rhs(self, expr_prec, lhs):
         """Parse the right-hand-side of a binary expression.
@@ -345,20 +379,14 @@ class Parser(object):
 
     # prototype ::= id '(' id* ')'
     def _parse_prototype(self):
-        if self.cur_tok.kind != TokenKind.IDENTIFIER:
-            raise ParseError('Expected function name in prototype')
         name = self.cur_tok.value
-        self._get_next_token()  # consume the name
+        self._match(TokenKind.IDENTIFIER)
-        if not self._cur_tok_is_operator('('):
+        self._match(TokenKind.OPERATOR, '(')
-            raise ParseError('Expected "(" in prototype')
-        self._get_next_token()  # consume '('
         argnames = []
         while self.cur_tok.kind == TokenKind.IDENTIFIER:
             argnames.append(self.cur_tok.value)
             self._get_next_token()
-        if not self._cur_tok_is_operator(')'):
+        self._match(TokenKind.OPERATOR, ')')
-            raise ParseError('Expected ")" in prototype')
-        self._get_next_token()  # consume ')'
         return PrototypeAST(name, argnames)
 
     # external ::= 'extern' prototype