From f6e5366d0077e9f171651f37282cb5c47629d1b6 Mon Sep 17 00:00:00 2001
From: zeertzjq <zeertzjq@outlook.com>
Date: Sat, 2 Dec 2023 07:55:44 +0800
Subject: [PATCH] refactor: free more reachable memory with EXITFREE (#26349)

Discovered using __sanitizer_print_memory_profile().
---
 src/nvim/api/ui.c              | 27 +++++++++++++++++++++------
 src/nvim/autocmd.c             |  5 +++++
 src/nvim/channel.c             | 31 ++++++++++++++++++++++++-------
 src/nvim/eval.c                |  4 ++++
 src/nvim/grid.c                |  5 +++++
 src/nvim/map_defs.h            | 17 +++++++++++++----
 src/nvim/marktree.c            |  1 -
 src/nvim/memory.c              | 25 +++++++++++++++++++++++++
 src/nvim/msgpack_rpc/channel.c |  5 +++++
 src/nvim/msgpack_rpc/helpers.c |  8 ++++++++
 src/nvim/option.c              |  3 +++
 src/nvim/os/env.c              | 11 +++++++++++
 src/nvim/os/input.c            |  7 +++++++
 src/nvim/ui.c                  |  2 ++
 src/nvim/ui_client.c           |  9 +++++++++
 src/nvim/ui_compositor.c       |  9 +++++++++
 16 files changed, 151 insertions(+), 18 deletions(-)

diff --git a/src/nvim/api/ui.c b/src/nvim/api/ui.c
index 7e64ce9cd1..b73c026d57 100644
--- a/src/nvim/api/ui.c
+++ b/src/nvim/api/ui.c
@@ -106,22 +106,37 @@ static void mpack_str(char **buf, const char *str)
   *buf += len;
 }
 
+static void remote_ui_destroy(UI *ui)
+  FUNC_ATTR_NONNULL_ALL
+{
+  UIData *data = ui->data;
+  kv_destroy(data->call_buf);
+  XFREE_CLEAR(ui->term_name);
+  xfree(ui);
+}
+
 void remote_ui_disconnect(uint64_t channel_id)
 {
   UI *ui = pmap_get(uint64_t)(&connected_uis, channel_id);
   if (!ui) {
     return;
   }
-  UIData *data = ui->data;
-  kv_destroy(data->call_buf);
   pmap_del(uint64_t)(&connected_uis, channel_id, NULL);
   ui_detach_impl(ui, channel_id);
-
-  // Destroy `ui`.
-  XFREE_CLEAR(ui->term_name);
-  xfree(ui);
+  remote_ui_destroy(ui);
 }
 
+#ifdef EXITFREE
+void remote_ui_free_all_mem(void)
+{
+  UI *ui;
+  map_foreach_value(&connected_uis, ui, {
+    remote_ui_destroy(ui);
+  });
+  map_destroy(uint64_t, &connected_uis);
+}
+#endif
+
 /// Wait until ui has connected on stdio channel if only_stdio
 /// is true, otherwise any channel.
 void remote_ui_wait_for_attach(bool only_stdio)
