diff --git a/starlark/ast.h b/starlark/ast.h index 03cdd83..b0e1e18 100644 --- a/starlark/ast.h +++ b/starlark/ast.h @@ -34,8 +34,17 @@ struct CallExpr; struct DictExpr; struct ListExpr; struct ListComp; -using Expression = - std::variant; +struct SliceExpr; + +using Expression = std::variant< + CallExpr, + StringLiteral, + IntLiteral, + Identifier, + ListComp, + ListExpr, + DictExpr, + SliceExpr>; struct Argument; @@ -51,6 +60,12 @@ struct DictExpr { constexpr bool operator==(DictExpr const &) const; }; +struct SliceExpr { + std::shared_ptr target; + std::shared_ptr index; + constexpr bool operator==(SliceExpr const &) const; +}; + // TODO(robinlinden): shared_ptr is silly here, but right now the ast has to be // copyable for some reason. struct ListComp { @@ -77,6 +92,10 @@ constexpr bool CallExpr::operator==(CallExpr const &o) const { constexpr bool DictExpr::operator==(DictExpr const &o) const { return entries == o.entries; } +constexpr bool SliceExpr::operator==(SliceExpr const &o) const { + return *target == *o.target && *index == *o.index; +} + constexpr bool ListComp::operator==(ListComp const &o) const { return *element == *o.element && iterator_var == o.iterator_var && *iterable == *o.iterable; } diff --git a/starlark/parser.h b/starlark/parser.h index 24c89b5..46278e2 100644 --- a/starlark/parser.h +++ b/starlark/parser.h @@ -292,9 +292,28 @@ class Parser { .target = std::make_shared(std::move(*operand)), .args = std::move(*args), }; - } + } else if (std::holds_alternative(*next)) { + next = next_token(); + if (!next) { + std::cerr << "Unexpected end of input after '['.\n"; + return std::nullopt; + } + + auto index_expr = parse_expression(*next); + if (!index_expr) { + std::cerr << "Failed to parse index expression.\n"; + return std::nullopt; + } - if (next) { + if (!expect_next_token(token::RBracket{})) { + return std::nullopt; + } + + return SliceExpr{ + .target = std::make_shared(std::move(*operand)), + .index = std::make_shared(std::move(*index_expr)), + }; + } else { reconsume(std::move(*next)); } diff --git a/starlark/parser_test.cc b/starlark/parser_test.cc index 0db7a51..44e875c 100644 --- a/starlark/parser_test.cc +++ b/starlark/parser_test.cc @@ -259,6 +259,23 @@ int main() { }, }, }, + { + "42[5]", + starlark::Program{ + .statements{ + starlark::ExpressionStmt{ + .expr{ + starlark::SliceExpr{ + .target = std::make_shared( + starlark::IntLiteral{42}), + .index = + std::make_shared(starlark::IntLiteral{5}), + }, + }, + }, + }, + }, + }, }); // TODO(robinlinden): Return error codes from parser and use that here.