patch 8.2.2034: Vim9: list unpack in for statement not compiled yetv8.2.2034

Problem:    Vim9: list unpack in for statement not compiled yet.
Solution:   Compile list unpack. (closes #7345)

8 files changed, 269 insertions, 35 deletions

diff --git a/src/errors.h b/src/errors.h
index c7aad127b..69450eb33 100644
--- a/src/errors.h
+++ b/src/errors.h
@@ -21,6 +21,10 @@ EXTERN char e_invalid_command[]
 #ifdef FEAT_EVAL
 EXTERN char e_invalid_command_str[]
 	INIT(= N_("E476: Invalid command: %s"));
+EXTERN char e_list_value_has_more_items_than_targets[]
+	INIT(= N_("E710: List value has more items than targets"));
+EXTERN char e_list_value_does_not_have_enough_items[]
+	INIT(= N_("E711: List value does not have enough items"));
 EXTERN char e_cannot_slice_dictionary[]
 	INIT(= N_("E719: Cannot slice a Dictionary"));
 EXTERN char e_assert_fails_second_arg[]
@@ -305,3 +309,5 @@ EXTERN char e_using_bool_as_number[]
 	INIT(= N_("E1138: Using a Bool as a Number"));
 EXTERN char e_missing_matching_bracket_after_dict_key[]
 	INIT(= N_("E1139: Missing matching bracket after dict key"));
+EXTERN char e_for_argument_must_be_sequence_of_lists[]
+	INIT(= N_("E1140: For argument must be a sequence of lists"));
diff --git a/src/eval.c b/src/eval.c
index 33976d2a8..6fea43605 100644
--- a/src/eval.c
+++ b/src/eval.c
@@ -1397,11 +1397,11 @@ set_var_lval(
 	    ++lp->ll_n1;
 	}
 	if (ri != NULL)
-	    emsg(_("E710: List value has more items than target"));
+	    emsg(_(e_list_value_has_more_items_than_targets));
 	else if (lp->ll_empty2
 		? (lp->ll_li != NULL && lp->ll_li->li_next != NULL)
 		: lp->ll_n1 != lp->ll_n2)
-	    emsg(_("E711: List value has not enough items"));
+	    emsg(_(e_list_value_does_not_have_enough_items));
     }
     else
     {
diff --git a/src/testdir/test_vim9_disassemble.vim b/src/testdir/test_vim9_disassemble.vim
index 12a3c8269..ba230a2a5 100644
--- a/src/testdir/test_vim9_disassemble.vim
+++ b/src/testdir/test_vim9_disassemble.vim
@@ -1026,6 +1026,40 @@ def Test_disassemble_for_loop_eval()
         instr)
 enddef
 
+def ForLoopUnpack()
+  for [x1, x2] in [[1, 2], [3, 4]]
+    echo x1 x2
+  endfor
+enddef
+
+def Test_disassemble_for_loop_unpack()
+  var instr = execute('disassemble ForLoopUnpack')
+  assert_match('ForLoopUnpack\_s*' ..
+        'for \[x1, x2\] in \[\[1, 2\], \[3, 4\]\]\_s*' ..
+        '\d\+ STORE -1 in $0\_s*' ..
+        '\d\+ PUSHNR 1\_s*' ..
+        '\d\+ PUSHNR 2\_s*' ..
+        '\d\+ NEWLIST size 2\_s*' ..
+        '\d\+ PUSHNR 3\_s*' ..
+        '\d\+ PUSHNR 4\_s*' ..
+        '\d\+ NEWLIST size 2\_s*' ..
+        '\d\+ NEWLIST size 2\_s*' ..
+        '\d\+ FOR $0 -> 16\_s*' ..
+        '\d\+ UNPACK 2\_s*' ..
+        '\d\+ STORE $1\_s*' ..
+        '\d\+ STORE $2\_s*' ..
+        'echo x1 x2\_s*' ..
+        '\d\+ LOAD $1\_s*' ..
+        '\d\+ LOAD $2\_s*' ..
+        '\d\+ ECHO 2\_s*' ..
+        'endfor\_s*' ..
+        '\d\+ JUMP -> 8\_s*' ..
+        '\d\+ DROP\_s*' ..
+        '\d\+ PUSHNR 0\_s*' ..
+        '\d\+ RETURN',
+        instr)
+enddef
+
 let g:number = 42
 
 def TypeCast()
