From 7da30807617c4cee31840cf6c0c3790754a6617e Mon Sep 17 00:00:00 2001
From: Scott Lahteine <sourcetree@thinkyhead.com>
Date: Wed, 22 Jun 2016 14:20:34 -0700
Subject: [PATCH] Simpler cleanup in G28 / G29

---
 Marlin/Marlin_main.cpp | 38 ++++++++++++++++----------------------
 1 file changed, 16 insertions(+), 22 deletions(-)

diff --git a/Marlin/Marlin_main.cpp b/Marlin/Marlin_main.cpp
index 768e49802..cc1bdc4a2 100644
--- a/Marlin/Marlin_main.cpp
+++ b/Marlin/Marlin_main.cpp
@@ -1607,16 +1607,17 @@ static void setup_for_endstop_move() {
   endstops.enable();
 }
 
+static void clean_up_after_endstop_or_probe_move() {
+  #if ENABLED(DEBUG_LEVELING_FEATURE)
+    if (DEBUGGING(LEVELING)) DEBUG_POS("clean_up_after_endstop_or_probe_move", current_position);
+  #endif
+  feedrate = saved_feedrate;
+  feedrate_multiplier = saved_feedrate_multiplier;
+  refresh_cmd_timeout();
+}
+
 #if HAS_BED_PROBE
 
-  static void clean_up_after_endstop_or_probe_move() {
-    #if ENABLED(DEBUG_LEVELING_FEATURE)
-      if (DEBUGGING(LEVELING)) DEBUG_POS("clean_up_after_endstop_or_probe_move", current_position);
-    #endif
-    feedrate = saved_feedrate;
-    feedrate_multiplier = saved_feedrate_multiplier;
-    refresh_cmd_timeout();
-  }
   static void clean_up_after_endstop_move() {
     clean_up_after_endstop_or_probe_move();
     endstops.not_homing();
@@ -3058,14 +3059,7 @@ inline void gcode_G28() {
 
   #endif // !DELTA (gcode_G28)
 
-  #if ENABLED(ENDSTOPS_ONLY_FOR_HOMING)
-    endstops.enable(false);
-    #if ENABLED(DEBUG_LEVELING_FEATURE)
-      if (DEBUGGING(LEVELING)) {
-        SERIAL_ECHOLNPGM("ENDSTOPS_ONLY_FOR_HOMING endstops.enable(false)");
-      }
-    #endif
-  #endif
+  endstops.not_homing();
 
   // Enable mesh leveling again
   #if ENABLED(MESH_BED_LEVELING)
@@ -3097,9 +3091,8 @@ inline void gcode_G28() {
     }
   #endif
 
-  feedrate = saved_feedrate;
-  feedrate_multiplier = saved_feedrate_multiplier;
-  refresh_cmd_timeout();
+  clean_up_after_endstop_or_probe_move();
+
   endstops.hit_on_purpose(); // clear endstop hit flags
 
   #if ENABLED(DEBUG_LEVELING_FEATURE)
@@ -3597,8 +3590,6 @@ inline void gcode_G28() {
         if (DEBUGGING(LEVELING)) DEBUG_POS("> probing complete", current_position);
       #endif
 
-      clean_up_after_endstop_or_probe_move();
-
       #if ENABLED(DELTA)
 
         if (!dryrun) extrapolate_unprobed_bed_level();
@@ -3725,7 +3716,7 @@ inline void gcode_G28() {
                                   ABL_PROBE_PT_3_Y + home_offset[Y_AXIS],
                                   current_position[Z_AXIS] + Z_RAISE_BETWEEN_PROBINGS,
                                   p3, verbose_level);
-      clean_up_after_endstop_or_probe_move();
+
       if (!dryrun) set_bed_level_equation_3pts(z_at_pt_1, z_at_pt_2, z_at_pt_3);
 
     #endif // !AUTO_BED_LEVELING_GRID
@@ -3807,6 +3798,9 @@ inline void gcode_G28() {
     // Stow the probe. Servo will raise if needed.
     stow_z_probe();
 
+    // Restore state after probing
+    clean_up_after_endstop_or_probe_move();
+
     #ifdef Z_PROBE_END_SCRIPT
       #if ENABLED(DEBUG_LEVELING_FEATURE)
         if (DEBUGGING(LEVELING)) {