kernel_sim.hpp

#pragma once

#include <vector>
#include <unordered_map>
#include <cstdint>
#include <chrono>
#include <string>
#include <stdexcept>
#include <memory>
#include <thread>
#include <mutex>
#include <condition_variable>
#include <cstdlib>
#include <cstdio>

//
// Page protections.
//
const uint32_t PAGE_READONLY  = 0x01;
const uint32_t PAGE_READWRITE = 0x02;
const uint32_t PAGE_GUARD     = 0x100;

//
// ReactOS-inspired MMPTE (from internal/mmtypes.h, amd64)
//
union MMPTE {
    uint64_t u_long;
    struct {  // Hardware PTE
        uint64_t valid : 1;
        uint64_t write : 1;
        uint64_t owner : 1;  // 1=user
        uint64_t write_through : 1;
        uint64_t cache_disable : 1;
        uint64_t accessed : 1;
        uint64_t dirty : 1;
        uint64_t large_page : 1;
        uint64_t global : 1;
        uint64_t copy_on_write : 1;
        uint64_t prototype : 1;
        uint64_t transition : 1;
        uint64_t page_frame_number : 36;
        uint64_t reserved : 12;
        uint64_t nx : 1;
    } hard;
    struct {  // Software PTE (invalid)
        uint64_t valid : 1;  // 0
        uint64_t page_file : 1;
        uint64_t protection : 5;
        uint64_t prototype : 1;
        uint64_t transition : 1;
        uint64_t used_page_table_entries : 10;
        uint64_t reserved : 10;
        uint64_t page_file_low : 4;
        uint64_t page_file_high : 32;
    } soft;
    MMPTE();
};

//
// Simple VAD node (inspired by ReactOS MMADDRESS_NODE)
//
struct VadNode {
    uint64_t start_va;
    uint64_t end_va;
    uint32_t protection;  // e.g., PAGE_READONLY
    std::shared_ptr<VadNode> left, right;
};

//
// 4-level Page Table (PML4 root)
//
class PageTable {
private:
    size_t entries_per_level;
    std::vector<MMPTE> pml4;  // Level 4
    std::vector<std::vector<MMPTE>> pdpts;  // Level 3
    std::vector<std::vector<std::vector<MMPTE>>> pds;  // Level 2
    std::vector<std::vector<std::vector<std::vector<MMPTE>>>> pts;  // Level 1
    std::shared_ptr<VadNode> vad_root;  // VAD tree

    // Shared "physical" view of PT pages for flaw simulation (vector of all PT memory)
    std::shared_ptr<std::vector<uint8_t>> pt_physical_mem;
    mutable std::mutex pt_mutex;

    // Helper for VAD insert
    std::shared_ptr<VadNode> insert_vad_helper(std::shared_ptr<VadNode> node, uint64_t start, uint64_t end, uint32_t prot);

    // Helper for VAD query
    uint32_t query_vad_prot_helper(const std::shared_ptr<VadNode>& node, uint64_t va) const;

public:
    PageTable();

    // Verify integrity (OS-level check)
    void verify_integrity();

    // Set PTE via VA (allocates tables if needed)
    void set_pte(uint64_t va, const MMPTE& pte);

    // Get PTE via VA
    MMPTE get_pte(uint64_t va);

    // Flaw interface: Get shared view of PT "physical" mem for user modification
    std::shared_ptr<std::vector<uint8_t>> get_flawed_pt_view();

    // VAD insert (simple BST)
    void insert_vad(uint64_t start, uint64_t end, uint32_t prot);

    // Query VAD protection
    uint32_t query_vad_prot(uint64_t va);
};

//
// Multi-level TLB sim (L1 4-way assoc, L2 unified)
//
class TLB {
private:
    static constexpr size_t L1_SETS = 64;
    static constexpr size_t L1_WAYS = 4;
    static constexpr size_t L2_ENTRIES = 1024;
    struct Entry { uint64_t vpn; uint64_t pfn; bool write, owner; };
    std::vector<std::vector<Entry>> l1_itlb, l1_dtlb;  // Instruction/Data L1
    std::unordered_map<uint64_t, Entry> l2_tlb;
    size_t l1_hits;
    size_t l1_accesses;
    size_t l2_hits;
    size_t l2_accesses;

public:
    TLB();

    // Lookup entry in TLB
    bool lookup(uint64_t vpn, uint64_t& pfn, bool& write, bool& owner, bool is_data);

    // Insert entry into TLB
    void insert(uint64_t vpn, uint64_t pfn, bool write, bool owner, bool is_data);

    // Invalidate TLB entry
    void invalidate(uint64_t vpn);

    // Get L1 hit rate
    double l1_hit_rate() const;

    // Get L2 hit rate
    double l2_hit_rate() const;
};

//
// MMU with page walk and fault handling (inspired by ReactOS amd64/page.c and ARM3/page.c)
//
class MMU {
private:
    TLB tlb;
    std::mutex mmu_mutex;

    // Page fault dispatcher (like MiDispatchFault)
    void dispatch_fault(const std::shared_ptr<PageTable>& pt, uint64_t va, bool is_write, bool is_user, bool is_data);

    // Full page walk (from amd64/page.c MiAddressSpaceWalk-inspired)
    uint64_t page_walk(const std::shared_ptr<PageTable>& pt, uint64_t va, bool is_write, bool is_user, bool is_data);

public:
    // Translate VA to PA
    uint64_t translate(const std::shared_ptr<PageTable>& pt, uint64_t va, bool is_write, bool is_user, bool is_data = true);

    // Invalidate TLB for VA
    void invalidate_tlb(uint64_t va);

    // Get L1 hit rate
    double get_l1_hit_rate() const;

    // Get L2 hit rate
    double get_l2_hit_rate() const;
};

//
// Physical Memory
//
class PhysicalMemory {
private:
    size_t size;  // Swapped order: size before mem
    std::vector<uint8_t> mem;
    std::mutex mem_mutex;

public:
    PhysicalMemory(size_t frames, size_t frame_size = 4096);

    // Write to physical address
    void write(uint64_t pa, const std::string& data);

    // Read from physical address
    std::string read(uint64_t pa, size_t len);
};

//
// Kernel thread function (manages sim)
//
void kernel_thread_func(std::shared_ptr<PageTable> pt, std::shared_ptr<PhysicalMemory> phys_mem, std::shared_ptr<MMU> mmu, std::condition_variable& cv, std::mutex& cv_mutex, bool& ready);