#include <Windows.h>
#include <intrin.h>
#include <cstddef>
#include <iostream>
#include <cstdint>
#include <cassert>
#include <cstdio>
#include <string>
#include <sstream>
#include "../common/xstddef.h"

using namespace std;
using namespace ns_OClasses;

struct TlbInfo
{
    static unsigned const LEVEL1    = 0,
                          LEVEL2    = 1,
                          PAGES_4KB = 0,
                          PAGES_4MB = 1;
    struct
    {
        struct
        {
            bool     fKnown;
            unsigned entries;
            int      associativity;
        } aPageSizes[2];
    } aLevels[2];
};

union CacheLineChain
{
    CacheLineChain volatile *pclcNext;
};

void AutoExit( bool fExit, char *message );
bool GetCacheInfo( unsigned level, size_t *pSize, unsigned *pAssociativity, unsigned *pLineSizeShift, unsigned *pSectors );
void GetTlbInfo( TlbInfo *pTlbInfo );

size_t InitializePointerChain( void *pBlock, unsigned blockBits, unsigned randBlockBits , unsigned lineBits );
string GetSizeString( unsigned blockBits );

int main()
{
    unsigned const  MAX_BLOCK_BITS = 28; // 256MB
    TlbInfo         ti;
    SYSTEM_INFO     si;
    unsigned        pageBits;
    unsigned        lineBits;
    unsigned        rbBits;
    void           *pBlock;

    AutoExit( !SetPriorityClass( GetCurrentProcess(), HIGH_PRIORITY_CLASS ), "can't set process-priority" );;
    AutoExit( !SetThreadPriority( GetCurrentThread(), THREAD_PRIORITY_HIGHEST ), "can't set thread-priority" );;
    AutoExit( SetThreadAffinityMask( GetCurrentThread(), 1 ) == 0, "can't set thread affinity");

    GetTlbInfo( &ti );
    AutoExit( !ti.aLevels[TlbInfo::LEVEL1].aPageSizes[TlbInfo::PAGES_4MB].fKnown, "can't determine l1 dtlb-size" );
    GetSystemInfo( &si );
    pageBits = floor_log2<DWORD>( si.dwPageSize );
    AutoExit( !GetCacheInfo( 1, nullptr, nullptr, &lineBits, nullptr ), "can't determine l1-cache parameters" );
    rbBits = floor_log2<unsigned>( ti.aLevels[TlbInfo::LEVEL1].aPageSizes[TlbInfo::PAGES_4MB].entries ) + pageBits;

    AutoExit( (pBlock = VirtualAlloc( NULL, (size_t)1 << MAX_BLOCK_BITS, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE )) == NULL, "can't allocate memory for test-block" );

    unsigned                 blockBits;
    size_t                   nCacheLines;
    DWORDLONG                dwlTick,
                             dwlFastestTick;
    unsigned                 turn;
    double                   ticksPerCacheline;

    for( blockBits = pageBits; blockBits <= MAX_BLOCK_BITS; blockBits++ )
    {
        nCacheLines = InitializePointerChain( pBlock, blockBits, rbBits, lineBits );

        for( dwlFastestTick = (DWORDLONG)(LONGLONG)-1, turn = 0; turn < 100; turn++ )
        {
            dwlTick        = __rdtsc();
            for( CacheLineChain volatile *pclc = (CacheLineChain *)pBlock; (pclc = pclc->pclcNext) != nullptr; );
            dwlTick        = __rdtsc() - dwlTick;
            dwlFastestTick = dwlTick < dwlFastestTick ? dwlTick : dwlFastestTick;
        }

        ticksPerCacheline = (double)dwlFastestTick / nCacheLines;
        cout << GetSizeString( blockBits ) << " " << ticksPerCacheline << endl;
    }

    return EXIT_SUCCESS;
}

uint64_t ReverseBits( uint64_t reverse, unsigned bits );

size_t InitializePointerChain( void *pBlock, unsigned blockBits, unsigned randBlockBits , unsigned lineBits )
{
    assert(blockBits > lineBits);

    union ClcPointer
    {
        BYTE                    *pbCl;
        CacheLineChain volatile *pclc;
    };

    CacheLineChain volatile *pclcFirst,
                            *pclcPrev;
    ClcPointer               cpRandBlock,
                             cp;
    size_t                   randBlock,
                             reverseCacheLine,
                             nRandBlocks;

    pclcFirst        = (CacheLineChain *)pBlock;
    cpRandBlock.pclc = pclcFirst;
    pclcPrev         = nullptr;
    nRandBlocks      = (size_t)1 << (blockBits >= randBlockBits ? blockBits - randBlockBits : 0);
    randBlock        = 0;
    do
    {
        reverseCacheLine = 0;
        do
        {
            cp.pbCl = &cpRandBlock.pbCl[(size_t)ReverseBits( reverseCacheLine, blockBits >= randBlockBits ? randBlockBits - lineBits : blockBits - lineBits ) << lineBits];
            if( pclcPrev )
                pclcPrev->pclcNext = cp.pclc;
            pclcPrev = cp.pclc;
        } while( ++reverseCacheLine < ((size_t)1 << (blockBits >= randBlockBits ? randBlockBits - lineBits : blockBits - lineBits)) );
        cpRandBlock.pbCl += (size_t)1 << randBlockBits;
    } while( ++randBlock < nRandBlocks );
    pclcPrev->pclcNext = nullptr;

    return (size_t)1 << (blockBits - lineBits);
}

