From 522c78b5e6f2faf23e6cbdf928362d784145c29d Mon Sep 17 00:00:00 2001
From: zeleps <39417467+zeleps@users.noreply.github.com>
Date: Sun, 17 Jan 2021 07:08:48 +0200
Subject: [PATCH] Fix Tramming Wizard behavior (#20796)

---
 Marlin/src/feature/bedlevel/bedlevel.h | 4 ++++
 Marlin/src/gcode/calibrate/G28.cpp     | 5 ++---
 Marlin/src/lcd/menu/menu_tramming.cpp  | 9 +++++++--
 Marlin/src/module/tool_change.cpp      | 2 +-
 4 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/Marlin/src/feature/bedlevel/bedlevel.h b/Marlin/src/feature/bedlevel/bedlevel.h
index 995e9d0dbc..a33f08ad0e 100644
--- a/Marlin/src/feature/bedlevel/bedlevel.h
+++ b/Marlin/src/feature/bedlevel/bedlevel.h
@@ -23,6 +23,10 @@
 
 #include "../../inc/MarlinConfigPre.h"
 
+#if EITHER(RESTORE_LEVELING_AFTER_G28, ENABLE_LEVELING_AFTER_G28)
+  #define G28_L0_ENSURES_LEVELING_OFF 1
+#endif
+
 #if ENABLED(PROBE_MANUALLY)
   extern bool g29_in_progress;
 #else
diff --git a/Marlin/src/gcode/calibrate/G28.cpp b/Marlin/src/gcode/calibrate/G28.cpp
index 7c13587a67..2de029a08b 100644
--- a/Marlin/src/gcode/calibrate/G28.cpp
+++ b/Marlin/src/gcode/calibrate/G28.cpp
@@ -241,7 +241,7 @@ void GcodeSuite::G28() {
 
   // Disable the leveling matrix before homing
   #if HAS_LEVELING
-    IF_ENABLED(RESTORE_LEVELING_AFTER_G28, const bool leveling_restore_state = planner.leveling_active);
+    const bool leveling_restore_state = parser.boolval('L', TERN(RESTORE_LEVELING_AFTER_G28, planner.leveling_active, ENABLED(ENABLE_LEVELING_AFTER_G28)));
     IF_ENABLED(PROBE_MANUALLY, g29_in_progress = false); // Cancel the active G29 session
     set_bed_leveling_enabled(false);
   #endif
@@ -440,8 +440,7 @@ void GcodeSuite::G28() {
     do_blocking_move_to_z(delta_clip_start_height);
   #endif
 
-  IF_ENABLED(RESTORE_LEVELING_AFTER_G28, set_bed_leveling_enabled(leveling_restore_state));
-  IF_ENABLED(ENABLE_LEVELING_AFTER_G28, set_bed_leveling_enabled(true));
+  TERN_(HAS_LEVELING, set_bed_leveling_enabled(leveling_restore_state));
 
   restore_feedrate_and_scaling();
 
diff --git a/Marlin/src/lcd/menu/menu_tramming.cpp b/Marlin/src/lcd/menu/menu_tramming.cpp
index a77709e108..da7afd86ef 100644
--- a/Marlin/src/lcd/menu/menu_tramming.cpp
+++ b/Marlin/src/lcd/menu/menu_tramming.cpp
@@ -42,6 +42,10 @@
 float z_measured[G35_PROBE_COUNT] = { 0 };
 static uint8_t tram_index = 0;
 
+#if HAS_LEVELING
+  #include "../../feature/bedlevel/bedlevel.h"
+#endif
+
 static bool probe_single_point() {
   do_blocking_move_to_z(TERN(BLTOUCH, Z_CLEARANCE_DEPLOY_PROBE, Z_CLEARANCE_BETWEEN_PROBES));
   // Stow after each point with BLTouch "HIGH SPEED" mode for push-pin safety
@@ -60,7 +64,7 @@ static void _menu_single_probe(const uint8_t point) {
   STATIC_ITEM(MSG_LEVEL_CORNERS, SS_LEFT);
   STATIC_ITEM(MSG_LAST_VALUE_SP, SS_LEFT, ftostr42_52(z_measured[0] - z_measured[point])); // Print diff
   ACTION_ITEM(MSG_UBL_BC_INSERT2, []{ if (probe_single_point()) ui.refresh(); });
-  ACTION_ITEM(MSG_BUTTON_DONE, []{ ui.goto_previous_screen_no_defer(); }); // Back
+  ACTION_ITEM(MSG_BUTTON_DONE, []{ ui.goto_previous_screen(); }); // Back
   END_MENU();
 }
 
@@ -88,7 +92,8 @@ void goto_tramming_wizard() {
 
   // Inject G28, wait for homing to complete,
   set_all_unhomed();
-  queue.inject_P(G28_STR);
+  queue.inject_P(TERN(G28_L0_ENSURES_LEVELING_OFF, PSTR("G28L0"), G28_STR));
+
   ui.goto_screen([]{
     _lcd_draw_homing();
     if (all_axes_homed())
diff --git a/Marlin/src/module/tool_change.cpp b/Marlin/src/module/tool_change.cpp
index 7f581131d8..95f32f2faa 100644
--- a/Marlin/src/module/tool_change.cpp
+++ b/Marlin/src/module/tool_change.cpp
@@ -918,7 +918,7 @@ void tool_change(const uint8_t new_tool, bool no_move/*=false*/) {
       DEBUG_ECHOLNPGM("No move (not homed)");
     }
 
-    TERN_(HAS_LCD_MENU, if (!no_move) ui.return_to_status());
+    TERN_(HAS_LCD_MENU, if (!no_move) ui.update());
 
     #if ENABLED(DUAL_X_CARRIAGE)
       const bool idex_full_control = dual_x_carriage_mode == DXC_FULL_CONTROL_MODE;