From a1f3d2f3cd9bf8bb4da16da0e55cdb5ff264a15d Mon Sep 17 00:00:00 2001
From: Scott Lahteine <github@thinkyhead.com>
Date: Sun, 31 May 2020 01:03:28 -0500
Subject: [PATCH] Move set_all_z_lock to Stepper

---
 Marlin/src/gcode/calibrate/G34_M422.cpp | 25 ++-----------------------
 Marlin/src/module/motion.cpp            |  8 ++++----
 Marlin/src/module/stepper.h             | 12 +++++++++++-
 3 files changed, 17 insertions(+), 28 deletions(-)

diff --git a/Marlin/src/gcode/calibrate/G34_M422.cpp b/Marlin/src/gcode/calibrate/G34_M422.cpp
index f5addd9cd6..0961cce59c 100644
--- a/Marlin/src/gcode/calibrate/G34_M422.cpp
+++ b/Marlin/src/gcode/calibrate/G34_M422.cpp
@@ -47,17 +47,6 @@
 #define DEBUG_OUT ENABLED(DEBUG_LEVELING_FEATURE)
 #include "../../core/debug_out.h"
 
-inline void set_all_z_lock(const bool lock) {
-  stepper.set_z_lock(lock);
-  stepper.set_z2_lock(lock);
-  #if NUM_Z_STEPPER_DRIVERS >= 3
-    stepper.set_z3_lock(lock);
-    #if NUM_Z_STEPPER_DRIVERS >= 4
-      stepper.set_z4_lock(lock);
-    #endif
-  #endif
-}
-
 /**
  * G34: Z-Stepper automatic alignment
  *
@@ -318,17 +307,7 @@ void GcodeSuite::G34() {
         if (DEBUGGING(LEVELING)) DEBUG_ECHOLNPAIR("> Z", int(zstepper + 1), " corrected by ", z_align_move);
 
         // Lock all steppers except one
-        set_all_z_lock(true);
-        switch (zstepper) {
-          case 0: stepper.set_z_lock(false); break;
-          case 1: stepper.set_z2_lock(false); break;
-          #if NUM_Z_STEPPER_DRIVERS >= 3
-            case 2: stepper.set_z3_lock(false); break;
-          #endif
-          #if NUM_Z_STEPPER_DRIVERS == 4
-            case 3: stepper.set_z4_lock(false); break;
-          #endif
-        }
+        stepper.set_all_z_lock(true, zstepper);
 
         #if DISABLED(Z_STEPPER_ALIGN_KNOWN_STEPPER_POSITIONS)
           // Decreasing accuracy was detected so move was inverted.
@@ -342,7 +321,7 @@ void GcodeSuite::G34() {
       } // for (zstepper)
 
       // Back to normal stepper operations
-      set_all_z_lock(false);
+      stepper.set_all_z_lock(false);
       stepper.set_separate_multi_axis(false);
 
       if (err_break) break;
diff --git a/Marlin/src/module/motion.cpp b/Marlin/src/module/motion.cpp
index 09f0f390c5..727bb15271 100644
--- a/Marlin/src/module/motion.cpp
+++ b/Marlin/src/module/motion.cpp
@@ -1644,9 +1644,9 @@ void homeaxis(const AxisEnum axis) {
 
           const float adj = ABS(endstops.z2_endstop_adj);
           if (adj) {
-            if (pos_dir ? (endstops.z2_endstop_adj > 0) : (endstops.z2_endstop_adj < 0)) stepper.set_z_lock(true); else stepper.set_z2_lock(true);
+            if (pos_dir ? (endstops.z2_endstop_adj > 0) : (endstops.z2_endstop_adj < 0)) stepper.set_z1_lock(true); else stepper.set_z2_lock(true);
             do_homing_move(axis, pos_dir ? -adj : adj);
-            stepper.set_z_lock(false);
+            stepper.set_z1_lock(false);
             stepper.set_z2_lock(false);
           }
 
@@ -1657,7 +1657,7 @@ void homeaxis(const AxisEnum axis) {
           typedef void (*adjustFunc_t)(const bool);
 
           adjustFunc_t lock[] = {
-            stepper.set_z_lock, stepper.set_z2_lock, stepper.set_z3_lock
+            stepper.set_z1_lock, stepper.set_z2_lock, stepper.set_z3_lock
             #if NUM_Z_STEPPER_DRIVERS >= 4
               , stepper.set_z4_lock
             #endif
@@ -1725,7 +1725,7 @@ void homeaxis(const AxisEnum axis) {
             do_homing_move(axis, adj[0] - adj[1]);
           }
 
-          stepper.set_z_lock(false);
+          stepper.set_z1_lock(false);
           stepper.set_z2_lock(false);
           stepper.set_z3_lock(false);
           #if NUM_Z_STEPPER_DRIVERS >= 4
diff --git a/Marlin/src/module/stepper.h b/Marlin/src/module/stepper.h
index 38af36db51..bc35f55ff2 100644
--- a/Marlin/src/module/stepper.h
+++ b/Marlin/src/module/stepper.h
@@ -484,7 +484,7 @@ class Stepper {
       FORCE_INLINE static void set_y2_lock(const bool state) { locked_Y2_motor = state; }
     #endif
     #if EITHER(Z_MULTI_ENDSTOPS, Z_STEPPER_AUTO_ALIGN)
-      FORCE_INLINE static void set_z_lock(const bool state) { locked_Z_motor = state; }
+      FORCE_INLINE static void set_z1_lock(const bool state) { locked_Z_motor = state; }
       FORCE_INLINE static void set_z2_lock(const bool state) { locked_Z2_motor = state; }
       #if NUM_Z_STEPPER_DRIVERS >= 3
         FORCE_INLINE static void set_z3_lock(const bool state) { locked_Z3_motor = state; }
@@ -492,6 +492,16 @@ class Stepper {
           FORCE_INLINE static void set_z4_lock(const bool state) { locked_Z4_motor = state; }
         #endif
       #endif
+      static inline void set_all_z_lock(const bool lock, const int8_t except=-1) {
+        set_z1_lock(lock ^ (except == 0));
+        set_z2_lock(lock ^ (except == 1));
+        #if NUM_Z_STEPPER_DRIVERS >= 3
+          set_z3_lock(lock ^ (except == 2));
+          #if NUM_Z_STEPPER_DRIVERS >= 4
+            set_z4_lock(lock ^ (except == 3));
+          #endif
+        #endif
+      }
     #endif
 
     #if ENABLED(BABYSTEPPING)