patch 8.2.1819: Vim9: Memory leak when using a closure

Problem:    Vim9: Memory leak when using a closure.
Solution:   Compute the mininal refcount in the funcstack.  Reenable disabled
            tests.

src/eval.c

@@ -3984,21 +3984,12 @@ partial_free(partial_T *pt)
     else
 	func_ptr_unref(pt->pt_func);
 
+    // Decrease the reference count for the context of a closure.  If down
+    // to the minimum it may be time to free it.
     if (pt->pt_funcstack != NULL)
     {
-	// Decrease the reference count for the context of a closure.  If down
+	--pt->pt_funcstack->fs_refcount;
-	// to zero free it and clear the variables on the stack.
+	funcstack_check_refcount(pt->pt_funcstack);
-	if (--pt->pt_funcstack->fs_refcount == 0)
-	{
-	    garray_T	*gap = &pt->pt_funcstack->fs_ga;
-	    typval_T	*stack = gap->ga_data;
-
-	    for (i = 0; i < gap->ga_len; ++i)
-		clear_tv(stack + i);
-	    ga_clear(gap);
-	    vim_free(pt->pt_funcstack);
-	}
-	pt->pt_funcstack = NULL;
     }
 
     vim_free(pt);
@@ -4011,8 +4002,16 @@ partial_free(partial_T *pt)
     void
 partial_unref(partial_T *pt)
 {
-    if (pt != NULL && --pt->pt_refcount <= 0)
+    if (pt != NULL)
+    {
+	if (--pt->pt_refcount <= 0)
 	    partial_free(pt);
+
+	// If the reference count goes down to one, the funcstack may be the
+	// only reference and can be freed if no other partials reference it.
+	else if (pt->pt_refcount == 1 && pt->pt_funcstack != NULL)
+	    funcstack_check_refcount(pt->pt_funcstack);
+    }
 }
 
 /*

src/proto/vim9execute.pro

@@ -1,5 +1,6 @@
 /* vim9execute.c */
 void to_string_error(vartype_T vartype);
+void funcstack_check_refcount(funcstack_T *funcstack);
 int call_def_function(ufunc_T *ufunc, int argc_arg, typval_T *argv, partial_T *partial, typval_T *rettv);
 void ex_disassemble(exarg_T *eap);
 int tv2bool(typval_T *tv);

src/structs.h

@@ -1869,8 +1869,11 @@ typedef struct funcstack_S
 				// - arguments
 				// - frame
 				// - local variables
+    int		fs_var_offset;	// count of arguments + frame size == offset to
+				// local variables
 
     int		fs_refcount;	// nr of closures referencing this funcstack
+    int		fs_min_refcount; // nr of closures on this funcstack
     int		fs_copyID;	// for garray_T collection
 } funcstack_T;
 

src/testdir/test_vim9_disassemble.vim

@@ -436,42 +436,42 @@ def Test_disassemble_call()
         res)
 enddef
 
