diff --git a/starlark/ast.h b/starlark/ast.h index e5521b1..eeab81f 100644 --- a/starlark/ast.h +++ b/starlark/ast.h @@ -70,6 +70,12 @@ constexpr bool ListComp::operator==(ListComp const &o) const { return *element == *o.element && iterator_var == o.iterator_var && *iterable == *o.iterable; } +struct AssignStmt { + Expression target; + Expression value; + constexpr bool operator==(AssignStmt const &) const = default; +}; + struct ExpressionStmt { Expression expr; constexpr bool operator==(ExpressionStmt const &) const = default; @@ -81,7 +87,7 @@ struct LoadStmt { constexpr bool operator==(LoadStmt const &) const = default; }; -using Statement = std::variant; +using Statement = std::variant; struct Program { std::vector statements; diff --git a/starlark/parser.h b/starlark/parser.h index 2a81627..9b8d1da 100644 --- a/starlark/parser.h +++ b/starlark/parser.h @@ -45,6 +45,27 @@ class Parser { } if (auto expr = parse_expression(token); expr.has_value()) { + auto next = next_token(); + if (next.has_value() && std::holds_alternative(*next)) { + auto rhs_token = next_token(); + if (!rhs_token) { + std::cerr << "Unexpected end of input in assignment.\n"; + return std::nullopt; + } + + auto rhs = parse_expression(*rhs_token); + if (!rhs) { + std::cerr << "Failed to parse right-hand side of assignment.\n"; + return std::nullopt; + } + + program.statements.push_back( + AssignStmt{.target = std::move(*expr), .value = std::move(*rhs)}); + continue; + } else if (next.has_value()) { + reconsume(std::move(*next)); + } + program.statements.push_back(ExpressionStmt{.expr = std::move(*expr)}); continue; } diff --git a/starlark/parser_test.cc b/starlark/parser_test.cc index 35626f8..84ec8b3 100644 --- a/starlark/parser_test.cc +++ b/starlark/parser_test.cc @@ -217,6 +217,17 @@ int main() { }, }, }, + { + "A = B", + starlark::Program{ + .statements{ + starlark::AssignStmt{ + .target = starlark::Identifier{"A"}, + .value = starlark::Identifier{"B"}, + }, + }, + }, + }, }); // TODO(robinlinden): Return error codes from parser and use that here. @@ -259,6 +270,12 @@ int main() { "[e for y in z", // Parse error in iterable expression. "[e for y in '", + + // AssignStmt + // Tokenization error in target. + "A = \"", + // Parse error in value. + "A = foo(", }); etest::Suite s{};