# HG changeset patch
# User Bram Moolenaar <Bram@vim.org>
# Date 1588435204 -7200
# Node ID e46e72aaff74a44210a1aec5bfeeaaa4a6fc3a73
# Parent  ef22504328019705b3a7cfc9fee26b3fe45a4f50
patch 8.2.0679: Vim9: incomplete support for closures

Commit: https://github.com/vim/vim/commit/bf67ea1af05cbb30cd8f0b665fb567c0dca79796
Author: Bram Moolenaar <Bram@vim.org>
Date:   Sat May 2 17:52:42 2020 +0200

    patch 8.2.0679: Vim9: incomplete support for closures

    Problem:    Vim9: incomplete support for closures.
    Solution:   At the end of a function copy arguments and local variables if
                they are still used by a referenced closure.

diff --git a/src/structs.h b/src/structs.h
--- a/src/structs.h
+++ b/src/structs.h
@@ -1563,8 +1563,6 @@ typedef struct
     int		uf_refcount;	// reference count, see func_name_refcount()
 
     funccall_T	*uf_scoped;	// l: local variables for closure
-    garray_T	*uf_ectx_stack;	// where compiled closure finds local vars
-    int		uf_ectx_frame;	// index of function frame in uf_ectx_stack
 
     char_u	*uf_name_exp;	// if "uf_name[]" starts with SNR the name with
 				// "<SNR>" as a string, otherwise NULL
