<commit_before>#pragma once
#ifndef NBT_GZ_PARSE_HPP_
#define NBT_GZ_PARSE_HPP_

#include <filesystem>
#include <nbt/nbt.hpp>
#include <stack>
#include <zlib.h>

#define _NTOHS(ptr) (int16_t(((ptr)[0] << 8) + (ptr)[1]))
#define _NTOHI(ptr)                                                            \
  ((uint32_t((ptr)[0]) << 24) + (uint32_t((ptr)[1]) << 16) +                   \
   (uint32_t((ptr)[2]) << 8) + uint32_t((ptr)[3]))
#define _NTOHL(ptr)                                                            \
  ((uint64_t((ptr)[0]) << 56) + (uint64_t((ptr)[1]) << 48) +                   \
   (uint64_t((ptr)[2]) << 40) + (uint64_t((ptr)[3]) << 32) +                   \
   (uint64_t((ptr)[4]) << 24) + (uint64_t((ptr)[5]) << 16) +                   \
   (uint64_t((ptr)[6]) << 8) + uint64_t((ptr)[7]))

// Max size of a single element to read in memory (string)
#define MAXELEMENTSIZE 65025

// Check if the context indicates a being in a list
#define LIST (context.size() && context.top().second < tag_type::tag_long_array)

union FloatTranslator {
  // See this union as an uint8_t[8] array. Put an integer in the buffer, get
  // its byte reinterpretation in floating point, WITHOUT THE WARNINGS !
  uint64_t buffer;
  float _float;
  double _double;

  FloatTranslator(uint32_t float_bytes) : buffer(float_bytes){};
  FloatTranslator(uint64_t double_bytes) : buffer(double_bytes){};
};

struct ByteStream {
  // Adapter for the matryoshkas to work with both files and memory buffers
  enum BufferType { MEMORY, GZFILE };

  union Source {
    gzFile file;
    std::pair<uint8_t *, size_t> array;

    Source(gzFile f) : file(f){};
    Source(uint8_t *address, size_t size) : array({address, size}){};
  };

  BufferType type;
  Source source;

  ByteStream(gzFile f) : type(GZFILE), source(f){};
  ByteStream(uint8_t *address, size_t size)
      : type(MEMORY), source(address, size){};

  void read(size_t num, uint8_t *buffer, bool *error) {
    switch (type) {
    case GZFILE: {
      if (size_t(gzread(source.file, buffer, num)) < num) {
        logger::error("Unexpected EOF\n");
        *error = true;
        memset(buffer, 0, num);
      }

      break;
    }

    case MEMORY: {
      if (source.array.second < num) {
        logger::error("Not enough data in memory buffer\n");
        *error = true;
        memset(buffer, 0, num);
      }

      memcpy(buffer, source.array.first, num);
      source.array.first += num;
      source.array.second -= num;

      break;
    }
    }

    *error = (*error || false);
  }
};