diff --git a/src/nvim/autocmd.c b/src/nvim/autocmd.c
index 7d635984b8..46a08c5706 100644
--- a/src/nvim/autocmd.c
+++ b/src/nvim/autocmd.c
@@ -2573,6 +2573,11 @@ void do_autocmd_uienter(uint64_t chanid, bool attached)
 {
   static bool recursive = false;
 
+#ifdef EXITFREE
+  if (entered_free_all_mem) {
+    return;
+  }
+#endif
   if (starting == NO_SCREEN) {
     return;  // user config hasn't been sourced yet
   }
diff --git a/src/nvim/channel.c b/src/nvim/channel.c
index 849b63ae31..24793bcb2a 100644
--- a/src/nvim/channel.c
+++ b/src/nvim/channel.c
@@ -53,13 +53,25 @@ static uint64_t next_chan_id = CHAN_STDERR + 1;
 /// Teardown the module
 void channel_teardown(void)
 {
-  Channel *channel;
-
-  map_foreach_value(&channels, channel, {
-    channel_close(channel->id, kChannelPartAll, NULL);
+  Channel *chan;
+  map_foreach_value(&channels, chan, {
+    channel_close(chan->id, kChannelPartAll, NULL);
   });
 }
 
+#ifdef EXITFREE
+void channel_free_all_mem(void)
+{
+  Channel *chan;
+  map_foreach_value(&channels, chan, {
+    channel_destroy(chan);
+  });
+  map_destroy(uint64_t, &channels);
+
+  callback_free(&on_print);
+}
+#endif
+
 /// Closes a channel
 ///
 /// @param id The channel id
@@ -260,9 +272,8 @@ void callback_reader_start(CallbackReader *reader, const char *type)
   reader->type = type;
 }
 
-static void free_channel_event(void **argv)
+static void channel_destroy(Channel *chan)
 {
-  Channel *chan = argv[0];
   if (chan->is_rpc) {
     rpc_free(chan);
   }
@@ -275,11 +286,17 @@ static void free_channel_event(void **argv)
   callback_reader_free(&chan->on_stderr);
   callback_free(&chan->on_exit);
 
-  pmap_del(uint64_t)(&channels, chan->id, NULL);
   multiqueue_free(chan->events);
   xfree(chan);
 }
 
+static void free_channel_event(void **argv)
+{
+  Channel *chan = argv[0];
+  pmap_del(uint64_t)(&channels, chan->id, NULL);
+  channel_destroy(chan);
+}
+
 static void channel_destroy_early(Channel *chan)
 {
   if ((chan->id != --next_chan_id)) {
diff --git a/src/nvim/eval.c b/src/nvim/eval.c
index 6786316b8e..ae34f5715f 100644
--- a/src/nvim/eval.c
+++ b/src/nvim/eval.c
@@ -500,6 +500,10 @@ static void evalvars_clear(void)
       p->vv_list = NULL;
     }
   }
+
+  partial_unref(vvlua_partial);
+  vimvars[VV_LUA].vv_partial = vvlua_partial = NULL;
+
   hash_clear(&vimvarht);
   hash_init(&vimvarht);    // garbage_collect() will access it
   hash_clear(&compat_hashtab);
diff --git a/src/nvim/grid.c b/src/nvim/grid.c
index 2ef89b778e..58884e7f72 100644
--- a/src/nvim/grid.c
+++ b/src/nvim/grid.c
@@ -879,15 +879,20 @@ void grid_free(ScreenGrid *grid)
   grid->line_offset = NULL;
 }
 
+#ifdef EXITFREE
 /// Doesn't allow reinit, so must only be called by free_all_mem!
 void grid_free_all_mem(void)
 {
   grid_free(&default_grid);
+  grid_free(&msg_grid);
+  XFREE_CLEAR(msg_grid.dirty_col);
   xfree(linebuf_char);
   xfree(linebuf_attr);
   xfree(linebuf_vcol);
   xfree(linebuf_scratch);
+  set_destroy(glyph, &glyph_cache);
 }
+#endif
 
 /// (Re)allocates a window grid if size changed while in ext_multigrid mode.
 /// Updates size, offsets and handle for the grid regardless.
diff --git a/src/nvim/map_defs.h b/src/nvim/map_defs.h
index 147c03327a..f3c4e4ea95 100644
--- a/src/nvim/map_defs.h
+++ b/src/nvim/map_defs.h
@@ -172,9 +172,14 @@ MAP_DECLS(ColorKey, ColorItem)
 #define set_put_ref(T, set, key, key_alloc) set_put_##T(set, key, key_alloc)
 #define set_put_idx(T, set, key, status) mh_put_##T(set, key, status)
 #define set_del(T, set, key) set_del_##T(set, key)
-#define set_destroy(T, set) (xfree((set)->keys), xfree((set)->h.hash))
-#define set_clear(T, set) mh_clear(&(set)->h)
 #define set_size(set) ((set)->h.size)
+#define set_clear(T, set) mh_clear(&(set)->h)
+#define set_destroy(T, set) \
+  do { \
+    xfree((set)->keys); \
+    xfree((set)->h.hash); \
+    *(set) = (Set(T)) SET_INIT; \
+  } while (0)
 
 #define map_get(T, U) map_get_##T##U
 #define map_has(T, map, key) set_has(T, &(map)->set, key)
@@ -182,9 +187,13 @@ MAP_DECLS(ColorKey, ColorItem)
 #define map_ref(T, U) map_ref_##T##U
 #define map_put_ref(T, U) map_put_ref_##T##U
 #define map_del(T, U) map_del_##T##U
-#define map_destroy(T, map) (set_destroy(T, &(map)->set), xfree((map)->values))
-#define map_clear(T, map) set_clear(T, &(map)->set)
 #define map_size(map) set_size(&(map)->set)
+#define map_clear(T, map) set_clear(T, &(map)->set)
+#define map_destroy(T, map) \
+  do { \
+    set_destroy(T, &(map)->set); \
+    XFREE_CLEAR((map)->values); \
+  } while (0)
 
 #define pmap_get(T) map_get(T, ptr_t)
 #define pmap_put(T) map_put(T, ptr_t)
diff --git a/src/nvim/marktree.c b/src/nvim/marktree.c
index b555b3b4ae..f14da1b605 100644
--- a/src/nvim/marktree.c
+++ b/src/nvim/marktree.c
@@ -1117,7 +1117,6 @@ void marktree_clear(MarkTree *b)
     b->root = NULL;
   }
   map_destroy(uint64_t, b->id2node);
