<commit_before>#include "assembly_reader.h"

#include <cstring>

#include "../common/ptr_util.h"
#include "opcodes.h"
#include "../pe/coded_index.h"
#include "../pe/pe_image_reader.h"
#include "assembly.h"

namespace stereo {
    namespace assemblies {

        AssemblyReader::AssemblyReader(const char* file_path)
            :
            logger_(std::make_unique<logging::ConsoleLogger>()),
            image_(pe::PEImageReader::load_image(file_path))
        {
        }

        u32 AssemblyReader::get_entry_point()
        {
            auto ept = pe::MetadataToken(image_->cli_header.entry_point_token);
            return ept.rid();
        }

        const ModuleDef* AssemblyReader::read_module_def()
        {
            if (module_ != nullptr)
                return module_.get();

            auto module = std::make_unique<ModuleDef>();

            auto table_row_ptr = get_table_row_ptr(pe::MetadataTable::Module, 1);

            // Skip the generation column since it's always zero
            table_row_ptr += 2;

            module->name = read_string(&table_row_ptr);

            // TODO: Read the mvid guid

            module_ = std::move(module);
            return module_.get();
        }

        const MethodDef* AssemblyReader::read_method_def(u32 rid)
        {
            if (already_read(method_defs_, rid))
                return method_defs_[get_index_from_rid(rid)].get();

            resize_if_needed(method_defs_, pe::MetadataTable::Method);

            auto table_row_ptr = get_table_row_ptr(pe::MetadataTable::Method, rid);
            auto method_def = std::make_unique<MethodDef>();

            // RVA (a 4-byte constant) 
            method_def->rva = ptrutil::read32(&table_row_ptr);

            // ImplFlags (a 2-byte bitmask of type MethodImplAttributes, II.23.1.10) 
            method_def->impl_attributes = static_cast<MethodImplAttributes>(ptrutil::read16(&table_row_ptr));

            // Flags (a 2-byte bitmask of type MethodAttributes, II.23.1.10) 
            method_def->attributes = static_cast<MethodAttributes>(ptrutil::read16(&table_row_ptr));

            // Name (an index into the String heap) 
            method_def->name = read_string(&table_row_ptr);

            // TODO: Read signature, parameters etc

            read_method_body(method_def.get());

            method_defs_[get_index_from_rid(rid)] = std::move(method_def);
            return method_defs_[get_index_from_rid(rid)].get();
        }

        const MemberRef* AssemblyReader::read_member_ref(u32 rid)
        {
            if (already_read(member_refs_, rid))
                return member_refs_[get_index_from_rid(rid)].get();

            resize_if_needed(member_refs_, pe::MetadataTable::MemberRef);

            auto table_row_ptr = get_table_row_ptr(pe::MetadataTable::MemberRef, rid);
            auto member_ref = std::make_unique<MemberRef>();

            // Class (an index into the MethodDef, ModuleRef, TypeDef, TypeRef, or TypeSpec tables; more precisely, a MemberRefParent(II.24.2.6) coded index)
            auto token = read_metadata_token(&table_row_ptr, pe::CodedIndexType::MemberRefParent);

            // Name (an index into the String heap) 
            member_ref->name = read_string(&table_row_ptr);

            // Signature (an index into the Blob heap) 
            auto signature = read_blob(&table_row_ptr);

            // TODO: Implement all token type cases
            if (token.type() == pe::MetadataTokenType::TypeRef)
            {
                member_ref->type_ref = read_type_ref(token.rid());
            }
            else
            {
                throw "read_member_ref -> unsupported token type";
            }

            member_refs_[get_index_from_rid(rid)] = std::move(member_ref);
            return member_refs_[get_index_from_rid(rid)].get();
        }

