Store load arg trailing comma in AST

Summary:
{F1082857697}

Used in the following diff D48941531.

Reviewed By: cjhopman

Differential Revision: D48941532

fbshipit-source-id: a2b4ae4bd84c8769b60503e0f01d6f2f59e0df84

Author: stepancheg

starlark/src/typing/fill_types_for_lint.rs

@@ -291,7 +291,7 @@ impl<'a, 'v> GlobalTypesBuilder<'a, 'v> {
     }
 
     fn load(&mut self, load: &LoadP<CstPayload>) -> Result<(), InternalError> {
-        for LoadArgP { local, their } in &load.args {
+        for LoadArgP { local, their, .. } in &load.args {
             let ty = load.payload.get(their).cloned().unwrap_or_else(Ty::any);
             self.assign_ident_value(local, GlobalValue::ty(ty))?;
         }

starlark_lsp/src/definition.rs

@@ -573,7 +573,7 @@ impl LspModule {
                             path: load.module.node.to_owned(),
                         })
                     } else {
-                        load.args.iter().find_map(|LoadArgP { local, their }| {
+                        load.args.iter().find_map(|LoadArgP { local, their, .. }| {
                             if local.span.contains(pos) || their.span.contains(pos) {
                                 Some(IdentifierDefinition::LoadedLocation {
                                     source: codemap.resolve_span(their.span),

starlark_lsp/src/server.rs

@@ -957,7 +957,9 @@ impl<T: LspContext> Backend<T> {
                 let load_span = ast.ast.codemap.resolve_span(*load_span);
                 let mut load_args: Vec<(&str, &str)> = previously_loaded_symbols
                     .iter()
-                    .map(|LoadArgP { local, their }| (local.ident.as_str(), their.node.as_str()))
+                    .map(|LoadArgP { local, their, .. }| {
+                        (local.ident.as_str(), their.node.as_str())
+                    })
                     .collect();
                 load_args.push((symbol, symbol));
                 load_args.sort_by(|(_, a), (_, b)| a.cmp(b));

starlark_syntax/src/syntax/ast.rs

@@ -52,6 +52,10 @@ impl AstPayload for AstNoPayload {
     type TypeExprPayload = ();
 }
 
+/// `,` token.
+#[derive(Copy, Clone, Dupe, Debug)]
+pub struct Comma;
+
 pub type Expr = ExprP<AstNoPayload>;
 pub type TypeExpr = TypeExprP<AstNoPayload>;
 pub type AssignTarget = AssignTargetP<AstNoPayload>;
@@ -246,6 +250,8 @@ pub struct LoadArgP<P: AstPayload> {
     pub local: AstAssignIdentP<P>,
     /// `"y" in `x="y"`.
     pub their: AstString,
+    /// Trailing comma.
+    pub comma: Option<Spanned<Comma>>,
 }
 
 impl<P: AstPayload> LoadArgP<P> {

starlark_syntax/src/syntax/grammar.lalrpop

@@ -191,9 +191,15 @@ AssignStmt_: Stmt = <lhs:TestList> <ty:Type> <op:AssignOp> <rhs:TestList>
 ExprStmt: AstStmt = ASTS<ExprStmt_>;
 ExprStmt_: Stmt = <Test> => Stmt::Expression(<>);
 
+Comma: Spanned<Comma> = <@L> "," <@R> => Comma.ast(<>);
+
 LoadStmt: AstStmt = ASTS<LoadStmt_>;
-LoadStmt_: Stmt = "load" "(" <module:string> <args:("," <LoadStmtSyms>)*> ","? ")"
-        => grammar_util::check_load(<>, state);
+LoadStmt_: Stmt = {
+    "load" "(" <module:string> ")"
+        => grammar_util::check_load_0(<>, state),
+    "load" "(" <module:string> Comma <args:(<LoadStmtSyms> <Comma>)*> <last:(<LoadStmtSyms>)?> ")"
+        => grammar_util::check_load(<>, state)
+};
 
 LoadStmtBindingName: AstString = <identifier> "=";
 

starlark_syntax/src/syntax/grammar_util.rs

@@ -40,6 +40,7 @@ use crate::syntax::ast::AstParameter;
 use crate::syntax::ast::AstStmt;
 use crate::syntax::ast::AstString;
 use crate::syntax::ast::AstTypeExpr;
+use crate::syntax::ast::Comma;
 use crate::syntax::ast::DefP;
 use crate::syntax::ast::Expr;
 use crate::syntax::ast::ExprP;
@@ -189,20 +190,42 @@ pub fn check_def(
     })
 }
 
+pub(crate) fn check_load_0(module: AstString, parser_state: &mut ParserState) -> Stmt {
+    parser_state.errors.push(EvalException::new(
+        GrammarUtilError::LoadRequiresAtLeastTwoArguments.into(),
+        module.span,
+        parser_state.codemap,
+    ));
+    Stmt::Load(LoadP {
+        module,
+        args: Vec::new(),
+        payload: (),
+    })
+}
+
 pub(crate) fn check_load(
     module: AstString,
-    args: Vec<(AstAssignIdent, AstString)>,
+    args: Vec<((AstAssignIdent, AstString), Spanned<Comma>)>,
+    last: Option<(AstAssignIdent, AstString)>,
     parser_state: &mut ParserState,
 ) -> Stmt {
-    if args.is_empty() {
-        parser_state.errors.push(EvalException::new(
-            GrammarUtilError::LoadRequiresAtLeastTwoArguments.into(),
-            module.span,
-            parser_state.codemap,
-        ));
+    if args.is_empty() && last.is_none() {
+        return check_load_0(module, parser_state);
     }
 
-    let args = args.into_map(|(local, their)| LoadArgP { local, their });
+    let args = args
+        .into_iter()
+        .map(|((local, their), comma)| LoadArgP {
+            local,
+            their,
+            comma: Some(comma),
+        })
+        .chain(last.map(|(local, their)| LoadArgP {
+            local,
+            their,
+            comma: None,
+        }))
+        .collect();
 
     Stmt::Load(LoadP {
         module,

starlark_syntax/src/syntax/module.rs

@@ -183,7 +183,7 @@ impl AstModule {
                     symbols: load
                         .args
                         .iter()
-                        .map(|LoadArgP { local, their }| {
+                        .map(|LoadArgP { local, their, .. }| {
                             (local.node.ident.as_str(), their.node.as_str())
                         })
                         .collect(),

starlark_syntax/src/syntax/payload_map.rs

@@ -54,10 +54,12 @@ impl<A: AstPayload> LoadArgP<A> {
         let LoadArgP {
             local,
             their: remote,
+            comma,
         } = self;
         LoadArgP {
             local: local.into_map_payload(f),
             their: remote,
+            comma,
         }
     }
 }