3 files changed, 27 insertions, 17 deletions

diff --git a/include/wlr/types/wlr_output.h b/include/wlr/types/wlr_output.h
index e096a3ce..93943f0d 100644
--- a/include/wlr/types/wlr_output.h
+++ b/include/wlr/types/wlr_output.h
@@ -571,6 +571,8 @@ void wlr_output_state_set_subpixel(struct wlr_output_state *state,
 	enum wl_output_subpixel subpixel);
 void wlr_output_state_set_buffer(struct wlr_output_state *state,
 	struct wlr_buffer *buffer);
+bool wlr_output_state_set_gamma_lut(struct wlr_output_state *state,
+	size_t ramp_size, const uint16_t *r, const uint16_t *g, const uint16_t *b);
 
 
 /**
diff --git a/types/output/output.c b/types/output/output.c
index e212928c..76ee54e6 100644
--- a/types/output/output.c
+++ b/types/output/output.c
@@ -945,23 +945,7 @@ void wlr_output_send_request_state(struct wlr_output *output,
 
 void wlr_output_set_gamma(struct wlr_output *output, size_t size,
 		const uint16_t *r, const uint16_t *g, const uint16_t *b) {
-	uint16_t *gamma_lut = NULL;
-	if (size > 0) {
-		gamma_lut = malloc(3 * size * sizeof(uint16_t));
-		if (gamma_lut == NULL) {
-			wlr_log_errno(WLR_ERROR, "Allocation failed");
-			return;
-		}
-		memcpy(gamma_lut, r, size * sizeof(uint16_t));
-		memcpy(gamma_lut + size, g, size * sizeof(uint16_t));
-		memcpy(gamma_lut + 2 * size, b, size * sizeof(uint16_t));
-	}
-
-	output_state_clear_gamma_lut(&output->pending);
-
-	output->pending.gamma_lut_size = size;
-	output->pending.gamma_lut = gamma_lut;
-	output->pending.committed |= WLR_OUTPUT_STATE_GAMMA_LUT;
+	wlr_output_state_set_gamma_lut(&output->pending, size, r, g, b);
 }
 
 size_t wlr_output_get_gamma_size(struct wlr_output *output) {
diff --git a/types/output/state.c b/types/output/state.c
index b54c251e..cbbdc5c0 100644
--- a/types/output/state.c
+++ b/types/output/state.c
@@ -1,4 +1,6 @@
 #include <stdlib.h>
+#include <string.h>
+#include <wlr/util/log.h>
 #include "types/wlr_output.h"
 
 void wlr_output_state_finish(struct wlr_output_state *state) {
@@ -73,3 +75,25 @@ void wlr_output_state_set_buffer(struct wlr_output_state *state,
 	wlr_buffer_unlock(state->buffer);
 	state->buffer = wlr_buffer_lock(buffer);
 }
+
+bool wlr_output_state_set_gamma_lut(struct wlr_output_state *state,
+		size_t ramp_size, const uint16_t *r, const uint16_t *g, const uint16_t *b) {
+	uint16_t *gamma_lut = NULL;
+	if (ramp_size > 0) {
+		gamma_lut = realloc(state->gamma_lut, 3 * ramp_size * sizeof(uint16_t));
+		if (gamma_lut == NULL) {
+			wlr_log_errno(WLR_ERROR, "Allocation failed");
+			return false;
+		}
+		memcpy(gamma_lut, r, ramp_size * sizeof(uint16_t));
+		memcpy(gamma_lut + ramp_size, g, ramp_size * sizeof(uint16_t));
+		memcpy(gamma_lut + 2 * ramp_size, b, ramp_size * sizeof(uint16_t));
+	} else {
+		free(state->gamma_lut);
+	}
+
+	state->committed |= WLR_OUTPUT_STATE_GAMMA_LUT;
+	state->gamma_lut_size = ramp_size;
+	state->gamma_lut = gamma_lut;
+	return true;
+}