        const TypeRef* AssemblyReader::read_type_ref(u32 rid)
        {
            if (already_read(type_refs_, rid))
                return type_refs_[get_index_from_rid(rid)].get();

            resize_if_needed(type_refs_, pe::MetadataTable::TypeRef);

            auto table_row_ptr = get_table_row_ptr(pe::MetadataTable::TypeRef, rid);

            auto type_ref = std::make_unique<TypeRef>();

            // ResolutionScope (an index into a Module, ModuleRef, AssemblyRef or TypeRef table, or null; more precisely, a ResolutionScope(II.24.2.6) coded index)
            auto res_scope = read_metadata_token(&table_row_ptr, pe::CodedIndexType::ResolutionScope);

            if (res_scope.type() == pe::MetadataTokenType::AssemblyRef)
            {
                type_ref->resolution_scope = read_assembly_ref(res_scope.rid());
            }
            else
            {
                throw "read_type_ref -> unsupported ResolutionScope type";
            }

            // TypeName(an index into the String heap)
            type_ref->name = read_string(&table_row_ptr);

            // TypeNamespace(an index into the String heap)
            type_ref->name_space = read_string(&table_row_ptr);

            type_refs_[get_index_from_rid(rid)] = std::make_unique<TypeRef>();
            return type_refs_[get_index_from_rid(rid)].get();
        }

