starlark: Support list comprehensions

starlark/ast.h

@@ -5,6 +5,7 @@
 #ifndef STARLARK_AST_H_
 #define STARLARK_AST_H_
 
+#include <memory>
 #include <optional>
 #include <string>
 #include <utility>
@@ -26,8 +27,8 @@ struct Identifier {
 struct CallExpr;
 struct DictExpr;
 struct ListExpr;
-
-using Expression = std::variant<CallExpr, StringLiteral, Identifier, ListExpr, DictExpr>;
+struct ListComp;
+using Expression = std::variant<CallExpr, StringLiteral, Identifier, ListComp, ListExpr, DictExpr>;
 
 struct Argument;
 
@@ -43,6 +44,15 @@ struct DictExpr {
     constexpr bool operator==(DictExpr const &) const;
 };
 
+// TODO(robinlinden): shared_ptr is silly here, but right now the ast has to be
+// copyable for some reason.
+struct ListComp {
+    std::shared_ptr<Expression> element;
+    Identifier iterator_var;
+    std::shared_ptr<Expression> iterable;
+    constexpr bool operator==(ListComp const &) const;
+};
+
 struct ListExpr {
     std::vector<Expression> elements;
     constexpr bool operator==(ListExpr const &) const = default;
@@ -56,6 +66,10 @@ struct Argument {
 
 constexpr bool DictExpr::operator==(DictExpr const &o) const { return entries == o.entries; }
 
+constexpr bool ListComp::operator==(ListComp const &o) const {
+    return *element == *o.element && iterator_var == o.iterator_var && *iterable == *o.iterable;
+}
+
 struct ExpressionStmt {
     Expression expr;
     constexpr bool operator==(ExpressionStmt const &) const = default;

starlark/parser.h

@@ -120,9 +120,53 @@ class Parser {
                     return std::nullopt;
                 }
 
+                // On the first iteration, we check if this is a list comprehension.
+                auto maybe_next = next_token();
+                if (elements.empty() && maybe_next.has_value() &&
+                    std::holds_alternative<token::For>(*maybe_next)) {
+                    auto var_token = next_token();
+                    if (!var_token || !std::holds_alternative<token::Identifier>(*var_token)) {
+                        std::cerr << "Expected identifier in list comprehension.\n";
+                        return std::nullopt;
+                    }
+
+                    auto in_token = next_token();
+                    if (!in_token || !std::holds_alternative<token::In>(*in_token)) {
+                        std::cerr << "Expected 'in' in list comprehension.\n";
+                        return std::nullopt;
+                    }
+
+                    auto iterable_token = next_token();
+                    if (!iterable_token) {
+                        std::cerr << "Unexpected end of input in list comprehension.\n";
+                        return std::nullopt;
+                    }
+
+                    auto iterable_expr = parse_expression(*iterable_token);
+                    if (!iterable_expr) {
+                        std::cerr << "Failed to parse iterable expression in list comprehension.\n";
+                        return std::nullopt;
+                    }
+
+                    auto maybe_closing = next_token();
+                    if (!maybe_closing ||
+                        !std::holds_alternative<token::RBracket>(*maybe_closing)) {
+                        std::cerr << "Expected closing ']' in list comprehension.\n";
+                        return std::nullopt;
+                    }
+
+                    // TODO(robinlinden): Handle optional 'if' clause.
+
+                    return ListComp{
+                        .element = std::make_shared<Expression>(std::move(*element_expr)),
+                        .iterator_var =
+                            Identifier{.name = std::get<token::Identifier>(*var_token).name},
+                        .iterable = std::make_shared<Expression>(std::move(*iterable_expr)),
+                    };
+                }
+
                 elements.push_back(std::move(*element_expr));
 
-                auto maybe_next = next_token();
                 if (!maybe_next) {
                     std::cerr << "Tokenization error in list expression.\n";
                     return std::nullopt;

starlark/parser_test.cc

@@ -176,6 +176,47 @@ int main() {
                 },
             },
         },
+        {
+            "[x for x in y]",
+            starlark::Program{
+                .statements{
+                    starlark::ExpressionStmt{
+                        .expr{
+                            starlark::ListComp{
+                                .element{std::make_shared<starlark::Expression>(
+                                    starlark::Identifier{"x"})},
+                                .iterator_var{"x"},
+                                .iterable = std::make_shared<starlark::Expression>(
+                                    starlark::Identifier{"y"}),
+                            },
+                        },
+                    },
+                },
+            },
+        },
+        {
+            R"([x for x in ["a", "b"]])",
+            starlark::Program{
+                .statements{
+                    starlark::ExpressionStmt{
+                        .expr{
+                            starlark::ListComp{
+                                .element{std::make_shared<starlark::Expression>(
+                                    starlark::Identifier{"x"})},
+                                .iterator_var{"x"},
+                                .iterable =
+                                    std::make_shared<starlark::Expression>(starlark::ListExpr{
+                                        .elements{
+                                            starlark::Expression{starlark::StringLiteral{"a"}},
+                                            starlark::Expression{starlark::StringLiteral{"b"}},
+                                        },
+                                    }),
+                            },
+                        },
+                    },
+                },
+            },
+        },
     });
 
     // TODO(robinlinden): Return error codes from parser and use that here.
@@ -209,6 +250,15 @@ int main() {
         R"({"key": "value" ")",
         // Unexpected token after entry.
         R"({"key": "value" foo})",
+
+        // ListComp
+        // Abrupt end of input.
+        "[e for",
+        "[e for y",
+        "[e for y in",
+        "[e for y in z",
+        // Parse error in iterable expression.
+        "[e for y in '",
     });
 
     etest::Suite s{};