fix(float): make `screen*()` functions respect floating windows

Resolves #19013.

3 files changed, 43 insertions, 23 deletions

diff --git a/src/nvim/eval.c b/src/nvim/eval.c
index a4f56b47e6..43d6ed558f 100644
--- a/src/nvim/eval.c
+++ b/src/nvim/eval.c
@@ -49,6 +49,7 @@
 #include "nvim/sign.h"
 #include "nvim/syntax.h"
 #include "nvim/ui.h"
+#include "nvim/ui_compositor.h"
 #include "nvim/undo.h"
 #include "nvim/version.h"
 #include "nvim/window.h"
@@ -6849,19 +6850,19 @@ void return_register(int regname, typval_T *rettv)
   rettv->vval.v_string = xstrdup(buf);
 }
 
-void screenchar_adjust_grid(ScreenGrid **grid, int *row, int *col)
+void screenchar_adjust(ScreenGrid **grid, int *row, int *col)
 {
   // TODO(bfredl): this is a hack for legacy tests which use screenchar()
   // to check printed messages on the screen (but not floats etc
   // as these are not legacy features). If the compositor is refactored to
   // have its own buffer, this should just read from it instead.
   msg_scroll_flush();
-  if (msg_grid.chars && msg_grid.comp_index > 0 && *row >= msg_grid.comp_row
-      && *row < (msg_grid.rows + msg_grid.comp_row)
-      && *col < msg_grid.cols) {
-    *grid = &msg_grid;
-    *row -= msg_grid.comp_row;
-  }
+
+  *grid = ui_comp_get_grid_at_coord(*row, *col);
+
+  // Make `row` and `col` relative to the grid
+  *row -= (*grid)->comp_row;
+  *col -= (*grid)->comp_col;
 }
 
 /// Set line or list of lines in buffer "buf".
diff --git a/src/nvim/eval/funcs.c b/src/nvim/eval/funcs.c
index 2225076a0a..56468190b5 100644
--- a/src/nvim/eval/funcs.c
+++ b/src/nvim/eval/funcs.c
@@ -8035,14 +8035,15 @@ static void f_screenattr(typval_T *argvars, typval_T *rettv, FunPtr fptr)
 {
   int c;
 
+  ScreenGrid *grid;
   int row = (int)tv_get_number_chk(&argvars[0], NULL) - 1;
   int col = (int)tv_get_number_chk(&argvars[1], NULL) - 1;
-  if (row < 0 || row >= default_grid.rows
-      || col < 0 || col >= default_grid.cols) {
+
+  screenchar_adjust(&grid, &row, &col);
+
+  if (row < 0 || row >= grid->rows || col < 0 || col >= grid->cols) {
     c = -1;
   } else {
-    ScreenGrid *grid = &default_grid;
-    screenchar_adjust_grid(&grid, &row, &col);
     c = grid->attrs[grid->line_offset[row] + col];
   }
   rettv->vval.v_number = c;
@@ -8053,14 +8054,15 @@ static void f_screenchar(typval_T *argvars, typval_T *rettv, FunPtr fptr)
 {
   int c;
 
+  ScreenGrid *grid;
   int row = tv_get_number_chk(&argvars[0], NULL) - 1;
   int col = tv_get_number_chk(&argvars[1], NULL) - 1;
-  if (row < 0 || row >= default_grid.rows
-      || col < 0 || col >= default_grid.cols) {
+
+  screenchar_adjust(&grid, &row, &col);
+
+  if (row < 0 || row >= grid->rows || col < 0 || col >= grid->cols) {
     c = -1;
   } else {
-    ScreenGrid *grid = &default_grid;
-    screenchar_adjust_grid(&grid, &row, &col);
     c = utf_ptr2char((char *)grid->chars[grid->line_offset[row] + col]);
   }
   rettv->vval.v_number = c;
@@ -8069,15 +8071,16 @@ static void f_screenchar(typval_T *argvars, typval_T *rettv, FunPtr fptr)
 /// "screenchars()" function
 static void f_screenchars(typval_T *argvars, typval_T *rettv, FunPtr fptr)
 {
+  ScreenGrid *grid;
   int row = tv_get_number_chk(&argvars[0], NULL) - 1;
   int col = tv_get_number_chk(&argvars[1], NULL) - 1;
-  if (row < 0 || row >= default_grid.rows
-      || col < 0 || col >= default_grid.cols) {
+
+  screenchar_adjust(&grid, &row, &col);
+
+  if (row < 0 || row >= grid->rows || col < 0 || col >= grid->cols) {
     tv_list_alloc_ret(rettv, 0);
     return;
   }
-  ScreenGrid *grid = &default_grid;
-  screenchar_adjust_grid(&grid, &row, &col);
   int pcc[MAX_MCO];
   int c = utfc_ptr2char(grid->chars[grid->line_offset[row] + col], pcc);
   int composing_len = 0;
@@ -8136,14 +8139,17 @@ static void f_screenstring(typval_T *argvars, typval_T *rettv, FunPtr fptr)
 {
   rettv->vval.v_string = NULL;
   rettv->v_type = VAR_STRING;
+
+  ScreenGrid *grid;
   int row = tv_get_number_chk(&argvars[0], NULL) - 1;
   int col = tv_get_number_chk(&argvars[1], NULL) - 1;
-  if (row < 0 || row >= default_grid.rows
-      || col < 0 || col >= default_grid.cols) {
+
+  screenchar_adjust(&grid, &row, &col);
+
+  if (row < 0 || row >= grid->rows || col < 0 || col >= grid->cols) {
     return;
   }
-  ScreenGrid *grid = &default_grid;
-  screenchar_adjust_grid(&grid, &row, &col);
+
   rettv->vval.v_string = (char *)vim_strsave(grid->chars[grid->line_offset[row] + col]);
 }
 
diff --git a/src/nvim/ui_compositor.c b/src/nvim/ui_compositor.c
index 3ce2b80ea9..5df70d0d8e 100644
--- a/src/nvim/ui_compositor.c
+++ b/src/nvim/ui_compositor.c
@@ -300,6 +300,19 @@ ScreenGrid *ui_comp_mouse_focus(int row, int col)
   return NULL;
 }
 
+/// Compute which grid is on top at supplied screen coordinates
+ScreenGrid *ui_comp_get_grid_at_coord(int row, int col)
+{
+  for (ssize_t i = (ssize_t)kv_size(layers) - 1; i > 0; i--) {
+    ScreenGrid *grid = kv_A(layers, i);
+    if (row >= grid->comp_row && row < grid->comp_row + grid->rows
+        && col >= grid->comp_col && col < grid->comp_col + grid->cols) {
+      return grid;
+    }
+  }
+  return &default_grid;
+}
+
 /// Baseline implementation. This is always correct, but we can sometimes
 /// do something more efficient (where efficiency means smaller deltas to
 /// the downstream UI.)