// SPDX-License-Identifier: GPL-2.0-only
/*
 * Copyright (C) 2024, SUSE LLC
 *
 * Authors: Enzo Matsumiya <ematsumiya@suse.de>
 *
 * Implementation of the LZ77 "plain" compression algorithm, as per MS-XCA spec.
 */
#include <linux/slab.h>
#include <linux/sizes.h>
#include <linux/count_zeros.h>
#include <linux/unaligned.h>

#include "lz77.h"

/*
 * Compression parameters.
 */
#define LZ77_MATCH_MIN_LEN      4
#define LZ77_MATCH_MIN_DIST     1
#define LZ77_MATCH_MAX_DIST     SZ_1K
#define LZ77_HASH_LOG           15
#define LZ77_HASH_SIZE          (1 << LZ77_HASH_LOG)
#define LZ77_STEP_SIZE          sizeof(u64)

static __always_inline u8 lz77_read8(const u8 *ptr)
{
        return get_unaligned(ptr);
}

static __always_inline u64 lz77_read64(const u64 *ptr)
{
        return get_unaligned(ptr);
}

static __always_inline void lz77_write8(u8 *ptr, u8 v)
{
        put_unaligned(v, ptr);
}

static __always_inline void lz77_write16(u16 *ptr, u16 v)
{
        put_unaligned_le16(v, ptr);
}

static __always_inline void lz77_write32(u32 *ptr, u32 v)
{
        put_unaligned_le32(v, ptr);
}

static __always_inline u32 lz77_match_len(const void *wnd, const void *cur, const void *end)
{
        const void *start = cur;
        u64 diff;

        /* Safe for a do/while because otherwise we wouldn't reach here from the main loop. */
        do {
                diff = lz77_read64(cur) ^ lz77_read64(wnd);
                if (!diff) {
                        cur += LZ77_STEP_SIZE;
                        wnd += LZ77_STEP_SIZE;

                        continue;
                }

                /* This computes the number of common bytes in @diff. */
                cur += count_trailing_zeros(diff) >> 3;

                return (cur - start);
        } while (likely(cur + LZ77_STEP_SIZE < end));

        while (cur < end && lz77_read8(cur++) == lz77_read8(wnd++))
                ;

        return (cur - start);
}

static __always_inline void *lz77_write_match(void *dst, void **nib, u32 dist, u32 len)
{
        len -= 3;
        dist--;
        dist <<= 3;

        if (len < 7) {
                lz77_write16(dst, dist + len);

                return dst + 2;
        }

        dist |= 7;
        lz77_write16(dst, dist);
        dst += 2;
        len -= 7;

        if (!*nib) {
                lz77_write8(dst, umin(len, 15));
                *nib = dst;
                dst++;
        } else {
                u8 *b = *nib;

                lz77_write8(b, *b | umin(len, 15) << 4);
                *nib = NULL;
        }

        if (len < 15)
                return dst;

        len -= 15;
        if (len < 255) {
                lz77_write8(dst, len);

                return dst + 1;
        }

        lz77_write8(dst, 0xff);
        dst++;
        len += 7 + 15;
        if (len <= 0xffff) {
                lz77_write16(dst, len);

                return dst + 2;
        }

        lz77_write16(dst, 0);
        dst += 2;
        lz77_write32(dst, len);

        return dst + 4;
}

noinline int lz77_compress(const void *src, u32 slen, void *dst, u32 *dlen)
{
        const void *srcp, *end;
        void *dstp, *nib, *flag_pos;
        u32 flag_count = 0;
        long flag = 0;
        u64 *htable;

        srcp = src;
        end = src + slen;
        dstp = dst;
        nib = NULL;
        flag_pos = dstp;
        dstp += 4;

        htable = kvcalloc(LZ77_HASH_SIZE, sizeof(*htable), GFP_KERNEL);
        if (!htable)
                return -ENOMEM;

        /* Main loop. */
        do {
                u32 dist, len = 0;
                const void *wnd;
                u64 hash;

                hash = ((lz77_read64(srcp) << 24) * 889523592379ULL) >> (64 - LZ77_HASH_LOG);
                wnd = src + htable[hash];
                htable[hash] = srcp - src;
                dist = srcp - wnd;

                if (dist && dist < LZ77_MATCH_MAX_DIST)
                        len = lz77_match_len(wnd, srcp, end);

                if (len < LZ77_MATCH_MIN_LEN) {
                        lz77_write8(dstp, lz77_read8(srcp));

                        dstp++;
                        srcp++;

                        flag <<= 1;
                        flag_count++;
                        if (flag_count == 32) {
                                lz77_write32(flag_pos, flag);
                                flag_count = 0;
                                flag_pos = dstp;
                                dstp += 4;
                        }

                        continue;
                }

                /*
                 * Bail out if @dstp reached >= 7/8 of @slen -- already compressed badly, not worth
                 * going further.
                 */
                if (unlikely(dstp - dst >= slen - (slen >> 3))) {
                        *dlen = slen;
                        goto out;
                }

                dstp = lz77_write_match(dstp, &nib, dist, len);
                srcp += len;

                flag = (flag << 1) | 1;
                flag_count++;
                if (flag_count == 32) {
                        lz77_write32(flag_pos, flag);
                        flag_count = 0;
                        flag_pos = dstp;
                        dstp += 4;
                }
        } while (likely(srcp + LZ77_STEP_SIZE < end));

        while (srcp < end) {
                u32 c = umin(end - srcp, 32 - flag_count);

                memcpy(dstp, srcp, c);

                dstp += c;
                srcp += c;

                flag <<= c;
                flag_count += c;
                if (flag_count == 32) {
                        lz77_write32(flag_pos, flag);
                        flag_count = 0;
                        flag_pos = dstp;
                        dstp += 4;
                }
        }

        flag <<= (32 - flag_count);
        flag |= (1 << (32 - flag_count)) - 1;
        lz77_write32(flag_pos, flag);

        *dlen = dstp - dst;
out:
        kvfree(htable);

        if (*dlen < slen)
                return 0;

        return -EMSGSIZE;
}