/* SPDX-License-Identifier: GPL-2.0 */
#ifndef ARCH_S390_ENTRY_PERCPU_H
#define ARCH_S390_ENTRY_PERCPU_H

#include <linux/kprobes.h>
#include <linux/percpu.h>
#include <asm/lowcore.h>
#include <asm/ptrace.h>
#include <asm/asm-offsets.h>

static __always_inline void percpu_entry(struct pt_regs *regs)
{
        struct lowcore *lc = get_lowcore();

        if (user_mode(regs))
                return;
        regs->cpu = lc->cpu_nr;
        regs->percpu_register = lc->percpu_register;
        lc->percpu_register = 0;
}

static __always_inline bool percpu_code_check(struct pt_regs *regs)
{
        unsigned int insn, disp;
        struct kprobe *p;

        if (likely(user_mode(regs) || !regs->percpu_register))
                return false;
        /*
         * Within a percpu code section - check if the percpu base register
         * needs to be updated. This is the case if the PSW does not point to
         * the ADD instruction within the section.
         * - AG %rx,percpu_offset_in_lowcore(%r0,%r0)
         * which adds the percpu offset to the percpu base register.
         */
        lockdep_assert_preemption_disabled();
again:
        insn = READ_ONCE(*(u16 *)psw_bits(regs->psw).ia);
        if (unlikely(insn == BREAKPOINT_INSTRUCTION)) {
                p = get_kprobe((void *)psw_bits(regs->psw).ia);
                /*
                 * If the kprobe is concurrently removed on a different CPU
                 * it might not be found anymore. However text must have
                 * been restored - try again.
                 */
                if (!p)
                        goto again;
                insn = p->opcode;
        }
        if ((insn & 0xff0f) != 0xe300)
                return true;
        disp = offsetof(struct lowcore, percpu_offset);
        if (machine_has_relocated_lowcore())
                disp += LOWCORE_ALT_ADDRESS;
        insn = (disp & 0xff000) >> 4 | (disp & 0x00fff) << 16 | 0x8;
        if (*(u32 *)(psw_bits(regs->psw).ia + 2) != insn)
                return true;
        return false;
}

static __always_inline void percpu_exit(struct pt_regs *regs, bool needs_fixup)
{
        struct lowcore *lc = get_lowcore();
        unsigned char reg;

        if (user_mode(regs))
                return;
        reg = regs->percpu_register;
        lc->percpu_register = reg;
        if (likely(!needs_fixup))
                return;
        /* Check if process has been migrated to a different CPU. */
        if (regs->cpu == lc->cpu_nr)
                return;
        /* Fixup percpu base register */
        regs->gprs[reg] -= __per_cpu_offset[regs->cpu];
        regs->gprs[reg] += lc->percpu_offset;
}

#endif