From e12c6813efcd8ab060ea464e96087c80051e3d41 Mon Sep 17 00:00:00 2001
From: Jeong YunWon <jeong@youknowone.org>
Date: Tue, 7 May 2019 04:41:31 +0900
Subject: [PATCH] Add dict unpacking support for literal

---
 parser/src/ast.rs         |  2 +-
 parser/src/python.lalrpop |  9 +++++++--
 tests/snippets/dict.py    |  7 +++++++
 vm/src/compile.rs         | 18 +++++++++++++++---
 vm/src/stdlib/ast.rs      |  6 +++++-
 vm/src/symboltable.rs     |  6 +++++-
 6 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/parser/src/ast.rs b/parser/src/ast.rs
index e0a050138..ec0677fe8 100644
--- a/parser/src/ast.rs
+++ b/parser/src/ast.rs
@@ -195,7 +195,7 @@ pub enum Expression {
         elements: Vec<Expression>,
     },
     Dict {
-        elements: Vec<(Expression, Expression)>,
+        elements: Vec<(Option<Expression>, Expression)>,
     },
     Set {
         elements: Vec<Expression>,
diff --git a/parser/src/python.lalrpop b/parser/src/python.lalrpop
index 02fd9e0b9..8d6dd7cac 100644
--- a/parser/src/python.lalrpop
+++ b/parser/src/python.lalrpop
@@ -833,8 +833,8 @@ TestListComp2: ast::Expression = {
     },
 };
 
-TestDict: Vec<(ast::Expression, ast::Expression)> = {
-    <elements:OneOrMore<DictEntry>> <_trailing_comma:","?> => elements,
+TestDict: Vec<(Option<ast::Expression>, ast::Expression)> = {
+    <elements:OneOrMore<DictElement>> <_trailing_comma:","?> => elements,
 };
 
 TestDictComp: ast::Expression = {
@@ -850,6 +850,11 @@ DictEntry: (ast::Expression, ast::Expression) = {
     <e1: Test> ":" <e2: Test> => (e1, e2),
 };
 
+DictElement: (Option<ast::Expression>, ast::Expression) = {
+    <e:DictEntry> => (Some(e.0), e.1),
+    "**" <e:Expression> => (None, e),
+};
+
 TestSet: Vec<ast::Expression> = {
     <e1:OneOrMore<Test>> ","? => e1
 };
diff --git a/tests/snippets/dict.py b/tests/snippets/dict.py
index 7d67668c9..270d2c07e 100644
--- a/tests/snippets/dict.py
+++ b/tests/snippets/dict.py
@@ -191,3 +191,10 @@ assert x['c'] is None
 
 assert {1: None, "b": None} == dict.fromkeys([1, "b"])
 assert {1: 0, "b": 0} == dict.fromkeys([1, "b"], 0)
+
+x = {'a': 1, 'b': 1, 'c': 1}
+y = {'b': 2, 'c': 2, 'd': 2}
+z = {'c': 3, 'd': 3, 'e': 3}
+
+w = {1: 1, **x, 2: 2, **y, 3: 3, **z, 4: 4}
+assert w == {1: 1, 'a': 1, 'b': 2, 'c': 3, 2: 2, 'd': 3, 3: 3, 'e': 3, 4: 4}
diff --git a/vm/src/compile.rs b/vm/src/compile.rs
index 21f1402a6..b1964ca55 100644
--- a/vm/src/compile.rs
+++ b/vm/src/compile.rs
@@ -1256,13 +1256,25 @@ impl Compiler {
             }
             ast::Expression::Dict { elements } => {
                 let size = elements.len();
+                let has_double_star = elements.iter().any(|e| e.0.is_none());
                 for (key, value) in elements {
-                    self.compile_expression(key)?;
-                    self.compile_expression(value)?;
+                    if let Some(key) = key {
+                        self.compile_expression(key)?;
+                        self.compile_expression(value)?;
+                        if has_double_star {
+                            self.emit(Instruction::BuildMap {
+                                size: 1,
+                                unpack: false,
+                            });
+                        }
+                    } else {
+                        // dict unpacking
+                        self.compile_expression(value)?;
+                    }
                 }
                 self.emit(Instruction::BuildMap {
                     size,
-                    unpack: false,
+                    unpack: has_double_star,
                 });
             }
             ast::Expression::Slice { elements } => {
diff --git a/vm/src/stdlib/ast.rs b/vm/src/stdlib/ast.rs
index f27b5eb40..2ac36c81e 100644
--- a/vm/src/stdlib/ast.rs
+++ b/vm/src/stdlib/ast.rs
@@ -358,7 +358,11 @@ fn expression_to_ast(vm: &VirtualMachine, expression: &ast::Expression) -> PyRes
             let mut keys = Vec::new();
             let mut values = Vec::new();
             for (k, v) in elements {
-                keys.push(expression_to_ast(vm, k)?.into_object());
+                if let Some(k) = k {
+                    keys.push(expression_to_ast(vm, k)?.into_object());
+                } else {
+                    keys.push(vm.ctx.none());
+                }
                 values.push(expression_to_ast(vm, v)?.into_object());
             }
 
diff --git a/vm/src/symboltable.rs b/vm/src/symboltable.rs
index 1d4a72ed9..ea9ae52a0 100644
--- a/vm/src/symboltable.rs
+++ b/vm/src/symboltable.rs
@@ -404,7 +404,11 @@ impl SymbolTableBuilder {
             }
             ast::Expression::Dict { elements } => {
                 for (key, value) in elements {
-                    self.scan_expression(key)?;
+                    if let Some(key) = key {
+                        self.scan_expression(key)?;
+                    } else {
+                        // dict unpacking marker
+                    }
                     self.scan_expression(value)?;
                 }
             }