/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 * All rights reserved.
 *
 * This source code is licensed under both the BSD-style license (found in the
 * LICENSE file in the root directory of this source tree) and the GPLv2 (found
 * in the COPYING file in the root directory of this source tree).
 * You may select, at your option, one of the above-listed licenses.
 */

/* The LDM algorithms and sequence-store encoding live in
 * rust/src/zstd_ldm.rs. This file retains private match-state extraction and
 * block-compressor dispatch only. */
#include "zstd_ldm.h"

#include "../common/debug.h"
#include "zstd_fast.h"
#include "zstd_double_fast.h"
#include "zstd_ldm_geartab.h"

typedef struct {
    U32 offBase;
    U16 litLength;
    U16 mlBase;
} ZSTD_rust_seq_def_layout;

typedef struct {
    SeqDef* sequencesStart;
    SeqDef* sequences;
    BYTE* litStart;
    BYTE* lit;
    BYTE* llCode;
    BYTE* mlCode;
    BYTE* ofCode;
    size_t maxNbSeq;
    size_t maxNbLit;
    int longLengthType;
    U32 longLengthPos;
} ZSTD_rust_seq_store_layout;

typedef struct {
    U32 offset;
    U32 checksum;
} ZSTD_rust_ldm_entry_layout;

typedef struct {
    U32 offset;
    U32 litLength;
    U32 matchLength;
} ZSTD_rust_raw_seq_layout;

typedef struct {
    rawSeq* seq;
    size_t pos;
    size_t posInSequence;
    size_t size;
    size_t capacity;
} ZSTD_rust_raw_seq_store_layout;

typedef struct {
    ZSTD_ParamSwitch_e enableLdm;
    U32 hashLog;
    U32 bucketSizeLog;
    U32 minMatchLength;
    U32 hashRateLog;
    U32 windowLog;
} ZSTD_rust_ldm_params_layout;

typedef struct {
    BYTE const* nextSrc;
    BYTE const* base;
    BYTE const* dictBase;
    U32 dictLimit;
    U32 lowLimit;
    U32 nbOverflowCorrections;
} ZSTD_rust_ldm_window_layout;

typedef char ZSTD_rust_ldm_entry_layout_check[
        sizeof(ldmEntry_t) == sizeof(ZSTD_rust_ldm_entry_layout) ? 1 : -1];
typedef char ZSTD_rust_seq_def_layout_check[
        sizeof(SeqDef) == sizeof(ZSTD_rust_seq_def_layout) ? 1 : -1];
typedef char ZSTD_rust_seq_store_layout_check[
        sizeof(SeqStore_t) == sizeof(ZSTD_rust_seq_store_layout) ? 1 : -1];
typedef char ZSTD_rust_raw_seq_layout_check[
        sizeof(rawSeq) == sizeof(ZSTD_rust_raw_seq_layout) ? 1 : -1];
typedef char ZSTD_rust_raw_seq_store_layout_check[
        sizeof(RawSeqStore_t) == sizeof(ZSTD_rust_raw_seq_store_layout) ? 1 : -1];
typedef char ZSTD_rust_ldm_params_layout_check[
        sizeof(ldmParams_t) == sizeof(ZSTD_rust_ldm_params_layout) ? 1 : -1];
typedef char ZSTD_rust_ldm_window_layout_check[
        sizeof(ZSTD_window_t) == sizeof(ZSTD_rust_ldm_window_layout) ? 1 : -1];
typedef char ZSTD_rust_param_switch_layout_check[
        sizeof(ZSTD_ParamSwitch_e) == sizeof(int) ? 1 : -1];

