diff --git a/starlark/ast.h b/starlark/ast.h index 3c610d1..e5521b1 100644 --- a/starlark/ast.h +++ b/starlark/ast.h @@ -5,6 +5,7 @@ #ifndef STARLARK_AST_H_ #define STARLARK_AST_H_ +#include #include #include #include @@ -26,8 +27,8 @@ struct Identifier { struct CallExpr; struct DictExpr; struct ListExpr; - -using Expression = std::variant; +struct ListComp; +using Expression = std::variant; 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 element; + Identifier iterator_var; + std::shared_ptr iterable; + constexpr bool operator==(ListComp const &) const; +}; + struct ListExpr { std::vector 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; diff --git a/starlark/parser.h b/starlark/parser.h index 6352664..2a81627 100644 --- a/starlark/parser.h +++ b/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(*maybe_next)) { + auto var_token = next_token(); + if (!var_token || !std::holds_alternative(*var_token)) { + std::cerr << "Expected identifier in list comprehension.\n"; + return std::nullopt; + } + + auto in_token = next_token(); + if (!in_token || !std::holds_alternative(*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(*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(std::move(*element_expr)), + .iterator_var = + Identifier{.name = std::get(*var_token).name}, + .iterable = std::make_shared(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; diff --git a/starlark/parser_test.cc b/starlark/parser_test.cc index 9caf0cb..35626f8 100644 --- a/starlark/parser_test.cc +++ b/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::Identifier{"x"})}, + .iterator_var{"x"}, + .iterable = std::make_shared( + starlark::Identifier{"y"}), + }, + }, + }, + }, + }, + }, + { + R"([x for x in ["a", "b"]])", + starlark::Program{ + .statements{ + starlark::ExpressionStmt{ + .expr{ + starlark::ListComp{ + .element{std::make_shared( + starlark::Identifier{"x"})}, + .iterator_var{"x"}, + .iterable = + std::make_shared(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{};