2 files changed, 8 insertions, 5 deletions

diff --git a/quantum/quantum.c b/quantum/quantum.c
index 5a2a73d4a..2662e5ef1 100644
--- a/quantum/quantum.c
+++ b/quantum/quantum.c
@@ -780,12 +780,14 @@ void set_single_persistent_default_layer(uint8_t default_layer) {
   default_layer_set(1U<<default_layer);
 }
 
+uint32_t update_tri_layer_state(uint32_t state, uint8_t layer1, uint8_t layer2, uint8_t layer3) {
+  uint32_t mask12 = (1UL << layer1) | (1UL << layer2);
+  uint32_t mask3 = 1UL << layer3;
+  return (state & mask12) == mask12 ? (state | mask3) : (state & ~mask3);
+}
+
 void update_tri_layer(uint8_t layer1, uint8_t layer2, uint8_t layer3) {
-  if (IS_LAYER_ON(layer1) && IS_LAYER_ON(layer2)) {
-    layer_on(layer3);
-  } else {
-    layer_off(layer3);
-  }
+  layer_state_set(update_tri_layer_state(layer_state, layer1, layer2, layer3));
 }
 
 void tap_random_base64(void) {
diff --git a/quantum/quantum.h b/quantum/quantum.h
index 80a606f09..195f578de 100644
--- a/quantum/quantum.h
+++ b/quantum/quantum.h
@@ -142,6 +142,7 @@ void send_char(char ascii_code);
 
 // For tri-layer
 void update_tri_layer(uint8_t layer1, uint8_t layer2, uint8_t layer3);
+uint32_t update_tri_layer_state(uint32_t state, uint8_t layer1, uint8_t layer2, uint8_t layer3);
 
 void set_single_persistent_default_layer(uint8_t default_layer);