From c1b900aae9d2b57c397d7687cb7ff1ed7edf9518 Mon Sep 17 00:00:00 2001
From: Chris Pepper <p3p@p3psoft.co.uk>
Date: Tue, 22 Dec 2020 11:59:25 +0000
Subject: [PATCH] Fix UBL mesh inset Z position (#20538)

---
 Marlin/src/feature/bedlevel/ubl/ubl.h         | 33 ++++++++-----
 .../src/feature/bedlevel/ubl/ubl_motion.cpp   | 49 ++++++++-----------
 2 files changed, 42 insertions(+), 40 deletions(-)

diff --git a/Marlin/src/feature/bedlevel/ubl/ubl.h b/Marlin/src/feature/bedlevel/ubl/ubl.h
index c90b1f7ac15..762becfb699 100644
--- a/Marlin/src/feature/bedlevel/ubl/ubl.h
+++ b/Marlin/src/feature/bedlevel/ubl/ubl.h
@@ -122,20 +122,29 @@ class unified_bed_leveling {
 
     FORCE_INLINE static void set_z(const int8_t px, const int8_t py, const float &z) { z_values[px][py] = z; }
 
+    static int8_t cell_index_x_raw(const float &x) {
+      return FLOOR((x - (MESH_MIN_X)) * RECIPROCAL(MESH_X_DIST));
+    }
+
+    static int8_t cell_index_y_raw(const float &y) {
+      return FLOOR((y - (MESH_MIN_Y)) * RECIPROCAL(MESH_Y_DIST));
+    }
+
+    static int8_t cell_index_x_valid(const float &x) {
+      return WITHIN(cell_index_x_raw(x), 0, (GRID_MAX_POINTS_X - 2));
+    }
+
+    static int8_t cell_index_y_valid(const float &y) {
+      return WITHIN(cell_index_y_raw(y), 0, (GRID_MAX_POINTS_Y - 2));
+    }
+
     static int8_t cell_index_x(const float &x) {
-      const int8_t cx = (x - (MESH_MIN_X)) * RECIPROCAL(MESH_X_DIST);
-      return constrain(cx, 0, (GRID_MAX_POINTS_X) - 1);   // -1 is appropriate if we want all movement to the X_MAX
-    }                                                     // position. But with this defined this way, it is possible
-                                                          // to extrapolate off of this point even further out. Probably
-                                                          // that is OK because something else should be keeping that from
-                                                          // happening and should not be worried about at this level.
+      return constrain(cell_index_x_raw(x), 0, (GRID_MAX_POINTS_X) - 2);
+    }
+
     static int8_t cell_index_y(const float &y) {
-      const int8_t cy = (y - (MESH_MIN_Y)) * RECIPROCAL(MESH_Y_DIST);
-      return constrain(cy, 0, (GRID_MAX_POINTS_Y) - 1);   // -1 is appropriate if we want all movement to the Y_MAX
-    }                                                     // position. But with this defined this way, it is possible
-                                                          // to extrapolate off of this point even further out. Probably
-                                                          // that is OK because something else should be keeping that from
-                                                          // happening and should not be worried about at this level.
+      return constrain(cell_index_y_raw(y), 0, (GRID_MAX_POINTS_Y) - 2);
+    }
 
     static inline xy_int8_t cell_indexes(const float &x, const float &y) {
       return { cell_index_x(x), cell_index_y(y) };
diff --git a/Marlin/src/feature/bedlevel/ubl/ubl_motion.cpp b/Marlin/src/feature/bedlevel/ubl/ubl_motion.cpp
index 010b5951be5..8b7cd15a3c6 100644
--- a/Marlin/src/feature/bedlevel/ubl/ubl_motion.cpp
+++ b/Marlin/src/feature/bedlevel/ubl/ubl_motion.cpp
@@ -56,39 +56,32 @@
     // A move within the same cell needs no splitting
     if (istart == iend) {
 
-      // For a move off the bed, use a constant Z raise
-      if (!WITHIN(iend.x, 0, GRID_MAX_POINTS_X - 1) || !WITHIN(iend.y, 0, GRID_MAX_POINTS_Y - 1)) {
-
-        // Note: There is no Z Correction in this case. We are off the grid and don't know what
-        // a reasonable correction would be.  If the user has specified a UBL_Z_RAISE_WHEN_OFF_MESH
-        // value, that will be used instead of a calculated (Bi-Linear interpolation) correction.
-
-        #ifdef UBL_Z_RAISE_WHEN_OFF_MESH
-          end.z += UBL_Z_RAISE_WHEN_OFF_MESH;
-        #endif
-        planner.buffer_segment(end, scaled_fr_mm_s, extruder);
-        current_position = destination;
-        return;
-      }
-
       FINAL_MOVE:
 
-      // The distance is always MESH_X_DIST so multiply by the constant reciprocal.
-      const float xratio = (end.x - mesh_index_to_xpos(iend.x)) * RECIPROCAL(MESH_X_DIST);
+      // When UBL_Z_RAISE_WHEN_OFF_MESH is disabled Z correction is extrapolated from the edge of the mesh
+      #ifdef UBL_Z_RAISE_WHEN_OFF_MESH
+        // For a move off the UBL mesh, use a constant Z raise
+        if (!cell_index_x_valid(end.x) || !cell_index_y_valid(end.y)) {
 
-      float z1, z2;
-      if (iend.x >= GRID_MAX_POINTS_X - 1)
-        z1 = z2 = 0.0;
-      else {
-        z1 = z_values[iend.x    ][iend.y    ] + xratio *
-            (z_values[iend.x + 1][iend.y    ] - z_values[iend.x][iend.y    ]),
-        z2 = z_values[iend.x    ][iend.y + 1] + xratio *
-            (z_values[iend.x + 1][iend.y + 1] - z_values[iend.x][iend.y + 1]);
-      }
+          // Note: There is no Z Correction in this case. We are off the mesh and don't know what
+          // a reasonable correction would be, UBL_Z_RAISE_WHEN_OFF_MESH will be used instead of
+          // a calculated (Bi-Linear interpolation) correction.
+
+          end.z += UBL_Z_RAISE_WHEN_OFF_MESH;
+          planner.buffer_segment(end, scaled_fr_mm_s, extruder);
+          current_position = destination;
+          return;
+        }
+      #endif
+
+      // The distance is always MESH_X_DIST so multiply by the constant reciprocal.
+      const float xratio = (end.x - mesh_index_to_xpos(iend.x)) * RECIPROCAL(MESH_X_DIST),
+                  yratio = (end.y - mesh_index_to_ypos(iend.y)) * RECIPROCAL(MESH_Y_DIST),
+                  z1 = z_values[iend.x][iend.y    ] + xratio * (z_values[iend.x + 1][iend.y    ] - z_values[iend.x][iend.y    ]),
+                  z2 = z_values[iend.x][iend.y + 1] + xratio * (z_values[iend.x + 1][iend.y + 1] - z_values[iend.x][iend.y + 1]);
 
       // X cell-fraction done. Interpolate the two Z offsets with the Y fraction for the final Z offset.
-      const float yratio = (end.y - mesh_index_to_ypos(iend.y)) * RECIPROCAL(MESH_Y_DIST),
-                  z0 = iend.y < GRID_MAX_POINTS_Y - 1 ? (z1 + (z2 - z1) * yratio) * planner.fade_scaling_factor_for_z(end.z) : 0.0;
+      const float z0 = (z1 + (z2 - z1) * yratio) * planner.fade_scaling_factor_for_z(end.z);
 
       // Undefined parts of the Mesh in z_values[][] are NAN.
       // Replace NAN corrections with 0.0 to prevent NAN propagation.