<commit_before>#include <iostream>
#include <boost/filesystem.hpp>
#include <boost/preprocessor/repetition/enum.hpp>

#include "shannon.hpp"
#include "output.hpp"
#include "graph.hpp"

#define SHANNON_IDENT(z, n, value) value

template <uint8_t Base>
constexpr double log(double value) {
    return log2(value) / log2(Base);
}

namespace fs = boost::filesystem;

using counter_t = std::array<uint64_t, 256>;
using allowed_t = std::array<bool, 256>;
constexpr auto DEFAULT_BLOCK_SIZE = 16 * 1024;

constexpr allowed_t ALLOWED_DATA{BOOST_PP_ENUM(256, SHANNON_IDENT, true)};
constexpr allowed_t ALLOWED_TEXT{BOOST_PP_ENUM(128, SHANNON_IDENT, true)};
constexpr allowed_t ALLOWED_BASE64{
    BOOST_PP_ENUM(43, SHANNON_IDENT, false),
    true, // '+'
    BOOST_PP_ENUM(3, SHANNON_IDENT, false),
    BOOST_PP_ENUM(11, SHANNON_IDENT, true), // '/' and 0-9
    BOOST_PP_ENUM(7, SHANNON_IDENT, false),
    BOOST_PP_ENUM(26, SHANNON_IDENT, true), // A-Z
    BOOST_PP_ENUM(6, SHANNON_IDENT, false),
    BOOST_PP_ENUM(26, SHANNON_IDENT, true), // a-z
};

double shannon_score(const counter_t& counts, std::size_t total_size,
                     DataFormat format) {
    const allowed_t* allowed = nullptr;
    switch (format) {
    case DataFormat::DATA:
        allowed = &ALLOWED_DATA;
        break;
    case DataFormat::TEXT:
        allowed = &ALLOWED_TEXT;
        break;
    case DataFormat::BASE64:
        allowed = &ALLOWED_BASE64;
        break;
    }

    double score = 0;
    for (std::size_t index = 0; index < counts.size(); ++index) {
        auto count = counts[index];
        if (count == 0 || !(*allowed)[index])
            continue;
        double p_i = count / static_cast<double>(total_size);
        score += p_i * log2(p_i);
    }

    switch (format) {
    case DataFormat::DATA:
        return std::abs(score) / 8.0;
    case DataFormat::TEXT:
        return std::abs(score) / 7.0;
    case DataFormat::BASE64:
        return std::abs(score) / 6.0;
    }
    return 0.0;
}

template <typename Iter>
void shannon_digest(Iter begin, Iter end, counter_t& counts) {
    for (auto iter = begin; iter != end; ++iter) {
        counts[*iter] += 1;
    }
}

class shannon_iterator
    : public boost::iterator_facade<shannon_iterator, double,
                                    boost::forward_traversal_tag, double> {
    std::istream* m_stream = nullptr;
    DataFormat m_format;
    uint64_t m_block_size;
    std::vector<uint8_t> m_block;
    bool m_clear_stats;
    counter_t m_counter;

public:
    shannon_iterator() = default;
    shannon_iterator(std::ifstream& stream, uint64_t block_size,
                     DataFormat format, bool clear_stats = true)
        : m_stream{&stream},
          m_format{format},
          m_block_size{block_size},
          m_block(block_size),
          m_clear_stats{clear_stats},
          m_counter{} {
        this->increment(); // Get meaningful data in the buffer
    }
    void increment() {
        assert(m_block.capacity() >= m_block_size);<commit_msg>Fix bug where shannon_iterator required ifstream instead of istream<commit_after>#include <iostream>
#include <boost/filesystem.hpp>
#include <boost/preprocessor/repetition/enum.hpp>

#include "shannon.hpp"
#include "output.hpp"
#include "graph.hpp"

#define SHANNON_IDENT(z, n, value) value

template <uint8_t Base>
constexpr double log(double value) {
    return log2(value) / log2(Base);
}

namespace fs = boost::filesystem;

using counter_t = std::array<uint64_t, 256>;
using allowed_t = std::array<bool, 256>;
constexpr auto DEFAULT_BLOCK_SIZE = 16 * 1024;

constexpr allowed_t ALLOWED_DATA{BOOST_PP_ENUM(256, SHANNON_IDENT, true)};
constexpr allowed_t ALLOWED_TEXT{BOOST_PP_ENUM(128, SHANNON_IDENT, true)};
constexpr allowed_t ALLOWED_BASE64{
    BOOST_PP_ENUM(43, SHANNON_IDENT, false),
    true, // '+'
    BOOST_PP_ENUM(3, SHANNON_IDENT, false),
    BOOST_PP_ENUM(11, SHANNON_IDENT, true), // '/' and 0-9
    BOOST_PP_ENUM(7, SHANNON_IDENT, false),
    BOOST_PP_ENUM(26, SHANNON_IDENT, true), // A-Z
    BOOST_PP_ENUM(6, SHANNON_IDENT, false),
    BOOST_PP_ENUM(26, SHANNON_IDENT, true), // a-z
};

double shannon_score(const counter_t& counts, std::size_t total_size,
                     DataFormat format) {
    const allowed_t* allowed = nullptr;
    switch (format) {
    case DataFormat::DATA:
        allowed = &ALLOWED_DATA;
        break;
    case DataFormat::TEXT:
        allowed = &ALLOWED_TEXT;
        break;
    case DataFormat::BASE64:
        allowed = &ALLOWED_BASE64;
        break;
    }

    double score = 0;
    for (std::size_t index = 0; index < counts.size(); ++index) {
        auto count = counts[index];
        if (count == 0 || !(*allowed)[index])
            continue;
        double p_i = count / static_cast<double>(total_size);
        score += p_i * log2(p_i);
    }

    switch (format) {
    case DataFormat::DATA:
        return std::abs(score) / 8.0;
    case DataFormat::TEXT:
        return std::abs(score) / 7.0;
    case DataFormat::BASE64:
        return std::abs(score) / 6.0;
    }
    return 0.0;
}

template <typename Iter>
void shannon_digest(Iter begin, Iter end, counter_t& counts) {
    for (auto iter = begin; iter != end; ++iter) {
        counts[*iter] += 1;
    }
}

class shannon_iterator
    : public boost::iterator_facade<shannon_iterator, double,
                                    boost::forward_traversal_tag, double> {
    std::istream* m_stream = nullptr;
    DataFormat m_format;
    uint64_t m_block_size;
    std::vector<uint8_t> m_block;
    bool m_clear_stats;
    counter_t m_counter;

public:
    shannon_iterator() = default;
    shannon_iterator(std::istream& stream, uint64_t block_size,
                     DataFormat format, bool clear_stats = true)
        : m_stream{&stream},
          m_format{format},
          m_block_size{block_size},
          m_block(block_size),
          m_clear_stats{clear_stats},
          m_counter{} {
        this->increment(); // Get meaningful data in the buffer
    }
    void increment() {
        assert(m_block.capacity() >= m_block_size);
{'fraction_non_alphanumeric': 0.11898110106820049, 'fraction_numerical': 0.02037797863599014, 'mean_word_length': 3.027796161482462, 'pattern_counts': {'":': 0, '<': 25, '<?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': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2056129', 'n_tokens_mistral': 2148, 'n_tokens_neox': 2059, 'n_words': 453}