starlark: Support parsing slice expressions

starlark/ast.h

@@ -34,8 +34,17 @@ struct CallExpr;
 struct DictExpr;
 struct ListExpr;
 struct ListComp;
-using Expression =
-    std::variant<CallExpr, StringLiteral, IntLiteral, Identifier, ListComp, ListExpr, DictExpr>;
+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<Expression> target;
+    std::shared_ptr<Expression> 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;
 }

starlark/parser.h

@@ -292,9 +292,28 @@ class Parser {
                 .target = std::make_shared<Expression>(std::move(*operand)),
                 .args = std::move(*args),
             };
-        }
+        } else if (std::holds_alternative<token::LBracket>(*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<Expression>(std::move(*operand)),
+                .index = std::make_shared<Expression>(std::move(*index_expr)),
+            };
+        } else {
             reconsume(std::move(*next));
         }
 

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::Expression>(
+                                    starlark::IntLiteral{42}),
+                                .index =
+                                    std::make_shared<starlark::Expression>(starlark::IntLiteral{5}),
+                            },
+                        },
+                    },
+                },
+            },
+        },
     });
 
     // TODO(robinlinden): Return error codes from parser and use that here.