string GetSizeString( unsigned blockBits )
{
    size_t size;
    stringstream ss;

    if( (size = (size_t)1 << blockBits) < 1024 )
        ss << size << "B";
    else if( size < (1024 * 1024) )
        ss << (size / 1024) << "kB";
    else if( size < (1024 * 1024 * 1024) )
        ss << (size / (1024 * 1024)) << "MB";
    else
        ss << (size / (1024 * 1024 * 1024)) << "GB";

    return ss.str();
}

void AutoExit( bool fExit, char *message )
{
    if( fExit )
    {
        cout << message << endl;
        getchar();
        ExitProcess( EXIT_FAILURE );
    }
}


uint64_t ReverseBits( uint64_t reverse, unsigned bits )
{
    assert(bits >= 1 && bits <= 64);
    reverse   = ((reverse & 0xFFFFFFFF00000000) >> 32) | ((reverse & 0x00000000FFFFFFFF) << 32);
    reverse   = ((reverse & 0xFFFF0000FFFF0000) >> 16) | ((reverse & 0x0000FFFF0000FFFF) << 16);
    reverse   = ((reverse & 0xFF00FF00FF00FF00) >>  8) | ((reverse & 0x00FF00FF00FF00FF) <<  8);
    reverse   = ((reverse & 0xF0F0F0F0F0F0F0F0) >>  4) | ((reverse & 0x0F0F0F0F0F0F0F0F) <<  4);
    reverse   = ((reverse & 0xCCCCCCCCCCCCCCCC) >>  2) | ((reverse & 0x3333333333333333) <<  2);
    reverse   = ((reverse & 0xAAAAAAAAAAAAAAAA) >>  1) | ((reverse & 0x5555555555555555) <<  1);
    reverse >>= 64 - bits;

    return reverse;
}

DWORD CpuId( DWORD dwCode, DWORD dwEcx2ndParameter, DWORD *pdwRegisters );

bool GetCacheInfo( unsigned level, size_t *pSize, unsigned *pAssociativity, unsigned *pLineSizeShift, unsigned *pSectors )
{
    DWORD adwCpuIdRetRegisters[4];

    if( (CpuId( 0, 0, adwCpuIdRetRegisters ),
         adwCpuIdRetRegisters[1] == 'uneG' &&
         adwCpuIdRetRegisters[3] == 'Ieni' &&
         adwCpuIdRetRegisters[2] == 'letn') )
        return false;

    if( adwCpuIdRetRegisters[1] == 'htuA' &&
        adwCpuIdRetRegisters[3] == 'itne' &&
        adwCpuIdRetRegisters[2] == 'DMAc' &&
        CpuId( 0x80000000u, 0, adwCpuIdRetRegisters ) >= (0x80000005u + (level > 1)) )
    {
        size_t   size;
        unsigned associativity;
        unsigned clShift;
        unsigned sectors;

        CpuId( 0x80000005u + (level > 1), 0, adwCpuIdRetRegisters );

        if( level == 1 )
            size          = (adwCpuIdRetRegisters[2] >> 24) * 1024,
            associativity = (unsigned)(signed char)(adwCpuIdRetRegisters[2] >> 16),
            clShift       = floor_log2<DWORD>( adwCpuIdRetRegisters[2] & 0x0FF ),
            sectors       = (unsigned)((adwCpuIdRetRegisters[2] >> 8) & 0x0FF);
        else if( level == 2 )
            size          = (adwCpuIdRetRegisters[2] >> 16) * 1024,
            associativity = (unsigned)(signed char)((adwCpuIdRetRegisters[2] >> 12) & 0x0F),
            clShift       = floor_log2<DWORD>( adwCpuIdRetRegisters[2] & 0x0FF ),
            sectors       = (unsigned)((adwCpuIdRetRegisters[2] >> 8) & 0x0F);
        else if( level == 3 )
            size          = ((adwCpuIdRetRegisters[3] >> 18) & 0x3FFF) * 512 * 1024,
            associativity = (unsigned)(signed char)((adwCpuIdRetRegisters[3] >> 12) & 0x0F),
            clShift       = floor_log2<DWORD>( adwCpuIdRetRegisters[3] & 0x0FF ),
            sectors       = (unsigned)((adwCpuIdRetRegisters[2] >> 8) & 0x0F);
        else
            return false;

        if( pSize )
            *pSize = size;

        if( pAssociativity )
            *pAssociativity = associativity;

        if( pLineSizeShift )
            *pLineSizeShift = clShift;

        if( pSectors )
            *pSectors = sectors;

        return true;
    }

    return false;
}

