/*
 * 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];
typedef char ZSTD_rust_ldm_strategy_values_check[
        (ZSTD_fast == 1 && ZSTD_dfast == 2 && ZSTD_btultra2 == 9) ? 1 : -1];

size_t ZSTD_rust_ldm_getTableSize(
        const void* params, int enableLdm, size_t redzoneSize);
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);
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 strategy,
        int useRowMatchFinder);
U32 ZSTD_rust_ldm_limitTableUpdate(U32 curr, U32 nextToUpdate);

enum {
    ZSTD_RUST_LDM_FAST_TABLES_NONE = 0,
    ZSTD_RUST_LDM_FAST_TABLES_FAST = 1,
    ZSTD_RUST_LDM_FAST_TABLES_DFAST = 2
};

int ZSTD_rust_ldm_fastTableKind(int strategy);

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);

typedef size_t (*ZSTD_rust_ldm_select_block_compressor_f)(
        void* blockContext, int useRowMatchFinder);

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

typedef struct {
    void* callbackContext;
    ZSTD_rust_ldm_select_block_compressor_f selectBlockCompressor;
    ZSTD_BlockCompressor_f blockCompressor;
} ZSTD_rust_ldm_block_context;
typedef char ZSTD_rust_ldm_block_context_layout[
        (offsetof(ZSTD_rust_ldm_block_context, callbackContext) == 0
         && offsetof(ZSTD_rust_ldm_block_context, selectBlockCompressor)
                == sizeof(void*)
         && offsetof(ZSTD_rust_ldm_block_context, blockCompressor)
                == 2 * sizeof(void*)
         && sizeof(ZSTD_rust_ldm_block_context) == 3 * sizeof(void*))
        ? 1 : -1];

/* Rust owns when block-compressor selection happens and how its status is
 * handled.  This callback retains the dictionary-mode inspection and the
 * private MatchState/compressor selection in C. */
static size_t ZSTD_rust_ldm_selectBlockCompressor(
        void* blockContext, int useRowMatchFinder)
{
    ZSTD_rust_ldm_block_context* const context =
            (ZSTD_rust_ldm_block_context*)blockContext;
    ZSTD_MatchState_t* const ms = (ZSTD_MatchState_t*)context->callbackContext;
    context->blockCompressor = ZSTD_selectBlockCompressor(
            ms->cParams.strategy, (ZSTD_ParamSwitch_e)useRowMatchFinder,
            ZSTD_matchState_dictMode(ms));
    assert(context->blockCompressor != NULL);
    return context->blockCompressor == NULL ? ERROR(GENERIC) : 0;
}

static void ZSTD_rust_ldm_fillFastTables(ZSTD_MatchState_t* ms, const BYTE* end)
{
    switch (ZSTD_rust_ldm_fastTableKind((int)ms->cParams.strategy)) {
    case ZSTD_RUST_LDM_FAST_TABLES_FAST:
        ZSTD_fillHashTable(ms, end, ZSTD_dtlm_fast, ZSTD_tfp_forCCtx);
        break;
    case ZSTD_RUST_LDM_FAST_TABLES_DFAST:
#ifndef ZSTD_EXCLUDE_DFAST_BLOCK_COMPRESSOR
        ZSTD_fillDoubleHashTable(ms, end, ZSTD_dtlm_fast, ZSTD_tfp_forCCtx);
#else
        assert(0);
#endif
        break;
    case ZSTD_RUST_LDM_FAST_TABLES_NONE:
        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;
    ZSTD_MatchState_t* const ms = (ZSTD_MatchState_t*)context->callbackContext;
    U32 const curr = (U32)((const BYTE*)anchor - ms->window.base);
    ms->nextToUpdate = ZSTD_rust_ldm_limitTableUpdate(curr, ms->nextToUpdate);
    ZSTD_rust_ldm_fillFastTables(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;
    ZSTD_MatchState_t* const ms = (ZSTD_MatchState_t*)context->callbackContext;
    assert(context->blockCompressor != NULL);
    return context->blockCompressor(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;
    ZSTD_MatchState_t* const ms = (ZSTD_MatchState_t*)context->callbackContext;
    ms->ldmSeqStore = (const RawSeqStore_t*)rawSeqStore;
}

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);
}

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);
}

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.callbackContext = ms;
    context.selectBlockCompressor = ZSTD_rust_ldm_selectBlockCompressor;
    context.blockCompressor = NULL;
    return ZSTD_rust_ldm_blockCompress(
            rawSeqStore, &context, seqStore, rep, src, srcSize,
            ms->cParams.minMatch, (int)ms->cParams.strategy,
            (int)useRowMatchFinder);
}