starlark: Support dict expressions

starlark/parser.h

@@ -29,9 +29,10 @@ struct Identifier {
 };
 
 struct CallExpr;
+struct DictExpr;
 struct ListExpr;
 
-using Expression = std::variant<CallExpr, StringLiteral, Identifier, ListExpr>;
+using Expression = std::variant<CallExpr, StringLiteral, Identifier, ListExpr, DictExpr>;
 
 struct Argument;
 
@@ -41,6 +42,12 @@ struct CallExpr {
     constexpr bool operator==(CallExpr const &) const = default;
 };
 
+struct DictExpr {
+    std::vector<std::pair<Expression, Expression>> entries;
+    // Clang 18 dies if this is defaulted.
+    constexpr bool operator==(DictExpr const &) const;
+};
+
 struct ListExpr {
     std::vector<Expression> elements;
     constexpr bool operator==(ListExpr const &) const = default;
@@ -52,6 +59,8 @@ struct Argument {
     constexpr bool operator==(Argument const &) const = default;
 };
 
+constexpr bool DictExpr::operator==(DictExpr const &o) const { return entries == o.entries; }
+
 struct ExpressionStmt {
     Expression expr;
     constexpr bool operator==(ExpressionStmt const &) const = default;
@@ -196,6 +205,71 @@ class Parser {
             return ListExpr{.elements = std::move(elements)};
         }
 
+        if (std::holds_alternative<token::LBrace>(token)) {
+            std::vector<std::pair<Expression, Expression>> entries;
+
+            while (true) {
+                auto maybe_token = next_token();
+                if (!maybe_token) {
+                    std::cerr << "Tokenization error in dict expression.\n";
+                    return std::nullopt;
+                }
+
+                auto &t = *maybe_token;
+
+                if (std::holds_alternative<token::RBrace>(t)) {
+                    break;
+                }
+
+                auto key_expr = parse_expression(t);
+                if (!key_expr) {
+                    std::cerr << "Failed to parse expression for dict key.\n";
+                    return std::nullopt;
+                }
+
+                auto colon_token = next_token();
+                if (!colon_token || !std::holds_alternative<token::Colon>(*colon_token)) {
+                    std::cerr << "Expected ':' after dict key.\n";
+                    return std::nullopt;
+                }
+
+                auto value_token = next_token();
+                if (!value_token) {
+                    std::cerr << "Unexpected end of input in dict expression.\n";
+                    return std::nullopt;
+                }
+
+                auto value_expr = parse_expression(*value_token);
+                if (!value_expr) {
+                    std::cerr << "Failed to parse expression for dict value.\n";
+                    return std::nullopt;
+                }
+
+                entries.emplace_back(std::move(*key_expr), std::move(*value_expr));
+
+                auto maybe_next = next_token();
+                if (!maybe_next) {
+                    std::cerr << "Tokenization error in dict expression.\n";
+                    return std::nullopt;
+                }
+
+                auto const &next_token = *maybe_next;
+                if (std::holds_alternative<token::RBrace>(next_token)) {
+                    break;
+                }
+
+                if (std::holds_alternative<token::Comma>(next_token)) {
+                    continue;
+                }
+
+                std::cerr << "Expected ',' or '}' in dict expression, got " << to_string(next_token)
+                          << ".\n";
+                return std::nullopt;
+            }
+
+            return DictExpr{.entries = std::move(entries)};
+        }
+
         std::cerr << "Unexpected token in expression: " << to_string(token) << ".\n";
         return std::nullopt;
     }

starlark/parser_test.cc

@@ -134,6 +134,46 @@ int main() {
                 },
             },
         },
+        {
+            R"({"foo": bar, baz(): "qux"})",
+            starlark::Program{
+                .statements{
+                    starlark::ExpressionStmt{
+                        .expr{
+                            starlark::DictExpr{
+                                .entries{
+                                    {
+                                        starlark::StringLiteral{"foo"},
+                                        starlark::Identifier{"bar"},
+                                    },
+                                    {
+                                        starlark::CallExpr{
+                                            .target = "baz",
+                                            .args{},
+                                        },
+                                        starlark::StringLiteral{"qux"},
+                                    },
+                                },
+                            },
+                        },
+                    },
+                },
+            },
+        },
+        {
+            "{}",
+            starlark::Program{
+                .statements{
+                    starlark::ExpressionStmt{
+                        .expr{
+                            starlark::DictExpr{
+                                .entries{},
+                            },
+                        },
+                    },
+                },
+            },
+        },
     });
 
     // TODO(robinlinden): Return error codes from parser and use that here.
@@ -151,6 +191,22 @@ int main() {
         R"(["foo" ")",
         // Unexpected token after element.
         R"(["foo" foo])",
+
+        // DictExpr
+        // Tokenization error.
+        R"({")",
+        // Parse error in key.
+        R"({not: "value"})",
+        // Missing colon after key.
+        R"({"key" "value"})",
+        // Tokenization error in value.
+        R"({"key": ")",
+        // Parse error in value.
+        R"({"key": not})",
+        // Tokenization error after entry.
+        R"({"key": "value" ")",
+        // Unexpected token after entry.
+        R"({"key": "value" foo})",
     });
 
     etest::Suite s{};