namespace nbt {

static bool format_check(ByteStream &b) {
  // Check the byte stream begins with a non-end tag and contains a valid UTF-8
  // name
  uint8_t buffer[MAXELEMENTSIZE];
  uint16_t name_length = 0;
  bool error = false;

  b.read(1, buffer, &error);
  if (error || !buffer[0] || buffer[0] > 13)
    return false;

  b.read(2, buffer, &error);
  if (error)
    return false;

  b.read((name_length = _NTOHS(buffer)), buffer, &error);
  if (error)
    return false;

  for (uint16_t i = 0; i < name_length; i++) {
    if (i < 0x21 || i > 0x7e)
      return false;
  }

  return true;
}

static bool matryoshka(ByteStream &b, NBT &destination) {
  bool error = false;

  uint8_t buffer[MAXELEMENTSIZE];

  NBT current;
  tag_type current_type = tag_type::tag_end, list_type;

  std::string current_name;
  uint32_t elements, list_elements, name_size;

  // The stack with open containers. When a container is found, it is pushed
  // on this stack; when an element is finished, it is pushed in the container
  // on top of the stack, or returned if the stack is empty.
  std::stack<NBT> opened_elements = {};

  // The context stack. All was well until the NBT lists came along that
  // changed the element format (no type/name). This stack tracks every
  // element on the stack. Its contents follow the format: <uint32_t
  // content_left, tag_type list_type> If content left > 0, it is assumed to
  // be a NBT list. This changes the behaviour of the parser accordingly. When
  // pushing/popping containers, the second element gets the type of the
  // current list.
  std::stack<std::pair<uint32_t, tag_type>> context = {};
<commit_msg>Functionnal format check<commit_after>#pragma once
#ifndef NBT_GZ_PARSE_HPP_
#define NBT_GZ_PARSE_HPP_

#include <filesystem>
#include <nbt/nbt.hpp>
#include <stack>
#include <zlib.h>

#define _NTOHS(ptr) (int16_t(((ptr)[0] << 8) + (ptr)[1]))
#define _NTOHI(ptr)                                                            \
  ((uint32_t((ptr)[0]) << 24) + (uint32_t((ptr)[1]) << 16) +                   \
   (uint32_t((ptr)[2]) << 8) + uint32_t((ptr)[3]))
#define _NTOHL(ptr)                                                            \
  ((uint64_t((ptr)[0]) << 56) + (uint64_t((ptr)[1]) << 48) +                   \
   (uint64_t((ptr)[2]) << 40) + (uint64_t((ptr)[3]) << 32) +                   \
   (uint64_t((ptr)[4]) << 24) + (uint64_t((ptr)[5]) << 16) +                   \
   (uint64_t((ptr)[6]) << 8) + uint64_t((ptr)[7]))

// Max size of a single element to read in memory (string)
#define MAXELEMENTSIZE 65025

// Check if the context indicates a being in a list
#define LIST (context.size() && context.top().second < tag_type::tag_long_array)

union FloatTranslator {
  // See this union as an uint8_t[8] array. Put an integer in the buffer, get
  // its byte reinterpretation in floating point, WITHOUT THE WARNINGS !
  uint64_t buffer;
  float _float;
  double _double;

  FloatTranslator(uint32_t float_bytes) : buffer(float_bytes){};
  FloatTranslator(uint64_t double_bytes) : buffer(double_bytes){};
};

struct ByteStream {
  // Adapter for the matryoshkas to work with both files and memory buffers
  enum BufferType { MEMORY, GZFILE };

  union Source {
    gzFile file;
    std::pair<uint8_t *, size_t> array;

    Source(gzFile f) : file(f){};
    Source(uint8_t *address, size_t size) : array({address, size}){};
  };

  BufferType type;
  Source source;

  ByteStream(gzFile f) : type(GZFILE), source(f){};
  ByteStream(uint8_t *address, size_t size)
      : type(MEMORY), source(address, size){};

  void read(size_t num, uint8_t *buffer, bool *error) {
    switch (type) {
    case GZFILE: {
      if (size_t(gzread(source.file, buffer, num)) < num) {
        logger::error("Unexpected EOF\n");
        *error = true;
        memset(buffer, 0, num);
      }

      break;
    }

    case MEMORY: {
      if (source.array.second < num) {
        logger::error("Not enough data in memory buffer\n");
        *error = true;
        memset(buffer, 0, num);
      }

      memcpy(buffer, source.array.first, num);
      source.array.first += num;
      source.array.second -= num;

      break;
    }
    }
  }
};

namespace nbt {

static bool format_check(ByteStream &b) {
  // Check the byte stream begins with a non-end tag and contains a valid UTF-8
  // name
  uint8_t buffer[MAXELEMENTSIZE];
  uint16_t name_length = 0;
  bool error = false;

  b.read(1, buffer, &error);
  if (error || !buffer[0] || buffer[0] > 13) {
    logger::deep_debug("NBT format check error: Invalid type read\n");
    return false;
  }

  b.read(2, buffer, &error);
  if (error) {
    logger::deep_debug("NBT format check error: Invalid name size read\n");
    return false;
  }

  b.read((name_length = _NTOHS(buffer)), buffer, &error);
  if (error) {
    logger::deep_debug("NBT format check error: Invalid name read\n");
    return false;
  }

  for (uint16_t i = 0; i < name_length; i++) {
    if (buffer[i] < 0x21 || buffer[i] > 0x7e) {
      logger::deep_debug(
          "NBT format check error: Invalid character read: {:02x}\n",
          buffer[i]);
      return false;
    }
  }

  return true;
}

static bool matryoshka(ByteStream &b, NBT &destination) {
  bool error = false;

  uint8_t buffer[MAXELEMENTSIZE];

  NBT current;
  tag_type current_type = tag_type::tag_end, list_type;

  std::string current_name;
  uint32_t elements, list_elements, name_size;

  // The stack with open containers. When a container is found, it is pushed
  // on this stack; when an element is finished, it is pushed in the container
  // on top of the stack, or returned if the stack is empty.
  std::stack<NBT> opened_elements = {};

  // The context stack. All was well until the NBT lists came along that
  // changed the element format (no type/name). This stack tracks every
  // element on the stack. Its contents follow the format: <uint32_t
  // content_left, tag_type list_type> If content left > 0, it is assumed to
  // be a NBT list. This changes the behaviour of the parser accordingly. When
  // pushing/popping containers, the second element gets the type of the
  // current list.
  std::stack<std::pair<uint32_t, tag_type>> context = {};

{'fraction_non_alphanumeric': 0.13477366255144033, 'fraction_numerical': 0.02423411065386374, 'mean_word_length': 2.808881149325207, 'pattern_counts': {'":': 0, '<': 75, '<?xml version=': 0, '>': 27, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3798694', 'n_tokens_mistral': 3087, 'n_tokens_neox': 2911, 'n_words': 865}