diff --git a/src/testdir/test_vim9_script.vim b/src/testdir/test_vim9_script.vim
index 970436e0d..a99560f4f 100644
--- a/src/testdir/test_vim9_script.vim
+++ b/src/testdir/test_vim9_script.vim
@@ -1862,6 +1862,44 @@ def Test_for_loop_fails()
   CheckDefFailure(['for i in range(3)', 'echo 3'], 'E170:')
 enddef
 
+def Test_for_loop_unpack()
+  var result = []
+  for [v1, v2] in [[1, 2], [3, 4]]
+    result->add(v1)
+    result->add(v2)
+  endfor
+  assert_equal([1, 2, 3, 4], result)
+
+  result = []
+  for [v1, v2; v3] in [[1, 2], [3, 4, 5, 6]]
+    result->add(v1)
+    result->add(v2)
+    result->add(v3)
+  endfor
+  assert_equal([1, 2, [], 3, 4, [5, 6]], result)
+
+  var lines =<< trim END
+      for [v1, v2] in [[1, 2, 3], [3, 4]]
+        echo v1 v2
+      endfor
+  END
+  CheckDefExecFailure(lines, 'E710:', 1)
+
+  lines =<< trim END
+      for [v1, v2] in [[1], [3, 4]]
+        echo v1 v2
+      endfor
+  END
+  CheckDefExecFailure(lines, 'E711:', 1)
+
+  lines =<< trim END
+      for [v1, v1] in [[1, 2], [3, 4]]
+        echo v1
+      endfor
+  END
+  CheckDefExecFailure(lines, 'E1017:', 1)
+enddef
+
 def Test_while_loop()
   var result = ''
   var cnt = 0
diff --git a/src/version.c b/src/version.c
index 380710166..a4ca493ca 100644
--- a/src/version.c
+++ b/src/version.c
@@ -751,6 +751,8 @@ static char *(features[]) =
 static int included_patches[] =
 {   /* Add new patch number below this line */
 /**/
+    2034,
+/**/
     2033,
 /**/
     2032,
diff --git a/src/vim9.h b/src/vim9.h
index 9685b9d57..1368b3843 100644
--- a/src/vim9.h
+++ b/src/vim9.h
@@ -146,6 +146,7 @@ typedef enum {
     ISN_CMDMOD,	    // set cmdmod
     ISN_CMDMOD_REV, // undo ISN_CMDMOD
 
+    ISN_UNPACK,	    // unpack list into items, uses isn_arg.unpack
     ISN_SHUFFLE,    // move item on stack up or down
     ISN_DROP	    // pop stack and discard value
 } isntype_T;
@@ -284,6 +285,12 @@ typedef struct {
     cmdmod_T	*cf_cmdmod;	// allocated
 } cmod_T;
 
+// arguments to ISN_UNPACK
+typedef struct {
+    int		unp_count;	// number of items to produce
+    int		unp_semicolon;	// last item gets list of remainder
+} unpack_T;
+
 /*
  * Instruction
  */
@@ -321,6 +328,7 @@ struct isn_S {
 	shuffle_T	    shuffle;
 	put_T		    put;
 	cmod_T		    cmdmod;
+	unpack_T	    unpack;
     } isn_arg;
 };
 
diff --git a/src/vim9compile.c b/src/vim9compile.c
index b870023ea..821242c05 100644
--- a/src/vim9compile.c
+++ b/src/vim9compile.c
@@ -1888,6 +1888,19 @@ generate_EXECCONCAT(cctx_T *cctx, int count)
     return OK;
 }
 