@@ -1591,7 +1589,7 @@ typedef struct
 #define FIXVAR_CNT	12	// number of fixed variables
 
 /*
- * structure to hold info for a function that is currently being executed.
+ * Structure to hold info for a function that is currently being executed.
  */
 struct funccall_S
 {
diff --git a/src/testdir/test_vim9_func.vim b/src/testdir/test_vim9_func.vim
--- a/src/testdir/test_vim9_func.vim
+++ b/src/testdir/test_vim9_func.vim
@@ -650,4 +650,16 @@ def Test_closure_simple()
   assert_equal('some more', RefFunc({s -> local .. s}))
 enddef
 
+def MakeRef()
+  let local = 'some '
+  g:Ref = {s -> local .. s}
+enddef
+
+def Test_closure_ref_after_return()
+  MakeRef()
+  assert_equal('some thing', g:Ref('thing'))
+  unlet g:Ref
+enddef
+
+
 " vim: ts=8 sw=2 sts=2 expandtab tw=80 fdm=marker
diff --git a/src/version.c b/src/version.c
--- a/src/version.c
+++ b/src/version.c
@@ -747,6 +747,8 @@ static char *(features[]) =
 static int included_patches[] =
 {   /* Add new patch number below this line */
 /**/
+    679,
+/**/
     678,
 /**/
     677,
diff --git a/src/vim9.h b/src/vim9.h
--- a/src/vim9.h
+++ b/src/vim9.h
@@ -71,7 +71,7 @@ typedef enum {
     ISN_PCALL,	    // call partial, use isn_arg.pfunc
     ISN_PCALL_END,  // cleanup after ISN_PCALL with cpf_top set
     ISN_RETURN,	    // return, result is on top of stack
-    ISN_FUNCREF,    // push a function ref to dfunc isn_arg.number
+    ISN_FUNCREF,    // push a function ref to dfunc isn_arg.funcref
 
     // expression operations
     ISN_JUMP,	    // jump if condition is matched isn_arg.jump
@@ -218,6 +218,12 @@ typedef struct {
     int		ul_forceit;	// forceit flag
 } unlet_T;
 
+// arguments to ISN_FUNCREF
+typedef struct {
+    int		fr_func;	// function index
+    int		fr_var_idx;	// variable to store partial
+} funcref_T;
+
 /*
  * Instruction
  */
@@ -249,10 +255,26 @@ struct isn_S {
 	loadstore_T	    loadstore;
 	script_T	    script;
 	unlet_T		    unlet;
+	funcref_T	    funcref;
     } isn_arg;
 };
 
 /*
+ * Structure to hold the context of a compiled function, used by closures
+ * defined in that function.
+ */
+typedef struct funcstack_S
+{
+    garray_T	fs_ga;		// contains the stack, with:
+				// - arguments
+				// - frame
+				// - local variables
+
+    int		fs_refcount;	// nr of closures referencing this funcstack
+    int		fs_copyID;	// for garray_T collection
+} funcstack_T;
+
+/*
  * Info about a function defined with :def.  Used in "def_functions".
  */
 struct dfunc_S {
@@ -264,10 +286,19 @@ struct dfunc_S {
     isn_T	*df_instr;	    // function body to be executed
     int		df_instr_count;
 
+    garray_T	*df_ectx_stack;	    // where compiled closure finds local vars
+    int		df_ectx_frame;	    // index of function frame in uf_ectx_stack
+    funcstack_T	*df_funcstack;	    // copy of stack for closure, used after
+				    // closure context function returns
+
     int		df_varcount;	    // number of local variables
+    int		df_closure_count;   // number of closures created
 };
 
 // Number of entries used by stack frame for a function call.
+// - function index
+// - instruction index
+// - previous frame index
 #define STACK_FRAME_SIZE 3
 
 
diff --git a/src/vim9compile.c b/src/vim9compile.c
--- a/src/vim9compile.c
+++ b/src/vim9compile.c
@@ -116,6 +116,9 @@ struct cctx_S {
     garray_T	ctx_locals;	    // currently visible local variables
     int		ctx_locals_count;   // total number of local variables
 
+    int		ctx_closure_count;  // number of closures created in the
+				    // function
+
     garray_T	ctx_imports;	    // imported items
 
     int		ctx_skip;	    // when TRUE skip commands, when FALSE skip
@@ -1254,7 +1257,8 @@ generate_FUNCREF(cctx_T *cctx, int dfunc
     RETURN_OK_IF_SKIP(cctx);
     if ((isn = generate_instr(cctx, ISN_FUNCREF)) == NULL)
 	return FAIL;
-    isn->isn_arg.number = dfunc_idx;
+    isn->isn_arg.funcref.fr_func = dfunc_idx;
+    isn->isn_arg.funcref.fr_var_idx = cctx->ctx_closure_count++;
 
     if (ga_grow(stack, 1) == FAIL)
 	return FAIL;
@@ -6395,6 +6399,7 @@ compile_def_function(ufunc_T *ufunc, int
 	dfunc->df_instr = instr->ga_data;
 	dfunc->df_instr_count = instr->ga_len;
 	dfunc->df_varcount = cctx.ctx_locals_count;
+	dfunc->df_closure_count = cctx.ctx_closure_count;
 	if (cctx.ctx_outer_used)
 	    ufunc->uf_flags |= FC_CLOSURE;
     }
@@ -6620,6 +6625,23 @@ delete_def_function_contents(dfunc_T *df
 	    delete_instr(dfunc->df_instr + idx);
 	VIM_CLEAR(dfunc->df_instr);
     }
+    if (dfunc->df_funcstack != NULL)
+    {
+	// Decrease the reference count for the context of a closure.  If down
+	// to zero free it and clear the variables on the stack.
+	if (--dfunc->df_funcstack->fs_refcount == 0)
+	{
+	    garray_T	*gap = &dfunc->df_funcstack->fs_ga;
+	    typval_T	*stack = gap->ga_data;
+	    int		i;
+
+	    for (i = 0; i < gap->ga_len; ++i)
+		clear_tv(stack + i);
+	    ga_clear(gap);
+	    vim_free(dfunc->df_funcstack);
+	}
+	dfunc->df_funcstack = NULL;
+    }
 
     dfunc->df_deleted = TRUE;
 }
diff --git a/src/vim9execute.c b/src/vim9execute.c
--- a/src/vim9execute.c
+++ b/src/vim9execute.c
@@ -24,7 +24,7 @@
 
 // Structure put on ec_trystack when ISN_TRY is encountered.
 typedef struct {
-    int	    tcd_frame;		// ec_frame when ISN_TRY was encountered
+    int	    tcd_frame_idx;	// ec_frame_idx when ISN_TRY was encountered
     int	    tcd_catch_idx;	// instruction of the first catch
     int	    tcd_finally_idx;	// instruction of the finally block
     int	    tcd_caught;		// catch block entered
@@ -56,7 +56,7 @@ typedef struct {
  */
 typedef struct {
     garray_T	ec_stack;	// stack of typval_T values
-    int		ec_frame;	// index in ec_stack: context of ec_dfunc_idx
+    int		ec_frame_idx;	// index in ec_stack: context of ec_dfunc_idx
 
     garray_T	*ec_outer_stack;    // stack used for closures
     int		ec_outer_frame;	    // stack frame in ec_outer_stack
@@ -202,7 +202,8 @@ call_dfunc(int cdf_idx, int argcount_arg
 	iemsg("Argument count wrong?");
 	return FAIL;
     }
-    if (ga_grow(&ectx->ec_stack, arg_to_add + 3 + dfunc->df_varcount) == FAIL)
+    if (ga_grow(&ectx->ec_stack, arg_to_add + 3
+		       + dfunc->df_varcount + dfunc->df_closure_count) == FAIL)
 	return FAIL;
 
     // Move the vararg-list to below the missing optional arguments.
@@ -215,17 +216,16 @@ call_dfunc(int cdf_idx, int argcount_arg
     ectx->ec_stack.ga_len += arg_to_add;
 
     // Store current execution state in stack frame for ISN_RETURN.
-    // TODO: If the actual number of arguments doesn't match what the called
-    // function expects things go bad.
     STACK_TV_BOT(0)->vval.v_number = ectx->ec_dfunc_idx;
     STACK_TV_BOT(1)->vval.v_number = ectx->ec_iidx;
-    STACK_TV_BOT(2)->vval.v_number = ectx->ec_frame;
-    ectx->ec_frame = ectx->ec_stack.ga_len;
+    STACK_TV_BOT(2)->vval.v_number = ectx->ec_frame_idx;
+    ectx->ec_frame_idx = ectx->ec_stack.ga_len;
 
     // Initialize local variables
-    for (idx = 0; idx < dfunc->df_varcount; ++idx)
+    for (idx = 0; idx < dfunc->df_varcount + dfunc->df_closure_count; ++idx)
 	STACK_TV_BOT(STACK_FRAME_SIZE + idx)->v_type = VAR_UNKNOWN;
-    ectx->ec_stack.ga_len += STACK_FRAME_SIZE + dfunc->df_varcount;
+    ectx->ec_stack.ga_len += STACK_FRAME_SIZE
+				+ dfunc->df_varcount + dfunc->df_closure_count;
 
     // Set execution state to the start of the called function.
     ectx->ec_dfunc_idx = cdf_idx;
@@ -233,8 +233,8 @@ call_dfunc(int cdf_idx, int argcount_arg
     estack_push_ufunc(ETYPE_UFUNC, dfunc->df_ufunc, 1);
 
     // used for closures
-    ectx->ec_outer_stack = ufunc->uf_ectx_stack;
-    ectx->ec_outer_frame = ufunc->uf_ectx_frame;
+    ectx->ec_outer_stack = dfunc->df_ectx_stack;
+    ectx->ec_outer_frame = dfunc->df_ectx_frame;
 
     // Decide where to start execution, handles optional arguments.
     init_instr_idx(ufunc, argcount, ectx);
@@ -246,34 +246,122 @@ call_dfunc(int cdf_idx, int argcount_arg
 #define STACK_TV(idx) (((typval_T *)ectx->ec_stack.ga_data) + idx)
 
 /*
+ * Used when returning from a function: Check if any closure is still
+ * referenced.  If so then move the arguments and variables to a separate piece
+ * of stack to be used when the closure is called.
+ * When "free_arguments" is TRUE the arguments are to be freed.
+ * Returns FAIL when out of memory.
+ */
+    static int
+handle_closure_in_use(ectx_T *ectx, int free_arguments)
+{
+    dfunc_T	*dfunc = ((dfunc_T *)def_functions.ga_data)
+							  + ectx->ec_dfunc_idx;
+    int		argcount = ufunc_argcount(dfunc->df_ufunc);
+    int		top = ectx->ec_frame_idx - argcount;
+    int		idx;
+    typval_T	*tv;
+    int		closure_in_use = FALSE;
+
+    // Check if any created closure is still in use.
+    for (idx = 0; idx < dfunc->df_closure_count; ++idx)
+    {
+	tv = STACK_TV(ectx->ec_frame_idx + STACK_FRAME_SIZE
+						   + dfunc->df_varcount + idx);
+	if (tv->v_type == VAR_PARTIAL && tv->vval.v_partial->pt_refcount > 1)
+	    closure_in_use = TRUE;
+    }
+
+    if (closure_in_use)
+    {
+	funcstack_T *funcstack = ALLOC_CLEAR_ONE(funcstack_T);
+	typval_T    *stack;
+
+	// A closure is using the arguments and/or local variables.
+	// Move them to the called function.
+	if (funcstack == NULL)
+	    return FAIL;
+	funcstack->fs_ga.ga_len = argcount + STACK_FRAME_SIZE
+							  + dfunc->df_varcount;
+	stack = ALLOC_CLEAR_MULT(typval_T, funcstack->fs_ga.ga_len);
+	funcstack->fs_ga.ga_data = stack;
+	if (stack == NULL)
+	{
+	    vim_free(funcstack);
+	    return FAIL;
+	}
+
+	// Move or copy the arguments.
+	for (idx = 0; idx < argcount; ++idx)
+	{
+	    tv = STACK_TV(top + idx);
+	    if (free_arguments)
+	    {
+		*(stack + idx) = *tv;
+		tv->v_type = VAR_UNKNOWN;
+	    }
+	    else
+		copy_tv(tv, stack + idx);
+	}
+	// Move the local variables.
+	for (idx = 0; idx < dfunc->df_varcount; ++idx)
+	{
+	    tv = STACK_TV(ectx->ec_frame_idx + STACK_FRAME_SIZE + idx);
+	    *(stack + argcount + STACK_FRAME_SIZE + idx) = *tv;
+	    tv->v_type = VAR_UNKNOWN;
+	}
+
+	for (idx = 0; idx < dfunc->df_closure_count; ++idx)
+	{
+	    tv = STACK_TV(ectx->ec_frame_idx + STACK_FRAME_SIZE
+						   + dfunc->df_varcount + idx);
+	    if (tv->v_type == VAR_PARTIAL
+					&& tv->vval.v_partial->pt_refcount > 1)
+	    {
+		dfunc_T	*pt_dfunc = ((dfunc_T *)def_functions.ga_data)
+				   + tv->vval.v_partial->pt_func->uf_dfunc_idx;
+		++funcstack->fs_refcount;
+		pt_dfunc->df_funcstack = funcstack;
+		pt_dfunc->df_ectx_stack = &funcstack->fs_ga;
+		pt_dfunc->df_ectx_frame = ectx->ec_frame_idx - top;
+	    }
+	}
+    }
+
+    return OK;
+}
+
+/*
  * Return from the current function.
  */
-    static void
+    static int
 func_return(ectx_T *ectx)
 {
     int		idx;
-    dfunc_T	*dfunc;
-    int		top;
+    dfunc_T	*dfunc = ((dfunc_T *)def_functions.ga_data)
+							  + ectx->ec_dfunc_idx;
+    int		argcount = ufunc_argcount(dfunc->df_ufunc);
+    int		top = ectx->ec_frame_idx - argcount;
 
     // execution context goes one level up
     estack_pop();
 
-    // Clear the local variables and temporary values, but not
-    // the return value.
-    for (idx = ectx->ec_frame + STACK_FRAME_SIZE;
+    if (handle_closure_in_use(ectx, TRUE) == FAIL)
+	return FAIL;
+
+    // Clear the arguments.
+    for (idx = top; idx < ectx->ec_frame_idx; ++idx)
+	clear_tv(STACK_TV(idx));
+
+    // Clear local variables and temp values, but not the return value.
+    for (idx = ectx->ec_frame_idx + STACK_FRAME_SIZE;
 					idx < ectx->ec_stack.ga_len - 1; ++idx)
 	clear_tv(STACK_TV(idx));
 
-    // Clear the arguments.
-    dfunc = ((dfunc_T *)def_functions.ga_data) + ectx->ec_dfunc_idx;
-    top = ectx->ec_frame - ufunc_argcount(dfunc->df_ufunc);
-    for (idx = top; idx < ectx->ec_frame; ++idx)
-	clear_tv(STACK_TV(idx));
-
     // Restore the previous frame.
-    ectx->ec_dfunc_idx = STACK_TV(ectx->ec_frame)->vval.v_number;
-    ectx->ec_iidx = STACK_TV(ectx->ec_frame + 1)->vval.v_number;
-    ectx->ec_frame = STACK_TV(ectx->ec_frame + 2)->vval.v_number;
+    ectx->ec_dfunc_idx = STACK_TV(ectx->ec_frame_idx)->vval.v_number;
+    ectx->ec_iidx = STACK_TV(ectx->ec_frame_idx + 1)->vval.v_number;
+    ectx->ec_frame_idx = STACK_TV(ectx->ec_frame_idx + 2)->vval.v_number;
     dfunc = ((dfunc_T *)def_functions.ga_data) + ectx->ec_dfunc_idx;
     ectx->ec_instr = dfunc->df_instr;
 
@@ -282,6 +370,8 @@ func_return(ectx_T *ectx)
     idx = ectx->ec_stack.ga_len - 1;
     ectx->ec_stack.ga_len = top + 1;
     *STACK_TV_BOT(-1) = *STACK_TV(idx);
+
+    return OK;
 }
 
 #undef STACK_TV
@@ -498,7 +588,7 @@ call_def_function(
 {
     ectx_T	ectx;		// execution context
     int		argc = argc_arg;
-    int		initial_frame_ptr;
+    int		initial_frame_idx;
     typval_T	*tv;
     int		idx;
     int		ret = FAIL;
@@ -513,7 +603,7 @@ call_def_function(
 #define STACK_TV_BOT(idx) (((typval_T *)ectx.ec_stack.ga_data) + ectx.ec_stack.ga_len + idx)
 
 // Get pointer to a local variable on the stack.  Negative for arguments.
-#define STACK_TV_VAR(idx) (((typval_T *)ectx.ec_stack.ga_data) + ectx.ec_frame + STACK_FRAME_SIZE + idx)
+#define STACK_TV_VAR(idx) (((typval_T *)ectx.ec_stack.ga_data) + ectx.ec_frame_idx + STACK_FRAME_SIZE + idx)
 
 // Like STACK_TV_VAR but use the outer scope
 #define STACK_OUT_TV_VAR(idx) (((typval_T *)ectx.ec_outer_stack->ga_data) + ectx.ec_outer_frame + STACK_FRAME_SIZE + idx)
@@ -562,8 +652,8 @@ call_def_function(
 	    ++ectx.ec_stack.ga_len;
 
     // Frame pointer points to just after arguments.
-    ectx.ec_frame = ectx.ec_stack.ga_len;
-    initial_frame_ptr = ectx.ec_frame;
+    ectx.ec_frame_idx = ectx.ec_stack.ga_len;
+    initial_frame_idx = ectx.ec_frame_idx;
 
     // dummy frame entries
     for (idx = 0; idx < STACK_FRAME_SIZE; ++idx)
@@ -573,13 +663,14 @@ call_def_function(
     }
 
     {
-	// Reserve space for local variables.
+	// Reserve space for local variables and closure references.
 	dfunc_T	*dfunc = ((dfunc_T *)def_functions.ga_data)
 							 + ufunc->uf_dfunc_idx;
+	int	count = dfunc->df_varcount + dfunc->df_closure_count;
 
-	for (idx = 0; idx < dfunc->df_varcount; ++idx)
+	for (idx = 0; idx < count; ++idx)
 	    STACK_TV_VAR(idx)->v_type = VAR_UNKNOWN;
-	ectx.ec_stack.ga_len += dfunc->df_varcount;
+	ectx.ec_stack.ga_len += count;
 
 	ectx.ec_instr = dfunc->df_instr;
     }
@@ -623,7 +714,7 @@ call_def_function(
 	    // the current function.
 	    if (trystack->ga_len > 0)
 		trycmd = ((trycmd_T *)trystack->ga_data) + trystack->ga_len - 1;
-	    if (trycmd != NULL && trycmd->tcd_frame == ectx.ec_frame)
+	    if (trycmd != NULL && trycmd->tcd_frame_idx == ectx.ec_frame_idx)
 	    {
 		// jump to ":catch" or ":finally"
 		ectx.ec_in_catch = TRUE;
@@ -632,7 +723,7 @@ call_def_function(
 	    else
 	    {
 		// not inside try or need to return from current functions.
-		if (ectx.ec_frame == initial_frame_ptr)
+		if (ectx.ec_frame_idx == initial_frame_idx)
 		{
 		    // At the toplevel we are done.  Push a dummy return value.
 		    if (ga_grow(&ectx.ec_stack, 1) == FAIL)
@@ -642,10 +733,13 @@ call_def_function(
 		    tv->vval.v_number = 0;
 		    ++ectx.ec_stack.ga_len;
 		    need_rethrow = TRUE;
+		    if (handle_closure_in_use(&ectx, FALSE) == FAIL)
+			goto failed;
 		    goto done;
 		}
 
-		func_return(&ectx);
+		if (func_return(&ectx) == FAIL)
+		    goto failed;
 	    }
 	    continue;
 	}
@@ -1073,8 +1167,7 @@ call_def_function(
 		    }
 
 		    --ectx.ec_stack.ga_len;
-		    di = find_var_in_ht(ht, 0,
-					       iptr->isn_arg.string + 2, TRUE);
+		    di = find_var_in_ht(ht, 0, iptr->isn_arg.string + 2, TRUE);
 		    if (di == NULL)
 			store_var(iptr->isn_arg.string, STACK_TV_BOT(0));
 		    else
@@ -1289,7 +1382,8 @@ call_def_function(
 		    if (trystack->ga_len > 0)
 			trycmd = ((trycmd_T *)trystack->ga_data)
 							+ trystack->ga_len - 1;
-		    if (trycmd != NULL && trycmd->tcd_frame == ectx.ec_frame
+		    if (trycmd != NULL
+				  && trycmd->tcd_frame_idx == ectx.ec_frame_idx
 			    && trycmd->tcd_finally_idx != 0)
 		    {
 			// jump to ":finally"
@@ -1300,10 +1394,15 @@ call_def_function(
 		    {
 			// Restore previous function. If the frame pointer
 			// is zero then there is none and we are done.
-			if (ectx.ec_frame == initial_frame_ptr)
+			if (ectx.ec_frame_idx == initial_frame_idx)
+			{
+			    if (handle_closure_in_use(&ectx, FALSE) == FAIL)
+				goto failed;
 			    goto done;
+			}
 
-			func_return(&ectx);
+			if (func_return(&ectx) == FAIL)
+			    goto failed;
 		    }
 		}
 		break;
@@ -1312,27 +1411,49 @@ call_def_function(
 	    case ISN_FUNCREF:
 		{
 		    partial_T   *pt = NULL;
-		    dfunc_T	*dfunc;
+		    dfunc_T	*pt_dfunc;
 
 		    pt = ALLOC_CLEAR_ONE(partial_T);
 		    if (pt == NULL)
 			goto failed;
-		    dfunc = ((dfunc_T *)def_functions.ga_data)
-							+ iptr->isn_arg.number;
-		    pt->pt_func = dfunc->df_ufunc;
+		    if (ga_grow(&ectx.ec_stack, 1) == FAIL)
+		    {
+			vim_free(pt);
+			goto failed;
+		    }
+		    pt_dfunc = ((dfunc_T *)def_functions.ga_data)
+					       + iptr->isn_arg.funcref.fr_func;
+		    pt->pt_func = pt_dfunc->df_ufunc;
 		    pt->pt_refcount = 1;
-		    ++dfunc->df_ufunc->uf_refcount;
+		    ++pt_dfunc->df_ufunc->uf_refcount;
+
+		    if (pt_dfunc->df_ufunc->uf_flags & FC_CLOSURE)
+		    {
+			dfunc_T	*dfunc = ((dfunc_T *)def_functions.ga_data)
+							   + ectx.ec_dfunc_idx;
 
-		    if (dfunc->df_ufunc->uf_flags & FC_CLOSURE)
-		    {
-			// Closure needs to find local variables in the current
-			// stack.
-			dfunc->df_ufunc->uf_ectx_stack = &ectx.ec_stack;
-			dfunc->df_ufunc->uf_ectx_frame = ectx.ec_frame;
+			// The closure needs to find arguments and local
+			// variables in the current stack.
+			pt_dfunc->df_ectx_stack = &ectx.ec_stack;
+			pt_dfunc->df_ectx_frame = ectx.ec_frame_idx;
+
+			// If this function returns and the closure is still
+			// used, we need to make a copy of the context
+			// (arguments and local variables). Store a reference
+			// to the partial so we can handle that.
+			++pt->pt_refcount;
+			tv = STACK_TV_VAR(dfunc->df_varcount
+					   + iptr->isn_arg.funcref.fr_var_idx);
+			if (tv->v_type == VAR_PARTIAL)
+			{
+			    // TODO: use a garray_T on ectx.
+			    emsg("Multiple closures not supported yet");
+			    goto failed;
+			}
+			tv->v_type = VAR_PARTIAL;
+			tv->vval.v_partial = pt;
 		    }
 
-		    if (ga_grow(&ectx.ec_stack, 1) == FAIL)
-			goto failed;
 		    tv = STACK_TV_BOT(0);
 		    ++ectx.ec_stack.ga_len;
 		    tv->vval.v_partial = pt;
@@ -1410,7 +1531,7 @@ call_def_function(
 						     + ectx.ec_trystack.ga_len;
 		    ++ectx.ec_trystack.ga_len;
 		    ++trylevel;
-		    trycmd->tcd_frame = ectx.ec_frame;
+		    trycmd->tcd_frame_idx = ectx.ec_frame_idx;
 		    trycmd->tcd_catch_idx = iptr->isn_arg.try.try_catch;
 		    trycmd->tcd_finally_idx = iptr->isn_arg.try.try_finally;
 		    trycmd->tcd_caught = FALSE;
@@ -1472,10 +1593,15 @@ call_def_function(
 			{
 			    // Restore previous function. If the frame pointer
 			    // is zero then there is none and we are done.
-			    if (ectx.ec_frame == initial_frame_ptr)
+			    if (ectx.ec_frame_idx == initial_frame_idx)
+			    {
+				if (handle_closure_in_use(&ectx, FALSE) == FAIL)
+				    goto failed;
 				goto done;
+			    }
 
-			    func_return(&ectx);
+			    if (func_return(&ectx) == FAIL)
+				goto failed;
 			}
 		    }
 		}
@@ -1949,12 +2075,15 @@ done:
 
 failed:
     // When failed need to unwind the call stack.
-    while (ectx.ec_frame != initial_frame_ptr)
+    while (ectx.ec_frame_idx != initial_frame_idx)
 	func_return(&ectx);
 failed_early:
     current_sctx.sc_version = save_sc_version;
+
+    // Free all local variables, but not arguments.
     for (idx = 0; idx < ectx.ec_stack.ga_len; ++idx)
 	clear_tv(STACK_TV(idx));
+
     vim_free(ectx.ec_stack.ga_data);
     vim_free(ectx.ec_trystack.ga_data);
     return ret;
@@ -2309,9 +2438,10 @@ ex_disassemble(exarg_T *eap)
 	    case ISN_FUNCREF:
 		{
 		    dfunc_T	*df = ((dfunc_T *)def_functions.ga_data)
-							+ iptr->isn_arg.number;
+					       + iptr->isn_arg.funcref.fr_func;
 
-		    smsg("%4d FUNCREF %s", current, df->df_ufunc->uf_name);
+		    smsg("%4d FUNCREF %s $%d", current, df->df_ufunc->uf_name,
+			   iptr->isn_arg.funcref.fr_var_idx + df->df_varcount);
 		}
 		break;