// SPDX-License-Identifier: GPL-2.0
/*
 * Copyright (C) 2019 MediaTek Inc.
 * Authors:
 *      Stanley Chu <stanley.chu@mediatek.com>
 *      Peter Wang <peter.wang@mediatek.com>
 */

#include <linux/arm-smccc.h>
#include <linux/bitfield.h>
#include <linux/of.h>
#include <linux/of_address.h>
#include <linux/of_device.h>
#include <linux/phy/phy.h>
#include <linux/platform_device.h>
#include <linux/regulator/consumer.h>
#include <linux/reset.h>
#include <linux/soc/mediatek/mtk_sip_svc.h>

#include "ufshcd.h"
#include "ufshcd-crypto.h"
#include "ufshcd-pltfrm.h"
#include "ufs_quirks.h"
#include "unipro.h"
#include "ufs-mediatek.h"

#define CREATE_TRACE_POINTS
#include "ufs-mediatek-trace.h"

#define ufs_mtk_smc(cmd, val, res) \
        arm_smccc_smc(MTK_SIP_UFS_CONTROL, \
                      cmd, val, 0, 0, 0, 0, 0, &(res))

#define ufs_mtk_va09_pwr_ctrl(res, on) \
        ufs_mtk_smc(UFS_MTK_SIP_VA09_PWR_CTRL, on, res)

#define ufs_mtk_crypto_ctrl(res, enable) \
        ufs_mtk_smc(UFS_MTK_SIP_CRYPTO_CTRL, enable, res)

#define ufs_mtk_ref_clk_notify(on, res) \
        ufs_mtk_smc(UFS_MTK_SIP_REF_CLK_NOTIFICATION, on, res)

#define ufs_mtk_device_reset_ctrl(high, res) \
        ufs_mtk_smc(UFS_MTK_SIP_DEVICE_RESET, high, res)

static struct ufs_dev_fix ufs_mtk_dev_fixups[] = {
        UFS_FIX(UFS_VENDOR_MICRON, UFS_ANY_MODEL,
                UFS_DEVICE_QUIRK_DELAY_AFTER_LPM),
        UFS_FIX(UFS_VENDOR_SKHYNIX, "H9HQ21AFAMZDAR",
                UFS_DEVICE_QUIRK_SUPPORT_EXTENDED_FEATURES),
        END_FIX
};

static const struct of_device_id ufs_mtk_of_match[] = {
        { .compatible = "mediatek,mt8183-ufshci" },
        {},
};

static bool ufs_mtk_is_boost_crypt_enabled(struct ufs_hba *hba)
{
        struct ufs_mtk_host *host = ufshcd_get_variant(hba);

        return !!(host->caps & UFS_MTK_CAP_BOOST_CRYPT_ENGINE);
}

static bool ufs_mtk_is_va09_supported(struct ufs_hba *hba)
{
        struct ufs_mtk_host *host = ufshcd_get_variant(hba);

        return !!(host->caps & UFS_MTK_CAP_VA09_PWR_CTRL);
}

static bool ufs_mtk_is_broken_vcc(struct ufs_hba *hba)
{
        struct ufs_mtk_host *host = ufshcd_get_variant(hba);

        return !!(host->caps & UFS_MTK_CAP_BROKEN_VCC);
}

static void ufs_mtk_cfg_unipro_cg(struct ufs_hba *hba, bool enable)
{
        u32 tmp;

        if (enable) {
                ufshcd_dme_get(hba,
                               UIC_ARG_MIB(VS_SAVEPOWERCONTROL), &tmp);
                tmp = tmp |
                      (1 << RX_SYMBOL_CLK_GATE_EN) |
                      (1 << SYS_CLK_GATE_EN) |
                      (1 << TX_CLK_GATE_EN);
                ufshcd_dme_set(hba,
                               UIC_ARG_MIB(VS_SAVEPOWERCONTROL), tmp);

                ufshcd_dme_get(hba,
                               UIC_ARG_MIB(VS_DEBUGCLOCKENABLE), &tmp);
                tmp = tmp & ~(1 << TX_SYMBOL_CLK_REQ_FORCE);
                ufshcd_dme_set(hba,
                               UIC_ARG_MIB(VS_DEBUGCLOCKENABLE), tmp);
        } else {
                ufshcd_dme_get(hba,
                               UIC_ARG_MIB(VS_SAVEPOWERCONTROL), &tmp);
                tmp = tmp & ~((1 << RX_SYMBOL_CLK_GATE_EN) |
                              (1 << SYS_CLK_GATE_EN) |
                              (1 << TX_CLK_GATE_EN));
                ufshcd_dme_set(hba,
                               UIC_ARG_MIB(VS_SAVEPOWERCONTROL), tmp);

                ufshcd_dme_get(hba,
                               UIC_ARG_MIB(VS_DEBUGCLOCKENABLE), &tmp);
                tmp = tmp | (1 << TX_SYMBOL_CLK_REQ_FORCE);
                ufshcd_dme_set(hba,
                               UIC_ARG_MIB(VS_DEBUGCLOCKENABLE), tmp);
        }
}