+    static int
+generate_UNPACK(cctx_T *cctx, int var_count, int semicolon)
+{
+    isn_T	*isn;
+
+    RETURN_OK_IF_SKIP(cctx);
+    if ((isn = generate_instr(cctx, ISN_UNPACK)) == NULL)
+	return FAIL;
+    isn->isn_arg.unpack.unp_count = var_count;
+    isn->isn_arg.unpack.unp_semicolon = semicolon;
+    return OK;
+}
+
 /*
  * Generate an instruction for any command modifiers.
  */
@@ -6323,12 +6336,12 @@ compile_endif(char_u *arg, cctx_T *cctx)
 }
 
 /*
- * compile "for var in expr"
+ * Compile "for var in expr":
  *
  * Produces instructions:
  *       PUSHNR -1
  *       STORE loop-idx		Set index to -1
- *       EVAL expr		Push result of "expr"
+ *       EVAL expr		result of "expr" on top of stack
  * top:  FOR loop-idx, end	Increment index, use list on bottom of stack
  *				- if beyond end, jump to "end"
  *				- otherwise get item from list and push it
@@ -6337,11 +6350,19 @@ compile_endif(char_u *arg, cctx_T *cctx)
  *       JUMP top		Jump back to repeat
  * end:	 DROP			Drop the result of "expr"
  *
+ * Compile "for [var1, var2] in expr" - as above, but instead of "STORE var":
+ *	 UNPACK 2		Split item in 2
+ *       STORE var1		Store item in "var1"
+ *       STORE var2		Store item in "var2"
  */
     static char_u *
-compile_for(char_u *arg, cctx_T *cctx)
+compile_for(char_u *arg_start, cctx_T *cctx)
 {
+    char_u	*arg;
+    char_u	*arg_end;
     char_u	*p;
+    int		var_count = 0;
+    int		semicolon = FALSE;
     size_t	varlen;
     garray_T	*instr = &cctx->ctx_instr;
     garray_T	*stack = &cctx->ctx_type_stack;
@@ -6349,18 +6370,12 @@ compile_for(char_u *arg, cctx_T *cctx)
     lvar_T	*loop_lvar;	// loop iteration variable
     lvar_T	*var_lvar;	// variable for "var"
     type_T	*vartype;
+    type_T	*item_type = &t_any;
+    int		idx;
 
-    // TODO: list of variables: "for [key, value] in dict"
-    // parse "var"
-    for (p = arg; eval_isnamec1(*p); ++p)
-	;
-    varlen = p - arg;
-    var_lvar = lookup_local(arg, varlen, cctx);
-    if (var_lvar != NULL)
-    {
-	semsg(_(e_variable_already_declared), arg);
-	return NULL;
-    }
+    p = skip_var_list(arg_start, TRUE, &var_count, &semicolon, FALSE);
+    if (var_count == 0)
+	var_count = 1;
 
     // consume "in"
     p = skipwhite(p);
@@ -6371,12 +6386,12 @@ compile_for(char_u *arg, cctx_T *cctx)
     }
     p = skipwhite(p + 2);
 
-
     scope = new_scope(cctx, FOR_SCOPE);
     if (scope == NULL)
 	return NULL;
 
-    // Reserve a variable to store the loop iteration counter.
+    // Reserve a variable to store the loop iteration counter and initialize it
+    // to -1.
     loop_lvar = reserve_local(cctx, (char_u *)"", 0, FALSE, &t_number);
     if (loop_lvar == NULL)
     {
@@ -6384,16 +6399,6 @@ compile_for(char_u *arg, cctx_T *cctx)
 	drop_scope(cctx);
 	return NULL;
     }
