/**
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

#pragma once

#include <stdint.h>
#include <vector>

// Utilities for bit packing and unpacking CPU non-interleaved and GPU
// interleaved by 32 encodings
namespace faiss {
namespace gpu {

// Unpacks arbitrary bitwidth codes to a whole number of bytes per code
// The layout of the input is (v0 d0)(v0 d1) ... (v0 dD)(v1 d0) ...
// (bit packed)
// The layout of the output is the same (byte packed to roundUp(bitsPerCode, 8)
// / 8 bytes)
std::vector<uint8_t> unpackNonInterleaved(
        std::vector<uint8_t> data,
        int numVecs,
        int dims,
        int bitsPerCode);

// Unpacks arbitrary bitwidth codes to a whole number of bytes per scalar code
// The layout of the input is (v0 d0)(v1 d0) ... (v31 d0)(v0 d1) ...
// (bit packed)
// The layout of the input is (v0 d0)(v0 d1) ... (v0 dD)(v1 d0) ...
// (byte packed)
std::vector<uint8_t> unpackInterleaved(
        std::vector<uint8_t> data,
        int numVecs,
        int dims,
        int bitsPerCode);

// Packs data in the byte packed non-interleaved form to bit packed
// non-interleaved form
std::vector<uint8_t> packNonInterleaved(
        std::vector<uint8_t> data,
        int numVecs,
        int dims,
        int bitsPerCode);

// Packs data in the byte packed non-interleaved form to bit packed
// interleaved form
std::vector<uint8_t> packInterleaved(
        std::vector<uint8_t> data,
        int numVecs,
        int dims,
        int bitsPerCode);

} // namespace gpu
} // namespace faiss