static void ufs_mtk_crypto_enable(struct ufs_hba *hba)
{
        struct arm_smccc_res res;

        ufs_mtk_crypto_ctrl(res, 1);
        if (res.a0) {
                dev_info(hba->dev, "%s: crypto enable failed, err: %lu\n",
                         __func__, res.a0);
                hba->caps &= ~UFSHCD_CAP_CRYPTO;
        }
}

static void ufs_mtk_host_reset(struct ufs_hba *hba)
{
        struct ufs_mtk_host *host = ufshcd_get_variant(hba);

        reset_control_assert(host->hci_reset);
        reset_control_assert(host->crypto_reset);
        reset_control_assert(host->unipro_reset);

        usleep_range(100, 110);

        reset_control_deassert(host->unipro_reset);
        reset_control_deassert(host->crypto_reset);
        reset_control_deassert(host->hci_reset);
}

static void ufs_mtk_init_reset_control(struct ufs_hba *hba,
                                       struct reset_control **rc,
                                       char *str)
{
        *rc = devm_reset_control_get(hba->dev, str);
        if (IS_ERR(*rc)) {
                dev_info(hba->dev, "Failed to get reset control %s: %ld\n",
                         str, PTR_ERR(*rc));
                *rc = NULL;
        }
}

static void ufs_mtk_init_reset(struct ufs_hba *hba)
{
        struct ufs_mtk_host *host = ufshcd_get_variant(hba);

        ufs_mtk_init_reset_control(hba, &host->hci_reset,
                                   "hci_rst");
        ufs_mtk_init_reset_control(hba, &host->unipro_reset,
                                   "unipro_rst");
        ufs_mtk_init_reset_control(hba, &host->crypto_reset,
                                   "crypto_rst");
}

static int ufs_mtk_hce_enable_notify(struct ufs_hba *hba,
                                     enum ufs_notify_change_status status)
{
        struct ufs_mtk_host *host = ufshcd_get_variant(hba);
        unsigned long flags;

        if (status == PRE_CHANGE) {
                if (host->unipro_lpm) {
                        hba->vps->hba_enable_delay_us = 0;
                } else {
                        hba->vps->hba_enable_delay_us = 600;
                        ufs_mtk_host_reset(hba);
                }

                if (hba->caps & UFSHCD_CAP_CRYPTO)
                        ufs_mtk_crypto_enable(hba);

                if (host->caps & UFS_MTK_CAP_DISABLE_AH8) {
                        spin_lock_irqsave(hba->host->host_lock, flags);
                        ufshcd_writel(hba, 0,
                                      REG_AUTO_HIBERNATE_IDLE_TIMER);
                        spin_unlock_irqrestore(hba->host->host_lock,
                                               flags);

                        hba->capabilities &= ~MASK_AUTO_HIBERN8_SUPPORT;
                        hba->ahit = 0;
                }
        }

        return 0;
}

static int ufs_mtk_bind_mphy(struct ufs_hba *hba)
{
        struct ufs_mtk_host *host = ufshcd_get_variant(hba);
        struct device *dev = hba->dev;
        struct device_node *np = dev->of_node;
        int err = 0;

        host->mphy = devm_of_phy_get_by_index(dev, np, 0);

        if (host->mphy == ERR_PTR(-EPROBE_DEFER)) {
                /*
                 * UFS driver might be probed before the phy driver does.
                 * In that case we would like to return EPROBE_DEFER code.
                 */
                err = -EPROBE_DEFER;
                dev_info(dev,
                         "%s: required phy hasn't probed yet. err = %d\n",
                        __func__, err);
        } else if (IS_ERR(host->mphy)) {
                err = PTR_ERR(host->mphy);
                if (err != -ENODEV) {
                        dev_info(dev, "%s: PHY get failed %d\n", __func__,
                                 err);
                }
        }

        if (err)
                host->mphy = NULL;
        /*
         * Allow unbound mphy because not every platform needs specific
         * mphy control.
         */
        if (err == -ENODEV)
                err = 0;

        return err;
}