-
-    // Reserve a variable to store "var"
-    var_lvar = reserve_local(cctx, arg, varlen, FALSE, &t_any);
-    if (var_lvar == NULL)
-    {
-	// out of memory or used as an argument
-	drop_scope(cctx);
-	return NULL;
-    }
-
     generate_STORENR(cctx, loop_lvar->lv_idx, -1);
 
     // compile "expr", it remains on the stack until "endfor"
@@ -6403,6 +6408,7 @@ compile_for(char_u *arg, cctx_T *cctx)
 	drop_scope(cctx);
 	return NULL;
     }
+    arg_end = arg;
 
     // Now that we know the type of "var", check that it is a list, now or at
     // runtime.
@@ -6412,16 +6418,78 @@ compile_for(char_u *arg, cctx_T *cctx)
 	drop_scope(cctx);
 	return NULL;
     }
+
     if (vartype->tt_type == VAR_LIST && vartype->tt_member->tt_type != VAR_ANY)
-	var_lvar->lv_type = vartype->tt_member;
+    {
+	if (var_count == 1)
+	    item_type = vartype->tt_member;
+	else if (vartype->tt_member->tt_type == VAR_LIST
+		      && vartype->tt_member->tt_member->tt_type != VAR_ANY)
+	    item_type = vartype->tt_member->tt_member;
+    }
 
     // "for_end" is set when ":endfor" is found
     scope->se_u.se_for.fs_top_label = instr->ga_len;
-
     generate_FOR(cctx, loop_lvar->lv_idx);
-    generate_STORE(cctx, ISN_STORE, var_lvar->lv_idx, NULL);
 