        const AssemblyRef* AssemblyReader::read_assembly_ref(u32 rid)
        {
            if (already_read(assembly_refs_, rid))
            {
                return assembly_refs_[get_index_from_rid(rid)].get();
            }

            resize_if_needed(assembly_refs_, pe::MetadataTable::AssemblyRef);

            auto table_row_ptr = get_table_row_ptr(pe::MetadataTable::AssemblyRef, rid);

            auto asm_ref = std::make_unique<AssemblyRef>();

            // MajorVersion, MinorVersion, BuildNumber, RevisionNumber (each being 2-byte constants)
            asm_ref->major_version = ptrutil::read16(&table_row_ptr);
            asm_ref->minor_version = ptrutil::read16(&table_row_ptr);<commit_msg>move the read type ref<commit_after>#include "assembly_reader.h"

#include <cstring>

#include "../common/ptr_util.h"
#include "opcodes.h"
#include "../pe/coded_index.h"
#include "../pe/pe_image_reader.h"
#include "assembly.h"

namespace stereo {
    namespace assemblies {

        AssemblyReader::AssemblyReader(const char* file_path)
            :
            logger_(std::make_unique<logging::ConsoleLogger>()),
            image_(pe::PEImageReader::load_image(file_path))
        {
        }

        const pe::PEImage * AssemblyReader::get_image()
        {
            return image_.get();
        }

        u32 AssemblyReader::get_entry_point()
        {
            auto ept = pe::MetadataToken(image_->cli_header.entry_point_token);
            return ept.rid();
        }

        const ModuleDef* AssemblyReader::read_module_def()
        {
            if (module_ != nullptr)
                return module_.get();

            auto module = std::make_unique<ModuleDef>();

            auto table_row_ptr = get_table_row_ptr(pe::MetadataTable::Module, 1);

            // Skip the generation column since it's always zero
            table_row_ptr += 2;

            module->name = read_string(&table_row_ptr);

            // TODO: Read the mvid guid

            module_ = std::move(module);
            return module_.get();
        }

        const MethodDef* AssemblyReader::read_method_def(u32 rid)
        {
            if (already_read(method_defs_, rid))
                return method_defs_[get_index_from_rid(rid)].get();

            resize_if_needed(method_defs_, pe::MetadataTable::Method);

            auto table_row_ptr = get_table_row_ptr(pe::MetadataTable::Method, rid);
            auto method_def = std::make_unique<MethodDef>();

            // RVA (a 4-byte constant) 
            method_def->rva = ptrutil::read32(&table_row_ptr);

            // ImplFlags (a 2-byte bitmask of type MethodImplAttributes, II.23.1.10) 
            method_def->impl_attributes = static_cast<MethodImplAttributes>(ptrutil::read16(&table_row_ptr));

            // Flags (a 2-byte bitmask of type MethodAttributes, II.23.1.10) 
            method_def->attributes = static_cast<MethodAttributes>(ptrutil::read16(&table_row_ptr));

            // Name (an index into the String heap) 
            method_def->name = read_string(&table_row_ptr);

            // TODO: Read signature, parameters etc

            read_method_body(method_def.get());

            method_defs_[get_index_from_rid(rid)] = std::move(method_def);
            return method_defs_[get_index_from_rid(rid)].get();
        }

        const MemberRef* AssemblyReader::read_member_ref(u32 rid)
        {
            if (already_read(member_refs_, rid))
                return member_refs_[get_index_from_rid(rid)].get();

            resize_if_needed(member_refs_, pe::MetadataTable::MemberRef);

            auto table_row_ptr = get_table_row_ptr(pe::MetadataTable::MemberRef, rid);
            auto member_ref = std::make_unique<MemberRef>();

            // Class (an index into the MethodDef, ModuleRef, TypeDef, TypeRef, or TypeSpec tables; more precisely, a MemberRefParent(II.24.2.6) coded index)
            auto token = read_metadata_token(&table_row_ptr, pe::CodedIndexType::MemberRefParent);

            // Name (an index into the String heap) 
            member_ref->name = read_string(&table_row_ptr);

            // Signature (an index into the Blob heap) 
            auto signature = read_blob(&table_row_ptr);

            // TODO: Implement all token type cases
            if (token.type() == pe::MetadataTokenType::TypeRef)
            {
                member_ref->type_ref = read_type_ref(token.rid());
            }
            else
            {
                throw "read_member_ref -> unsupported token type";
            }

            member_refs_[get_index_from_rid(rid)] = std::move(member_ref);
            return member_refs_[get_index_from_rid(rid)].get();
        }

        const TypeRef* AssemblyReader::read_type_ref(u32 rid)
        {
            if (already_read(type_refs_, rid))
                return type_refs_[get_index_from_rid(rid)].get();

            resize_if_needed(type_refs_, pe::MetadataTable::TypeRef);

            auto table_row_ptr = get_table_row_ptr(pe::MetadataTable::TypeRef, rid);

            auto type_ref = std::make_unique<TypeRef>();

            // ResolutionScope (an index into a Module, ModuleRef, AssemblyRef or TypeRef table, or null; more precisely, a ResolutionScope(II.24.2.6) coded index)
            auto res_scope = read_metadata_token(&table_row_ptr, pe::CodedIndexType::ResolutionScope);

            if (res_scope.type() == pe::MetadataTokenType::AssemblyRef)
            {
                type_ref->resolution_scope = read_assembly_ref(res_scope.rid());
            }
            else
            {
                throw "read_type_ref -> unsupported ResolutionScope type";
            }

            // TypeName(an index into the String heap)
            type_ref->name = read_string(&table_row_ptr);

            // TypeNamespace(an index into the String heap)
            type_ref->name_space = read_string(&table_row_ptr);

            type_refs_[get_index_from_rid(rid)] = std::move(type_ref);
            return type_refs_[get_index_from_rid(rid)].get();
        }

        const AssemblyRef* AssemblyReader::read_assembly_ref(u32 rid)
        {
            if (already_read(assembly_refs_, rid))
            {
                return assembly_refs_[get_index_from_rid(rid)].get();
            }

            resize_if_needed(assembly_refs_, pe::MetadataTable::AssemblyRef);

            auto table_row_ptr = get_table_row_ptr(pe::MetadataTable::AssemblyRef, rid);

            auto asm_ref = std::make_unique<AssemblyRef>();

            // MajorVersion, MinorVersion, BuildNumber, RevisionNumber (each being 2-byte constants)
            asm_ref->major_version = ptrutil::read16(&table_row_ptr);
            asm_ref->minor_version = ptrutil::read16(&table_row_ptr);
{'fraction_non_alphanumeric': 0.10816002672680197, 'fraction_numerical': 0.0073498705420529525, 'mean_word_length': 2.368213783403657, 'pattern_counts': {'":': 0, '<': 22, '<?xml version=': 0, '>': 52, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5982421', 'n_tokens_mistral': 3579, 'n_tokens_neox': 3325, 'n_words': 630}