-" TODO: fix memory leak and enable again
+
-"def s:CreateRefs()
+def s:CreateRefs()
-"  var local = 'a'
+  var local = 'a'
-"  def Append(arg: string)
+  def Append(arg: string)
-"    local ..= arg
+    local ..= arg
-"  enddef
+  enddef
-"  g:Append = Append
+  g:Append = Append
-"  def Get(): string
+  def Get(): string
-"    return local
+    return local
-"  enddef
+  enddef
-"  g:Get = Get
+  g:Get = Get
-"enddef
+enddef
-"
+
-"def Test_disassemble_closure()
+def Test_disassemble_closure()
-"  CreateRefs()
+  CreateRefs()
-"  var res = execute('disass g:Append')
+  var res = execute('disass g:Append')
-"  assert_match('<lambda>\d\_s*' ..
+  assert_match('<lambda>\d\_s*' ..
-"        'local ..= arg\_s*' ..
+        'local ..= arg\_s*' ..
-"        '\d LOADOUTER $0\_s*' ..
+        '\d LOADOUTER $0\_s*' ..
-"        '\d LOAD arg\[-1\]\_s*' ..
+        '\d LOAD arg\[-1\]\_s*' ..
-"        '\d CONCAT\_s*' ..
+        '\d CONCAT\_s*' ..
-"        '\d STOREOUTER $0\_s*' ..
+        '\d STOREOUTER $0\_s*' ..
-"        '\d PUSHNR 0\_s*' ..
+        '\d PUSHNR 0\_s*' ..
-"        '\d RETURN',
+        '\d RETURN',
-"        res)
+        res)
-"
+
-"  res = execute('disass g:Get')
+  res = execute('disass g:Get')
-"  assert_match('<lambda>\d\_s*' ..
+  assert_match('<lambda>\d\_s*' ..
-"        'return local\_s*' ..
+        'return local\_s*' ..
-"        '\d LOADOUTER $0\_s*' ..
+        '\d LOADOUTER $0\_s*' ..
-"        '\d RETURN',
+        '\d RETURN',
-"        res)
+        res)
-"
+
-"  unlet g:Append
+  unlet g:Append
-"  unlet g:Get
+  unlet g:Get
-"enddef
+enddef
 
 
 def EchoArg(arg: string): string

src/testdir/test_vim9_func.vim

@@ -1330,32 +1330,31 @@ def Test_closure_using_argument()
   unlet g:UseVararg
 enddef
 
-" TODO: reenable after fixing memory leak
+def MakeGetAndAppendRefs()
-"def MakeGetAndAppendRefs()
+  var local = 'a'
-"  var local = 'a'
+
-"
+  def Append(arg: string)
-"  def Append(arg: string)
+    local ..= arg
-"    local ..= arg
+  enddef
-"  enddef
+  g:Append = Append
-"  g:Append = Append
+
-"
+  def Get(): string
-"  def Get(): string
+    return local
-"    return local
+  enddef
-"  enddef
+  g:Get = Get
-"  g:Get = Get
+enddef
-"enddef
+
-"
+def Test_closure_append_get()
-"def Test_closure_append_get()
+  MakeGetAndAppendRefs()
-"  MakeGetAndAppendRefs()
+  g:Get()->assert_equal('a')
-"  g:Get()->assert_equal('a')
+  g:Append('-b')
-"  g:Append('-b')
+  g:Get()->assert_equal('a-b')
-"  g:Get()->assert_equal('a-b')
+  g:Append('-c')
-"  g:Append('-c')
+  g:Get()->assert_equal('a-b-c')
-"  g:Get()->assert_equal('a-b-c')
+
-"
+  unlet g:Append
-"  unlet g:Append
+  unlet g:Get
-"  unlet g:Get
+enddef
-"enddef
 
 def Test_nested_closure()
   var local = 'text'
@@ -1389,20 +1388,19 @@ def Test_double_closure_fails()
   CheckScriptSuccess(lines)
 enddef
 
-" TODO: reenable after fixing memory leak
+def Test_nested_closure_used()
-"def Test_nested_closure_used()
+  var lines =<< trim END
-"  var lines =<< trim END
+      vim9script
-"      vim9script
+      def Func()
-"      def Func()
+        var x = 'hello'
-"        var x = 'hello'
+        var Closure = {-> x}
-"        var Closure = {-> x}
+        g:Myclosure = {-> Closure()}
-"        g:Myclosure = {-> Closure()}
+      enddef
-"      enddef
+      Func()
-"      Func()
+      assert_equal('hello', g:Myclosure())
-"      assert_equal('hello', g:Myclosure())
+  END
-"  END
+  CheckScriptSuccess(lines)
-"  CheckScriptSuccess(lines)
+enddef
-"enddef
 
 def Test_nested_closure_fails()
   var lines =<< trim END

src/version.c