-    return arg;
+    arg = arg_start;
+    if (var_count > 1)
+    {
+	generate_UNPACK(cctx, var_count, semicolon);
+	arg = skipwhite(arg + 1);	// skip white after '['
+
+	// the list item is replaced by a number of items
+	if (ga_grow(stack, var_count - 1) == FAIL)
+	{
+	    drop_scope(cctx);
+	    return NULL;
+	}
+	--stack->ga_len;
+	for (idx = 0; idx < var_count; ++idx)
+	{
+	    ((type_T **)stack->ga_data)[stack->ga_len] =
+				(semicolon && idx == 0) ? vartype : item_type;
+	    ++stack->ga_len;
+	}
+    }
+
+    for (idx = 0; idx < var_count; ++idx)
+    {
+	// TODO: use skip_var_one, also assign to @r, $VAR, etc.
+	p = arg;
+	while (eval_isnamec(*p))
+	    ++p;
+	varlen = p - arg;
+	var_lvar = lookup_local(arg, varlen, cctx);
+	if (var_lvar != NULL)
+	{
+	    semsg(_(e_variable_already_declared), arg);
+	    drop_scope(cctx);
+	    return NULL;
+	}
+
+	// Reserve a variable to store "var".
+	// TODO: check for type
+	var_lvar = reserve_local(cctx, arg, varlen, FALSE, &t_any);
+	if (var_lvar == NULL)
+	{
+	    // out of memory or used as an argument
+	    drop_scope(cctx);
+	    return NULL;
+	}
+
+	if (semicolon && idx == var_count - 1)
+	    var_lvar->lv_type = vartype;
+	else
+	    var_lvar->lv_type = item_type;
+	generate_STORE(cctx, ISN_STORE, var_lvar->lv_idx, NULL);
+
+	if (*p == ',' || *p == ';')
+	    ++p;
+	arg = skipwhite(p);
+    }
+
+    return arg_end;
 }
 
 /*
@@ -7957,6 +8025,7 @@ delete_instr(isn_T *isn)
 	case ISN_STRSLICE:
 	case ISN_THROW:
 	case ISN_TRY:
+	case ISN_UNPACK:
 	    // nothing allocated
 	    break;
     }
diff --git a/src/vim9execute.c b/src/vim9execute.c
index 8b614bfdc..323f2926b 100644
--- a/src/vim9execute.c
+++ b/src/vim9execute.c
@@ -2877,11 +2877,84 @@ call_def_function(
 		restore_cmdmod = FALSE;
 		break;
 
+	    case ISN_UNPACK:
+		{
+		    int		count = iptr->isn_arg.unpack.unp_count;
+		    int		semicolon = iptr->isn_arg.unpack.unp_semicolon;
+		    list_T	*l;
+		    listitem_T	*li;
+		    int		i;
+
+		    // Check there is a valid list to unpack.
+		    tv = STACK_TV_BOT(-1);
+		    if (tv->v_type != VAR_LIST)
+		    {
+			SOURCING_LNUM = iptr->isn_lnum;
+			emsg(_(e_for_argument_must_be_sequence_of_lists));
+			goto on_error;
+		    }
+		    l = tv->vval.v_list;
+		    if (l == NULL
+				|| l->lv_len < (semicolon ? count - 1 : count))
+		    {
+			SOURCING_LNUM = iptr->isn_lnum;
+			emsg(_(e_list_value_does_not_have_enough_items));
+			goto on_error;
+		    }
+		    else if (!semicolon && l->lv_len > count)
+		    {
+			SOURCING_LNUM = iptr->isn_lnum;
+			emsg(_(e_list_value_has_more_items_than_targets));
+			goto on_error;
+		    }
+
+		    CHECK_LIST_MATERIALIZE(l);
+		    if (GA_GROW(&ectx.ec_stack, count - 1) == FAIL)
+			goto failed;
+		    ectx.ec_stack.ga_len += count - 1;
+
+		    // Variable after semicolon gets a list with the remaining
+		    // items.
+		    if (semicolon)
+		    {
+			list_T	*rem_list =
+				  list_alloc_with_items(l->lv_len - count + 1);
+
+			if (rem_list == NULL)
+			    goto failed;
+			tv = STACK_TV_BOT(-count);
+			tv->vval.v_list = rem_list;
+			++rem_list->lv_refcount;
+			tv->v_lock = 0;
+			li = l->lv_first;
+			for (i = 0; i < count - 1; ++i)
+			    li = li->li_next;
+			for (i = 0; li != NULL; ++i)
+			{
+			    list_set_item(rem_list, i, &li->li_tv);
+			    li = li->li_next;
+			}
+			--count;
+		    }
+
+		    // Produce the values in reverse order, first item last.
+		    li = l->lv_first;
+		    for (i = 0; i < count; ++i)
+		    {
+			tv = STACK_TV_BOT(-i - 1);
+			copy_tv(&li->li_tv, tv);
+			li = li->li_next;
+		    }
+
+		    list_unref(l);
+		}
+		break;
+
 	    case ISN_SHUFFLE:
 		{
-		    typval_T	    tmp_tv;
-		    int		    item = iptr->isn_arg.shuffle.shfl_item;
-		    int		    up = iptr->isn_arg.shuffle.shfl_up;
+		    typval_T	tmp_tv;
+		    int		item = iptr->isn_arg.shuffle.shfl_item;
+		    int		up = iptr->isn_arg.shuffle.shfl_up;
 
 		    tmp_tv = *STACK_TV_BOT(-item);
 		    for ( ; up > 0 && item > 1; --up)
@@ -3606,6 +3679,10 @@ ex_disassemble(exarg_T *eap)
 		}
 	    case ISN_CMDMOD_REV: smsg("%4d CMDMOD_REV", current); break;
 
+	    case ISN_UNPACK: smsg("%4d UNPACK %d%s", current,
+			iptr->isn_arg.unpack.unp_count,
+			iptr->isn_arg.unpack.unp_semicolon ? " semicolon" : "");
+			      break;
 	    case ISN_SHUFFLE: smsg("%4d SHUFFLE %d up %d", current,
 					 iptr->isn_arg.shuffle.shfl_item,
 					 iptr->isn_arg.shuffle.shfl_up);