#define ZSTD_RUST_LDM_OFFSET_CHECK(name, c_type, c_field, rust_type, rust_field) \
    typedef char name[offsetof(c_type, c_field) == offsetof(rust_type, rust_field) ? 1 : -1]

ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_ldm_entry_offset_check,
        ldmEntry_t, offset,
        ZSTD_rust_ldm_entry_layout, offset);
ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_ldm_entry_checksum_check,
        ldmEntry_t, checksum,
        ZSTD_rust_ldm_entry_layout, checksum);

ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_seq_def_off_base_check,
        SeqDef, offBase,
        ZSTD_rust_seq_def_layout, offBase);
ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_seq_def_lit_length_check,
        SeqDef, litLength,
        ZSTD_rust_seq_def_layout, litLength);
ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_seq_def_ml_base_check,
        SeqDef, mlBase,
        ZSTD_rust_seq_def_layout, mlBase);

ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_seq_store_sequences_start_check,
        SeqStore_t, sequencesStart,
        ZSTD_rust_seq_store_layout, sequencesStart);
ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_seq_store_sequences_check,
        SeqStore_t, sequences,
        ZSTD_rust_seq_store_layout, sequences);
ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_seq_store_lit_start_check,
        SeqStore_t, litStart,
        ZSTD_rust_seq_store_layout, litStart);
ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_seq_store_lit_check,
        SeqStore_t, lit,
        ZSTD_rust_seq_store_layout, lit);
ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_seq_store_ll_code_check,
        SeqStore_t, llCode,
        ZSTD_rust_seq_store_layout, llCode);
ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_seq_store_ml_code_check,
        SeqStore_t, mlCode,
        ZSTD_rust_seq_store_layout, mlCode);
ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_seq_store_of_code_check,
        SeqStore_t, ofCode,
        ZSTD_rust_seq_store_layout, ofCode);
ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_seq_store_max_nb_seq_check,
        SeqStore_t, maxNbSeq,
        ZSTD_rust_seq_store_layout, maxNbSeq);
ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_seq_store_max_nb_lit_check,
        SeqStore_t, maxNbLit,
        ZSTD_rust_seq_store_layout, maxNbLit);
ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_seq_store_long_length_type_check,
        SeqStore_t, longLengthType,
        ZSTD_rust_seq_store_layout, longLengthType);
ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_seq_store_long_length_pos_check,
        SeqStore_t, longLengthPos,
        ZSTD_rust_seq_store_layout, longLengthPos);

ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_raw_seq_offset_check,
        rawSeq, offset,
        ZSTD_rust_raw_seq_layout, offset);
ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_raw_seq_lit_length_check,
        rawSeq, litLength,
        ZSTD_rust_raw_seq_layout, litLength);
ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_raw_seq_match_length_check,
        rawSeq, matchLength,
        ZSTD_rust_raw_seq_layout, matchLength);

ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_raw_seq_store_seq_check,
        RawSeqStore_t, seq,
        ZSTD_rust_raw_seq_store_layout, seq);
ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_raw_seq_store_pos_check,
        RawSeqStore_t, pos,
        ZSTD_rust_raw_seq_store_layout, pos);
ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_raw_seq_store_pos_in_sequence_check,
        RawSeqStore_t, posInSequence,
        ZSTD_rust_raw_seq_store_layout, posInSequence);
ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_raw_seq_store_size_check,
        RawSeqStore_t, size,
        ZSTD_rust_raw_seq_store_layout, size);
ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_raw_seq_store_capacity_check,
        RawSeqStore_t, capacity,
        ZSTD_rust_raw_seq_store_layout, capacity);

ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_ldm_params_enable_check,
        ldmParams_t, enableLdm,
        ZSTD_rust_ldm_params_layout, enableLdm);
ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_ldm_params_hash_log_check,
        ldmParams_t, hashLog,
        ZSTD_rust_ldm_params_layout, hashLog);
ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_ldm_params_bucket_size_log_check,
        ldmParams_t, bucketSizeLog,
        ZSTD_rust_ldm_params_layout, bucketSizeLog);
ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_ldm_params_min_match_length_check,
        ldmParams_t, minMatchLength,
        ZSTD_rust_ldm_params_layout, minMatchLength);
ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_ldm_params_hash_rate_log_check,
        ldmParams_t, hashRateLog,
        ZSTD_rust_ldm_params_layout, hashRateLog);
ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_ldm_params_window_log_check,
        ldmParams_t, windowLog,
        ZSTD_rust_ldm_params_layout, windowLog);

ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_ldm_window_next_src_check,
        ZSTD_window_t, nextSrc,
        ZSTD_rust_ldm_window_layout, nextSrc);
ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_ldm_window_base_check,
        ZSTD_window_t, base,
        ZSTD_rust_ldm_window_layout, base);
ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_ldm_window_dict_base_check,
        ZSTD_window_t, dictBase,
        ZSTD_rust_ldm_window_layout, dictBase);
ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_ldm_window_dict_limit_check,
        ZSTD_window_t, dictLimit,
        ZSTD_rust_ldm_window_layout, dictLimit);
ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_ldm_window_low_limit_check,
        ZSTD_window_t, lowLimit,
        ZSTD_rust_ldm_window_layout, lowLimit);