static int ufs_mtk_setup_ref_clk(struct ufs_hba *hba, bool on)
{
        struct ufs_mtk_host *host = ufshcd_get_variant(hba);
        struct arm_smccc_res res;
        ktime_t timeout, time_checked;
        u32 value;

        if (host->ref_clk_enabled == on)
                return 0;

        if (on) {
                ufs_mtk_ref_clk_notify(on, res);
                ufshcd_delay_us(host->ref_clk_ungating_wait_us, 10);
                ufshcd_writel(hba, REFCLK_REQUEST, REG_UFS_REFCLK_CTRL);
        } else {
                ufshcd_writel(hba, REFCLK_RELEASE, REG_UFS_REFCLK_CTRL);
        }

        /* Wait for ack */
        timeout = ktime_add_us(ktime_get(), REFCLK_REQ_TIMEOUT_US);
        do {
                time_checked = ktime_get();
                value = ufshcd_readl(hba, REG_UFS_REFCLK_CTRL);

                /* Wait until ack bit equals to req bit */
                if (((value & REFCLK_ACK) >> 1) == (value & REFCLK_REQUEST))
                        goto out;

                usleep_range(100, 200);
        } while (ktime_before(time_checked, timeout));

        dev_err(hba->dev, "missing ack of refclk req, reg: 0x%x\n", value);

        ufs_mtk_ref_clk_notify(host->ref_clk_enabled, res);

        return -ETIMEDOUT;

out:
        host->ref_clk_enabled = on;
        if (!on) {
                ufshcd_delay_us(host->ref_clk_gating_wait_us, 10);
                ufs_mtk_ref_clk_notify(on, res);
        }

        return 0;
}

static void ufs_mtk_setup_ref_clk_wait_us(struct ufs_hba *hba,
                                          u16 gating_us, u16 ungating_us)
{
        struct ufs_mtk_host *host = ufshcd_get_variant(hba);

        if (hba->dev_info.clk_gating_wait_us) {
                host->ref_clk_gating_wait_us =
                        hba->dev_info.clk_gating_wait_us;
        } else {
                host->ref_clk_gating_wait_us = gating_us;
        }

        host->ref_clk_ungating_wait_us = ungating_us;
}

static int ufs_mtk_wait_link_state(struct ufs_hba *hba, u32 state,
                                   unsigned long max_wait_ms)
{
        ktime_t timeout, time_checked;
        u32 val;

        timeout = ktime_add_ms(ktime_get(), max_wait_ms);
        do {
                time_checked = ktime_get();
                ufshcd_writel(hba, 0x20, REG_UFS_DEBUG_SEL);
                val = ufshcd_readl(hba, REG_UFS_PROBE);
                val = val >> 28;

                if (val == state)
                        return 0;

                /* Sleep for max. 200us */
                usleep_range(100, 200);
        } while (ktime_before(time_checked, timeout));

        if (val == state)
                return 0;

        return -ETIMEDOUT;
}

static int ufs_mtk_mphy_power_on(struct ufs_hba *hba, bool on)
{
        struct ufs_mtk_host *host = ufshcd_get_variant(hba);
        struct phy *mphy = host->mphy;
        struct arm_smccc_res res;
        int ret = 0;

        if (!mphy || !(on ^ host->mphy_powered_on))
                return 0;

        if (on) {
                if (ufs_mtk_is_va09_supported(hba)) {
                        ret = regulator_enable(host->reg_va09);
                        if (ret < 0)
                                goto out;
                        /* wait 200 us to stablize VA09 */
                        usleep_range(200, 210);
                        ufs_mtk_va09_pwr_ctrl(res, 1);
                }
                phy_power_on(mphy);
        } else {
                phy_power_off(mphy);
                if (ufs_mtk_is_va09_supported(hba)) {
                        ufs_mtk_va09_pwr_ctrl(res, 0);
                        ret = regulator_disable(host->reg_va09);
                        if (ret < 0)
                                goto out;
                }
        }
out:
        if (ret) {
                dev_info(hba->dev,
                         "failed to %s va09: %d\n",
                         on ? "enable" : "disable",
                         ret);
        } else {
                host->mphy_powered_on = on;
        }

        return ret;
}

static int ufs_mtk_get_host_clk(struct device *dev, const char *name,
                                struct clk **clk_out)
{
        struct clk *clk;
        int err = 0;

        clk = devm_clk_get(dev, name);
        if (IS_ERR(clk))
                err = PTR_ERR(clk);
        else
                *clk_out = clk;

        return err;
}

static void ufs_mtk_boost_crypt(struct ufs_hba *hba, bool boost)
{
        struct ufs_mtk_host *host = ufshcd_get_variant(hba);
        struct ufs_mtk_crypt_cfg *cfg;
        struct regulator *reg;
        int volt, ret;

        if (!ufs_mtk_is_boost_crypt_enabled(hba))
                return;

        cfg = host->crypt;
        volt = cfg->vcore_volt;
        reg = cfg->reg_vcore;

        ret = clk_prepare_enable(cfg->clk_crypt_mux);
        if (ret) {
                dev_info(hba->dev, "clk_prepare_enable(): %d\n",
                         ret);
                return;
        }

        if (boost) {
                ret = regulator_set_voltage(reg, volt, INT_MAX);
                if (ret) {
                        dev_info(hba->dev,
                                 "failed to set vcore to %d\n", volt);
                        goto out;
                }

                ret = clk_set_parent(cfg->clk_crypt_mux,
                                     cfg->clk_crypt_perf);
                if (ret) {
                        dev_info(hba->dev,
                                 "failed to set clk_crypt_perf\n");
                        regulator_set_voltage(reg, 0, INT_MAX);
                        goto out;
                }
        } else {
                ret = clk_set_parent(cfg->clk_crypt_mux,
                                     cfg->clk_crypt_lp);
                if (ret) {
                        dev_info(hba->dev,
                                 "failed to set clk_crypt_lp\n");
                        goto out;
                }

                ret = regulator_set_voltage(reg, 0, INT_MAX);
                if (ret) {
                        dev_info(hba->dev,
                                 "failed to set vcore to MIN\n");
                }
        }
out:
        clk_disable_unprepare(cfg->clk_crypt_mux);
}