void GetTlbInfo( TlbInfo *pti )
{
    DWORD adwCpuidRetRegisters[4];
    DWORD dwMaxExtendedCpuId;

    pti->aLevels[TlbInfo::LEVEL1].aPageSizes[TlbInfo::PAGES_4KB].fKnown = false;
    pti->aLevels[TlbInfo::LEVEL1].aPageSizes[TlbInfo::PAGES_4MB].fKnown = false;
    pti->aLevels[TlbInfo::LEVEL2].aPageSizes[TlbInfo::PAGES_4KB].fKnown = false;
    pti->aLevels[TlbInfo::LEVEL2].aPageSizes[TlbInfo::PAGES_4MB].fKnown = false;

    if( (CpuId( 0, 0, adwCpuidRetRegisters ),
         adwCpuidRetRegisters[1] == 'uneG' &&
         adwCpuidRetRegisters[3] == 'Ieni' &&
         adwCpuidRetRegisters[2] == 'letn') )
        return;

    if( (CpuId( 0, 0, adwCpuidRetRegisters ),
         adwCpuidRetRegisters[1] == 'htuA' &&
         adwCpuidRetRegisters[3] == 'itne' &&
         adwCpuidRetRegisters[2] == 'DMAc') &&
        (dwMaxExtendedCpuId = CpuId( 0x80000000u, 0, adwCpuidRetRegisters )) >= 0x80000005u )
    {
        int associativity;

        CpuId( 0x80000005u, 0, adwCpuidRetRegisters );
        pti->aLevels[TlbInfo::LEVEL1].aPageSizes[TlbInfo::PAGES_4KB].fKnown        = true;
        pti->aLevels[TlbInfo::LEVEL1].aPageSizes[TlbInfo::PAGES_4KB].entries       =         (unsigned)((adwCpuidRetRegisters[1] & 0x00FF0000u) >> 16);
        pti->aLevels[TlbInfo::LEVEL1].aPageSizes[TlbInfo::PAGES_4KB].associativity = (int)(signed char)((adwCpuidRetRegisters[1] & 0xFF000000u) >> 24);
        pti->aLevels[TlbInfo::LEVEL1].aPageSizes[TlbInfo::PAGES_4MB].fKnown        = true;
        pti->aLevels[TlbInfo::LEVEL1].aPageSizes[TlbInfo::PAGES_4MB].entries       =         (unsigned)((adwCpuidRetRegisters[0] & 0x00FF0000u) >> 16);
        pti->aLevels[TlbInfo::LEVEL1].aPageSizes[TlbInfo::PAGES_4MB].associativity = (int)(signed char)((adwCpuidRetRegisters[0] & 0xFF000000u) >> 24);

        if( dwMaxExtendedCpuId >= 0x80000006u )
            CpuId( 0x80000006u, 0, adwCpuidRetRegisters ),
            pti->aLevels[TlbInfo::LEVEL2].aPageSizes[TlbInfo::PAGES_4KB].fKnown        = true,
            pti->aLevels[TlbInfo::LEVEL2].aPageSizes[TlbInfo::PAGES_4KB].entries       = (unsigned)((adwCpuidRetRegisters[1] & 0x0FFF0000u) >> 16),
            associativity                                                              = (int)((adwCpuidRetRegisters[1] & 0xF0000000u) >> 28),
            pti->aLevels[TlbInfo::LEVEL2].aPageSizes[TlbInfo::PAGES_4KB].associativity = (associativity < 0xF) ? associativity : -1,
            pti->aLevels[TlbInfo::LEVEL2].aPageSizes[TlbInfo::PAGES_4MB].fKnown        = true,
            pti->aLevels[TlbInfo::LEVEL2].aPageSizes[TlbInfo::PAGES_4MB].entries       = (unsigned)((adwCpuidRetRegisters[0] & 0x0FFF0000u) >> 16),
            associativity                                                              = (int)((adwCpuidRetRegisters[0] & 0xF0000000u) >> 28),
            pti->aLevels[TlbInfo::LEVEL2].aPageSizes[TlbInfo::PAGES_4MB].associativity = (associativity < 0xF) ? associativity : -1;

    }
}

DWORD CpuId( DWORD dwCode, DWORD dwEcx2ndParameter, DWORD *pdwRegisters )
{
    int aRegs[4];

    __cpuidex( aRegs, (int)dwCode, (int)dwEcx2ndParameter );
    for( int i = 0; i < 4; pdwRegisters[i] = (DWORD)aRegs[i], i++ );

    return (DWORD)aRegs[0];
}