API: Context: save/restore

8 files changed, 218 insertions, 28 deletions

diff --git a/runtime/doc/eval.txt b/runtime/doc/eval.txt
index 01bd2a0c65..cfcfa8e929 100644
--- a/runtime/doc/eval.txt
+++ b/runtime/doc/eval.txt
@@ -9110,9 +9110,9 @@ functions.
 					*:fu* *:function* *E128* *E129* *E123*
 :fu[nction]		List all functions and their arguments.
 
-:fu[nction] {name}	List function {name}.
-			{name} can also be a |Dictionary| entry that is a
-			|Funcref|: >
+:fu[nction][!] {name}	List function {name}, annotated with line numbers
+			unless "!" is given.
+			{name} may be a |Dictionary| |Funcref| entry: >
 				:function dict.init
 
 :fu[nction] /{pattern}	List functions with a name matching {pattern}.
diff --git a/src/nvim/api/vim.c b/src/nvim/api/vim.c
index 3fd50ffaf8..ed6a28bcda 100644
--- a/src/nvim/api/vim.c
+++ b/src/nvim/api/vim.c
@@ -1293,6 +1293,10 @@ Dictionary nvim_get_context(Array types)
           int_types |= kCtxBuflist;
         } else if (strequal(current, "gvars")) {
           int_types |= kCtxGVars;
+        } else if (strequal(current, "sfuncs")) {
+          int_types |= kCtxSFuncs;
+        } else if (strequal(current, "funcs")) {
+          int_types |= kCtxFuncs;
         }
       }
     }
diff --git a/src/nvim/context.c b/src/nvim/context.c
index 09f2d7bd76..b2a2fd3fd9 100644
--- a/src/nvim/context.c
+++ b/src/nvim/context.c
@@ -5,20 +5,23 @@
 
 #include "nvim/context.h"
 #include "nvim/eval/encode.h"
+#include "nvim/ex_docmd.h"
 #include "nvim/option.h"
 #include "nvim/shada.h"
+#include "nvim/api/vim.h"
 #include "nvim/api/private/helpers.h"
 
 #ifdef INCLUDE_GENERATED_DECLARATIONS
 # include "context.c.generated.h"
 #endif
 
-int kCtxAll = (kCtxRegs | kCtxJumps | kCtxBuflist | kCtxGVars);
+int kCtxAll = (kCtxRegs | kCtxJumps | kCtxBuflist | kCtxGVars | kCtxSFuncs
+               | kCtxFuncs);
 
 static ContextVec ctx_stack = KV_INITIAL_VALUE;
 
 /// Clears and frees the context stack
-void free_ctx_stack(void)
+void ctx_free_all(void)
 {
   for (size_t i = 0; i < kv_size(ctx_stack); i++) {
     ctx_free(&kv_A(ctx_stack, i));
@@ -60,6 +63,9 @@ void ctx_free(Context *ctx)
   if (ctx->gvars.data) {
     msgpack_sbuffer_destroy(&ctx->gvars);
   }
+  if (ctx->funcs.items) {
+    api_free_array(ctx->funcs);
+  }
 }
 
 /// Saves the editor state to a context.
@@ -79,15 +85,24 @@ void ctx_save(Context *ctx, const int flags)
   if (flags & kCtxRegs) {
     ctx_save_regs(ctx);
   }
+
   if (flags & kCtxJumps) {
     ctx_save_jumps(ctx);
   }
+
   if (flags & kCtxBuflist) {
     ctx_save_buflist(ctx);
   }
+
   if (flags & kCtxGVars) {
     ctx_save_gvars(ctx);
   }
+
+  if (flags & kCtxFuncs) {
+    ctx_save_funcs(ctx, false);
+  } else if (flags & kCtxSFuncs) {
+    ctx_save_funcs(ctx, true);
+  }
 }
 
 /// Restores the editor state from a context.
@@ -117,16 +132,23 @@ bool ctx_restore(Context *ctx, const int flags)
   if (flags & kCtxRegs) {
     ctx_restore_regs(ctx);
   }
+
   if (flags & kCtxJumps) {
     ctx_restore_jumps(ctx);
   }
+
   if (flags & kCtxBuflist) {
     ctx_restore_buflist(ctx);
   }
+
   if (flags & kCtxGVars) {
     ctx_restore_gvars(ctx);
   }
 
+  if (flags & kCtxFuncs) {
+    ctx_restore_funcs(ctx);
+  }
+
   if (free_ctx) {
     ctx_free(ctx);
   }
@@ -213,6 +235,46 @@ static inline void ctx_restore_gvars(Context *ctx)
   shada_read_sbuf(&ctx->gvars, kShaDaWantInfo | kShaDaForceit);
 }
 