static int ufs_mtk_init_host_clk(struct ufs_hba *hba, const char *name,
                                 struct clk **clk)
{
        int ret;

        ret = ufs_mtk_get_host_clk(hba->dev, name, clk);
        if (ret) {
                dev_info(hba->dev, "%s: failed to get %s: %d", __func__,
                         name, ret);
        }

        return ret;
}

static void ufs_mtk_init_boost_crypt(struct ufs_hba *hba)
{
        struct ufs_mtk_host *host = ufshcd_get_variant(hba);
        struct ufs_mtk_crypt_cfg *cfg;
        struct device *dev = hba->dev;
        struct regulator *reg;
        u32 volt;

        host->crypt = devm_kzalloc(dev, sizeof(*(host->crypt)),
                                   GFP_KERNEL);
        if (!host->crypt)
                goto disable_caps;

        reg = devm_regulator_get_optional(dev, "dvfsrc-vcore");
        if (IS_ERR(reg)) {
                dev_info(dev, "failed to get dvfsrc-vcore: %ld",
                         PTR_ERR(reg));
                goto disable_caps;
        }

        if (of_property_read_u32(dev->of_node, "boost-crypt-vcore-min",
                                 &volt)) {
                dev_info(dev, "failed to get boost-crypt-vcore-min");
                goto disable_caps;
        }

        cfg = host->crypt;
        if (ufs_mtk_init_host_clk(hba, "crypt_mux",
                                  &cfg->clk_crypt_mux))
                goto disable_caps;

        if (ufs_mtk_init_host_clk(hba, "crypt_lp",
                                  &cfg->clk_crypt_lp))
                goto disable_caps;

        if (ufs_mtk_init_host_clk(hba, "crypt_perf",
                                  &cfg->clk_crypt_perf))
                goto disable_caps;

        cfg->reg_vcore = reg;
        cfg->vcore_volt = volt;
        host->caps |= UFS_MTK_CAP_BOOST_CRYPT_ENGINE;

disable_caps:
        return;
}

static void ufs_mtk_init_va09_pwr_ctrl(struct ufs_hba *hba)
{
        struct ufs_mtk_host *host = ufshcd_get_variant(hba);

        host->reg_va09 = regulator_get(hba->dev, "va09");
        if (!host->reg_va09)
                dev_info(hba->dev, "failed to get va09");
        else
                host->caps |= UFS_MTK_CAP_VA09_PWR_CTRL;
}

static void ufs_mtk_init_host_caps(struct ufs_hba *hba)
{
        struct ufs_mtk_host *host = ufshcd_get_variant(hba);
        struct device_node *np = hba->dev->of_node;

        if (of_property_read_bool(np, "mediatek,ufs-boost-crypt"))
                ufs_mtk_init_boost_crypt(hba);

        if (of_property_read_bool(np, "mediatek,ufs-support-va09"))
                ufs_mtk_init_va09_pwr_ctrl(hba);

        if (of_property_read_bool(np, "mediatek,ufs-disable-ah8"))
                host->caps |= UFS_MTK_CAP_DISABLE_AH8;

        if (of_property_read_bool(np, "mediatek,ufs-broken-vcc"))
                host->caps |= UFS_MTK_CAP_BROKEN_VCC;

        dev_info(hba->dev, "caps: 0x%x", host->caps);
}

static void ufs_mtk_scale_perf(struct ufs_hba *hba, bool up)
{
        struct ufs_mtk_host *host = ufshcd_get_variant(hba);

        ufs_mtk_boost_crypt(hba, up);
        ufs_mtk_setup_ref_clk(hba, up);

        if (up)
                phy_power_on(host->mphy);
        else
                phy_power_off(host->mphy);
}

/**
 * ufs_mtk_setup_clocks - enables/disable clocks
 * @hba: host controller instance
 * @on: If true, enable clocks else disable them.
 * @status: PRE_CHANGE or POST_CHANGE notify
 *
 * Returns 0 on success, non-zero on failure.
 */
static int ufs_mtk_setup_clocks(struct ufs_hba *hba, bool on,
                                enum ufs_notify_change_status status)
{
        struct ufs_mtk_host *host = ufshcd_get_variant(hba);
        bool clk_pwr_off = false;
        int ret = 0;

        /*
         * In case ufs_mtk_init() is not yet done, simply ignore.
         * This ufs_mtk_setup_clocks() shall be called from
         * ufs_mtk_init() after init is done.
         */
        if (!host)
                return 0;