@@ -750,6 +750,8 @@ static char *(features[]) =
 
 static int included_patches[] =
 {   /* Add new patch number below this line */
+/**/
+    1819,
 /**/
     1818,
 /**/

src/vim9execute.c

@@ -349,8 +349,8 @@ handle_closure_in_use(ectx_T *ectx, int free_arguments)
 	// Move them to the called function.
 	if (funcstack == NULL)
 	    return FAIL;
-	funcstack->fs_ga.ga_len = argcount + STACK_FRAME_SIZE
+	funcstack->fs_var_offset = argcount + STACK_FRAME_SIZE;
-							  + dfunc->df_varcount;
+	funcstack->fs_ga.ga_len = funcstack->fs_var_offset + dfunc->df_varcount;
 	stack = ALLOC_CLEAR_MULT(typval_T, funcstack->fs_ga.ga_len);
 	funcstack->fs_ga.ga_data = stack;
 	if (stack == NULL)
@@ -376,29 +376,22 @@ handle_closure_in_use(ectx_T *ectx, int free_arguments)
 	{
 	    tv = STACK_TV(ectx->ec_frame_idx + STACK_FRAME_SIZE + idx);
 
-	    // Do not copy a partial created for a local function.
+	    // A partial created for a local function, that is also used as a
-	    // TODO: This won't work if the closure actually uses it.  But when
+	    // local variable, has a reference count for the variable, thus
-	    // keeping it it gets complicated: it will create a reference cycle
+	    // will never go down to zero.  When all these refcounts are one
-	    // inside the partial, thus needs special handling for garbage
+	    // then the funcstack is unused.  We need to count how many we have
-	    // collection.
+	    // so we need when to check.
-	    // For now, decide on the reference count.
 	    if (tv->v_type == VAR_PARTIAL && tv->vval.v_partial != NULL)
 	    {
 		int	    i;
 
 		for (i = 0; i < closure_count; ++i)
-		{
+		    if (tv->vval.v_partial == ((partial_T **)gap->ga_data)[
-		    partial_T *pt = ((partial_T **)gap->ga_data)[gap->ga_len
+					      gap->ga_len - closure_count + i])
-							  - closure_count + i];
+			++funcstack->fs_min_refcount;
-
-		    if (tv->vval.v_partial == pt && pt->pt_refcount < 2)
-			break;
-		}
-		if (i < closure_count)
-		    continue;
 	    }
 
-	    *(stack + argcount + STACK_FRAME_SIZE + idx) = *tv;
+	    *(stack + funcstack->fs_var_offset + idx) = *tv;
 	    tv->v_type = VAR_UNKNOWN;
 	}
 
@@ -426,6 +419,43 @@ handle_closure_in_use(ectx_T *ectx, int free_arguments)
     return OK;
 }
 
+/*
+ * Called when a partial is freed or its reference count goes down to one.  The
+ * funcstack may be the only reference to the partials in the local variables.
+ * Go over all of them, the funcref and can be freed if all partials
+ * referencing the funcstack have a reference count of one.
+ */
+    void
+funcstack_check_refcount(funcstack_T *funcstack)
+{
+    int		    i;
+    garray_T	    *gap = &funcstack->fs_ga;
+    int		    done = 0;
+
+    if (funcstack->fs_refcount > funcstack->fs_min_refcount)
+	return;
+    for (i = funcstack->fs_var_offset; i < gap->ga_len; ++i)
+    {
+	typval_T *tv = ((typval_T *)gap->ga_data) + i;
+
+	if (tv->v_type == VAR_PARTIAL && tv->vval.v_partial != NULL
+		&& tv->vval.v_partial->pt_funcstack == funcstack
+		&& tv->vval.v_partial->pt_refcount == 1)
+	    ++done;
+    }
+    if (done == funcstack->fs_min_refcount)
+    {
+	typval_T	*stack = gap->ga_data;
+
+	// All partials referencing the funcstack have a reference count of
+	// one, thus the funcstack is no longer of use.
+	for (i = 0; i < gap->ga_len; ++i)
+	    clear_tv(stack + i);
+	vim_free(stack);
+	vim_free(funcstack);
+    }
+}
+
 /*
  * Return from the current function.
  */