+/// Saves functions to a context.
+///
+/// @param  ctx         Save to this context.
+/// @param  scriptonly  Save script-local (s:) functions only.
+static inline void ctx_save_funcs(Context *ctx, bool scriptonly)
+  FUNC_ATTR_NONNULL_ALL
+{
+  ctx->funcs = (Array)ARRAY_DICT_INIT;
+  Error err = ERROR_INIT;
+
+  HASHTAB_ITER(&func_hashtab, hi, {
+    const char_u *const name = hi->hi_key;
+    bool islambda = (STRNCMP(name, "<lambda>", 8) == 0);
+    bool isscript = (name[0] == K_SPECIAL);
+
+    if (!islambda && (!scriptonly || isscript)) {
+      size_t cmd_len = sizeof("func! ") + STRLEN(name);
+      char *cmd = xmalloc(cmd_len);
+      snprintf(cmd, cmd_len, "func! %s", name);
+      String func_body = nvim_command_output(cstr_as_string(cmd), &err);
+      xfree(cmd);
+      if (!ERROR_SET(&err)) {
+        ADD(ctx->funcs, STRING_OBJ(func_body));
+      }
+      api_clear_error(&err);
+    }
+  });
+}
+
+/// Restores functions from a context.
+///
+/// @param  ctx  Restore from this context.
+static inline void ctx_restore_funcs(Context *ctx)
+  FUNC_ATTR_NONNULL_ALL
+{
+  for (size_t i = 0; i < ctx->funcs.size; i++) {
+    do_cmdline_cmd(ctx->funcs.items[i].data.string.data);
+  }
+}
+
 /// Convert msgpack_sbuffer to readfile()-style array.
 ///
 /// @param[in]  sbuf  msgpack_sbuffer to convert.
@@ -277,6 +339,7 @@ Dictionary ctx_to_dict(Context *ctx)
   PUT(rv, "jumps", ARRAY_OBJ(sbuf_to_array(ctx->jumps)));
   PUT(rv, "buflist", ARRAY_OBJ(sbuf_to_array(ctx->buflist)));
   PUT(rv, "gvars", ARRAY_OBJ(sbuf_to_array(ctx->gvars)));
+  PUT(rv, "funcs", ARRAY_OBJ(copy_array(ctx->funcs)));
 
   return rv;
 }
@@ -310,6 +373,9 @@ int ctx_from_dict(Dictionary dict, Context *ctx)
     } else if (strequal(item.key.data, "gvars")) {
       types |= kCtxGVars;
       ctx->gvars = array_to_sbuf(item.value.data.array);
+    } else if (strequal(item.key.data, "funcs")) {
+      types |= kCtxFuncs;
+      ctx->funcs = copy_object(item.value).data.array;
     }
   }
 
diff --git a/src/nvim/context.h b/src/nvim/context.h
index 4e641adeda..328e12c6a6 100644
--- a/src/nvim/context.h
+++ b/src/nvim/context.h
@@ -10,6 +10,7 @@ typedef struct {
   msgpack_sbuffer jumps;    ///< Jumplist.
   msgpack_sbuffer buflist;  ///< Buffer list.
   msgpack_sbuffer gvars;    ///< Global variables.
+  Array funcs;              ///< Functions.
 } Context;
 typedef kvec_t(Context) ContextVec;
 
@@ -24,6 +25,7 @@ typedef kvec_t(Context) ContextVec;
   .jumps = MSGPACK_SBUFFER_INIT, \
   .buflist = MSGPACK_SBUFFER_INIT, \
   .gvars = MSGPACK_SBUFFER_INIT, \
+  .funcs = ARRAY_DICT_INIT, \
 }
 
 typedef enum {
@@ -31,6 +33,8 @@ typedef enum {
   kCtxJumps = 2,      ///< Jumplist
   kCtxBuflist = 4,    ///< Buffer list
   kCtxGVars = 8,      ///< Global variables
+  kCtxSFuncs = 16,    ///< Script functions
+  kCtxFuncs = 32,     ///< Functions
 } ContextTypeFlags;
 
 extern int kCtxAll;
diff --git a/src/nvim/eval.c b/src/nvim/eval.c
index 33e3cd64f2..cd447d3e0f 100644
--- a/src/nvim/eval.c
+++ b/src/nvim/eval.c
@@ -8064,6 +8064,10 @@ static void f_ctxpush(typval_T *argvars, typval_T *rettv, FunPtr fptr)
           types |= kCtxBuflist;
         } else if (strequal((char *)tv_li->vval.v_string, "gvars")) {
           types |= kCtxGVars;
+        } else if (strequal((char *)tv_li->vval.v_string, "sfuncs")) {
+          types |= kCtxSFuncs;
+        } else if (strequal((char *)tv_li->vval.v_string, "funcs")) {
+          types |= kCtxFuncs;
         }
       }
     });