        if (!on && status == PRE_CHANGE) {
                if (ufshcd_is_link_off(hba)) {
                        clk_pwr_off = true;
                } else if (ufshcd_is_link_hibern8(hba) ||
                         (!ufshcd_can_hibern8_during_gating(hba) &&
                         ufshcd_is_auto_hibern8_enabled(hba))) {
                        /*
                         * Gate ref-clk and poweroff mphy if link state is in
                         * OFF or Hibern8 by either Auto-Hibern8 or
                         * ufshcd_link_state_transition().
                         */
                        ret = ufs_mtk_wait_link_state(hba,
                                                      VS_LINK_HIBERN8,
                                                      15);
                        if (!ret)
                                clk_pwr_off = true;
                }

                if (clk_pwr_off)
                        ufs_mtk_scale_perf(hba, false);
        } else if (on && status == POST_CHANGE) {
                ufs_mtk_scale_perf(hba, true);
        }

        return ret;
}

static void ufs_mtk_get_controller_version(struct ufs_hba *hba)
{
        struct ufs_mtk_host *host = ufshcd_get_variant(hba);
        int ret, ver = 0;

        if (host->hw_ver.major)
                return;

        /* Set default (minimum) version anyway */
        host->hw_ver.major = 2;

        ret = ufshcd_dme_get(hba, UIC_ARG_MIB(PA_LOCALVERINFO), &ver);
        if (!ret) {
                if (ver >= UFS_UNIPRO_VER_1_8) {
                        host->hw_ver.major = 3;
                        /*
                         * Fix HCI version for some platforms with
                         * incorrect version
                         */
                        if (hba->ufs_version < ufshci_version(3, 0))
                                hba->ufs_version = ufshci_version(3, 0);
                }
        }
}

static u32 ufs_mtk_get_ufs_hci_version(struct ufs_hba *hba)
{
        return hba->ufs_version;
}

/**
 * ufs_mtk_init - find other essential mmio bases
 * @hba: host controller instance
 *
 * Binds PHY with controller and powers up PHY enabling clocks
 * and regulators.
 *
 * Returns -EPROBE_DEFER if binding fails, returns negative error
 * on phy power up failure and returns zero on success.
 */
static int ufs_mtk_init(struct ufs_hba *hba)
{
        const struct of_device_id *id;
        struct device *dev = hba->dev;
        struct ufs_mtk_host *host;
        int err = 0;

        host = devm_kzalloc(dev, sizeof(*host), GFP_KERNEL);
        if (!host) {
                err = -ENOMEM;
                dev_info(dev, "%s: no memory for mtk ufs host\n", __func__);
                goto out;
        }

        host->hba = hba;
        ufshcd_set_variant(hba, host);

        id = of_match_device(ufs_mtk_of_match, dev);
        if (!id) {
                err = -EINVAL;
                goto out;
        }

        /* Initialize host capability */
        ufs_mtk_init_host_caps(hba);

        err = ufs_mtk_bind_mphy(hba);
        if (err)
                goto out_variant_clear;

        ufs_mtk_init_reset(hba);

        /* Enable runtime autosuspend */
        hba->caps |= UFSHCD_CAP_RPM_AUTOSUSPEND;

        /* Enable clock-gating */
        hba->caps |= UFSHCD_CAP_CLK_GATING;

        /* Enable inline encryption */
        hba->caps |= UFSHCD_CAP_CRYPTO;

        /* Enable WriteBooster */
        hba->caps |= UFSHCD_CAP_WB_EN;
        hba->quirks |= UFSHCI_QUIRK_SKIP_MANUAL_WB_FLUSH_CTRL;
        hba->vps->wb_flush_threshold = UFS_WB_BUF_REMAIN_PERCENT(80);

        if (host->caps & UFS_MTK_CAP_DISABLE_AH8)
                hba->caps |= UFSHCD_CAP_HIBERN8_WITH_CLK_GATING;

        /*
         * ufshcd_vops_init() is invoked after
         * ufshcd_setup_clock(true) in ufshcd_hba_init() thus
         * phy clock setup is skipped.
         *
         * Enable phy clocks specifically here.
         */
        ufs_mtk_mphy_power_on(hba, true);
        ufs_mtk_setup_clocks(hba, true, POST_CHANGE);

        goto out;

out_variant_clear:
        ufshcd_set_variant(hba, NULL);
out:
        return err;
}

static int ufs_mtk_pre_pwr_change(struct ufs_hba *hba,
                                  struct ufs_pa_layer_attr *dev_max_params,
                                  struct ufs_pa_layer_attr *dev_req_params)
{
        struct ufs_mtk_host *host = ufshcd_get_variant(hba);
        struct ufs_dev_params host_cap;
        int ret;

        ufshcd_init_pwr_dev_param(&host_cap);
        host_cap.hs_rx_gear = UFS_HS_G4;
        host_cap.hs_tx_gear = UFS_HS_G4;