-  *b->id2node = (PMap(uint64_t)) MAP_INIT;
   b->n_keys = 0;
   assert(b->n_nodes == 0);
 }
diff --git a/src/nvim/memory.c b/src/nvim/memory.c
index 35ae6afde7..52fdf9f923 100644
--- a/src/nvim/memory.c
+++ b/src/nvim/memory.c
@@ -9,9 +9,12 @@
 #include <time.h>
 
 #include "nvim/api/extmark.h"
+#include "nvim/api/private/helpers.h"
+#include "nvim/api/ui.h"
 #include "nvim/arglist.h"
 #include "nvim/ascii_defs.h"
 #include "nvim/buffer_updates.h"
+#include "nvim/channel.h"
 #include "nvim/context.h"
 #include "nvim/decoration_provider.h"
 #include "nvim/drawline.h"
@@ -23,14 +26,19 @@
 #include "nvim/insexpand.h"
 #include "nvim/lua/executor.h"
 #include "nvim/main.h"
+#include "nvim/map_defs.h"
 #include "nvim/mapping.h"
 #include "nvim/memfile.h"
 #include "nvim/memory.h"
 #include "nvim/message.h"
 #include "nvim/option_vars.h"
+#include "nvim/os/input.h"
 #include "nvim/sign.h"
 #include "nvim/state_defs.h"
+#include "nvim/statusline.h"
 #include "nvim/ui.h"
+#include "nvim/ui_client.h"
+#include "nvim/ui_compositor.h"
 #include "nvim/usercmd.h"
 #include "nvim/vim_defs.h"
 
@@ -670,6 +678,7 @@ char *arena_memdupz(Arena *arena, const char *buf, size_t size)
 # include "nvim/grid.h"
 # include "nvim/mark.h"
 # include "nvim/msgpack_rpc/channel.h"
+# include "nvim/msgpack_rpc/helpers.h"
 # include "nvim/ops.h"
 # include "nvim/option.h"
 # include "nvim/os/os.h"
@@ -738,6 +747,7 @@ void free_all_mem(void)
   free_all_marks();
   alist_clear(&global_alist);
   free_homedir();
+  free_envmap();
   free_users();
   free_search_patterns();
   free_old_sub();
@@ -792,6 +802,7 @@ void free_all_mem(void)
     }
   }
 
+  channel_free_all_mem();
   eval_clear();
   api_extmark_free_all_mem();
   ctx_free_all();
@@ -815,8 +826,14 @@ void free_all_mem(void)
     buf = bufref_valid(&bufref) ? nextbuf : firstbuf;
   }
 
+  map_destroy(int, &buffer_handles);
+  map_destroy(int, &window_handles);
+  map_destroy(int, &tabpage_handles);
+
   // free screenlines (can't display anything now!)
   grid_free_all_mem();
+  stl_clear_click_defs(tab_page_click_defs, tab_page_click_defs_size);
+  xfree(tab_page_click_defs);
 
   clear_hl_tables(false);
 
@@ -824,10 +841,18 @@ void free_all_mem(void)
 
   decor_free_all_mem();
   drawline_free_all_mem();
+  input_free_all_mem();
 
+  if (ui_client_channel_id) {
+    ui_client_free_all_mem();
+  }
+
+  remote_ui_free_all_mem();
   ui_free_all_mem();
+  ui_comp_free_all_mem();
   nlua_free_all_mem();
   rpc_free_all_mem();
+  msgpack_rpc_helpers_free_all_mem();
 
   // should be last, in case earlier free functions deallocates arenas
   arena_free_reuse_blks();
diff --git a/src/nvim/msgpack_rpc/channel.c b/src/nvim/msgpack_rpc/channel.c
index acc21bbf7e..2336853609 100644
--- a/src/nvim/msgpack_rpc/channel.c
+++ b/src/nvim/msgpack_rpc/channel.c
@@ -751,6 +751,7 @@ const char *get_client_info(Channel *chan, const char *key)
   return NULL;
 }
 
+#ifdef EXITFREE
 void rpc_free_all_mem(void)
 {
   cstr_t key;
@@ -758,4 +759,8 @@ void rpc_free_all_mem(void)
     xfree((void *)key);
   });
   set_destroy(cstr_t, &event_strings);
