1 files changed, 9 insertions, 15 deletions

diff --git a/drivers/thermal/testing/zone.c b/drivers/thermal/testing/zone.c
index fb83d2501284..2619d3df5bd9 100644
--- a/drivers/thermal/testing/zone.c
+++ b/drivers/thermal/testing/zone.c
@@ -305,6 +305,9 @@ static void tt_put_tt_zone(struct tt_thermal_zone *tt_zone)
 	tt_zone->refcount--;
 }
 
+DEFINE_FREE(put_tt_zone, struct tt_thermal_zone *,
+	    if (!IS_ERR_OR_NULL(_T)) tt_put_tt_zone(_T))
+
 static void tt_zone_add_trip_work_fn(struct work_struct *work)
 {
 	struct tt_work *tt_work = tt_work_of_work(work);
@@ -327,9 +330,9 @@ static void tt_zone_add_trip_work_fn(struct work_struct *work)
 
 int tt_zone_add_trip(const char *arg)
 {
+	struct tt_thermal_zone *tt_zone __free(put_tt_zone) = NULL;
 	struct tt_work *tt_work __free(kfree);
 	struct tt_trip *tt_trip __free(kfree);
-	struct tt_thermal_zone *tt_zone;
 	int id;
 
 	tt_work = kzalloc(sizeof(*tt_work), GFP_KERNEL);
@@ -345,10 +348,8 @@ int tt_zone_add_trip(const char *arg)
 		return PTR_ERR(tt_zone);
 
 	id = ida_alloc(&tt_zone->ida, GFP_KERNEL);
-	if (id < 0) {
-		tt_put_tt_zone(tt_zone);
+	if (id < 0)
 		return id;
-	}
 
 	tt_trip->trip.type = THERMAL_TRIP_ACTIVE;
 	tt_trip->trip.temperature = THERMAL_TEMP_INVALID;
@@ -361,7 +362,7 @@ int tt_zone_add_trip(const char *arg)
 	tt_zone->num_trips++;
 
 	INIT_WORK(&tt_work->work, tt_zone_add_trip_work_fn);
-	tt_work->tt_zone = tt_zone;
+	tt_work->tt_zone = no_free_ptr(tt_zone);
 	tt_work->tt_trip = no_free_ptr(tt_trip);
 	schedule_work(&(no_free_ptr(tt_work)->work));
 
@@ -420,23 +421,18 @@ static int tt_zone_register_tz(struct tt_thermal_zone *tt_zone)
 
 int tt_zone_reg(const char *arg)
 {
-	struct tt_thermal_zone *tt_zone;
-	int ret;
+	struct tt_thermal_zone *tt_zone __free(put_tt_zone);
 
 	tt_zone = tt_get_tt_zone(arg);
 	if (IS_ERR(tt_zone))
 		return PTR_ERR(tt_zone);
 
-	ret = tt_zone_register_tz(tt_zone);
-
-	tt_put_tt_zone(tt_zone);
-
-	return ret;
+	return tt_zone_register_tz(tt_zone);
 }
 
 int tt_zone_unreg(const char *arg)
 {
-	struct tt_thermal_zone *tt_zone;
+	struct tt_thermal_zone *tt_zone __free(put_tt_zone);
 
 	tt_zone = tt_get_tt_zone(arg);
 	if (IS_ERR(tt_zone))
@@ -444,8 +440,6 @@ int tt_zone_unreg(const char *arg)
 
 	tt_zone_unregister_tz(tt_zone);
 
-	tt_put_tt_zone(tt_zone);
-
 	return 0;
 }