        ret = ufshcd_get_pwr_dev_param(&host_cap,
                                       dev_max_params,
                                       dev_req_params);
        if (ret) {
                pr_info("%s: failed to determine capabilities\n",
                        __func__);
        }

        if (host->hw_ver.major >= 3) {
                ret = ufshcd_dme_configure_adapt(hba,
                                           dev_req_params->gear_tx,
                                           PA_INITIAL_ADAPT);
        }

        return ret;
}

static int ufs_mtk_pwr_change_notify(struct ufs_hba *hba,
                                     enum ufs_notify_change_status stage,
                                     struct ufs_pa_layer_attr *dev_max_params,
                                     struct ufs_pa_layer_attr *dev_req_params)
{
        int ret = 0;

        switch (stage) {
        case PRE_CHANGE:
                ret = ufs_mtk_pre_pwr_change(hba, dev_max_params,
                                             dev_req_params);
                break;
        case POST_CHANGE:
                break;
        default:
                ret = -EINVAL;
                break;
        }

        return ret;
}

static int ufs_mtk_unipro_set_lpm(struct ufs_hba *hba, bool lpm)
{
        int ret;
        struct ufs_mtk_host *host = ufshcd_get_variant(hba);

        ret = ufshcd_dme_set(hba,
                             UIC_ARG_MIB_SEL(VS_UNIPROPOWERDOWNCONTROL, 0),
                             lpm ? 1 : 0);
        if (!ret || !lpm) {
                /*
                 * Forcibly set as non-LPM mode if UIC commands is failed
                 * to use default hba_enable_delay_us value for re-enabling
                 * the host.
                 */
                host->unipro_lpm = lpm;
        }

        return ret;
}

static int ufs_mtk_pre_link(struct ufs_hba *hba)
{
        int ret;
        u32 tmp;

        ufs_mtk_get_controller_version(hba);

        ret = ufs_mtk_unipro_set_lpm(hba, false);
        if (ret)
                return ret;

        /*
         * Setting PA_Local_TX_LCC_Enable to 0 before link startup
         * to make sure that both host and device TX LCC are disabled
         * once link startup is completed.
         */
        ret = ufshcd_disable_host_tx_lcc(hba);
        if (ret)
                return ret;

        /* disable deep stall */
        ret = ufshcd_dme_get(hba, UIC_ARG_MIB(VS_SAVEPOWERCONTROL), &tmp);
        if (ret)
                return ret;

        tmp &= ~(1 << 6);

        ret = ufshcd_dme_set(hba, UIC_ARG_MIB(VS_SAVEPOWERCONTROL), tmp);

        return ret;
}

static void ufs_mtk_setup_clk_gating(struct ufs_hba *hba)
{
        unsigned long flags;
        u32 ah_ms;

        if (ufshcd_is_clkgating_allowed(hba)) {
                if (ufshcd_is_auto_hibern8_supported(hba) && hba->ahit)
                        ah_ms = FIELD_GET(UFSHCI_AHIBERN8_TIMER_MASK,
                                          hba->ahit);
                else
                        ah_ms = 10;
                spin_lock_irqsave(hba->host->host_lock, flags);
                hba->clk_gating.delay_ms = ah_ms + 5;
                spin_unlock_irqrestore(hba->host->host_lock, flags);
        }
}

static int ufs_mtk_post_link(struct ufs_hba *hba)
{
        /* enable unipro clock gating feature */
        ufs_mtk_cfg_unipro_cg(hba, true);

        /* will be configured during probe hba */
        if (ufshcd_is_auto_hibern8_supported(hba))
                hba->ahit = FIELD_PREP(UFSHCI_AHIBERN8_TIMER_MASK, 10) |
                        FIELD_PREP(UFSHCI_AHIBERN8_SCALE_MASK, 3);

        ufs_mtk_setup_clk_gating(hba);

        return 0;
}

static int ufs_mtk_link_startup_notify(struct ufs_hba *hba,
                                       enum ufs_notify_change_status stage)
{
        int ret = 0;

        switch (stage) {
        case PRE_CHANGE:
                ret = ufs_mtk_pre_link(hba);
                break;
        case POST_CHANGE:
                ret = ufs_mtk_post_link(hba);
                break;
        default:
                ret = -EINVAL;
                break;
        }

        return ret;
}

static int ufs_mtk_device_reset(struct ufs_hba *hba)
{
        struct arm_smccc_res res;

        /* disable hba before device reset */
        ufshcd_hba_stop(hba);

        ufs_mtk_device_reset_ctrl(0, res);

        /*
         * The reset signal is active low. UFS devices shall detect
         * more than or equal to 1us of positive or negative RST_n
         * pulse width.
         *
         * To be on safe side, keep the reset low for at least 10us.
         */
        usleep_range(10, 15);

        ufs_mtk_device_reset_ctrl(1, res);

        /* Some devices may need time to respond to rst_n */
        usleep_range(10000, 15000);

        dev_info(hba->dev, "device reset done\n");

        return 0;
}

