From 867fecd157f45f591c11ae7e40c646f1111695d3 Mon Sep 17 00:00:00 2001
From: Paul Monigatti <paulmonigatti@users.noreply.github.com>
Date: Thu, 14 Oct 2021 08:59:52 +1300
Subject: [PATCH] Fix: Light flash not restoring previous LightState (#2383)

* Update light state when transformer has finished

* Revert writing direct to output

* Correct handling of zero-length light transformers

* Allow transformers to handle zero-length transitions, and check more boundary conditions when transitioning back to start state

* Removed log.h

* Fixed race condition between LightFlashTransformer.apply() and is_finished()

* clang-format

* Step progress from 0.0f to 1.0f at t=start_time for zero-length transforms to avoid divide-by-zero
---
 .../components/light/addressable_light.cpp    |  2 +-
 esphome/components/light/light_transformer.h  | 10 ++++-
 esphome/components/light/transformers.h       | 37 ++++++++++---------
 3 files changed, 30 insertions(+), 19 deletions(-)

diff --git a/esphome/components/light/addressable_light.cpp b/esphome/components/light/addressable_light.cpp
index f3e6c0ef1d..a8e0c7b762 100644
--- a/esphome/components/light/addressable_light.cpp
+++ b/esphome/components/light/addressable_light.cpp
@@ -79,7 +79,7 @@ optional<LightColorValues> AddressableLightTransformer::apply() {
   // dynamically-calculated alpha values to match the look.
 
   float denom = (1.0f - smoothed_progress);
-  float alpha = denom == 0.0f ? 0.0f : (smoothed_progress - this->last_transition_progress_) / denom;
+  float alpha = denom == 0.0f ? 1.0f : (smoothed_progress - this->last_transition_progress_) / denom;
 
   // We need to use a low-resolution alpha here which makes the transition set in only after ~half of the length
   // We solve this by accumulating the fractional part of the alpha over time.
diff --git a/esphome/components/light/light_transformer.h b/esphome/components/light/light_transformer.h
index dd904d0eed..35b045d5b4 100644
--- a/esphome/components/light/light_transformer.h
+++ b/esphome/components/light/light_transformer.h
@@ -39,7 +39,15 @@ class LightTransformer {
 
  protected:
   /// The progress of this transition, on a scale of 0 to 1.
-  float get_progress_() { return clamp((millis() - this->start_time_) / float(this->length_), 0.0f, 1.0f); }
+  float get_progress_() {
+    uint32_t now = esphome::millis();
+    if (now < this->start_time_)
+      return 0.0f;
+    if (now >= this->start_time_ + this->length_)
+      return 1.0f;
+
+    return clamp((now - this->start_time_) / float(this->length_), 0.0f, 1.0f);
+  }
 
   uint32_t start_time_;
   uint32_t length_;
diff --git a/esphome/components/light/transformers.h b/esphome/components/light/transformers.h
index c22846ceb1..a557bd39b1 100644
--- a/esphome/components/light/transformers.h
+++ b/esphome/components/light/transformers.h
@@ -73,9 +73,7 @@ class LightFlashTransformer : public LightTransformer {
     if (this->transition_length_ * 2 > this->length_)
       this->transition_length_ = this->length_ / 2;
 
-    // do not create transition if length is 0
-    if (this->transition_length_ == 0)
-      return;
+    this->begun_lightstate_restore_ = false;
 
     // first transition to original target
     this->transformer_ = this->state_.get_output()->create_default_transition();
@@ -83,40 +81,45 @@ class LightFlashTransformer : public LightTransformer {
   }
 
   optional<LightColorValues> apply() override {
-    // transition transformer does not handle 0 length as progress returns nan
-    if (this->transition_length_ == 0)
-      return this->target_values_;
+    optional<LightColorValues> result = {};
+
+    if (this->transformer_ == nullptr && millis() > this->start_time_ + this->length_ - this->transition_length_) {
+      // second transition back to start value
+      this->transformer_ = this->state_.get_output()->create_default_transition();
+      this->transformer_->setup(this->state_.current_values, this->get_start_values(), this->transition_length_);
+      this->begun_lightstate_restore_ = true;
+    }
 
     if (this->transformer_ != nullptr) {
-      if (!this->transformer_->is_finished()) {
-        return this->transformer_->apply();
-      } else {
+      result = this->transformer_->apply();
+
+      if (this->transformer_->is_finished()) {
         this->transformer_->stop();
         this->transformer_ = nullptr;
       }
     }
 
-    if (millis() > this->start_time_ + this->length_ - this->transition_length_) {
-      // second transition back to start value
-      this->transformer_ = this->state_.get_output()->create_default_transition();
-      this->transformer_->setup(this->state_.current_values, this->get_start_values(), this->transition_length_);
-    }
-
-    // once transition is complete, don't change states until next transition
-    return optional<LightColorValues>();
+    return result;
   }
 
   // Restore the original values after the flash.
   void stop() override {
+    if (this->transformer_ != nullptr) {
+      this->transformer_->stop();
+      this->transformer_ = nullptr;
+    }
     this->state_.current_values = this->get_start_values();
     this->state_.remote_values = this->get_start_values();
     this->state_.publish_state();
   }
 
+  bool is_finished() override { return this->begun_lightstate_restore_ && LightTransformer::is_finished(); }
+
  protected:
   LightState &state_;
   uint32_t transition_length_;
   std::unique_ptr<LightTransformer> transformer_{nullptr};
+  bool begun_lightstate_restore_;
 };
 
 }  // namespace light