ZSTD_RUST_LDM_OFFSET_CHECK(
        ZSTD_rust_ldm_window_overflow_corrections_check,
        ZSTD_window_t, nbOverflowCorrections,
        ZSTD_rust_ldm_window_layout, nbOverflowCorrections);

#undef ZSTD_RUST_LDM_OFFSET_CHECK
typedef char ZSTD_rust_ldm_rep_count[(ZSTD_REP_NUM == 3) ? 1 : -1];

void ZSTD_rust_ldm_adjustParameters(
        void* params, U32 windowLog, int strategy,
        U32 hashLogMax, U32 bucketSizeLogMax, int btultra);
size_t ZSTD_rust_ldm_getTableSize(
        const void* params, int enableLdm, size_t redzoneSize);
size_t ZSTD_rust_ldm_getMaxNbSeq(
        const void* params, int enableLdm, size_t maxChunkSize);
void ZSTD_rust_ldm_fillHashTable(
        void* hashTable, BYTE* bucketOffsets, const BYTE* base,
        const BYTE* ip, const BYTE* iend, const void* params);
size_t ZSTD_rust_ldm_generateSequences(
        void* hashTable, BYTE* bucketOffsets, void* window, U32* loadedDictEnd,
        void* rawSeqStore, const void* params,
        const void* src, size_t srcSize, int overflowCorrectFrequently);
void ZSTD_rust_ldm_skipSequences(void* rawSeqStore, size_t srcSize, U32 minMatch);
void ZSTD_rust_ldm_skipRawSeqStoreBytes(void* rawSeqStore, size_t nbBytes);
size_t ZSTD_rust_ldm_blockCompress(
        void* rawSeqStore, void* blockContext, void* seqStore, U32 rep[ZSTD_REP_NUM],
        const void* src, size_t srcSize, U32 minMatch, int useOptimalParser);
U32 ZSTD_rust_ldm_limitTableUpdate(U32 curr, U32 nextToUpdate);

const U64* ZSTD_ldm_rust_gearTable(void);
void ZSTD_ldm_rust_prepareBlock(void* blockContext, const void* anchor);
size_t ZSTD_ldm_rust_compressLiterals(
        void* blockContext, void* seqStore, U32 rep[ZSTD_REP_NUM],
        const void* src, size_t srcSize);
void ZSTD_ldm_rust_setLdmSeqStore(void* blockContext, const void* rawSeqStore);

const U64* ZSTD_ldm_rust_gearTable(void)
{
    return ZSTD_ldm_gearTab;
}

typedef struct {
    ZSTD_MatchState_t* ms;
    ZSTD_BlockCompressor_f blockCompressor;
} ZSTD_rust_ldm_block_context;

static void ZSTD_rust_ldm_fillFastTables(ZSTD_MatchState_t* ms, const BYTE* end)
{
    switch (ms->cParams.strategy) {
    case ZSTD_fast:
        ZSTD_fillHashTable(ms, end, ZSTD_dtlm_fast, ZSTD_tfp_forCCtx);
        break;
    case ZSTD_dfast:
#ifndef ZSTD_EXCLUDE_DFAST_BLOCK_COMPRESSOR
        ZSTD_fillDoubleHashTable(ms, end, ZSTD_dtlm_fast, ZSTD_tfp_forCCtx);
#else
        assert(0);
#endif
        break;
    case ZSTD_greedy:
    case ZSTD_lazy:
    case ZSTD_lazy2:
    case ZSTD_btlazy2:
    case ZSTD_btopt:
    case ZSTD_btultra:
    case ZSTD_btultra2:
        break;
    default:
        assert(0);
    }
}

void ZSTD_ldm_rust_prepareBlock(void* blockContext, const void* anchor)
{
    ZSTD_rust_ldm_block_context* const context = (ZSTD_rust_ldm_block_context*)blockContext;
    U32 const curr = (U32)((const BYTE*)anchor - context->ms->window.base);
    context->ms->nextToUpdate = ZSTD_rust_ldm_limitTableUpdate(
            curr, context->ms->nextToUpdate);
    ZSTD_rust_ldm_fillFastTables(context->ms, (const BYTE*)anchor);
}