static int ufs_mtk_link_set_hpm(struct ufs_hba *hba)
{
        int err;

        err = ufshcd_hba_enable(hba);
        if (err)
                return err;

        err = ufs_mtk_unipro_set_lpm(hba, false);
        if (err)
                return err;

        err = ufshcd_uic_hibern8_exit(hba);
        if (!err)
                ufshcd_set_link_active(hba);
        else
                return err;

        err = ufshcd_make_hba_operational(hba);
        if (err)
                return err;

        return 0;
}

static int ufs_mtk_link_set_lpm(struct ufs_hba *hba)
{
        int err;

        err = ufs_mtk_unipro_set_lpm(hba, true);
        if (err) {
                /* Resume UniPro state for following error recovery */
                ufs_mtk_unipro_set_lpm(hba, false);
                return err;
        }

        return 0;
}

static void ufs_mtk_vreg_set_lpm(struct ufs_hba *hba, bool lpm)
{
        if (!hba->vreg_info.vccq2 || !hba->vreg_info.vcc)
                return;

        if (lpm && !hba->vreg_info.vcc->enabled)
                regulator_set_mode(hba->vreg_info.vccq2->reg,
                                   REGULATOR_MODE_IDLE);
        else if (!lpm)
                regulator_set_mode(hba->vreg_info.vccq2->reg,
                                   REGULATOR_MODE_NORMAL);
}

static int ufs_mtk_suspend(struct ufs_hba *hba, enum ufs_pm_op pm_op)
{
        int err;
        struct arm_smccc_res res;

        if (ufshcd_is_link_hibern8(hba)) {
                err = ufs_mtk_link_set_lpm(hba);
                if (err)
                        goto fail;
        }

        if (!ufshcd_is_link_active(hba)) {
                /*
                 * Make sure no error will be returned to prevent
                 * ufshcd_suspend() re-enabling regulators while vreg is still
                 * in low-power mode.
                 */
                ufs_mtk_vreg_set_lpm(hba, true);
                err = ufs_mtk_mphy_power_on(hba, false);
                if (err)
                        goto fail;
        }

        if (ufshcd_is_link_off(hba))
                ufs_mtk_device_reset_ctrl(0, res);

        return 0;
fail:
        /*
         * Set link as off state enforcedly to trigger
         * ufshcd_host_reset_and_restore() in ufshcd_suspend()
         * for completed host reset.
         */
        ufshcd_set_link_off(hba);
        return -EAGAIN;
}

static int ufs_mtk_resume(struct ufs_hba *hba, enum ufs_pm_op pm_op)
{
        int err;

        err = ufs_mtk_mphy_power_on(hba, true);
        if (err)
                goto fail;

        ufs_mtk_vreg_set_lpm(hba, false);

        if (ufshcd_is_link_hibern8(hba)) {
                err = ufs_mtk_link_set_hpm(hba);
                if (err)
                        goto fail;
        }

        return 0;
fail:
        return ufshcd_link_recovery(hba);
}

static void ufs_mtk_dbg_register_dump(struct ufs_hba *hba)
{
        ufshcd_dump_regs(hba, REG_UFS_REFCLK_CTRL, 0x4, "Ref-Clk Ctrl ");

        ufshcd_dump_regs(hba, REG_UFS_EXTREG, 0x4, "Ext Reg ");

        ufshcd_dump_regs(hba, REG_UFS_MPHYCTRL,
                         REG_UFS_REJECT_MON - REG_UFS_MPHYCTRL + 4,

                         "MPHY Ctrl ");

        /* Direct debugging information to REG_MTK_PROBE */
        ufshcd_writel(hba, 0x20, REG_UFS_DEBUG_SEL);
        ufshcd_dump_regs(hba, REG_UFS_PROBE, 0x4, "Debug Probe ");
}

static int ufs_mtk_apply_dev_quirks(struct ufs_hba *hba)
{
        struct ufs_dev_info *dev_info = &hba->dev_info;
        u16 mid = dev_info->wmanufacturerid;

        if (mid == UFS_VENDOR_SAMSUNG)
                ufshcd_dme_set(hba, UIC_ARG_MIB(PA_TACTIVATE), 6);

        /*
         * Decide waiting time before gating reference clock and
         * after ungating reference clock according to vendors'
         * requirements.
         */
        if (mid == UFS_VENDOR_SAMSUNG)
                ufs_mtk_setup_ref_clk_wait_us(hba, 1, 1);
        else if (mid == UFS_VENDOR_SKHYNIX)
                ufs_mtk_setup_ref_clk_wait_us(hba, 30, 30);
        else if (mid == UFS_VENDOR_TOSHIBA)
                ufs_mtk_setup_ref_clk_wait_us(hba, 100, 32);

        return 0;
}

