starlark: Support assign statements

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<LoadStmt, ExpressionStmt>;
+using Statement = std::variant<AssignStmt, LoadStmt, ExpressionStmt>;
 
 struct Program {
     std::vector<Statement> statements;

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<token::Equals>(*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;
             }

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{};