size_t ZSTD_ldm_rust_compressLiterals(
        void* blockContext, void* seqStore, U32 rep[ZSTD_REP_NUM],
        const void* src, size_t srcSize)
{
    ZSTD_rust_ldm_block_context* const context = (ZSTD_rust_ldm_block_context*)blockContext;
    return context->blockCompressor(context->ms, (SeqStore_t*)seqStore, rep, src, srcSize);
}

void ZSTD_ldm_rust_setLdmSeqStore(void* blockContext, const void* rawSeqStore)
{
    ZSTD_rust_ldm_block_context* const context = (ZSTD_rust_ldm_block_context*)blockContext;
    context->ms->ldmSeqStore = (const RawSeqStore_t*)rawSeqStore;
}

void ZSTD_ldm_adjustParameters(ldmParams_t* params,
                               const ZSTD_compressionParameters* cParams)
{
    ZSTD_rust_ldm_adjustParameters(
            params, cParams->windowLog, (int)cParams->strategy,
            ZSTD_HASHLOG_MAX, ZSTD_LDM_BUCKETSIZELOG_MAX, (int)ZSTD_btultra);
}

size_t ZSTD_ldm_getTableSize(ldmParams_t params)
{
#if ZSTD_ADDRESS_SANITIZER && !defined(ZSTD_ASAN_DONT_POISON_WORKSPACE)
    size_t const redzoneSize = ZSTD_CWKSP_ASAN_REDZONE_SIZE;
#else
    size_t const redzoneSize = 0;
#endif
    return ZSTD_rust_ldm_getTableSize(
            &params, params.enableLdm == ZSTD_ps_enable, redzoneSize);
}

size_t ZSTD_ldm_getMaxNbSeq(ldmParams_t params, size_t maxChunkSize)
{
    return ZSTD_rust_ldm_getMaxNbSeq(
            &params, params.enableLdm == ZSTD_ps_enable, maxChunkSize);
}

void ZSTD_ldm_fillHashTable(
        ldmState_t* ldmState, const BYTE* ip,
        const BYTE* iend, ldmParams_t const* params)
{
    ZSTD_rust_ldm_fillHashTable(
            ldmState->hashTable, ldmState->bucketOffsets, ldmState->window.base,
            ip, iend, params);
}

size_t ZSTD_ldm_generateSequences(
        ldmState_t* ldmState, RawSeqStore_t* sequences,
        ldmParams_t const* params, void const* src, size_t srcSize)
{
    assert(ldmState->window.nextSrc >= (const BYTE*)src + srcSize);
    assert(sequences->pos <= sequences->size);
    assert(sequences->size <= sequences->capacity);
    return ZSTD_rust_ldm_generateSequences(
            ldmState->hashTable, ldmState->bucketOffsets, &ldmState->window,
            &ldmState->loadedDictEnd, sequences, params, src, srcSize,
            ZSTD_WINDOW_OVERFLOW_CORRECT_FREQUENTLY);
}

void ZSTD_ldm_skipSequences(RawSeqStore_t* rawSeqStore, size_t srcSize,
                            U32 const minMatch)
{
    ZSTD_rust_ldm_skipSequences(rawSeqStore, srcSize, minMatch);
}

void ZSTD_ldm_skipRawSeqStoreBytes(RawSeqStore_t* rawSeqStore, size_t nbBytes)
{
    ZSTD_rust_ldm_skipRawSeqStoreBytes(rawSeqStore, nbBytes);
}

size_t ZSTD_ldm_blockCompress(RawSeqStore_t* rawSeqStore,
    ZSTD_MatchState_t* ms, SeqStore_t* seqStore, U32 rep[ZSTD_REP_NUM],
    ZSTD_ParamSwitch_e useRowMatchFinder,
    void const* src, size_t srcSize)
{
    ZSTD_rust_ldm_block_context context;
    context.ms = ms;
    context.blockCompressor = ZSTD_selectBlockCompressor(
            ms->cParams.strategy, useRowMatchFinder, ZSTD_matchState_dictMode(ms));
    assert(context.blockCompressor != NULL);
    return ZSTD_rust_ldm_blockCompress(
            rawSeqStore, &context, seqStore, rep, src, srcSize,
            ms->cParams.minMatch, ms->cParams.strategy >= ZSTD_btopt);
}