static void ufs_mtk_fixup_dev_quirks(struct ufs_hba *hba)
{
        ufshcd_fixup_dev_quirks(hba, ufs_mtk_dev_fixups);

        if (ufs_mtk_is_broken_vcc(hba) && hba->vreg_info.vcc &&
            (hba->dev_quirks & UFS_DEVICE_QUIRK_DELAY_AFTER_LPM)) {
                hba->vreg_info.vcc->always_on = true;
                /*
                 * VCC will be kept always-on thus we don't
                 * need any delay during regulator operations
                 */
                hba->dev_quirks &= ~(UFS_DEVICE_QUIRK_DELAY_BEFORE_LPM |
                        UFS_DEVICE_QUIRK_DELAY_AFTER_LPM);
        }
}

static void ufs_mtk_event_notify(struct ufs_hba *hba,
                                 enum ufs_event_type evt, void *data)
{
        unsigned int val = *(u32 *)data;

        trace_ufs_mtk_event(evt, val);
}

/*
 * struct ufs_hba_mtk_vops - UFS MTK specific variant operations
 *
 * The variant operations configure the necessary controller and PHY
 * handshake during initialization.
 */
static const struct ufs_hba_variant_ops ufs_hba_mtk_vops = {
        .name                = "mediatek.ufshci",
        .init                = ufs_mtk_init,
        .get_ufs_hci_version = ufs_mtk_get_ufs_hci_version,
        .setup_clocks        = ufs_mtk_setup_clocks,
        .hce_enable_notify   = ufs_mtk_hce_enable_notify,
        .link_startup_notify = ufs_mtk_link_startup_notify,
        .pwr_change_notify   = ufs_mtk_pwr_change_notify,
        .apply_dev_quirks    = ufs_mtk_apply_dev_quirks,
        .fixup_dev_quirks    = ufs_mtk_fixup_dev_quirks,
        .suspend             = ufs_mtk_suspend,
        .resume              = ufs_mtk_resume,
        .dbg_register_dump   = ufs_mtk_dbg_register_dump,
        .device_reset        = ufs_mtk_device_reset,
        .event_notify        = ufs_mtk_event_notify,
};

/**
 * ufs_mtk_probe - probe routine of the driver
 * @pdev: pointer to Platform device handle
 *
 * Return zero for success and non-zero for failure
 */
static int ufs_mtk_probe(struct platform_device *pdev)
{
        int err;
        struct device *dev = &pdev->dev;
        struct device_node *reset_node;
        struct platform_device *reset_pdev;
        struct device_link *link;

        reset_node = of_find_compatible_node(NULL, NULL,
                                             "ti,syscon-reset");
        if (!reset_node) {
                dev_notice(dev, "find ti,syscon-reset fail\n");
                goto skip_reset;
        }
        reset_pdev = of_find_device_by_node(reset_node);
        if (!reset_pdev) {
                dev_notice(dev, "find reset_pdev fail\n");
                goto skip_reset;
        }
        link = device_link_add(dev, &reset_pdev->dev,
                DL_FLAG_AUTOPROBE_CONSUMER);
        if (!link) {
                dev_notice(dev, "add reset device_link fail\n");
                goto skip_reset;
        }
        /* supplier is not probed */
        if (link->status == DL_STATE_DORMANT) {
                err = -EPROBE_DEFER;
                goto out;
        }

skip_reset:
        /* perform generic probe */
        err = ufshcd_pltfrm_init(pdev, &ufs_hba_mtk_vops);

out:
        if (err)
                dev_info(dev, "probe failed %d\n", err);

        of_node_put(reset_node);
        return err;
}

/**
 * ufs_mtk_remove - set driver_data of the device to NULL
 * @pdev: pointer to platform device handle
 *
 * Always return 0
 */
static int ufs_mtk_remove(struct platform_device *pdev)
{
        struct ufs_hba *hba =  platform_get_drvdata(pdev);

        pm_runtime_get_sync(&(pdev)->dev);
        ufshcd_remove(hba);
        return 0;
}

static const struct dev_pm_ops ufs_mtk_pm_ops = {
        SET_SYSTEM_SLEEP_PM_OPS(ufshcd_system_suspend, ufshcd_system_resume)
        SET_RUNTIME_PM_OPS(ufshcd_runtime_suspend, ufshcd_runtime_resume, NULL)
        .prepare         = ufshcd_suspend_prepare,
        .complete        = ufshcd_resume_complete,
};

static struct platform_driver ufs_mtk_pltform = {
        .probe      = ufs_mtk_probe,
        .remove     = ufs_mtk_remove,
        .shutdown   = ufshcd_pltfrm_shutdown,
        .driver = {
                .name   = "ufshcd-mtk",
                .pm     = &ufs_mtk_pm_ops,
                .of_match_table = ufs_mtk_of_match,
        },
};

MODULE_AUTHOR("Stanley Chu <stanley.chu@mediatek.com>");
MODULE_AUTHOR("Peter Wang <peter.wang@mediatek.com>");
MODULE_DESCRIPTION("MediaTek UFS Host Driver");
MODULE_LICENSE("GPL v2");

module_platform_driver(ufs_mtk_pltform);