@@ -20981,7 +20985,7 @@ void ex_function(exarg_T *eap)
             continue;
           }
           if (!func_name_refcount(fp->uf_name)) {
-            list_func_head(fp, false);
+            list_func_head(fp, false, false);
           }
         }
       }
@@ -21012,7 +21016,7 @@ void ex_function(exarg_T *eap)
             fp = HI2UF(hi);
             if (!isdigit(*fp->uf_name)
                 && vim_regexec(&regmatch, fp->uf_name, 0))
-              list_func_head(fp, FALSE);
+              list_func_head(fp, false, false);
           }
         }
         vim_regfree(regmatch.regprog);
@@ -21062,9 +21066,12 @@ void ex_function(exarg_T *eap)
   saved_did_emsg = did_emsg;
   did_emsg = FALSE;
 
-  /*
-   * ":function func" with only function name: list function.
-   */
+  //
+  // ":function func" with only function name: list function.
+  // If bang is given:
+  //  - include "!" in function head
+  //  - exclude line numbers from function body
+  //
   if (!paren) {
     if (!ends_excmd(*skipwhite(p))) {
       EMSG(_(e_trailing));
@@ -21076,17 +21083,20 @@ void ex_function(exarg_T *eap)
     if (!eap->skip && !got_int) {
       fp = find_func(name);
       if (fp != NULL) {
-        list_func_head(fp, TRUE);
-        for (int j = 0; j < fp->uf_lines.ga_len && !got_int; ++j) {
-          if (FUNCLINE(fp, j) == NULL)
+        list_func_head(fp, !eap->forceit, eap->forceit);
+        for (int j = 0; j < fp->uf_lines.ga_len && !got_int; j++) {
+          if (FUNCLINE(fp, j) == NULL) {
             continue;
-          msg_putchar('\n');
-          msg_outnum((long)j + 1);
-          if (j < 9) {
-            msg_putchar(' ');
           }
-          if (j < 99) {
-            msg_putchar(' ');
+          msg_putchar('\n');
+          if (!eap->forceit) {
+            msg_outnum((long)j + 1);
+            if (j < 9) {
+              msg_putchar(' ');
+            }
+            if (j < 99) {
+              msg_putchar(' ');
+            }
           }
           msg_prt_line(FUNCLINE(fp, j), false);
           ui_flush();                  // show a line at a time
@@ -21094,7 +21104,7 @@ void ex_function(exarg_T *eap)
         }
         if (!got_int) {
           msg_putchar('\n');
-          msg_puts("   endfunction");
+          msg_puts(eap->forceit ? "endfunction" : "   endfunction");
         }
       } else
         emsg_funcname(N_("E123: Undefined function: %s"), name);
@@ -21784,15 +21794,17 @@ static inline bool eval_fname_sid(const char *const name)
   return *name == 's' || TOUPPER_ASC(name[2]) == 'I';
 }
 
-/*
- * List the head of the function: "name(arg1, arg2)".
- */
-static void list_func_head(ufunc_T *fp, int indent)
+/// List the head of the function: "name(arg1, arg2)".
+///
+/// @param[in]  fp      Function pointer.
+/// @param[in]  indent  Indent line.
+/// @param[in]  force   Include bang "!" (i.e.: "function!").
+static void list_func_head(ufunc_T *fp, int indent, bool force)
 {
   msg_start();
   if (indent)
     MSG_PUTS("   ");
-  MSG_PUTS("function ");
+  MSG_PUTS(force ? "function! " : "function ");
   if (fp->uf_name[0] == K_SPECIAL) {
     MSG_PUTS_ATTR("<SNR>", HL_ATTR(HLF_8));
     msg_puts((const char *)fp->uf_name + 3);
diff --git a/src/nvim/main.c b/src/nvim/main.c
index 216343d8f4..9c342e62c0 100644
--- a/src/nvim/main.c
+++ b/src/nvim/main.c
@@ -14,7 +14,6 @@
 #include "nvim/main.h"
 #include "nvim/buffer.h"
 #include "nvim/charset.h"
-#include "nvim/context.h"
 #include "nvim/diff.h"
 #include "nvim/eval.h"
 #include "nvim/ex_cmds.h"
@@ -673,8 +672,6 @@ void getout(int exitval)
     garbage_collect(false);
   }
 
-  free_ctx_stack();
-
   mch_exit(exitval);
 }
 
diff --git a/src/nvim/memory.c b/src/nvim/memory.c
index 1384aa177b..64aae71433 100644
--- a/src/nvim/memory.c
+++ b/src/nvim/memory.c
@@ -9,6 +9,7 @@
 #include <stdbool.h>
 
 #include "nvim/vim.h"
+#include "nvim/context.h"
 #include "nvim/eval.h"
 #include "nvim/highlight.h"
 #include "nvim/memfile.h"
@@ -671,6 +672,7 @@ void free_all_mem(void)
 
   eval_clear();
   api_vim_free_all_mem();
+  ctx_free_all();
 
   // Free all buffers.  Reset 'autochdir' to avoid accessing things that
   // were freed already.
diff --git a/test/functional/eval/ctx_functions_spec.lua b/test/functional/eval/ctx_functions_spec.lua
index 3699af0793..35133e2341 100644
--- a/test/functional/eval/ctx_functions_spec.lua
+++ b/test/functional/eval/ctx_functions_spec.lua
@@ -10,6 +10,8 @@ local feed = helpers.feed
 local map = helpers.map
 local nvim = helpers.nvim
 local parse_context = helpers.parse_context
+local redir_exec = helpers.redir_exec
+local source = helpers.source
 local trim = helpers.trim
 local write_file = helpers.write_file
 
@@ -126,6 +128,109 @@ describe('context functions', function()
       eq({1, 2 ,3}, eval('[g:one, g:Two, g:THREE]'))
     end)
 
+    it('saves and restores script functions properly', function()
+      source([[
+      function s:greet(name)
+        echom 'Hello, '.a:name.'!'
+      endfunction
+
+      function s:greet_all(name, ...)
+        echom 'Hello, '.a:name.'!'
+        for more in a:000
+          echom 'Hello, '.more.'!'
+        endfor
+      endfunction
+
+      function Greet(name)
+        call call('s:greet', [a:name])
+      endfunction
+
+      function GreetAll(name, ...)
+        call call('s:greet_all', extend([a:name], a:000))
+      endfunction
+
+      function SaveSFuncs()
+        call ctxpush(['sfuncs'])
+      endfunction
+
+      function DeleteSFuncs()
+        delfunction s:greet
+        delfunction s:greet_all
+      endfunction
+
+      function RestoreFuncs()
+        call ctxpop()
+      endfunction
+      ]])
+
+      eq('\nHello, World!', redir_exec([[call Greet('World')]]))
+      eq('\nHello, World!'..
+         '\nHello, One!'..
+         '\nHello, Two!'..
+         '\nHello, Three!',
+         redir_exec([[call GreetAll('World', 'One', 'Two', 'Three')]]))
+
+      call('SaveSFuncs')
+      call('DeleteSFuncs')
+
+      eq('\nError detected while processing function Greet:'..
+         '\nline    1:'..
+         '\nE117: Unknown function: s:greet',
+         redir_exec([[call Greet('World')]]))
+      eq('\nError detected while processing function GreetAll:'..
+         '\nline    1:'..
+         '\nE117: Unknown function: s:greet_all',
+         redir_exec([[call GreetAll('World', 'One', 'Two', 'Three')]]))
+
+      call('RestoreFuncs')
+
+      eq('\nHello, World!', redir_exec([[call Greet('World')]]))
+      eq('\nHello, World!'..
+         '\nHello, One!'..
+         '\nHello, Two!'..
+         '\nHello, Three!',
+         redir_exec([[call GreetAll('World', 'One', 'Two', 'Three')]]))
+    end)
+
+    it('saves and restores functions properly', function()
+      source([[
+      function Greet(name)
+        echom 'Hello, '.a:name.'!'
+      endfunction
+
+      function GreetAll(name, ...)
+        echom 'Hello, '.a:name.'!'
+        for more in a:000
+          echom 'Hello, '.more.'!'
+        endfor
+      endfunction
+      ]])
+
+      eq('\nHello, World!', redir_exec([[call Greet('World')]]))
+      eq('\nHello, World!'..
+         '\nHello, One!'..
+         '\nHello, Two!'..
+         '\nHello, Three!',
+         redir_exec([[call GreetAll('World', 'One', 'Two', 'Three')]]))
+
+      call('ctxpush', {'funcs'})
+      command('delfunction Greet')
+      command('delfunction GreetAll')
+
+      expect_err('Vim:E117: Unknown function: Greet', call, 'Greet', 'World')
+      expect_err('Vim:E117: Unknown function: Greet', call, 'GreetAll',
+                 'World', 'One', 'Two', 'Three')
+
+      call('ctxpop')
+
+      eq('\nHello, World!', redir_exec([[call Greet('World')]]))
+      eq('\nHello, World!'..
+         '\nHello, One!'..
+         '\nHello, Two!'..
+         '\nHello, Three!',
+         redir_exec([[call GreetAll('World', 'One', 'Two', 'Three')]]))
+    end)
+
     it('errors out when context stack is empty', function()
       local err = 'Vim:Context stack is empty'
       expect_err(err, call, 'ctxpop')