This patch should fix shutdown issues on Rock Pi 4 and probably
all other boards using RK808 PMIC.

This is an adaptation of former works:
mfd: rk808: power off system in syscore shutdown
https://github.com/rockchip-linux/kernel/commit/94a7fc2cbdd8b32f6a385e4110b4e1476f684ae4

mfd: rk808: Add RK817 and RK809 support
https://github.com/torvalds/linux/commit/586c1b4125b3c7bf5b482fcafab5d568b8a3c285

mfd: rk808: Check pm_power_off pointer
https://github.com/torvalds/linux/commit/76304994645028accc0cfe287652344b696f4470

It will have to be adapted after switching to Linux Kernel 5.3.
Some of the changes will already be merged there but unfortunately only for RK809 and RK817.

diff --git a/drivers/mfd/rk808.c b/drivers/mfd/rk808.c
index 908c1f45e..969745685 100644
--- a/drivers/mfd/rk808.c
+++ b/drivers/mfd/rk808.c
@@ -19,6 +19,7 @@
 #include <linux/module.h>
 #include <linux/of_device.h>
 #include <linux/regmap.h>
+#include <linux/syscore_ops.h>

 struct rk808_reg_data {
 	int addr;
@@ -415,6 +416,25 @@ static void rk818_device_shutdown(void)
 		dev_err(&rk808_i2c_client->dev, "power off error!\n");
 }

+static void rk8xx_syscore_shutdown(void)
+{
+	struct rk808 *rk808 = i2c_get_clientdata(rk808_i2c_client);
+
+	if (system_state == SYSTEM_POWER_OFF) {
+		switch(rk808->variant) {
+			case RK808_ID:
+				rk808_device_shutdown();
+				break;
+			default:
+				break;
+		}
+	}
+}
+
+static struct syscore_ops rk808_syscore_ops = {
+	.shutdown = rk8xx_syscore_shutdown,
+};
+
 static const struct of_device_id rk808_of_match[] = {
 	{ .compatible = "rockchip,rk805" },
 	{ .compatible = "rockchip,rk808" },
@@ -430,7 +450,6 @@ static int rk808_probe(struct i2c_client *client,
 	struct rk808 *rk808;
 	const struct rk808_reg_data *pre_init_reg;
 	const struct mfd_cell *cells;
-	void (*pm_pwroff_fn)(void);
 	int nr_pre_init_regs;
 	int nr_cells;
 	int pm_off = 0, msb, lsb;
@@ -467,7 +486,7 @@ static int rk808_probe(struct i2c_client *client,
 		nr_pre_init_regs = ARRAY_SIZE(rk805_pre_init_reg);
 		cells = rk805s;
 		nr_cells = ARRAY_SIZE(rk805s);
-		pm_pwroff_fn = rk805_device_shutdown;
+		rk808->pm_pwroff_fn = rk805_device_shutdown;
 		break;
 	case RK808_ID:
 		rk808->regmap_cfg = &rk808_regmap_config;
@@ -476,7 +495,9 @@ static int rk808_probe(struct i2c_client *client,
 		nr_pre_init_regs = ARRAY_SIZE(rk808_pre_init_reg);
 		cells = rk808s;
 		nr_cells = ARRAY_SIZE(rk808s);
-		pm_pwroff_fn = rk808_device_shutdown;
+		rk808->pm_pwroff_fn = rk818_device_shutdown;
+		register_syscore_ops(&rk808_syscore_ops);
+		rk808->registered_syscore_ops = true;
 		break;
 	case RK818_ID:
 		rk808->regmap_cfg = &rk818_regmap_config;
@@ -485,7 +506,7 @@ static int rk808_probe(struct i2c_client *client,
 		nr_pre_init_regs = ARRAY_SIZE(rk818_pre_init_reg);
 		cells = rk818s;
 		nr_cells = ARRAY_SIZE(rk818s);
-		pm_pwroff_fn = rk818_device_shutdown;
+		rk808->pm_pwroff_fn = rk818_device_shutdown;
 		break;
 	default:
 		dev_err(&client->dev, "Unsupported RK8XX ID %lu\n",
@@ -495,6 +516,7 @@ static int rk808_probe(struct i2c_client *client,

 	rk808->i2c = client;
 	i2c_set_clientdata(client, rk808);
+	rk808_i2c_client = client;

 	rk808->regmap = devm_regmap_init_i2c(client, rk808->regmap_cfg);
 	if (IS_ERR(rk808->regmap)) {
@@ -538,9 +560,13 @@ static int rk808_probe(struct i2c_client *client,

 	pm_off = of_property_read_bool(np,
 				"rockchip,system-power-controller");
+	/**
+	 * Below condition is probably never met because PSCI has already claimed
+	 * the value of pm_power_off and that is why boards with rk808 (e.g. rock pi 4)
+	 * never fully shut down and heat up while being "off"
+	 */
 	if (pm_off && !pm_power_off) {
-		rk808_i2c_client = client;
-		pm_power_off = pm_pwroff_fn;
+		pm_power_off = rk808->pm_pwroff_fn;
 	}

 	return 0;
@@ -555,7 +581,16 @@ static int rk808_remove(struct i2c_client *client)
 	struct rk808 *rk808 = i2c_get_clientdata(client);

 	regmap_del_irq_chip(client->irq, rk808->irq_data);
-	pm_power_off = NULL;
+
+	/**
+	 * pm_power_off may point to a function from another module.
+	 * Check if the pointer is set by us and only then overwrite it.
+	 */
+	if (rk808->pm_pwroff_fn && pm_power_off == rk808->pm_pwroff_fn)
+		pm_power_off = NULL;
+
+	if (rk808->registered_syscore_ops)
+		unregister_syscore_ops(&rk808_syscore_ops);

 	return 0;
 }
diff --git a/include/linux/mfd/rk808.h b/include/linux/mfd/rk808.h
index 1d831c722..4d45d8dcc 100644
--- a/include/linux/mfd/rk808.h
+++ b/include/linux/mfd/rk808.h
@@ -445,5 +445,7 @@ struct rk808 {
 	long				variant;
 	const struct regmap_config	*regmap_cfg;
 	const struct regmap_irq_chip	*regmap_irq_chip;
+	void				(*pm_pwroff_fn)(void);
+	bool				registered_syscore_ops;
 };
 #endif /* __LINUX_REGULATOR_RK808_H */