+
+  msgpack_sbuffer_destroy(&out_buffer);
+  multiqueue_free(ch_before_blocking_events);
 }
+#endif
diff --git a/src/nvim/msgpack_rpc/helpers.c b/src/nvim/msgpack_rpc/helpers.c
index d2be321e7a..3162269df6 100644
--- a/src/nvim/msgpack_rpc/helpers.c
+++ b/src/nvim/msgpack_rpc/helpers.c
@@ -27,6 +27,14 @@ void msgpack_rpc_helpers_init(void)
   msgpack_sbuffer_init(&sbuffer);
 }
 
+#ifdef EXITFREE
+void msgpack_rpc_helpers_free_all_mem(void)
+{
+  msgpack_zone_destroy(&zone);
+  msgpack_sbuffer_destroy(&sbuffer);
+}
+#endif
+
 typedef struct {
   const msgpack_object *mobj;
   Object *aobj;
diff --git a/src/nvim/option.c b/src/nvim/option.c
index 882722a575..7a7cda2fa0 100644
--- a/src/nvim/option.c
+++ b/src/nvim/option.c
@@ -583,6 +583,9 @@ void free_all_options(void)
   }
   free_operatorfunc_option();
   free_tagfunc_option();
+  XFREE_CLEAR(fenc_default);
+  XFREE_CLEAR(p_term);
+  XFREE_CLEAR(p_ttytype);
 }
 #endif
 
diff --git a/src/nvim/os/env.c b/src/nvim/os/env.c
index b1e680e469..6c14f3d593 100644
--- a/src/nvim/os/env.c
+++ b/src/nvim/os/env.c
@@ -510,6 +510,17 @@ void free_homedir(void)
   xfree(homedir);
 }
 
+void free_envmap(void)
+{
+  cstr_t name;
+  ptr_t e;
+  map_foreach(&envmap, name, e, {
+    xfree((char *)name);
+    xfree(e);
+  });
+  map_destroy(cstr_t, &envmap);
+}
+
 #endif
 
 /// Call expand_env() and store the result in an allocated string.
diff --git a/src/nvim/os/input.c b/src/nvim/os/input.c
index 7c404aa736..69d328754b 100644
--- a/src/nvim/os/input.c
+++ b/src/nvim/os/input.c
@@ -74,6 +74,13 @@ void input_stop(void)
   stream_close(&read_stream, NULL, NULL);
 }
 
+#ifdef EXITFREE
+void input_free_all_mem(void)
+{
+  rbuffer_free(input_buffer);
+}
+#endif
+
 static void cursorhold_event(void **argv)
 {
   event_T event = State & MODE_INSERT ? EVENT_CURSORHOLDI : EVENT_CURSORHOLD;
diff --git a/src/nvim/ui.c b/src/nvim/ui.c
index b068847e85..cb4ebb5c3b 100644
--- a/src/nvim/ui.c
+++ b/src/nvim/ui.c
@@ -134,6 +134,8 @@ void ui_free_all_mem(void)
     free_ui_event_callback(event_cb);
   })
   map_destroy(uint32_t, &ui_event_cbs);
+
+  multiqueue_free(resize_events);
 }
 #endif
 
diff --git a/src/nvim/ui_client.c b/src/nvim/ui_client.c
index a7a1c5912a..eb32c16881 100644
--- a/src/nvim/ui_client.c
+++ b/src/nvim/ui_client.c
@@ -210,3 +210,12 @@ void ui_client_event_raw_line(GridLineEvent *g)
   tui_raw_line(tui, grid, row, startcol, endcol, clearcol, g->cur_attr, lineflags,
                (const schar_T *)grid_line_buf_char, grid_line_buf_attr);
 }
+
+#ifdef EXITFREE
+void ui_client_free_all_mem(void)
+{
+  tui_free_all_mem(tui);
+  xfree(grid_line_buf_char);
+  xfree(grid_line_buf_attr);
+}
+#endif
diff --git a/src/nvim/ui_compositor.c b/src/nvim/ui_compositor.c
index b698e017dc..c4078d6f63 100644
--- a/src/nvim/ui_compositor.c
+++ b/src/nvim/ui_compositor.c
@@ -60,6 +60,15 @@ void ui_comp_init(void)
   curgrid = &default_grid;
 }
 
+#ifdef EXITFREE
+void ui_comp_free_all_mem(void)
+{
+  kv_destroy(layers);
+  xfree(linebuf);
+  xfree(attrbuf);
+}
+#endif
+
 void ui_comp_syn_init(void)
 {
   dbghl_normal = syn_check_group(S_LEN("RedrawDebugNormal"));