Refactor addressable light and fix partition issue (#512)

src/esphome/application.cpp

@@ -1281,7 +1281,7 @@ stepper::ULN2003 *Application::make_uln2003(const GPIOOutputPin &pin_a, const GP
 #ifdef USE_LIGHT
 Application::MakePartitionLight Application::make_partition_light(
     const std::string &name, const std::vector<light::AddressableSegment> &segments) {
-  auto *part = new PartitionLightOutput(segments);
+  auto *part = App.register_component(new PartitionLightOutput(segments));
   auto make = this->make_light_for_light_output(name, part);
 
   return MakePartitionLight{

src/esphome/light/addressable_light.cpp

@@ -132,7 +132,45 @@ void AddressableLight::set_effect_active(bool effect_active) {
   this->effect_active_ = effect_active;
 }
 void AddressableLight::write_state(LightState *state) {
-  this->write_state(state, 0, this->size());
+  LightColorValues value = state->get_current_values();
+  uint8_t max_brightness = roundf(value.get_brightness() * value.get_state() * 255.0f);
+  this->correction_.set_local_brightness(max_brightness);
+
+  if (this->is_effect_active())
+    return;
+
+  auto val = state->get_current_values();
+  // don't use LightState helper, gamma correction+brightness is handled by ESPColorView
+  ESPColor color = ESPColor(
+      uint8_t(roundf(val.get_red() * 255.0f)),
+      uint8_t(roundf(val.get_green() * 255.0f)),
+      uint8_t(roundf(val.get_blue() * 255.0f)),
+      uint8_t(roundf(val.get_white() * 255.0f))
+  );
+
+  for (int i = 0; i < this->size(); i++) {
+    (*this)[i] = color;
+  }
+}
+void AddressableLight::set_correction(float red, float green, float blue, float white) {
+  this->correction_.set_max_brightness(ESPColor(
+      uint8_t(roundf(red * 255.0f)),
+      uint8_t(roundf(green * 255.0f)),
+      uint8_t(roundf(blue * 255.0f)),
+      uint8_t(roundf(white * 255.0f))
+  ));
+}
+void AddressableLight::setup_state(LightState *state) {
+  this->correction_.calculate_gamma_table(state->get_gamma_correct());
+}
+void AddressableLight::schedule_show() {
+  this->next_show_ = true;
+}
+bool AddressableLight::should_show_() const {
+  return this->effect_active_ || this->next_show_;
+}
+void AddressableLight::mark_shown_() {
+  this->next_show_ = false;
 }
 
 int32_t PartitionLightOutput::size() const {
@@ -156,6 +194,7 @@ ESPColorView PartitionLightOutput::operator[](int32_t index) const {
   }
   auto seg = this->segments_[lo];
   auto view = (*seg.get_src())[index - seg.get_src_offset()];
+  view.set_color_correction_(&this->correction_);
   return view;
 }
 void PartitionLightOutput::clear_effect_data() {
@@ -166,24 +205,21 @@ void PartitionLightOutput::clear_effect_data() {
 LightTraits PartitionLightOutput::get_traits() {
   return this->segments_[0].get_src()->get_traits();
 }
-void PartitionLightOutput::write_state(LightState *state, int32_t begin, int32_t end) {
-  for (auto seg : this->segments_) {
-    int32_t dbegin = seg.get_dst_offset();
-    if (begin > dbegin)
-      continue;
-
-    int32_t sbegin = seg.get_src_offset();
-    int32_t send = begin + std::min(seg.get_size(), end - dbegin);
-    seg.get_src()->write_state(state, sbegin, send);
-  }
-}
 PartitionLightOutput::PartitionLightOutput(const std::vector<AddressableSegment> &segments) : segments_(segments) {
   int32_t off = 0;
   for (auto seg : this->segments_) {
     seg.set_dst_offset(off);
     off += seg.get_size();
   }
 }
+void PartitionLightOutput::loop() {
+  if (this->should_show_()) {
+    for (auto seg : this->segments_) {
+      seg.get_src()->schedule_show();
+    }
+    this->mark_shown_();
+  }
+}
 
 AddressableSegment::AddressableSegment(LightState *src, int32_t src_offset, int32_t size)
     : src_(static_cast<AddressableLight *>(src->get_output())), src_offset_(src_offset), size_(size) {

src/esphome/light/addressable_light.h

@@ -127,13 +127,14 @@ class ESPColorView {
   inline uint8_t get_blue() const ALWAYS_INLINE;
   inline uint8_t get_white() const ALWAYS_INLINE;
   inline uint8_t get_effect_data() const ALWAYS_INLINE;
+  inline void set_color_correction_(const ESPColorCorrection *color_correction) ALWAYS_INLINE;
  protected:
   uint8_t *const red_;
   uint8_t *const green_;
   uint8_t *const blue_;
   uint8_t *const white_;
   uint8_t *const effect_data_;
-  const ESPColorCorrection *const color_correction_;
+  const ESPColorCorrection *color_correction_;
 };
 
 class AddressableLight : public LightOutput {
@@ -145,9 +146,16 @@ class AddressableLight : public LightOutput {
   bool is_effect_active() const;
   void set_effect_active(bool effect_active);
   void write_state(LightState *state) override;
-  virtual void write_state(LightState *state, int32_t begin, int32_t end) = 0;
+  void set_correction(float red, float green, float blue, float white = 1.0f);
+  void setup_state(LightState *state) override;
+  void schedule_show();
  protected:
+  bool should_show_() const;
+  void mark_shown_();
+
   bool effect_active_{false};
+  bool next_show_{true};
+  ESPColorCorrection correction_{};
 };
 
 class AddressableSegment {
@@ -167,17 +175,18 @@ class AddressableSegment {
   int32_t dst_offset_;
 };
 
-class PartitionLightOutput : public AddressableLight {
+class PartitionLightOutput : public AddressableLight, public Component {
  public:
   PartitionLightOutput(const std::vector<AddressableSegment> &segments);
   int32_t size() const override;
   ESPColorView operator[](int32_t index) const override;
   void clear_effect_data() override;
   LightTraits get_traits() override;
-  void write_state(LightState *state, int32_t begin, int32_t end) override;
+  void loop() override;
 
  protected:
   std::vector<AddressableSegment> segments_;
+  ESPColorCorrection correction_{};
 };
 
 } // namespace light

src/esphome/light/addressable_light.tcc

@@ -230,6 +230,9 @@ const ESPColorView &ESPColorView::operator=(const ESPHSVColor &rhs) const {
   this->set(rhs);
   return *this;
 }
+void ESPColorView::set_color_correction_(const ESPColorCorrection *color_correction) {
+  this->color_correction_ = color_correction;
+}
 
 ESPColor ESPColorCorrection::color_correct(ESPColor color) const {
   // corrected = (uncorrected * max_brightness * local_brightness) ^ gamma

src/esphome/light/fast_led_light_output.cpp

@@ -14,28 +14,6 @@ static const char *TAG = "light.fast_led";
 LightTraits FastLEDLightOutputComponent::get_traits() {
   return {true, true, false, false};
 }
-void FastLEDLightOutputComponent::write_state(LightState *state, int32_t begin, int32_t end) {
-  LightColorValues value = state->get_current_values();
-  uint8_t max_brightness = roundf(value.get_brightness() * value.get_state() * 255.0f);
-  this->correction_.set_local_brightness(max_brightness);
-
-  if (this->is_effect_active())
-    return;
-
-  auto val = state->get_current_values();
-  // don't use LightState helper, gamma correction+brightness is handled by ESPColorView
-  ESPColor color = ESPColor(
-      uint8_t(roundf(val.get_red() * 255.0f)),
-      uint8_t(roundf(val.get_green() * 255.0f)),
-      uint8_t(roundf(val.get_blue() * 255.0f))
-  );
-
-  for (int i = begin; i < end; i++) {
-    (*this)[i] = color;
-  }
-
-  this->schedule_show();
-}
 void FastLEDLightOutputComponent::setup() {
   ESP_LOGCONFIG(TAG, "Setting up FastLED light...");
   this->controller_->init();
@@ -51,7 +29,7 @@ void FastLEDLightOutputComponent::dump_config() {
   ESP_LOGCONFIG(TAG, "  Max refresh rate: %u", *this->max_refresh_rate_);
 }
 void FastLEDLightOutputComponent::loop() {
-  if (!this->next_show_ && !this->is_effect_active())
+  if (!this->should_show_())
     return;
 
   uint32_t now = micros();
@@ -60,7 +38,7 @@ void FastLEDLightOutputComponent::loop() {
     return;
   }
   this->last_refresh_ = now;
-  this->next_show_ = false;
+  this->mark_shown_();
 
   ESP_LOGVV(TAG, "Writing RGB values to bus...");
 
@@ -86,9 +64,6 @@ void FastLEDLightOutputComponent::loop() {
 #endif
   this->controller_->showLeds();
 }
-void FastLEDLightOutputComponent::schedule_show() {
-  this->next_show_ = true;
-}
 CLEDController &FastLEDLightOutputComponent::add_leds(CLEDController *controller, int num_leds) {
   this->controller_ = controller;
   this->num_leds_ = num_leds;
@@ -112,30 +87,20 @@ float FastLEDLightOutputComponent::get_setup_priority() const {
 void FastLEDLightOutputComponent::set_power_supply(PowerSupplyComponent *power_supply) {
   this->power_supply_ = power_supply;
 }
-void FastLEDLightOutputComponent::set_correction(float red, float green, float blue) {
-  this->correction_.set_max_brightness(ESPColor(
-      uint8_t(roundf(red * 255.0f)),
-      uint8_t(roundf(green * 255.0f)),
-      uint8_t(roundf(blue * 255.0f)),
-      0
-  ));
-}
+#endif
+
 ESPColorView FastLEDLightOutputComponent::operator[](int32_t index) const {
   return ESPColorView(&this->leds_[index].r, &this->leds_[index].g, &this->leds_[index].b, nullptr,
                       &this->effect_data_[index], &this->correction_);
 }
 int32_t FastLEDLightOutputComponent::size() const {
   return this->num_leds_;
 }
-void FastLEDLightOutputComponent::setup_state(LightState *state) {
-  this->correction_.calculate_gamma_table(state->get_gamma_correct());
-}
 void FastLEDLightOutputComponent::clear_effect_data() {
   for (int i = 0; i < this->size(); i++)
     this->effect_data_[i] = 0;
 }
 
-#endif
 
 } // namespace light
 

src/esphome/light/fast_led_light_output.h

@@ -36,8 +36,6 @@ namespace light {
  */
 class FastLEDLightOutputComponent : public Component, public AddressableLight {
  public:
-  void schedule_show();
-
   /// Only for custom effects: Get the internal controller.
   CLEDController *get_controller() const;
 
@@ -48,11 +46,9 @@ class FastLEDLightOutputComponent : public Component, public AddressableLight {
   /// Set a maximum refresh rate in µs as some lights do not like being updated too often.
   void set_max_refresh_rate(uint32_t interval_us);
 
+#ifdef USE_OUTPUT
   void set_power_supply(PowerSupplyComponent *power_supply);
-
-  void set_correction(float red, float green, float blue);
-
-  void setup_state(LightState *state) override;
+#endif
 
   /// Add some LEDS, can only be called once.
   CLEDController &add_leds(CLEDController *controller, int num_leds);
@@ -325,7 +321,6 @@ class FastLEDLightOutputComponent : public Component, public AddressableLight {
   // ========== INTERNAL METHODS ==========
   // (In most use cases you won't need these)
   LightTraits get_traits() override;
-  void write_state(LightState *state, int32_t begin, int32_t end) override;
   void setup() override;
   void dump_config() override;
   void loop() override;
@@ -336,7 +331,6 @@ class FastLEDLightOutputComponent : public Component, public AddressableLight {
  protected:
   CLEDController *controller_{nullptr};
   CRGB *leds_{nullptr};
-  ESPColorCorrection correction_{};
   uint8_t *effect_data_{nullptr};
   int num_leds_{0};
   uint32_t last_refresh_{0};

src/esphome/light/neo_pixel_bus_light_output.h

@@ -54,29 +54,21 @@ enum class ESPNeoPixelOrder {
 template<typename T_METHOD, typename T_COLOR_FEATURE>
 class NeoPixelBusLightOutputBase : public Component, public AddressableLight {
  public:
-  void schedule_show();
-
 #ifdef USE_OUTPUT
   void set_power_supply(PowerSupplyComponent *power_supply);
 #endif
 
   NeoPixelBus<T_COLOR_FEATURE, T_METHOD> *get_controller() const;
 
-  void set_correction(float red, float green, float blue, float white = 0.0f);
-
   void clear_effect_data() override;
 
-  void setup_state(LightState *state) override;
-
   /// Add some LEDS, can only be called once.
   void add_leds(uint16_t count_pixels, uint8_t pin);
   void add_leds(uint16_t count_pixels, uint8_t pin_clock, uint8_t pin_data);
   void add_leds(uint16_t count_pixels);
   void add_leds(NeoPixelBus<T_COLOR_FEATURE, T_METHOD> *controller);
 
   // ========== INTERNAL METHODS ==========
-  void write_state(LightState *state, int32_t begin, int32_t end) override;
-
   void setup() override;
 
   void dump_config() override;
@@ -91,8 +83,6 @@ class NeoPixelBusLightOutputBase : public Component, public AddressableLight {
 
  protected:
   NeoPixelBus<T_COLOR_FEATURE, T_METHOD> *controller_{nullptr};
-  bool next_show_{true};
-  ESPColorCorrection correction_{};
   uint8_t *effect_data_{nullptr};
   uint8_t rgb_offsets_[4];
 #ifdef USE_OUTPUT