Fix dict spreading in dict literal

Author: harupy

compiler/ast/src/ast_gen.rs

@@ -195,7 +195,7 @@ pub enum ExprKind<U = ()> {
         orelse: Box<Expr<U>>,
     },
     Dict {
-        keys: Vec<Expr<U>>,
+        keys: Vec<Option<Expr<U>>>,
         values: Vec<Expr<U>>,
     },
     Set {

compiler/parser/python.lalrpop

@@ -1138,42 +1138,42 @@ Atom<Goal>: ast::Expr = {
     <location:@L> "{" <e:DictLiteralValues?> "}" <end_location:@R> => {
         let pairs = e.unwrap_or_default();
 
-        let (keys, values) = match pairs.iter().position(|(k,_)| k.is_none()) {
-            Some(unpack_idx) => {
-                let mut pairs = pairs;
-                let (keys, mut values): (_, Vec<_>) = pairs.drain(..unpack_idx).map(|(k, v)| (*k.unwrap(), v)).unzip();
-
-                fn build_map(items: &mut Vec<(ast::Expr, ast::Expr)>) -> ast::Expr {
-                    let location = items[0].0.location;
-                    let end_location = items[0].0.end_location;
-                    let (keys, values) = items.drain(..).unzip();
-                    ast::Expr {
-                        location,
-                        end_location,
-                        custom: (),
-                        node: ast::ExprKind::Dict { keys, values }
-                    }
-                }
-
-                let mut items = Vec::new();
-                for (key, value) in pairs.into_iter() {
-                    if let Some(key) = key {
-                        items.push((*key, value));
-                        continue;
-                    }
-                    if !items.is_empty() {
-                        values.push(build_map(&mut items));
-                    }
-                    values.push(value);
-                }
-                if !items.is_empty() {
-                    values.push(build_map(&mut items));
-                }
-                (keys, values)
-            },
-            None => pairs.into_iter().map(|(k, v)| (*k.unwrap(), v)).unzip()
-        };
-
+        // let (keys, values) = match pairs.iter().position(|(k,_)| k.is_none()) {
+        //     Some(unpack_idx) => {
+        //         let mut pairs = pairs;
+        //         let (keys, mut values): (_, Vec<_>) = pairs.drain(..unpack_idx).map(|(k, v)| (*k.unwrap(), v)).unzip();
+        //
+        //         fn build_map(items: &mut Vec<(ast::Expr, ast::Expr)>) -> ast::Expr {
+        //             let location = items[0].0.location;
+        //             let end_location = items[0].0.end_location;
+        //             let (keys, values) = items.drain(..).unzip();
+        //             ast::Expr {
+        //                 location,
+        //                 end_location,
+        //                 custom: (),
+        //                 node: ast::ExprKind::Dict { keys, values }
+        //             }
+        //         }
+        //
+        //         let mut items = Vec::new();
+        //         for (key, value) in pairs.into_iter() {
+        //             if let Some(key) = key {
+        //                 items.push((*key, value));
+        //                 continue;
+        //             }
+        //             if !items.is_empty() {
+        //                 values.push(build_map(&mut items));
+        //             }
+        //             values.push(value);
+        //         }
+        //         if !items.is_empty() {
+        //             values.push(build_map(&mut items));
+        //         }
+        //         (keys, values)
+        //     },
+        //     None => pairs.into_iter().map(|(k, v)| (k, v)).unzip()
+        // };
+        let (keys, values) = pairs.into_iter().map(|(k, v)| (k.map(|x| *x), v)).unzip();
         ast::Expr {
             location,
             end_location: Some(end_location),

compiler/parser/src/parser.rs

@@ -309,4 +309,10 @@ with (0 as a, 1 as b,): pass
             assert!(parse_program(source, "<test>").is_err());
         }
     }
+
+    #[test]
+    fn test_dict_containing_spread() {
+        let parse_ast = parse_expression(r#"{"k": "v", **d}"#, "<test>").unwrap();
+        insta::assert_debug_snapshot!(parse_ast);
+    }
 }